[GDN] Auto-select FlashInfer GDN prefill on validated SM100 configs (#29734)

python/sglang/srt/layers/attention/attention_registry.py

@@ -274,7 +274,10 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
                 HybridLinearAttnBackend,
                 Mamba2AttnBackend,
             )
-            from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend
+            from sglang.srt.layers.attention.linear.gdn_backend import (
+                GDNAttnBackend,
+                maybe_set_default_flashinfer_gdn_prefill,
+            )
         else:
             from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import (
                 AscendGDNAttnBackend as GDNAttnBackend,
@@ -287,6 +290,8 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
             )
 
         check_environments()
+        if runner.hybrid_gdn_config is not None and not is_npu():
+            maybe_set_default_flashinfer_gdn_prefill(runner)
         initialize_linear_attn_config(runner.server_args)
         hybrid_backend_cls = HybridLinearAttnBackend
         if runner.hybrid_gdn_config is not None:

python/sglang/srt/layers/attention/linear/gdn_backend.py

@@ -55,6 +55,47 @@ elif is_cpu():
     fused_gdn_gating = torch.ops.sgl_kernel.fused_gdn_gating_cpu
 
 
+def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None:
+    """Use FlashInfer for the narrow SM100 GDN prefill domain we validated."""
+    args = model_runner.server_args
+    if (
+        args.linear_attn_prefill_backend is not None
+        or args.linear_attn_backend != "triton"
+        or args.enable_page_major_kv_layout
+        or not is_cuda()
+        or torch.cuda.get_device_capability()[0] != 10
+    ):
+        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 = model_runner.hybrid_gdn_config
+    if (
+        cuda_version is None
+        or int(cuda_version.split(".", 1)[0]) < 13
+        or args.enable_dynamic_chunking
+        or chunk_size is None
+        or not 1 <= chunk_size <= 8192
+        or getattr(config, "linear_key_head_dim", None) != 128
+        or getattr(config, "linear_value_head_dim", None) != 128
+        or model_runner.req_to_token_pool.mamba_pool.mamba_cache.temporal.dtype
+        != torch.bfloat16
+    ):
+        return
+
+    from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+        is_flashinfer_gdn_prefill_available,
+    )
+
+    if is_flashinfer_gdn_prefill_available():
+        args.linear_attn_prefill_backend = "flashinfer"
+        rank0_log("Defaulting SM100 GDN prefill backend to FlashInfer.")
+
+
 class GDNKernelDispatcher:
     """Dispatches GDN kernel calls to the appropriate backend per mode."""
 
@@ -64,6 +105,7 @@ class GDNKernelDispatcher:
         prefill_backend: LinearAttnKernelBackend,
     ):
         triton_kernel = TritonGDNKernel()
+        self.tree_verify_kernel = triton_kernel
 
         cutedsl_kernel = None
         if decode_backend.is_triton():
@@ -251,7 +293,15 @@ class GDNKernelDispatcher:
         query_start_loc: torch.Tensor,
         **kwargs,
     ) -> torch.Tensor:
-        return self.verify_kernel.target_verify(
+        # FlashInfer verify supports a linear MTP chain. Tree-shaped drafts
+        # carry parent indices and must use Triton even when decode/prefill use
+        # FlashInfer.
+        verify_kernel = (
+            self.tree_verify_kernel
+            if kwargs.get("retrieve_parent_token") is not None
+            else self.verify_kernel
+        )
+        return verify_kernel.target_verify(
             A_log=A_log,
             dt_bias=dt_bias,
             q=q,

python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py

@@ -72,6 +72,12 @@ def _get_flashinfer_gdn_kernels():
     )
 
 
+def is_flashinfer_gdn_prefill_available() -> bool:
+    """Return whether the kernel loader can construct the prefill path."""
+    available, prefill_fn, *_ = _get_flashinfer_gdn_kernels()
+    return bool(available and prefill_fn is not None)
+
+
 # ---------------------------------------------------------------------------
 # Kernel implementation
 # ---------------------------------------------------------------------------

python/sglang/srt/server_args.py

@@ -2011,7 +2011,7 @@ class ServerArgs:
     linear_attn_prefill_backend: A[
         Optional[str],
         Arg(
-            help="Override the kernel backend for linear attention prefill/extend. If not set, uses --linear-attn-backend.",
+            help="Override the kernel backend for linear attention prefill/extend. If not set, uses --linear-attn-backend; compatible SM100 GDN models may automatically select FlashInfer.",
             choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
         ),
     ] = None