minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[GDN] Amortize ReplaySSM checkpoint materialization (#35544)

Browse Source
This commit is contained in:
YAMY
2026-09-04 15:13:20 -07:00
committed by GitHub
parent e3eeabbbfa
commit db89f639ef
9 changed files with 853 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py
+649 -116
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/arg_groups/attention_hook.py
+7 -9
View File
@@ -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,
)
python/sglang/srt/configs/mamba_utils.py
+22 -16
View File
@@ -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
python/sglang/srt/layers/attention/linear/gdn_backend.py
+17 -17
View File
@@ -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)
python/sglang/srt/mem_cache/kv_cache_configurator.py
+19 -19
View File
@@ -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
python/sglang/srt/mem_cache/mamba_radix_cache.py
+21 -3
View File
@@ -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)
python/sglang/srt/mem_cache/memory_pool.py
+64 -83
View File
@@ -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:
python/sglang/srt/speculative/spec_utils.py
+43 -20
View File
@@ -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