125 lines
6.1 KiB
Diff
125 lines
6.1 KiB
Diff
From e6f5849032deddd8eee331d99242d3ab646198f1 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
Date: Thu, 1 Oct 2026 20:14:32 +0200
|
|
Subject: [PATCH 079/103] recoverssm: GDN: commit forward with the verify
|
|
step's own decay, bit-identical to per-token (port of upstream #58863,
|
|
514102a1db)
|
|
|
|
Port of the upstream vLLM change #58863 (Apache-2.0). The verify kernel records the decay it applied and the commit replays the
|
|
accepted tokens forward with the same update, so the committed state equals the per-token state bit for bit.
|
|
---
|
|
tests/kernels/mamba/test_recoverssm_gdn.py | 14 ++++-------
|
|
.../layers/mamba/gdn/recoverssm_gdn.py | 24 ++++++++++---------
|
|
2 files changed, 18 insertions(+), 20 deletions(-)
|
|
|
|
diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py
|
|
index 724a41d..59ddb7f 100644
|
|
--- a/tests/kernels/mamba/test_recoverssm_gdn.py
|
|
+++ b/tests/kernels/mamba/test_recoverssm_gdn.py
|
|
@@ -111,7 +111,7 @@ def test_commit_matches_native_per_token_state(accepted):
|
|
ctx.commit(acc, inp["state_indices"], inp["query_start_loc"])
|
|
for i, blk in enumerate(inp["state_indices"].tolist()):
|
|
tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1
|
|
- torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6)
|
|
+ assert torch.equal(ckpt[blk], states[tok])
|
|
others = [b for b in range(ckpt.shape[0]) if b not in inp["state_indices"].tolist()]
|
|
assert torch.equal(ckpt[others], inp["checkpoint_state"][others])
|
|
|
|
@@ -301,7 +301,7 @@ def test_commit_maps_noncontiguous_requests_in_align_mode():
|
|
c0 = int(nc[r])
|
|
final_blk = int(bt[r, max(c0 + n - 1, 0) // bs])
|
|
tok = int(inp["query_start_loc"][i]) + n - 1
|
|
- torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6)
|
|
+ assert torch.equal(ckpt[final_blk], states[tok])
|
|
|
|
|
|
# (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on
|
|
@@ -343,12 +343,8 @@ def test_align_boundary_and_final_state_match_native(nc, accepted, bs):
|
|
if (
|
|
nc + accepted >= next_boundary
|
|
): # crosses: the boundary state goes into the source block
|
|
- torch.testing.assert_close(
|
|
- ckpt[int(src)], states[next_boundary - nc - 1], rtol=1e-5, atol=1e-6
|
|
- )
|
|
- torch.testing.assert_close(
|
|
- ckpt[final_blk], states[accepted - 1], rtol=1e-5, atol=1e-6
|
|
- )
|
|
+ assert torch.equal(ckpt[int(src)], states[next_boundary - nc - 1])
|
|
+ assert torch.equal(ckpt[final_blk], states[accepted - 1])
|
|
|
|
|
|
@pytest.mark.parametrize("gate", list(GATES))
|
|
@@ -367,7 +363,7 @@ def test_commit_matches_native_under_extreme_gates(gate, seed):
|
|
ctx.commit(acc, inp["state_indices"], inp["query_start_loc"])
|
|
for i, blk in enumerate(inp["state_indices"].tolist()):
|
|
tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1
|
|
- torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6)
|
|
+ assert torch.equal(ckpt[blk], states[tok])
|
|
|
|
|
|
@pytest.mark.parametrize("batch", [1, 2, 4])
|
|
diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py
|
|
index 1404edc..8adc094 100644
|
|
--- a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py
|
|
+++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py
|
|
@@ -9,7 +9,7 @@ GDN port of the Kimi-K3 KDA RecoverSSM path
|
|
scalar per value head (g = -exp(A_log) * softplus(a + dt_bias)), q/k heads are grouped
|
|
over value heads (HV // H), and q/k are L2-normalised inside the kernel when requested.
|
|
The verify step caches, per token, the delta-rule correction c_t = beta_t * (v_t - S'_t
|
|
-k_t) (V floats), the normalised key (K floats) and the log-decay g_t (1 float) in one
|
|
+k_t) (V floats), the normalised key (K floats) and the decay exp(g_t) (1 float) in one
|
|
fp32 "replay" record. Recovery folds the accepted prefix: S_n = exp(sum_t g_t) S_0 +
|
|
sum_t exp(sum_{s>t} g_s) c_t k_t^T and writes the checkpoint once (plus the
|
|
block-boundary state in align mode)."""
|
|
@@ -150,7 +150,8 @@ def _gdn_recoverssm_verify_kernel(
|
|
q = q * tl.rsqrt(tl.sum(q * q) + 1e-6)
|
|
k = k * tl.rsqrt(tl.sum(k * k) + 1e-6)
|
|
q = q * scale
|
|
- new_state = state * tl.exp(g)
|
|
+ decay = tl.exp(g)
|
|
+ new_state = state * decay
|
|
corr = (v - tl.sum(new_state * k[None, :], 1)) * beta
|
|
new_state = new_state + corr[:, None] * k[None, :]
|
|
state = tl.where(valid, new_state, state)
|
|
@@ -164,7 +165,8 @@ def _gdn_recoverssm_verify_kernel(
|
|
tl.store(rec + offs_v * stride_replay_dim, corr, mask=valid & mask_v)
|
|
if pid_v == 0:
|
|
tl.store(rec + (V + offs_k) * stride_replay_dim, k, mask=valid & mask_k)
|
|
- tl.store(rec + (V + K) * stride_replay_dim, g, mask=valid)
|
|
+ # the decay as multiplied here, so the commit replays this exact float
|
|
+ tl.store(rec + (V + K) * stride_replay_dim, decay, mask=valid)
|
|
|
|
|
|
def gdn_recoverssm_verify(
|
|
@@ -399,19 +401,19 @@ def _commit_gdn_state_kernel(
|
|
mask=mask_state,
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
- decay = 1.0
|
|
- corr = tl.zeros([BV, BK], tl.float32)
|
|
- for r in range(n):
|
|
- t = n - r - 1
|
|
+ # Forward over the accepted tokens with the verify step's own update, so the committed
|
|
+ # state is the per-token state bit for bit (the decay and corr stored are the floats
|
|
+ # the verify kernel used).
|
|
+ h = s0
|
|
+ for t in range(n):
|
|
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)
|
|
- corr += (c[:, None] * kk[None, :]) * decay
|
|
- decay *= tl.exp(g)
|
|
+ decay = tl.load(rec + (V + K) * stride_replay_dim)
|
|
+ h = tl.fma(c[:, None], kk[None, :], h * decay)
|
|
tl.store(
|
|
state_ptr + dst_idx * sbs + pid_h * stride_state_head + sp,
|
|
- (s0 * decay + corr).to(state_ref_ptr.dtype.element_ty),
|
|
+ h.to(state_ref_ptr.dtype.element_ty),
|
|
mask=mask_state,
|
|
)
|
|
|
|
--
|
|
2.43.0
|
|
|