[NVIDIA] [GDN] Add FlashInfer prefill support for SM100+ (Blackwell) (#22921)

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

@@ -3,7 +3,7 @@
 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+): decode-only with bf16 state.  More support on the way.
+SM100+ (Blackwell+): decode and prefill with bf16 state.  MTP verify on the way.
 
 Requires flashinfer >= 0.6.4 (SM90) or >= 0.6.5 (SM100+).
 """
@@ -74,8 +74,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
     """FlashInfer kernel for GDN with K-last SSM state layout.
 
     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).
+    SM100+ (Blackwell+): decode and prefill supported; MTP verify not yet supported.
 
     Requires flashinfer >= 0.6.4 (SM90) or >= 0.6.5 (SM100+).
     """
@@ -97,7 +96,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
             raise RuntimeError("FlashInfer GDN decode kernel is unavailable.")
 
         sm_major = torch.cuda.get_device_capability()[0]
-        self.use_state_pool = sm_major != 9
+        self.use_state_pool = sm_major >= 10
         self.supports_target_verify = sm_major == 9
 
         if sm_major == 9:
@@ -187,13 +186,6 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
         query_start_loc: torch.Tensor,
         **kwargs,
     ) -> tuple:
-        if self.use_state_pool:
-            raise NotImplementedError(
-                "FlashInfer GDN prefill is not supported on SM100+. "
-                "Use --linear-attn-prefill-backend triton."
-            )
-
-        # SM90: chunked prefill using FlashInfer GDN prefill kernel.
         from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
 
         total_seq_len = q.shape[1]
@@ -208,30 +200,50 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
         alpha_fi = torch.exp(g[0].to(torch.float32))
         beta_fi = beta[0].to(torch.float32)
 
-        cu_seqlens_fi = query_start_loc.to(torch.int64)
-
-        # Remap negative padding indices to sentinel slot
-        ssm_cache_indices = torch.where(
-            cache_indices >= 0,
-            cache_indices,
-            ssm_states.shape[0] - 1,
-        ).to(torch.int64)
-
-        # FlashInfer requires float32 initial state, K-last layout [B, HV, V, K]
-        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=cu_seqlens_fi,
-            use_qk_l2norm_in_kernel=False,
-        )
+        if self.use_state_pool:
+            # Negative indices (e.g. -1) are padding markers for slots not yet
+            # assigned to a real sequence; clamp them to 0 (the reserved dummy
+            # 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,
+            )
+        else:
+            # SM90: preserve original negative-index handling (remap to last slot).
+            ssm_cache_indices = torch.where(
+                cache_indices >= 0,
+                cache_indices,
+                ssm_states.shape[0] - 1,
+            ).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,
+            )
 
         # Write back state to pool
         ssm_states.index_copy_(

python/sglang/srt/server_args.py

@@ -3099,6 +3099,22 @@ class ServerArgs:
                 f"got {self.mamba_ssm_dtype!r}"
             )
 
+        # SM100+ FlashInfer GDN prefill requires CUDA 13+ (CuTe DSL kernel)
+        # for correctness and best performance.
+        prefill = self.linear_attn_prefill_backend or self.linear_attn_backend
+        cuda_version = torch.version.cuda
+        cuda_major = int(cuda_version.split(".")[0]) if cuda_version is not None else 0
+        if (
+            prefill == "flashinfer"
+            and torch.cuda.is_available()
+            and torch.cuda.get_device_capability()[0] >= 10
+            and cuda_major < 13
+        ):
+            raise ValueError(
+                "--linear-attn-prefill-backend flashinfer on SM100+ requires CUDA 13+, "
+                f"got CUDA {cuda_version or 'unknown'}"
+            )
+
     def _handle_context_parallelism(self):
         if (
             self.enable_prefill_context_parallel