From a3a2ebe06c078c5fd311bcd74903b787866c4759 Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 4 Jun 2026 14:19:52 +0800 Subject: [PATCH 1/2] [None][fix] Stabilize Mamba replay state update Signed-off-by: qgai --- .../_torch/modules/mamba/mamba2_mixer.py | 6 +++++- .../mamba/replay_selective_state_update.py | 17 +++++++++-------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 0c75aa58bd77..36aed732cd85 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -349,7 +349,11 @@ def forward( num_prefills] has_initial_states_p = has_initial_states[:num_prefills] - conv_states[state_indices_p[~has_initial_states_p]].zero_() + # Explicitly zero fresh context slots. Recycled slots may still + # carry a prior request's SSM/conv state, which kernels read. + fresh_p = state_indices_p[~has_initial_states_p] + conv_states[fresh_p] = 0 + ssm_states[fresh_p] = 0 # Fused kernel to avoid expensive .contiguous() call in causal_conv1d_fn. xbc_p_t = extract_transpose_xbc_prefill(zxbcdt, num_prefill_tokens, self.tp_d_inner, diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 9c510f8f4956..547cc3153c5f 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -226,7 +226,7 @@ def _replay_precompute_kernel( other=0.0, ) - # Compute raw CB once — shared across all heads in this block + # Compute raw CB once, shared across all heads in this block. raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) # Store B to cache (once per group, only if this block covers the first heads) @@ -458,17 +458,18 @@ def _replay_state_update_kernel( # two back). coeff is all-zero (offs_t < 0), total_decay is 1.0, so the # replay leaves `state` unchanged — cache contents don't matter on step 0. coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - coeff = tl.where(offs_t < prev_num_accepted_tokens, coeff, 0.0) + accepted_mask = t_mask & (offs_t < prev_num_accepted_tokens) + coeff = tl.where(accepted_mask, coeff, 0.0) - # Load old_x: (BLOCK_SIZE_T, BLOCK_SIZE_M) — single-buffered + # Zero stale rows beyond PNAT to prevent Inf/NaN from reaching tl.dot. old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head old_x_all = tl.load( old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=t_mask[:, None] & m_mask[None, :], + mask=accepted_mask[:, None] & m_mask[None, :], other=0.0, ) - # Load old_B from READ buffer: (BLOCK_SIZE_T, BLOCK_SIZE_DSTATE) + # Apply the same accepted-row mask to old_B. old_B_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache @@ -477,7 +478,7 @@ def _replay_state_update_kernel( ) old_B_all = tl.load( old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], + mask=accepted_mask[:, None] & n_mask[None, :], other=0.0, ).to(tl.float32) @@ -488,7 +489,7 @@ def _replay_state_update_kernel( total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) state *= total_decay - # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate) + # tl.dot fast-forward: old_x^T @ dB_scaled -> (M, dstate) state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) # Write post-replay state @@ -771,7 +772,7 @@ def replay_selective_state_update( device = x.device BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) - # Allocate precomputed intermediates (per-call, not cached) + # Allocate precomputed intermediates (per-call, not cached). cb_scaled = torch.empty( batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_T, device=device, dtype=torch.float32 ) From 945c5327ae12bd96d7105cd8d796203e501fbeb5 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 10 Jun 2026 03:06:22 -0700 Subject: [PATCH 2/2] Clarify fresh Mamba prefill state handling Signed-off-by: qgai --- tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py | 2 +- tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py | 7 +------ 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py b/tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py index 83467a06a0d6..cccc3dfc872c 100644 --- a/tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py +++ b/tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py @@ -158,7 +158,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] col0 = tl.load(conv_states_ptrs, mask_w, 0.0) else: - # prior-tokens are zeros + # No cached prefix: start the convolution window from zeros. if KERNEL_WIDTH >= 2: # STRATEGY1 # first chunk and does not have prior-token, so just set to 0 col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 36aed732cd85..b7d24e72f460 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -348,12 +348,6 @@ def forward( has_initial_states = mamba_metadata.has_initial_states[: num_prefills] - has_initial_states_p = has_initial_states[:num_prefills] - # Explicitly zero fresh context slots. Recycled slots may still - # carry a prior request's SSM/conv state, which kernels read. - fresh_p = state_indices_p[~has_initial_states_p] - conv_states[fresh_p] = 0 - ssm_states[fresh_p] = 0 # Fused kernel to avoid expensive .contiguous() call in causal_conv1d_fn. xbc_p_t = extract_transpose_xbc_prefill(zxbcdt, num_prefill_tokens, self.tp_d_inner, @@ -380,6 +374,7 @@ def forward( initial_states = None if mamba_metadata.use_initial_states: + # Rows without cached prefix state start SSM from zero. initial_states = torch.where( has_initial_states[:, None, None, None], ssm_states[state_indices_p], 0)