[https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch - #13566
Conversation
📝 WalkthroughWalkthroughThe changes refactor Mamba metadata computation by introducing a new Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py (1)
66-77: Hardencompute_extra_chunks_cpuagainst accidental CUDA inputs.This helper is specifically for host-side computation; a CUDA tensor here can reintroduce sync risk. Consider a fail-fast guard.
Proposed patch
-def compute_extra_chunks_cpu(seq_lens, num_seqs: int, chunk_size: int) -> int: +def compute_extra_chunks_cpu( + seq_lens: torch.Tensor | list[int], num_seqs: int, chunk_size: int +) -> int: """Count extra chunks caused by misaligned sequence boundaries. Computes from CPU seq_lens to avoid GPU->CPU synchronization. """ + if isinstance(seq_lens, torch.Tensor) and seq_lens.is_cuda: + raise ValueError("seq_lens must be a CPU tensor") + cumsum = 0 extra = 0 for i in range(num_seqs - 1): cumsum += int(seq_lens[i]) if cumsum % chunk_size != 0: extra += 1 return extra🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py` around lines 66 - 77, The function compute_extra_chunks_cpu must fail fast on CUDA inputs to avoid accidental GPU->CPU syncs: at the start of compute_extra_chunks_cpu detect if seq_lens is a torch Tensor on CUDA (e.g., has .device and .device.type == "cuda") and raise a clear TypeError telling callers to pass a CPU tensor/ndarray or call .cpu()/.tolist() first; ensure the check is a simple predicate so Python lists/ndarrays remain supported and keep the rest of the logic unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py`:
- Around line 66-77: The function compute_extra_chunks_cpu must fail fast on
CUDA inputs to avoid accidental GPU->CPU syncs: at the start of
compute_extra_chunks_cpu detect if seq_lens is a torch Tensor on CUDA (e.g., has
.device and .device.type == "cuda") and raise a clear TypeError telling callers
to pass a CPU tensor/ndarray or call .cpu()/.tolist() first; ensure the check is
a simple predicate so Python lists/ndarrays remain supported and keep the rest
of the logic unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7335ea2c-1415-46be-b1a0-fff61987ad73
📒 Files selected for processing (2)
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.pytensorrt_llm/_torch/modules/mamba/mamba2_metadata.py
|
[Repair Bot] Attempted automated rebase/maintenance, but could not complete it safely. Leaving the PR unchanged. Reason: The bot will try again on a later cycle. Manual rebase is also fine. |
|
@tcherckez-nvidia : any updates here. Have we checked if it fixes https://nvbugs/6120981? |
|
@suyoggupta Looked at it, and I can understand the fix, but I'm not that familiar with the code. |
eb1231c to
e6a83b5
Compare
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #48370 [ run ] triggered by Bot. Commit: |
|
PR_Github #48370 [ run ] completed with state
|
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #48391 [ run ] triggered by Bot. Commit: |
|
PR_Github #48391 [ run ] completed with state
|
c507cb9 to
0546861
Compare
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #48741 [ run ] triggered by Bot. Commit: |
|
PR_Github #48741 [ run ] completed with state
|
|
PR_Github #48838 [ run ] triggered by Bot. Commit: |
|
PR_Github #48838 [ run ] completed with state
|
|
/bot run |
|
PR_Github #48900 [ run ] triggered by Bot. Commit: |
|
PR_Github #48900 [ run ] completed with state
|
0546861 to
1ed51c6
Compare
|
/bot run |
|
PR_Github #49223 [ run ] triggered by Bot. Commit: |
|
PR_Github #49223 [ run ] completed with state
|
1ed51c6 to
4c68b9a
Compare
|
/bot run |
1 similar comment
|
/bot run |
4c68b9a to
84f3c9f
Compare
…M metadata to prevent EP deadlock Replace cu_seqlens_to_chunk_indices_offsets (which iterates GPU tensor elements causing implicit cudaStreamSynchronize) with cu_seqlens_to_chunk_indices_offsets_triton. Pre-compute total_seqlens and extra_chunks from CPU-side batch_info_host and seq_len_host tensors. Add output_size to repeat_interleave to avoid its implicit sync. This prevents deadlocks when NCCL all-to-all collectives from MoE expert-parallel layers are pending, as the GPU->CPU sync would block waiting for the collective while other ranks are still executing MoE layers. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
84f3c9f to
fa1ceae
Compare
|
/bot run |
|
PR_Github #50081 [ run ] triggered by Bot. Commit: |
|
PR_Github #50081 [ run ] completed with state
|
|
/bot run |
|
PR_Github #50143 [ run ] triggered by Bot. Commit: |
|
PR_Github #50143 [ run ] completed with state |
…ffsets_triton with total_seqlens/extra_ch (NVIDIA#13566) Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
…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>
…ffsets_triton with total_seqlens/extra_ch (NVIDIA#13566) Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
…etadata to pre-#13566 state (#14640) Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com> Co-authored-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>
4 participants
Summary
Test plan
Links
Summary by CodeRabbit
Bug Fixes
Performance