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

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

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

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+).
+Both SM90 and SM100+ use the same pool layout: [pool, HV, V, K] (K-last).
-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.
 
-Requires ``flashinfer`` with GDN kernel support to be installed, e.g.
+SM90 (Hopper): full support — decode, prefill, MTP.  State dtype: fp32.
-``pip install -e ".[cutlass]"`` from the FlashInfer repo.
+SM100+ (Blackwell+): decode-only with bf16 state.  More support on the way.
 
-NOTE: FlashInfer >= 0.6.4 includes a fix (PR#2509) that caches
+Requires flashinfer >= 0.6.4 (SM90) or >= 0.6.5 (SM100+).
-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.
 """
 
 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(
+                logger.info("FlashInfer GDN kernels loaded successfully")
-                    "FlashInfer GDN kernels (prefill + decode + MTP) 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
+    SM90 (Hopper): decode uses gather/scatter; prefill and MTP verify supported.
-    target_verify (MTP).  Requires SM90+ and FlashInfer with GDN support.
+    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(
+        sm_major = torch.cuda.get_device_capability()[0]
-            "K-last mode: Using FlashInfer GDN prefill, decode and MTP (verify) kernels"
+        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
+        if self.use_state_pool:
-        state_batch = ssm_states[cache_indices]
+            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].
+        # SM90: chunked prefill using FlashInfer GDN prefill kernel.
-        q and k are L2-normalized before calling the kernel (the kernel is called
-        with ``use_qk_l2norm_in_kernel=False``).
-        """
         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
+        # h=None since FlashInfer doesn't provide intermediate states.
-        # (prefix caching for K-last prefill is not supported).
         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, "
+                "FlashInfer GDN MTP kernel requires a, b, A_log, dt_bias."
-                "dt_bias to be passed via kwargs."
             )
 
         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)

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

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,
         )
 

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):
+class TestQwen35FP4(unittest.TestCase):
-    @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)
-
     def test_gsm8k(self):
-        args = SimpleNamespace(
+        base_args = [
-            model=self.model,
+            "--tp-size",
-            eval_name="gsm8k",
+            "4",
-            num_shots=5,
+            "--chunked-prefill-size",
-            num_examples=200,
+            "2048",
-            max_tokens=16000,
+            "--mamba-scheduler-strategy",
-            num_threads=128,
+            "extra_buffer",
-            repeat=1,
+            "--mamba-track-interval",
-            temperature=0.6,
+            "128",
-            top_p=0.95,
+            "--mamba-ssm-dtype",
-            top_k=20,
+            "bfloat16",
-            base_url=self.base_url,
+            "--max-running-requests",
-            host="http://127.0.0.1",
+            "128",
-            port=int(self.base_url.split(":")[-1]),
+            "--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):