[Bug Report] SuperV3 MTP regression on main introduced by #13566 (NVBUG 6221483) #14595

Summary

After PR #13566 ([fix][bug 6120981] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_chunks, squash-merged as 998f41855d), the NemotronSuperV3 MTP accuracy tests fail on main. Tracking issue for NVBUG 6221483.

The original suspect was PR #13478 (squash-merged as 92c5030ff4), because the bug only became visible at CUDA-graph-capture time after #13478's IR-canonical NemotronH model landed. Bisection (see comment below) showed #13478 only shifts the crash earlier: with #13478 the assert fires during compile_model/capture_graph for bs=8; without #13478 the same assert fires during the first MTP decode iteration. Reverting only #13478 does not make the test pass; reverting only #13566 does.

NVBUG

• https://nvbugspro.nvidia.com/bug/6221483

Failing tests

• accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm]
• accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm]

Reproduction (confirmed empirically)

git checkout main                              # FAIL (Indexing.cu:1515 indexSelectSmallIndex during cuda graph capture for bs=8)
git revert 998f41855d && retest                # PASS  -> #13566 IS the root cause
git revert 92c5030ff4 && retest                # still FAIL (shifted to first-decode crash) -> #13478 is NOT the root cause

Root cause

_mamba_ssm_prepare_metadata in tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py (added by #13566) trusts CPU-side batch_info_host.get_num_tokens() and seq_len_host for three shortcut kwargs that avoid GPU->CPU sync:

• cu_seqlens_to_chunk_indices_offsets_triton(..., total_seqlens=num_prefill_tokens, extra_chunks=_extra)
• torch.repeat_interleave(..., output_size=num_prefill_tokens)

During CUDA-graph capture warmup the CPU host tensors carry dummy values that do not match GPU-side seq_len/cu_seqlen, so PyTorch's internal index_select inside repeat_interleave walks off the end of the source tensor and fires Indexing.cu:1515 indexSelectSmallIndex: srcIndex < srcSelectDimSize. #13478's IR-canonical model exposes the warmup path earlier, which is why the symptom moved from "first decode" to "capture time" when #13478 landed.

Why CI did not catch it

Both PRs (#13566 and #13478) were merged via /bot skip after their L0 stages reported unrelated failures. Each was defensible against its own L0 but neither L0 exercised the post-merge interaction:

• [https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566 (the actual root cause) merged earlier; its L0 did not run the SuperV3 MTP cases.
• [TRTLLM-13429][feat] Switch DeepSeek/NemotronH/Qwen3/Qwen3.5-MoE to sharding-IR canonical models #13478 (the trigger) merged later; its L0 did run the SuperV3 MTP cases but against a stale base that pre-dated [https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566's effect on the warmup path.

A "stale-base gate" has now been added to the local pr-monitor skill to prevent a repeat.

Investigation plan (completed)

1. Reproduce locally on main -> confirmed: Indexing.cu:1515 indexSelectSmallIndex during compile_model/capture_graph for bs=8.
2. Bisect within [https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566's diff -> see bisection comment below; only the full file revert of mamba_backend_common.py passes, even keeping just the seq_len_host signature re-triggers the bug.
3. Apply the minimal fix: revert tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py in full to its pre-[https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566 state; leave mamba2_metadata.py and waives.txt parts of [https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566 untouched.
4. Local validation: full repro (world_size=4, real FP8 weights, full model, cuda_graph_batch_sizes=[1,2,4,8]) -> 10/10 prompts generated, 0 asserts.
5. CI validation: DGX_H100-4_GPUs-AutoDeploy-1 (owns fp8_ws4_80gb) and DGX_B200-4_GPUs-AutoDeploy-1 (owns nvfp4_ws4_80gb) -- triggered on the fixing PR.

Scope

• Surface: tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py
• Affected component: AutoDeploy SSM prepare-metadata custom op (NemotronH-family models with MoE EP)

Risk

library-visible. The fix re-introduces the GPU->CPU sync that #13566 was originally trying to remove for EP-deadlock prevention (NVBUG 6120981). A follow-up that restores that optimization safely (e.g. by gating on torch.cuda.is_current_stream_capturing()) is a separate task.

Links

• Offending PR (actual root cause): [https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566 -- squash-merge 998f41855d
• Exposing PR (kept as additional context): [TRTLLM-13429][feat] Switch DeepSeek/NemotronH/Qwen3/Qwen3.5-MoE to sharding-IR canonical models #13478 -- squash-merge 92c5030ff4
• Test-restructuring PR in the gap: [#14173][tests] move autodeploy accuracy tests to post merge and use model registry #14352 (a8cd4fffb3)
• Fixing PR: [https://nvbugs/6221483][fix] Revert auto_deploy _mamba_ssm_prepare_metadata to pre-#13566 state #14640