From f5b942010a1cce6c0807bcd9cbcc36c1f6bec792 Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 29 Sep 2026 20:19:14 -0700 Subject: [PATCH 076/103] recoverssm: one fold per program in the GDN RecoverSSM state commit (port of upstream #58863, af0d2b45d7) Port of the upstream vLLM change #58863 (Apache-2.0). Each program folds either the final or the block-boundary state, not both, which halves the live accumulator tiles. --- .../layers/mamba/gdn/recoverssm_gdn.py | 51 +++++++++---------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py index d347bef..a95d579 100644 --- a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py @@ -353,6 +353,10 @@ def _commit_gdn_state_kernel( NUM_HEADS: tl.constexpr, ALIGN_MODE: tl.constexpr, ): + # In align mode the second half of axis 0 folds the block-boundary state + # (the window's first boundary_len tokens) instead of the final state, so every + # program runs one fold and keeps one tile live. + NUM_V_TILES: tl.constexpr = (V + BV - 1) // BV pid_v = tl.program_id(0) pid_b = tl.program_id(1) pid_lh = tl.program_id(2) @@ -365,10 +369,17 @@ def _commit_gdn_state_kernel( if n == 0: return final_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) - boundary_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) - boundary_len = tl.load(boundary_recovery_lens_ptr + pid_b) if final_idx <= null_block_id: return + dst_idx = final_idx + # ALIGN_MODE is a constexpr: kept apart so non-align builds drop the branch. + if ALIGN_MODE: # noqa: SIM102 + if pid_v >= NUM_V_TILES: + pid_v -= NUM_V_TILES + dst_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + if dst_idx <= null_block_id: + return + n = tl.load(boundary_recovery_lens_ptr + pid_b) state_ptr = tl.load(state_base_addrs_ptr + pid_l).to( tl.pointer_type(state_ref_ptr.dtype.element_ty) ) @@ -389,35 +400,19 @@ def _commit_gdn_state_kernel( mask=mask_state, other=0.0, ).to(tl.float32) - fin_decay = 1.0 - fin_corr = tl.zeros([BV, BK], tl.float32) - bnd_decay = 1.0 - bnd_corr = tl.zeros([BV, BK], tl.float32) + decay = 1.0 + corr = tl.zeros([BV, BK], tl.float32) for r in range(n): t = n - r - 1 rec = rec_base + t * stride_replay_pos c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0) kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0) g = tl.load(rec + (V + K) * stride_replay_dim) - upd = c[:, None] * kk[None, :] - dec = tl.exp(g) - fin_corr += upd * fin_decay - fin_decay *= dec - if ALIGN_MODE: - before = t < boundary_len - bnd_corr += tl.where(before, upd * bnd_decay, 0.0) - bnd_decay *= tl.where(before, dec, 1.0) - if ALIGN_MODE: - bs = s0 * bnd_decay + bnd_corr - tl.store( - state_ptr + boundary_idx * sbs + pid_h * stride_state_head + sp, - bs.to(state_ref_ptr.dtype.element_ty), - mask=mask_state & (boundary_idx > null_block_id), - ) - fs = s0 * fin_decay + fin_corr + corr += (c[:, None] * kk[None, :]) * decay + decay *= tl.exp(g) tl.store( - state_ptr + final_idx * sbs + pid_h * stride_state_head + sp, - fs.to(state_ref_ptr.dtype.element_ty), + state_ptr + dst_idx * sbs + pid_h * stride_state_head + sp, + (s0 * decay + corr).to(state_ref_ptr.dtype.element_ty), mask=mask_state, ) @@ -659,9 +654,13 @@ class GDNRecoverSSMCommitContext: ref = self.checkpoints[0] _, HV, V, K = ref.shape BK = triton.next_power_of_2(K) - BV = min(triton.next_power_of_2(V), 32) + # 16-row tiles: the commit streams each state once, and smaller tiles + # keep more of them in flight (about 13% more bandwidth than 32 on H200). + BV = min(triton.next_power_of_2(V), 16) rr = self.replays[0] - _commit_gdn_state_kernel[(triton.cdiv(V, BV), batch, num_layers * HV)]( + align = block_table is not None + num_v_tiles = triton.cdiv(V, BV) * (2 if align else 1) + _commit_gdn_state_kernel[(num_v_tiles, batch, num_layers * HV)]( ref, self.state_base_addrs, self.state_block_strides, -- 2.43.0