109 lines
4.6 KiB
Diff
109 lines
4.6 KiB
Diff
From f5b942010a1cce6c0807bcd9cbcc36c1f6bec792 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
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
|
|
|