From e6f5849032deddd8eee331d99242d3ab646198f1 Mon Sep 17 00:00:00 2001 From: halt95 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