From 4a1cd68da883ff5c0eb50b21d064fd1a027f2704 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Tue, 28 Apr 2026 07:50:12 -0700 Subject: [PATCH 1/2] [nvbugs/6120981][fix] Eliminate GPU->CPU syncs in AutoDeploy Mamba SSM 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> --- .../custom_ops/mamba/mamba_backend_common.py | 26 ++++++++++++++++--- .../_torch/modules/mamba/mamba2_metadata.py | 25 +++++++++++------- 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py index a91095fe7aa1..45feb2d4ea4b 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py @@ -19,7 +19,10 @@ from torch._ops import OpOverloadPacket from torch.fx import Node -from tensorrt_llm._torch.modules.mamba.mamba2_metadata import cu_seqlens_to_chunk_indices_offsets +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + compute_extra_chunks_cpu, + cu_seqlens_to_chunk_indices_offsets_triton, +) from tensorrt_llm._torch.modules.mamba.ssd_combined import mamba_chunk_scan_combined from ..._compat import KvCacheConfig @@ -41,6 +44,7 @@ def _mamba_ssm_prepare_metadata( position_ids: torch.Tensor, batch_info_host: torch.Tensor, seq_len: torch.Tensor, + seq_len_host: torch.Tensor, cu_seqlen: torch.Tensor, # EXTRA METADATA PROVIDED BY THE DESCRIPTOR chunk_size: int, @@ -48,6 +52,10 @@ def _mamba_ssm_prepare_metadata( """Prepare metadata for cached SSM transform. Returns a tuple of (chunk_indices, chunk_offsets, seq_idx_prefill). + + Uses seq_len_host (CPU tensor) and batch_info_host to derive total_seqlens + and extra_chunks without GPU->CPU synchronization, preventing deadlocks + when NCCL collectives from MoE expert-parallel layers are pending. """ device = cu_seqlen.device batch_info = BatchInfo(batch_info_host) @@ -55,11 +63,20 @@ def _mamba_ssm_prepare_metadata( num_prefill, _, _ = batch_info.get_num_sequences() if num_prefill > 0: - chunk_indices, chunk_offsets = cu_seqlens_to_chunk_indices_offsets( - cu_seqlen[: num_prefill + 1], chunk_size + num_prefill_tokens, _, _ = batch_info.get_num_tokens() + + _extra = compute_extra_chunks_cpu(seq_len_host, num_prefill, chunk_size) + + chunk_indices, chunk_offsets = cu_seqlens_to_chunk_indices_offsets_triton( + cu_seqlen[: num_prefill + 1], + chunk_size, + total_seqlens=num_prefill_tokens, + extra_chunks=_extra, ) seq_idx_prefill = torch.repeat_interleave( - torch.arange(num_prefill, device=device, dtype=torch.int32), seq_len[:num_prefill] + torch.arange(num_prefill, device=device, dtype=torch.int32), + seq_len[:num_prefill], + output_size=num_prefill_tokens, ).view(1, -1) else: chunk_indices = torch.empty(0, dtype=torch.int32, device=device) @@ -75,6 +92,7 @@ def _mamba_ssm_prepare_metadata_fake( position_ids: torch.Tensor, batch_info_host: torch.Tensor, seq_len: torch.Tensor, + seq_len_host: torch.Tensor, cu_seqlen: torch.Tensor, # EXTRA METADATA PROVIDED BY THE DESCRIPTOR chunk_size: int, diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 3d1316e40519..4ddd67707bb2 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -63,6 +63,20 @@ def _cu_seqlens_triton_kernel( tl.store(chunk_offsets_ptr + offsets, chunk_offsets.to(tl.int32), mask=mask) +def compute_extra_chunks_cpu(seq_lens, 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. + """ + 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 + + def cu_seqlens_to_chunk_indices_offsets_triton( cu_seqlens: torch.Tensor, chunk_size: int, @@ -327,15 +341,8 @@ def prepare(self, attn_metadata: AttentionMetadata): self.has_initial_states_cpu[:num_contexts].any()) if self.use_initial_states: - # Compute extra_chunks using pure Python arithmetic on CPU - # seq_lens to avoid any GPU->CPU sync point. - _cs = self.chunk_size - _cumsum = 0 - _extra = 0 - for i in range(num_contexts - 1): - _cumsum += int(attn_metadata.seq_lens[i]) - if _cumsum % _cs != 0: - _extra += 1 + _extra = compute_extra_chunks_cpu(attn_metadata.seq_lens, + num_contexts, self.chunk_size) self.chunk_indices, self.chunk_offsets = cu_seqlens_to_chunk_indices_offsets_triton( self.cu_seqlens[:num_contexts + 1], From fa1ceae2dfdc9cf136a303a71e1899d5b2d5f630 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Thu, 21 May 2026 04:21:40 -0700 Subject: [PATCH 2/2] [nvbugs/6120981][chore] Remove stale waiver after fix Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index a471e8386e10..66c737d57d54 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -15,7 +15,6 @@ accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 SKIP (https://nvbu accuracy/test_llm_api_autodeploy.py::TestGemmaE2B::test_gemma4_e2b_it SKIP (https://nvbugs/6194934) accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (https://nvbugs/6158397) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-trtllm] SKIP (https://nvbugs/6200112) -accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-4-trtllm] SKIP (https://nvbugs/6120981) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[bf16] SKIP (https://nvbugs/6162114) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[fp8] SKIP (https://nvbugs/6162114) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SKIP (https://nvbugs/6158397)