huggingface/transformers · #50311
OpenVINO exporter: let OpenVINO fuse attention, RoPE, the KV cache and compressed matmuls
docs/source/en/exporters.md1 + / 0 −
@@ -840,6 +840,7 @@ if it's using something else. - `grouped_mm` traces fine through `DynamoExporter` and is auto-translated for `OnnxExporter`. For `ExecutorchExporter` with either XNNPACK or MLX, the exporter swaps MoE experts to `batched_mm` before export.+- Sliding-window layers (Mistral, Gemma 3) export with a full key-value cache; the sliding mask still applies the window. ## Next steps src/transformers/exporters/exporter_dynamo.py16 + / 1 −
@@ -59,7 +59,7 @@ import torch from torch.export import ExportedProgram - from ..cache_utils import Cache+ from ..cache_utils import Cache, DynamicLayer from ..modeling_utils import PreTrainedModel @@ -214,6 +214,21 @@ def _flat_forward(**kwargs): # applied to every chunked-vision attention class — see the long list below). +@register_patch(+ "dynamo",+ "transformers.cache_utils.DynamicSlidingWindowLayer.update",+ "transformers.cache_utils.DynamicSlidingWindowLayer.get_mask_sizes",+ "transformers.cache_utils.DynamicSlidingWindowLayer.get_seq_length",+)+def _patch_sliding_window_layer(original):+ """Trace sliding-window cache layers as a full `DynamicLayer`; the sliding mask still applies the window.++ A trace bakes in the sliding layer's Python token counter and its cut to the window, so the exported decode breaks+ past the window. Temporary: it goes once huggingface/transformers#49542 keeps sliding-window decode on one graph.+ """+ return getattr(DynamicLayer, original.__name__)++ @register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier") def _patch_classifier_cast(_original): """Disable classifier dtype cast in nllb-moe (not traceable)."""src/transformers/exporters/exporter_openvino.py301 + / 79 −
@@ -31,7 +31,8 @@ 4. **Conversion**: the prepared module is decoded via ``TorchFXPythonDecoder`` and handed to ``openvino.convert_model`` together with the custom ``ConversionExtension``\\ s for ops without a built-in OV lowering. Ports are then renamed to their dotted leaf paths and- non-tensor inputs repaired.+ non-tensor inputs repaired. A bf16/fp16 model is traced with f32 activations and 16-bit+ weights (see "16-bit models" below). 5. **Stateful transformation** (``OpenVINOConfig.stateful``, on by default): fold round-tripped state tensors (KV cache, SSM states, …) into internal OV variables with a fused ``beam_idx`` reorder. Optionally written to disk via ``openvino.save_model`` when@@ -40,22 +41,27 @@ from __future__ import annotations +import contextlib+import copy import math import operator import re from collections.abc import Mapping, MutableMapping-from typing import TYPE_CHECKING, Any+from typing import Any +from ..integrations.sdpa_attention import repeat_kv from ..utils import logging-from ..utils.import_utils import is_openvino_available, is_torch_available+from ..utils.import_utils import is_openvino_available, is_openvino_greater_or_equal, is_torch_available from .configs import OpenVINOConfig from .exporter_dynamo import DynamoExporter, is_cache_object from .exporter_onnx import disambiguate_io_names, patch_model_outputs from .utils import ( apply_fx_node_fixes, apply_fx_program_fixes, apply_patches,+ cast_leaf_tensors, get_leaf_tensors,+ module_device, register_fx_node_fix, register_fx_program_fix, register_patch,@@ -67,6 +73,9 @@ from torch.export import ExportedProgram from torch.export.decomp_utils import CustomDecompTable + from ..integrations.moe import _default_apply_gate+ from ..modeling_utils import PreTrainedModel+ if is_openvino_available(): import numpy as np@@ -75,11 +84,10 @@ from openvino._offline_transformations import apply_make_stateful_transformation from openvino.frontend.pytorch import ConversionExtension from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder- from openvino.frontend.pytorch.torchdynamo.export_decompositions import ops_to_not_decompose --if TYPE_CHECKING:- from ..modeling_utils import PreTrainedModel+ # OpenVINO 2026.5 converts exported programs without decomposing them, and drops its decomposition table.+ if not is_openvino_greater_or_equal("2026.5", accept_dev=True):+ from openvino.frontend.pytorch.torchdynamo.export_decompositions import ops_to_not_decompose logger = logging.get_logger(__name__)@@ -113,11 +121,16 @@ def export( elif type(config) is not OpenVINOConfig: raise TypeError(f"Expected config to be an OpenVINOConfig or dict, got {type(config)}") - with patch_model_outputs(model) as (inputs_names, outputs_names), apply_patches("openvino"):+ # From here on the inputs match the graph's dtypes: calibration and the state folding read them too.+ with (+ apply_patches("openvino"),+ patch_16bit_model(model, sample_inputs) as sample_inputs,+ patch_model_outputs(model) as (inputs_names, outputs_names),+ ): exported_program: ExportedProgram = super().export(model, sample_inputs, config=config)-- exported_program, graph_module = _fix_exported_program(exported_program)- ov_model = _convert_to_openvino(graph_module)+ # Converted before the model gets its dtypes back: the program's constants are its own tensors.+ exported_program, graph_module = _fix_exported_program(exported_program)+ ov_model = _convert_to_openvino(graph_module) inputs_names = [name for name in inputs_names if name in get_leaf_tensors(sample_inputs)] inputs_names, outputs_names = disambiguate_io_names(inputs_names, outputs_names)@@ -292,6 +305,100 @@ def _fix_non_tensor_inputs(ov_model: openvino.Model, graph_module) -> None: ov_model.validate_nodes_and_infer_types() +# ── 16-bit models ──────────────────────────────────────────────────────────+# OV's CPU plugin fuses its kernels on f32-typed graphs and picks the inference precision itself+# (bf16 on CPUs with AMX or AVX-512 BF16): a bf16-typed graph misses some of them, notably the+# compressed-weight matmuls of an NNCF int8/int4 graph. So, as OV's own `__make_16bit_traceable`+# does, a bf16/fp16 model is traced with f32 activations while its linear and embedding weights stay+# 16-bit: the patched `nn.Linear` / `nn.Embedding` / MoE experts forwards cast them to f32 inside the graph, so they+# land in the OV model as a 16-bit constant and a `Convert`, which the plugin decompresses inside its+# matmuls, without the weights ever being copied.+++@contextlib.contextmanager+def patch_16bit_model(model: PreTrainedModel, sample_inputs: MutableMapping[str, Any]):+ """Trace a bf16/fp16 model with f32 activations and 16-bit weights, yielding f32 copies of the inputs. The model gets+ its dtypes back afterwards."""+ half = (torch.bfloat16, torch.float16)++ set_aside = []+ for module in model.modules():+ for name in _16bit_parameter_names(module):+ parameter = getattr(module, name, None)+ if parameter is not None and parameter.dtype in half:+ set_aside.append((module, name, parameter))+ set_aside_ids = {id(parameter) for _, _, parameter in set_aside}++ # The other 16-bit tensors (norms, ...) are small: upcast for the trace.+ upcast = []+ for tensor in [*model.parameters(), *model.buffers()]:+ if tensor.dtype in half and id(tensor) not in set_aside_ids:+ upcast.append((tensor, tensor.dtype))++ # The configs too: masks and caches the forward builds take their dtype from them.+ configs = []+ for config in _configs(model):+ if getattr(config, "dtype", None) in half:+ configs.append((config, config.dtype))++ if not set_aside and not upcast:+ yield sample_inputs+ return++ for module, name, parameter in set_aside:+ module.register_parameter(f"_16bit_{name}", parameter)+ # An f32 placeholder broadcasting a single zero: code reading `weight.dtype` sees f32, at no memory cost.+ placeholder = torch.zeros((), device=parameter.device).expand(parameter.shape)+ setattr(module, name, torch.nn.Parameter(placeholder, requires_grad=False))+ for tensor, _ in upcast:+ # In place, so every module sharing the tensor sees it.+ tensor.data = tensor.data.float()+ for config, _ in configs:+ config.dtype = torch.float32++ try:+ # A copy: casting rewrites dicts and caches in place, and the caller's inputs stay theirs.+ yield cast_leaf_tensors(copy.deepcopy(sample_inputs), torch.float32, module_device(model))+ finally:+ for module, name, parameter in set_aside:+ setattr(module, name, parameter)+ del module._parameters[f"_16bit_{name}"]+ for tensor, dtype in upcast:+ tensor.data = tensor.data.to(dtype)+ for config, dtype in configs:+ config.dtype = dtype+++def _16bit_parameter_names(module: torch.nn.Module) -> tuple[str, ...]:+ """The parameters of `module` that stay 16-bit through [`patch_16bit_model`]: the patched forwards cast them."""+ if isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):+ return ("weight", "bias")+ if hasattr(module, "has_gate") and hasattr(module, "is_transposed"): # `@use_experts_implementation` experts+ return ("gate_up_proj", "up_proj", "down_proj", "gate_up_proj_bias", "up_proj_bias", "down_proj_bias")+ return ()+++def _16bit_parameter(module: torch.nn.Module, name: str) -> torch.Tensor | None:+ """`module.<name>`, or the 16-bit parameter [`patch_16bit_model`] set aside in its place."""+ return module._parameters.get(f"_16bit_{name}", getattr(module, name, None))+++def _configs(model: PreTrainedModel):+ """Every config `model` holds, once: its pretrained models' (some, like T5's stacks, hold their own copy) and their+ sub-configs (a multimodal model's text and vision ones)."""+ pending = [module.config for module in model.modules() if isinstance(module, PreTrainedModel)]+ seen = set()+ while pending:+ config = pending.pop()+ if id(config) in seen:+ continue+ seen.add(id(config))+ yield config+ for name in getattr(config, "sub_configs", {}):+ if getattr(config, name, None) is not None:+ pending.append(getattr(config, name))++ # ── Stateful transformation ───────────────────────────────────────────────── # Folds round-tripped state tensors (KV cache, SSM conv/ssm states, …) into internal OV # ``ReadValue``/``Assign`` variables so the runtime carries them across ``infer()`` calls@@ -354,7 +461,8 @@ def _state_init_dims( outputs_names: list[str], ) -> dict[str, list]: """Compute a per-dim init spec for every state input: ``"batch"`` (follow the main input's- batch at runtime), ``0`` (the growing dim — empty on first inference), or a concrete length.+ batch at runtime), ``"grow"`` (the growing dim — empty on first inference), or a concrete length (a NoPE+ layer's empty rotary cache keeps its ``0``). The growing dim is identified from the ``ExportedProgram``'s symbolic shapes: a state input dim whose SymInt expression differs from the paired output's (e.g. ``s2`` vs ``s2 + s3``)@@ -391,7 +499,7 @@ def _state_init_dims( elif str(in_val.shape[axis]) == str(out_val.shape[axis]): dims.append(int(sample.shape[axis])) else:- dims.append(0)+ dims.append("grow") # ``apply_make_stateful_transformation`` names each variable by concatenating the input # and output tensor names — key the specs the same way for the ReadValue lookup. init_dims[f"{input_name}{output_name}"] = dims@@ -463,7 +571,7 @@ def _build_state_initializers(ov_model: openvino.Model, init_dims: dict[str, lis if update_is_static: dims = [update_shape[axis].get_length() if d == "batch" else d for axis, d in enumerate(dims)] info = variables[op.get_variable_id()].get_info()- info.data_shape = openvino.PartialShape([-1 if d in ("batch", 0) else d for d in dims])+ info.data_shape = openvino.PartialShape([-1 if d in ("batch", "grow") else d for d in dims]) variables[op.get_variable_id()].update(info) # The growing dim's zero length is emitted as ``batch - batch`` rather than a literal # ``[0]``: the CPU plugin fuses single-consumer init subgraphs into@@ -472,7 +580,7 @@ def _build_state_initializers(ov_model: openvino.Model, init_dims: dict[str, lis zero = ov_ops.subtract(batch, batch) shape = ov_ops.concat( [- batch if d == "batch" else zero if d == 0 else ov_ops.constant(np.array([d], dtype=np.int64))+ batch if d == "batch" else zero if d == "grow" else ov_ops.constant(np.array([d], dtype=np.int64)) for d in dims ], axis=0,@@ -551,13 +659,14 @@ def _pin_state_update_shapes(ov_model: openvino.Model) -> None: """Reconcile each ``Assign``'s update shape with its variable's shape. The CPU plugin refuses to reorder dynamic descriptors into the state memory, so the two- sides must agree. When the update is fully static and the variable is not (a static trace- bakes the batch into the update while the variable was declared batch-dynamic), the variable- is pinned to the update's shape. Conversely, when shape inference leaves the update- under-specified against a static variable (olmo_hybrid's rolled conv state comes out+ sides must agree. When the update fixes dims the variable leaves dynamic (a static trace bakes+ the batch into the update, or a batch-1 sample's size-1 batch gets specialized), the variable+ takes those dims: a Reshape between the KV-cache concat and its ``Assign`` would keep the CPU+ plugin from fusing the cache into its attention. Conversely, when shape inference leaves the+ update under-specified against a static variable dim (olmo_hybrid's rolled conv state comes out ``[?,32,2..]`` against a ``[?,32,3]`` variable), a ``special_zero`` Reshape pins the- statically-known dims and copies the dynamic ones from the input. Pinning a variable- refines shapes downstream (other updates read from it), so this iterates to a fixpoint.+ statically-known dims and copies the dynamic ones from the input. Pinning a variable refines+ shapes downstream (other updates read from it), so this iterates to a fixpoint. """ variables = {variable.get_info().variable_id: variable for variable in ov_model.get_variables()} read_values = {op.get_variable_id(): op for op in ov_model.get_ordered_ops() if op.get_type_name() == "ReadValue"}@@ -572,29 +681,31 @@ def _pin_state_update_shapes(ov_model: openvino.Model) -> None: update_shape = update.get_partial_shape() if update_shape == variable_shape: continue- if update_shape.is_static:+ # Unknown or different ranks: nothing to reconcile dim by dim.+ if variable_shape.rank.is_dynamic or update_shape.rank.is_dynamic:+ continue+ if len(update_shape) != len(variable_shape):+ continue+ dims = list(zip(variable_shape, update_shape))+ if any(v.is_dynamic and u.is_static for v, u in dims): info = variable.get_info()- info.data_shape = update_shape+ info.data_shape = openvino.PartialShape([u if v.is_dynamic and u.is_static else v for v, u in dims]) variable.update(info) # The variable's shape must relax its init expression's, so the init gets the- # same static shape (its dims were runtime-derived but numerically identical).+ # same dims (they were runtime-derived but numerically identical). read_value = read_values.get(op.get_variable_id()) if read_value is not None and read_value.get_input_size() > 0:- target = np.array([dim.get_length() for dim in update_shape], dtype=np.int64)+ target = [dim.get_length() if dim.is_static else 0 for dim in info.data_shape] pinned_init = ov_ops.reshape(- read_value.input_value(0), ov_ops.constant(target), special_zero=False+ read_value.input_value(0), ov_ops.constant(np.array(target, dtype=np.int64)), special_zero=True ) read_value.set_arguments([pinned_init]) changed = True- continue- if variable_shape.rank.is_dynamic:- continue- target = [dim.get_length() if dim.is_static else 0 for dim in variable_shape]- if all(t == 0 for t in target):- continue- pinned = ov_ops.reshape(update, ov_ops.constant(np.array(target, dtype=np.int64)), special_zero=True)- op.input(0).replace_source_output(pinned.output(0))- changed = True+ elif any(v.is_static and u.is_dynamic for v, u in dims):+ target = [dim.get_length() if dim.is_static else 0 for dim in variable_shape]+ pinned = ov_ops.reshape(update, ov_ops.constant(np.array(target, dtype=np.int64)), special_zero=True)+ op.input(0).replace_source_output(pinned.output(0))+ changed = True if not changed: break ov_model.validate_nodes_and_infer_types()@@ -636,8 +747,11 @@ def _run_openvino_decompositions(exported_program: ExportedProgram) -> ExportedP Decomposing up front lets the FX node fixes and the bare-name renames below operate on the graph OV actually decodes — ``convert_model(exported_program)`` would re-run decompositions- internally, regenerating node names and silently discarding those fixes.+ internally, regenerating node names and silently discarding those fixes. Since OpenVINO 2026.5+ the frontend converts exported programs without decomposing them, and so does this. """+ if is_openvino_greater_or_equal("2026.5", accept_dev=True):+ return exported_program decomp_table = CustomDecompTable() for op in ops_to_not_decompose(): decomp_table.pop(op, None)@@ -738,6 +852,25 @@ def _fix_sym_float(gm, node): return True +@register_fx_node_fix("openvino")+def _fix_mm_dtype(gm, node):+ """Rewrite ``aten.mm.dtype`` (``torch.mm(..., out_dtype=...)``, which OV can't convert) as an ``mm`` on its+ inputs cast to ``out_dtype``, the dtype the CUDA kernel accumulates in."""+ if node.target is not torch.ops.aten.mm.dtype:+ return False+ *matrices, out_dtype = node.args+ with gm.graph.inserting_before(node):+ casts = [+ gm.graph.call_function(torch.ops.aten._to_copy.default, args=(matrix,), kwargs={"dtype": out_dtype})+ for matrix in matrices+ ]+ result = gm.graph.call_function(torch.ops.aten.mm.default, args=tuple(casts))+ result.meta.update(node.meta)+ node.replace_all_uses_with(result)+ gm.graph.erase_node(node)+ return True++ @register_fx_node_fix("openvino") def _fix_sym_min_max(gm, node): """Rewrite ``torch.sym_min``/``torch.sym_max`` to the built-in ``min``/``max``.@@ -965,9 +1098,13 @@ def _is_empty(n): # which is always correct (the empty-cat optimisation just doesn't fire). return guard_or_false(val.numel() == 0) - if _is_empty(operands[0]):+ first_empty, second_empty = _is_empty(operands[0]), _is_empty(operands[1])+ # a NoPE layer's zero-width rotary cache: empty with the same rank on both sides, the cat still grows it+ if first_empty and second_empty and operands[0].meta["val"].ndim == operands[1].meta["val"].ndim:+ return False+ if first_empty: keep = operands[1]- elif _is_empty(operands[1]):+ elif second_empty: keep = operands[0] else: return False@@ -1162,6 +1299,123 @@ def _fix_index_put_bool_mask(gm, node): # To add a new patch: define a `_patch_*` factory and decorate it. +@register_patch("openvino", "torch.Tensor.__neg__", "torch.neg", "torch.Tensor.neg")+def _patch_neg(original):+ """Negate as a multiply by ``-1``, the form OV's CPU plugin fuses ``rotate_half`` (so RoPE) from.++ The frontend lowers ``aten.neg`` to a multiply by a ``Convert``ed ``-1`` that it marks+ ``DisableConstantFolding``, which keeps the plugin's RoPE fusion from matching until the IR is+ serialized and read back: the model this exporter returns would run RoPE decomposed.+ """++ def patch(input):+ return input * -1++ return patch+++@register_patch("openvino", "torch.nn.Linear.forward")+def _patch_linear_forward(original):+ """Cast the 16-bit weight [`patch_16bit_model`] set aside to the f32 activations, inside the graph."""++ def patch(self, input):+ if "_16bit_weight" not in self._parameters:+ return original(self, input)+ weight = self._16bit_weight.to(input.dtype)+ bias = _16bit_parameter(self, "bias")+ if bias is not None:+ bias = bias.to(input.dtype)+ return torch.nn.functional.linear(input, weight, bias)++ return patch+++@register_patch("openvino", "torch.nn.Embedding.forward")+def _patch_embedding_forward(original):+ """Look up the 16-bit embedding [`patch_16bit_model`] set aside, cast to f32 inside the graph.++ The cast comes before the lookup, as in the linear patch: a tied `lm_head` then reads the same `Convert`, and+ stays a compressed-weight matmul once NNCF quantizes the shared weight.+ """++ def patch(self, input):+ if "_16bit_weight" not in self._parameters:+ return original(self, input)+ return torch.nn.functional.embedding(+ input,+ self._16bit_weight.float(),+ self.padding_idx,+ self.max_norm,+ self.norm_type,+ self.scale_grad_by_freq,+ self.sparse,+ )++ return patch+++@register_patch("openvino", "transformers.integrations.moe.ExpertsInterface.get_interface")+def _patch_experts_interface(original):+ """Trace every `@use_experts_implementation` experts module with [`_tiled_experts_forward`]."""++ def patch(self, experts_implementation, default):+ return _tiled_experts_forward++ return patch+++def _tiled_experts_forward(self, hidden_states, top_k_index, top_k_weights):+ """`batched_mm_experts_forward`, with every expert run on every token and weighted by the dense routing weights.++ This is the form OV's CPU plugin turns into `GatherMatmul`s that only compute the routed experts+ (`ConvertTiledMoeBlockTo{2,3}GatherMatmuls`): the hidden states tiled over the experts, one batched matmul per+ projection against the stacked 16-bit expert weights (gate and up apart for the default gating), scaled by the+ routing weights and summed over the experts. The per-token expert gather of the other forwards doesn't fuse.+ """+ num_tokens, hidden_dim = hidden_states.shape+ dtype = hidden_states.dtype++ def projection(name):+ return _16bit_parameter(self, name), _16bit_parameter(self, f"{name}_bias")++ def batched_linear(input, weight, bias=None):+ # (num_experts, num_tokens, input_dim) -> (num_experts, num_tokens, output_dim)+ weight = weight.to(dtype)+ output = torch.bmm(input, weight if self.is_transposed else weight.transpose(1, 2))+ return output + bias.to(dtype)[:, None] if bias is not None else output++ # Every token, once per expert: (num_experts, num_tokens, hidden_dim)+ selected_hidden_states = hidden_states.repeat(self.num_experts, 1).view(self.num_experts, num_tokens, hidden_dim)++ # --- Up projection per expert (batched), then gating or activation ---+ if not self.has_gate:+ proj_out = self.act_fn(batched_linear(selected_hidden_states, *projection("up_proj")))+ elif type(self)._apply_gate is _default_apply_gate and self.is_concatenated:+ # The default gating, with gate and up projected apart: the form OV fuses.+ gate_up_proj, gate_up_proj_bias = projection("gate_up_proj")+ gate_proj, up_proj = gate_up_proj.chunk(2, dim=-1 if self.is_transposed else -2)+ gate_bias, up_bias = (None, None) if gate_up_proj_bias is None else gate_up_proj_bias.chunk(2, dim=-1)+ gate = batched_linear(selected_hidden_states, gate_proj, gate_bias)+ up = batched_linear(selected_hidden_states, up_proj, up_bias)+ proj_out = self.act_fn(gate) * up+ else:+ proj_out = self._apply_gate(batched_linear(selected_hidden_states, *projection("gate_up_proj")))++ # --- Down projection per expert (batched) ---+ proj_out = batched_linear(proj_out, *projection("down_proj"))+ proj_out = proj_out.view(self.num_experts, num_tokens, hidden_dim)++ # Apply optional post expert norm+ if self.has_post_expert_norm:+ proj_out = self.post_expert_norm(proj_out)++ # Apply routing weights: zero for the experts a token isn't routed to+ routing_weights = torch.zeros(num_tokens, self.num_experts, dtype=dtype, device=hidden_states.device)+ routing_weights = routing_weights.scatter(1, top_k_index, top_k_weights.to(dtype))+ weighted_out = proj_out * routing_weights.transpose(0, 1)[..., None]+ return weighted_out.sum(dim=0)++ @register_patch("openvino", "torch.nn.functional.layer_norm") def _patch_layer_norm(original): """Substitute identity ``weight=ones``/``bias=zeros`` when either is ``None``.@@ -1228,37 +1482,36 @@ def patch(input, size=None, scale_factor=None, mode="nearest", align_corners=Non def _patch_sdpa(original): """Pre-expand K/V to Q's head count before calling SDPA. - Also clamps the additive mask so fully-masked rows stay finite, and zeroes those rows to match- the fused kernels torch runs them through.+ Also clamps the additive mask so fully-masked rows stay finite: such a row then gets the uniform+ average the mask describes, as transformers' eager attention gives it with a dtype-minimum mask. OV's ``opset13::ScaledDotProductAttention`` op rejects GQA shapes (e.g. Q=[B,4,T,D],- K/V=[B,2,T,D]) with ``Key input shape not compatible with other inputs``. Repeating K/V via- ``repeat_interleave`` on the head axis keeps the math identical and gives OV matching shapes.+ K/V=[B,2,T,D]) with ``Key input shape not compatible with other inputs``. Repeating K/V with+ ``repeat_kv`` (unsqueeze, expand, reshape) keeps the math identical and gives OV matching+ shapes, in the form OV's CPU plugin fuses with a stateful KV cache: a ``repeat_interleave``+ (a ``Gather``) between the cache and the attention leaves every attention decomposed. """ def patch(query, key, value, attn_mask=None, *args, **kwargs): # OV's SDPA diverges from aten on a *boolean* mask and returns NaN for fully masked rows # (legit under left padding + causal), poisoning the batch entry # (https://github.com/openvinotoolkit/openvino/issues/31630). Pass an additive mask with a # finite dtype minimum instead, like optimum-intel does.- unattended = None if attn_mask is not None: masked_value = torch.finfo(query.dtype).min if attn_mask.dtype == torch.bool:- unattended = ~attn_mask.any(dim=-1, keepdim=True) attn_mask = torch.where( attn_mask, torch.zeros((), dtype=query.dtype, device=attn_mask.device), torch.full((), masked_value, dtype=query.dtype, device=attn_mask.device), ) else: attn_mask = attn_mask.clamp_min(masked_value)- unattended = attn_mask.amax(dim=-1, keepdim=True) <= masked_value q_heads, k_heads = query.shape[-3], key.shape[-3] if q_heads != k_heads and q_heads % k_heads == 0: reps = q_heads // k_heads- key = key.repeat_interleave(reps, dim=-3)- value = value.repeat_interleave(reps, dim=-3)+ key = repeat_kv(key, reps)+ value = repeat_kv(value, reps) # OV's SDPA op takes no `scale` from the traced graph (the aten node's scalar `scale` isn't # surfaced as a frontend input), so it always applies its own default ``head_dim**-0.5``. Models # with a non-default scale (T5-family fold ``1/sqrt(d)`` into init and pass ``scale=1.0``) would be@@ -1268,38 +1521,7 @@ def patch(query, key, value, attn_mask=None, *args, **kwargs): default_scale = query.shape[-1] ** -0.5 if scale != default_scale: query = query * (scale / default_scale)- attn_output = original(query, key, value, attn_mask, *args, **kwargs)- # A row that masks every key has no defined value: OV returns the uniform average the mask- # describes, torch's fused kernels write zeros. Zero them so the export answers like eager.- if unattended is not None:- attn_output = torch.where(- unattended, torch.zeros((), dtype=attn_output.dtype, device=attn_output.device), attn_output- )- return attn_output-- return patch---@register_patch(- "openvino",- "transformers.integrations.sdpa_attention.repeat_kv",- "transformers.integrations.eager_paged.repeat_kv",- "transformers.integrations.flex_attention.repeat_kv",-)-def _patch_repeat_kv(original):- """Expand GQA K/V heads via ``repeat_interleave`` instead of a 5-D ``expand`` + ``reshape``.-- The stock ``repeat_kv`` unsqueezes to ``[b, kv_heads, 1, kv_seq, head_dim]`` then- ``expand``s the new axis to ``n_rep``. OV's frontend can't keep the ``kv_seq`` axis dynamic- through that 5-D broadcast in stateful decode (``Broadcast Check 'input_shape[j] == 1'`` on- dim 3), baking it to the traced length. ``repeat_interleave`` on the head axis is the exact- equivalent (per ``repeat_kv``'s own docstring) and lowers to an OV op that stays dynamic.- """-- def patch(hidden_states, n_rep):- if n_rep == 1:- return hidden_states- return hidden_states.repeat_interleave(n_rep, dim=1)+ return original(query, key, value, attn_mask, *args, **kwargs) return patch src/transformers/utils/import_utils.py15 + / 0 −
@@ -1069,6 +1069,21 @@ def is_openvino_available() -> bool: return _is_package_available("openvino")[0] +@lru_cache+def is_openvino_greater_or_equal(library_version: str, accept_dev: bool = False) -> bool:+ """+ Accepts a library version and returns True if the current version of the library is greater than or equal to the+ given version. If `accept_dev` is True, it will also accept development versions (e.g. 2026.5.0.dev20260901+ matches 2026.5.0).+ """+ available, openvino_version = _is_package_available("openvino", return_version=True)+ if not available:+ return False+ if accept_dev:+ return version.parse(version.parse(openvino_version).base_version) >= version.parse(library_version)+ return version.parse(openvino_version) >= version.parse(library_version)++ @lru_cache def is_nncf_available() -> bool: return _is_package_available("nncf")[0]tests/exporters/test_export.py15 + / 3 −
@@ -447,8 +447,13 @@ "RTDetrV2ForObjectDetection": "Same tied-`topk` selection as `DFineModel`.", "PPDocLayoutV2ForObjectDetection": "Same tied-`topk` selection as `DFineModel`.", "PPDocLayoutV3ForObjectDetection": "Same tied-`topk` selection as `DFineModel`.",+ "PPDocLayoutV4ForObjectDetection": "Same tied-`topk` selection as `DFineModel`.", "MMGroundingDinoModel": "Same tied-`topk` box selection as `DFineModel`.", "MMGroundingDinoForObjectDetection": "Same tied-`topk` box selection as `DFineModel`.",+ "TimesFmModelForPrediction": (+ "The left-padding patches attend to no key: OpenVINO gives them the uniform average, torch's fused "+ "kernels zeros. The real patches and the predictions match."+ ), }, # OpenVINO, generate path only. "openvino.generate": {},@@ -724,11 +729,12 @@ def _state_path(state): results = request.infer(feed) outputs = {}+ declared = {name for output in ov_model.outputs for name in output.get_names()} for port in compiled.outputs:- # Compilation may merge a named output tensor with an intermediate that kept its- # numeric id — prefer the human-readable alias over ``get_any_name``'s sorted-first.+ # Compilation may merge an output's tensor names with the ones of the intermediates it fused (a numeric id,+ # `add_4:0`): take the name the exported model gave it. names = sorted(port.get_names())- name = next((n for n in names if not n.isdigit()), names[0])+ name = next((n for n in names if n in declared), names[0]) outputs[re.sub(r"^output\.", "", name)] = torch.as_tensor(results[port]) # Folded state tensors are outputs too — read them back so the returned dict covers the@@ -989,6 +995,12 @@ def _check_outputs_close(self, actual, expected, atol, rtol, check_device=True, attend to -- so each runtime fills it differently and no caller reads it; comparing those positions measures nothing. """+ # exports keep a full cache on sliding-window layers, eager only the window: skip them+ sliding = tuple(+ key.removesuffix("_sliding_window_tensor") for key in expected if key.endswith("._sliding_window_tensor")+ )+ actual = {key: value for key, value in actual.items() if not key.startswith(sliding)}+ expected = {key: value for key, value in expected.items() if not key.startswith(sliding)} # a model with per-layer masks passes a `dict` here, and a flex-attention one a `BlockMask`; # only a plain 2D `(batch, seq)` tensor maps onto output positions. NaViT-style vision models # (siglip2) mark their valid patches with `pixel_attention_mask` instead.tests/exporters/test_fusions.pyadded168 + / 0 −
@@ -0,0 +1,168 @@+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+# http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+"""Fusion tests, one class per backend.++A backend only runs its fused kernels on the patterns it recognizes, and a graph that misses them still runs, decomposed+and slower. Each class exports the tiny model of each decoder family's own tests and counts, per layer, the fused ops+the backend runs, and that none of the ops they replace are left.+"""++import collections+import unittest++from parameterized import parameterized++from tests.models.gemma3.test_modeling_gemma3 import Gemma3TextModelTester+from tests.models.granite.test_modeling_granite import GraniteModelTester+from tests.models.llama.test_modeling_llama import LlamaModelTester+from tests.models.mistral.test_modeling_mistral import MistralModelTester+from tests.models.mixtral.test_modeling_mixtral import MixtralModelTester+from tests.models.olmo2.test_modeling_olmo2 import Olmo2ModelTester+from tests.models.phi3.test_modeling_phi3 import Phi3ModelTester+from tests.models.qwen2.test_modeling_qwen2 import Qwen2ModelTester+from tests.models.qwen3.test_modeling_qwen3 import Qwen3ModelTester+from tests.models.qwen3_moe.test_modeling_qwen3_moe import Qwen3MoeModelTester+from tests.test_memory_cleanup_mixin import MemoryCleanupMixin+from transformers import AutoModelForCausalLM, GenerationConfig, set_seed+from transformers.exporters import OpenVINOConfig, OpenVINOExporter+from transformers.testing_utils import require_nncf, require_openvino, require_torch, slow+from transformers.utils import is_torch_available+++if is_torch_available():+ import torch+++# model type: the tester of its tiny model+TESTERS = {+ "llama": LlamaModelTester,+ "qwen2": Qwen2ModelTester,+ "qwen3": Qwen3ModelTester, # per-head query and key norms+ "phi3": Phi3ModelTester,+ "olmo2": Olmo2ModelTester,+ "granite": GraniteModelTester,+ "mistral": MistralModelTester, # sliding-window attention+ "gemma3": Gemma3TextModelTester, # sliding-window and full attention layers+ "mixtral": MixtralModelTester, # mixture of experts+ "qwen3_moe": Qwen3MoeModelTester, # mixture of experts+}+MOE_MODELS = {"mixtral", "qwen3_moe"}+++def _model(model_type, dtype=None, **config):+ """The tiny causal LM the model's own tests use, with `config` overriding its config."""+ set_seed(0)+ model_config = TESTERS[model_type](parent=None).get_config()+ for name, value in config.items():+ setattr(model_config, name, value)+ return AutoModelForCausalLM.from_config(model_config, dtype=dtype or torch.float32).eval()+++# The fused kernels OpenVINO's CPU plugin must compile each layer of a decode graph into. `MatMul` counts the matmuls+# left decomposed (an unfused attention or expert), so it must stay at 0.+OPENVINO_FUSIONS = {"ScaledDotProductAttention": 1, "RoPE": 2, "MatMul": 0} # RoPE for the queries and the keys+OPENVINO_MOE_FUSIONS = {**OPENVINO_FUSIONS, "GatherMatmul": 3} # gate, up and down, computing only the routed experts+++def _decode(model, multi_token_decode=True, quantizer=None):+ inputs = {"input_ids": torch.tensor([[1, 5, 7, 9, 11]]), "attention_mask": torch.ones(1, 5, dtype=torch.long)}+ exported = OpenVINOExporter().export_for_generation(+ model,+ inputs,+ OpenVINOConfig(dynamic=True, quantizer=quantizer),+ generation_config=GenerationConfig(max_new_tokens=3),+ multi_token_decode=multi_token_decode,+ )+ return exported["decode"]+++def _compiled_ops(ov_model):+ """`(layer type, runtime precision)` of every op the CPU plugin compiled `ov_model` into."""+ import openvino++ # f32 inference: the fusions must not depend on the CPU having bf16 (AMX or AVX-512 BF16)+ compiled = openvino.Core().compile_model(ov_model, "CPU", {"INFERENCE_PRECISION_HINT": "f32"})+ return [+ (op.get_rt_info()["layerType"].astype(str), op.get_rt_info()["runtimePrecision"].astype(str))+ for op in compiled.get_runtime_model().get_ordered_ops()+ ]+++@slow+@require_torch+@require_openvino+class OpenVINOFusionTest(MemoryCleanupMixin, unittest.TestCase):+ """The kernels OpenVINO's CPU plugin compiles the decode graph into: the attention fused with its KV cache, RoPE+ and compressed-weight matmuls."""++ def assert_decode_fused(self, model, decode, fusions):+ num_layers = model.config.num_hidden_layers+ # The cache update feeds its state as is: anything between them keeps the plugin from fusing the cache into+ # its attention.+ assign_sources = [+ op.input_value(0).get_node().get_type_name()+ for op in decode.get_ordered_ops()+ if op.get_type_name() == "Assign"+ ]+ self.assertEqual(assign_sources, ["Concat"] * 2 * num_layers)+ layers = collections.Counter(layer for layer, _ in _compiled_ops(decode))+ self.assertEqual({op: layers[op] for op in fusions}, {op: n * num_layers for op, n in fusions.items()})++ # bf16, as checkpoints ship: MoE experts only fuse with 16-bit weights+ @parameterized.expand(list(TESTERS))+ def test_decode_fuses(self, model_type):+ model = _model(model_type, torch.bfloat16)+ fusions = OPENVINO_MOE_FUSIONS if model_type in MOE_MODELS else OPENVINO_FUSIONS+ self.assert_decode_fused(model, _decode(model), fusions)++ def test_float32_decode_fuses(self):+ model = _model("llama")+ self.assert_decode_fused(model, _decode(model), OPENVINO_FUSIONS)++ def test_single_token_decode_fuses(self):+ model = _model("llama", torch.bfloat16)+ self.assert_decode_fused(model, _decode(model, multi_token_decode=False), OPENVINO_FUSIONS)++ def test_16bit_model_exports_float32_graph_with_16bit_weights(self):+ import openvino++ model = _model("llama", torch.bfloat16)+ decode = _decode(model)+ self.assertTrue(all(p.dtype == torch.bfloat16 for p in model.parameters()), "the model keeps its dtype")+ self.assertTrue(all(output.get_element_type() == openvino.Type.f32 for output in decode.outputs))+ weights = [+ op+ for op in decode.get_ordered_ops()+ if op.get_type_name() == "Constant"+ and len(op.get_output_shape(0)) == 2+ and op.get_output_shape(0)[1] == model.config.hidden_size+ ]+ self.assertTrue(weights)+ self.assertTrue(all(op.get_element_type() == openvino.Type.bf16 for op in weights))+ self.assertIn("FullyConnected", {layer for layer, _ in _compiled_ops(decode)})++ @require_nncf+ def test_nncf_int8_weights_fuse_compressed_matmuls(self):+ import nncf++ from transformers.exporters import NNCFOpenVINOQuantizer, NNCFWeightCompressionConfig++ quantizer = NNCFOpenVINOQuantizer(NNCFWeightCompressionConfig(mode=nncf.CompressWeightsMode.INT8_ASYM))+ model = _model("llama", torch.bfloat16)+ decode = _decode(model, quantizer=quantizer)+ precisions = collections.Counter(+ precision for layer, precision in _compiled_ops(decode) if layer == "FullyConnected"+ )+ # q, k, v, o, gate, up and down per layer read u8 weights in the fused matmul instead of decompressing them+ self.assertGreaterEqual(precisions["u8"], 7 * model.config.num_hidden_layers)