transformers

72b487d1 - OpenVINO exporter: let OpenVINO fuse attention, RoPE, the KV cache and compressed matmuls (#50311)

OpenVINO exporter: let OpenVINO fuse attention, RoPE, the KV cache and compressed matmuls (#50311) * OpenVINO exporter: repeat GQA heads with `repeat_kv` so OpenVINO fuses its KV-cache attention The exporter swapped `repeat_kv` for `repeat_interleave`, and its SDPA patch repeated K/V the same way. `repeat_interleave` lowers to a `Gather` between the stateful KV cache and the attention, which OpenVINO's CPU plugin can't fuse: every attention ran decomposed into MatMul/Softmax and copied the cache each step. Keep transformers' `repeat_kv` (unsqueeze, expand, reshape) in both places. The frozen KV length the `repeat_kv` patch worked around no longer reproduces on OpenVINO 2026.4: Qwen2.5-0.5B and Qwen3-0.6B decode a left-padded batch of two, 48 tokens past the prompt, token-for-token like eager. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: drop the SDPA patch's fully-masked-row zeroing The patch zeroed the attention output of rows that mask every key, with a select on a mask-derived condition after each SDPA. OpenVINO's CPU plugin miscomputes that pattern: EmbeddingGemma2's fp32 export returned NaN (layer 2 onwards) with no row fully masked. Without it the fp32 export matches eager (token cosine 1.0000), and a fully masked row gets the uniform average its dtype-minimum mask describes, as transformers' eager attention gives it. Generated tokens are unaffected: Qwen2.5-0.5B and Qwen3-0.6B still decode a left-padded batch of two like eager. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: trace bf16/fp16 models as f32 graphs OpenVINO'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: every weight was decompressed on each call. Trace and convert in f32, then give each parameter and buffer, and each config, its own dtype back. With an f32 inference hint the f32 graph matches eager f32 on the same weights token for token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: give KV-cache variables the update's dims instead of reshaping the update A batch-1 sample specializes the batch of the cache concat to 1 while its variable stays batch-dynamic, so `_pin_state_update_shapes` put a Reshape between the concat and its `Assign`. That Reshape keeps the CPU plugin from fusing the cache into its stateful SDPA. The variable now takes the dims the update fixes, and a pinning Reshape is only added where the variable fixes a dim the update leaves dynamic. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: fold the frontend's marked scalar casts so RoPE fuses The PyTorch frontend translates `aten.neg` as a multiply by a `Convert`ed -1, with the Convert marked DisableConstantFolding and the multiply marked as a dequantization node. Until those marks are dropped, which only happens on serialization, the CPU plugin's RoPE fusion doesn't match `rotate_half`. Fold such scalar Converts into constants and drop the marks on their consumers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: store a 16-bit model's weights as 16-bit constants plus a Convert After the f32 trace, restore bf16/fp16 weights as 16-bit constants followed by a Convert to f32. The CPU plugin decompresses this pattern inside its FullyConnected, and the IR keeps the checkpoint's exact values at their original size. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test that OpenVINO fuses an exported decoder's attention, RoPE, KV cache and compressed matmuls Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: give the 16-bit model helpers their own section Moves the f32 trace and the 16-bit weight restore into a "16-bit models" section and folds the scalar casts inside `_convert_to_openvino`, so `export()` reads as on main. `_pin_state_update_shapes` handles the variable refinement inline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Move the OpenVINO fusion tests to test_openvino.py and run them on every decoder family Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: keep linear and embedding weights 16-bit during the f32 trace, as optimum-intel does Upcasting every weight in place doubled the model's memory during export (5.4 GiB peak for a 1.1 GiB bf16 Qwen3-0.6B, against 1.6 GiB on main) and needed a pass restoring 16-bit constants afterwards. Like OV's `__make_16bit_traceable`, linear and embedding weights now stay 16-bit behind zero-cost f32 placeholders, and patched `nn.Linear` / `nn.Embedding` forwards cast them to f32 inside the graph: they land as a 16-bit constant and a `Convert` directly, with no copy and no restore pass. The embedding casts before the lookup, so a tied `lm_head` stays a compressed-weight matmul after NNCF quantization. `aten.neg` is traced as a multiply by -1 instead of folding the frontend's marked scalar casts after conversion, which gives OV's RoPE fusion the same pattern. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Mark the OpenVINO fusion tests slow, like the other export tests Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: trace MoE experts as a tiled block so OpenVINO fuses them into GatherMatmuls Every `@use_experts_implementation` experts module is traced with a forward that runs every expert on every token and weights the outputs by the dense routing weights, the block OpenVINO's CPU plugin rewrites into `GatherMatmul`s that only compute the routed experts (`ConvertTiledMoeBlockTo{2,3}GatherMatmuls`). Its expert weights stay 16-bit through the f32 trace, like linear weights: the pass only matches 16-bit weights with an in-graph `Convert`, and upcasting them doubled a 16-bit MoE model's expert memory during export. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Declare each OpenVINO fusion test model's expected fused kernels in MODELS Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the OpenVINO fusion test models' config overrides: only Phi-3 needs a pad token inside the tiny vocabulary Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Build the OpenVINO fusion test models from each model's own tester Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: give sliding-window cache layers a full cache so their attention fuses A `DynamicSlidingWindowLayer` slices its cache to the window on every update, which puts a `Slice` between the KV-cache concat and its state and keeps OpenVINO's CPU plugin from fusing the cache into its attention (Mistral, Gemma3's sliding layers). Under the OpenVINO trace these layers take `DynamicLayer`'s update and mask sizes: the sliding mask still applies the window on absolute positions, and their cache grows with the sequence. Mistral and Gemma3 join the fusion tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: support OpenVINO 2026.5, which no longer decomposes exported programs OpenVINO 2026.5 converts exported programs without `run_decompositions` and removes `openvino.frontend.pytorch.torchdynamo.export_decompositions`, whose `ops_to_not_decompose` the exporter imports at module level, so the exporter failed to import. The import is now optional: with OpenVINO < 2026.5 the exporter decomposes with OpenVINO's table as before, and with 2026.5 it converts the exported program as is, like OpenVINO's own frontend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: enter the 16-bit model patch with the other export patches `_traced_in_float32` becomes the public `patch_16bit_model`, entered in the same `with` as `apply_patches("openvino")` and `patch_model_outputs`, with the trace and the conversion inside. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: give sliding-window cache layers a full cache on every backend A `DynamicSlidingWindowLayer` sizes its mask from a Python token counter and slices its cache to the window, so a trace fixes both at the sample's length and the exported decode fails past the window: the ONNX decode of a tiny Mistral with a 4-token window errors on its first step. Moving the full-cache patch from the OpenVINO patches to the shared dynamo ones fixes that for every backend (16 tokens past the window equal to eager with ONNX Runtime), and OpenVINO still gets the cache shape it fuses. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Rename the OpenVINO fusion tests to tests/exporters/test_fusions.py One file for every backend's fusion tests, one class each, the way test_export.py and test_quantization.py hold every backend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Describe tests/exporters/test_fusions.py as every backend's fusion tests Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Share the fusion tests' model table across backends `TESTERS` and `_model` are the same in the ONNX Runtime fusion tests, so both backends' classes run on the same tiny models; OpenVINO's expected kernels move to `OPENVINO_FUSIONS` / `OPENVINO_MOE_FUSIONS`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Cache: let sliding-window layers of a DynamicCache keep every token `DynamicCache(full_sliding_layers=True)` builds sliding-window and chunked attention layers as full ones: the attention mask still applies the window, so outputs are unchanged, and the cache grows with the sequence. Some runtimes need the cache update to feed their state as is to fuse it into attention (OpenVINO, like optimum-intel's exports). `generate` passes `GenerationConfig.cache_config` to the `DynamicCache` it builds, as the field documents ("arguments used in the key-value cache class"), so `GenerationConfig(cache_config={"full_sliding_layers": True})` sets it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO: keep sliding-window caches full through the generation config The `"dynamo"` patch that gave every `DynamicSlidingWindowLayer` a `DynamicLayer`'s methods during the trace is gone: it changed every backend's exports, and nothing in the generation config said so, while the ONNX, ExecuTorch and torch.export runtimes rebuild the cache from it. `OpenVINOExporter.export_for_generation` now captures a stateful export with `GenerationConfig(cache_config={"full_sliding_layers": True})` when the generation config sets no `cache_config`: OpenVINO's CPU plugin only fuses a KV cache into its attention when the update feeds the state as is, and a stateful model keeps its cache inside, so nothing rebuilds it at runtime. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Revert "OpenVINO: keep sliding-window caches full through the generation config" This reverts commit f2ca057867216fcc325152cd3e95df327ac7b2c7. * Revert "Cache: let sliding-window layers of a DynamicCache keep every token" This reverts commit 18ba237cf591184826d8a93d4c8d61bd1fac658d. * Note that the sliding-window export patch goes with #49542 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: address review - `is_openvino_greater_or_equal` gates the 2026.5 decomposition change instead of an ImportError. - `patch_16bit_model`: loops, a shorter docstring, `_16bit_parameter` (it returns biases too). - `_configs` walks the pretrained models only. - `_tiled_experts_forward` follows `batched_mm_experts_forward`'s layout and names, with a `projection` helper. - `_pin_state_update_shapes`: one early exit per case. - The sliding-window patch's docstring is shorter. - The fusion tests build their models in the dtype (`from_config(dtype=...)`), not cast after, so float32 buffers stay float32. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: OpenVINO 16-bit exports, the full sliding-window cache, and OpenVINO 2026.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: free the OpenVINO fusion tests' memory with MemoryCleanupMixin Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: keep the OpenVINO notes to one line each Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: drop the OpenVINO 16-bit note Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: read each OpenVINO output under the name the export gave it Fused RoPE merges the replaced nodes' tensor names (`add_4:0`, numeric ids) onto the key-cache outputs of partial-rotary models (gpt_neox, glm4_moe, bamba, ...): the first non-numeric name was no longer the declared one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: convert `torch.mm(..., out_dtype=...)` OV has no conversion rule for `aten.mm.dtype` (Kolibri1's router on CUDA): rewrite it as an `mm` on its inputs cast to `out_dtype`, as the ONNX exporter translates it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OpenVINO exporter: keep a zero-width cache state static and growing Mistral Large 4's NoPE layers cache a rotary part of width 0. The state init spec used `0` both for the growing dim and for such a static zero-size dim, so the width went dynamic and `ReadValue` rejected the variable; the growing dim is now marked `"grow"`. `_fix_empty_cat` also dropped the decode cat of two empty caches, so the state never grew: it now keeps a cat whose operands are both empty with the same rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: skip sliding-window caches in export comparisons, exactness for TimesFM and PP-DocLayoutV4 Exports keep a full cache on sliding-window layers (the sliding mask still applies the window) while eager keeps only the window, so their cache outputs aren't compared. TimesFM's left-padding patches attend to no key and PP-DocLayoutV4's tiny config ties every `topk` score, so neither has a defined value to compare on OpenVINO. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>