From c41c573ce9ed2053d38aa711582e19a67a5e9755 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Mon, 20 Jul 2026 22:06:30 +0800 Subject: [PATCH] [GDN] Support ReplaySSM Ring Spec-Verify (#28695) Co-authored-by: luoyuan.luo Co-authored-by: vincentzed <207368749+vincentzed@users.noreply.github.com> --- .../fla/gdn_replayssm_spec_decode.py | 1021 +++++++++++++++++ .../attention/hybrid_linear_attn_backend.py | 29 +- .../layers/attention/linear/gdn_backend.py | 159 ++- .../srt/mem_cache/kv_cache_configurator.py | 8 + python/sglang/srt/mem_cache/memory_pool.py | 224 +++- python/sglang/srt/server_args.py | 117 ++ python/sglang/srt/speculative/spec_utils.py | 67 ++ 7 files changed, 1577 insertions(+), 48 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py diff --git a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py new file mode 100644 index 000000000..d578cd422 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py @@ -0,0 +1,1021 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM speculative-decode verify kernel for GDN (Gated DeltaNet). + +Ported from the vLLM ReplaySSM reference implementation +(github.com/Johnny-Liou/ReplaySSM, commit ``3c85112``, +``vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py``) and the Dao +AI Lab ReplaySSM blog (dao-lab.ai/blog/2026/replayssm). Part B of SGLang RFC +#28511. + +Adapted to SGLang's GDN verify path: the recurrent intermediate-state snapshot +(full ``[V, K]`` state written per draft token to ``intermediate_ssm``) is +replaced by a per-slot **circular cache** of the last ``L`` committed steps' +``(d, k, g)`` records plus a frozen checkpoint ``h0``. The verify output for the +whole draft window is reconstructed *output-only* (never materialising the state) +via the chunked delta-rule ``(I + A)^{-1}`` UT-transform, and the full state is +flushed back only every ``L`` committed tokens. Rejected drafts roll back by a +pointer move on the cursors -- no state write-back. + +Closed-loop exact fold (the state / output error split): + The scheme has two error paths with opposite structure. The OUTPUT is one-shot + (computed per step, consumed by the sampler, discarded) -- its error never + compounds. The STATE accumulates: folding the stored ``d`` records open-loop + (as the vLLM reference does) feeds the chunked transform's cancellation error + -- amplified up to ``2^(BS-1)`` by ``(I + A)^{-1}`` -- plus their storage + quantization into every future window, undamped, so the error grows with + generation length. Therefore the flush here does NOT fold ``d``: instead + :func:`gdn_replayssm_exact_fold_kernel` sequentially replays the committed + window from rings of the RAW inputs (``v`` and pre-norm ``k``, both born in + the activation dtype and hence stored losslessly, plus fp32 ``g`` / ``beta``), + mirroring ``fused_sigmoid_gating_delta_rule_update_kernel``'s fp32 op order + exactly. The delta-rule recurrence is contractive (per step the perturbation + gain is ``exp(g) * |1 - beta| < 1`` along ``k`` and ``exp(g) < 1`` elsewhere), + so given identical inputs the replayed checkpoint is bit-identical to the + recurrent baseline's committed state and carries NO length-dependent error. + The chunked transform is kept only for the non-accumulating output. Its + dots therefore need only stay below the bf16 OUTPUT-cast floor (eps ~ 2^-8 + ~ 4e-3 relative), not match fp32 exactly: ``DOT_PRECISION`` defaults to + ``"tf32"`` (~5e-4, tensor-core path; worst case through the (I+A)^{-1} + amplification 2^(BS-1) still lands at the floor). ``"ieee"`` / ``"tf32x3"`` + remain selectable for ablations. The committed state is untouched by any of + these dots (requires the fp32 SSM checkpoint; enforced in server_args). + +Differences from the vLLM reference: + * SGLang passes **split** ``q`` / ``k`` / ``v`` tensors (already split + post + causal-conv1d in the GDN backend) rather than a packed ``mixed_qkv``. + * State / cache layouts and ``(I + A)^{-1}`` math are unchanged + (SGLang ``ssm_states`` is ``[slots, HV, V, K]`` with ``K`` contiguous, which + matches the reference checkpoint exactly). + +Linear-chain only: the intra-window interaction uses a strictly-lower causal +mask, so this kernel is valid only for linear draft chains +(``speculative_eagle_topk <= 1``, i.e. MTP / frozen-KV-MTP). Tree verify +(topk > 1), KDA, and NPU/CPU fall back to the recurrent verify kernel. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def gdn_replayssm_spec_circular_kernel( + q, # [total_tokens, H, K] + k, # [total_tokens, H, K] + v, # [total_tokens, HV, V] + a, # [total_tokens, HV] + b, # [total_tokens, HV] + A_log, # [HV] fp32 + dt_bias, # [HV] fp32 + o, # [total_tokens, HV, V] preallocated output + h0, # [num_slots, HV, V, K] checkpoint (read-only; folded by the exact-fold kernel) + d_cache, # [num_slots, HV, L, V] chunked deltas (output reconstruction only) + k_cache, # [num_slots, H, L, K] L2-normalized keys (output reconstruction only) + g_cache, # [num_slots, HV, L] fp32 + rawv_cache, # [num_slots, HV, L, V] raw v (exact-fold replay) + rawk_cache, # [num_slots, H, L, K] raw pre-norm k (exact-fold replay) + beta_cache, # [num_slots, HV, L] fp32 beta (exact-fold replay) + query_start_loc, # [B+1] int packed cu_seqlens + ssm_state_indices, # [B] int physical block per request + write_pos, # [num_slots] int32 block-keyed + cache_base, # [num_slots] int32 block-keyed circular origin + is_flush_flags, # [num_slots] int8 block-keyed + scale, + stride_q_t: tl.constexpr, # per-token stride of q (= H*K) + stride_k_t: tl.constexpr, # per-token stride of k (= H*K) + stride_v_t: tl.constexpr, # per-token stride of v (= HV*V) + stride_a_t: tl.constexpr, + stride_b_t: tl.constexpr, + stride_o_t: tl.constexpr, # per-token stride of o (= HV*V) + stride_state_slot: tl.constexpr, + stride_d_slot: tl.constexpr, + stride_k_slot: tl.constexpr, + stride_g_slot: tl.constexpr, + stride_rawv_slot: tl.constexpr, + stride_rawk_slot: tl.constexpr, + stride_beta_slot: tl.constexpr, + stride_qsl: tl.constexpr, + stride_indices: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BS: tl.constexpr, + BC: tl.constexpr, + NK: tl.constexpr, + BKT: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_FLUSH: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_v = tl.program_id(0) + i_n = tl.program_id(1) + i_hv = tl.program_id(2) + i_h = i_hv // (HV // H) + + o_v = i_v * BV + tl.arange(0, BV) + o_s = tl.arange(0, BS) + o_c = tl.arange(0, BC) + mask_v = o_v < V + + # --- per-request packed window --- + bos = tl.load(query_start_loc + i_n * stride_qsl).to(tl.int64) + eos = tl.load(query_start_loc + (i_n + 1) * stride_qsl).to(tl.int64) + spec_len = eos - bos # full window length + # Clamp the token index for padding lanes (o_s >= spec_len, which exist when + # BS = max(bs_min, npow2(max_spec_len)) > spec_len) so masked loads/stores + # never form out-of-bounds token addresses into the packed varlen buffers. + # mask_s still zeroes the *values* of those lanes, so this is numerically a + # no-op -- it only keeps the computed pointers in-bounds. + mask_s = o_s < spec_len + o_s_safe = tl.where(mask_s, o_s, 0) + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + + # output pointer (packed): token (bos + o_s), value-head i_hv, dim o_v + p_o = o + (bos + o_s_safe[:, None]) * stride_o_t + i_hv * V + o_v[None, :] + + if IS_FLUSH: + if state_idx <= NULL_BLOCK_ID: + return + b_is_flush = tl.load(is_flush_flags + state_idx) != 0 + if not b_is_flush: + return + else: + if state_idx <= NULL_BLOCK_ID: + full_mask = (o_s < spec_len)[:, None] & mask_v[None, :] + tl.store( + p_o, + tl.zeros([BS, BV], dtype=tl.float32).to(p_o.dtype.element_ty), + mask=full_mask, + ) + return + b_is_flush = tl.load(is_flush_flags + state_idx) != 0 + if b_is_flush: + return + + b_write_pos = tl.load(write_pos + state_idx).to(tl.int64) + b_cache_base = tl.load(cache_base + state_idx).to(tl.int32) + + out_mask = mask_s[:, None] & mask_v[None, :] + + b_wp_i = b_write_pos.to(tl.int32) + cache_valid = o_c < b_write_pos + + # CIRCULAR physical slots (addresses only; masks/cumsums stay logical). + phys_c = (b_cache_base + o_c) & (MAX_CACHE_LEN - 1) # [BC] history + phys_spec = (b_cache_base + b_wp_i + o_s) & (MAX_CACHE_LEN - 1) # [BS] spec + + # ------------------------------------------------------------------ + # Block 0: gates / beta / local cumsum + committed-history replay decay. + # ------------------------------------------------------------------ + A_log_val = tl.load(A_log + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) + a_s = tl.load(a + (bos + o_s_safe) * stride_a_t + i_hv, mask=mask_s, other=0.0).to( + tl.float32 + ) + b_s = tl.load(b + (bos + o_s_safe) * stride_b_t + i_hv, mask=mask_s, other=0.0).to( + tl.float32 + ) + x = a_s + dt_bias_val + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + g_s = tl.where(mask_s, -tl.exp(A_log_val) * softplus_x, 0.0) + # Explicit sigmoid expression (NOT tl.sigmoid): bitwise clone of + # fused_sigmoid_gating_delta_rule_update_kernel's beta so the fp32 value + # stored to beta_cache is exactly what the recurrent baseline would compute. + beta_s = tl.where(mask_s, 1.0 / (1.0 + tl.exp(-b_s)), 0.0) + G_s = tl.cumsum(g_s, axis=0) + expG_s = tl.exp(G_s) + + if not IS_FLUSH: + # Committed-history replay decay from cached g (history loads -> phys_c). + # Output reconstruction only. On flush steps the history is folded into + # the checkpoint by the exact-fold kernel (launched first), so none of + # this is needed there. + p_g_main = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys_c + b_g_all = tl.load(p_g_main, mask=cache_valid, other=0.0).to(tl.float32) + b_g_prefix = tl.cumsum(b_g_all, axis=0) + b_g_total = tl.sum(b_g_all, axis=0) + b_replay_decay = tl.where(cache_valid, tl.exp(b_g_total - b_g_prefix), 0.0) + b_total_decay = tl.exp(b_g_total) + + p_d_main = d_cache + ( + state_idx * stride_d_slot + + (i_hv * MAX_CACHE_LEN + phys_c[None, :]) * V + + o_v[:, None] + ) + b_d_all = tl.load( + p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0.0 + ).to(tl.float32) + b_d_scaled = (b_d_all * b_replay_decay[None, :]).to(h0.dtype.element_ty) + + if USE_QK_L2NORM_IN_KERNEL: + qnorm_acc = tl.zeros([BS], dtype=tl.float32) + knorm_acc = tl.zeros([BS], dtype=tl.float32) + for kk in range(NK): + o_kt = kk * BKT + tl.arange(0, BKT) + mask_kt = o_kt < K + ld = mask_s[:, None] & mask_kt[None, :] + qn = tl.load( + q + (bos + o_s_safe[:, None]) * stride_q_t + i_h * K + o_kt[None, :], + mask=ld, + other=0.0, + ).to(tl.float32) + knn = tl.load( + k + (bos + o_s_safe[:, None]) * stride_k_t + i_h * K + o_kt[None, :], + mask=ld, + other=0.0, + ).to(tl.float32) + qnorm_acc += tl.sum(qn * qn, axis=1) + knorm_acc += tl.sum(knn * knn, axis=1) + q_rnorm = tl.where(mask_s, 1.0 / tl.sqrt(qnorm_acc + 1e-6), 0.0) + k_rnorm = tl.where(mask_s, 1.0 / tl.sqrt(knorm_acc + 1e-6), 0.0) + else: + q_rnorm = tl.where(mask_s, 1.0, 0.0) + k_rnorm = tl.where(mask_s, 1.0, 0.0) + + # ------------------------------------------------------------------ + # K-Tiled Fused Projection and Intra-Spec Matrices (+ flush) + # ------------------------------------------------------------------ + hw_q = tl.zeros([BV, BS], dtype=tl.float32) + hw_k = tl.zeros([BV, BS], dtype=tl.float32) + if not IS_FLUSH: + scores_q = tl.zeros([BC, BS], dtype=tl.float32) + scores_k = tl.zeros([BC, BS], dtype=tl.float32) + kk_mat = tl.zeros([BS, BS], dtype=tl.float32) + kq_mat = tl.zeros([BS, BS], dtype=tl.float32) + + write_k = (i_v == 0) and (i_hv == i_h * (HV // H)) + + for kk in range(NK): + o_kt = kk * BKT + tl.arange(0, BKT) + mask_kt = o_kt < K + # NOTE: do NOT mask the token (o_s) dimension on these loads. With the + # o_s_safe clamp the padding-row addresses are in-bounds (they read + # token `bos`), and `q_rnorm`/`k_rnorm` are already 0 on padding rows + # (line above), so the multiply below zeroes them. Masking the token + # dimension here instead leaves the tl.dot shared-memory staging of the + # padding rows uninitialised -- on this Triton/MMA path that stale + # staging leaks into the *valid* output columns (manifests at + # spec_len=2). Loading the full (token-unmasked) tile keeps staging + # fully written; the result is numerically identical. + q_tile = tl.load( + q + (bos + o_s_safe[:, None]) * stride_q_t + i_h * K + o_kt[None, :], + mask=mask_kt[None, :], + other=0.0, + ).to(tl.float32) + k_tile = tl.load( + k + (bos + o_s_safe[:, None]) * stride_k_t + i_h * K + o_kt[None, :], + mask=mask_kt[None, :], + other=0.0, + ).to(tl.float32) + # Keep the raw (pre-norm) key tile for the exact-fold ring: the replay + # recomputes the L2 norm in fp32 exactly as the recurrent kernel does + # (its division form differs bitwise from the reciprocal-multiply below). + k_raw_tile = k_tile + q_tile = (q_tile * (q_rnorm * scale)[:, None]).to(h0.dtype.element_ty) + k_tile = (k_tile * k_rnorm[:, None]).to(h0.dtype.element_ty) + + p_h0 = ( + h0 + + state_idx * stride_state_slot + + i_hv * V * K + + o_v[:, None] * K + + o_kt[None, :] + ) + sc_tile = tl.load(p_h0, mask=mask_v[:, None] & mask_kt[None, :], other=0.0).to( + h0.dtype.element_ty + ) + + qT = tl.trans(q_tile) + kT = tl.trans(k_tile) + kk_mat += tl.dot(k_tile, kT, input_precision=DOT_PRECISION) + kq_mat += tl.dot(k_tile, qT, input_precision=DOT_PRECISION) + + # Checkpoint projection. On flush steps the exact-fold kernel (launched + # first) has already folded the committed history into h0, so this is + # the whole non-window contribution and no d-fold happens here anymore. + hw_q += tl.dot(sc_tile, qT, input_precision=DOT_PRECISION) + hw_k += tl.dot(sc_tile, kT, input_precision=DOT_PRECISION) + + if not IS_FLUSH: + # cached-key history load -> phys_c (output reconstruction only) + p_k = k_cache + ( + state_idx * stride_k_slot + + (i_h * MAX_CACHE_LEN + phys_c[:, None]) * K + + o_kt[None, :] + ) + khist_tile = tl.load( + p_k, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0 + ).to(h0.dtype.element_ty) + scores_q += tl.dot(khist_tile, qT, input_precision=DOT_PRECISION) + scores_k += tl.dot(khist_tile, kT, input_precision=DOT_PRECISION) + + if write_k: + spec_kt_mask = ( + mask_s[:, None] + & mask_kt[None, :] + & ((b_write_pos + o_s[:, None]) < MAX_CACHE_LEN) + ) + # raw pre-norm key for the exact-fold replay -> phys_spec (circular). + # Born in the activation dtype, so the store round-trips losslessly. + p_cur_rawk = rawk_cache + ( + state_idx * stride_rawk_slot + + (i_h * MAX_CACHE_LEN + phys_spec[:, None]) * K + + o_kt[None, :] + ) + tl.store( + p_cur_rawk, + k_raw_tile.to(p_cur_rawk.dtype.element_ty), + mask=spec_kt_mask, + ) + # normalized spec key store -> phys_spec (circular) + p_cur_k = k_cache + ( + state_idx * stride_k_slot + + (i_h * MAX_CACHE_LEN + phys_spec[:, None]) * K + + o_kt[None, :] + ) + tl.store( + p_cur_k, + k_tile, + mask=spec_kt_mask, + ) + + if not IS_FLUSH: + hw_q = b_total_decay * hw_q + tl.dot( + b_d_scaled, scores_q.to(b_d_scaled.dtype), input_precision=DOT_PRECISION + ) + hw_k = b_total_decay * hw_k + tl.dot( + b_d_scaled, scores_k.to(b_d_scaled.dtype), input_precision=DOT_PRECISION + ) + + # ------------------------------------------------------------------ + # strictly-lower A and T = (I + A)^{-1}. + # ------------------------------------------------------------------ + lower = (o_s[:, None] > o_s[None, :]) & mask_s[:, None] & mask_s[None, :] + diff_ij = G_s[:, None] - G_s[None, :] + A_mat = tl.where(lower, beta_s[:, None] * tl.exp(diff_ij) * kk_mat, 0.0) + + b_Ai = -A_mat + for ii in range(2, BS): + row = tl.sum(tl.where((o_s == ii)[:, None], -A_mat, 0.0), axis=0) + row = tl.where(o_s < ii, row, 0.0) + row = row + tl.sum(row[:, None] * b_Ai, axis=0) + b_Ai = tl.where((o_s == ii)[:, None], row, b_Ai) + T_mat = b_Ai + (o_s[:, None] == o_s[None, :]).to(tl.float32) + + # ------------------------------------------------------------------ + # R and D_spec = R @ T^T. + # ------------------------------------------------------------------ + p_v = v + (bos + o_s_safe[None, :]) * stride_v_t + i_hv * V + o_v[:, None] + v_tile = tl.load(p_v, mask=mask_v[:, None] & mask_s[None, :], other=0.0).to( + tl.float32 + ) + R_mat = beta_s[None, :] * (v_tile - expG_s[None, :] * hw_k) + D_spec = tl.zeros([BV, BS], dtype=tl.float32) + for j in tl.static_range(BS): + Rj = tl.sum(tl.where((o_s == j)[None, :], R_mat, 0.0), axis=1) + Tj = tl.sum(tl.where((o_s == j)[None, :], T_mat, 0.0), axis=1) + D_spec += Rj[:, None] * Tj[None, :] + + # ------------------------------------------------------------------ + # outputs. + # ------------------------------------------------------------------ + causalF = (o_s[:, None] <= o_s[None, :]) & mask_s[:, None] & mask_s[None, :] + diff_ji = G_s[None, :] - G_s[:, None] + F_mat = tl.where(causalF, tl.exp(diff_ji) * kq_mat, 0.0) + DF = tl.zeros([BV, BS], dtype=tl.float32) + for j in tl.static_range(BS): + Dj = tl.sum(tl.where((o_s == j)[None, :], D_spec, 0.0), axis=1) + Fj = tl.sum(tl.where((o_s == j)[:, None], F_mat, 0.0), axis=0) + DF += Dj[:, None] * Fj[None, :] + O_tile = expG_s[None, :] * hw_q + DF + + tl.store(p_o, tl.trans(O_tile).to(p_o.dtype.element_ty), mask=out_mask) + + # ------------------------------------------------------------------ + # write speculative d / raw-v / g / beta at circular positions (phys_spec). + # ------------------------------------------------------------------ + spec_pos = b_write_pos + o_s + spec_store_mask = mask_s & (spec_pos < MAX_CACHE_LEN) + p_cur_d = d_cache + ( + state_idx * stride_d_slot + + (i_hv * MAX_CACHE_LEN + phys_spec[None, :]) * V + + o_v[:, None] + ) + tl.store( + p_cur_d, + D_spec.to(p_cur_d.dtype.element_ty), + mask=mask_v[:, None] & spec_store_mask[None, :], + ) + # raw v for the exact-fold replay (activation dtype: lossless round-trip) + p_cur_v = rawv_cache + ( + state_idx * stride_rawv_slot + + (i_hv * MAX_CACHE_LEN + phys_spec[None, :]) * V + + o_v[:, None] + ) + tl.store( + p_cur_v, + v_tile.to(p_cur_v.dtype.element_ty), + mask=mask_v[:, None] & spec_store_mask[None, :], + ) + if i_v == 0: + p_cur_g = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys_spec + tl.store(p_cur_g, g_s, mask=spec_store_mask) + p_cur_beta = ( + beta_cache + state_idx * stride_beta_slot + i_hv * MAX_CACHE_LEN + phys_spec + ) + tl.store(p_cur_beta, beta_s, mask=spec_store_mask) + + +@triton.jit +def gdn_replayssm_exact_fold_kernel( + h0, # [num_slots, HV, V, K] fp32 checkpoint (in-place) + rawv_cache, # [num_slots, HV, L, V] raw v + rawk_cache, # [num_slots, H, L, K] raw pre-norm k + g_cache, # [num_slots, HV, L] fp32 + beta_cache, # [num_slots, HV, L] fp32 + ssm_state_indices, # [B] int physical block per request + write_pos, # [num_slots] int32 block-keyed + cache_base, # [num_slots] int32 block-keyed + is_flush_flags, # [num_slots] int8 block-keyed + stride_state_slot: tl.constexpr, + stride_rawv_slot: tl.constexpr, + stride_rawk_slot: tl.constexpr, + stride_g_slot: tl.constexpr, + stride_beta_slot: tl.constexpr, + stride_indices: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + """Closed-loop exact fold: sequentially replay the committed ring window + into the fp32 checkpoint on flush rows. + + BITWISE CLONE of ``fused_sigmoid_gating_delta_rule_update_kernel``'s state + recurrence (fused_sigmoid_gating_recurrent.py) -- same [BK, BV] register + tile with the K axis as rows, same 1-D full-K reductions, the same + division-form L2 norm with eps inside the sqrt, and the same + decay -> delta -> rank-1-update op order. Given identical inputs the + committed state matches the recurrent baseline bit-for-bit, which is the + whole point: do NOT "optimize" this into tl.dot / reciprocal-multiply / + reordered expressions, and keep the launch config (BV=32, num_warps=1) + identical to the recurrent kernel so the reduction trees agree. + + Replaces the previous open-loop d-fold, whose per-flush error (chunked + (I+A)^{-1} cancellation + d-storage quantization) fed forward across + flushes undamped and grew with generation length. + """ + i_v = tl.program_id(0) + i_n = tl.program_id(1) + i_hv = tl.program_id(2) + i_h = i_hv // (HV // H) + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + if state_idx <= NULL_BLOCK_ID: + return + if tl.load(is_flush_flags + state_idx) == 0: + return + b_write_pos = tl.load(write_pos + state_idx).to(tl.int32) + if b_write_pos <= 0: + return + b_cache_base = tl.load(cache_base + state_idx).to(tl.int32) + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + # [BK, BV] tile, K rows / V columns -- the recurrent kernel's layout + # (memory offset v * K + k, K contiguous). + p_h0 = ( + h0 + + state_idx * stride_state_slot + + i_hv * V * K + + o_v[None, :] * K + + o_k[:, None] + ) + b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for t in range(0, b_write_pos): + phys = ((b_cache_base + t) & (MAX_CACHE_LEN - 1)).to(tl.int64) + b_k = tl.load( + rawk_cache + + state_idx * stride_rawk_slot + + (i_h * MAX_CACHE_LEN + phys) * K + + o_k, + mask=mask_k, + other=0.0, + ).to(tl.float32) + b_v = tl.load( + rawv_cache + + state_idx * stride_rawv_slot + + (i_hv * MAX_CACHE_LEN + phys) * V + + o_v, + mask=mask_v, + other=0.0, + ).to(tl.float32) + b_g = tl.load( + g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys + ).to(tl.float32) + b_beta = tl.load( + beta_cache + state_idx * stride_beta_slot + i_hv * MAX_CACHE_LEN + phys + ).to(tl.float32) + + # --- verbatim recurrent update (see the clone note above) --- + if USE_QK_L2NORM_IN_KERNEL: + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) + b_h *= tl.exp(b_g) + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= b_beta + b_h += b_k[:, None] * b_v[None, :] + + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +@triton.jit +def _advance_gdn_spec_cursors_kernel( + write_pos_ptr, + cache_base_ptr, + is_flush_ptr, + num_accepted_ptr, + state_batch_indices_ptr, + n_rows, + stride_sbi: tl.constexpr, + stride_na: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + MAX_SPEC_LEN: tl.constexpr, + CACHE_BUF_LEN: tl.constexpr, + BLOCK: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + row_mask = offs < n_rows + blk = tl.load( + state_batch_indices_ptr + offs * stride_sbi, mask=row_mask, other=NULL_BLOCK_ID + ).to(tl.int64) + valid = row_mask & (blk > NULL_BLOCK_ID) + + write_pos = tl.load(write_pos_ptr + blk, mask=valid, other=0).to(tl.int32) + cache_base = tl.load(cache_base_ptr + blk, mask=valid, other=0).to(tl.int32) + is_flush_cur = tl.load(is_flush_ptr + blk, mask=valid, other=0).to(tl.int32) + num_acc = tl.load(num_accepted_ptr + offs * stride_na, mask=valid, other=0).to( + tl.int32 + ) + + total_commit = num_acc + flush_now = (total_commit > 0) & (is_flush_cur != 0) + + new_base = tl.where( + flush_now, (cache_base + write_pos) & (CACHE_BUF_LEN - 1), cache_base + ) + new_wp = tl.where(is_flush_cur != 0, total_commit, write_pos + total_commit).to( + tl.int32 + ) + # EARLY-FLUSH (margin = 2 * max_spec_len, strict '>'): flush one window early + # so that on every verify step write_pos + spec_len <= max_cache_len holds, + # i.e. the spec window NEVER overflows the circular cache. This is required + # for e2e correctness: the proposer (n-gram / MTP) cannot be told to cap its + # draft count, and the rejection sampler reads logits for EVERY window + # position -- so an overflowing position would feed the sampler an + # uninitialized-`out` garbage logit (emitting a wrong token) and desync the + # committed state. The strict '>' uses the buffer exactly (max write_pos at a + # flush step = max_cache_len - max_spec_len, zero headroom); usable committed + # history is max_cache_len - 2*max_spec_len + 1; raise the cache length + # for more. Config enforces max_cache_len >= 2 * max_spec_len. + next_is_flush = ((new_wp + 2 * MAX_SPEC_LEN) > MAX_CACHE_LEN).to(tl.int8) + + tl.store(write_pos_ptr + blk, new_wp, mask=valid) + tl.store(cache_base_ptr + blk, new_base, mask=valid) + tl.store(is_flush_ptr + blk, next_is_flush, mask=valid) + + +@triton.jit +def _reset_gdn_replayssm_spec_cursors_kernel( + write_pos_ptr, + cache_base_ptr, + is_flush_ptr, + do_reset_ptr, + state_batch_indices_ptr, + n_rows, + stride_sbi: tl.constexpr, + stride_reset: tl.constexpr, + INIT_FLUSH: tl.constexpr, + BLOCK: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + row_mask = offs < n_rows + blk = tl.load( + state_batch_indices_ptr + offs * stride_sbi, mask=row_mask, other=NULL_BLOCK_ID + ).to(tl.int64) + do_reset = tl.load(do_reset_ptr + offs * stride_reset, mask=row_mask, other=0).to( + tl.int32 + ) + do = row_mask & (blk > NULL_BLOCK_ID) & (do_reset != 0) + + tl.store(write_pos_ptr + blk, tl.zeros_like(blk).to(tl.int32), mask=do) + tl.store(cache_base_ptr + blk, tl.zeros_like(blk).to(tl.int32), mask=do) + tl.store( + is_flush_ptr + blk, + tl.full([BLOCK], INIT_FLUSH, dtype=tl.int8), + mask=do, + ) + + +# --------------------------------------------------------------------------- +# Python wrappers. +# --------------------------------------------------------------------------- +def _launch_gdn_spec( + q, + k, + v, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + rawv_cache, + rawk_cache, + beta_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + max_cache_len, + max_spec_len, + use_qk_l2norm_in_kernel, + is_flush_kernel, + block_v, + num_warps, + num_stages, + nk, + bs_min, + null_block_id, + dot_precision, +): + num_slots, HV, V, K = checkpoint_state.shape + H = k.shape[1] + B = query_start_loc.shape[0] - 1 + assert ( + max_cache_len & (max_cache_len - 1) == 0 + ), "circular cache requires power-of-two max_cache_len" + assert d_cache.shape[2] == max_cache_len + + BK = triton.next_power_of_2(K) + if triton.cdiv(K, BK) != 1: + raise ValueError(f"only NK_global=1 supported (K={K}, BK={BK}).") + if BK % nk != 0: + raise ValueError(f"nk={nk} must divide BK={BK}.") + BKT = BK // nk + if BKT < 16: + raise ValueError(f"BKT={BKT} must be >=16 for tl.dot.") + BV = block_v if block_v is not None else min(triton.next_power_of_2(V), 64) + BS = max(bs_min, triton.next_power_of_2(max_spec_len)) + BC = max(16, triton.next_power_of_2(max_cache_len)) + + grid = (triton.cdiv(V, BV), B, HV) + gdn_replayssm_spec_circular_kernel[grid]( + q, + k, + v, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + rawv_cache, + rawk_cache, + beta_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + q.stride(0), + k.stride(0), + v.stride(0), + a.stride(0), + b.stride(0), + out.stride(0), + checkpoint_state.stride(0), + d_cache.stride(0), + k_cache.stride(0), + g_cache.stride(0), + rawv_cache.stride(0), + rawk_cache.stride(0), + beta_cache.stride(0), + query_start_loc.stride(0), + ssm_state_indices.stride(0), + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + BS=BS, + BC=BC, + NK=nk, + BKT=BKT, + MAX_CACHE_LEN=max_cache_len, + SOFTPLUS_THRESHOLD=20.0, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + IS_FLUSH=is_flush_kernel, + NULL_BLOCK_ID=null_block_id, + DOT_PRECISION=dot_precision, + num_warps=num_warps, + num_stages=num_stages, + ) + + +def _launch_gdn_exact_fold( + checkpoint_state, + rawv_cache, + rawk_cache, + g_cache, + beta_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + max_cache_len, + use_qk_l2norm_in_kernel, + null_block_id, + num_k_heads, +): + """Launch the closed-loop exact fold (flush rows only; device-routed). + + Tiling clones the recurrent kernel exactly (full-K rows, BV = min(np2(V), 32) + columns, num_warps=1, num_stages=3) so every reduction tree matches + ``fused_sigmoid_gating_delta_rule_update`` and the folded checkpoint is + bit-identical to the recurrent baseline's committed state. + """ + num_slots, HV, V, K = checkpoint_state.shape + B = query_start_loc.shape[0] - 1 + BK = triton.next_power_of_2(K) + BV = min(triton.next_power_of_2(V), 32) + grid = (triton.cdiv(V, BV), B, HV) + gdn_replayssm_exact_fold_kernel[grid]( + checkpoint_state, + rawv_cache, + rawk_cache, + g_cache, + beta_cache, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + checkpoint_state.stride(0), + rawv_cache.stride(0), + rawk_cache.stride(0), + g_cache.stride(0), + beta_cache.stride(0), + ssm_state_indices.stride(0), + H=num_k_heads, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + MAX_CACHE_LEN=max_cache_len, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + NULL_BLOCK_ID=null_block_id, + num_warps=1, + num_stages=3, + ) + + +def gdn_replayssm_spec_decode( + q: torch.Tensor, # [total_tokens, H, K] post-conv + k: torch.Tensor, # [total_tokens, H, K] post-conv + v: torch.Tensor, # [total_tokens, HV, V] post-conv + a: torch.Tensor, # [total_tokens, HV] + b: torch.Tensor, # [total_tokens, HV] + A_log: torch.Tensor, # [HV] fp32 + dt_bias: torch.Tensor, # [HV] fp32 + checkpoint_state: torch.Tensor, # [num_slots, HV, V, K] fp32 (folded in-place) + d_cache: torch.Tensor, # [num_slots, HV, L, V] + k_cache: torch.Tensor, # [num_slots, H, L, K] + g_cache: torch.Tensor, # [num_slots, HV, L] fp32 + rawv_cache: torch.Tensor, # [num_slots, HV, L, V] raw v (exact fold) + rawk_cache: torch.Tensor, # [num_slots, H, L, K] raw pre-norm k (exact fold) + beta_cache: torch.Tensor, # [num_slots, HV, L] fp32 beta (exact fold) + out: torch.Tensor, # [total_tokens, HV, V] preallocated + query_start_loc: torch.Tensor, # [B+1] int + ssm_state_indices: torch.Tensor, # [B] int physical block per request + write_pos: torch.Tensor, # [num_slots] int32 block-keyed + cache_base: torch.Tensor, # [num_slots] int32 block-keyed + is_flush: torch.Tensor, # [num_slots] int8 block-keyed + max_cache_len: int, + max_spec_len: int, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + null_block_id: int = 0, + block_v: int = 64, + num_warps: int = 1, + num_stages: int = 2, + nk: int = 2, + bs_min: int = 4, + block_v_flush: int = 64, + num_warps_flush: int = 1, + num_stages_flush: int = 2, + nk_flush: int = 2, + launch_mode: str = "both", + dot_precision: str = "tf32", +): + """GDN cached speculative-decode on a CIRCULAR ring cache (split-qkv varlen). + + Three launches, all device-routed per row so the step stays CUDA-graph + capturable: + 1. exact fold (flush rows): sequentially replay the committed ring window + (raw v / raw k / g / beta) into the fp32 checkpoint, bit-identical to + the recurrent baseline (closed loop -- no accumulating state error); + 2. verify (non-flush rows): chunked output from checkpoint + d/k history; + 3. flush-output (flush rows): chunked output from the freshly folded + checkpoint (statically empty history). + Cursors are block-keyed (indexed by ``ssm_state_indices``) and advanced + out-of-kernel by :func:`commit_gdn_replayssm_spec`. + """ + if scale is None: + scale = checkpoint_state.shape[-1] ** -0.5 + if is_flush.dtype != torch.int8: + is_flush = is_flush.to(torch.int8) + + if launch_mode in ("both", "flush"): + # Must precede the flush-output launch: it reads the folded checkpoint. + _launch_gdn_exact_fold( + checkpoint_state, + rawv_cache, + rawk_cache, + g_cache, + beta_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + max_cache_len, + use_qk_l2norm_in_kernel, + null_block_id, + k.shape[1], + ) + if launch_mode in ("both", "verify"): + _launch_gdn_spec( + q, + k, + v, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + rawv_cache, + rawk_cache, + beta_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + max_cache_len, + max_spec_len, + use_qk_l2norm_in_kernel, + False, + block_v, + num_warps, + num_stages, + nk, + bs_min, + null_block_id, + dot_precision, + ) + if launch_mode in ("both", "flush"): + _launch_gdn_spec( + q, + k, + v, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + rawv_cache, + rawk_cache, + beta_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + max_cache_len, + max_spec_len, + use_qk_l2norm_in_kernel, + True, + block_v_flush, + num_warps_flush, + num_stages_flush, + nk_flush, + bs_min, + null_block_id, + dot_precision, + ) + return out + + +def commit_gdn_replayssm_spec( + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, + num_accepted: torch.Tensor, # [n_rows] int (already includes the bonus token) + state_batch_indices: torch.Tensor, # [n_rows] int physical block per row + max_cache_len: int, + max_spec_len: int, + cache_buf_len: int | None = None, + null_block_id: int = 0, +): + """Advance the block-keyed cursors once per decode step (device-only).""" + if cache_buf_len is None: + cache_buf_len = max_cache_len + n_rows = state_batch_indices.shape[0] + BLOCK = triton.next_power_of_2(max(1, n_rows)) + _advance_gdn_spec_cursors_kernel[(1,)]( + write_pos, + cache_base, + is_flush, + num_accepted, + state_batch_indices, + n_rows, + stride_sbi=state_batch_indices.stride(0), + stride_na=num_accepted.stride(0), + MAX_CACHE_LEN=max_cache_len, + MAX_SPEC_LEN=max_spec_len, + CACHE_BUF_LEN=cache_buf_len, + BLOCK=BLOCK, + NULL_BLOCK_ID=null_block_id, + ) + + +def reset_gdn_replayssm_spec_cursors( + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, + do_reset: torch.Tensor, # [n_rows] int/bool 1 for first-decode rows + state_batch_indices: torch.Tensor, # [n_rows] int + max_cache_len: int, + max_spec_len: int, + null_block_id: int = 0, +): + """Reset the cursors of first-decode rows (prefill->decode handoff).""" + n_rows = state_batch_indices.shape[0] + BLOCK = triton.next_power_of_2(max(1, n_rows)) + # Early-flush margin (2 * max_spec_len, strict '>'): mirror _advance_gdn_spec_cursors_kernel. + init_flush = 1 if 2 * max_spec_len > max_cache_len else 0 + _reset_gdn_replayssm_spec_cursors_kernel[(1,)]( + write_pos, + cache_base, + is_flush, + do_reset, + state_batch_indices, + n_rows, + stride_sbi=state_batch_indices.stride(0), + stride_reset=do_reset.stride(0), + INIT_FLUSH=init_flush, + BLOCK=BLOCK, + NULL_BLOCK_ID=null_block_id, + ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 19119bcf3..6e9e57adf 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -94,10 +94,16 @@ class MambaAttnBackendBase(AttentionBackend): ) # The ring cursor is a per-slot decode counter shared by all GDN layers; # manage it once here (snapshot, hand to layers, advance mod L), not per-layer. + # Gate on the linear_replayssm FLAG, not on cursor-tensor presence: the + # spec-verify ring (--enable-gdn-replayssm-spec) shares the write_pos + # allocation but owns it exclusively via commit_gdn_replayssm_spec + # (advance-by-accept-count once per verify step). Advancing it here as + # well inserts one phantom/stale ring entry per step and cumulatively + # poisons the reconstruction (degenerate repetition at 10k+ tokens). mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None) write_pos_buf = ( - getattr(mamba_pool, "replayssm_write_pos", None) - if mamba_pool is not None + mamba_pool.replayssm_write_pos + if mamba_pool is not None and mamba_pool.enable_linear_replayssm else None ) if write_pos_buf is not None: @@ -323,12 +329,20 @@ class MambaAttnBackendBase(AttentionBackend): ) def _replayssm_enabled(self) -> bool: - """True iff --enable-linear-replayssm allocated the ring cursor - (MambaPool.replayssm_write_pos doubles as the on/off gate).""" + """True iff --enable-linear-replayssm is on for this pool. + + Gate on the FLAG, not on ``replayssm_write_pos is not None``: the + spec-verify ring (--enable-gdn-replayssm-spec) also allocates the + cursor tensor but owns it exclusively via commit_gdn_replayssm_spec. + The decode-ring metadata machinery gated here (per-bs static cursor + buffers, the per-replay snapshot + advance-by-one in _replay_metadata, + and the decode-kernel ring rerouting downstream) must stay fully + dormant for the spec ring. + """ mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None) if mamba_pool is None: return False - return getattr(mamba_pool, "replayssm_write_pos", None) is not None + return bool(mamba_pool.enable_linear_replayssm) def _replayssm_track_flush_mask( self, seq_lens_cpu: torch.Tensor, bs: int @@ -544,7 +558,10 @@ class MambaAttnBackendBase(AttentionBackend): static_ff.copy_(force_flush_dev) else: static_ff.zero_() - if not in_capture: + # Defense in depth: the decode-ring advance is only meaningful for + # decode/idle forwards (mirrors the eager path's gating). A + # TARGET_VERIFY replay must never advance the cursor. + if not in_capture and forward_mode.is_decode_or_idle(): L = mamba_pool.linear_replayssm_cache_len # Advance only valid (non-padded) slots; a forced flush empties # the ring -> next write_pos 0, like the natural L-1 wrap. diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index ac5525ec1..b535666d6 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -565,22 +565,62 @@ class GDNAttnBackend(MambaAttnBackendBase): value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim) if is_target_verify: - core_attn_out = self.kernel_dispatcher.target_verify( - A_log=layer.A_log, - dt_bias=layer.dt_bias, - q=query, - k=key, - v=value, - a=a, - b=b, - ssm_states=ssm_states, - cache_indices=cache_indices, - query_start_loc=query_start_loc, - intermediate_states_buffer=intermediate_state_cache, - intermediate_state_indices=intermediate_state_indices, - cache_steps=forward_batch.spec_info.draft_token_num, - retrieve_parent_token=retrieve_parent_token, + # ReplaySSM spec-verify (Part B of #28511): when the per-slot ring is + # allocated (--enable-gdn-replayssm-spec, GDN + linear-chain topk<=1), + # reconstruct the verify output for the whole draft window from the + # frozen checkpoint (`temporal`) + the per-slot circular (d, k, g) ring + # instead of the recurrent verify that snapshots a full state per draft + # token. The cursors are advanced once per decode step by the worker + # (commit_gdn_replayssm_spec in spec_utils). GDN-only: KDA (per-K gate) + # routes through kda_backend and never reaches here; we additionally + # guard on `not replayssm_is_kda` for safety. Falls back to the + # recurrent verify when the ring is absent. + mamba_pool = self.req_to_token_pool.mamba_pool + use_replayssm_spec = ( + mamba_cache_params.replayssm_d is not None + and getattr(mamba_pool, "replayssm_cache_base", None) is not None + and not getattr(mamba_pool, "replayssm_is_kda", False) ) + if use_replayssm_spec: + core_attn_out = self._replayssm_target_verify( + layer=layer, + query=query, + key=key, + value=value, + a=a, + b=b, + mamba_pool=mamba_pool, + layer_cache=mamba_cache_params, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + draft_token_num=forward_batch.spec_info.draft_token_num, + ) + else: + # The recurrent fallback needs the per-draft snapshots, which + # the pool gates OFF under --enable-gdn-replayssm-spec (the + # same flag that makes `use_replayssm_spec` true above), so + # this branch is unreachable with a None buffer by + # construction -- keep it loud rather than silently frozen. + assert intermediate_state_cache is not None, ( + "recurrent target_verify fallback requires intermediate_ssm, " + "which is not allocated under --enable-gdn-replayssm-spec" + ) + core_attn_out = self.kernel_dispatcher.target_verify( + A_log=layer.A_log, + dt_bias=layer.dt_bias, + q=query, + k=key, + v=value, + a=a, + b=b, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + intermediate_states_buffer=intermediate_state_cache, + intermediate_state_indices=intermediate_state_indices, + cache_steps=forward_batch.spec_info.draft_token_num, + retrieve_parent_token=retrieve_parent_token, + ) else: g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend( @@ -612,3 +652,92 @@ class GDNAttnBackend(MambaAttnBackendBase): ) return core_attn_out + + def _replayssm_target_verify( + self, + *, + layer: RadixLinearAttention, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + mamba_pool: MambaPool, + layer_cache: "MambaPool.SpeculativeState", + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + draft_token_num: int, + ) -> torch.Tensor: + """ReplaySSM GDN spec-verify (Part B of #28511). + + Reconstructs the verify output for the whole draft window from the frozen + checkpoint (``temporal``) + the per-slot circular ``(d, k, g)`` ring, and + appends this window's drafts to the rings (chunked ``d`` for output + reconstruction; raw ``v`` / pre-norm ``k`` / fp32 ``beta`` for the + closed-loop exact fold that replays the recurrent update into the fp32 + checkpoint at flush). The rings are PER-LAYER + (sliced via ``mamba2_layer_cache``), while the cursors (write_pos, + cache_base, is_flush) are PER-SLOT pool attributes shared by all GDN layers + of the step; the cursors persist across steps and are advanced once per step + by the worker (commit_gdn_replayssm_spec) -- here we only read them and + write this step's ring entries. GDN has K == V, so ``temporal`` + ([slots, HV, K, V]) is consumed directly as the kernel's [slots, HV, V, K] + checkpoint. + """ + from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_decode import ( + gdn_replayssm_spec_decode, + ) + + H, K = layer.num_k_heads, layer.head_k_dim + HV, V = layer.num_v_heads, layer.head_v_dim + # q/k/v may be [1, seq, *] (fallback split) or [seq, *] (fused split); + # derive the packed token count from numel so both layouts flatten. + seq_len = query.numel() // (H * K) + q = query.reshape(seq_len, H, K) + k = key.reshape(seq_len, H, K) + v = value.reshape(seq_len, HV, V) + a = a.reshape(seq_len, HV) + b = b.reshape(seq_len, HV) + d_cache = layer_cache.replayssm_d # [slots, HV, L, V] + max_cache_len = d_cache.shape[-2] # ring length L + out = q.new_empty(seq_len, HV, V) + gdn_replayssm_spec_decode( + q=q, + k=k, + v=v, + a=a, + b=b, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + checkpoint_state=layer_cache.temporal, + d_cache=d_cache, + k_cache=layer_cache.replayssm_k, + g_cache=layer_cache.replayssm_g, + # Closed-loop exact-fold rings: raw v / raw pre-norm k / fp32 beta. + # The flush replays these through the recurrent update (bit-identical + # to the recurrent baseline) instead of folding `d` open-loop. + rawv_cache=layer_cache.replayssm_rawv, + rawk_cache=layer_cache.replayssm_rawk, + beta_cache=layer_cache.replayssm_beta, + out=out, + query_start_loc=query_start_loc, + ssm_state_indices=cache_indices, + # Per-slot cursors live on the pool (shared across all GDN layers), + # NOT in forward_metadata: the verify kernel reads/writes them + # block-keyed via ssm_state_indices and must NOT advance write_pos + # (the worker does that after acceptance), so the decode-path + # forward_metadata.replayssm_write_pos snapshot is not used here. + write_pos=mamba_pool.replayssm_write_pos, + cache_base=mamba_pool.replayssm_cache_base, + is_flush=mamba_pool.replayssm_is_flush, + max_cache_len=max_cache_len, + max_spec_len=draft_token_num, + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + # SGLang marks invalid/padding requests with a negative mamba slot + # index (valid slots start at 0), so the kernel's "null block" + # sentinel is -1, not the vLLM default of 0. + null_block_id=-1, + ) + # Match the recurrent target_verify output shape (== value.shape). + return out.reshape(value.shape) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 5b406127b..526098cdc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -723,6 +723,14 @@ class KVCacheConfigurator: enable_linear_replayssm=self.server_args.enable_linear_replayssm, linear_replayssm_cache_len=self.server_args.linear_replayssm_cache_len, mamba_envelope_layout=self.server_args.enable_page_major_kv_layout, + # ReplaySSM spec-verify is GDN-only: activate the pool machinery + # (rings + cursors + the intermediate_ssm gate) only for GDN-hybrid + # models, so any other mamba-ish model (Mamba2/Nemotron, lightning, + # ...) run with the flag set stays byte-identical to flag-off. + enable_gdn_replayssm_spec=( + self.server_args.enable_gdn_replayssm_spec + and self.hybrid_gdn_config is not None + ), ) return req_to_token_pool diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 65faa9dad..5c951f811 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -329,9 +329,20 @@ class MambaPool: # replayssm_d: [num_layers, num_slots, HV, L, V] # replayssm_k: [num_layers, num_slots, H, L, K] # replayssm_g: [num_layers, num_slots, HV, L] (fp32) + # replayssm_rawv: [num_layers, num_slots, HV, L, V] (conv/activation dtype) + # replayssm_rawk: [num_layers, num_slots, H, L, K] (conv/activation dtype) + # replayssm_beta: [num_layers, num_slots, HV, L] (fp32) + # The raw rings + beta exist only under --enable-gdn-replayssm-spec: the + # closed-loop exact fold sequentially replays them through the recurrent + # update at flush -- bit-identical to the recurrent baseline -- instead + # of folding the chunked `d` records open-loop (which accumulates error + # across flushes). See fla/gdn_replayssm_spec_decode.py. replayssm_d: Optional[torch.Tensor] = None replayssm_k: Optional[torch.Tensor] = None replayssm_g: Optional[torch.Tensor] = None + replayssm_rawv: Optional[torch.Tensor] = None + replayssm_rawk: Optional[torch.Tensor] = None + replayssm_beta: Optional[torch.Tensor] = None def at_layer_idx(self, layer: int): kwargs = {} @@ -357,7 +368,10 @@ class MambaPool: @dataclass(frozen=True, kw_only=True) class SpeculativeState(State): - intermediate_ssm: torch.Tensor + # None under --enable-gdn-replayssm-spec: the spec ring owns rollback + # (verify writes ring records, commit moves cursors), so the per-draft + # full-state snapshots are never produced or consumed. + intermediate_ssm: Optional[torch.Tensor] intermediate_conv_window: List[torch.Tensor] def _allocate_deduplicated_conv_window( @@ -414,6 +428,7 @@ class MambaPool: enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, envelope_layout: bool = False, + enable_gdn_replayssm_spec: bool = False, ): conv_state_shape = cache_params.shape.conv temporal_state_shape = cache_params.shape.temporal @@ -429,6 +444,13 @@ class MambaPool: self.debug_memory_pool = envs.SGLANG_DEBUG_MEMORY_POOL.get() self.enable_linear_replayssm = enable_linear_replayssm self.linear_replayssm_cache_len = linear_replayssm_cache_len + # ReplaySSM spec-verify (Part B of #28511) REUSES the linear_replayssm ring + # (replayssm_d/k/g + write_pos) and ADDS two per-slot cursors + # (replayssm_cache_base + replayssm_is_flush). Enabling the spec-verify path + # therefore implies the ring, so the d/k/g + write_pos allocation gates on + # `_replayssm_on` (either flag). GDN-only is enforced upstream + below. + self.enable_gdn_replayssm_spec = enable_gdn_replayssm_spec + _replayssm_on = enable_linear_replayssm or enable_gdn_replayssm_spec # for disagg with nvlink self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( @@ -504,22 +526,34 @@ class MambaPool: # GDN ReplaySSM ring buffers (slice 1a). Allocated only when the # flag is on; otherwise left as None so the legacy State is - # byte-identical. temporal_state_shape == (HV, V, K). + # byte-identical. temporal_state_shape == (HV, V, K). Either the decode + # ring (--enable-linear-replayssm) or the spec-verify ring + # (--enable-gdn-replayssm-spec) shares this allocation. replayssm_d = replayssm_k = replayssm_g = None - if enable_linear_replayssm: + replayssm_rawv = replayssm_rawk = replayssm_beta = None + if _replayssm_on: hv, v_dim, k_dim = temporal_state_shape h_k = getattr(cache_params.shape, "num_k_heads_per_tp", hv) L = linear_replayssm_cache_len num_slots = size + 1 - # Ring records live in the SSM dtype (bf16/fp32) except g (fp32). + # Ring dtype. DECODE ring (--enable-linear-replayssm): records + # follow the SSM dtype -- its flush folds `d` directly into the + # state. SPEC-verify ring (--enable-gdn-replayssm-spec): d/k feed + # ONLY the one-shot output reconstruction (the closed-loop exact + # fold replays the raw rings for state instead), so their + # quantization noise stays below the bf16 output cast; keep them + # in the conv/activation dtype instead of the (fp32-enforced) + # SSM dtype to halve the ring traffic. g stays fp32 everywhere + # (exact-fold input). The two flags are mutually exclusive. + ring_dtype = conv_dtype if enable_gdn_replayssm_spec else ssm_dtype replayssm_d = torch.zeros( size=(num_mamba_layers, num_slots, hv, L, v_dim), - dtype=ssm_dtype, + dtype=ring_dtype, device=device, ) replayssm_k = torch.zeros( size=(num_mamba_layers, num_slots, h_k, L, k_dim), - dtype=ssm_dtype, + dtype=ring_dtype, device=device, ) # The log-decay gate ring (fp32): per-head SCALAR for the GDN @@ -535,6 +569,42 @@ class MambaPool: dtype=torch.float32, device=device, ) + # Closed-loop exact-fold rings (spec-verify only). Raw v / raw + # pre-norm k live in the conv (activation) dtype -- they are born + # there, so storage round-trips losslessly -- beta in fp32. The + # flush replays these through the recurrent update sequentially + # (bit-identical to the recurrent baseline) instead of folding + # the chunked `d` records open-loop. + if enable_gdn_replayssm_spec: + # Backstop for the spec-verify ring invariants; this pool + # is sized with the final adaptive-aware draft maximum. + if L & (L - 1) != 0: + raise ValueError( + f"spec-verify ring length must be a power of two, got {L}" + ) + if ( + speculative_num_draft_tokens is not None + and L < 2 * speculative_num_draft_tokens + ): + raise ValueError( + f"spec-verify ring too small: {L} < " + f"2 * {speculative_num_draft_tokens} (early-flush margin)" + ) + replayssm_rawv = torch.zeros( + size=(num_mamba_layers, num_slots, hv, L, v_dim), + dtype=conv_dtype, + device=device, + ) + replayssm_rawk = torch.zeros( + size=(num_mamba_layers, num_slots, h_k, L, k_dim), + dtype=conv_dtype, + device=device, + ) + replayssm_beta = torch.zeros( + size=(num_mamba_layers, num_slots, hv, L), + dtype=torch.float32, + device=device, + ) if speculative_num_draft_tokens is not None: if _is_npu: @@ -546,18 +616,30 @@ class MambaPool: ) # Cache intermediate SSM states per draft token during target verify # Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V] - intermediate_ssm_state_cache = torch.zeros( - size=( - num_mamba_layers, - spec_state_size + 1, - speculative_num_draft_tokens, - temporal_state_shape[0], - temporal_state_shape[1], - temporal_state_shape[2], - ), - dtype=ssm_dtype, - device="cuda", - ) + # + # ReplaySSM spec-verify owns rollback via the ring + cursors (the + # verify kernel never writes per-draft snapshots; the commit never + # reads them), so this buffer -- the dominant spec scratch, ~46x + # the conv state -- is dead weight there and is skipped. The conv + # intermediate windows below STAY (conv rollback consumes them). + # The recurrent-verify fallback cannot be reached under the flag + # (GDN + linear chain + triton enforced in server_args; the + # backend asserts loudly if it ever is). + if enable_gdn_replayssm_spec: + intermediate_ssm_state_cache = None + else: + intermediate_ssm_state_cache = torch.zeros( + size=( + num_mamba_layers, + spec_state_size + 1, + speculative_num_draft_tokens, + temporal_state_shape[0], + temporal_state_shape[1], + temporal_state_shape[2], + ), + dtype=ssm_dtype, + device="cuda", + ) # Cache intermediate conv windows (last K-1 inputs) per draft token # during target verify. # @@ -625,13 +707,21 @@ class MambaPool: replayssm_d=replayssm_d, replayssm_k=replayssm_k, replayssm_g=replayssm_g, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_beta=replayssm_beta, + ) + intermediate_ssm_gb = ( + get_tensor_size_bytes(intermediate_ssm_state_cache) / GB + if intermediate_ssm_state_cache is not None + else 0.0 ) logger.info( f"Mamba Cache is allocated. " f"max_mamba_cache_size: {size}, " f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, " f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB " - f"intermediate_ssm_state_cache size: {get_tensor_size_bytes(intermediate_ssm_state_cache) / GB:.2f}GB " + f"intermediate_ssm_state_cache size: {intermediate_ssm_gb:.2f}GB " # Report the deduplicated PHYSICAL conv-window buffers (the view # over-reports its logical, un-deduplicated size). f"intermediate_conv_window_cache size: {get_tensor_size_bytes(self._intermediate_conv_window_phys) / GB:.2f}GB " @@ -643,6 +733,9 @@ class MambaPool: replayssm_d=replayssm_d, replayssm_k=replayssm_k, replayssm_g=replayssm_g, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_beta=replayssm_beta, ) logger.info( f"Mamba Cache is allocated. " @@ -650,26 +743,48 @@ class MambaPool: f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, " f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB " ) - if enable_linear_replayssm: + if _replayssm_on: logger.info( f"GDN ReplaySSM ring buffers allocated (L=" f"{linear_replayssm_cache_len}): " f"d={get_tensor_size_bytes(replayssm_d) / GB:.3f}GB, " f"k={get_tensor_size_bytes(replayssm_k) / GB:.3f}GB, " f"g={get_tensor_size_bytes(replayssm_g) / GB:.3f}GB " + + ( + f"rawv={get_tensor_size_bytes(replayssm_rawv) / GB:.3f}GB, " + f"rawk={get_tensor_size_bytes(replayssm_rawk) / GB:.3f}GB, " + f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB " + if enable_gdn_replayssm_spec + else "" + ) ) # Gate granularity of the linear-attn layers (drives the kernel's # IS_KDA path + the g_cache layout). Read by the backend metadata to # decide the per-K (KDA) vs scalar (GDN) flush/advance handling. - self.replayssm_is_kda = bool( - enable_linear_replayssm and cache_params.is_kda - ) + self.replayssm_is_kda = bool(_replayssm_on and cache_params.is_kda) # Persistent per-slot decode-position cursor for ReplaySSM. Shared # across all linear-attn layers; advanced once per decode forward by - # the backend metadata build. Index 0..size; reset on slot (re)alloc. + # the backend metadata build (decode ring) or once per verify step by + # the worker (spec-verify ring). Index 0..size; reset on slot (re)alloc. self.replayssm_write_pos = ( torch.zeros((size + 1,), dtype=torch.int32, device=device) - if enable_linear_replayssm + if _replayssm_on + else None + ) + # ReplaySSM spec-verify (Part B of #28511) extra per-slot cursors. The + # circular ring's rolling origin (cache_base) + the per-slot flush flag + # (is_flush). Block-keyed (indexed by the physical mamba slot), shared by + # all GDN layers of one verify step; advanced by commit_gdn_replayssm_spec. + # Only allocated for the spec-verify ring (the decode ring does not use + # a circular buffer); None otherwise. + self.replayssm_cache_base = ( + torch.zeros((size + 1,), dtype=torch.int32, device=device) + if enable_gdn_replayssm_spec + else None + ) + self.replayssm_is_flush = ( + torch.zeros((size + 1,), dtype=torch.int8, device=device) + if enable_gdn_replayssm_spec else None ) mem_usage_bytes = self.mamba_cache.mem_usage_bytes() @@ -804,6 +919,12 @@ class MambaPool: ] if self.replayssm_write_pos is not None: self.replayssm_write_pos[dst_indices] = 0 + # ReplaySSM spec-verify ring: a copied checkpoint has no pending ring + # entries, so its rolling origin + flush flag reset alongside write_pos. + if self.replayssm_cache_base is not None: + self.replayssm_cache_base[dst_indices] = 0 + if self.replayssm_is_flush is not None: + self.replayssm_is_flush[dst_indices] = 0 def get_cpu_copy(self, indices): current_platform.synchronize() @@ -814,17 +935,46 @@ class MambaPool: temporal_cpu = self.mamba_cache.temporal[:, indices].to( "cpu", non_blocking=True ) + # ReplaySSM spec-verify ring: round-trip the per-slot cursors with the + # checkpoint so a restored slot reconstructs exactly. Only the spec ring + # adds the 3rd tuple element; every other config keeps the legacy 2-tuple + # so those paths stay byte-identical. + if self.replayssm_cache_base is not None: + cursors_cpu = ( + self.replayssm_write_pos[indices].to("cpu", non_blocking=True), + self.replayssm_cache_base[indices].to("cpu", non_blocking=True), + self.replayssm_is_flush[indices].to("cpu", non_blocking=True), + ) + current_platform.synchronize() + return conv_cpu, temporal_cpu, cursors_cpu current_platform.synchronize() return conv_cpu, temporal_cpu def load_cpu_copy(self, mamba_cache_cpu, indices): - conv_cpu, temporal_cpu = mamba_cache_cpu + # Accept both the legacy 2-tuple (conv, temporal) and the 3-tuple that also + # carries the ReplaySSM spec-verify cursors. + if len(mamba_cache_cpu) == 3: + conv_cpu, temporal_cpu, cursors_cpu = mamba_cache_cpu + else: + conv_cpu, temporal_cpu = mamba_cache_cpu + cursors_cpu = None current_platform.synchronize() for i, conv in enumerate(self.mamba_cache.conv): conv[:, indices] = conv_cpu[i].to(conv.device, non_blocking=True) self.mamba_cache.temporal[:, indices] = temporal_cpu.to( self.mamba_cache.temporal.device, non_blocking=True ) + if cursors_cpu is not None and self.replayssm_cache_base is not None: + wp_cpu, cb_cpu, fl_cpu = cursors_cpu + self.replayssm_write_pos[indices] = wp_cpu.to( + self.replayssm_write_pos.device, non_blocking=True + ) + self.replayssm_cache_base[indices] = cb_cpu.to( + self.replayssm_cache_base.device, non_blocking=True + ) + self.replayssm_is_flush[indices] = fl_cpu.to( + self.replayssm_is_flush.device, non_blocking=True + ) current_platform.synchronize() def get_contiguous_buf_infos(self): @@ -841,7 +991,14 @@ class MambaPool: continue # Skip GDN ReplaySSM ring buffers: they are derived/transient decode # scratch, not part of the persistent transferable state. - if field in ("replayssm_d", "replayssm_k", "replayssm_g"): + if field in ( + "replayssm_d", + "replayssm_k", + "replayssm_g", + "replayssm_rawv", + "replayssm_rawk", + "replayssm_beta", + ): continue value = getattr(self.mamba_cache, field) if value is None: @@ -882,6 +1039,9 @@ class MambaPool: "replayssm_d", "replayssm_k", "replayssm_g", + "replayssm_rawv", + "replayssm_rawk", + "replayssm_beta", ): continue value = getattr(self.mamba_cache, field) @@ -964,6 +1124,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, mamba_envelope_layout: bool = False, + enable_gdn_replayssm_spec: bool = False, ): super().__init__( size=size, @@ -990,6 +1151,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm=enable_linear_replayssm, linear_replayssm_cache_len=linear_replayssm_cache_len, mamba_envelope_layout=mamba_envelope_layout, + enable_gdn_replayssm_spec=enable_gdn_replayssm_spec, ) def _init_mamba_pool( @@ -1005,6 +1167,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, mamba_envelope_layout: bool = False, + enable_gdn_replayssm_spec: bool = False, ): self.mamba_pool = self.mamba_pool_cls( size=mamba_size, @@ -1018,6 +1181,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm=enable_linear_replayssm, linear_replayssm_cache_len=linear_replayssm_cache_len, envelope_layout=mamba_envelope_layout, + enable_gdn_replayssm_spec=enable_gdn_replayssm_spec, ) self.mamba_allocator = MambaSlotAllocator( size=mamba_size, @@ -1116,6 +1280,12 @@ class HybridReqToTokenPool(ReqToTokenPool): # (the post-prefill state that prefill wrote into this slot). if self.mamba_pool.replayssm_write_pos is not None: self.mamba_pool.replayssm_write_pos[req.mamba_pool_idx] = 0 + # ReplaySSM spec-verify ring: an empty ring also resets the + # circular origin + flush flag so the first verify step on this + # freshly-prefilled slot reconstructs from the checkpoint alone. + if self.mamba_pool.replayssm_cache_base is not None: + self.mamba_pool.replayssm_cache_base[req.mamba_pool_idx] = 0 + self.mamba_pool.replayssm_is_flush[req.mamba_pool_idx] = 0 mamba_indices.append(req.mamba_pool_idx) if self.enable_mamba_extra_buffer: if req.mamba_ping_pong_track_buffer is None: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 150373a6b..2f93fcd4b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2164,6 +2164,16 @@ class ServerArgs: int, "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps.", ] = 16 + # ReplaySSM spec-verify (Part B of RFC #28511): GDN linear-chain target-verify + # via a per-slot circular (d, k, g) ring + periodic flush instead of per-draft + # full-state snapshots. GDN only; linear-chain (topk <= 1) only. Reuses the + # `linear_replayssm` ring (replayssm_d/k/g + write_pos) and adds two per-slot + # cursors (cache_base, is_flush); the ring length reuses + # `linear_replayssm_cache_len`. + enable_gdn_replayssm_spec: A[ + bool, + "Enable the ReplaySSM GDN spec-verify kernel (Part B of RFC #28511): a per-slot circular (d, k, g) ring + periodic flush replacing the recurrent verify's per-draft full-state snapshots. GDN only, linear-chain (--speculative-eagle-topk in {None, 1}) only. Reuses --linear-replayssm-cache-len for the ring length.", + ] = False # ------------------------------------------------------------------------- # Hierarchical cache @@ -2998,6 +3008,9 @@ class ServerArgs: handle_speculative_decoding(self) + # Needs the draft-token count derived just above. + self._validate_gdn_replayssm_spec_ring() + # Validate the CuteDSL A2A token budget now that num_tokens_per_req is final. self._validate_cutedsl_a2a_token_budget() @@ -5307,6 +5320,110 @@ class ServerArgs: f"{self.linear_replayssm_cache_len}." ) + # ReplaySSM spec-verify (Part B of #28511): GDN-only, linear-chain target + # verify. Reuses the `linear_replayssm` ring (replayssm_d/k/g + write_pos) + # plus two extra per-slot cursors (cache_base, is_flush) and the chunked + # (I+A)^-1 reconstruction verify kernel. The intra-window interaction uses a + # strictly-lower causal mask, so it is valid ONLY for a linear draft chain + # (speculative_eagle_topk in {None, 1}, i.e. NEXTN / MTP); EAGLE tree verify + # (topk > 1) must fall back to the recurrent verify. GDN-only is enforced at + # runtime (KDA routes through kda_backend, which never enters this path; the + # pool gate also checks `not cache_params.is_kda`). The ring length reuses + # --linear-replayssm-cache-len (no separate flag). + if self.enable_gdn_replayssm_spec: + if self.speculative_eagle_topk not in (None, 1): + raise ValueError( + "--enable-gdn-replayssm-spec requires a linear draft chain " + "(--speculative-eagle-topk in {None, 1}); the chunked verify " + "kernel uses a strictly-lower causal mask and is invalid for " + "EAGLE tree verify. Got " + f"--speculative-eagle-topk={self.speculative_eagle_topk!r}." + ) + if decode != "triton": + raise ValueError( + "--enable-gdn-replayssm-spec requires the Triton linear-attn " + "decode backend, got " + f"--linear-attn-decode-backend={decode!r}." + ) + if self.enable_mamba_extra_buffer(): + # The spec-verify path does not yet implement the device-side + # force-flush needed to keep `temporal` consistent with the ring at + # radix mamba-track boundaries, so it is incompatible with + # extra_buffer (radix prefix caching). + raise ValueError( + "--enable-gdn-replayssm-spec is not yet compatible with mamba " + "extra_buffer (radix prefix caching); use --disable-radix-cache " + "or --mamba-radix-cache-strategy no_buffer." + ) + if self.disaggregation_mode != "null": + raise ValueError( + "--enable-gdn-replayssm-spec is not supported under PD " + "disaggregation yet (follow-up). Got " + f"--disaggregation-mode={self.disaggregation_mode!r}." + ) + if self.linear_replayssm_cache_len < 1: + raise ValueError( + "--linear-replayssm-cache-len must be >= 1, got " + f"{self.linear_replayssm_cache_len}." + ) + if self.enable_linear_replayssm: + raise ValueError( + "--enable-gdn-replayssm-spec and --enable-linear-replayssm are " + "mutually exclusive: they share the ring storage but drive it " + "with incompatible cursor protocols (per-decode-forward vs " + "per-verify-commit advance)." + ) + ring_len = self.linear_replayssm_cache_len + if ring_len & (ring_len - 1) != 0: + raise ValueError( + "--linear-replayssm-cache-len must be a power of two for the " + f"circular spec-verify ring, got {ring_len}." + ) + # ring_len >= 2 * max drafts is checked in + # _validate_gdn_replayssm_spec_ring() (draft tokens not derived yet). + # Closed-loop exact fold: the flush replays raw ring inputs through + # the recurrent update into the checkpoint, bit-identical to the + # recurrent baseline -- which keeps its state in fp32. A 16-bit + # checkpoint would re-quantize the exactly-folded state every flush + # and become the dominant residual error source, so require fp32. + if self.mamba_ssm_dtype is None: + logger.info( + "--enable-gdn-replayssm-spec: setting --mamba-ssm-dtype " + "float32 (the closed-loop exact fold requires the fp32 SSM " + "checkpoint for recurrent-parity)." + ) + self.mamba_ssm_dtype = "float32" + elif self.mamba_ssm_dtype != "float32": + raise ValueError( + "--enable-gdn-replayssm-spec requires --mamba-ssm-dtype " + f"float32, got {self.mamba_ssm_dtype!r}. The closed-loop " + "exact fold keeps the committed state bit-identical to the " + "recurrent baseline, which is only meaningful against the " + "fp32 checkpoint; a 16-bit checkpoint would re-quantize it " + "every flush." + ) + + def _validate_gdn_replayssm_spec_ring(self): + """Enforce ring_len >= 2 * max draft tokens for the spec-verify ring. + + Early-flush margin: write_pos + spec_len <= ring_len must hold on every + verify step (see _advance_gdn_spec_cursors_kernel). Runs after + handle_speculative_decoding() so the (adaptive-aware) max is final; + MambaPool re-checks at ring allocation as a backstop. + """ + if not self.enable_gdn_replayssm_spec: + return + max_drafts = self.max_speculative_num_draft_tokens + if max_drafts is None: + return + ring_len = self.linear_replayssm_cache_len + if ring_len < 2 * max_drafts: + raise ValueError( + "--linear-replayssm-cache-len must be >= 2 * the maximum " + "speculative draft-token count for the spec-verify ring " + f"(early-flush margin), got {ring_len} < {2 * max_drafts}." + ) + def _handle_legacy_cp_arguments(self): legacy_mode_to_strategy = { "in-seq-split": "zigzag", diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 235952647..39c7ef31b 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -696,6 +696,73 @@ def commit_mamba_states_after_verify( model_runner = target_worker.model_runner if mambaish_config(model_runner.model_config) is None: return + + # ReplaySSM spec-verify path (Part B of #28511): the accepted drafts already + # live in the per-slot circular ring (written during verify). Instead of + # scattering an intermediate full SSM state into `temporal`, advance the + # block-keyed cursors by the accepted count (the ring owns the SSM state; the + # verify/flush kernel folds it into `temporal` periodically). The CONV state + # still needs its usual accept-rollback, so we keep the conv-window scatter and + # skip only the SSM scatter. GDN-only + linear-chain (topk<=1) -- the runtime + # ring is allocated only then; KDA never allocates the cursors. + req_pool = model_runner.req_to_token_pool + mamba_pool = getattr(req_pool, "mamba_pool", None) + if ( + mamba_pool is not None + and getattr(mamba_pool, "replayssm_cache_base", None) is not None + and not getattr(mamba_pool, "replayssm_is_kda", False) + ): + if batch.forward_mode.is_idle() or accept_index.numel() == 0: + return + from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_decode import ( + commit_gdn_replayssm_spec, + ) + from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( + fused_conv_window_scatter_with_mask, + ) + + spec_state = req_pool.get_speculative_mamba2_params_all_layers() + bs = accept_lens.shape[0] + state_batch_indices = req_pool.get_mamba_indices(batch.req_pool_indices) + # Advance the per-slot circular cursors by the accepted count (incl. the + # bonus token). max_cache_len = ring length L = replayssm_d.shape[-2]. + commit_gdn_replayssm_spec( + write_pos=mamba_pool.replayssm_write_pos, + cache_base=mamba_pool.replayssm_cache_base, + is_flush=mamba_pool.replayssm_is_flush, + num_accepted=accept_lens, # [bs], includes the bonus token + state_batch_indices=state_batch_indices, + max_cache_len=spec_state.replayssm_d.shape[-2], + max_spec_len=draft_token_num, + null_block_id=-1, # SGLang: valid slots >= 0, padding == -1 + ) + # Roll back / commit the conv state to the last accepted draft step + # (same logic as the recurrent commit, but conv-only). + accept_indices_offset = torch.arange( + 0, + bs * draft_token_num, + step=draft_token_num, + dtype=accept_lens.dtype, + device=accept_lens.device, + ) + req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device) + last_correct_step_indices = ( + accept_index[req_idx, (accept_lens - 1).to(torch.int64)] + - accept_indices_offset + ) + fused_conv_window_scatter_with_mask( + spec_state.conv[0], + spec_state.intermediate_conv_window[0], + state_batch_indices, + last_correct_step_indices, + ) + # NOTE: radix mamba prefix-caching (mamba_track / extra_buffer) would need + # a device-side force-flush so `temporal` reflects the ring before a + # snapshot; not wired for Part B (server_args forbids extra_buffer with + # --enable-gdn-replayssm-spec), so the per-track scatters are intentionally + # skipped here. + return + attn_backend = model_runner.attn_backend bs = accept_lens.shape[0]