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], 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)