[Bug] NemotronH --mamba-scheduler-strategy extra_buffer accuracy drop on AIME26 (Nemotron-3-Super-120B) #31833

Checklist

• I searched related issues but found no solution.
• The bug persists in the latest version.
• Issues without environment info and a minimal reproducible demo are hard to resolve and may receive no feedback.
• If this is not a bug report but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed.
• Please use English. Otherwise, it will be closed.

Bug Feature Tracking for NemotronH --mamba-scheduler-strategy extra_buffer:

Describe the bug

Serving nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 (NemotronH, hybrid Mamba2 + attention) with --mamba-scheduler-strategy extra_buffer produces a visible accuracy drop on AIME26 vs. the default no_buffer strategy, with everything else (model, TP/EP layout, parsers, eval config) held identical. extra_buffer is meant to be a pure scheduling/caching optimization (overlap schedule + mamba-state radix-cache branching) with no expected effect on generation quality.

Root cause is suspected to be in the mamba-state tracking/checkpointing math in hybrid_linear_attn_backend.py::_init_track_ssm_indices, used by Mamba2AttnBackend for this model's SSM layers:

_init_track_ssm_indices computes a per-request "how many packed intermediate h states did this request produce" count as extend_seq_lens // mamba_cache_chunk_size and cumulative-sums it across requests to find each request's offset into the packed intermediate_states tensor returned by mamba_chunk_scan_combined. That tensor is actually indexed by a global physical chunk grid over the whole flattened extend batch (ssd_state_passing.py's _state_passing_fwd_kernel loops for c in range(nchunks) over the shared grid, writing one state per physical chunk regardless of request boundaries), not by a per-request-local contiguous count. The cumulative local-count offset is only correct if every earlier request in the same tracked prefill batch has a chunk-aligned length — not guaranteed in general.

For this model (chunk_size=128 in config.json, i.e. native mamba_chunk_size=128) served at --page-size 64 (the effective default when radix cache + extra_buffer + trtllm_mha are combined), half of all page-aligned tracked prefill lengths are not chunk-128-aligned, so the affected _init_track_ssm_indices "unaligned" branch fires routinely — this is not an edge case.

The net effect: a wrong (but same-shaped) SSM state gets read from intermediate_states and written into a request's mamba radix-cache slot whenever this happens, silently corrupting later prefix-cache-hit continuations for that node — no crash, no assertion, just wrong logits from that point on. This matches the observed AIME26 score degradation: extra_buffer should be numerically a no-op vs. no_buffer, and any drop attributable to it (beyond noise/seed variance) is evidence of this corruption path.

(Score for the extra_buffer run is pending — will update this issue once the run completes.)

Reproduction

Server (baseline — accuracy OK, no_buffer is the resolved default):

python3 -m sglang.launch_server \
  --model-path nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 \
  --trust-remote-code \
  --tp 4 \
  --ep 4 \
  --reasoning-parser nemotron_3 \
  --tool-call-parser qwen3_coder

Server (suspected buggy — extra_buffer mamba radix-cache strategy):

python3 -m sglang.launch_server \
  --model-path nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 \
  --trust-remote-code \
  --tp 4 \
  --ep 4 \
  --reasoning-parser nemotron_3 \
  --tool-call-parser qwen3_coder \
  --mamba-scheduler-strategy extra_buffer \
  --attention-backend trtllm_mha

Eval (identical for both):

sgl-eval run aime26 --base-url http://localhost:30000/v1 --n-repeats 16 --max-tokens 64000 --temperature 1.0 --top-p 0.95 --num-threads 416

Baseline (no_buffer) result:

== aime26 ==
30 examples x 16 repeats  |  1384.2s  |  7824 tok/s  |  10.8M tokens

* pass@1[avg-of-16]  =  84.38% +/- 3.59% (SEM 0.90%)
  pass@16            =  93.33%
  majority@16        =  90.00%
  no_answer          =  11.88%  [warn: consider --max-tokens]
  stop_rate          =  88.12%
  truncated_rate     =  11.88%  [warn: hitting max_tokens]
  error_rate         =  0.00%

extra_buffer result: TBD

Environment

Python: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0]
CUDA available: True
GPU 0,1,2,3,4,5,6,7: NVIDIA B300 SXM6 AC
GPU 0,1,2,3,4,5,6,7 Compute Capability: 10.3
CUDA_HOME: /usr/local/cuda
NVCC: Cuda compilation tools, release 13.0, V13.0.88
CUDA Driver Version: 610.43.02
PyTorch: 2.11.0+cu130
sglang: 0.5.15.dev1093+ge856eae92
sglang-kernel: 0.4.5
flashinfer_python: 0.6.14
flashinfer_cubin: 0.6.14
flashinfer_jit_cache: 0.6.14+cu130
triton: 3.6.0
transformers: 5.12.1