From 4cc19862efded3a4aff44aafa8dcd87be8cff5c2 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Tue, 17 Mar 2026 22:11:18 -0700 Subject: [PATCH] [NVIDIA] Integrate FlashInfer decode kernel (Blackwell) for Qwen3.5 (#19150) Co-authored-by: Claude Sonnet 4.6 --- .../layers/attention/linear/gdn_backend.py | 5 +- .../linear/kernels/gdn_flashinfer.py | 137 +++++++++--------- python/sglang/srt/server_args.py | 31 ++++ python/sglang/test/accuracy_test_runner.py | 8 +- .../4-gpu-models/test_qwen35_models.py | 108 +++++++------- 5 files changed, 161 insertions(+), 128 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index d15ad8fe5..1e3cefa47 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index 32ffd9664..f0bf4a04a 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e39f4810b..590de8594 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/python/sglang/test/accuracy_test_runner.py b/python/sglang/test/accuracy_test_runner.py index ebdd0a106..0cf007220 100644 --- a/python/sglang/test/accuracy_test_runner.py +++ b/python/sglang/test/accuracy_test_runner.py @@ -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, ) diff --git a/test/registered/4-gpu-models/test_qwen35_models.py b/test/registered/4-gpu-models/test_qwen35_models.py index d67a298b7..c457361e6 100644 --- a/test/registered/4-gpu-models/test_qwen35_models.py +++ b/test/registered/4-gpu-models/test_qwen35_models.py @@ -7,15 +7,18 @@ import requests from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree +from sglang.test.accuracy_test_runner import AccuracyTestParams from sglang.test.ci.ci_register import register_cuda_ci # This eval harness applies the chat_template, which is critical for qwen3.5 # to get good accuracy on gsm8k +from sglang.test.run_combined_tests import run_combined_tests from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, + ModelLaunchSettings, popen_launch_server, ) @@ -26,62 +29,59 @@ QWEN35_27B_MODEL = "Qwen/Qwen3.5-27B" ACC_THRESHOLDS = {QWEN35_FP4_MODEL: {"gsm8k": 0.95}, QWEN35_27B_MODEL: {"gsm8k": 0.8}} -class TestQwen35FP4(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = QWEN35_FP4_MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--tp-size", - "4", - "--chunked-prefill-size", - "2048", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - "128", - "--mamba-ssm-dtype", - "bfloat16", - "--max-running-requests", - "128", - "--reasoning-parser", - "qwen3", - "--attention-backend", - "trtllm_mha", - "--quantization", - "modelopt_fp4", - "--model-loader-extra-config", - '{"enable_multithread_load": true,"num_threads": 64}', - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - +class TestQwen35FP4(unittest.TestCase): def test_gsm8k(self): - args = SimpleNamespace( - model=self.model, - eval_name="gsm8k", - num_shots=5, - num_examples=200, - max_tokens=16000, - num_threads=128, - repeat=1, - temperature=0.6, - top_p=0.95, - top_k=20, - base_url=self.base_url, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), + base_args = [ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + "--mamba-ssm-dtype", + "bfloat16", + "--max-running-requests", + "128", + "--reasoning-parser", + "qwen3", + "--attention-backend", + "trtllm_mha", + "--quantization", + "modelopt_fp4", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', + ] + + variants = [ + ModelLaunchSettings( + QWEN35_FP4_MODEL, + extra_args=base_args, + variant="Triton", + ), + ModelLaunchSettings( + QWEN35_FP4_MODEL, + extra_args=base_args + ["--linear-attn-decode-backend", "flashinfer"], + variant="FlashInfer", + ), + ] + + run_combined_tests( + models=variants, + test_name="Qwen3.5-397B-A17B-NVFP4", + accuracy_params=AccuracyTestParams( + dataset="gsm8k", + baseline_accuracy=ACC_THRESHOLDS[QWEN35_FP4_MODEL]["gsm8k"], + num_examples=200, + num_threads=128, + max_tokens=16000, + thinking_mode="qwen3", + temperature=0.6, + top_p=0.95, + top_k=20, + ), ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]) class TestQwen35FP4MTP(CustomTestCase):