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 index aa1e53122..d484b396b 100644 --- a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py +++ b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py @@ -16,30 +16,13 @@ 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 (fp32 SSM checkpoint is the server_args default; a 16-bit - checkpoint is allowed with a warning, unvalidated for GDN). +Compact materialization keeps the update vector and normalized key as high/low +activation-dtype parts. The commit kernel reconstructs both parts before the +update reaches the recurrent checkpoint, while the scalar decay stays fp32. +For output reconstruction, a fp32 checkpoint keeps the pre-existing fp32/TF32 +dot operands; a 16-bit checkpoint uses the activation tensor-core path plus the +normalized-key residual terms. ``DOT_PRECISION`` remains selectable for +ablations. Differences from the vLLM reference: * SGLang passes **split** ``q`` / ``k`` / ``v`` tensors (already split + post @@ -71,18 +54,19 @@ def gdn_replayssm_spec_circular_kernel( 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) + h0, # [num_slots, HV, V, K] checkpoint (read-only in verify) 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) + rawv_cache, # exact raw v or compact D low part + rawk_cache, # exact raw k or normalized-K low part + beta_cache, # fp32 beta for exact 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 + replay_indices, # [B] int active request slot per request + write_pos, # [num_replay_slots] int32 + cache_base, # [num_replay_slots] int32 + is_flush_flags, # [num_replay_slots] int8 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) @@ -99,6 +83,7 @@ def gdn_replayssm_spec_circular_kernel( stride_beta_slot: tl.constexpr, stride_qsl: tl.constexpr, stride_indices: tl.constexpr, + stride_replay_indices: tl.constexpr, H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, @@ -112,6 +97,10 @@ def gdn_replayssm_spec_circular_kernel( MAX_CACHE_LEN: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + STATE_FP32: tl.constexpr, + STORE_EXACT_RAW: tl.constexpr, + STORE_D_RESIDUAL: tl.constexpr, + STORE_K_RESIDUAL: tl.constexpr, IS_FLUSH: tl.constexpr, NULL_BLOCK_ID: tl.constexpr, DOT_PRECISION: tl.constexpr, @@ -139,6 +128,7 @@ def gdn_replayssm_spec_circular_kernel( o_s_safe = tl.where(mask_s, o_s, 0) state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + replay_idx = tl.load(replay_indices + i_n * stride_replay_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, :] @@ -146,7 +136,7 @@ def gdn_replayssm_spec_circular_kernel( if IS_FLUSH: if state_idx <= NULL_BLOCK_ID: return - b_is_flush = tl.load(is_flush_flags + state_idx) != 0 + b_is_flush = tl.load(is_flush_flags + replay_idx) != 0 if not b_is_flush: return else: @@ -158,12 +148,12 @@ def gdn_replayssm_spec_circular_kernel( mask=full_mask, ) return - b_is_flush = tl.load(is_flush_flags + state_idx) != 0 + b_is_flush = tl.load(is_flush_flags + replay_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) + b_write_pos = tl.load(write_pos + replay_idx).to(tl.int64) + b_cache_base = tl.load(cache_base + replay_idx).to(tl.int32) out_mask = mask_s[:, None] & mask_v[None, :] @@ -200,7 +190,7 @@ def gdn_replayssm_spec_circular_kernel( # 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 + p_g_main = g_cache + replay_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) @@ -208,14 +198,20 @@ def gdn_replayssm_spec_circular_kernel( b_total_decay = tl.exp(b_g_total) p_d_main = d_cache + ( - state_idx * stride_d_slot + replay_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) + b_d_scaled_fp32 = b_d_all * b_replay_decay[None, :] + if STATE_FP32: + # Preserve the pre-existing fp32-checkpoint path: these dots use + # fp32 operands with DOT_PRECISION (TF32 by default). + b_d_scaled = b_d_scaled_fp32 + else: + b_d_scaled = b_d_scaled_fp32.to(d_cache.dtype.element_ty) if USE_QK_L2NORM_IN_KERNEL: qnorm_acc = tl.zeros([BS], dtype=tl.float32) @@ -277,12 +273,20 @@ def gdn_replayssm_spec_circular_kernel( 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) + q_precise = q_tile * (q_rnorm * scale)[:, None] + k_precise = k_tile * k_rnorm[:, None] + k_store_tile = k_precise.to(k_cache.dtype.element_ty) + k_residual_tile = (k_precise - k_store_tile.to(tl.float32)).to( + k.dtype.element_ty + ) + if STATE_FP32: + # Do not lower the fp32-checkpoint output path to activation dtype. + q_tile = q_precise + k_tile = k_precise + else: + q_tile = q_precise.to(q.dtype.element_ty) + k_tile = k_store_tile p_h0 = ( h0 @@ -291,31 +295,49 @@ def gdn_replayssm_spec_circular_kernel( + 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 - ) + sc_tile = tl.load(p_h0, mask=mask_v[:, None] & mask_kt[None, :], other=0.0) 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) + if STORE_K_RESIDUAL and not STATE_FP32: + k_residual_T = tl.trans(k_residual_tile) + kk_mat += tl.dot(k_tile, k_residual_T, input_precision=DOT_PRECISION) + kk_mat += tl.dot(k_residual_tile, kT, input_precision=DOT_PRECISION) + kk_mat += tl.dot( + k_residual_tile, k_residual_T, input_precision=DOT_PRECISION + ) + kq_mat += tl.dot(k_residual_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 STATE_FP32: + sc_tile = sc_tile.to(tl.float32) + hw_q += tl.dot(sc_tile, qT, input_precision=DOT_PRECISION) + hw_k += tl.dot(sc_tile, kT, input_precision=DOT_PRECISION) + else: + sc_tile = sc_tile.to(k_tile.dtype) + hw_q += tl.dot(sc_tile, qT) + hw_k += tl.dot(sc_tile, kT) + if STORE_K_RESIDUAL: + hw_k += tl.dot(sc_tile, k_residual_T) if not IS_FLUSH: # cached-key history load -> phys_c (output reconstruction only) p_k = k_cache + ( - state_idx * stride_k_slot + replay_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) + ) + if STATE_FP32: + khist_tile = khist_tile.to(tl.float32) + else: + khist_tile = khist_tile.to(k_cache.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) @@ -325,27 +347,29 @@ def gdn_replayssm_spec_circular_kernel( & 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, - ) + if STORE_EXACT_RAW or STORE_K_RESIDUAL: + p_cur_rawk = rawk_cache + ( + replay_idx * stride_rawk_slot + + (i_h * MAX_CACHE_LEN + phys_spec[:, None]) * K + + o_kt[None, :] + ) + if STORE_EXACT_RAW: + tl.store( + p_cur_rawk, + k_raw_tile.to(p_cur_rawk.dtype.element_ty), + mask=spec_kt_mask, + ) + else: + tl.store(p_cur_rawk, k_residual_tile, mask=spec_kt_mask) # normalized spec key store -> phys_spec (circular) p_cur_k = k_cache + ( - state_idx * stride_k_slot + replay_idx * stride_k_slot + (i_h * MAX_CACHE_LEN + phys_spec[:, None]) * K + o_kt[None, :] ) tl.store( p_cur_k, - k_tile, + k_store_tile, mask=spec_kt_mask, ) @@ -407,33 +431,47 @@ def gdn_replayssm_spec_circular_kernel( 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 + replay_idx * stride_d_slot + (i_hv * MAX_CACHE_LEN + phys_spec[None, :]) * V + o_v[:, None] ) + d_hi = D_spec.to(p_cur_d.dtype.element_ty) tl.store( p_cur_d, - D_spec.to(p_cur_d.dtype.element_ty), + d_hi, 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 + if STORE_EXACT_RAW or STORE_D_RESIDUAL: + p_cur_v = rawv_cache + ( + replay_idx * stride_rawv_slot + + (i_hv * MAX_CACHE_LEN + phys_spec[None, :]) * V + + o_v[:, None] ) - tl.store(p_cur_beta, beta_s, mask=spec_store_mask) + if STORE_EXACT_RAW: + tl.store( + p_cur_v, + v_tile.to(p_cur_v.dtype.element_ty), + mask=mask_v[:, None] & spec_store_mask[None, :], + ) + else: + tl.store( + p_cur_v, + (D_spec - d_hi.to(tl.float32)).to(p_cur_v.dtype.element_ty), + mask=mask_v[:, None] & spec_store_mask[None, :], + ) + if i_v == 0: + p_cur_g = ( + g_cache + replay_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys_spec + ) + tl.store(p_cur_g, g_s, mask=spec_store_mask) + if STORE_EXACT_RAW: + p_cur_beta = ( + beta_cache + + replay_idx * stride_beta_slot + + i_hv * MAX_CACHE_LEN + + phys_spec + ) + tl.store(p_cur_beta, beta_s, mask=spec_store_mask) @triton.jit @@ -444,15 +482,17 @@ def gdn_replayssm_exact_fold_kernel( 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 + replay_indices, # [B] int active request slot per request + write_pos, + cache_base, + is_flush_flags, 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, + stride_replay_indices: tl.constexpr, H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, @@ -486,14 +526,15 @@ def gdn_replayssm_exact_fold_kernel( i_h = i_hv // (HV // H) state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + replay_idx = tl.load(replay_indices + i_n * stride_replay_indices).to(tl.int64) if state_idx <= NULL_BLOCK_ID: return - if tl.load(is_flush_flags + state_idx) == 0: + if tl.load(is_flush_flags + replay_idx) == 0: return - b_write_pos = tl.load(write_pos + state_idx).to(tl.int32) + b_write_pos = tl.load(write_pos + replay_idx).to(tl.int32) if b_write_pos <= 0: return - b_cache_base = tl.load(cache_base + state_idx).to(tl.int32) + b_cache_base = tl.load(cache_base + replay_idx).to(tl.int32) o_k = tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) @@ -514,6 +555,134 @@ def gdn_replayssm_exact_fold_kernel( 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 + + replay_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 + + replay_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 + replay_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys + ).to(tl.float32) + b_beta = tl.load( + beta_cache + replay_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 gdn_replayssm_circular_commit_kernel( + h0, # [num_layers, num_slots, HV, V, K] + rawv_cache, # [num_layers, num_slots, HV, L, V] + rawk_cache, # [num_layers, num_slots, H, L, K] + g_cache, # [num_layers, num_slots, HV, L] + beta_cache, # [num_layers, num_slots, HV, L] + ssm_state_indices, # [B] + write_pos, # [num_slots], already advanced by the accepted length + cache_base, # [num_slots] + is_flush_flags, # [num_slots] + accept_lens, # [B] + mamba_track_indices, # [B] + mamba_steps_to_track, # [B] + 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_state_layer: tl.constexpr, + stride_rawv_layer: tl.constexpr, + stride_rawk_layer: tl.constexpr, + stride_g_layer: tl.constexpr, + stride_beta_layer: tl.constexpr, + stride_indices: tl.constexpr, + stride_accept: tl.constexpr, + stride_track: tl.constexpr, + stride_steps: 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, + HAS_TRACK: tl.constexpr, +): + """Materialize circular history only at capacity or radix-track boundaries.""" + i_v = tl.program_id(0) + i_n = tl.program_id(1) + i_hvl = tl.program_id(2) + i_layer = (i_hvl // HV).to(tl.int64) + i_hv = i_hvl % HV + 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 + n_history = tl.load(write_pos + state_idx).to(tl.int32) + fold_active = tl.load(is_flush_flags + state_idx) != 0 + + if HAS_TRACK: + track_idx = tl.load(mamba_track_indices + i_n * stride_track).to(tl.int64) + track_step = tl.load(mamba_steps_to_track + i_n * stride_steps).to(tl.int32) + has_track = (track_idx > NULL_BLOCK_ID) & (track_step >= 0) + n_accepted = tl.load(accept_lens + i_n * stride_accept).to(tl.int32) + previous_history = n_history - n_accepted + track_pos = previous_history + track_step + 1 + else: + track_idx = NULL_BLOCK_ID + has_track = False + track_pos = 0 + + if (not fold_active) and (not has_track): + return + replay_len = tl.maximum(tl.where(fold_active, n_history, 0), track_pos) + if replay_len <= 0: + return + + h0 = h0 + i_layer * stride_state_layer + rawv_cache = rawv_cache + i_layer * stride_rawv_layer + rawk_cache = rawk_cache + i_layer * stride_rawk_layer + g_cache = g_cache + i_layer * stride_g_layer + beta_cache = beta_cache + i_layer * stride_beta_layer + 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, :] + 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.0).to(tl.float32) + + for t in range(0, replay_len): + phys = ((base + t) & (MAX_CACHE_LEN - 1)).to(tl.int64) b_k = tl.load( rawk_cache + state_idx * stride_rawk_slot @@ -536,16 +705,242 @@ def gdn_replayssm_exact_fold_kernel( 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_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) + if HAS_TRACK: + if has_track and (t + 1 == track_pos): + tl.store( + h0 + + track_idx * stride_state_slot + + i_hv * V * K + + o_v[None, :] * K + + o_k[:, None], + b_h.to(h0.dtype.element_ty), + mask=mask_h, + ) + + if fold_active: + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +@triton.jit +def gdn_replayssm_compact_commit_kernel( + h0, + d_cache, + k_cache, + g_cache, + d_residual_cache, + k_residual_cache, + ssm_state_indices, + replay_indices, + write_pos, + cache_base, + is_flush_flags, + accept_lens, + mamba_track_indices, + mamba_steps_to_track, + stride_state_slot: tl.constexpr, + stride_d_slot: tl.constexpr, + stride_k_slot: tl.constexpr, + stride_g_slot: tl.constexpr, + stride_d_residual_slot: tl.constexpr, + stride_k_residual_slot: tl.constexpr, + stride_state_layer: tl.constexpr, + stride_d_layer: tl.constexpr, + stride_k_layer: tl.constexpr, + stride_g_layer: tl.constexpr, + stride_d_residual_layer: tl.constexpr, + stride_k_residual_layer: tl.constexpr, + stride_indices: tl.constexpr, + stride_replay_indices: tl.constexpr, + stride_accept: tl.constexpr, + stride_track: tl.constexpr, + stride_steps: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, + HAS_TRACK: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, +): + """Materialize the checkpoint directly from compact ReplaySSM D/K/G.""" + i_v = tl.program_id(0) + i_n = tl.program_id(1) + i_hvl = tl.program_id(2) + i_layer = (i_hvl // HV).to(tl.int64) + i_hv = i_hvl % HV + i_h = i_hv // (HV // H) + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + replay_idx = tl.load(replay_indices + i_n * stride_replay_indices).to(tl.int64) + if state_idx <= NULL_BLOCK_ID: + return + n_history = tl.load(write_pos + replay_idx).to(tl.int32) + fold_active = tl.load(is_flush_flags + replay_idx) != 0 + if HAS_TRACK: + track_idx = tl.load(mamba_track_indices + i_n * stride_track).to(tl.int64) + track_step = tl.load(mamba_steps_to_track + i_n * stride_steps).to(tl.int32) + has_track = (track_idx > NULL_BLOCK_ID) & (track_step >= 0) + n_accepted = tl.load(accept_lens + i_n * stride_accept).to(tl.int32) + track_len = n_history - n_accepted + track_step + 1 + else: + track_idx = NULL_BLOCK_ID + has_track = False + track_len = 0 + if (not fold_active) and (not has_track): + return + + h0 += i_layer * stride_state_layer + d_cache += i_layer * stride_d_layer + k_cache += i_layer * stride_k_layer + g_cache += i_layer * stride_g_layer + d_residual_cache += i_layer * stride_d_residual_layer + k_residual_cache += i_layer * stride_k_residual_layer + base = tl.load(cache_base + replay_idx).to(tl.int32) + o_t = tl.arange(0, MAX_CACHE_LEN) + phys = (base + o_t) & (MAX_CACHE_LEN - 1) + o_k = tl.arange(0, K) + o_v = i_v * BV + tl.arange(0, BV) + mask_v = o_v < V + + p_h0 = ( + h0 + + state_idx * stride_state_slot + + i_hv * V * K + + o_v[None, :] * K + + o_k[:, None] + ) + old_state = tl.load(p_h0, mask=mask_v[None, :], other=0.0).to(tl.float32) + keys = tl.load( + k_cache + + replay_idx * stride_k_slot + + (i_h * MAX_CACHE_LEN + phys[:, None]) * K + + o_k[None, :] + ) + if HAS_RESIDUAL: + key_residual = tl.load( + k_residual_cache + + replay_idx * stride_k_residual_slot + + (i_h * MAX_CACHE_LEN + phys[:, None]) * K + + o_k[None, :] + ) + updates = tl.load( + d_cache + + replay_idx * stride_d_slot + + (i_hv * MAX_CACHE_LEN + phys[:, None]) * V + + o_v[None, :], + mask=mask_v[None, :], + other=0.0, + ).to(tl.float32) + if HAS_RESIDUAL: + updates += tl.load( + d_residual_cache + + replay_idx * stride_d_residual_slot + + (i_hv * MAX_CACHE_LEN + phys[:, None]) * V + + o_v[None, :], + mask=mask_v[None, :], + other=0.0, + ).to(tl.float32) + gates = tl.load( + g_cache + replay_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys + ).to(tl.float32) + + if fold_active: + active_mask = o_t < n_history + active_g = tl.where(active_mask, gates, 0.0) + active_prefix = tl.cumsum(active_g, axis=0) + active_total = tl.sum(active_g, axis=0) + active_decay = tl.where(active_mask, tl.exp(active_total - active_prefix), 0.0) + active_updates = updates * active_decay[:, None] + active_hi = active_updates.to(k_cache.dtype.element_ty) + active_lo = (active_updates - active_hi.to(tl.float32)).to( + k_cache.dtype.element_ty + ) + active_delta = tl.dot(tl.trans(keys), active_hi) + active_delta += tl.dot(tl.trans(keys), active_lo) + if HAS_RESIDUAL: + active_delta += tl.dot(tl.trans(key_residual), active_hi) + active_delta += tl.dot(tl.trans(key_residual), active_lo) + tl.store( + p_h0, + (old_state * tl.exp(active_total) + active_delta).to(p_h0.dtype.element_ty), + mask=mask_v[None, :], + ) + + if HAS_TRACK: + if has_track: + track_mask = o_t < track_len + track_g = tl.where(track_mask, gates, 0.0) + track_prefix = tl.cumsum(track_g, axis=0) + track_total = tl.sum(track_g, axis=0) + track_decay = tl.where(track_mask, tl.exp(track_total - track_prefix), 0.0) + track_updates = updates * track_decay[:, None] + track_hi = track_updates.to(k_cache.dtype.element_ty) + track_lo = (track_updates - track_hi.to(tl.float32)).to( + k_cache.dtype.element_ty + ) + track_delta = tl.dot(tl.trans(keys), track_hi) + track_delta += tl.dot(tl.trans(keys), track_lo) + if HAS_RESIDUAL: + track_delta += tl.dot(tl.trans(key_residual), track_hi) + track_delta += tl.dot(tl.trans(key_residual), track_lo) + tl.store( + h0 + + track_idx * stride_state_slot + + i_hv * V * K + + o_v[None, :] * K + + o_k[:, None], + (old_state * tl.exp(track_total) + track_delta).to(h0.dtype.element_ty), + mask=mask_v[None, :], + ) + + +@triton.jit +def _finish_gdn_replayssm_circular_fold_kernel( + write_pos, + cache_base, + is_flush, + state_batch_indices, + replay_indices, + n_rows, + stride_indices: tl.constexpr, + stride_replay_indices: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + BLOCK: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + row_mask = offs < n_rows + state_idx = tl.load( + state_batch_indices + offs * stride_indices, + mask=row_mask, + other=NULL_BLOCK_ID, + ).to(tl.int64) + replay_idx = tl.load( + replay_indices + offs * stride_replay_indices, + mask=row_mask, + other=0, + ).to(tl.int64) + valid = row_mask & (state_idx > NULL_BLOCK_ID) + do_fold = tl.load(is_flush + replay_idx, mask=valid, other=0) != 0 + n_history = tl.load(write_pos + replay_idx, mask=valid, other=0).to(tl.int32) + base = tl.load(cache_base + replay_idx, mask=valid, other=0).to(tl.int32) + mask = valid & do_fold + tl.store( + cache_base + replay_idx, + (base + n_history) & (MAX_CACHE_LEN - 1), + mask=mask, + ) + tl.store(write_pos + replay_idx, 0, mask=mask) + tl.store(is_flush + replay_idx, 0, mask=mask) @triton.jit @@ -554,20 +949,23 @@ def _advance_gdn_spec_cursors_kernel( cache_base_ptr, is_flush_ptr, num_accepted_ptr, - state_batch_indices_ptr, + replay_indices_ptr, n_rows, - stride_sbi: tl.constexpr, + stride_replay_indices: tl.constexpr, stride_na: tl.constexpr, MAX_CACHE_LEN: tl.constexpr, MAX_SPEC_LEN: tl.constexpr, CACHE_BUF_LEN: tl.constexpr, + FOLD_EVERY_COMMIT: 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 + replay_indices_ptr + offs * stride_replay_indices, + mask=row_mask, + other=NULL_BLOCK_ID, ).to(tl.int64) valid = row_mask & (blk > NULL_BLOCK_ID) @@ -598,7 +996,10 @@ def _advance_gdn_spec_cursors_kernel( # 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) + if FOLD_EVERY_COMMIT: + next_is_flush = (new_wp > 0).to(tl.int8) + else: + 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) @@ -659,6 +1060,7 @@ def _launch_gdn_spec( beta_cache, query_start_loc, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -694,6 +1096,15 @@ def _launch_gdn_spec( 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)) + store_exact_raw = ( + rawv_cache is not None and rawk_cache is not None and beta_cache is not None + ) + store_residual = ( + rawv_cache is not None and rawk_cache is not None and beta_cache is None + ) + rawv_buf = rawv_cache if rawv_cache is not None else d_cache + rawk_buf = rawk_cache if rawk_cache is not None else k_cache + beta_buf = beta_cache if beta_cache is not None else g_cache grid = (triton.cdiv(V, BV), B, HV) gdn_replayssm_spec_circular_kernel[grid]( @@ -709,11 +1120,12 @@ def _launch_gdn_spec( d_cache, k_cache, g_cache, - rawv_cache, - rawk_cache, - beta_cache, + rawv_buf, + rawk_buf, + beta_buf, query_start_loc, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -728,11 +1140,12 @@ def _launch_gdn_spec( d_cache.stride(0), k_cache.stride(0), g_cache.stride(0), - rawv_cache.stride(0), - rawk_cache.stride(0), - beta_cache.stride(0), + rawv_buf.stride(0), + rawk_buf.stride(0), + beta_buf.stride(0), query_start_loc.stride(0), ssm_state_indices.stride(0), + replay_indices.stride(0), H=H, HV=HV, K=K, @@ -746,6 +1159,10 @@ def _launch_gdn_spec( MAX_CACHE_LEN=max_cache_len, SOFTPLUS_THRESHOLD=20.0, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + STATE_FP32=checkpoint_state.dtype == torch.float32, + STORE_EXACT_RAW=store_exact_raw, + STORE_D_RESIDUAL=store_residual, + STORE_K_RESIDUAL=store_residual, IS_FLUSH=is_flush_kernel, NULL_BLOCK_ID=null_block_id, DOT_PRECISION=dot_precision, @@ -762,6 +1179,7 @@ def _launch_gdn_exact_fold( beta_cache, query_start_loc, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -789,6 +1207,7 @@ def _launch_gdn_exact_fold( g_cache, beta_cache, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -798,6 +1217,7 @@ def _launch_gdn_exact_fold( g_cache.stride(0), beta_cache.stride(0), ssm_state_indices.stride(0), + replay_indices.stride(0), H=num_k_heads, HV=HV, K=K, @@ -824,15 +1244,16 @@ def gdn_replayssm_spec_decode( 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) + rawv_cache: torch.Tensor | None, + rawk_cache: torch.Tensor | None, + beta_cache: torch.Tensor | None, 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 + replay_indices: torch.Tensor, # [B] int active request slot per request + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, max_cache_len: int, max_spec_len: int, scale: float | None = None, @@ -860,8 +1281,8 @@ def gdn_replayssm_spec_decode( 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`. + The circular history and cursors are request-keyed by ``replay_indices``; + the checkpoint remains keyed by physical ``ssm_state_indices``. """ if scale is None: scale = checkpoint_state.shape[-1] ** -0.5 @@ -869,6 +1290,8 @@ def gdn_replayssm_spec_decode( is_flush = is_flush.to(torch.int8) if launch_mode in ("both", "flush"): + if rawv_cache is None or rawk_cache is None or beta_cache is None: + raise ValueError("raw replay buffers are required for exact-fold launch") # Must precede the flush-output launch: it reads the folded checkpoint. _launch_gdn_exact_fold( checkpoint_state, @@ -878,6 +1301,7 @@ def gdn_replayssm_spec_decode( beta_cache, query_start_loc, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -905,6 +1329,7 @@ def gdn_replayssm_spec_decode( beta_cache, query_start_loc, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -940,6 +1365,7 @@ def gdn_replayssm_spec_decode( beta_cache, query_start_loc, ssm_state_indices, + replay_indices, write_pos, cache_base, is_flush, @@ -964,34 +1390,141 @@ def commit_gdn_replayssm_spec( 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 + replay_indices: torch.Tensor, # [n_rows] int active request row max_cache_len: int, max_spec_len: int, + fold_every_commit: bool = False, cache_buf_len: int | None = None, null_block_id: int = 0, ): - """Advance the block-keyed cursors once per decode step (device-only).""" + """Advance request-keyed cursors once per verify step (device-only).""" if cache_buf_len is None: cache_buf_len = max_cache_len - n_rows = state_batch_indices.shape[0] + n_rows = replay_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, + replay_indices, n_rows, - stride_sbi=state_batch_indices.stride(0), + stride_replay_indices=replay_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, + FOLD_EVERY_COMMIT=fold_every_commit, BLOCK=BLOCK, NULL_BLOCK_ID=null_block_id, ) +def commit_gdn_replayssm_circular( + *, + checkpoint_state: torch.Tensor, + d_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + d_residual_cache: torch.Tensor | None, + k_residual_cache: torch.Tensor | None, + state_batch_indices: torch.Tensor, + replay_indices: torch.Tensor, + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, + accept_lens: torch.Tensor, + mamba_track_indices: torch.Tensor | None = None, + mamba_steps_to_track: torch.Tensor | None = None, + null_block_id: int = -1, +) -> None: + """Materialize circular history only for capacity and track rows. + + ``write_pos`` must already include this step's accepted prefix. Capacity + rows update their active checkpoint and reset the ring; track rows snapshot + the exact state at the requested accepted step without disturbing the active + checkpoint. + """ + num_layers, _, HV, V, K = checkpoint_state.shape + H = k_cache.shape[2] + max_cache_len = d_cache.shape[-2] + B = state_batch_indices.shape[0] + has_residual = d_residual_cache is not None and k_residual_cache is not None + d_residual_buf = d_residual_cache if has_residual else d_cache + k_residual_buf = k_residual_cache if has_residual else k_cache + BV = min(triton.next_power_of_2(V), 64) + has_track = mamba_track_indices is not None and mamba_steps_to_track is not None + if has_track: + track_indices = mamba_track_indices + track_steps = mamba_steps_to_track + stride_track = track_indices.stride(0) + stride_steps = track_steps.stride(0) + else: + track_indices = state_batch_indices + track_steps = accept_lens + stride_track = 0 + stride_steps = 0 + grid = (triton.cdiv(V, BV), B, HV * num_layers) + gdn_replayssm_compact_commit_kernel[grid]( + checkpoint_state, + d_cache, + k_cache, + g_cache, + d_residual_buf, + k_residual_buf, + state_batch_indices, + replay_indices, + write_pos, + cache_base, + is_flush, + accept_lens, + track_indices, + track_steps, + checkpoint_state.stride(1), + d_cache.stride(1), + k_cache.stride(1), + g_cache.stride(1), + d_residual_buf.stride(1), + k_residual_buf.stride(1), + checkpoint_state.stride(0), + d_cache.stride(0), + k_cache.stride(0), + g_cache.stride(0), + d_residual_buf.stride(0), + k_residual_buf.stride(0), + state_batch_indices.stride(0), + replay_indices.stride(0), + accept_lens.stride(0), + stride_track, + stride_steps, + H=H, + HV=HV, + K=K, + V=V, + BV=BV, + MAX_CACHE_LEN=max_cache_len, + NULL_BLOCK_ID=null_block_id, + HAS_TRACK=has_track, + HAS_RESIDUAL=has_residual, + num_warps=4, + num_stages=2, + ) + block = triton.next_power_of_2(max(1, B)) + _finish_gdn_replayssm_circular_fold_kernel[(1,)]( + write_pos, + cache_base, + is_flush, + state_batch_indices, + replay_indices, + B, + stride_indices=state_batch_indices.stride(0), + stride_replay_indices=replay_indices.stride(0), + MAX_CACHE_LEN=max_cache_len, + BLOCK=block, + NULL_BLOCK_ID=null_block_id, + ) + + def reset_gdn_replayssm_spec_cursors( write_pos: torch.Tensor, cache_base: torch.Tensor, diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index d9bfab8c2..f2b0f7073 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -371,9 +371,9 @@ def handle_linear_attn_backend(server_args: Any): ) # ReplaySSM spec-verify (Part B of #28511): linear-chain target verify via - # fold-every-commit -- the verify stores each draft step's raw inputs into - # the per-slot (rawv, rawk, g, beta) window and the commit replays the - # accepted prefix into the fp32 checkpoint. The intra-window interaction + # compact cached replay. Verify stores normalized keys, update vectors, + # and fp32 log-decays; accepted BF16 windows are materialized with + # compensated hi/lo accumulation. 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. @@ -439,8 +439,8 @@ def handle_linear_attn_backend(server_args: Any): if cfg.mamba_ssm_dtype is None: logger.info( "--enable-linear-replayssm-spec: setting --mamba-ssm-dtype " - "float32 (the closed-loop exact fold keeps the SSM checkpoint " - "bit-identical to the recurrent baseline)." + "float32 (cached replay uses compensated checkpoint " + "projection and materialization)." ) declare_resolution( server_args, @@ -450,10 +450,8 @@ def handle_linear_attn_backend(server_args: Any): elif cfg.mamba_ssm_dtype != "float32": logger.warning( "--enable-linear-replayssm-spec with --mamba-ssm-dtype=%s: the " - "closed-loop fold re-quantizes the committed state each " - "commit/flush (fp32 keeps it bit-exact to the fp32 recurrent " - "baseline), so it may drift over long sequences. Validate " - "accuracy for your model.", + "compact checkpoint is materialized after each accepted " + "verify window; validate long-sequence accuracy and throughput.", cfg.mamba_ssm_dtype, ) diff --git a/python/sglang/srt/configs/mamba_utils.py b/python/sglang/srt/configs/mamba_utils.py index 4a9591d54..54689773d 100644 --- a/python/sglang/srt/configs/mamba_utils.py +++ b/python/sglang/srt/configs/mamba_utils.py @@ -125,29 +125,35 @@ class BaseLinearStateParams(ABC): ) * len(self.layers) def replayssm_ring_bytes_per_req(self, record_len: int) -> int: - """Per-slot bytes of the ReplaySSM spec-verify fold window (all - layers). Not part of ``mamba_cache_per_req``, so the memory solver - must charge it separately. MUST mirror the ``MambaPool`` allocation: - raw v/k in the conv dtype + fp32 beta, plus the fp32 gate ring - (per-head scalar for GDN, per-K vector for KDA). KDA additionally - keeps the chunked d/k rings under spec (its forward_decode routes on - their presence), also in the conv dtype.""" + """ReplaySSM spec-verify scratch bytes across all layers. + + GDN keeps compact d/k/g plus low parts for the activation-dtype d/k + rings. KDA keeps its raw-input fold window and d/k rings. + """ hv, v_dim, k_dim = self.shape.temporal h_k = self.shape.num_k_heads_per_tp conv_b = self.dtype.conv.itemsize fp32_b = 4 - per_layer = ( - hv * record_len * v_dim * conv_b # rawv - + h_k * record_len * k_dim * conv_b # rawk - + hv * record_len * fp32_b # beta (fp32) - # g (fp32): GDN per-head scalar, KDA per-K vector - + hv * record_len * (k_dim if self.is_kda else 1) * fp32_b - ) if self.is_kda: - per_layer += ( - hv * record_len * v_dim * conv_b # d + per_layer = ( + hv * record_len * v_dim * conv_b # rawv + + h_k * record_len * k_dim * conv_b # rawk + + hv * record_len * fp32_b # beta + + hv * record_len * k_dim * fp32_b # vector g + + hv * record_len * v_dim * conv_b # d + h_k * record_len * k_dim * conv_b # k ) + else: + per_layer = ( + hv * record_len * v_dim * conv_b # d + + h_k * record_len * k_dim * conv_b # normalized k + + hv * record_len * fp32_b # scalar g + ) + if self.dtype.conv != torch.float32: + per_layer += ( + hv * record_len * v_dim * conv_b # d low part + + h_k * record_len * k_dim * conv_b # normalized-k low part + ) return per_layer * len(self.layers) @property diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index d9be3424c..faca9e1d0 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -853,6 +853,7 @@ class GDNAttnBackend(MambaAttnBackendBase): mamba_pool=mamba_pool, layer_cache=mamba_cache_params, cache_indices=cache_indices, + replay_indices=forward_batch.req_pool_indices, query_start_loc=query_start_loc, draft_token_num=forward_batch.spec_info.draft_token_num, ) @@ -1017,6 +1018,7 @@ class GDNAttnBackend(MambaAttnBackendBase): mamba_pool: MambaPool, layer_cache: "MambaPool.SpeculativeState", cache_indices: torch.Tensor, + replay_indices: torch.Tensor, query_start_loc: torch.Tensor, draft_token_num: int, ) -> torch.Tensor: @@ -1024,15 +1026,13 @@ class GDNAttnBackend(MambaAttnBackendBase): 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`` + appends this window's drafts to the compact ``(d, k, g)`` rings. BF16 + checkpoints also keep low parts for compensated materialization. The rings + are per-layer (sliced via ``mamba2_layer_cache``), while the cursors + (write_pos, cache_base, is_flush) are request-slot pool attributes shared + by all GDN layers of the step. The cursors advance once per accepted step; + 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. """ @@ -1065,21 +1065,18 @@ class GDNAttnBackend(MambaAttnBackendBase): 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. + # BF16 compact D/K low parts; None for fp32 checkpoints. 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 + replay_indices=replay_indices, + # Request-slot cursors live on the pool (shared across all GDN layers) + # and advance only after acceptance, so the decode-path # forward_metadata.replayssm_write_pos snapshot is not used here. - write_pos=mamba_pool.replayssm_write_pos, + write_pos=mamba_pool.replayssm_spec_write_pos, cache_base=mamba_pool.replayssm_cache_base, is_flush=mamba_pool.replayssm_is_flush, max_cache_len=max_cache_len, @@ -1090,6 +1087,9 @@ class GDNAttnBackend(MambaAttnBackendBase): # 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, + # Capacity folds are committed once across every GDN layer after + # acceptance; the active path is therefore a single launch/layer. + launch_mode="verify", ) # 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 6ac0bfe1f..92d1007e2 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2276,22 +2276,14 @@ class KVCacheConfigurator: # no longer reserves the (1 + D/ratio) intermediate factor -- the whole # budget goes to persistent slots (K sized like non-spec), which is how the # freed ~9GB turns into higher max_running. - # The ring is allocated per slot but is not part of mamba_cache_per_req; - # the solve must charge it too or num_slots is over-provisioned. + # The ring is not part of mamba_cache_per_req. GDN replay is fixed-size + # request scratch; KDA replay remains attached to each mamba slot. replayssm_active = get_exec().mamba.enable_linear_replayssm_spec and ( self.hybrid_gdn_config is not None or kimi_linear_config(self.model_config) is not None ) if replayssm_active: - # GDN sizes the fold window to the draft maximum; the KDA ring - # stays --linear-replayssm-cache-len long (mirrors MambaPool). - max_draft_tokens = max_speculative_num_draft_tokens() - if kimi_linear_config(self.model_config) is not None: - record_len = get_exec().mamba.linear_replayssm_cache_len - elif max_draft_tokens is not None: - record_len = max_draft_tokens - else: - record_len = get_exec().mamba.linear_replayssm_cache_len + record_len = get_exec().mamba.linear_replayssm_cache_len replayssm_ring_per_req = ( config.mamba2_cache_params.replayssm_ring_bytes_per_req( record_len=record_len @@ -2300,6 +2292,15 @@ class KVCacheConfigurator: else: replayssm_ring_per_req = 0 replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale) + if replayssm_active and kimi_linear_config(self.model_config) is None: + replay_req_slots = ( + get_schedule().max_running_requests // self.ps.attn_dp_size + 1 + ) + replayssm_fixed_bytes = replayssm_ring_per_req * replay_req_slots + replayssm_ring_per_slot = 0 + else: + replayssm_fixed_bytes = 0 + replayssm_ring_per_slot = replayssm_ring_per_req if has_spec_dec: assert get_spec().speculative_num_draft_tokens is not None assert get_schedule().max_running_requests is not None @@ -2380,11 +2381,12 @@ class KVCacheConfigurator: intermediate_size = per_req * (capped_reqs + 1) * D total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30)) else: - per_slot = per_req + replayssm_ring_per_req + per_slot = per_req + replayssm_ring_per_slot get_context().override( "mamba_pool.memory_budget", max_mamba_cache_size=int( - (mamba_budget_bytes - per_slot) // per_slot + (mamba_budget_bytes - replayssm_fixed_bytes - per_slot) + // per_slot ), ) @@ -2404,14 +2406,12 @@ class KVCacheConfigurator: f"(4) use GPUs with more memory." ) - # +1: the pool's padding slot is allocated alongside the request slots. - # ReplaySSM ring rides on every slot too (replayssm_ring_per_req is 0 when - # the ring is not allocated). + # +1 accounts for each pool's padding slot. mamba_state_memory = ( (get_schedule().max_mamba_cache_size + 1) - * (stage_per_req + replayssm_ring_per_req) - / (1 << 30) - ) + * (stage_per_req + replayssm_ring_per_slot) + + replayssm_fixed_bytes + ) / (1 << 30) return total_rest_memory - mamba_state_memory diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 0847748d0..b170775e6 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -568,10 +568,23 @@ class MambaRadixCache(BasePrefixCache): # donate to the last flush boundary (where temporal is current) # and reset the cursor, keeping the donated checkpoint consistent # with its key length. page_size is asserted == 1, so no realign. - write_pos_buf = self.req_to_token_pool.mamba_pool.replayssm_write_pos + mamba_pool = self.req_to_token_pool.mamba_pool + write_pos_buf = mamba_pool.replayssm_write_pos + cursor_idx = req.kv.mamba_pool_idx + if write_pos_buf is None: + write_pos_buf = getattr( + mamba_pool, "replayssm_spec_write_pos", None + ) + cursor_idx = req.kv.req_pool_idx if write_pos_buf is not None: - cache_len -= int(write_pos_buf[req.kv.mamba_pool_idx].item()) - write_pos_buf[req.kv.mamba_pool_idx] = 0 + cache_len -= int(write_pos_buf[cursor_idx].item()) + write_pos_buf[cursor_idx] = 0 + if ( + getattr(mamba_pool, "replayssm_spec_write_pos", None) + is not None + ): + mamba_pool.replayssm_cache_base[cursor_idx] = 0 + mamba_pool.replayssm_is_flush[cursor_idx] = 0 if cache_len is None: cache_len = 0 if cache_len != len(token_ids): @@ -688,6 +701,11 @@ class MambaRadixCache(BasePrefixCache): if self.enable_mamba_extra_buffer else len(token_ids) ) + spec_write_pos = getattr( + self.req_to_token_pool.mamba_pool, "replayssm_spec_write_pos", None + ) + if not self.enable_mamba_extra_buffer and spec_write_pos is not None: + cache_len -= int(spec_write_pos[req.kv.req_pool_idx].item()) if self.disable or cache_len is None: return _skip_cache_unfinished_req(req) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 14ae9cdc5..59cec0fb5 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -383,14 +383,8 @@ 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-linear-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. + # Under GDN spec verify, rawv/rawk hold compact D/K low parts and beta is + # None. KDA uses the same fields for its raw-input fold window. replayssm_d: Optional[torch.Tensor] = None replayssm_k: Optional[torch.Tensor] = None replayssm_g: Optional[torch.Tensor] = None @@ -529,12 +523,13 @@ class MambaPool: self.linear_replayssm_cache_len = linear_replayssm_cache_len # ReplaySSM: the decode ring (--enable-linear-replayssm) allocates the # chunked (d, k) records + write_pos; the spec-verify flag - # (--enable-linear-replayssm-spec) always uses fold-every-commit and - # allocates only the raw (v, k, g, beta) window -- no chunked records, - # no cursors (KDA additionally keeps d/k, see the allocation below). - # The shared g allocation gates on `_replayssm_on`. + # (--enable-linear-replayssm-spec) uses compact replay for GDN and raw + # fold-every-commit for KDA. The shared g allocation gates on + # `_replayssm_on`. self.enable_linear_replayssm_spec = enable_linear_replayssm_spec - self.replayssm_spec_fold = bool(enable_linear_replayssm_spec) + self.replayssm_spec_fold = bool( + enable_linear_replayssm_spec and cache_params.is_kda + ) _replayssm_on = enable_linear_replayssm or enable_linear_replayssm_spec # for disagg with nvlink @@ -623,16 +618,15 @@ class MambaPool: 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 dtype. DECODE ring (--enable-linear-replayssm): records - # follow the SSM dtype -- its flush folds `d` directly into the - # state. SPEC-verify ring (--enable-linear-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. + # GDN speculative replay is request-lifetime scratch. Size it by + # active requests instead of every persistent radix-cache slot. + num_slots = ( + spec_state_size + 1 + if enable_linear_replayssm_spec and not cache_params.is_kda + else size + 1 + ) + # Decode records follow the SSM dtype. Spec-verify compact d/k + # records follow the activation dtype; g stays fp32. ring_dtype = conv_dtype if enable_linear_replayssm_spec else ssm_dtype # Fold-every-commit: one verify window, no chunked (d, k) # records. KDA is the exception on both counts: its window @@ -674,14 +668,10 @@ 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_linear_replayssm_spec: - if cache_params.is_kda: + # KDA still uses raw-input fold-every-commit. GDN materializes + # its compact d/k/g history directly and needs no duplicate ring. + if enable_linear_replayssm_spec and cache_params.is_kda: + if cache_params.is_kda or not self.replayssm_spec_fold: # Backstop for the KDA ring invariants; this pool is # sized with the final adaptive-aware draft maximum. if L & (L - 1) != 0: @@ -711,6 +701,20 @@ class MambaPool: dtype=torch.float32, device=device, ) + elif enable_linear_replayssm_spec and ring_dtype != torch.float32: + # Low parts of compact D and normalized K. The rings follow + # the activation dtype regardless of checkpoint dtype, so + # materialization always needs both parts. + replayssm_rawv = torch.zeros( + size=(num_mamba_layers, num_slots, hv, record_len, v_dim), + dtype=conv_dtype, + device=device, + ) + replayssm_rawk = torch.zeros( + size=(num_mamba_layers, num_slots, h_k, record_len, k_dim), + dtype=conv_dtype, + device=device, + ) if speculative_num_draft_tokens is not None: if _is_npu: @@ -875,8 +879,12 @@ class MambaPool: + ( 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_linear_replayssm_spec + + ( + f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB " + if replayssm_beta is not None + else "" + ) + if replayssm_rawv is not None else "" ) ) @@ -884,26 +892,24 @@ class MambaPool: # 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(_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 (decode ring) or once per verify step by - # the worker (spec-verify ring). Index 0..size; reset on slot (re)alloc. + # Decode ReplaySSM remains keyed by persistent mamba slots. self.replayssm_write_pos = ( torch.zeros((size + 1,), dtype=torch.int32, device=device) - if _replayssm_on and not self.replayssm_spec_fold + if enable_linear_replayssm + else None + ) + self.replayssm_spec_write_pos = ( + torch.zeros((spec_state_size + 1,), dtype=torch.int32, device=device) + if enable_linear_replayssm_spec and not self.replayssm_spec_fold 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. self.replayssm_cache_base = ( - torch.zeros((size + 1,), dtype=torch.int32, device=device) + torch.zeros((spec_state_size + 1,), dtype=torch.int32, device=device) if enable_linear_replayssm_spec and not self.replayssm_spec_fold else None ) self.replayssm_is_flush = ( - torch.zeros((size + 1,), dtype=torch.int8, device=device) + torch.zeros((spec_state_size + 1,), dtype=torch.int8, device=device) if enable_linear_replayssm_spec and not self.replayssm_spec_fold else None ) @@ -1031,12 +1037,6 @@ 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() @@ -1047,46 +1047,22 @@ 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): - # Accept both the legacy 2-tuple (conv, temporal) and the 3-tuple that also - # carries the ReplaySSM spec-verify cursors. + # Accept historical 3-tuples, but request-keyed replay scratch is not + # restored with a physical checkpoint slot. if len(mamba_cache_cpu) == 3: - conv_cpu, temporal_cpu, cursors_cpu = mamba_cache_cpu + conv_cpu, temporal_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() _NON_TRANSFER_STATE_FIELDS = frozenset( @@ -1349,10 +1325,21 @@ class HybridReqToTokenPool(ReqToTokenPool): # For chunk prefill req, we do not need to allocate mamba cache, # We could use allocated mamba cache instead. def alloc(self, reqs: List[Req]) -> Optional[List[int]]: + fresh_req_rows = [req.kv.req_pool_idx is None for req in reqs] select_index = super().alloc(reqs) if select_index is None: return None + spec_write_pos = getattr(self.mamba_pool, "replayssm_spec_write_pos", None) + if spec_write_pos is not None: + fresh_indices = [ + idx for idx, fresh in zip(select_index, fresh_req_rows) if fresh + ] + if fresh_indices: + spec_write_pos[fresh_indices] = 0 + self.mamba_pool.replayssm_cache_base[fresh_indices] = 0 + self.mamba_pool.replayssm_is_flush[fresh_indices] = 0 + mamba_indices: list[torch.Tensor] = [] mamba_ping_pong_track_buffers: list[torch.Tensor] = [] for req in reqs: @@ -1371,12 +1358,6 @@ 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.kv.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.kv.mamba_pool_idx] = 0 - self.mamba_pool.replayssm_is_flush[req.kv.mamba_pool_idx] = 0 mamba_indices.append(req.kv.mamba_pool_idx) if self.enable_mamba_extra_buffer: if req.kv.mamba_ping_pong_track_buffer is None: diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 95dfb753f..941a94535 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -913,6 +913,7 @@ def commit_mamba_states_after_verify( 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_circular, commit_gdn_replayssm_spec, ) from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( @@ -922,37 +923,59 @@ def commit_mamba_states_after_verify( 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). - last_correct_step_indices, _ = _verify_commit_step_indices( + replay_indices = batch.req_pool_indices + last_correct_step_indices, mamba_steps_to_track = _verify_commit_step_indices( batch=batch, accept_index=accept_index, accept_lens=accept_lens, draft_token_num=draft_token_num, ) + # Advance the per-request 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_spec_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 + replay_indices=replay_indices, + max_cache_len=spec_state.replayssm_d.shape[-2], + max_spec_len=draft_token_num, + fold_every_commit=spec_state.temporal.dtype != torch.float32, + null_block_id=-1, # SGLang: valid slots >= 0, padding == -1 + ) + # Capacity rows fold all layers in one launch; track rows snapshot the + # exact crossing state without disturbing the active circular history. + commit_gdn_replayssm_circular( + checkpoint_state=spec_state.temporal, + d_cache=spec_state.replayssm_d, + k_cache=spec_state.replayssm_k, + g_cache=spec_state.replayssm_g, + d_residual_cache=spec_state.replayssm_rawv, + k_residual_cache=spec_state.replayssm_rawk, + state_batch_indices=state_batch_indices, + replay_indices=replay_indices, + write_pos=mamba_pool.replayssm_spec_write_pos, + cache_base=mamba_pool.replayssm_cache_base, + is_flush=mamba_pool.replayssm_is_flush, + accept_lens=accept_lens, + mamba_track_indices=batch.mamba_track_indices, + mamba_steps_to_track=mamba_steps_to_track, + null_block_id=-1, + ) + # Roll back active conv state and snapshot its interval-crossing window. 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-linear-replayssm-spec), so the per-track scatters are intentionally - # skipped here. + if batch.mamba_track_indices is not None: + fused_conv_window_scatter_with_mask( + spec_state.conv[0], + spec_state.intermediate_conv_window[0], + batch.mamba_track_indices, + mamba_steps_to_track, + ) return # KDA ReplaySSM (fold-every-commit): KDA keeps its own recurrent verify kernel diff --git a/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py b/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py index 578bc0b1f..4696e5ff1 100644 --- a/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py +++ b/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py @@ -28,10 +28,8 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") # temporal = (hv=4, v_dim=8, k_dim=8), num_k_heads_per_tp = 4, record_len = 8, # 2 layers. conv bf16 (2B), fp32 gate/beta (4B). Ring tensors (per slot, per # layer): -# rawv hv*RL*v_dim, rawk h_k*RL*k_dim -> conv dtype -# g hv*RL (GDN) / hv*RL*k_dim (KDA) -> fp32 -# beta hv*RL -> fp32 -# d/k like rawv/rawk -> conv dtype (KDA only) +# GDN stores compact d/k plus scalar g. KDA stores raw v/k, vector g, beta, +# and its existing d/k rings. DTYPE = Mamba2StateDType(conv=torch.bfloat16, temporal=torch.float32) RL = 8 LAYERS = [0, 1] @@ -44,7 +42,7 @@ def _kda_params(): return KimiLinearCacheParams(shape=shape, dtype=DTYPE, layers=LAYERS) -def _gdn_params(): +def _gdn_params(temporal_dtype=torch.float32): # Only shape.temporal and shape.num_k_heads_per_tp are read here; the rest # are dummy (the accounting does not depend on them). shape = Mamba2StateShape( @@ -59,15 +57,20 @@ def _gdn_params(): conv_kernel=0, num_k_heads_per_tp=4, ) - return Mamba2CacheParams(shape=shape, dtype=DTYPE, layers=LAYERS) + dtype = Mamba2StateDType(conv=torch.bfloat16, temporal=temporal_dtype) + return Mamba2CacheParams(shape=shape, dtype=dtype, layers=LAYERS) class TestReplaySSMRingAccounting(CustomTestCase): def test_gdn_fold(self): - # fold window: rawv 512 + rawk 512 + g(scalar, 4*8*4) 128 + beta 128 = 1280 + # d 512 + normalized k 512 + scalar g 128 + d/k low parts 1024 = 2176 self.assertEqual( _gdn_params().replayssm_ring_bytes_per_req(record_len=RL), - 1280 * len(LAYERS), + 2176 * len(LAYERS), + ) + self.assertEqual( + _gdn_params(torch.bfloat16).replayssm_ring_bytes_per_req(record_len=RL), + 2176 * len(LAYERS), ) def test_kda_fold(self):