Incorrect inter-chunk recurrence in `NemotronHMamba2Mixer.torch_forward` (slow path). the Mamba2 fix from #35154 was never propagated to `nemotron_h`
Author: Jeronymous
Pull request: no
Labels: name: bug

### System Info

- `transformers` versions affected: **5.6.0 through 5.13.0, and current `main`** (the `nemotron_h` model was refactored into the library after ~5.3.0).
- Platform: any (bug is in pure-PyTorch code; independent of OS/hardware).
- Only the **slow / no-kernel path** is affected (`use_mamba_kernels=False`, or `mamba-ssm` / `causal-conv1d` not installed, or CPU). The CUDA-kernel path (`cuda_kernels_forward`) is correct.


### Who can help?

The equivalent bug in the canonical Mamba2 model was fixed by @vasqu in #35154 (reported by @HanGuo97 in #34817). This report is that the same fix was not carried over to the `nemotron_h` port. cc anyone maintaining SSM/Mamba models.


### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Self-contained, no model download needed — it isolates the exact tensor op and compares it against a brute-force reference and against the correct canonical Mamba2 form:

```python
import torch


def segment_sum(input_tensor):  # verbatim from modeling_nemotron_h.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))))

    # nemotron_h current (buggy)
    sp = states.permute(0, 2, 1, 3, 4)
    buggy = (decay_chunk[..., None, None] * sp[:, :, None, ...]).sum(dim=2).permute(0, 2, 1, 3, 4)
    # canonical mamba2 (correct)
    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
```

The correct form matches the reference exactly in all cases; the current `nemotron_h` form diverges except in the single masked case below.

### Impact / why it hasn't been caught

The wrong contraction only affects the `.sum` result when `states` varies over the reduced axis, so it is **masked in exactly one case: a fresh prefill (empty SSM cache) with sequence length ≤ `chunk_size`** (a single chunk, where `previous_states` is zero — the first row of the diff above). That is the regime most fast unit tests run in, so they pass.

It **does** produce wrong numbers whenever:
- the input is longer than `chunk_size` (default 128) tokens on the slow path (prefill), **or**
- generation is continued from a non-empty SSM cache on the slow path,

corrupting both the layer output (via `Y_off`) and the SSM state written back to the cache (so subsequent decode steps are wrong too). GPU inference with the Mamba CUDA kernels is unaffected because it uses `cuda_kernels_forward`.

### Related

- #34817 — original report of the same bug in `Mamba2`
- #35154 — the fix for `Mamba2` (supersedes #34901); this report asks to propagate it to `nemotron_h`

### Secondary observation (same method, lower severity)

In the same `torch_forward`, `dt` is clamped with `torch.clamp(dt, self.time_step_min)` whereas the canonical Mamba2 (and this model's own `cuda_kernels_forward`, via `dt_limit`) clamps with `time_step_limit`. For configs where `time_step_limit == (0.0, inf)` this introduces a `0.001` floor on the slow path that the kernel path does not apply — a minor slow-path/kernel-path inconsistency worth aligning while fixing the above.


### Expected behavior

### Description

`NemotronHMamba2Mixer.torch_forward` computes the inter-chunk SSM recurrence ("step 3 — produces correct SSM states at chunk boundaries") with an incorrect reduction dimension. It sums over the wrong axis, which collapses the recurrence `new_state[j] = Σ_i decay_chunk[j, i] · state[i]` into `new_state[q] = state[q] · Σ_p decay_chunk[p, q]` — i.e. each chunk state is merely scaled by a column-sum of the decay matrix instead of being mixed across chunk boundaries.

This is **exactly** the bug that was reported for the canonical `Mamba2` model in #34817 and fixed in #35154 (originally #34901). The canonical `mamba2/modeling_mamba2.py` now uses the correct `decay_chunk.transpose(1, 3)` + `.sum(dim=1)` form, but `nemotron_h/modeling_nemotron_h.py` still carries the pre-fix `.permute(...)` + `.sum(dim=2)` form.

### Offending code

`src/transformers/models/nemotron_h/modeling_nemotron_h.py`, in `NemotronHMamba2Mixer.torch_forward` (e.g. v5.13.0 lines 554–556; unchanged on `main`):

```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)
```

After `states = torch.cat([previous_states, states], dim=1)`, `states` has shape `(b, c+1, h, d, n)` and `decay_chunk` has shape `(b, h, c+1, c+1)` (lower-triangular, `[..., target, source]`). In the broadcast product `(b, h, z, z, d, n)`, the states' chunk index lands on axis 3 (shared with `decay_chunk`'s *source* axis), so `.sum(dim=2)` reduces over the *target* axis while `states` is broadcast over it — pulling `states` out of the sum. The reduction must be over the *source* axis instead.

For reference, the correct canonical Mamba2 form (`mamba2/modeling_mamba2.py`, `main`) is:

```python
decay_chunk = decay_chunk.transpose(1, 3)
new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
```

### Suggested fix

Adopt the canonical Mamba2 form (recommended, keeps the two implementations in sync):

```diff
-        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)
+        decay_chunk = decay_chunk.transpose(1, 3)
+        new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
```

(Equivalently, keeping the `permute` structure, changing `.sum(dim=2)` to `.sum(dim=3)` also produces the correct contraction — see #34817.)


Comments:
We really should start refactoring Mamba2 to have some core functions 😢 if you'd like, can you open a PR to fix? 

Sure: https://github.com/huggingface/transformers/pull/47250