chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NVIDIA] Integrate FlashInfer decode kernel (Blackwell) for Qwen3.5 (#19150)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Kaixi Hou
2026-03-17 22:11:18 -07:00
committed by GitHub
parent da24478913
commit 4cc19862ef
5 changed files with 161 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/layers/attention/linear/gdn_backend.py
View File

@@ -72,7 +72,7 @@ class GDNKernelDispatcher:
self.decode_kernel = CuteDSLGDNKernel()
elif decode_backend.is_flashinfer():
if not is_cuda():
raise ValueError("FlashInfer backend requires CUDA")
raise ValueError("FlashInfer GDN backend requires CUDA")
from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
FlashInferGDNKernel,
)
@@ -91,7 +91,7 @@ class GDNKernelDispatcher:
)
elif prefill_backend.is_flashinfer():
if not is_cuda():
raise ValueError("FlashInfer backend requires CUDA")
raise ValueError("FlashInfer GDN backend requires CUDA")
# Reuse the FlashInfer kernel if already created for decode
if decode_backend.is_flashinfer():
self.extend_kernel = flashinfer_kernel
@@ -399,6 +399,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
cache_indices=cache_indices,
query_start_loc=query_start_loc,
)
if (is_npu() or is_cpu()) and last_recurrent_state is not None:
last_recurrent_state = last_recurrent_state.to(
ssm_states.dtype, copy=False

137
python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py
View File

@@ -1,16 +1,11 @@
"""FlashInfer-based kernels for GDN (Gated Delta Network) linear attention.
Provides K-last SSM layout support using FlashInfer CUTLASS kernels (SM90+).
The K-last layout stores SSM states as [pool, HV, V, K] instead of V-last
[pool, HV, K, V], enabling more efficient memory access patterns for decode.
Both SM90 and SM100+ use the same pool layout: [pool, HV, V, K] (K-last).
Requires ``flashinfer`` with GDN kernel support to be installed, e.g.
``pip install -e ".[cutlass]"`` from the FlashInfer repo.
SM90 (Hopper): full support — decode, prefill, MTP. State dtype: fp32.
SM100+ (Blackwell+): decode-only with bf16 state. More support on the way.
NOTE: FlashInfer >= 0.6.4 includes a fix (PR#2509) that caches
cudaGetDeviceProperties in the GDN prefill JIT launcher, eliminating ~80ms
of CPU overhead per prefill. Upgrading from 0.6.3 to 0.6.4 recovers ~50%
prefill throughput regression observed with stock FlashInfer.
Requires flashinfer >= 0.6.4 (SM90) or >= 0.6.5 (SM100+).
"""
import logging
@@ -52,17 +47,12 @@ def _get_flashinfer_gdn_kernels():
_flashinfer_chunk_gated_delta_rule = chunk_gated_delta_rule
_flashinfer_gated_delta_rule_mtp = gated_delta_rule_mtp
# Use pretranspose (K-last / V-major) decode kernel to match
# the K-last pool layout [pool, HV, V, K]
_flashinfer_gated_delta_rule_decode = gated_delta_rule_decode_pretranspose
# SM90+ required for FlashInfer GDN CUTLASS kernels
_flashinfer_gdn_available = (
torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
)
if _flashinfer_gdn_available:
logger.info(
"FlashInfer GDN kernels (prefill + decode + MTP) loaded successfully"
)
logger.info("FlashInfer GDN kernels loaded successfully")
except (ImportError, RuntimeError) as e:
logger.warning(f"FlashInfer GDN kernels not available: {e}")
_flashinfer_gdn_available = False
@@ -81,10 +71,13 @@ def _get_flashinfer_gdn_kernels():
class FlashInferGDNKernel(LinearAttnKernelBase):
"""FlashInfer CUTLASS kernel for GDN with K-last SSM state layout.
"""FlashInfer kernel for GDN with K-last SSM state layout.
Supports decode (pooled pretranspose), extend (chunked prefill) and
target_verify (MTP). Requires SM90+ and FlashInfer with GDN support.
SM90 (Hopper): decode uses gather/scatter; prefill and MTP verify supported.
SM100+ (Blackwell+): decode uses pool API (initial_state_indices); prefill
and MTP verify are not supported (use Triton backend for those).
Requires flashinfer >= 0.6.4 (SM90) or >= 0.6.5 (SM100+).
"""
def __init__(self):
@@ -100,16 +93,19 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
"FlashInfer GDN kernels are not available. "
"Requires SM90+ and FlashInfer with GDN kernel support."
)
if self._prefill_fn is None:
raise RuntimeError("FlashInfer GDN prefill kernel is unavailable.")
if self._mtp_fn is None:
raise RuntimeError("FlashInfer GDN MTP (verify) kernel is unavailable.")
if self._decode_fn is None:
raise RuntimeError("FlashInfer GDN decode kernel is unavailable.")
logger.info(
"K-last mode: Using FlashInfer GDN prefill, decode and MTP (verify) kernels"
)
sm_major = torch.cuda.get_device_capability()[0]
self.use_state_pool = sm_major != 9
if sm_major == 9:
if self._prefill_fn is None:
raise RuntimeError("FlashInfer GDN prefill kernel is unavailable.")
if self._mtp_fn is None:
raise RuntimeError("FlashInfer GDN MTP (verify) kernel is unavailable.")
logger.info("Using FlashInfer GDN kernels")
# ---- decode ----
@@ -128,13 +124,6 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
"""K-last decode using FlashInfer pretranspose kernel (stock, no pool indexing).
TODO: Once FlashInfer PR#2521 is merged and released, switch back
to pool-indexed decode (passing state_indices directly) to avoid
the gather/scatter overhead (~7-9% decode regression).
https://github.com/flashinfer-ai/flashinfer/pull/2521
"""
batch_size = cache_indices.shape[0]
num_heads = q.shape[2]
head_k_dim = q.shape[3]
@@ -147,27 +136,39 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
a_fi = a.view(batch_size, 1, num_v_heads)
b_fi = b.view(batch_size, 1, num_v_heads)
# Gather states from pool
state_batch = ssm_states[cache_indices]
if self.use_state_pool:
output_fi, _ = self._decode_fn(
q=query_fi,
k=key_fi,
v=value_fi,
state=None,
A_log=A_log.detach().float(),
a=a_fi,
dt_bias=dt_bias.detach(),
b=b_fi,
use_qk_l2norm=True,
initial_state=ssm_states,
initial_state_indices=cache_indices,
)
else:
# TODO: Once FlashInfer PR#2521 is merged for SM90, gather/scatter
# will no longer be needed here.
state_batch = ssm_states[cache_indices]
output_fi, new_state = self._decode_fn(
q=query_fi,
k=key_fi,
v=value_fi,
state=state_batch,
A_log=A_log.detach(),
a=a_fi,
dt_bias=dt_bias.detach(),
b=b_fi,
scale=None,
output=None,
use_qk_l2norm=True,
)
ssm_states[cache_indices] = new_state
output_fi, new_state = self._decode_fn(
q=query_fi,
k=key_fi,
v=value_fi,
state=state_batch,
A_log=A_log.detach(),
a=a_fi,
dt_bias=dt_bias.detach(),
b=b_fi,
scale=None,
output=None,
use_qk_l2norm=True,
)
# Scatter updated states back to pool
ssm_states[cache_indices] = new_state
# [bs, 1, HV, V] -> [1, bs, HV, V]
return output_fi.view(1, batch_size, num_v_heads, head_v_dim)
# ---- extend (prefill) ----
@@ -185,16 +186,15 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
query_start_loc: torch.Tensor,
**kwargs,
) -> tuple:
"""K-last chunked prefill using FlashInfer GDN prefill kernel.
if self.use_state_pool:
raise NotImplementedError(
"FlashInfer GDN prefill is not supported on SM100+. "
"Use --linear-attn-prefill-backend triton."
)
The FlashInfer kernel natively supports K-last state layout [N, H, V, K].
q and k are L2-normalized before calling the kernel (the kernel is called
with ``use_qk_l2norm_in_kernel=False``).
"""
# SM90: chunked prefill using FlashInfer GDN prefill kernel.
from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
# q, k: [1, seq, H, K] -> [seq, H, K]
# v: [1, seq, HV, V] -> [seq, HV, V]
total_seq_len = q.shape[1]
num_v_heads = v.shape[2]
head_v_dim = v.shape[3]
@@ -243,8 +243,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
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
# (prefix caching for K-last prefill is not supported).
# h=None since FlashInfer doesn't provide intermediate states.
return core_attn_out, None, None
# ---- target_verify (MTP) ----
@@ -268,18 +267,18 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
retrieve_parent_token: torch.Tensor,
**kwargs,
) -> torch.Tensor:
"""K-last MTP verify using FlashInfer GDN MTP kernel.
if self.use_state_pool:
raise NotImplementedError(
"FlashInfer GDN MTP verify is not yet supported on SM100+."
)
Only supports topk=1 (retrieve_parent_token must be None).
"""
# SM90: MTP verify using FlashInfer gated_delta_rule_mtp kernel.
if retrieve_parent_token is not None:
raise RuntimeError(
"FlashInfer GDN verify kernel only supports topk=1 "
"(retrieve_parent_token must be None)."
)
# Recover batch_size and draft_token_num from g shape
# g: [1, seq_len, HV] where seq_len = batch_size * draft_token_num
seq_len = q.shape[1]
batch_size = query_start_loc.shape[0] - 1
draft_token_num = seq_len // batch_size
@@ -289,16 +288,13 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
num_v_heads = v.shape[2]
head_v_dim = v.shape[3]
# Reshape [1, seq, H, D] -> [B, T, H, D] for FlashInfer MTP
query_mtp = q.view(batch_size, draft_token_num, num_heads, head_k_dim)
key_mtp = k.view(batch_size, draft_token_num, num_heads, head_k_dim)
value_mtp = v.view(batch_size, draft_token_num, num_v_heads, head_v_dim)
# a, b from g/beta: [1, seq, HV] -> [B, T, HV]
if a is None or b is None or A_log is None or dt_bias is None:
raise RuntimeError(
"FlashInfer GDN MTP kernel requires a_raw, b_raw, A_log, "
"dt_bias to be passed via kwargs."
"FlashInfer GDN MTP kernel requires a, b, A_log, dt_bias."
)
a_mtp = a.view(batch_size, draft_token_num, num_v_heads)
@@ -321,5 +317,4 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
use_qk_l2norm=True,
)
# [B, T, HV, V] -> [1, seq_len, HV, V]
return output_fi.view(1, seq_len, num_v_heads, head_v_dim)

31
python/sglang/srt/server_args.py
View File

@@ -771,6 +771,7 @@ class ServerArgs:
self._handle_sampling_backend()
self._handle_attention_backend_compatibility()
self._handle_mamba_backend()
self._handle_linear_attn_backend()
self._handle_kv4_compatibility()
self._handle_page_size()
self._handle_amd_specifics()
@@ -2036,6 +2037,19 @@ class ServerArgs:
not self.enable_mamba_extra_buffer()
), f"mamba extra_buffer is not supported for {model_arch} model"
# FlashInfer GDN decode is incompatible with no_buffer scheduling.
# See https://github.com/sgl-project/sglang/issues/20791
if (
self.linear_attn_decode_backend == "flashinfer"
and self.mamba_scheduler_strategy == "no_buffer"
):
raise ValueError(
"FlashInfer GDN decode (--linear-attn-decode-backend flashinfer) is not "
"compatible with --mamba-scheduler-strategy no_buffer. "
"Please use --mamba-scheduler-strategy extra_buffer instead. "
"See https://github.com/sgl-project/sglang/issues/20791"
)
if self.enable_mamba_extra_buffer(): # extra_buffer
if self.disable_radix_cache:
raise ValueError(
@@ -2444,6 +2458,23 @@ class ServerArgs:
"FlashInfer mamba module not available, please check flashinfer installation."
)
def _handle_linear_attn_backend(self):
# SM100+ FlashInfer GDN decode requires bf16 state; SM90 uses float32.
import torch
decode = self.linear_attn_decode_backend or self.linear_attn_backend
if (
decode == "flashinfer"
and self.mamba_ssm_dtype != "bfloat16"
and torch.cuda.is_available()
and torch.cuda.get_device_capability()[0] >= 10
):
raise ValueError(
"--linear-attn-decode-backend flashinfer on SM100+ requires "
"--mamba-ssm-dtype bfloat16, "
f"got {self.mamba_ssm_dtype!r}"
)
def _handle_context_parallelism(self):
if self.attn_cp_size > 1:
# The tp_size is the world size, not the real tensor parallel size

8
python/sglang/test/accuracy_test_runner.py
View File

@@ -27,6 +27,7 @@ class AccuracyTestParams:
thinking_mode: Optional[str] = None # e.g., "deepseek-v3"
temperature: Optional[float] = None
top_p: Optional[float] = None
top_k: Optional[int] = None
repeat: Optional[int] = None
@@ -83,6 +84,7 @@ def _run_simple_eval(
thinking_mode: Optional[str] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
repeat: Optional[int] = None,
) -> Tuple[bool, Optional[str], Optional[dict]]:
"""Run evaluation using simple_eval backend (run_eval.py).
@@ -123,6 +125,9 @@ def _run_simple_eval(
if top_p is not None:
args.top_p = top_p
if top_k is not None:
args.top_k = top_k
if repeat is not None:
args.repeat = repeat
@@ -223,7 +228,7 @@ def run_accuracy_test(
# Use simple_eval when any extended params are set that few_shot_eval doesn't support.
has_extended_params = any(
getattr(params, field) is not None
for field in ("thinking_mode", "temperature", "top_p", "repeat")
for field in ("thinking_mode", "temperature", "top_p", "top_k", "repeat")
)
if params.dataset == "gsm8k" and not has_extended_params:
success, error, metrics = _run_few_shot_eval(
@@ -244,6 +249,7 @@ def run_accuracy_test(
thinking_mode=params.thinking_mode,
temperature=params.temperature,
top_p=params.top_p,
top_k=params.top_k,
repeat=params.repeat,
)