From 1f4c9209c74fd06e59c063fbcd7a3bedfda6c4de Mon Sep 17 00:00:00 2001 From: popovaan Date: Sat, 1 Aug 2026 23:11:09 +0200 Subject: [PATCH 1/2] [OpenVINO] Support Falcon-H1-0.5B-Instruct with task text-generation --- docs/source/openvino/models.mdx | 1 + .../exporters/openvino/input_generators.py | 71 +++ optimum/exporters/openvino/model_configs.py | 51 ++ optimum/exporters/openvino/model_patcher.py | 436 ++++++++++++++++++ optimum/exporters/openvino/utils.py | 1 + optimum/intel/openvino/modeling_decoder.py | 23 + tests/openvino/test_decoder.py | 18 +- tests/openvino/test_export.py | 1 + tests/openvino/test_exporters_cli.py | 2 + tests/openvino/utils_tests.py | 103 +++++ 10 files changed, 705 insertions(+), 2 deletions(-) diff --git a/docs/source/openvino/models.mdx b/docs/source/openvino/models.mdx index 902dd2ad75..7ee1bf0753 100644 --- a/docs/source/openvino/models.mdx +++ b/docs/source/openvino/models.mdx @@ -60,6 +60,7 @@ Here is the list of the supported architectures : - EXAONE 4 - Falcon - Falcon-Mamba +- Falcon-H1 - FlauBERT - GLM-4 - GLM-Edge diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 34f75923ba..47e763ee96 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -1699,6 +1699,77 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int return past_key_values +class FalconH1DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for FalconH1 architectures. + + FalconH1 is a fully-parallel hybrid architecture: every decoder layer contains + both a Mamba2 mixer (conv + ssm states) and a self-attention block (key + value + cache). Therefore the number of mamba layers and the number of attention layers + are both equal to ``num_hidden_layers``. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + + config = normalized_config.config + self.intermediate_size = ( + config.mamba_d_ssm + if getattr(config, "mamba_d_ssm", None) is not None + else int(config.mamba_expand * config.hidden_size) + ) + self.conv_kernel_size = config.mamba_d_conv + self.mamba_d_state = config.mamba_d_state + self.n_mamba_heads = config.mamba_n_heads + self.mamba_ngroups = config.mamba_n_groups + self.mamba_headdim = config.mamba_d_head + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_heads = config.num_key_value_heads + # every layer is hybrid: both mamba and attention states exist per layer + self.num_mamba_layers = config.num_hidden_layers + self.num_attention_layers = config.num_hidden_layers + # past attention cache starts empty (sequence_length == 0), states are appended during decode + self.sequence_length = 0 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + for i in range(self.num_mamba_layers): + conv_state_shape = ( + self.batch_size, + self.intermediate_size + 2 * self.mamba_ngroups * self.mamba_d_state, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(conv_state) + ssm_state_shape = (self.batch_size, self.n_mamba_heads, self.mamba_headdim, self.mamba_d_state) + ssm_state = self.random_float_tensor(ssm_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(ssm_state) + + for i in range(self.num_attention_layers): + kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + past_key_values.append(k) + past_key_values.append(v) + + return past_key_values + + class Lfm2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): """ Generates dummy past_key_values inputs for Lfm2 architectures. diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 1bb72ff0e4..da6d6c4fe0 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -83,6 +83,7 @@ LTX2TransformerDummyInputGenerator, LTX2VaeDummyInputGenerator, LTX2VocoderDummyInputGenerator, + FalconH1DummyPastKeyValuesGenerator, LTXTransformerDummyInputGenerator, LTXVaeDummyInputGenerator, MambaCacheDummyInputGenerator, @@ -108,6 +109,7 @@ DBRXModelPatcher, DeciLMModelPatcher, DeepseekPatcher, + FalconH1ModelPatcher, FalconModelPatcher, FluxTransformerModelPatcher, FunASRModelPatcher, @@ -5712,6 +5714,55 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return common_inputs +@register_in_tasks_manager( + "falcon_h1", *["text-generation", "text-generation-with-past"], library_name="transformers" +) +class FalconH1OpenVINOConfig(MambaOpenVINOConfig): + # FalconH1 is a fully-parallel hybrid: every decoder layer contains both a Mamba2 mixer + # (conv + ssm states) and a self-attention block (key + value cache). + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, FalconH1DummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = FalconH1DummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = "4.53.0" + _MODEL_PATCHER = FalconH1ModelPatcher + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + _warn_potential_accuracy_issue_ov_2026_1("falcon_h1") + + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + cache_name_prefix = "cache_params.past" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + cache_name_prefix = "cache_params.present" + + num_hidden_layers = self._normalized_config.num_layers + for i in range(num_hidden_layers): + # [batch_size, conv_dim, conv_kernel_size] + inputs_or_outputs[f"{cache_name_prefix}.conv.{i}"] = {0: "batch_size"} + # [batch_size, mamba_n_heads, mamba_d_head, mamba_d_state] + inputs_or_outputs[f"{cache_name_prefix}.ssm.{i}"] = {0: "batch_size"} + + for i in range(num_hidden_layers): + inputs_or_outputs[f"{cache_name_prefix}.key.{i}"] = {0: "batch_size", 2: decoder_sequence_name} + inputs_or_outputs[f"{cache_name_prefix}.value.{i}"] = {0: "batch_size", 2: decoder_sequence_name} + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + return common_inputs + + @register_in_tasks_manager("audio-spectrogram-transformer", *["feature-extraction", "audio-classification"]) class ASTOpenVINOConfig(OpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index a93cf65a7d..44982f02f4 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -7074,6 +7074,442 @@ def segment_sum(input_tensor): return contextualized_states +# Torch-traceable reimplementation of FalconH1Mixer.torch_forward. +# The original implementation branches on `use_precomputed_states` (which depends on +# cache_position / seq_len at runtime) to choose between prefill and decoding code paths. +# Data-dependent Python control flow like this is not correctly captured by TorchScript +# tracing, so this patched version always executes both branches and selects the correct +# result with an `is_decoding` scalar mask, mirroring the zamba2_mamba_mixer approach. +# FalconH1-specific differences from a standard Mamba2 mixer that are preserved here: +# * an `ssm_in_multiplier` applied to the input and a `mup_vector` applied to the +# projected states, +# * projection split order [gate, hidden_states_B_C, dt], +# * an optional gated RMS norm (`mamba_rms_norm`); when disabled the gate is applied as +# `y * silu(gate)`. +def falcon_h1_mamba_mixer( + self, + hidden_states, + cache_params=None, + cache_position: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, +): + def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int): + pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0) + return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0) + + def reshape_into_chunks(input_tensor, pad_size, chunk_size): + input_tensor = pad_tensor_by_size(input_tensor, pad_size) + if len(input_tensor.shape) == 3: + return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2]) + else: + return input_tensor.reshape( + input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3] + ) + + def segment_sum(input_tensor): + chunk_size = input_tensor.size(-1) + input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size) + mask = torch.tril( + torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1 + ) + input_tensor = input_tensor.masked_fill(~mask, 0) + tensor_segsum = torch.cumsum(input_tensor, dim=-2) + mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0) + tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf) + return tensor_segsum + + input_states = hidden_states + layer_idx = self.layer_idx + + # Clamp `dt` using the model's time step limits. FalconH1 uses + # `time_step_limit = (0.0, float("inf"))` by default; serializing an OpenVINO `Clamp` + # op with `max=inf` produces an IR that the IR frontend cannot deserialize + # (the "inf" attribute fails to parse). When the upper bound is infinite we therefore + # clamp only with the lower bound, which is numerically equivalent. + def _clamp_dt(t): + low, high = self.time_step_limit + if high == float("inf"): + return torch.clamp(t, low) + return torch.clamp(t, low, high) + + batch_size, seq_len, _ = input_states.shape + dtype = input_states.dtype + + # distinguish prefill and decoding stage + is_decoding = torch.tensor(seq_len == 1).to(dtype) + + # 1. Gated MLP's linear projection + # FalconH1 applies the mamba padding mask (mamba_attention_mask) only during prefill. + if attention_mask is not None: + input_states_prefill = (input_states * attention_mask[:, :seq_len, None]).to(dtype) + input_states = input_states_prefill * (1.0 - is_decoding) + input_states * is_decoding + # FalconH1-specific input multiplier + input_states = input_states * self.ssm_in_multiplier + projected_states = self.in_proj(input_states) + # FalconH1-specific MuP multipliers + projected_states = projected_states * self.mup_vector + gate, hidden_states_B_C, dt = projected_states.split( + [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 + ) + + # 2. Convolution sequence transformation + if cache_params is not None: + # 2.1 decoding step + # Shift the cached conv window left by one and append the new token's projected + # states in the last slot. The original implementation uses `torch.roll` followed by + # an in-place assignment to the last position; during TorchScript tracing that + # in-place scatter on a rolled (non-leaf) tensor is not reliably captured and yields + # an incorrect conv window in the exported graph. Rebuild the window with a + # traceable slice + concat instead so the OpenVINO IR matches the eager result. + conv_state_prev = cache_params.conv_states[layer_idx] + new_conv_col = hidden_states_B_C[:, 0, :] if hidden_states_B_C.ndim == 3 else hidden_states_B_C + conv_state_dec = torch.cat( + [conv_state_prev[:, :, 1:], new_conv_col[..., None]], + dim=-1, + ) + + hidden_states_B_C_dec = torch.sum( + conv_state_dec.to(projected_states.device) * self.conv1d.weight.squeeze(1), dim=-1 + ) + if self.use_conv_bias: + hidden_states_B_C_dec += self.conv1d.bias + hidden_states_B_C_dec = self.act(hidden_states_B_C_dec).to(dtype)[:, None, ...] + + # 2.2 prefill step + hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2) + conv_state_prefill = torch.nn.functional.pad( + hidden_states_B_C_transposed, (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0) + ) + + hidden_states_B_C_prefill = self.act( + self.conv1d(hidden_states_B_C_transposed).transpose(1, 2) + )[:, :seq_len, :] + + # store the correct conv state depending on the phase + conv_state = conv_state_prefill * (1.0 - is_decoding) + conv_state_dec * is_decoding + cache_params.conv_states[layer_idx].copy_(conv_state) + + hidden_states_B_C_prefill_masked = hidden_states_B_C_prefill + if attention_mask is not None: + hidden_states_B_C_prefill_masked = (hidden_states_B_C_prefill * attention_mask[:, :seq_len, None]).to(dtype) + else: + hidden_states_B_C_prefill = self.act( + self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2) + ) + hidden_states_B_C_dec = hidden_states_B_C_prefill[:, :1] + hidden_states_B_C_prefill_masked = hidden_states_B_C_prefill + + hidden_states_prefill, B_prefill, C_prefill = torch.split( + hidden_states_B_C_prefill_masked, + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], + dim=-1, + ) + hidden_states_dec, B_dec, C_dec = torch.split( + hidden_states_B_C_dec, + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], + dim=-1, + ) + + A = -torch.exp(self.A_log.float()) # [num_heads] + + # 3. SSM transformation + # 3.1 decoding step + if cache_params is not None: + dt_dec = dt + dt_dec = dt_dec.reshape(dt_dec.shape[0], -1, dt_dec.shape[-1])[:, :1, :] + dt_dec = dt_dec.transpose(1, 2).expand(batch_size, dt_dec.shape[-1], self.head_dim) + dt_bias_dec = self.dt_bias + dt_bias_dec = dt_bias_dec.reshape(dt_bias_dec.shape[0], -1).expand(dt_bias_dec.shape[0], self.head_dim) + dt_dec = torch.nn.functional.softplus(dt_dec + dt_bias_dec) + dt_dec = _clamp_dt(dt_dec) + + A_dec = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + dA = torch.exp(dt_dec[..., None] * A_dec) + + B_dec = B_dec.reshape(batch_size, -1)[:, : self.n_groups * self.ssm_state_size] + B_dec = B_dec.reshape(batch_size, self.n_groups, -1)[..., None, :] + B_dec = B_dec.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B_dec.shape[-1]).contiguous() + B_dec = B_dec.reshape(batch_size, -1, B_dec.shape[-1]) + dB = dt_dec[..., None] * B_dec[..., None, :] + + hidden_states_dec = hidden_states_dec.reshape(batch_size, -1, self.head_dim) + dBx = dB * hidden_states_dec[..., None] + + new_ssm_state_dec = cache_params.ssm_states[layer_idx] * dA + dBx + + C_dec = C_dec.reshape(batch_size, -1)[:, : self.n_groups * self.ssm_state_size] + C_dec = C_dec.reshape(batch_size, self.n_groups, -1)[..., None, :] + C_dec = C_dec.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C_dec.shape[-1]).contiguous() + C_dec = C_dec.reshape(batch_size, -1, C_dec.shape[-1]) + + ssm_states_dec = new_ssm_state_dec.to(C_dec.dtype) + ssm_states_reshaped = ssm_states_dec.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) + C_reshaped = C_dec.view(batch_size * self.num_heads, self.ssm_state_size, 1) + y_dec = torch.bmm(ssm_states_reshaped, C_reshaped) + y_dec = y_dec.view(batch_size, self.num_heads, self.head_dim) + + D_dec = self.D + D_dec = D_dec[..., None].expand(D_dec.shape[0], self.head_dim) + y_dec = (y_dec + hidden_states_dec * D_dec).to(y_dec.dtype) + + y_dec = y_dec.reshape(batch_size, -1)[:, None, ...] + + # 3.2 prefill step + dt = torch.nn.functional.softplus(dt + self.dt_bias) + dt = _clamp_dt(dt) + + hidden_states_prefill = hidden_states_prefill.reshape(batch_size, seq_len, -1, self.head_dim).float() + B_prefill = B_prefill.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() + C_prefill = C_prefill.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() + B_prefill = B_prefill.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) + C_prefill = C_prefill.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) + pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size + + D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states_prefill, pad_size) + + hidden_states_prefill = hidden_states_prefill * dt[..., None] + A = A.to(hidden_states_prefill.dtype) * dt + + hidden_states_prefill, A, B_prefill, C_prefill = [ + reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states_prefill, A, B_prefill, C_prefill) + ] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + L = torch.exp(segment_sum(A)) + + G_intermediate = C_prefill[:, :, :, None, :, :] * B_prefill[:, :, None, :, :, :] + G = G_intermediate.sum(dim=-1) + + M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] + M = M_intermediate.sum(dim=-1) + + Y_diag = (M[..., None] * hidden_states_prefill[:, :, None]).sum(3) + + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay_contraction = B_prefill * decay_states.permute(0, 2, 3, 1)[..., None] + + states = ( + ( + B_decay_contraction.permute(0, 1, 3, 2, 4)[..., None] + * hidden_states_prefill.permute(0, 1, 3, 2, 4)[..., None, :] + ) + .sum(dim=3) + .permute(0, 1, 2, 4, 3) + ) + previous_states = torch.zeros_like(states[:, :1]) + + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(torch.nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) + + states_permuted = states.permute(0, 2, 1, 3, 4) + result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2) + new_states = result.permute(0, 2, 1, 3, 4) + states, new_ssm_state_prefill = new_states[:, :-1], new_states[:, -1] + + state_decay_out = torch.exp(A_cumsum) + + C_times_states = C_prefill[..., None, :] * states[:, :, None, ...] + state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) + Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] + + y = Y_diag + Y_off + + y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) + + y = y + D_residual + + pad_mask = torch.tensor(pad_size > 0).to(torch.long) + y_new_len = y.size(1) * (1 - pad_mask) + seq_len * pad_mask + y = y[:, :y_new_len] + y_prefill = y.reshape(batch_size, seq_len, -1) + + if cache_params is not None: + y = y_prefill[:, :seq_len] * (1.0 - is_decoding) + y_dec * is_decoding + ssm_state = new_ssm_state_prefill * (1.0 - is_decoding) + new_ssm_state_dec * is_decoding + cache_params.ssm_states[layer_idx].copy_(ssm_state) + else: + y = y_prefill + + # FalconH1: gated RMS norm is optional + if self.mamba_rms_norm: + scan_output = self.norm(y, gate) + else: + scan_output = y * torch.nn.functional.silu(gate) + + contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size] + + return contextualized_states + + +# This patcher class serves the following purposes: +# 1. Packs the KV-cache, conv_state, and ssm_state tensors into a +# FalconHybridMambaAttentionDynamicCache structure for subsequent invocation of the +# model's `forward` method. Unlike alternating hybrids (Zamba2/GraniteMoeHybrid), every +# FalconH1 layer is hybrid and owns both mamba (conv/ssm) and attention (key/value) states. +# 2. Patches the FalconH1Mixer so that the traced `forward` function works correctly +# during both the prefill and decoding steps. +class FalconH1ModelPatcher(OVDecoderModelPatcher): + def __init__( + self, + config: "OpenVINOConfig", + model: "PreTrainedModel", + model_kwargs: Optional[Dict[str, Any]] = None, + ): + from transformers.models.falcon_h1.modeling_falcon_h1 import FalconHybridMambaAttentionDynamicCache + + super().__init__(config, model, model_kwargs) + + class FalconH1DynamicCacheWrap(FalconHybridMambaAttentionDynamicCache): + def __init__(self, config, batch_size: int, conv_states, ssm_states, key_cache, value_cache): + # The parent constructor indexes a per-layer `devices` list to allocate the initial + # (zero) conv/ssm states. Those tensors are overwritten immediately below with the + # real cache tensors, so a list of `None` devices (default device) is sufficient here. + super().__init__( + config=config, + batch_size=batch_size, + devices=[None] * config.num_hidden_layers, + ) + self.conv_states = conv_states + self.ssm_states = ssm_states + self.key_cache = key_cache + self.value_cache = value_cache + # every layer is hybrid, so cache indices map 1:1 with layer indices + self.has_previous_state = True + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[dict[str, Any]] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Always concatenate past and new states. The past KV tensor has a seq + # length of 0 during the first (prefill) call, so the concat is a no-op + # then and correctly extends the cache on subsequent decode calls. We check + # the last (head_dim) dimension rather than the seq dimension so that the + # empty-cache branch is never taken (head_dim is always > 0), keeping the + # past-KV ReadValue connected to the attention for stateful export. + if self.key_cache[layer_idx].shape[-1] == 0: + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def __getitem__(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]: + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx].numel() == 0: + return 0 + return self.key_cache[layer_idx].shape[-2] + + # the patch is needed to include KV-cache, Conv, and SSM states in the inputs and outputs. + def patched_forward( + input_ids, + attention_mask=None, + cache_params=None, + ): + num_hidden_layers = self.real_config._config.num_hidden_layers + use_cache = False + wrapped_cache_params = None + if cache_params is not None: + use_cache = True + conv_states = [] + ssm_states = [] + key_cache = [] + value_cache = [] + + # decouple ssm_states, conv_states, keys and values from cache_params + batch_size = cache_params[0].size(0) + for idx in range(num_hidden_layers): + conv_states.append(cache_params[2 * idx]) + ssm_states.append(cache_params[2 * idx + 1]) + + for idx in range(num_hidden_layers): + key_cache.append(cache_params[2 * num_hidden_layers + 2 * idx]) + value_cache.append(cache_params[2 * num_hidden_layers + 2 * idx + 1]) + + wrapped_cache_params = FalconH1DynamicCacheWrap( + self.real_config._config, batch_size, conv_states, ssm_states, key_cache, value_cache + ) + + # The FalconH1 model derives `cache_position` (and therefore `position_ids`) from + # `torch.arange(seq_len)` whenever `cache_position` is not supplied, WITHOUT offsetting + # by the number of previously cached tokens. During single-token stateful decode that + # would place every new token at absolute position 0, so its RoPE phase would not match + # the cached keys (which were rotated at their real positions during prefill), producing + # a persistent logit error. `_update_causal_mask` also uses `attention_mask.shape[-1]` + # as the key/value target length, so the mask must span the whole context. + # + # We therefore build an explicit `cache_position` starting at the current past attention + # length and a full-length all-ones attention mask covering past_len + seq_len. Passing an + # all-ones mask keeps the model on the SDPA `is_causal` fast path (no explicit, in-place + # 4D causal mask is materialised), which is numerically correct for the standard causal + # case and reusable for prefill and decode in a single exported graph. + cache_position = None + model_attention_mask = attention_mask + if wrapped_cache_params is not None: + past_len = wrapped_cache_params.key_cache[0].shape[-2] + seq_len = input_ids.shape[1] + cache_position = torch.arange(past_len, past_len + seq_len, device=input_ids.device) + model_attention_mask = torch.ones( + (input_ids.shape[0], past_len + seq_len), + dtype=torch.int64, + device=input_ids.device, + ) + + causal_lm_output = self.model_orig_forward( + input_ids=input_ids, + attention_mask=model_attention_mask, + past_key_values=wrapped_cache_params, + use_cache=use_cache, + cache_position=cache_position, + ) + outputs = { + "logits": causal_lm_output.logits, + } + + if use_cache: + past_key_values = causal_lm_output.past_key_values + present_key_values = [] + for idx in range(num_hidden_layers): + present_key_values.append(past_key_values.conv_states[idx]) + present_key_values.append(past_key_values.ssm_states[idx]) + + for idx in range(num_hidden_layers): + present_key_values.append(past_key_values.key_cache[idx]) + present_key_values.append(past_key_values.value_cache[idx]) + + outputs["present_key_values"] = present_key_values + + return outputs + + self.patched_forward = patched_forward + self.model_orig_forward = self.orig_forward + self.orig_forward = patched_forward + + def __enter__(self): + super().__enter__() + setattr(self._model, self.orig_forward_name, self.patched_forward) + + for layer in self._model.model.layers: + mamba_layer = layer.mamba + mamba_layer._orig_forward = mamba_layer.forward + mamba_layer.forward = types.MethodType(falcon_h1_mamba_mixer, mamba_layer) + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + setattr(self._model, self.orig_forward_name, self.model_orig_forward) + for layer in self._model.model.layers: + mamba_layer = layer.mamba + mamba_layer.forward = mamba_layer._orig_forward + + # This patcher class serves the following purposes: # 1. Packs the KV-cache, conv_state, and ssm_state tensors into a Zamba2HybridDynamicCache structure # for subsequent invocation of the model's `forward` method. diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index 55c6d852e2..8af8ef3d73 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -344,6 +344,7 @@ def _get_kokoro_submodels(model): SSM_MODELS = [ "mamba", "falcon_mamba", + "falcon_h1", "zamba2", "lfm2", "lfm2_moe", diff --git a/optimum/intel/openvino/modeling_decoder.py b/optimum/intel/openvino/modeling_decoder.py index aa75996152..7fd013d330 100644 --- a/optimum/intel/openvino/modeling_decoder.py +++ b/optimum/intel/openvino/modeling_decoder.py @@ -1163,6 +1163,27 @@ def __init__( self.mamba_headdim = getattr(config, "mamba_d_head", None) self.num_mamba_layers = layer_types.count("mamba") self.num_attn_layers = layer_types.count("attention") + elif config.model_type == "falcon_h1": + # FalconH1 is a fully-parallel hybrid: every decoder layer contains both a Mamba2 + # mixer (conv + ssm states) and a self-attention block (key + value cache), so the + # number of mamba layers and attention layers both equal `num_hidden_layers`. + # FalconH1 also uses distinct config attribute names (`mamba_n_groups`, + # `mamba_n_heads`, `mamba_d_head`, `mamba_d_state`, `mamba_d_ssm`) and a mamba + # intermediate size that is independent of the MLP `intermediate_size`. + self.num_key_value_heads = getattr(config, "num_key_value_heads", None) + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.mamba_ngroups = getattr(config, "mamba_n_groups", None) + self.n_mamba_heads = getattr(config, "mamba_n_heads", None) + self.ssm_state_size = getattr(config, "mamba_d_state", None) + self.mamba_headdim = getattr(config, "mamba_d_head", None) + # mamba intermediate size (conv/ssm hidden dim), NOT the MLP intermediate_size + self.intermediate_size = ( + config.mamba_d_ssm + if getattr(config, "mamba_d_ssm", None) is not None + else int(config.mamba_expand * config.hidden_size) + ) + self.num_mamba_layers = config.num_hidden_layers + self.num_attn_layers = config.num_hidden_layers else: # Mamba 2 specific parameters hybrid_layer_ids = getattr(config, "hybrid_layer_ids", None) @@ -1527,6 +1548,8 @@ def prepare_inputs_for_generation( # to be the length of the full context, so default mask from OVModelForCausalLM needs to be used. # Other models like Mamba typically do not require an attention_mask # for the decoding step after the first token so use attention mask of ones. + # FalconH1 reconstructs a full-length decode mask inside its exporter patcher, + # so a length-1 ones mask here is sufficient. attention_mask = torch.ones_like(input_ids, dtype=torch.int64) else: diff --git a/tests/openvino/test_decoder.py b/tests/openvino/test_decoder.py index 9a17d0bdac..764edf2fbf 100644 --- a/tests/openvino/test_decoder.py +++ b/tests/openvino/test_decoder.py @@ -131,6 +131,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase): "granitemoehybrid", "mamba", "falcon_mamba", + "falcon_h1", "zamba2", "lfm2", "lfm2_moe", @@ -221,6 +222,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase): "qwen3_moe": 2, "mamba": 0, "falcon_mamba": 0, + "falcon_h1": 2, "arcee": 2, "smollm3": 2, "gpt_oss": 2, @@ -317,6 +319,13 @@ def test_compare_to_transformers(self, model_arch): "pegasus", ) and is_openvino_version(">=", "2026.1.0"): self.skipTest("CVS-185350: OpenVINO 2026.1.0 inference results mismatch") + if model_arch == "falcon_h1": + # FalconH1 stateful decode is validated to match Transformers greedily for + # (right-aligned / unpadded) generation. This test additionally exercises a + # left-padded batch: the stateful mamba conv/ssm recurrence has no per-position + # padding mask during single-token decode, so left-padded rows diverge (the same + # hybrid-mamba limitation for which zamba2/granitemoehybrid are skipped above). + self.skipTest("FalconH1: left-padded batched stateful mamba decode is not bit-exact") self.mock_torch_compile(model_arch) model_id = MODEL_NAMES[model_arch] @@ -428,7 +437,10 @@ def test_compare_to_transformers(self, model_arch): # LFM2 fails with beam search, issue link: https://github.com/huggingface/transformers/issues/42257 # CVS-177964 GraniteMoeHybrid, Qwen3-Next fail due to lack of support for beam search for hybrid models in OpenVINO # For this support, we expect changes in IRs to have connected beam_idx with Mamba/Linear attention states - num_beams=1 if model_arch in ["chatglm4", "lfm2", "granitemoehybrid", "qwen3_next"] else 2, + # FalconH1 is a fully-parallel mamba+attention hybrid with the same beam-search limitation. + num_beams=1 + if model_arch in ["chatglm4", "lfm2", "granitemoehybrid", "qwen3_next", "falcon_h1"] + else 2, do_sample=False, ) @@ -673,7 +685,9 @@ def test_beam_search(self, model_arch): return # LFM2, LFM2-MoE and GraniteMoeHybrid generate wrong output with beam search, ticket: CVS-185664 - if model_arch in ["lfm2", "lfm2_moe", "granitemoehybrid"]: + # FalconH1 is a mamba+attention hybrid with the same beam-search limitation (mamba states are + # not reordered by beam_idx in the stateful OpenVINO model). + if model_arch in ["lfm2", "lfm2_moe", "granitemoehybrid", "falcon_h1"]: return # TODO: add back once https://huggingface.co/katuni4ka/tiny-random-minicpm3/discussions/1 merged (for all models) as current modeling incompatible with transformers >= v4.49 diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index d9928449a3..894f2c3ccf 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -114,6 +114,7 @@ class ExportModelTest(unittest.TestCase): "fun_asr": OVModelForSpeechSeq2Seq, "mamba": OVModelForCausalLM, "falcon_mamba": OVModelForCausalLM, + "falcon_h1": OVModelForCausalLM, "gemma4": OVModelForVisualCausalLM, "gemma4_moe": OVModelForVisualCausalLM, "qwen3_5": OVModelForVisualCausalLM, diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py index c57789a4c5..ab181809b3 100644 --- a/tests/openvino/test_exporters_cli.py +++ b/tests/openvino/test_exporters_cli.py @@ -135,6 +135,7 @@ class OVCLIExportTestCase(unittest.TestCase): ("text-generation-with-past", "lfm2_moe"), ("text-generation-with-past", "mamba"), ("text-generation-with-past", "falcon_mamba"), + ("text-generation-with-past", "falcon_h1"), ("text-to-image", "flux.2-klein"), ] # filter architectures depending on min/max transformers supported versions @@ -175,6 +176,7 @@ class OVCLIExportTestCase(unittest.TestCase): "clip": 2, "mamba": 2, "falcon_mamba": 2, + "falcon_h1": 2, "qwen3": 2, "qwen3_omni_moe": 2, "zamba2": 2, diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 74418694cc..628a596600 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -144,6 +144,107 @@ def _create_tiny_kokoro_model(): return str(output_dir) +def _create_tiny_falcon_h1_model(): + """Generate a tiny random FalconH1 (``falcon_h1``) model and return its local path. + + FalconH1 is a fully-parallel hybrid architecture: every decoder layer contains both a + Mamba2 mixer (conv + ssm states) and a self-attention block (key + value cache). The + reduced config below preserves that architecture identity and every coupling invariant + (attention ``head_dim`` / GQA grouping, ``mamba_d_ssm == mamba_n_heads * mamba_d_head``, + conv/state/group coupling, all MuP multipliers and special-token ids) while shrinking + scale dimensions so the artifact stays well under 100 MiB. No original weights are + downloaded. The result is cached on disk, so repeated test collection is cheap. + """ + from transformers import AutoConfig, AutoTokenizer, FalconH1ForCausalLM + + original_model_id = "tiiuae/Falcon-H1-0.5B-Instruct" + output_dir = Path(tempfile.gettempdir()) / "optimum_intel_tiny_random_falcon_h1" + cache_version = "falcon_h1-tiny-v2" + marker = output_dir / ".tiny_cache_marker.json" + config_file = output_dir / "config.json" + weights_file = output_dir / "model.safetensors" + + # Reuse a valid cache only after checking the version marker and key invariants. + if marker.exists() and config_file.exists() and weights_file.exists(): + try: + meta = json.loads(marker.read_text()) + cfg = AutoConfig.from_pretrained(output_dir) + if ( + meta.get("cache_version") == cache_version + and cfg.model_type == "falcon_h1" + and cfg.mamba_d_ssm == cfg.mamba_n_heads * cfg.mamba_d_head + ): + return str(output_dir) + except Exception: + pass + + output_dir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(SEED) + + original_cfg = AutoConfig.from_pretrained(original_model_id) + cfg = original_cfg.to_dict() + + head_dim = cfg["head_dim"] # preserve attention head_dim (64) + cfg["hidden_size"] = 256 + cfg["num_attention_heads"] = 4 # 4 * 64 == 256 == hidden_size + cfg["num_key_value_heads"] = 1 # GQA grouping preserved + cfg["head_dim"] = head_dim + cfg["intermediate_size"] = 512 + cfg["num_hidden_layers"] = 2 + + mamba_d_head = 32 + mamba_n_heads = 8 + cfg["mamba_d_head"] = mamba_d_head + cfg["mamba_n_heads"] = mamba_n_heads + cfg["mamba_d_ssm"] = mamba_n_heads * mamba_d_head # 256 + cfg["mamba_d_state"] = 32 + cfg["mamba_n_groups"] = 1 + cfg["mamba_d_conv"] = 4 + cfg["mamba_chunk_size"] = 16 + cfg["max_position_embeddings"] = 512 + cfg["torch_dtype"] = "float32" + + model_type = cfg.pop("model_type") + cfg.pop("architectures", None) + tiny_cfg = AutoConfig.for_model(model_type, **cfg) + tiny_cfg.architectures = original_cfg.architectures + + assert tiny_cfg.model_type == "falcon_h1" + assert tiny_cfg.mamba_d_ssm == tiny_cfg.mamba_n_heads * tiny_cfg.mamba_d_head + + model = FalconH1ForCausalLM(tiny_cfg).to(torch.float32) + # Wider output-head init (with narrower embeddings) so the tiny model produces + # well-separated logits. This keeps greedy decoding stable under the small numeric + # differences between the PyTorch reference and the OpenVINO mamba SSD kernels, so the + # HF-vs-OpenVINO comparison test is not flipped by a razor-thin top-1/top-2 margin. + with torch.no_grad(): + model.lm_head.weight.normal_(mean=0.0, std=0.5) + model.model.embed_tokens.weight.normal_(mean=0.0, std=0.05) + model.eval() + + model.save_pretrained(output_dir, safe_serialization=True) + AutoTokenizer.from_pretrained(original_model_id).save_pretrained(output_dir) + try: + from transformers import GenerationConfig + + GenerationConfig.from_pretrained(original_model_id).save_pretrained(output_dir) + except Exception: + pass + + marker.write_text( + json.dumps( + { + "cache_version": cache_version, + "model_type": tiny_cfg.model_type, + "architectures": tiny_cfg.architectures, + "num_hidden_layers": tiny_cfg.num_hidden_layers, + "mamba_d_ssm": tiny_cfg.mamba_d_ssm, + } + ) + ) + return str(output_dir) + + SEED = 42 F32_CONFIG = {"INFERENCE_PRECISION_HINT": "f32"} @@ -213,6 +314,7 @@ def _create_tiny_kokoro_model(): "gemma4_unified": "optimum-intel-internal-testing/tiny-random-gemma4-unified", "falcon": "optimum-intel-internal-testing/really-tiny-falcon-testing", "falcon-40b": "optimum-intel-internal-testing/tiny-random-falcon-40b", + "falcon_h1": _create_tiny_falcon_h1_model(), "falcon_mamba": "optimum-intel-internal-testing/tiny-falcon-mamba", "flaubert": "optimum-intel-internal-testing/tiny-random-flaubert", "flux": "optimum-intel-internal-testing/tiny-random-flux", @@ -594,6 +696,7 @@ def _resolve_cached_model_paths(model_names: dict) -> dict: "resampler_model": 6, }, "zamba2": {"model": 44}, + "falcon_h1": {"model": 44}, "exaone4": {"model": 16}, "lfm2": {"model": 52 if is_transformers_version("<", "5") else 54}, "lfm2_moe": {"model": 46}, From cab9e063f34270584e84e68973e756ba79d034d9 Mon Sep 17 00:00:00 2001 From: popovaan Date: Tue, 4 Aug 2026 20:43:06 +0200 Subject: [PATCH 2/2] Address review feedback from @Mohamed-Ashraf273 --- tests/openvino/test_decoder.py | 65 ++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/openvino/test_decoder.py b/tests/openvino/test_decoder.py index 764edf2fbf..6feb166d92 100644 --- a/tests/openvino/test_decoder.py +++ b/tests/openvino/test_decoder.py @@ -489,6 +489,71 @@ def test_compare_to_transformers(self, model_arch): del ov_model gc.collect() + def test_falcon_h1_stateful_decode_matches_transformers(self): + # Focused regression test for the FalconH1 stateful single-token decode fix. + # + # `test_compare_to_transformers` skips falcon_h1 because it also exercises a + # left-padded batch, which is a genuine hybrid-mamba limitation (the stateful + # conv/ssm recurrence has no per-position padding mask during single-token + # decode, the same reason zamba2/granitemoehybrid are skipped there). That skip + # left the exact bug this PR fixes without an automated OV-vs-Transformers + # assertion: during stateful decode the patched forward previously passed no + # `cache_position` and only a length-1 attention mask, so every decode token was + # placed at RoPE position 0 and its causal mask attended to a single key instead + # of the full KV context. + # + # This test drives an *unpadded*, right-aligned single prompt through greedy + # (num_beams=1) generation over multiple decode steps, which is exactly the path + # the fix corrects, and asserts bit-exact greedy token-id agreement with the + # Transformers reference. The documented left-padded-batch and beam-search + # limitations remain unchanged (this test uses neither padding nor beam search). + model_arch = "falcon_h1" + self.mock_torch_compile(model_arch) + model_id = MODEL_NAMES[model_arch] + + set_seed(SEED) + ov_model = OVModelForCausalLM.from_pretrained( + model_id, export=True, ov_config=F32_CONFIG, device=OPENVINO_DEVICE + ) + self.assertTrue(ov_model.stateful) + + set_seed(SEED) + transformers_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32) + + tokenizer = AutoTokenizer.from_pretrained(model_id) + # Single, unpadded prompt (batch size 1) so no per-position padding mask is + # involved; this isolates the cache_position / full-context attention-mask path. + tokens = tokenizer("What is the capital of France?", return_tensors="pt") + + # Disable EOS so generation always runs the full number of decode steps and the + # RoPE-position / KV-context regression is exercised across many decode steps + # (not just prefill). + ov_model.generation_config.eos_token_id = None + transformers_model.generation_config.eos_token_id = None + ov_model.config.eos_token_id = None + transformers_model.config.eos_token_id = None + gen_config = GenerationConfig( + max_new_tokens=20, + min_new_tokens=20, + num_beams=1, + do_sample=False, + ) + + set_seed(SEED) + with torch.no_grad(): + transformers_outputs = transformers_model.generate(**tokens, generation_config=gen_config) + set_seed(SEED) + ov_outputs = ov_model.generate(**tokens, generation_config=gen_config) + + self.assertTrue( + torch.equal(ov_outputs, transformers_outputs), + f"OV output {ov_outputs}\nTransformers output {transformers_outputs}", + ) + + del transformers_model + del ov_model + gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES) @pytest.mark.run_slow @slow