iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

recoverssm: one fold per program in the GDN RecoverSSM state commit (port of upstream #58863, af0d2b45d7)

Browse Source 
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.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-09 20:08:55 +07:00
co-authored by Hinata
parent eb68fb9105
commit 2c6efaabe8
1 changed files with 25 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py
+25 -26
View File
@@ -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,