fix(nemotron_h): correct inter-chunk SSM recurrence axis and dt clamping in torch_forward

[adeev-mardia]

Fixes #47246.

What this PR does

Overrides NemotronHMamba2Mixer.torch_forward in modular_nemotron_h.py to fix two bugs in the slow (no-kernel) path that were inherited from Zamba2MambaMixer and never corrected for NemotronH.

---

Bug 1 — wrong reduction axis in the inter-chunk SSM recurrence

The inherited Zamba2 code computes step 3 (inter-chunk A-terms) by permuting states (swapping the batch/num-chunks axes) and reducing over dim=2:

# BEFORE — buggy (Zamba2 / inherited NemotronH)
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)

This collapses the wrong axis. The correct form — which matches canonical Mamba2 (fixed in #35154 for issue #34817) — transposes decay_chunk instead and reduces over dim=1:

# AFTER — correct
decay_chunk = decay_chunk.transpose(1, 3)
new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)

Impact: The bug produces incorrect layer outputs and corrupts the SSM cache state in every case except a fresh prefill with sequence length ≤ chunk_size (default 128). That single masked case is exactly what fast unit tests exercise, so the bug went undetected. It affects:

• Any slow-path prefill longer than 128 tokens
• Any slow-path generation step continued from a non-empty SSM cache

The CUDA-kernel path (cuda_kernels_forward) is correct and unaffected.

---

Bug 2 — dt clamped with only a lower bound on the slow path

The inherited code calls:

dt = torch.clamp(dt, self.time_step_min)   # one-sided — no upper bound

cuda_kernels_forward passes dt_limit (a 2-tuple), so the slow path silently applies a 0.001 floor that the kernel path does not. Fixed to:

dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])

---

Testing

The reproduction script from #47246 confirms the fix (runs on CPU, no model download needed):

Case                                              buggy_vs_ref  fixed_vs_ref
seq <= chunk_size, fresh prefill (empty cache)      0.00000      0.00000000
seq <= chunk_size, continuing from SSM cache        1.92688      0.00000000
seq  > chunk_size, fresh prefill                    5.26675      0.00000000
seq  > chunk_size, continuing from SSM cache       10.68568      0.00000000

The fixed form matches the brute-force reference exactly in all cases.

Related

• Issue Incorrect inter-chunk recurrence in NemotronHMamba2Mixer.torch_forward (slow path). the Mamba2 fix from #35154 was never propagated to nemotron_h #47246 — original report for NemotronH
• Issue Mamba2 torch_forward reduction dimension possibly incorrect? #34817 — original report for canonical Mamba2
• PR [Mamba2] Fix caching, slow path, and multi-gpu #35154 — the fix for canonical Mamba2 (this PR propagates it to NemotronH)

[github-actions]

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

run-slow: nemotron_h

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 30065575747
Result: failure | Grafana metrics are not available yet.

[Rocketknight1]

This is just raw code agent output and there was another PR at #47250