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

[GDN] Support FlashInfer GDN prefill with extra-buffer radix cache (#29735)

Browse Source
This commit is contained in:
YAMY
2026-07-30 00:47:35 +08:00
committed by GitHub
parent 50029f05a3
commit fddfc1fb5e
9 changed files with 282 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+5 -3
View File
@@ -724,16 +724,18 @@ class MambaAttnBackendBase(AttentionBackend):
def _track_mamba_state_extend(
self,
forward_batch: ForwardBatch,
h: torch.Tensor,
h: Optional[torch.Tensor],
ssm_states: torch.Tensor,
forward_metadata: ForwardMetadata,
):
"""Copy extend SSM state at the last chunk boundary to track slots (source
depends on chunk alignment; see `_init_track_ssm_indices`)."""
if forward_metadata.has_mamba_track_mask:
h = h.squeeze(0)
# 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)
python/sglang/srt/layers/attention/linear/gdn_backend.py
+20 -5
View File
@@ -75,10 +75,6 @@ def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None:
):
return
# Extra-buffer strategies need intermediate state checkpoints.
if args.uses_mamba_radix_cache and args.mamba_radix_cache_strategy != "no_buffer":
return
cuda_version = torch.version.cuda
chunk_size = args.chunked_prefill_size
config = hybrid_gdn_config(model_runner.model_config)
@@ -204,6 +200,10 @@ class GDNKernelDispatcher:
f"packed_decode={self.supports_packed_decode}"
)
@property
def extend_uses_state_checkpoints(self) -> bool:
return self.extend_kernel.uses_state_checkpoints
def packed_decode(
self,
mixed_qkv: torch.Tensor,
@@ -362,6 +362,14 @@ class GDNAttnBackend(MambaAttnBackendBase):
self.forward_metadata.mamba_track_mask_indices
]
)
if self.kernel_dispatcher.extend_uses_state_checkpoints:
from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
maybe_build_flashinfer_checkpoint_plan,
)
maybe_build_flashinfer_checkpoint_plan(
forward_batch, self.forward_metadata, self.device
)
def forward_decode(
self,
@@ -649,6 +657,13 @@ class GDNAttnBackend(MambaAttnBackendBase):
ssm_states=ssm_states_contig,
cache_indices=state_cache_indices,
query_start_loc=query_start_loc,
state_checkpoint_cu_starts=(
forward_metadata.state_checkpoint_cu_starts
),
num_state_checkpoints=forward_metadata.num_state_checkpoints,
state_checkpoint_every_n_tokens=(
forward_metadata.state_checkpoint_every_n_tokens
),
)
if is_npu() and last_recurrent_state is not None:
@@ -663,7 +678,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
conv_states[cache_indices] = conv_states_contig
ssm_states[cache_indices] = ssm_states_contig
if h is not None:
if forward_metadata.has_mamba_track_mask:
self._track_mamba_state_extend(
forward_batch, h, ssm_states, forward_metadata
)
python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py
+86 -34
View File
@@ -5,20 +5,27 @@ Both SM90 and SM100 use the same pool layout: [pool, HV, V, K] (K-last).
SM90 (Hopper): full support — decode, prefill, MTP. State dtype: fp32.
SM100 (Blackwell): full support — decode, prefill, MTP.
Requires flashinfer >= 0.6.7.
Requires flashinfer >= 0.6.14.
"""
from __future__ import annotations
import logging
import os
from typing import Optional
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
LinearAttnKernelBase,
)
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils import is_cuda
if TYPE_CHECKING:
from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
@@ -31,6 +38,47 @@ _flashinfer_gated_delta_rule_decode = None
_flashinfer_gated_delta_rule_mtp_bf16 = None
def maybe_build_flashinfer_checkpoint_plan(
forward_batch: ForwardBatch,
forward_metadata: ForwardMetadata,
device: str,
) -> None:
"""Populate packed FlashInfer checkpoint metadata when tracking requires it."""
if (
forward_metadata.track_ssm_h_src is None
or forward_metadata.track_ssm_h_src.numel() == 0
):
return
checkpoint_every_n_tokens = get_server_args().mamba_cache_chunk_size
extend_seq_lens = forward_batch.extend_seq_lens.to(device="cpu", dtype=torch.int64)
track_mask = forward_batch.mamba_track_mask.to(device="cpu", dtype=torch.bool)
relative_track_lens = forward_batch.mamba_track_seqlens.to(
device="cpu", dtype=torch.int64
) - forward_batch.extend_prefix_lens.to(device="cpu", dtype=torch.int64)
checkpoint_counts = extend_seq_lens // checkpoint_every_n_tokens
checkpoint_cu_starts = torch.zeros(checkpoint_counts.numel() + 1, dtype=torch.int64)
checkpoint_cu_starts[1:] = torch.cumsum(checkpoint_counts, dim=0)
use_checkpoint = track_mask & (relative_track_lens % checkpoint_every_n_tokens != 0)
track_checkpoint_src = checkpoint_cu_starts[:-1][use_checkpoint] + (
relative_track_lens[use_checkpoint] // checkpoint_every_n_tokens - 1
)
if track_checkpoint_src.numel() and track_checkpoint_src.min() < 0:
raise ValueError("Tracked GDN state precedes the first FlashInfer checkpoint.")
assert track_checkpoint_src.numel() == forward_metadata.track_ssm_h_dst.numel()
forward_metadata.track_ssm_h_src = track_checkpoint_src.to(
device, non_blocking=True
)
forward_metadata.state_checkpoint_cu_starts = checkpoint_cu_starts.to(
device, non_blocking=True
)
forward_metadata.num_state_checkpoints = int(checkpoint_cu_starts[-1])
forward_metadata.state_checkpoint_every_n_tokens = checkpoint_every_n_tokens
def _get_flashinfer_gdn_kernels():
"""Lazy import for FlashInfer GDN prefill, decode and verify (MTP) kernels.
@@ -89,9 +137,11 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
SM90 (Hopper): decode uses gather/scatter; prefill and MTP verify supported.
SM100 (Blackwell): decode uses gather/scatter; prefill and MTP verify supported.
Requires flashinfer >= 0.6.7.
Requires flashinfer >= 0.6.14.
"""
uses_state_checkpoints = True
def __init__(self):
(
available,
@@ -231,6 +281,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
state_checkpoint_cu_starts: Optional[torch.Tensor] = None,
num_state_checkpoints: int = 0,
state_checkpoint_every_n_tokens: int = 0,
**kwargs,
) -> tuple:
from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
@@ -253,23 +306,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
# slot) so the FlashInfer kernel never reads out-of-bounds state.
ssm_cache_indices = cache_indices.clamp(min=0).to(torch.int64)
initial_state_fi = ssm_states[ssm_cache_indices].contiguous()
# Pre-allocate bf16 output_state so the kernel compiles and writes the
# bf16 state path directly, avoiding a fp32 allocation and a subsequent
# fp32->bf16 conversion in the scatter step.
output_state_fi = torch.empty_like(initial_state_fi)
output_fi, output_state_fi = self._prefill_fn(
q=q_fi,
k=k_fi,
v=v_fi,
g=alpha_fi,
beta=beta_fi,
scale=None,
initial_state=initial_state_fi,
output_final_state=True,
cu_seqlens=query_start_loc, # already int32
use_qk_l2norm_in_kernel=False,
output_state=output_state_fi,
)
cu_seqlens = query_start_loc # already int32
else:
# SM90: preserve original negative-index handling (remap to last slot).
ssm_cache_indices = torch.where(
@@ -279,18 +316,33 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
).to(torch.int64)
# State must be float32; kernel requires int64 cu_seqlens.
initial_state_fi = ssm_states[ssm_cache_indices].to(torch.float32)
output_fi, output_state_fi = self._prefill_fn(
q=q_fi,
k=k_fi,
v=v_fi,
g=alpha_fi,
beta=beta_fi,
scale=None,
initial_state=initial_state_fi,
output_final_state=True,
cu_seqlens=query_start_loc.to(torch.int64),
use_qk_l2norm_in_kernel=False,
cu_seqlens = query_start_loc.to(torch.int64)
# Keep final state and checkpoints in the same kernel state dtype.
output_state_fi = torch.empty_like(initial_state_fi)
state_checkpoints = (
initial_state_fi.new_empty(
(num_state_checkpoints, *initial_state_fi.shape[1:])
)
if num_state_checkpoints > 0
else None
)
output_fi, output_state_fi = self._prefill_fn(
q=q_fi,
k=k_fi,
v=v_fi,
g=alpha_fi,
beta=beta_fi,
scale=None,
initial_state=initial_state_fi,
output_final_state=True,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=False,
output_state=output_state_fi,
state_checkpoints=state_checkpoints,
checkpoint_cu_starts=state_checkpoint_cu_starts,
checkpoint_every_n_tokens=state_checkpoint_every_n_tokens,
)
# Write back state to pool
ssm_states.index_copy_(
@@ -302,9 +354,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
# Output: [seq, HV, V] -> [1, seq, HV, V]
core_attn_out = output_fi.view(1, total_seq_len, num_v_heads, head_v_dim)
# Return (output, last_recurrent_state, h) to match Triton kernel interface.
# h=None since FlashInfer doesn't provide intermediate states.
return core_attn_out, None, None
# Match Triton's [1, checkpoints, H, V, K] intermediate-state layout.
h = state_checkpoints.unsqueeze(0) if state_checkpoints is not None else None
return core_attn_out, None, h
# ---- target_verify (MTP) ----
python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py
+2
View File
@@ -10,6 +10,8 @@ class LinearAttnKernelBase(ABC):
and provides decode/extend/target_verify methods with a unified interface.
"""
uses_state_checkpoints: bool = False
@abstractmethod
def decode(
self,
python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
+3
View File
@@ -55,6 +55,9 @@ 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
state_checkpoint_cu_starts: Optional[torch.Tensor] = None
num_state_checkpoints: int = 0
state_checkpoint_every_n_tokens: int = 0
is_target_verify: bool = False
draft_token_num: int = 1
python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py
+17 -2
View File
@@ -55,6 +55,7 @@ class GDNAttentionCase:
page_size: int
prefix_lens: tuple[int, ...]
extend_lens: tuple[int, ...] = ()
linear_attn_prefill_backend: str | None = None
@property
def batch_size(self) -> int:
@@ -245,7 +246,7 @@ class MockGDNModelRunner(ModelRunner):
enable_mis=False,
linear_attn_backend="triton",
linear_attn_decode_backend=None,
linear_attn_prefill_backend=None,
linear_attn_prefill_backend=case.linear_attn_prefill_backend,
max_running_requests=None,
revision=None,
speculative_algorithm=None,
@@ -268,10 +269,16 @@ class MockGDNModelRunner(ModelRunner):
state_size=head_k_dim,
conv_kernel=2,
)
temporal_state_dtype = (
dtype
if case.linear_attn_prefill_backend == "flashinfer"
and torch.cuda.get_device_capability()[0] >= 10
else torch.float32
)
cache_params = Mamba2CacheParams(
shape=cache_shape,
layers=[0],
dtype=Mamba2StateDType(conv=dtype, temporal=torch.float32),
dtype=Mamba2StateDType(conv=dtype, temporal=temporal_state_dtype),
)
self.req_to_token_pool = HybridReqToTokenPool(
size=pool_batch_size,
@@ -591,6 +598,14 @@ def build_gdn_attention_fixture(
initialize_linear_attn_config(runner.server_args)
linear_backend = GDNAttnBackend(runner)
if case.linear_attn_prefill_backend == "flashinfer":
from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
FlashInferGDNKernel,
)
testcase.assertIsInstance(
linear_backend.kernel_dispatcher.extend_kernel, FlashInferGDNKernel
)
backend = HybridLinearAttnBackend(full_backend, linear_backend, full_attn_layers=[])
actual_module = ProjectedGDNAttention(
num_k_heads=case.num_k_heads,