Incorrect inter-chunk recurrence in `NemotronHMamba2Mixer.torch_forward` (slow path). the Mamba2 fix from #35154 was never propagated to `nemotron_h`

[Jeronymous]

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:

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

• Mamba2 torch_forward reduction dimension possibly incorrect? #34817 — original report of the same bug in Mamba2
• [Mamba2] Fix caching, slow path, and multi-gpu #35154 — the fix for Mamba2 (supersedes [Mamba2] Fix slow path #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):

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:

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

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

[vasqu]

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

[Jeronymous]

Sure: #47250

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.