[https://nvbugs/6221483][fix] Revert auto_deploy _mamba_ssm_prepare_metadata to pre-#13566 state by greg-kwasniewski1 · Pull Request #14640 · NVIDIA/TensorRT-LLM

greg-kwasniewski1 · 2026-05-27T18:17:36Z

Summary

Fixes SuperV3 MTP regression (NVBUG 6221483) by reverting tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py to its pre-[https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566 contents.
Root cause: the new _mamba_ssm_prepare_metadata trusts stale CPU host tensors (batch_info_host.get_num_tokens() and seq_len_host) during CUDA-graph capture warmup; the repeat_interleave(..., output_size=k) shortcut hits Indexing.cu:1515 indexSelectSmallIndex: srcIndex < srcSelectDimSize when the host shadow disagrees with the GPU-side seq_len/cu_seqlen dummies emitted by warmup.
Minimum fix: bisection on upstream/main showed every smaller revert (S1-only, S2-only, S3-only, S1+S2, S1+S2+S3, function-swap-only, keep-signature-only) still crashes; only the full file revert passes (see bisection comment on [Bug Report] SuperV3 MTP regression on main introduced by #13566 (NVBUG 6221483) #14595).

Changes

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py: restore pre-[https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch #13566 contents (-22 / +4 lines), reverting the imports, the seq_len_host signature additions, and the body changes (_extra = compute_extra_chunks_cpu(...), the total_seqlens=/extra_chunks= kwargs, and the output_size=num_prefill_tokens kwarg). The 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 are intentionally left untouched.

Test plan

Local: build_and_run_ad.py --yaml-extra repro_nvbug6221483.yaml (full repro mirroring TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm]) on upstream/main with this fix applied: 10/10 prompts generated, 0 Indexing.cu asserts, clean MPI shutdown.
CI: DGX_H100-4_GPUs-AutoDeploy-1 (owns fp8_ws4_80gb-trtllm) and DGX_B200-4_GPUs-AutoDeploy-1 (owns nvfp4_ws4_80gb-trtllm) — triggered below.

Risk

library-visible — affects the AutoDeploy SSM prepare-metadata custom op. Re-introduces the GPU->CPU sync that #13566 was trying to remove for EP-deadlock prevention (nvbug 6120981). A follow-up that restores that optimization safely (e.g. by branching on torch.cuda.is_current_stream_capturing()) is a separate task.

Fixes #14595

Summary by CodeRabbit

Release Notes

Refactor
- Optimized Mamba SSM metadata computation by streamlining the custom operation signature and removing redundant processing paths, reducing computational overhead during inference.

…etadata to pre-NVIDIA#13566 state PR NVIDIA#13566 (998f418, "Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_chunks") introduced a CUDA-graph-capture-time crash on NemotronSuperV3 MTP. The new `_mamba_ssm_prepare_metadata` trusts CPU-side `batch_info_host.get_num_tokens()` and `seq_len_host` for the `total_seqlens`, `extra_chunks`, and `output_size` shortcuts that were added to avoid GPU->CPU sync. During CUDA-graph warmup the CPU host tensors carry dummy values that do not match GPU-side `seq_len`/`cu_seqlen`, so the internal `index_select` in `repeat_interleave(..., output_size=k)` walks off the source tensor and fires `Indexing.cu:1515 indexSelectSmallIndex: srcIndex < srcSelectDimSize`. Bisection on upstream/main + full repro_nvbug6221483.yaml showed every smaller revert (S1-only, S2-only, S3-only, S1+S2, S1+S2+S3, function-swap-only, keep signature only) still crashes. Reverting mamba_backend_common.py in full to its pre-NVIDIA#13566 contents passes the repro (10/10 prompts, 0 asserts). mamba2_metadata.py and waives.txt parts of NVIDIA#13566 are left untouched. This re-introduces the GPU->CPU sync that NVIDIA#13566 was trying to remove (nvbug 6120981 EP-deadlock prevention). A follow-up can restore that optimization by branching on torch.cuda.is_current_stream_capturing(). Fixes NVIDIA#14595 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

coderabbitai · 2026-05-27T18:21:05Z

📝 Walkthrough

Walkthrough

This PR simplifies the cached SSM metadata preparation in the Mamba backend by removing the seq_len_host parameter from the _mamba_ssm_prepare_metadata custom op and eliminating the CPU/TRITON extra-chunk computation path. The implementation now derives chunk indices and offsets directly from GPU cu_seqlen and chunk_size using cu_seqlens_to_chunk_indices_offsets, and updates the fake op registration accordingly.

Changes

SSM metadata computation simplification

Layer / File(s)	Summary
Custom op signature and metadata computation `tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py`	Module import updated to use `cu_seqlens_to_chunk_indices_offsets` directly, removing CPU/TRITON variants. The `_mamba_ssm_prepare_metadata` custom op `num_prefill > 0` path now computes chunk indices/offsets from GPU `cu_seqlen` tensor and `chunk_size` without `seq_len_host` or extra-chunk CPU logic. Fake op registration input list updated to remove `seq_len_host` parameter.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#13566: Prior PR that modified the same _mamba_ssm_prepare_metadata function to adjust chunk indices/offsets computation and seq_len_host parameter handling.

Suggested reviewers

shaharmor98
nzmora-nvidia
2ez4bz
tcherckez-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: reverting mamba_backend_common.py to pre-#13566 state to fix NVBUG 6221483, directly matching the primary objective.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description clearly explains the issue (regression fix), root cause (stale CPU tensors during CUDA-graph capture), changes made (full file revert), test coverage (local and CI), and risk assessment.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

greg-kwasniewski1 · 2026-05-27T18:21:47Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1"

tensorrt-cicd · 2026-05-27T18:27:52Z

PR_Github #50597 [ run ] triggered by Bot. Commit: b70ca26 Link to invocation

…ixed by previous commit Removes the two SKIP entries from tests/integration/test_lists/waives.txt that were tagged with NVBUG 6221483: - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-4] Both share the same mamba_backend_common.py code path that was reverted to its pre-NVIDIA#13566 state in the previous commit, so the fix should cover both. The fp8 variant (TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm]) was never in waives.txt - it was running and failing on DGX_H100-4_GPUs stage already; the AutoDeploy CI stages triggered on this PR exercise it. Fixes NVIDIA#14595 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

greg-kwasniewski1 · 2026-05-27T19:25:51Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1"

tensorrt-cicd · 2026-05-27T23:31:18Z

PR_Github #50597 [ run ] completed with state FAILURE. Commit: b70ca26
/LLM/main/L0_MergeRequest_PR pipeline #40092 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>

galagam · 2026-05-28T03:56:43Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1"

tensorrt-cicd · 2026-05-28T04:03:05Z

PR_Github #50687 [ run ] triggered by Bot. Commit: ec99a9b Link to invocation

tensorrt-cicd · 2026-05-28T06:34:34Z

PR_Github #50687 [ run ] completed with state SUCCESS. Commit: ec99a9b
/LLM/main/L0_MergeRequest_PR pipeline #40175 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

galagam · 2026-05-28T07:03:35Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1"

tensorrt-cicd · 2026-05-28T07:10:53Z

PR_Github #50723 [ run ] triggered by Bot. Commit: ec99a9b Link to invocation

tensorrt-cicd · 2026-05-28T08:22:28Z

PR_Github #50723 [ run ] completed with state SUCCESS. Commit: ec99a9b
/LLM/main/L0_MergeRequest_PR pipeline #40206 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

galagam · 2026-05-28T09:02:20Z

/bot run

tensorrt-cicd · 2026-05-28T09:08:14Z

PR_Github #50748 [ run ] triggered by Bot. Commit: ec99a9b Link to invocation

tensorrt-cicd · 2026-05-28T11:44:45Z

PR_Github #50748 [ run ] completed with state SUCCESS. Commit: ec99a9b
/LLM/main/L0_MergeRequest_PR pipeline #40227 completed with status: 'SUCCESS'

CI Report

Link to invocation

PR NVIDIA#14640 (merged 2026-05-28) fixed the mamba-revert bug that caused these AutoDeploy 4-GPU tests to be moved to post_merge. The "Move to post-merge due to https://nvbugspro.nvidia.com/bug/6221483" gate is no longer needed. Move back to pre_merge AutoDeploy 4-GPU blocks in both DGX_B200 and DGX_H100: TestNemotronSuperV3::test_accuracy[bf16-4-attn_dp_off-trtllm] TestNemotronSuperV3::test_accuracy[bf16-4-attn_dp_on-trtllm] (h100 only) TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_on-trtllm] TestNemotronSuperV3::test_accuracy[nvfp4-4-attn_dp_on-trtllm] (b200 only) TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm] (b200 only) TestNemotronNanoV3::test_accuracy[nvfp4-4-trtllm] (b200 only) Importantly, this brings the only end-to-end MoE all-to-all + attention-DP coverage (TestNemotronSuperV3::test_accuracy[*-attn_dp_on-*]) into pre-merge CI for PR NVIDIA#13723. The post_merge B200 block keeps only the bf16_ws4_180gb MTP variants and the perf-sanity tests, which are gated for separate hardware-availability reasons. Drops a duplicate UltraV3 nvfp4-4 entry that already lived in pre_merge B200. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

greg-kwasniewski1 requested a review from a team as a code owner May 27, 2026 18:17

greg-kwasniewski1 requested a review from bmarimuthu-nv May 27, 2026 18:17

github-actions Bot assigned greg-kwasniewski1 May 27, 2026

greg-kwasniewski1 mentioned this pull request May 27, 2026

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

Closed

greg-kwasniewski1 requested a review from tcherckez-nvidia May 27, 2026 18:27

Merge branch 'main' into gk/fix-nvbug-6221483

ec99a9b

Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>

galagam mentioned this pull request May 28, 2026

[#12359][feat] AutoDeploy: Support SSM replay kernel for MTP with FlashInfer #13725

Merged

1 task

galagam approved these changes May 28, 2026

View reviewed changes

galagam merged commit 6484b71 into NVIDIA:main May 28, 2026
9 of 10 checks passed

Uh oh!

Conversation

greg-kwasniewski1 commented May 27, 2026 • edited by galagam Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Test plan

Risk

Summary by CodeRabbit

Release Notes

Uh oh!

coderabbitai Bot commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

greg-kwasniewski1 commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

greg-kwasniewski1 commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

galagam commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

galagam commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

greg-kwasniewski1 commented May 27, 2026 •

edited by galagam

Loading

coderabbitai Bot commented May 27, 2026 •

edited

Loading