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

Store mamba prefix-cache checkpoints at the configured SSM state dtype (#34820)

Browse Source
This commit is contained in:
Yuhao Yang
2026-09-09 15:25:57 +08:00
committed by GitHub
parent daf66f6670
commit 13469c16d3
17 changed files with 740 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/chunk_delta_h.py
+49
View File
@@ -63,6 +63,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
stride_init_state,
cu_seqlens,
chunk_offsets,
track_state,
track_chunk_idx,
stride_track_state,
T,
H: tl.constexpr,
Hg: tl.constexpr,
@@ -78,6 +81,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
IS_VARLEN: tl.constexpr,
NT_BUCKET: tl.constexpr,
USE_EXP2: tl.constexpr,
TRACK_STATE: tl.constexpr,
):
i_v, i_nh = tl.program_id(0), tl.program_id(1)
i_n, i_h = i_nh // H, i_nh % H
@@ -130,6 +134,15 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
if INPLACE_UPDATE:
ht = ht + i_h * V * K
if TRACK_STATE:
i_track = tl.load(track_chunk_idx + i_n).to(tl.int32)
p_track_base = track_state + (i_n * stride_track_state + i_h * V * K).to(
tl.int64
)
else:
i_track = -1
p_track_base = track_state
# load initial state
if USE_INITIAL_STATE and valid_state:
p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
@@ -172,6 +185,27 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
)
tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
if TRACK_STATE and i_t == i_track:
p_t1 = tl.make_block_ptr(
p_track_base, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)
)
tl.store(p_t1, b_h1, boundary_check=(0, 1))
if K > 64:
p_t2 = tl.make_block_ptr(
p_track_base, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
)
tl.store(p_t2, b_h2, boundary_check=(0, 1))
if K > 128:
p_t3 = tl.make_block_ptr(
p_track_base, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
)
tl.store(p_t3, b_h3, boundary_check=(0, 1))
if K > 192:
p_t4 = tl.make_block_ptr(
p_track_base, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
)
tl.store(p_t4, b_h4, boundary_check=(0, 1))
p_w = tl.make_block_ptr(
w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)
)
@@ -326,10 +360,21 @@ def chunk_gated_delta_rule_fwd_h(
cu_seqlens: Optional[torch.LongTensor] = None,
chunk_indices: Optional[torch.LongTensor] = None,
use_exp2: bool = False,
track_state: Optional[torch.Tensor] = None,
track_chunk_idx: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
assert not (use_exp2 and g is not None), (
"use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
)
assert (track_state is None) == (track_chunk_idx is None), (
"track_state and track_chunk_idx must be passed together"
)
if track_state is not None:
# The caller rounds once to the pool dtype; a narrower buffer would
# silently double-round the snapshot.
assert track_state.dtype == torch.float32, (
f"track_state must be fp32, got {track_state.dtype}"
)
B, T, Hg, K, V = *k.shape, u.shape[-1]
H = u.shape[-2]
BT = CHUNK_SIZE
@@ -369,6 +414,9 @@ def chunk_gated_delta_rule_fwd_h(
stride_init_state=(initial_state.stride(0) if initial_state is not None else 0),
cu_seqlens=cu_seqlens,
chunk_offsets=chunk_offsets,
track_state=track_state,
track_chunk_idx=track_chunk_idx,
stride_track_state=(track_state.stride(0) if track_state is not None else 0),
T=T,
H=H,
Hg=Hg,
@@ -383,5 +431,6 @@ def chunk_gated_delta_rule_fwd_h(
IS_VARLEN=cu_seqlens is not None,
NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)),
USE_EXP2=use_exp2,
TRACK_STATE=track_state is not None,
)
return h, v_new
python/sglang/kernels/ops/attention/fla/kda.py
+8
View File
@@ -1094,6 +1094,8 @@ def chunk_kda_fwd(
dt_bias: Optional[torch.Tensor] = None,
lower_bound: Optional[float] = None,
output_intermediate_states: bool = False,
track_state: Optional[torch.Tensor] = None,
track_chunk_idx: Optional[torch.Tensor] = None,
):
chunk_size = 64
# Pre-compute chunk indices once and thread through all downstream kernels.
@@ -1169,6 +1171,8 @@ def chunk_kda_fwd(
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=True,
track_state=track_state,
track_chunk_idx=track_chunk_idx,
)
del w, u, kg
@@ -1210,6 +1214,8 @@ def chunk_kda(
dt_bias: Optional[torch.Tensor] = None,
lower_bound: Optional[float] = None,
output_intermediate_states: bool = False,
track_state: Optional[torch.Tensor] = None,
track_chunk_idx: Optional[torch.Tensor] = None,
beta_is_raw: bool = False,
**kwargs,
):
@@ -1238,4 +1244,6 @@ def chunk_kda(
dt_bias=dt_bias,
lower_bound=lower_bound,
output_intermediate_states=output_intermediate_states,
track_state=track_state,
track_chunk_idx=track_chunk_idx,
)
python/sglang/kernels/ops/attention/helion/kda_prefill.py
+93 -2
View File
@@ -1024,6 +1024,40 @@ _STATE_VARLEN_SMALL_HEAD_CONFIG = helion.Config(
range_unroll_factors=[0, 0],
)
# Track variants of the two varlen configs. The fp32 track snapshot adds one
# load + one store mid-body, which would shift the positional indexing /
# eviction lists above; separate kernels keep the non-track configs untouched.
# Tracked batches are rare (prefix-cache checkpointing only), so these trade
# the hand-tuned positional lists for plainly correct settings.
_STATE_VARLEN_TRACK_CONFIG = helion.Config(
atomic_indexing=[],
block_sizes=[64],
indexing="pointer",
l2_groupings=[4],
loop_orders=[[0, 2, 1]],
num_stages=2,
num_warps=4,
pid_type="flat",
range_flattens=[None, None],
range_multi_buffers=[None, False],
range_num_stages=[],
range_unroll_factors=[0, 2],
)
_STATE_VARLEN_SMALL_HEAD_TRACK_CONFIG = helion.Config(
atomic_indexing=[],
block_sizes=[32],
indexing="pointer",
l2_groupings=[1],
loop_orders=[[1, 2, 0]],
num_stages=3,
num_warps=8,
pid_type="flat",
range_flattens=[None, None],
range_multi_buffers=[None, True],
range_num_stages=[],
range_unroll_factors=[0, 0],
)
@helion.kernel(
static_shapes=False,
@@ -1040,6 +1074,9 @@ def _chunk_state(
cu_seqlens: torch.Tensor,
chunk_indices: torch.Tensor,
chunk_offsets: torch.Tensor,
track_state: torch.Tensor,
track_chunk_idx: torch.Tensor,
has_track: hl.constexpr, # pyrefly: ignore[bad-function-definition]
is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition]
) -> tuple[torch.Tensor, torch.Tensor]:
"""Propagate KDA state between chunks and update the state pool in place."""
@@ -1070,6 +1107,11 @@ def _chunk_state(
initial_state.stride(2),
initial_state.stride(3),
initial_state_indices.stride(0),
track_state.stride(0),
track_state.stride(1),
track_state.stride(2),
track_state.stride(3),
track_chunk_idx.stride(0),
)
)
@@ -1108,6 +1150,9 @@ def _chunk_state(
:,
].float()
if has_track:
i_track = track_chunk_idx[tile_sequence.id]
for token_tile in hl.tile(sequence_length, block_size=64):
global_chunk = output_offset + token_tile.id
h_rows[
@@ -1115,6 +1160,17 @@ def _chunk_state(
tile_v,
:,
] = state.to(h.dtype)
if has_track:
# Snapshot the fp32 accumulator at the tracked chunk boundary
# (the h store above rounds to the activation dtype; -1 marks
# untracked sequences and never matches a chunk id).
if token_tile.id == i_track:
track_state[
tile_sequence.id,
tile_h.id,
tile_v.index,
:,
] = state
token = begin + token_tile.index
valid = token < end
row = token * H + tile_h.id
@@ -1163,13 +1219,29 @@ _chunk_state_varlen_small_head = helion.kernel(
config=_STATE_VARLEN_SMALL_HEAD_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)(_chunk_state.fn)
_chunk_state_varlen_track = helion.kernel(
static_shapes=False,
config=_STATE_VARLEN_TRACK_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)(_chunk_state.fn)
_chunk_state_varlen_small_head_track = helion.kernel(
static_shapes=False,
config=_STATE_VARLEN_SMALL_HEAD_TRACK_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)(_chunk_state.fn)
def _select_state_kernel(*, is_varlen: bool, num_heads: int) -> helion.Kernel:
def _select_state_kernel(
*, is_varlen: bool, num_heads: int, has_track: bool
) -> helion.Kernel:
if not is_varlen:
return _chunk_state
if num_heads <= _PREFILL_SMALL_HEAD_THRESHOLD:
if has_track:
return _chunk_state_varlen_small_head_track
return _chunk_state_varlen_small_head
if has_track:
return _chunk_state_varlen_track
return _chunk_state_varlen
@@ -1314,6 +1386,8 @@ def chunk_kda(
dt_bias: torch.Tensor | None = None,
lower_bound: float | None = None,
output_intermediate_states: bool = False,
track_state: torch.Tensor | None = None,
track_chunk_idx: torch.Tensor | None = None,
beta_is_raw: bool = False,
**kwargs: object,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
@@ -1322,6 +1396,13 @@ def chunk_kda(
scale = k.shape[-1] ** -0.5
if initial_state is None or initial_state_indices is None:
raise ValueError("KDA prefill requires an indexed initial-state pool")
assert (track_state is None) == (track_chunk_idx is None), (
"track_state and track_chunk_idx must be passed together"
)
if track_state is not None:
assert track_state.dtype == torch.float32, (
f"track_state must be fp32, got {track_state.dtype}"
)
num_tokens = q.shape[1]
if g.shape[1] < num_tokens or beta.shape[1] < num_tokens:
@@ -1349,6 +1430,8 @@ def chunk_kda(
dt_bias=dt_bias,
lower_bound=lower_bound,
output_intermediate_states=output_intermediate_states,
track_state=track_state,
track_chunk_idx=track_chunk_idx,
)
q = q.contiguous()
@@ -1395,7 +1478,10 @@ def chunk_kda(
else:
metadata = torch.empty(0, device=q.device, dtype=torch.int32)
chunk_offsets = torch.empty(0, device=q.device, dtype=torch.long)
state_kernel = _select_state_kernel(is_varlen=is_varlen, num_heads=q.size(2))
has_track = track_state is not None
state_kernel = _select_state_kernel(
is_varlen=is_varlen, num_heads=q.size(2), has_track=has_track
)
h, v_new = state_kernel(
kg,
w,
@@ -1410,6 +1496,11 @@ def chunk_kda(
else torch.empty(0, 2, device=q.device, dtype=torch.long)
),
chunk_offsets,
# Unused when has_track is False; bind in-place tensors as stand-ins so
# the kernel signature always sees real tensors.
track_state if has_track else initial_state,
track_chunk_idx if has_track else initial_state_indices,
has_track,
is_varlen,
)
if chunk_indices is None:
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+36 -7
View File
@@ -114,7 +114,9 @@ class MambaAttnBackendBase(AttentionBackend):
retrieve_parent_token = None
track_conv_indices = None
track_ssm_h_src = None
track_chunk_idx = None
track_ssm_h_dst = None
track_ssm_h_batch_src = None
track_ssm_final_src = None
track_ssm_final_dst = None
track_ssm_seq_idx = None
@@ -250,8 +252,10 @@ class MambaAttnBackendBase(AttentionBackend):
)
(
track_chunk_idx,
track_ssm_h_src,
track_ssm_h_dst,
track_ssm_h_batch_src,
track_ssm_final_src,
track_ssm_final_dst,
track_ssm_seq_idx,
@@ -273,8 +277,10 @@ class MambaAttnBackendBase(AttentionBackend):
track_conv_indices=track_conv_indices,
track_ssm_h_src=track_ssm_h_src,
track_ssm_h_dst=track_ssm_h_dst,
track_ssm_h_batch_src=track_ssm_h_batch_src,
track_ssm_final_src=track_ssm_final_src,
track_ssm_final_dst=track_ssm_final_dst,
track_chunk_idx=track_chunk_idx,
track_ssm_seq_idx=track_ssm_seq_idx,
track_ssm_end_locs=track_ssm_end_locs,
track_ssm_recompute_dst=track_ssm_recompute_dst,
@@ -358,7 +364,9 @@ class MambaAttnBackendBase(AttentionBackend):
):
"""src/dst indices to track SSM states for prefix caching: aligned seqs
cache last_recurrent_state, unaligned cache intermediate `h` at the last
chunk boundary."""
chunk boundary. Also returns ``track_ssm_h_batch_src``: the batch rows of
the unaligned tracked seqs, used to integer-index the fp32 snapshot
buffer on the KDA path so the copy stays free of GPU syncs."""
state_chunk_size = self.mamba_chunk_size
# CPU to avoid kernel launches for the masking ops
mamba_track_mask = forward_batch.mamba_track_mask.cpu()
@@ -416,9 +424,17 @@ class MambaAttnBackendBase(AttentionBackend):
def to_device(t):
return None if t is None else t.to(self.device, non_blocking=True)
track_chunk_idx = torch.full((lens_to_track.shape[0],), -1, dtype=torch.int32)
tracked_seqs = mamba_track_mask.nonzero(as_tuple=True)[0][not_aligned]
track_chunk_idx[tracked_seqs] = (
lens_masked[not_aligned] // state_chunk_size
).to(torch.int32)
return (
to_device(track_chunk_idx),
to_device(track_ssm_h_src),
to_device(track_ssm_h_dst),
to_device(tracked_seqs),
to_device(track_ssm_final_src),
to_device(track_ssm_final_dst),
to_device(track_ssm_seq_idx),
@@ -887,18 +903,31 @@ class MambaAttnBackendBase(AttentionBackend):
ssm_states: torch.Tensor,
forward_metadata: ForwardMetadata,
track_states: Optional[torch.Tensor] = None,
*,
h_track_buf: Optional[torch.Tensor] = None,
):
"""Copy extend SSM state at the last chunk boundary to track slots (source
depends on chunk alignment; see `_init_track_ssm_indices`)."""
depends on chunk alignment; see `_init_track_ssm_indices`).
Unaligned rows read the fp32 ``h_track_buf`` snapshot written in-kernel
when given (its rows follow the batch, selected by the integer index
``track_ssm_h_batch_src`` — a boolean mask would nonzero() and sync the
stream once per layer); otherwise they fall back to the per-chunk
states ``h`` (already rounded to the activation dtype)."""
if forward_metadata.has_mamba_track_mask:
# Triton always returns h; FlashInfer returns it only when checkpoints
# were requested. Aligned-only tracking reads the final state below.
if forward_metadata.track_ssm_h_src.numel() > 0:
assert h is not None
h = h.squeeze(0)
ssm_states[forward_metadata.track_ssm_h_dst] = h[
forward_metadata.track_ssm_h_src
].to(ssm_states.dtype, copy=False)
if h_track_buf is not None:
ssm_states[forward_metadata.track_ssm_h_dst] = h_track_buf[
forward_metadata.track_ssm_h_batch_src
].to(ssm_states.dtype, copy=False)
else:
assert h is not None
h = h.squeeze(0)
ssm_states[forward_metadata.track_ssm_h_dst] = h[
forward_metadata.track_ssm_h_src
].to(ssm_states.dtype, copy=False)
if (
forward_metadata.track_ssm_recompute_dst is not None
and forward_metadata.track_ssm_recompute_dst.numel() > 0
python/sglang/srt/layers/attention/linear/kda_backend.py
+44 -6
View File
@@ -330,6 +330,14 @@ class KDAKernelDispatcher:
**kwargs,
)
def effective_extend_kernel(self, lower_bound: Optional[float]):
"""The kernel ``extend`` will actually run: safe-gate models reroute
kernels without ``supports_safe_gate`` to Triton."""
kernel = self.extend_kernel
if lower_bound is not None and not getattr(kernel, "supports_safe_gate", True):
kernel = self.triton_kernel
return kernel
def extend(
self,
q: torch.Tensor,
@@ -343,11 +351,7 @@ class KDAKernelDispatcher:
query_start_loc: torch.Tensor,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor | None]:
kernel = self.extend_kernel
if kwargs.get("lower_bound") is not None and not getattr(
kernel, "supports_safe_gate", True
):
kernel = self.triton_kernel
kernel = self.effective_extend_kernel(kwargs.get("lower_bound"))
return kernel.extend(
q,
k,
@@ -857,6 +861,34 @@ class KDAAttnBackend(MambaAttnBackendBase):
a = a.unflatten(-1, (-1, layer.head_k_dim))
track_ssm = self.forward_metadata.has_mamba_track_mask
track_chunk_idx = self.forward_metadata.track_chunk_idx
h_track_buf = None
if (
track_ssm
and track_chunk_idx is not None
# Same rows as track_ssm_h_batch_src, but known without a GPU sync.
and self.forward_metadata.track_ssm_h_src.numel() > 0
):
# fp32 scratch the kernel snapshots the tracked chunk-boundary
# states into (rows follow the batch; untracked rows stay unread).
# A kernel that does not declare support would leave the buffer
# unwritten and corrupt prefix-cache restores — fail loudly here.
# Check the kernel the dispatcher will actually run (safe-gate
# reroute included), not just the configured one.
extend_kernel = self.kernel_dispatcher.effective_extend_kernel(
layer.lower_bound
)
assert extend_kernel.supports_track_state_snapshot, (
f"{type(extend_kernel).__name__} cannot write the fp32 track "
f"snapshot required by the mamba track path; use "
f"--linear-attn-prefill-backend triton or "
f"--mamba-radix-cache-strategy no_buffer"
)
h_track_buf = torch.empty(
(track_chunk_idx.shape[0], *ssm_states.shape[1:]),
dtype=torch.float32,
device=ssm_states.device,
)
core_attn_out = self.kernel_dispatcher.extend(
q=q,
k=k,
@@ -881,6 +913,8 @@ class KDAAttnBackend(MambaAttnBackendBase):
track_ssm_h_src=(
self.forward_metadata.track_ssm_h_src if track_ssm else None
),
track_state=h_track_buf,
track_chunk_idx=(track_chunk_idx if h_track_buf is not None else None),
)
if track_ssm:
# Snapshot the SSM state at the last track-aligned chunk boundary
@@ -888,7 +922,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
# ping-pong track slots (see _init_track_ssm_indices).
core_attn_out, h = core_attn_out
self._track_mamba_state_extend(
forward_batch, h, ssm_states, self.forward_metadata
forward_batch,
h,
ssm_states,
self.forward_metadata,
h_track_buf=h_track_buf,
)
if logical_num_tokens < physical_num_tokens:
python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py
+10
View File
@@ -41,6 +41,8 @@ def _triton_fallback(
lower_bound=None,
beta_is_raw=False,
return_intermediate_states=False,
track_state=None,
track_chunk_idx=None,
):
"""Fall back to the Triton chunk_kda kernel (handles all preprocessing).
@@ -67,6 +69,8 @@ def _triton_fallback(
lower_bound=lower_bound,
beta_is_raw=beta_is_raw,
output_intermediate_states=return_intermediate_states,
track_state=track_state,
track_chunk_idx=track_chunk_idx,
)
@@ -83,6 +87,10 @@ class FlashKDAKernel(LinearAttnKernelBase):
Requires an SM90+ GPU with the ``flash_kda`` package.
"""
# Tracked batches always take the Triton fallback, which forwards the
# fp32 snapshot arguments (see _triton_fallback).
supports_track_state_snapshot: bool = True
def decode(
self,
q: torch.Tensor,
@@ -141,6 +149,8 @@ class FlashKDAKernel(LinearAttnKernelBase):
lower_bound=lower_bound,
beta_is_raw=beta_is_raw,
return_intermediate_states=return_intermediate_states,
track_state=kwargs.get("track_state"),
track_chunk_idx=kwargs.get("track_chunk_idx"),
)
return (
python/sglang/srt/layers/attention/linear/kernels/kda_helion.py
+3
View File
@@ -20,6 +20,7 @@ class HelionKDAKernel(LinearAttnKernelBase):
"""
supports_packed_decode = True
supports_track_state_snapshot: bool = True
def __init__(
self,
@@ -188,4 +189,6 @@ class HelionKDAKernel(LinearAttnKernelBase):
dt_bias=dt_bias,
lower_bound=lower_bound,
output_intermediate_states=return_intermediate_states,
track_state=kwargs.get("track_state"),
track_chunk_idx=kwargs.get("track_chunk_idx"),
)
python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py
+4
View File
@@ -76,6 +76,10 @@ def _from_nvidia_kda_state_layout(
class NvidiaKDAKernel(LinearAttnKernelBase):
# Tracked batches route to the embedded Triton fallback, which forwards
# the fp32 snapshot arguments (see _triton_extend).
supports_track_state_snapshot: bool = True
def __init__(self):
# This kernel uses tcgen05 + TMEM, which are available on datacenter
# Blackwell (SM100/SM103, reported as capability major 10), but not on
python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py
+16 -3
View File
@@ -9,9 +9,11 @@ serving shape: K = V = 128, chunk 64.
Scope: ordinary extend batches satisfying the kernel's fixed tensor contract.
Correctness-sensitive cases stay on Triton:
- track batches receive dense intermediate SSM states directly from the kernel
when the cache checkpoint stride is also 64 tokens. Other interior snapshots
stay on Triton; boundary-only tracking can still use the final state;
- track batches carrying the fp32 snapshot buffer (``track_state``, the mamba
extra_buffer track path) stay on Triton — the kernel cannot write it.
Interior snapshots consumed as dense ``h`` still come from the kernel when
the cache checkpoint stride is also 64 tokens; boundary-only tracking uses
the final state either way;
- spec-decode extends, which must stay rollback-able.
Single-sequence token counts that are not a multiple of the kernel's 64-token
@@ -49,6 +51,12 @@ _PAD_GATE = -1000.0
class PtxKDAKernel(LinearAttnKernelBase):
# Batches carrying the fp32 track snapshot buffer (track_state) route to
# the embedded Triton fallback, which forwards the snapshot arguments
# (the track_state check in extend -> _triton_extend); boundary-only
# tracking stays native.
supports_track_state_snapshot: bool = True
def __init__(self):
# tcgen05 + TMEM with sm_103a-only encodings: GB300 (SM103) only.
self.supports_prefill = torch.cuda.is_available() and (
@@ -217,6 +225,11 @@ class PtxKDAKernel(LinearAttnKernelBase):
)
eligible = (
not kwargs.get("is_spec_decode")
# The native kernel cannot write the fp32 track snapshot buffer;
# a batch carrying one must take the Triton fallback, which
# forwards the snapshot arguments (see _triton_extend). Leaving
# the buffer unwritten would corrupt prefix-cache track slots.
and kwargs.get("track_state") is None
and intermediate_stride_supported
and shape_known
and supported_shape
python/sglang/srt/layers/attention/linear/kernels/kda_triton.py
+3
View File
@@ -28,6 +28,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
# the same fallback CPU/NPU use. Batched decode is handled via query_start_loc.
supports_packed_decode: bool = not is_cpu() and not is_npu() and not is_xpu()
supports_fused_chain_verify: bool = not is_cpu() and not is_npu()
supports_track_state_snapshot: bool = True
def packed_decode(
self,
@@ -248,4 +249,6 @@ class TritonKDAKernel(LinearAttnKernelBase):
lower_bound=lower_bound,
beta_is_raw=beta_is_raw,
output_intermediate_states=return_intermediate_states,
track_state=kwargs.get("track_state"),
track_chunk_idx=kwargs.get("track_chunk_idx"),
)
python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py
+8
View File
@@ -13,6 +13,14 @@ class LinearAttnKernelBase(ABC):
uses_state_checkpoints: bool = False
supports_fused_chain_verify: bool = False
# True when extend() honors the fp32 track snapshot (track_state /
# track_chunk_idx), natively or by routing tracked batches to a kernel
# that does. KDAAttnBackend asserts this before allocating the snapshot
# buffer: a kernel that silently ignores those arguments leaves the buffer
# unwritten and corrupts prefix-cache restores. Kernels that reject
# tracked batches loudly (NotImplementedError) keep the default False.
supports_track_state_snapshot: bool = False
@abstractmethod
def decode(
self,
python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
+5
View File
@@ -55,6 +55,11 @@ class ForwardMetadata:
track_ssm_h_dst: Optional[torch.Tensor] = None
track_ssm_final_src: Optional[torch.Tensor] = None
track_ssm_final_dst: Optional[torch.Tensor] = None
track_chunk_idx: Optional[torch.Tensor] = None
# Batch rows of the chunk-unaligned tracked seqs; indexes the fp32
# h_track_buf snapshot (KDA path) with plain integer indexing, so the
# copy into the track slots does not nonzero()-sync the stream.
track_ssm_h_batch_src: Optional[torch.Tensor] = None
state_checkpoint_cu_starts: Optional[torch.Tensor] = None
num_state_checkpoints: int = 0
state_checkpoint_every_n_tokens: int = 0