🚨 Modeling changes for export, compile, and hybrid-attention standardization
Author: IlyasMoutawwakil
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46738)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46738)
<!-- ci-dashboard-badge:end -->

# Modeling changes for export, compile, and hybrid-attention standardization

A broad sweep of modeling-side changes that make a large set of models cleanly exportable (ONNX / torch.export / ExecuTorch), fullgraph-compilable, and consistent in how they declare their layer types and build masks/caches. Standalone from the exporter machinery PR — this is the *model code* side.

## What's in here

### Hybrid SSM / linear-attention standardization

- **Single `"linear_attention"` layer type** in `ALLOWED_LAYER_TYPES` (collapsed from the per-backbone variants iterated on during review). The recurrent backbone is determined by `config.model_type`, so per-layer suffixes were redundant.
- **`_LEGACY_LAYER_TYPE_REMAP`** (privatized) — `__post_init__` remaps legacy Hub config strings (`"mamba"` → `"linear_attention"`, `"attention"` → `"full_attention"`) transparently so existing checkpoints load unchanged.
- **`create_recurrent_padding_mask`** in `masking_utils.py` + `LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING` dispatch. `"linear_attention"` and `"conv"` (lfm2/lfm2_moe short-convolution layers — same recurrent cache shape, different mechanism) are both wired to the recurrent padding mask.
- **Standard `causal_mask_mapping[layer_type]` pattern** applied across `bamba`, `granitemoehybrid`, `nemotron_h`, `qwen3_5`, `qwen3_next`, `olmo_hybrid`, `jamba`, `falcon_h1`, `lfm2`, `lfm2_moe` — replaces the various ad-hoc `causal_mask` + `mamba_mask` / `linear_attn_mask` patterns with the gemma2-style dispatch. Lets `generate` pre-build a static-size mask dict via `create_masks_for_generate` and pass it straight through (avoids dynamo recompiles on growing 2D `attention_mask` under static-cache compile).

### Audio / legacy attention dispatch refactor

`moshi` / `mimi` / `kyutai_speech_to_text` / `nemotron` / `diffllama` now use the standard `ALL_ATTENTION_FUNCTIONS.get_interface(...)` dispatch with `_supports_sdpa` / `_supports_flash_attn` / `_supports_flex_attn` / `_supports_attention_backend` flags set, instead of bespoke per-backend `*FlashAttention2` / `*SdpaAttention` / `*ATTENTION_CLASSES` selectors.

- **`diffllama`** specifically: differential attention `(softmax(Q1·K1ᵀ) - λ·softmax(Q2·K2ᵀ))·V` is now realized as **two standard attention calls** sharing Q/K and using the two halves of V. The V-doubling shortcut (`V` of shape `(B, H, S, 2D)`) was slower in practice — asymmetric V trips PyTorch's SDPA backend selector (Flash/cuDNN require `head_dim_v == head_dim_q`) and forces the math kernel. Benchmarked: 2-call is ~30% faster on prefill / training / long-context. Docstring on `DiffLlamaAttention` explains the trap.

### Granite-MoE family refactor (`granitemoe` / `granitemoeshared` / `granitemoehybrid`)

- Split `*MoE` block into `*TopKRouter` (flat `nn.Parameter(num_experts, hidden_size)`, returns `top_k_index` + `top_k_weights`) + `@use_experts_implementation`-decorated `*Experts` with direct `gate_up_proj` / `down_proj` parameters.
- Dropped the JetMoe inheritance + `expert_size = expert_size.tolist()` data-dependent path that blocked fullgraph compile.
- Default `experts_implementation` is `grouped_mm` / `batched_mm` (compilable); eager fallback for users who explicitly opt in.
- **Weight conversion entry** in `conversion_mapping.py` remaps legacy `block_sparse_moe.input_linear.weight` / `output_linear.weight` / `router.layer.weight` → new `experts.gate_up_proj` / `experts.down_proj` / `router.weight`. Existing IBM Hub checkpoints (`ibm/PowerMoE-3b`, `ibm-granite/granite-4.0-h-tiny`, …) load with zero user action.

### Fullgraph-compile opt-ins

`_can_compile_fullgraph = True` on models whose codepath was verified clean: `mamba2`, `granitemoe`, `granitemoeshared`, `granitemoehybrid`, `lfm2`, `lfm2_moe`, on top of the original `bamba` / `jamba` / `nemotron_h` / `qwen3_5` / `qwen3_5_moe` / `qwen3_next` from the SSM-compile work. The `is_linear_attention_model` tolerance-relaxation in `test_generate_with_static_cache` was verified unnecessary and dropped along with the helper. `falcon_h1` / `zamba2` / `olmo_hybrid` remain `False` with documented blockers in code.

### Export-friendly modeling rewrites

The bulk of the diff is rewriting modeling code that historically blocked export:

- **Inlined helper functions** that wrapped trivial ops behind reshape gymnastics — `big_bird` / `bigbird_pegasus` `torch_bmm_nd` / `torch_bmm_nd_transpose` removed entirely (`torch.matmul` handles arbitrary batch dims natively), 14 call sites inlined.
- **Vectorized window-attention masks** — `swin` / `swinv2` / `sew` / `sew_d` / `mlcd` / `donut_swin` / `clap` / `maskformer_swin` cyclic-shift masks rewritten from Python-loop-over-slice-tuples to `arange`-based vectorized construction so the graph is static.
- **5D → 4D attention** in `granite_speech` / `granite_speech_plus` conformer block (eliminates a special branch in the ONNX exporter's SDPA patch).
- **`@can_return_tuple`** refactor in `patchtsmixer` — drops the manual `return_dict` plumbing in 6 sub-classes' forwards.
- **Vectorized BOS/EOS padding** in `clvp` (replaces Python loop with `gather` + `masked_fill`).
- **Vectorized `ngram_attention_bias`** in `prophetnet`.
- **`sew` / `sew_d` attention mask pooling** — replaces the `torch.arange(max_encoder_length)` indexing (which bakes seq length into the exported ONNX graph) with `max_pool1d`.
- **Per-model ONNX-blocker fixes**: explicit slice stops (avoids `Unsqueeze(None,…)` translations), `torch_compilable_check` standardization in `reformer`, NLLB classifier cast paths, granite_speech masked-fill via strided slice view, audio-model rewrites in `clap` / `clvp` / `dia`, vision rewrites in `flava` / `git` / `mlcd` / `sam2` / `vitdet` / `efficientloftr` / `videoprism` / `paddleocr_vl` / `minicpmv4_6`, etc.

### Vision utilities

`vision_utils.py` gains `get_image_patch_grid_thw` and related helpers used by MiniCPMV4.6 (and other NaViT-packed vision encoders) to precompute data-dependent vision tensors before export.

### Test infrastructure

- **`tests/generation/test_utils.py`**:
  - `test_generate_compile_model_forward_fullgraph`: skip the "default cache is non-compileable" assertion for models whose default `DynamicCache` is compileable by construction (fully-recurrent / hybrid SSM — `LinearAttentionLayer` is `is_compileable=True`).
  - `test_generate_compilation_all_outputs`: gate `output_attentions` on `self.has_attentions` (Mamba2 doesn't define `attentions` on its output).
  - `_check_attentions_for_generate`: for hybrid static caches, look up `get_max_cache_shape()` on the first layer that has a real max length — layer 0 may be a linear-attention layer that returns `-1`.
- **`tests/utils/test_masking_utils.py::test_create_masks_for_generate_defers_for_unmapped_layer_types`**: switched the canonical "unmapped" example from `"linear_attention"` (now mapped to `create_recurrent_padding_mask`) to `"moe"` (used by Nemotron-H and intentionally not mapped — MoE is a sub-layer composition, not an attention layer).
- **`tests/models/nemotron_h/`**: 4 hardcoded legacy-name assertions updated to the standardized names.
- **`tests/models/granitemoehybrid/`**: `test_attention_only_forward` — direct post-construction `config.layers_block_type = ["attention"] * N` was bypassing the remap; updated to `["full_attention"] * N`.
- **`tests/models/blt/`**: 4 new skips for BLT's hash-embedding / EncoderDecoderCache cross-attention limitations (in line with the existing two `inputs_embeds` skips). Tests: `test_generate_continue_from_inputs_embeds`, `test_assisted_decoding_matches_greedy_search_{0_random,1_same}`, `test_assisted_decoding_sample`.

### Small review-comment fixes

- Explicit `== "linear_attention"` everywhere instead of `.startswith("linear_attention")`.
- `nemotron_h` mtp default `["attention", "moe"]` → `["full_attention", "moe"]`.
- `smolvlm` uses `@merge_with_config_defaults` instead of `use_cache = use_cache if … else self.config.use_cache`.
- `tests/generation/test_utils.py::is_hybrid_mamba_model` → `is_linear_attention_model` (then removed entirely after the tolerance investigation showed the predicate wasn't needed).
- Removed stale `max_cache_len` comment in `bamba`.
- Updated `intermediate_size=37` → `40` in granitemoe / granitemoeshared test configs for `torch._grouped_mm` 16-byte stride alignment.
- `videoprism` modular: `+ 1` → `+ 1.0` (avoid `prim.device` graph node).

## Stats

147 files changed (`+2292 / −3567`) — net code reduction despite adding new infrastructure, thanks to dropped helpers and inheritance simplifications.

## Test commands run

```bash
# Hybrid mask / config / forward tests across the touched families
pytest tests/models/{bamba,jamba,granitemoehybrid,nemotron_h,qwen3_5,qwen3_5_moe,qwen3_next,olmo_hybrid,minimax,falcon_h1,mamba2,lfm2,lfm2_moe}/

# Audio / legacy attention dispatch refactor
pytest tests/models/{moshi,mimi,kyutai_speech_to_text,nemotron,diffllama}/

# Granite-MoE integration on real Hub checkpoints
RUN_SLOW=1 pytest tests/models/granitemoe/::GraniteMoeIntegrationTest \
                  tests/models/granitemoeshared/::GraniteMoeSharedIntegrationTest \
                  tests/models/granitemoehybrid/::GraniteMoeHybridIntegrationTest
# → 6/6 passed in 164s on ibm/PowerMoE-3b + ibm-granite/granite-4.0-h-tiny

# Static-cache equivalence (the test that gates _can_compile_fullgraph)
pytest -k test_generate_with_static_cache tests/models/{bamba,jamba,nemotron_h,qwen3_5,qwen3_5_moe,qwen3_next,mamba2,lfm2,lfm2_moe,granitemoe,granitemoeshared,granitemoehybrid}/

# Fullgraph-compile under static cache, no recompiles allowed
pytest -k test_generate_compile_model_forward_fullgraph tests/models/{mamba2,lfm2,lfm2_moe,bamba,nemotron_h,qwen3_next}/

# Patchtsmixer return_dict refactor
pytest tests/models/patchtsmixer/

# Granite-speech 4D attention refactor + ONNX export
pytest tests/models/granite_speech/ tests/models/granite_speech_plus/ -k "test_onnx_export or test_model_forward"

# Big_bird / bigbird_pegasus matmul inlining
pytest tests/models/big_bird/ tests/models/bigbird_pegasus/

# Swin family vectorized masks
pytest tests/models/{swin,swinv2,donut,mlcd,sam2,clap,maskformer}/

# Generation + masking infra tests
pytest tests/generation/ tests/utils/test_masking_utils.py
```

## Notes for review

- The `"linear_attention"` collapse means existing Hub configs with `["mamba", ...]` or `["linear_attention_*", ...]` (post-rename) all map to the same name. Legacy remap is unconditional in the model `__post_init__` — no opt-in flag.
- `"conv"` is kept as a distinct layer type for lfm2 / lfm2_moe (not folded into `"linear_attention"`) — `Lfm2ShortConv` is a 1D causal convolution, not an attention mechanism. It happens to share the recurrent cache shape with linear attention, same as `"moe"` / `"mlp"` in Nemotron-H share the cache shape without being attention layers.
- Granite-MoE weight rename goes through the `conversion_mapping.py` API. No `_keys_to_ignore_on_load_missing` hacks.
- Vectorized window-mask rewrites in `swin` / `swinv2` / `sew` / etc. are functionally equivalent (verified by existing tests + numerical-equivalence check on saved outputs).
- DiffLlama's two-call differential attention is bit-equivalent to the V-doubled form (linearity of softmax-weighted V combination). Verified with eager equivalence tests.

🤖 AI-assisted; reviewed and tested end-to-end.

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46738)
**Latest run:** [28429203342:1](https://github.com/huggingface/transformers/actions/runs/28429203342)
**Result:** `success` | **Jobs:** 14 | **Tests:** 73,495 | **Failures:** 0 | **Duration:** 7h 17m

<!-- ci-dashboard-recap:end -->

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46738). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

run-slow: big_bird, bigbird_pegasus, canine, clap, clvp, donut_swin, maskformer_swin, minicpmv4_6, paddleocr_vl, patchtsmixer, prophetnet, reformer, sam2, sew, swin, video_llava

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27813252959)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/big_bird", "models/bigbird_pegasus", "models/canine", "models/clap", "models/clvp", "models/minicpmv4_6", "models/paddleocr_vl", "models/patchtsmixer", "models/prophetnet", "models/reformer", "models/sam2", "models/sew", "models/swin", "models/video_llava"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27813252959)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [b830d13c](https://github.com/huggingface/transformers/commit/b830d13c0bb501296f306620091c692609e2959c) | workflow commit (merge commit) |
| PR | [ff33d8b3](https://github.com/huggingface/transformers/commit/ff33d8b32d58a2f6d30ba0884b995ab0af8ba8b6) | branch commit (from PR) |
| main | [b4f9cd23](https://github.com/huggingface/transformers/commit/b4f9cd238e74deedac274e0a5b01c33f3a4b49a0) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


slow tests are good, there are some relevant failures in fast ci, fixing them now

We also need to rerun slow ci https://huggingface.slack.com/archives/C01NE71C4F7/p1781877684393659 - it was broken :cry: 

run_slow: bamba, qwen3_5, big_bird, big_bird_pegasus, canine, clap, clvp, dia, difflama, donut_swin, edgetam_video, falcon_h1, flava, granite_speech, granite_speech_plus, granitemoe, granitemoehybrid, jamba, kyutai_speech_to_text, lfm2, maskformer_swin, mimi, minicpm4_6, minimax, mlcd, moshi, nemotron_h, paddleocr_vl, prophetnet, reformer, sam2, sew, sew_d, splinter, swin, timesfm, video_llava

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27960146503)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/bamba", "models/big_bird", "models/canine", "models/clap", "models/clvp", "models/dia", "models/edgetam_video", "models/falcon_h1", "models/flava", "models/granite_speech", "models/granite_speech_plus", "models/granitemoe", "models/granitemoehybrid", "models/jamba", "models/kyutai_speech_to_text", "models/lfm2", "models/mimi", "models/minimax", "models/mlcd", "models/moshi", "models/nemotron_h", "models/paddleocr_vl", "models/prophetnet", "models/qwen3_5", "models/reformer", "models/sam2", "models/sew", "models/sew_d", "models/splinter", "models/swin", "models/timesfm", "models/video_llava"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27960146503)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [bf41f2f9](https://github.com/huggingface/transformers/commit/bf41f2f97b451b0a283c6a0288e646c3a9bb1414) | workflow commit (merge commit) |
| PR | [5e7bdeee](https://github.com/huggingface/transformers/commit/5e7bdeeef5471ee03641f430b0081c3a3f7d60a0) | branch commit (from PR) |
| main | [f007617e](https://github.com/huggingface/transformers/commit/f007617ef2d9ad9bde89dce2e54deab17656ef21) | base commit (on `main`) |

⚠️ No test being reported (jobs are skipped or cancelled)!


run_slow: bamba, qwen3_5, big_bird, big_bird_pegasus, canine, clap, clvp, dia, difflama, donut_swin, edgetam_video, falcon_h1, flava, granite_speech, granite_speech_plus, granitemoe, granitemoehybrid, jamba, kyutai_speech_to_text, lfm2, maskformer_swin, mimi, minicpm4_6, minimax, mlcd, moshi, nemotron_h, paddleocr_vl, prophetnet, reformer, sam2, sew, sew_d, splinter, swin, timesfm, video_llava

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27961220494)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/bamba", "models/big_bird", "models/canine", "models/clap", "models/clvp", "models/dia", "models/edgetam_video", "models/falcon_h1", "models/flava", "models/granite_speech", "models/granite_speech_plus", "models/granitemoe", "models/granitemoehybrid", "models/jamba", "models/kyutai_speech_to_text", "models/lfm2", "models/mimi", "models/minimax", "models/mlcd", "models/moshi", "models/nemotron_h", "models/paddleocr_vl", "models/prophetnet", "models/qwen3_5", "models/reformer", "models/sam2", "models/sew", "models/sew_d", "models/splinter", "models/swin", "models/timesfm", "models/video_llava"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27961220494)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [dbf5a05c](https://github.com/huggingface/transformers/commit/dbf5a05c0ab8db363039489c61d0c210da2b47ed) | workflow commit (merge commit) |
| PR | [707df61d](https://github.com/huggingface/transformers/commit/707df61d7a3fcd07a1785d6bbf3d5441a9d80739) | branch commit (from PR) |
| main | [f007617e](https://github.com/huggingface/transformers/commit/f007617ef2d9ad9bde89dce2e54deab17656ef21) | base commit (on `main`) |

### Model CI Report

❌ **[7 new failed tests from this PR](https://huggingface.co/datasets/hf-internal-testing/transformers_pr_ci/raw/0f3561ef0e8261e0439502af38a9e1606d052b6a/2026-06-22/runs/36708-27961220494/ci_results_run_models_gpu/new_failures_with_bad_commit.json)** 😭

- [canine](https://github.com/huggingface/transformers/actions/runs/27961220494/job/82744530074):
    tests/models/canine/test_modeling_canine.py::CanineModelIntegrationTest::test_inference_no_head (✅ ⟹ ❌)

- [moshi](https://github.com/huggingface/transformers/actions/runs/27961220494/job/82744530250):
    tests/models/moshi/test_modeling_moshi.py::MoshiTest::test_eager_matches_sdpa_generate (✅ ⟹ ❌)
    tests/models/moshi/test_modeling_moshi.py::MoshiTest::test_flex_attention_with_grads (✅ ⟹ ❌)
    tests/models/moshi/test_modeling_moshi.py::MoshiIntegrationTests::test_moshika_greedy_unconditional_fp16 (❌ ⟹ ❌)
    tests/models/moshi/test_modeling_moshi.py::MoshiIntegrationTests::test_moshiko_greedy_unconditional_fp16 (❌ ⟹ ❌)
    tests/models/moshi/test_modeling_moshi.py::MoshiIntegrationTests::test_moshiko_greedy_unconditional_fp16_eager (❌ ⟹ ❌)

- [nemotron_h](https://github.com/huggingface/transformers/actions/runs/27961220494/job/82744530257):
    tests/models/nemotron_h/test_modeling_nemotron_h.py::NemotronHModelTest::test_generate_compile_model_forward_fullgraph (✅ ⟹ ❌)



run_slow: bamba, qwen3_5, big_bird, big_bird_pegasus, canine, clap, clvp, dia, difflama, donut_swin, edgetam_video, falcon_h1, flava, granite_speech, granite_speech_plus, granitemoe, granitemoehybrid, jamba, kyutai_speech_to_text, lfm2, maskformer_swin, mimi, minicpm4_6, minimax, mlcd, moshi, nemotron_h, paddleocr_vl, prophetnet, reformer, sam2, sew, sew_d, splinter, swin, timesfm, video_llava

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27983840724)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/bamba", "models/big_bird", "models/canine", "models/clap", "models/clvp", "models/dia", "models/edgetam_video", "models/falcon_h1", "models/flava", "models/granite_speech", "models/granite_speech_plus", "models/granitemoe", "models/granitemoehybrid", "models/jamba", "models/kyutai_speech_to_text", "models/lfm2", "models/mimi", "models/minimax", "models/mlcd", "models/moshi", "models/nemotron_h", "models/paddleocr_vl", "models/prophetnet", "models/qwen3_5", "models/reformer", "models/sam2", "models/sew", "models/sew_d", "models/splinter", "models/swin", "models/timesfm", "models/video_llava"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/27983840724)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [1d36f3da](https://github.com/huggingface/transformers/commit/1d36f3daeddb8c4ccf40723fb63401c45293f8d6) | workflow commit (merge commit) |
| PR | [25873ea2](https://github.com/huggingface/transformers/commit/25873ea28384d5f8474b8fabf9abab6c97316487) | branch commit (from PR) |
| main | [8fea7d39](https://github.com/huggingface/transformers/commit/8fea7d39b790cee717732069c89bba458cfddcdb) | base commit (on `main`) |

### Model CI Report

❌ **[3 new failed tests from this PR](https://huggingface.co/datasets/hf-internal-testing/transformers_pr_ci/raw/42a80d620fdde08b6557246f99665376d55dd899/2026-06-22/runs/36722-27983840724/ci_results_run_models_gpu/new_failures_with_bad_commit.json)** 😭

- [moshi](https://github.com/huggingface/transformers/actions/runs/27983840724/job/82825749506):
    tests/models/moshi/test_modeling_moshi.py::MoshiIntegrationTests::test_moshika_greedy_unconditional_fp16 (❌ ⟹ ❌)
    tests/models/moshi/test_modeling_moshi.py::MoshiIntegrationTests::test_moshiko_greedy_unconditional_fp16 (❌ ⟹ ❌)
    tests/models/moshi/test_modeling_moshi.py::MoshiIntegrationTests::test_moshiko_greedy_unconditional_fp16_eager (❌ ⟹ ❌)



moshi seems to have some failing integration tests on main, my changes only changed the failing line's number

yeah i was actually split between a class attr and an instance attr, standardizing

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: afmoe, bamba, bart, big_bird, bigbird_pegasus, blt, bros, canine, clap, clvp, cohere_asr, dia, diffllama, diffusion_gemma, donut, dpt