Fix inter-chunk recurrence in the Zamba2 / Nemotron-H Mamba2 slow path
Author: Jeronymous
Pull request: yes

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

### What does this PR do?

Fixes an incorrect reduction dimension in the inter-chunk SSM recurrence of the pure-PyTorch (no-kernel) `torch_forward` path shared by `Zamba2MambaMixer` and `NemotronHMamba2Mixer`.

The current code

```python
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)
```

sums over the *target*-chunk axis while `states` is broadcast over it, so `states` factors out of the sum. This collapses the intended recurrence

```
new_state[j] = Σ_i decay_chunk[j, i] · state[i]
```

into

```
new_state[q] = state[q] · Σ_p decay_chunk[p, q]
```

i.e. the per-chunk states are no longer propagated across chunk boundaries.

This is the same bug that was reported for `mamba2` in #34817 and fixed in #35154; the fix was never propagated to `zamba2`, and therefore to `nemotron_h`, whose mixer inherits `torch_forward` from `Zamba2MambaMixer`. All other Mamba2 descendants (`bamba`, `falcon_h1`, `granitemoehybrid`) already use the corrected form. The change aligns these two models with `mamba2`:

```python
decay_chunk = decay_chunk.transpose(1, 3)
new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
```

It was applied to `modular_zamba2.py` and mirrored into the generated `modeling_zamba2.py` and `modeling_nemotron_h.py` (identical lines, no name substitution), so `check_modular_conversion` stays green.

### Scope of impact

Only the slow path is affected (no `mamba-ssm` / `causal-conv1d` kernels — e.g. CPU, or `use_mamba_kernels=False`); the CUDA kernel path is correct. The bug is masked for a single chunk with an empty cache (`seq_len <= chunk_size`, fresh prefill), which is why existing short-sequence tests pass. It produces wrong outputs and wrong cached SSM states when:

- the input is longer than `chunk_size` (prefill spanning ≥ 2 chunks), or
- generation is continued from a populated SSM cache.

### Reproduction

<details>
<summary>Standalone numerical check (no model download)</summary>

```python
import torch


def segment_sum(input_tensor):  # verbatim from modeling_*.py
    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, 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, dtype=torch.bool), diagonal=0)
    return tensor_segsum.masked_fill(~mask, -torch.inf)


def run(num_chunks, zero_previous_state):
    b, h, d, n = 1, 2, 4, 5
    A_cumsum_last = torch.randn(b, h, num_chunks)          # == A_cumsum[:, :, :, -1]
    states = torch.randn(b, num_chunks + 1, h, d, n)       # cat([previous_states, per-chunk states])
    if zero_previous_state:
        states[:, :1] = 0                                  # fresh prefill => no cached state
    decay_chunk = torch.exp(segment_sum(torch.nn.functional.pad(A_cumsum_last, (1, 0))))

    # current (buggy): sum over dim=2
    sp = states.permute(0, 2, 1, 3, 4)
    buggy = (decay_chunk[..., None, None] * sp[:, :, None, ...]).sum(dim=2).permute(0, 2, 1, 3, 4)
    # fixed (mamba2 form): transpose(1, 3) + sum over dim=1
    dc_t = decay_chunk.transpose(1, 3)
    fixed = (dc_t[..., None, None] * states[:, :, None, ...]).sum(dim=1)
    # brute-force reference: new[j] = sum_i decay_chunk[h, j, i] * states[i]
    ref = torch.einsum("bhji,bihdn->bjhdn", decay_chunk, states)
    return (buggy - ref).abs().max().item(), (fixed - ref).abs().max().item()


torch.manual_seed(0)
for label, c, zero in [
    ("seq <= chunk_size, fresh prefill (empty cache)", 1, True),
    ("seq <= chunk_size, continuing from SSM cache",   1, False),
    ("seq  > chunk_size, fresh prefill",               3, True),
    ("seq  > chunk_size, continuing from SSM cache",   3, False),
]:
    bug, ok = run(c, zero)
    print(f"{label:48s}  buggy_vs_ref={bug:9.5f}  fixed_vs_ref={ok:9.5f}")
```

Output:

```
seq <= chunk_size, fresh prefill (empty cache)    buggy_vs_ref=  0.00000  fixed_vs_ref=  0.00000
seq <= chunk_size, continuing from SSM cache      buggy_vs_ref=  1.92688  fixed_vs_ref=  0.00000
seq  > chunk_size, fresh prefill                  buggy_vs_ref=  5.26675  fixed_vs_ref=  0.00000
seq  > chunk_size, continuing from SSM cache      buggy_vs_ref= 10.68568  fixed_vs_ref=  0.00000
```

</details>

Fixes #47246

### Who can review?

@vasqu (fixed the equivalent `mamba2` bug in #35154)

Comments:
<!-- 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=47250)
**Latest run:** [29114117332:2](https://github.com/huggingface/transformers/actions/runs/29407388784)
**Result:** `success` | **Jobs:** 5 | **Tests:** 646 | **Failures:** 0 | **Duration:** 5m 16s

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

Thank you for your contribution 🤗!

## CI Security Gate — automatic approval blocked

This PR was **not** automatically approved for CI because the security gate failed.

**Possible reasons:**
- The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
- A changed file is outside the allowed directories (`src/`, `tests/`, `docs/`, `utils/`), has a disallowed extension (only `.py`, `.txt`, `.md` permitted outside `tests/` and `docs/`), or is not `.md`/`.yml` inside `docs/`
- A new high-severity security issue was detected in the changed Python files (Bandit check)

See the [workflow run](https://github.com/huggingface/transformers/actions/runs/29425664968) for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

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

run-slow: nemotron_h, zamba2

run-slow: nemotron_h, zamba2

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

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

models: ["models/nemotron_h", "models/zamba2"]
quantizations: []

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

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [4dbe8928](https://github.com/huggingface/transformers/commit/4dbe89285ef648ab11a5276bb5468657ebf987ab) | workflow commit (merge commit) |
| PR | [65abbe6b](https://github.com/Jeronymous/transformers/commit/65abbe6b5316222a275013debfb94c04174f8ee0) | branch commit (from PR) |
| main | [e52d0fd6](https://github.com/huggingface/transformers/commit/e52d0fd6fa9eb874f7c2da048198276b04c919b9) | base commit (on `main`) |

⚠️ **Model CI failed to report results**

The test failure analysis could not be completed. Please check the [workflow run](https://github.com/huggingface/transformers/actions/runs/29432103851) for details.