[KDA] Add FlashKDA prefill backend for safe-gate KDA linear attention (#29472)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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

@@ -60,6 +60,12 @@ class KDAKernelDispatcher:
 
         if prefill_backend.is_triton():
             self.extend_kernel = triton_kernel
+        elif prefill_backend.is_flashkda():
+            from sglang.srt.layers.attention.linear.kernels.kda_flashkda import (
+                FlashKDAKernel,
+            )
+
+            self.extend_kernel = FlashKDAKernel()
         elif prefill_backend.is_cutedsl():
             if not is_cuda():
                 raise ValueError("KDA CuTe DSL backend requires CUDA")
@@ -81,7 +87,8 @@ class KDAKernelDispatcher:
         else:
             raise ValueError(
                 f"Unsupported KDA prefill backend: {prefill_backend}. "
-                "KDA supports 'triton' or 'cutedsl' (cutedsl prefill needs SM100)."
+                "KDA supports 'triton', 'flashkda', or 'cutedsl' "
+                "(cutedsl prefill needs SM100)."
             )
 
         self.supports_packed_decode = getattr(
@@ -367,6 +374,14 @@ class KDAAttnBackend(MambaAttnBackendBase):
             A_log=layer.A_log,
             dt_bias=layer.dt_bias,
             lower_bound=getattr(layer, "lower_bound", None),
+            extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            # target_verify / draft_extend_v2 also reach forward_extend; they must
+            # stay rollback-able, so a kernel that commits state in place (e.g.
+            # FlashKDA) must not run for them.
+            is_spec_decode=(
+                forward_batch.forward_mode.is_target_verify()
+                or forward_batch.forward_mode.is_draft_extend_v2()
+            ),
         )
 
         return core_attn_out

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

@@ -0,0 +1,257 @@
+from typing import Optional
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+
+# FlashKDA chunk size. Sequences shorter than this fall back to Triton.
+_FLASHKDA_CHUNK_SIZE = 64
+
+# FlashKDA's max sequence length, Batches whose longest sequence exceeds this
+# fall back to Triton for the whole batch.
+_FLASHKDA_MAX_SEQ_LEN = 2048
+
+
+def _load_flash_kda():
+    """Import the optional ``flash_kda`` CUTLASS module."""
+    try:
+        import flash_kda
+    except ImportError as e:
+        raise ImportError(
+            "The 'flashkda' KDA prefill backend requires the flash_kda module, "
+            "which is not installed. Install it from source:\n"
+            "    pip install git+https://github.com/MoonshotAI/FlashKDA.git"
+        ) from e
+    return flash_kda
+
+
+def _triton_fallback(
+    q,
+    k,
+    v,
+    g,
+    beta,
+    ssm_states,
+    cache_indices,
+    query_start_loc,
+    A_log=None,
+    dt_bias=None,
+    lower_bound=None,
+):
+    """Fall back to the Triton chunk_kda kernel (handles all preprocessing).
+
+    `g` is the RAW gate; chunk_kda applies the gate activation internally when
+    A_log is provided, so A_log/dt_bias/lower_bound must be threaded through too
+    -- otherwise the fallback silently skips activation. chunk_kda updates the
+    ssm state in-place via cache_indices and returns only the output tensor.
+    """
+    from sglang.srt.layers.attention.fla.kda import chunk_kda
+
+    return chunk_kda(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=ssm_states,
+        initial_state_indices=cache_indices,
+        use_qk_l2norm_in_kernel=True,
+        cu_seqlens=query_start_loc,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        lower_bound=lower_bound,
+    )
+
+
+class FlashKDAKernel(LinearAttnKernelBase):
+    """FlashKDA (MoonshotAI) fully-fused CUTLASS KDA prefill backend.
+
+    Wraps the external ``flash_kda`` package (https://github.com/MoonshotAI/FlashKDA).
+
+    FlashKDA fuses q/k L2 norm, beta sigmoid, and the KDA gate *inside* the
+    kernel, so we pass RAW tensors plus ``A_log``/``dt_bias``/``lower_bound``.
+    It is prefill-only, bf16, K == V == 128, HV == H (no GVA), and requires the
+    safe (bounded) gate (``lower_bound`` set). The non-safe path and sequences
+    outside [chunk_size, max_seq_len] fall back to Triton ``chunk_kda``.
+    Requires an SM90+ GPU with the ``flash_kda`` package.
+    """
+
+    def decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        *,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        raise NotImplementedError("FlashKDAKernel only supports prefill (extend)")
+
+    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: Optional[torch.Tensor] = None,
+        dt_bias: Optional[torch.Tensor] = None,
+        lower_bound: Optional[float] = None,
+        extend_seq_lens_cpu: Optional[list] = None,
+        is_spec_decode: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        if self._should_fall_back(
+            lower_bound, is_spec_decode, query_start_loc, extend_seq_lens_cpu
+        ):
+            return _triton_fallback(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states,
+                cache_indices,
+                query_start_loc,
+                A_log=A_log,
+                dt_bias=dt_bias,
+                lower_bound=lower_bound,
+            )
+
+        return self._flashkda_extend(
+            q,
+            k,
+            v,
+            g,
+            beta,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+        )
+
+    @staticmethod
+    def _should_fall_back(
+        lower_bound: Optional[float],
+        is_spec_decode: bool,
+        query_start_loc: torch.Tensor,
+        extend_seq_lens_cpu: Optional[list],
+    ) -> bool:
+        """Whether to use the Triton chunk_kda path instead of the fused kernel."""
+        # Safe-gate only: the fused kernel does not support the unbounded gate
+        # (-exp(A_log)*softplus); those models leave lower_bound unset.
+        if lower_bound is None:
+            return True
+        # FlashKDA writes the committed recurrent state back in place, so it is
+        # unsafe for speculative verify / draft-extend forwards (which must stay
+        # rollback-able). Those reach this backend through forward_extend, so
+        # gate them here rather than relying on the decode/target_verify stubs.
+        if is_spec_decode:
+            return True
+        # Short sequences (< chunk size) and long sequences (> the crossover
+        # where Triton's chunked prefill wins) are faster on Triton. Read the
+        # per-request lengths from the CPU-side extend_seq_lens to avoid a
+        # GPU->CPU sync on every layer; derive from query_start_loc (one sync)
+        # only if they are unavailable.
+        if extend_seq_lens_cpu is not None:
+            if torch.is_tensor(extend_seq_lens_cpu):
+                lo = int(extend_seq_lens_cpu.min())
+                hi = int(extend_seq_lens_cpu.max())
+            else:
+                lo = min(extend_seq_lens_cpu)
+                hi = max(extend_seq_lens_cpu)
+        else:
+            seq_lens = query_start_loc[1:] - query_start_loc[:-1]
+            lo_t, hi_t = torch.aminmax(seq_lens)
+            lo, hi = int(lo_t), int(hi_t)
+        return lo < _FLASHKDA_CHUNK_SIZE or hi > _FLASHKDA_MAX_SEQ_LEN
+
+    def _flashkda_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: Optional[torch.Tensor] = None,
+        dt_bias: Optional[torch.Tensor] = None,
+        lower_bound: Optional[float] = None,
+    ) -> torch.Tensor:
+        flash_kda = _load_flash_kda()
+
+        # Input shapes (varlen, B == 1, matching chunk_kda's contract):
+        #   q, k = [1, packed_seq, H, K]   v = [1, packed_seq, HV, V]
+        #   g    = [1, packed_seq, HV, K]  beta = [1, packed_seq, H]
+        # flash_kda wants these 4D tensors directly and RAW (it fuses l2norm /
+        # beta sigmoid / gate activation in-kernel).
+        num_heads = q.shape[2]
+        head_dim = q.shape[3]
+        scale = head_dim**-0.5
+
+        q = q.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
+        g = g.contiguous()
+
+        # KimiDeltaAttention.forward already applies sigmoid to beta on the
+        # prefill path, but flash_kda expects beta LOGITS (it sigmoids
+        # internally). Invert back so the kernel recovers the intended value:
+        # sigmoid(logit(p)) == p. (triton/cuLA consume the post-sigmoid beta.)
+        beta = torch.logit(beta.float().clamp_(1e-7, 1.0 - 1e-7)).to(torch.bfloat16)
+        beta = beta.contiguous()
+
+        # flash_kda wants A_log [H] fp32 and dt_bias [H, K] fp32. The model
+        # stores A_log as [1, 1, H, 1] and dt_bias as 1D [H*K], so reshape both.
+        A_log = A_log.reshape(-1).float().contiguous()
+        if dt_bias is not None:
+            dt_bias = dt_bias.reshape(num_heads, -1).float().contiguous()
+
+        # cu_seqlens must be int64 for flash_kda (FLA casts to long).
+        cu_seqlens = query_start_loc.to(torch.int64)
+
+        # flash_kda varlen state is [N, H, V, K] -- the SAME layout as sglang's
+        # KDA pool, so no transpose is needed. Advanced indexing copies, so the
+        # final state is written back in-place below (matching chunk_kda).
+        initial_state = ssm_states[cache_indices].contiguous()
+
+        out_buf = torch.empty_like(v)
+        final_state = torch.empty_like(initial_state)
+
+        flash_kda.fwd(
+            q,
+            k,
+            v,
+            g,
+            beta,
+            scale,
+            out_buf,
+            A_log,
+            dt_bias,
+            lower_bound,
+            initial_state=initial_state,
+            final_state=final_state,
+            cu_seqlens=cu_seqlens,
+        )
+
+        ssm_states[cache_indices] = final_state
+
+        # out_buf is already [1, packed_seq, HV, V].
+        return out_buf

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

@@ -16,6 +16,7 @@ class LinearAttnKernelBackend(Enum):
     TRITON = "triton"
     CUTEDSL = "cutedsl"
     FLASHINFER = "flashinfer"
+    FLASHKDA = "flashkda"
     CUSTOM = "custom"
 
     @classmethod
@@ -31,6 +32,9 @@ class LinearAttnKernelBackend(Enum):
     def is_flashinfer(self):
         return self == LinearAttnKernelBackend.FLASHINFER
 
+    def is_flashkda(self):
+        return self == LinearAttnKernelBackend.FLASHKDA
+
     def is_custom(self):
         return self == LinearAttnKernelBackend.CUSTOM
 

python/sglang/srt/server_args.py

@@ -315,7 +315,7 @@ MAMBA_RADIX_CACHE_STRATEGY_CHOICES = [
 
 MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"]
 
-LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl", "flashinfer"]
+LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl", "flashinfer", "flashkda"]
 
 
 # Allow external code to add more choices
@@ -5133,6 +5133,26 @@ class ServerArgs:
 
         # SM100+ FlashInfer GDN decode requires bf16 state; SM90 uses float32.
         decode = self.linear_attn_decode_backend or self.linear_attn_backend
+
+        # FlashKDA is a prefill-only KDA kernel (no decode kernel) but shares the
+        # backend choice list, so guard it from being selected for decode: error
+        # on an explicit --linear-attn-decode-backend flashkda, and fall back to
+        # triton decode when it was only inherited from base=flashkda (prefill
+        # keeps FlashKDA).
+        if decode == "flashkda":
+            if self.linear_attn_decode_backend == "flashkda":
+                raise ValueError(
+                    "--linear-attn-decode-backend flashkda is not supported: "
+                    "FlashKDA is prefill-only. Use "
+                    "--linear-attn-prefill-backend flashkda (decode stays on triton)."
+                )
+            self.linear_attn_decode_backend = "triton"
+            decode = "triton"
+            logger.info(
+                "FlashKDA is prefill-only; using triton for KDA decode "
+                "(FlashKDA stays on prefill)."
+            )
+
         if (
             decode == "flashinfer"
             and self.mamba_ssm_dtype != "bfloat16"