From 8283b15b1099bd721516190cc23a2087611c13ef Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Thu, 18 Jun 2026 11:32:18 -0700 Subject: [PATCH 1/2] [nvbugs/6317600][fix] Avoid CPU sync in gdn_mixer; cap CppMamba dummy reqs by per-window free blocks The Qwen3-Next-80B-A3B-Thinking tp4ep4 hang has two contributing pieces: 1. gdn_mixer.forward used boolean-mask indexing (state_indices_p[~has_initial_states_p]) on a CUDA bool tensor for the prefix-cache state-reset block. That forces a GPU->CPU sync per prefill step (twice per layer, for ssm_states and conv_states) so PyTorch can read the mask reduction count and allocate the output. Combined with TP=4 + EP=4 + the overlap scheduler, the variable per-rank latency of this sync was enough to desync subsequent TP/EP collectives and deadlock the forward pass mid-MMLU on Qwen3-Next-80B-A3B-Thinking. Replace mask indexing with the same pattern already used in mamba2_mixer.py: gate on the host-side use_initial_states flag and fall through to torch.where + index_copy_ for the mixed-batch case. Both paths preserve the original semantic (zero rows whose request has no prior mamba state, keep rows that resume from prefix cache) and have output shapes that do not depend on tensor contents, so no implicit CPU sync is introduced. 2. CppMambaHybridCacheManager.add_dummy_requests delegated straight to the base KVCacheManager without checking the recurrent-states window in the unified C++ KV pool. CudaGraphConfig.batch_sizes (with max_batch_size=720 and enable_padding=True) generates capture batches that can exceed the most-constrained window. _create_cuda_graph_warmup_request only checks get_num_free_blocks (full-attention window), so the recurrent-states window underflows and add_sequence_batch raises 'No free block found' from the C++ side, leaving collectives in an incomplete state. Add an upfront guard over min(num_free_blocks_per_window_size.values()) so oversized warmup batches return None, matching the 'if requests is None: return None' contract expected by _create_cuda_graph_warmup_request in model_engine.py. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mamba/gdn_mixer.py | 17 +++++++++++------ .../_torch/pyexecutor/mamba_cache_manager.py | 10 ++++++++++ 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index aa6146620ead..8a0d2351f345 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -892,12 +892,17 @@ def forward( if num_prefills > 0: # PyExecutor guarantees prefill requests are placed before decode requests has_initial_states_p = has_initial_states[:num_prefills] - ssm_states[state_indices_p[~has_initial_states_p]] = torch.zeros( - (), dtype=ssm_states.dtype, device=ssm_states.device - ) - conv_states[state_indices_p[~has_initial_states_p]] = torch.zeros( - (), dtype=conv_states.dtype, device=conv_states.device - ) + if not mamba_metadata.use_initial_states: + # All prefills are fresh — zero every slot unconditionally. + for state in (ssm_states, conv_states): + state[state_indices_p] = 0 + else: + # Use torch.where so the output shape is data-independent; + # boolean-mask indexing on a CUDA tensor would force a CPU sync. + for state in (ssm_states, conv_states): + kept = state[state_indices_p] + mask = has_initial_states_p.view(-1, *([1] * (kept.ndim - 1))) + state.index_copy_(0, state_indices_p, torch.where(mask, kept, 0)) is_target_verify = ( num_decodes > 0 diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 7f1af74d4c83..d2e37cc12dff 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1676,6 +1676,16 @@ def add_dummy_requests( num_extra_decoding_steps: int = 0, draft_kv_cache_manager: Optional[KVCacheManager] = None, ) -> List[LlmRequest]: + # The caller's get_num_free_blocks check sees only the full-attention + # window, but the unified C++ pool also holds the recurrent-states + # window used by mamba layers. If a CUDA-graph warmup batch exceeds + # that smaller window, add_sequence_batch raises "No free block found" + # from C++ and deadlocks ranks already past the collective. Return + # None so the caller skips this batch (the documented contract). + per_window_free = self.impl.get_kv_cache_stats( + ).num_free_blocks_per_window_size + if per_window_free and len(request_ids) > min(per_window_free.values()): + return None requests = super().add_dummy_requests( request_ids=request_ids, token_nums=token_nums, From 52ed7a1874e2fbc527d74049f65a9885a284cae2 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Thu, 18 Jun 2026 17:14:04 -0700 Subject: [PATCH 2/2] [nvbugs/6317600][fix] Skip TRTLLM-Gen FMHA JIT warmup for oversized engine configs TestQwen3NextThinking::test_auto_dtype[tp4ep4] hangs during engine startup in _run_attention_warmup. The C++ TRTLLM-Gen FMHA JIT warmup enumerates a (batchSize x seqLenKv) cartesian grid sized by engine maxima. For Qwen3-Next-80B-A3B-Thinking with max_batch_size=2048 and max_seq_len=262144, the densified grid pushes warmup TMA descriptor shapes past the flashinfer 2^32 limit and hangs engine startup. Skip the warmup whenever the maxima product exceeds 256 * 16384 (the pre-PR #15305 effective grid size). Any kernel not pre-warmed will JIT-compile lazily on its first real request - correct in all cases, only slightly slower for that first request. Same approach as the sister fix for GPT-OSS-120B (nvbugs/6316980 / nvbugs/6275959). Verified passing: MMLU 85.79% (threshold 84.18%), GSM8K 85.10% (threshold 78.37%), 1 passed in 534s on B200 tp4ep4. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mamba/gdn_mixer.py | 4 ++-- tensorrt_llm/_torch/pyexecutor/model_engine.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 8a0d2351f345..24c990498297 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -890,8 +890,6 @@ def forward( state_indices_p, state_indices_d = torch.split(state_indices, batch_split_size) if num_prefills > 0: - # PyExecutor guarantees prefill requests are placed before decode requests - has_initial_states_p = has_initial_states[:num_prefills] if not mamba_metadata.use_initial_states: # All prefills are fresh — zero every slot unconditionally. for state in (ssm_states, conv_states): @@ -899,6 +897,8 @@ def forward( else: # Use torch.where so the output shape is data-independent; # boolean-mask indexing on a CUDA tensor would force a CPU sync. + # PyExecutor guarantees prefill requests are placed before decode requests. + has_initial_states_p = has_initial_states[:num_prefills] for state in (ssm_states, conv_states): kept = state[state_indices_p] mask = has_initial_states_p.view(-1, *([1] * (kept.ndim - 1))) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4e5cac179c02..0a9c8de66506 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1128,6 +1128,22 @@ def _run_attention_warmup(self, if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata): return + # The C++ TRTLLM-Gen FMHA JIT warmup enumerates a (batchSize x seqLenKv) + # cartesian grid sized by engine maxima. For long-context configs such + # as Qwen3-Next-80B-A3B-Thinking tp4ep4 (max_batch_size=2048, + # max_seq_len=262144), the densified grid pushes warmup TMA descriptor + # shapes past the flashinfer 2^32 limit and hangs engine startup. Skip + # the warmup whenever the maxima product is too large; any kernel not + # pre-warmed JIT-compiles lazily on first request, which is correct + # (just slower for that one request). The threshold matches the + # pre-PR #15305 effective grid size. + if self.batch_size * self.max_seq_len > 256 * 16384: + logger.info( + f"Skipping TRTLLM-Gen FMHA JIT warmup: engine config " + f"(max_batch_size={self.batch_size}, max_seq_len={self.max_seq_len}) " + f"would produce too many warmup grid points") + return + @contextlib.contextmanager def trtllm_gen_fmha_jit_warmup(): previous = self._trtllm_gen_jit_warmup