feat(gdn): add FlashInfer K-last SSM layout support for GDN prefill and decode for Hopper

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

@@ -68,6 +68,15 @@ def __init__(
             )
 
             self.decode_kernel = CuteDSLGDNKernel()
+        elif decode_backend.is_flashinfer():
+            if not is_cuda():
+                raise ValueError("FlashInfer backend requires CUDA")
+            from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+                FlashInferGDNKernel,
+            )
+
+            flashinfer_kernel = FlashInferGDNKernel()
+            self.decode_kernel = flashinfer_kernel
         else:
             raise ValueError(f"Unsupported GDN decode backend: {decode_backend}")
 
@@ -78,10 +87,27 @@ def __init__(
                 "CuTe DSL backend only supports decode, not prefill. "
                 "Use --linear-attn-prefill-backend triton instead."
             )
+        elif prefill_backend.is_flashinfer():
+            if not is_cuda():
+                raise ValueError("FlashInfer backend requires CUDA")
+            # Reuse the FlashInfer kernel if already created for decode
+            if decode_backend.is_flashinfer():
+                self.extend_kernel = flashinfer_kernel
+            else:
+                from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+                    FlashInferGDNKernel,
+                )
+
+                flashinfer_kernel = FlashInferGDNKernel()
+                self.extend_kernel = flashinfer_kernel
         else:
             raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}")
 
-        self.verify_kernel = triton_kernel
+        # Verify kernel: use FlashInfer if either decode or prefill selected it
+        if decode_backend.is_flashinfer() or prefill_backend.is_flashinfer():
+            self.verify_kernel = flashinfer_kernel
+        else:
+            self.verify_kernel = triton_kernel
 
         rank0_log(
             f"GDN kernel dispatcher: decode={self.decode_kernel.__class__.__name__}, "
@@ -354,6 +380,11 @@ def forward_extend(
                 intermediate_state_indices=intermediate_state_indices,
                 cache_steps=forward_batch.spec_info.draft_token_num,
                 retrieve_parent_token=retrieve_parent_token,
+                # Pass raw pre-gating values for FlashInfer MTP kernel
+                a_raw=a,
+                b_raw=b,
+                A_log=layer.A_log,
+                dt_bias=layer.dt_bias,
             )
         else:
             core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
@@ -366,14 +397,15 @@ def forward_extend(
                 cache_indices=cache_indices,
                 query_start_loc=query_start_loc,
             )
-            if is_npu() or is_cpu():
+            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
                 )
                 ssm_states[cache_indices] = last_recurrent_state
 
-            self._track_mamba_state_extend(
-                forward_batch, h, ssm_states, forward_metadata
-            )
+            if h is not None:
+                self._track_mamba_state_extend(
+                    forward_batch, h, ssm_states, forward_metadata
+                )
 
         return core_attn_out