Fix Mamba2 chunked-prefill / speculative decoding for Zamba2, Nemotron-H, Bamba, FalconH1 and GraniteMoeHybrid
Author: Sunt-ing
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46741)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46741)
<!-- ci-dashboard-badge:end -->

# What does this PR do?

The Mamba2 mixer that Zamba2, Nemotron-H, Bamba, FalconH1 and GraniteMoeHybrid each carry decided whether to take the single-step decode path from whether the cache held a previous state, without checking `seq_len`:

```python
use_precomputed_state = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
```

So any forward with `seq_len > 1` and a non-empty cache (chunked prefill, or speculative / assisted decoding verifying several draft tokens at once) was routed into the single-token branch: the conv ran a single `update_conv_state` step, `dt = dt[:, 0, :]` kept only the first token, and the SSM did one recurrent step. Every token after the first came out wrong and the cached conv/recurrent state was corrupted, silently degrading the rest of generation. Bamba / FalconH1 / GraniteMoeHybrid already gated the single-step branch on `seq_len == 1`, but their multi-token cached branch still started from a zero conv context and zero SSM state, so chunked-prefill was history-less there too.

This is the same bug and fix as mamba2 (#46032, fixed in #46084), but #46084 only touches `modeling_mamba2.py`; these models carry their own Mamba2 mixer. The fix gates the single-step branch on `seq_len == 1`, and for the multi-token cached case runs the full path seeded from the cache: prepend the cached conv left-context (dropped after the conv), and pass the cached recurrent state as the SSM initial state. Both `torch_forward` and `cuda_kernels_forward` are fixed. Nemotron-H's `NemotronHMamba2Mixer` subclasses `Zamba2MambaMixer`, and GraniteMoeHybrid reuses Bamba's mixer, so the fix lives in the `modular_*.py` sources and `make fix-repo` propagates it. Jamba is left out: it uses the Mamba1 path, not this Mamba2 chunk-scan path.

Each model gets a chunked-prefill regression test adapted from the one merged in #46084: a model-level first-token causal check (the first position of a multi-token cached forward must equal the single-token cached forward), with `_cpu` and `@require_torch_accelerator @require_kernels` `_torch_device` variants. Each fails on `main` and passes here. Nemotron-H has no CPU variant because its MoE layers use `grouped_mm`, which has no CPU kernel; its torch path is the same code as Zamba2's via the modular inheritance and is covered by the Zamba2 CPU test.

<details>
<summary>fp32 reproduction (last-token logits, full forward vs prefill(1) + chunk(S-1))</summary>

Built each model from its own `ModelTester.get_config()` (legitimate hybrid configs, no hand-rolled configs), float32 to rule out precision, comparing per-position last-hidden / logits of a full forward against a 1-token prefill followed by a single `seq_len = S-1` cached chunk. `all-1` (token-by-token incremental decoding) is the negative control, since it exercises the genuine single-step path.

Environment: `transformers` 5.13.0.dev0, `torch` 2.8.0+cu128, Python 3.12.3, Linux, single RTX 4090 (Zamba2 / Bamba / FalconH1 on CPU; Nemotron-H / GraniteMoeHybrid on GPU because of their MoE layers).

| model | chunked-prefill per-position max abs diff (before → after) | all-1 control |
|---|---|---|
| Zamba2 | 1.2e-1 → 5.96e-8 | 3.7e-8 (unchanged) |
| Nemotron-H | 3.726e-1 → 1.19e-7 | 8.2e-8 (unchanged) |

```python
import torch
from transformers import AutoModelForCausalLM
# Zamba2 example (CPU); Nemotron-H is identical but needs a GPU because of its MoE layers.
from tests.models.zamba2.test_modeling_zamba2 import Zamba2ModelTester

class P:  # ModelTester only needs a parent for asserts
    def __getattr__(self, k): return lambda *a, **kw: None

cfg = Zamba2ModelTester(P()).get_config()
m = AutoModelForCausalLM.from_config(cfg).to(torch.float32).eval()
ids = torch.randint(5, 100, (1, 20))
with torch.no_grad():
    ref = m(ids, use_cache=False).logits[0]
    o1 = m(ids[:, :1], use_cache=True)
    o2 = m(ids[:, 1:], past_key_values=o1.past_key_values, use_cache=True)
print((ref[1:] - o2.logits[0]).abs().max())   # before: ~1e-1, after: ~1e-7
```

</details>

<details>
<summary>per-model regression test, fails on main, passes here (position-0 max abs diff)</summary>

| model | no fix → with fix | path |
|---|---|---|
| Zamba2 | 0.32 → < 1e-4 | CPU |
| Nemotron-H | 1.92 → < 1e-4 | GPU |
| Bamba | 0.50 → < 1e-4 | CPU |
| FalconH1 | 0.013 → < 1e-4 | CPU |
| GraniteMoeHybrid | 1.67 → < 1e-4 | CPU |

Full files green: `tests/models/{zamba2,bamba,falcon_h1,granitemoehybrid}/test_modeling_*.py` 506 passed on CPU, `tests/models/nemotron_h/test_modeling_nemotron_h.py` 145 passed on GPU, both including `test_prompt_lookup_decoding_matches_greedy_search`. The CUDA fast path (`kernels`) was verified separately on a 4090: before, the multi-token cached forward crashes in `causal_conv1d_update`; after, cuda chunked-prefill matches a single full forward exactly.

</details>

## Code Agent Policy

- [x] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request), Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? See #46032.
- [ ] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?

## Who can review?

@vasqu

<!-- 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=46741)
**Latest run:** [28296890636:2](https://github.com/huggingface/transformers/actions/runs/28378208509)
**Result:** `success` | **Jobs:** 5 | **Tests:** 1,501 | **Failures:** 0 | **Duration:** 5m 37s

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

Comments:
Hey, Thanks for picking up on the other models. in my PR vasqu asked me to prepend the full conv_state (torch.cat([conv_state, hidden_states], dim=-1), matching the qwen3-next as reference. Refer mine to keep with consistency.

@Ramshankar07 Thanks, switched both mixers to prepend the full cached `conv_states` to match #46084 and qwen3-next. fp32 chunked-prefill check stays at ~1e-7 for both Zamba2 and Nemotron-H.

@vasqu Done. Fixed `cuda_kernels_forward` for both mixers and verified on a 4090.

**Fix** (same shape as the torch path, mirrors #46084): gate the single-step branch on `seq_len == 1`, and for the multi-token cached case (chunked prefill / speculative verify) run the full kernels seeded from the cache — prepend the cached conv left-context and drop it after the conv, and pass the cached recurrent state as `initial_states` to `mamba_chunk_scan_combined`.

**Verified** on a 4090 (`pip install kernels`, fast path confirmed taken):
- before: the multi-token cached forward crashes in `causal_conv1d_update` (`weight must have shape (dim, width)`)
- after: cuda chunked-prefill and token-by-token both match a single full forward exactly (max abs diff `0.0`)

Two things I noticed checking the rest of the family on `main`:
- bamba / falcon_h1 / granitemoehybrid only gate the single-step branch in both paths — their multi-token cached branch still starts from a zero state (no conv prepend, no `initial_states`), so their chunked-prefill is history-less too. Happy to do that as a coordinated follow-up.
- `kernels-community/mamba-ssm` currently fails to import on `main` (it does `from transformers.generation import GreedySearchDecoderOnlyOutput`, which has been removed), so the mamba fast path is effectively dormant on `main` until that build is updated. I worked around it locally to run the verification.


@Sunt-ing Looks pretty solid now! I would first wait on #46084 (which will be merged in a sec) and we adapt the test from there then

1. Yes we should propogate to all mamba2 related models imo so bamba etc
2. Can you share your stack trace and torch / kernels version. Kernels community should no longer do that tbh so I'm a bit surprised. I know that the main mamba-ssm package might have done that on earlier versions but it was also patched there on main iirc

Thanks @vasqu, all three addressed in two follow-up commits.

**Tests adapted from #46084.** Now that #46084 is merged, I reshaped the regression test to its final form: a model-level first-token causal check (the first position of a multi-token cached forward must equal the single-token cached forward), with separate `_cpu` and `@require_torch_accelerator @require_kernels` `_torch_device` variants. Each one fails on `main` (without the fix) and passes with it:

| model | max abs diff at position 0, no fix → with fix |
|---|---|
| Zamba2 (CPU) | 0.32 → < 1e-4 |
| Nemotron-H (GPU) | 1.92 → < 1e-4 |
| Bamba (CPU) | 0.50 → < 1e-4 |
| FalconH1 (CPU) | 0.013 → < 1e-4 |
| GraniteMoeHybrid (CPU) | 1.67 → < 1e-4 |

Nemotron-H has no `_cpu` variant: its MoE layers use `grouped_mm`, which has no CPU kernel, so that test runs on an accelerator only. Its torch path is otherwise the same code as Zamba2's (via the modular inheritance) and is covered by the Zamba2 CPU test.

**Propagated to Bamba, FalconH1, and GraniteMoeHybrid.** Same shape as the Zamba2/Nemotron-H fix and #46084: gate the single-step branch on `seq_len == 1`, and for the multi-token cached case prepend the cached conv left-context (dropping it after the conv) and pass the cached recurrent state as the SSM initial state. I left Jamba out; it uses the Mamba1 path, not this Mamba2 chunk-scan path.

**Kernels import issue.** I was on `torch==2.8.0+cu128`, `kernels==0.12.3`, `transformers==5.13.0.dev0`, with neither `causal-conv1d` nor `mamba-ssm` installed. The cached `kernels-community/mamba-ssm` build fails to import because it pulls `GreedySearchDecoderOnlyOutput` from `transformers.generation`, which has been removed. I only aliased those names locally to run the CUDA verification.

<details>
<summary>versions and stack trace</summary>

```text
python==3.12.3
torch==2.8.0+cu128
transformers==5.13.0.dev0
kernels==0.12.3
causal-conv1d: not installed
mamba-ssm: not installed
```

```pytb
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File ".../site-packages/kernels/utils.py", line 323, in get_kernel
    return _import_from_path(package_name, variant_path)
  File ".../site-packages/kernels/utils.py", line 160, in _import_from_path
    spec.loader.exec_module(module)
  File ".../models--kernels-community--mamba-ssm/snapshots/.../build/torch28-cxx11-cu128-x86_64-linux/__init__.py", line 6, in <module>
    from .models.mixer_seq_simple import MambaLMHeadModel
  File ".../build/torch28-cxx11-cu128-x86_64-linux/models/mixer_seq_simple.py", line 20, in <module>
    from ..utils.generation import GenerationMixin
  File ".../build/torch28-cxx11-cu128-x86_64-linux/utils/generation.py", line 14, in <module>
    from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer
ImportError: cannot import name 'GreedySearchDecoderOnlyOutput' from 'transformers.generation'
```

</details>


> Kernels import issue

Yea this is known, the issue is that recent kernels relies on specific ABI stable features on torch side so having more previous torch versions covered would be indeed nice.

> Nemotron-H has no _cpu variant

Hmm, I thought we had a cpu fallback for grouped mm tbh but maybe I'm remembering wrong

@vasqu  

**Nit (extract state):** done across all five mixers. The cached `conv_states` / `recurrent_states` are now bound to locals right after `use_precomputed_states` and reused in the cat / `initial_states` / single-step paths, matching mamba2. Pure refactor, no behavior change.

**CPU variant for nemotron_h:** added. It does run on CPU, just from torch 2.11: the MoE experts dispatch through `integrations/moe.py` to `torch._grouped_mm`, which only gained a CPU kernel in 2.11. The new `test_nemotron_h_chunked_prefill_cpu` is guarded with `require_torch_greater_or_equal("2.11")`, which is the torch the run-slow image (`transformers-all-latest-gpu`) ships, so it runs on CI and skips on older torch. Verified: passes on 2.11, `NotImplementedError` on 2.8.


run-slow: bamba, falcon_h1, granitemoehybrid, nemotron_h, zamba2

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

run-slow: bamba, falcon_h1, granitemoehybrid, nemotron_h, zamba2

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

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

models: ["models/bamba", "models/falcon_h1", "models/granitemoehybrid", "models/nemotron_h", "models/zamba2"]
quantizations: []

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

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

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [59352305](https://github.com/huggingface/transformers/commit/59352305a0135210e9e1be65a93ea30d8858bbde) | workflow commit (merge commit) |
| PR | [7bd3cada](https://github.com/Sunt-ing/transformers/commit/7bd3cadaf1ffb38ce485348d59b73e62d4f4dd40) | branch commit (from PR) |
| main | [b779eb6c](https://github.com/huggingface/transformers/commit/b779eb6c49a707d13ff178cfc00acfd97d12a789) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !