feat(kda): integrate FlashInfer prefill checkpoints

python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py

@@ -319,20 +319,15 @@ def _prepare_extend_gate_inputs(
         Optional[torch.Tensor],
         Optional[torch.Tensor],
     ]:
-        """Apply the Ascend prefill gate contract.
-
-        The checkpoint was validated with FP32 gate activation before
-        ``chunk_kda``. Keeping this platform override here leaves the shared
-        GPU model/backend paths unchanged.
-        """
+        """Activate raw gate and beta projections in FP32 for Ascend prefill."""
         preactivated_g = fused_kda_gate_npu(
             g.flatten(-2),
             layer.A_log,
             layer.head_k_dim,
             gate_bias=layer.dt_bias,
             lower_bound=layer.lower_bound,
         )
-        return preactivated_g, beta, None, None
+        return preactivated_g, beta.float().sigmoid(), None, None
 
     def _forward_target_verify(
         self,

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

@@ -34,6 +34,7 @@
 
     causal_conv1d_update = causal_conv1d_update_cpu
 
+from sglang.srt.model_executor.cuda_graph_config import Backend as CudaGraphBackend
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.runtime_context import (
@@ -45,6 +46,24 @@
 )
 
 
+def _validate_flashinfer_kda_prefill(
+    *,
+    chunk_size: int,
+    enable_two_batch_overlap: bool,
+    prefill_cuda_graph_backend: CudaGraphBackend,
+) -> None:
+    if chunk_size <= 0 or chunk_size % 32:
+        raise ValueError(
+            "FlashInfer KDA checkpoint interval must be positive and divisible by 32"
+        )
+    if enable_two_batch_overlap:
+        raise ValueError("FlashInfer KDA prefill does not support two-batch overlap")
+    if prefill_cuda_graph_backend == CudaGraphBackend.FULL:
+        raise ValueError(
+            "FlashInfer KDA prefill requires eager linear attention; use disabled or breakable prefill CUDA graphs"
+        )
+
+
 class KDAKernelDispatcher:
     """Dispatches KDA kernel calls to the appropriate backend per mode."""
 
@@ -55,6 +74,7 @@ def __init__(
         verify_backend: LinearAttnKernelBackend,
     ):
         self.verify_backend = verify_backend
+        self.prefill_backend = prefill_backend
         triton_kernel = TritonKDAKernel()
         self.triton_kernel = triton_kernel
         helion_kernel = None
@@ -86,7 +106,6 @@ def __init__(
             self.decode_kernel = CuteDSLKDAKernel()
         elif decode_backend.is_flashinfer():
             # FlashInfer recurrent_kda: SM100 decode + MTP (target_verify).
-            # Prefill stays on Triton / CuTe DSL (FlashInfer has no KDA chunk kernel).
             if not is_cuda():
                 raise ValueError("KDA FlashInfer backend requires CUDA")
             from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
@@ -145,6 +164,14 @@ def __init__(
             )
 
             self.extend_kernel = FlashKDAKernel()
+        elif prefill_backend.is_flashinfer():
+            if not is_cuda():
+                raise ValueError("KDA FlashInfer prefill backend requires CUDA")
+            from sglang.srt.layers.attention.linear.kernels.kda_flashinfer_prefill import (
+                FlashInferKDAPrefillKernel,
+            )
+
+            self.extend_kernel = FlashInferKDAPrefillKernel()
         elif prefill_backend.is_cutedsl():
             if not is_cuda():
                 raise ValueError("KDA CuTe DSL backend requires CUDA")
@@ -195,7 +222,7 @@ def __init__(
         else:
             raise ValueError(
                 f"Unsupported KDA prefill backend: {prefill_backend}. "
-                "KDA supports 'triton', 'helion', 'flashkda', 'cutedsl', "
+                "KDA supports 'triton', 'helion', 'flashinfer', 'flashkda', 'cutedsl', "
                 "'nvidia_kda', or 'ptx_kda' (cutedsl/nvidia_kda prefill need "
                 "SM100, ptx_kda SM100 or SM103)."
             )
@@ -332,10 +359,11 @@ def target_verify(
             **kwargs,
         )
 
-    def effective_extend_kernel(self, lower_bound: Optional[float]):
-        """The kernel ``extend`` will actually run: safe-gate models reroute
-        kernels without ``supports_safe_gate`` to Triton."""
+    def effective_extend_kernel(self, lower_bound: Optional[float], num_tokens: int):
+        """The packed FlashInfer prefill kernel requires more than one token."""
         kernel = self.extend_kernel
+        if self.prefill_backend.is_flashinfer() and num_tokens <= 1:
+            return self.triton_kernel
         if lower_bound is not None and not getattr(kernel, "supports_safe_gate", True):
             kernel = self.triton_kernel
         return kernel
@@ -353,7 +381,7 @@ def extend(
         query_start_loc: torch.Tensor,
         **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        kernel = self.effective_extend_kernel(kwargs.get("lower_bound"))
+        kernel = self.effective_extend_kernel(kwargs.get("lower_bound"), q.shape[1])
         return kernel.extend(
             q,
             k,
@@ -429,6 +457,12 @@ def __init__(self, model_runner: ModelRunner):
         self.kernel_dispatcher = KDAKernelDispatcher(
             decode_backend, prefill_backend, verify_backend
         )
+        if prefill_backend.is_flashinfer():
+            _validate_flashinfer_kda_prefill(
+                chunk_size=self.mamba_chunk_size,
+                enable_two_batch_overlap=get_exec().overlap.enable_two_batch_overlap,
+                prefill_cuda_graph_backend=get_exec().graph.cuda_graph_config.prefill.backend,
+            )
         # One-shot; emitted at the first fused-decode interception below.
         self._fused_override_notice = (
             "K3 fused KDA decode engaged: --linear-attn-decode-backend "
@@ -543,6 +577,25 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                     self.forward_metadata.mamba_track_mask_indices
                 ]
             )
+            if (
+                self.kernel_dispatcher.extend_kernel.uses_state_checkpoints
+                and forward_batch.forward_mode.is_extend_without_speculative()
+                and self.forward_metadata.track_ssm_h_src is not None
+                and self.forward_metadata.track_ssm_h_src.numel() > 0
+            ):
+                from sglang.srt.layers.attention.linear.kernels.kda_flashinfer_prefill import (
+                    build_flashinfer_kda_checkpoint_plan,
+                )
+
+                assert self._has_cpu_prefill_track_metadata(forward_batch), (
+                    "FlashInfer KDA checkpoints require host prefill tracking metadata"
+                )
+                build_flashinfer_kda_checkpoint_plan(
+                    forward_batch,
+                    self.forward_metadata,
+                    self.device,
+                    self.mamba_chunk_size,
+                )
 
     def forward_decode(
         self,
@@ -872,8 +925,7 @@ def forward_extend(
         k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
         v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
 
-        gate_was_flat = a.ndim == 3
-        if gate_was_flat:
+        if a.ndim == 3:
             a = a.unflatten(-1, (-1, layer.head_k_dim))
 
         track_ssm = self.forward_metadata.has_mamba_track_mask
@@ -892,7 +944,7 @@ def forward_extend(
             # Check the kernel the dispatcher will actually run (safe-gate
             # reroute included), not just the configured one.
             extend_kernel = self.kernel_dispatcher.effective_extend_kernel(
-                layer.lower_bound
+                layer.lower_bound, q.shape[1]
             )
             assert extend_kernel.supports_track_state_snapshot, (
                 f"{type(extend_kernel).__name__} cannot write the fp32 track "
@@ -917,7 +969,7 @@ def forward_extend(
             A_log=layer.A_log,
             dt_bias=layer.dt_bias,
             lower_bound=layer.lower_bound,
-            beta_is_raw=gate_was_flat,
+            beta_is_raw=True,
             extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
             extend_prefix_lens=forward_batch.extend_prefix_lens,
             layer_id=layer.layer_id,
@@ -933,6 +985,13 @@ def forward_extend(
             ),
             track_state=h_track_buf,
             track_chunk_idx=(track_chunk_idx if h_track_buf is not None else None),
+            state_checkpoint_cu_starts=self.forward_metadata.state_checkpoint_cu_starts,
+            num_state_checkpoints=self.forward_metadata.num_state_checkpoints,
+            state_checkpoint_every_n_tokens=(
+                self.forward_metadata.state_checkpoint_every_n_tokens
+            ),
+            state_checkpoint_indices=self.forward_metadata.state_checkpoint_indices,
+            track_ssm_h_batch_src=self.forward_metadata.track_ssm_h_batch_src,
         )
         if track_ssm:
             # Snapshot the SSM state at the last track-aligned chunk boundary

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

@@ -1,7 +1,7 @@
 """FlashInfer KDA decode/verify wrapper.
 
-Wraps ``flashinfer.kda_decode.recurrent_kda`` (SM100 / Blackwell). FlashInfer has
-no KDA prefill kernel, so ``extend`` stays on Triton / CuTe DSL.
+Wraps ``flashinfer.kda_decode.recurrent_kda`` (SM100 / Blackwell). KDA prefill
+uses a separate adapter in ``kda_flashinfer_prefill``.
 
 Contract with the Triton KDA reference:
   - raw per-K gate ``a`` is activated in-kernel as
@@ -85,11 +85,7 @@ def build_fused_accept_indices(
 
 
 class FlashInferKDAKernel(LinearAttnKernelBase):
-    """FlashInfer KDA kernel: SM100 decode + MTP (target_verify), topk=1.
-
-    Prefill (``extend``) is intentionally not implemented -- FlashInfer ships no
-    KDA chunk kernel; the dispatcher keeps prefill on Triton / CuTe DSL.
-    """
+    """FlashInfer KDA kernel: SM100 decode + MTP (target_verify), topk=1."""
 
     def __init__(self):
         available, self._recurrent_kda = _get_flashinfer_kda_kernel()
@@ -381,5 +377,6 @@ def target_verify(
 
     def extend(self, *args, **kwargs):
         raise NotImplementedError(
-            "FlashInferKDAKernel has no prefill kernel; keep prefill on Triton / CuTe DSL."
+            "FlashInferKDAKernel is decode/verify-only; use "
+            "FlashInferKDAPrefillKernel for prefill."
         )

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

@@ -0,0 +1,148 @@
+"""FlashInfer 0.7 CuTe DSL KDA prefill with radix-cache checkpoints."""
+
+from __future__ import annotations
+
+from itertools import accumulate
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def build_flashinfer_kda_checkpoint_plan(
+    forward_batch: ForwardBatch,
+    metadata: ForwardMetadata,
+    device: torch.device,
+    chunk_size: int,
+) -> None:
+    """Plan scheduler-generated tracked boundaries, asserting their validity.
+
+    Host track metadata follows build_prefill_track_plan's unaligned-row order.
+    """
+    extend_lens = forward_batch.extend_seq_lens_cpu
+    track_lens = forward_batch.mamba_track_seqlens_cpu
+    prefix_lens = forward_batch.extend_prefix_lens_cpu
+    track_mask = forward_batch.mamba_prefill_track_mask_cpu
+
+    checkpoint_counts = [length // chunk_size for length in extend_lens]
+    checkpoint_starts = list(accumulate(checkpoint_counts, initial=0))
+    # NOTE(kpham-sgl): FlashInfer GDN writes all periodic checkpoints for SGLang
+    # to select from; FlashInfer KDA writes only mapped destinations (-1 skips).
+    checkpoint_destinations = [-1] * checkpoint_starts[-1]
+    num_tracked_checkpoints = 0
+    for row, tracked in enumerate(track_mask):
+        if not tracked:
+            continue
+        relative_track_len = track_lens[row] - prefix_lens[row]
+        if relative_track_len % chunk_size == 0:
+            continue
+        completed_chunks = relative_track_len // chunk_size
+        assert 0 < completed_chunks <= checkpoint_counts[row], (
+            f"Invalid KDA checkpoint metadata: row={row}, "
+            f"relative_track_len={relative_track_len}, extend_len={extend_lens[row]}, "
+            f"chunk_size={chunk_size}"
+        )
+        checkpoint_destinations[checkpoint_starts[row] + completed_chunks - 1] = (
+            num_tracked_checkpoints
+        )
+        num_tracked_checkpoints += 1
+
+    assert num_tracked_checkpoints == metadata.track_ssm_h_batch_src.numel(), (
+        "KDA checkpoint rows disagree with the prefill tracking metadata"
+    )
+    metadata.state_checkpoint_cu_starts = torch.tensor(
+        checkpoint_starts, dtype=torch.int64, device=device
+    )
+    metadata.num_state_checkpoints = num_tracked_checkpoints
+    metadata.state_checkpoint_every_n_tokens = chunk_size
+    metadata.state_checkpoint_indices = torch.tensor(
+        checkpoint_destinations, dtype=torch.int32, device=device
+    )
+
+
+class FlashInferKDAPrefillKernel(LinearAttnKernelBase):
+    """Execute KDA prefill; FlashInfer validates the tensor contract."""
+
+    uses_state_checkpoints = True
+    supports_track_state_snapshot = True
+    supports_safe_gate = True
+
+    def __init__(self):
+        if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)):
+            raise ValueError("FlashInfer KDA prefill requires SM100 or SM103")
+        from flashinfer.kda import recurrent_kda
+
+        self._recurrent_kda = recurrent_kda
+
+    def decode(self, *args, **kwargs):
+        raise NotImplementedError("FlashInferKDAPrefillKernel is prefill-only")
+
+    def extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        lower_bound: float,
+        return_intermediate_states: bool = False,
+        state_checkpoint_cu_starts: Optional[torch.Tensor] = None,
+        num_state_checkpoints: int = 0,
+        state_checkpoint_every_n_tokens: int = 0,
+        state_checkpoint_indices: Optional[torch.Tensor] = None,
+        track_ssm_h_batch_src: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        needs_checkpoint = (
+            return_intermediate_states and kwargs.get("track_state") is not None
+        )
+        checkpoints = (
+            ssm_states.new_empty((num_state_checkpoints, *ssm_states.shape[1:]))
+            if needs_checkpoint
+            else None
+        )
+        result = self._recurrent_kda(
+            q=q.contiguous(),
+            k=k.contiguous(),
+            v=v.contiguous(),
+            g=g[:, : q.shape[1]].contiguous(),
+            beta=beta[:, : q.shape[1]].contiguous(),
+            A_log=A_log.reshape(-1).float().contiguous(),
+            dt_bias=dt_bias.reshape(q.shape[2], 128).float().contiguous(),
+            initial_state=ssm_states,
+            output_final_state=False,
+            use_qk_l2norm_in_kernel=True,
+            use_gate_in_kernel=True,
+            lower_bound=lower_bound,
+            cu_seqlens=query_start_loc,
+            ssm_state_indices=cache_indices.to(torch.int32),
+            beta_is_logit=True,
+            state_checkpoints=checkpoints,
+            checkpoint_cu_starts=(
+                state_checkpoint_cu_starts if needs_checkpoint else None
+            ),
+            checkpoint_state_indices=(
+                state_checkpoint_indices if needs_checkpoint else None
+            ),
+            checkpoint_every_n_tokens=(
+                state_checkpoint_every_n_tokens if needs_checkpoint else 0
+            ),
+            backend="cute-dsl",
+        )
+        if needs_checkpoint:
+            kwargs["track_state"][track_ssm_h_batch_src] = checkpoints.float()
+        output = result[0]
+        return (output, None) if return_intermediate_states else output

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

@@ -172,6 +172,7 @@ def extend(
         dt_bias: torch.Tensor | None = None,
         lower_bound: float | None = None,
         return_intermediate_states: bool = False,
+        beta_is_raw: bool = False,
         **kwargs,
     ) -> torch.Tensor:
         assert self._chunk_kda is not None
@@ -188,6 +189,7 @@ def extend(
             A_log=A_log,
             dt_bias=dt_bias,
             lower_bound=lower_bound,
+            beta_is_raw=beta_is_raw,
             output_intermediate_states=return_intermediate_states,
             track_state=kwargs.get("track_state"),
             track_chunk_idx=kwargs.get("track_chunk_idx"),