[KDA] Add FlashInfer SM100 KDA decode + MTP (target_verify) backend (#30113)

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

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

@@ -52,12 +52,32 @@ class KDAKernelDispatcher:
             )
 
             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 (
+                FlashInferKDAKernel,
+            )
+
+            self.decode_kernel = FlashInferKDAKernel()
         else:
             raise ValueError(
                 f"Unsupported KDA decode backend: {decode_backend}. "
-                "KDA currently only supports 'triton'."
+                "KDA supports 'triton', 'cutedsl', or 'flashinfer'."
             )
 
+        # target_verify (MTP / speculative decode) kernel: each decode backend
+        # verifies with its own kernel. FlashInfer decode uses recurrent_kda (SM100,
+        # chain only); Triton -- and CuTe DSL, which has no verify of its own -- use
+        # the Triton fused KDA verify, which handles chain + tree
+        # (retrieve_parent_token) and per-step checkpointing and is the reference the
+        # KDA backend correctness tests assert against.
+        self.verify_kernel = (
+            self.decode_kernel if decode_backend.is_flashinfer() else triton_kernel
+        )
+
         if prefill_backend.is_triton():
             self.extend_kernel = triton_kernel
         elif prefill_backend.is_flashkda():
@@ -97,6 +117,7 @@ class KDAKernelDispatcher:
 
         rank0_log(
             f"KDA kernel dispatcher: decode={self.decode_kernel.__class__.__name__}, "
+            f"verify={self.verify_kernel.__class__.__name__}, "
             f"extend={self.extend_kernel.__class__.__name__} "
             f"packed_decode={self.supports_packed_decode}"
         )
@@ -163,6 +184,45 @@ class KDAKernelDispatcher:
             **kwargs,
         )
 
+    def target_verify(
+        self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        intermediate_states_buffer: torch.Tensor,
+        intermediate_state_indices: torch.Tensor,
+        cache_steps: int,
+        retrieve_parent_token: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        """MTP / speculative-decode verify, routed to ``self.verify_kernel``
+        (FlashInfer decode -> recurrent_kda; Triton / CuTe DSL decode -> the Triton
+        fused KDA verify)."""
+        return self.verify_kernel.target_verify(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            intermediate_states_buffer=intermediate_states_buffer,
+            intermediate_state_indices=intermediate_state_indices,
+            cache_steps=cache_steps,
+            retrieve_parent_token=retrieve_parent_token,
+        )
+
     def extend(
         self,
         q: torch.Tensor,
@@ -196,7 +256,23 @@ class KDAAttnBackend(MambaAttnBackendBase):
         super().__init__(model_runner)
         decode_backend = get_linear_attn_decode_backend()
         prefill_backend = get_linear_attn_prefill_backend()
+        # KDA FlashInfer speculative decode (target_verify) is linear-chain only --
+        # recurrent_kda has no tree-ancestor traversal. Reject EAGLE tree verify
+        # (topk > 1) early at setup instead of deep in the per-step verify call.
+        # (The kernel keeps a per-call retrieve_parent_token guard as a backstop; it
+        # also covers ngram tree, which this topk field does not.)
+        speculative_topk = model_runner.server_args.speculative_eagle_topk or 1
+        if decode_backend.is_flashinfer() and speculative_topk > 1:
+            raise ValueError(
+                "KDA FlashInfer speculative decoding only supports topk=1 "
+                "(EAGLE tree verify / retrieve_parent_token is unsupported)."
+            )
         self.kernel_dispatcher = KDAKernelDispatcher(decode_backend, prefill_backend)
+        # Per-request row index into the speculative `intermediate_ssm` scratch,
+        # used by the MTP / target_verify path (mirrors GDNAttnBackend).
+        self.verify_intermediate_state_indices = torch.arange(
+            self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device
+        )
 
     def forward_decode(
         self,
@@ -212,18 +288,8 @@ class KDAAttnBackend(MambaAttnBackendBase):
         query_start_loc = self.forward_metadata.query_start_loc
         cache_indices = self.forward_metadata.mamba_cache_indices
 
-        # ReplaySSM ring: per-layer ring slices + the once-per-forward per-row
-        # write cursor. All None unless --enable-linear-replayssm, so packed_decode
-        # falls through to the byte-identical legacy KDA path. KDA ships WITHOUT
-        # radix coordination for now, so force_flush is None/zeroed (the ring
-        # flushes only at the natural write_pos == L-1 wrap; set in the shared
-        # HybridLinearAttn metadata, which zeroes force_flush for KDA models).
-        # NOTE: ReplaySSM decode is a GDN (scalar-gate) bandwidth win; on KDA the
-        # per-K g_cache is K x larger and the reconstruction refolds the per-K
-        # decay every step, so it is correct but SLOWER than packed (a measured
-        # decode regression). Kept wired for correctness + the spec-decode path;
-        # not recommended for KDA decode. Revisit on Blackwell (more tensor-core
-        # throughput may flip the compute/bandwidth tradeoff).
+        # ReplaySSM is mostly a GDN bandwidth optimization. It remains wired for
+        # KDA correctness paths, but packed decode is faster for KDA today.
         replayssm_write_pos = getattr(
             self.forward_metadata, "replayssm_write_pos", None
         )
@@ -243,16 +309,8 @@ class KDAAttnBackend(MambaAttnBackendBase):
             conv_state_indices=cache_indices,
         )
 
-        # Skip split + reshape by consuming the packed mixed_qkv directly in a
-        # single fused Triton kernel (KDA per-K gate variant of GDN PR #20627).
-        #
-        # The packed kernel hard-assumes one token per sequence (T=1): it has no
-        # query_start_loc / per-sequence loop. forward_decode is only entered in
-        # decode mode (see HybridLinearAttnBackend.forward dispatch), where each
-        # request contributes exactly one token, so #tokens == #requests. Multi-
-        # token-per-seq speculative paths (target_verify / draft_extend) go
-        # through forward_extend instead. Assert the invariant so a future
-        # routing change fails loudly rather than silently corrupting state.
+        # The packed kernel assumes one token per request. Assert the dispatch
+        # invariant before taking the fused path.
         if self.kernel_dispatcher.supports_packed_decode:
             assert qkv.shape[0] == cache_indices.shape[0], (
                 "KDA packed decode requires one token per sequence (T=1): "
@@ -303,6 +361,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
         b: torch.Tensor,
         **kwargs,
     ):
+        # MTP / speculative-decode verify is a multi-token-per-seq path with
+        # per-step state checkpointing + central rollback; handled separately.
+        if forward_batch.forward_mode.is_target_verify():
+            return self._forward_target_verify(layer, forward_batch, mixed_qkv, a, b)
+
         query_start_loc = self.forward_metadata.query_start_loc
         cache_indices = self.forward_metadata.mamba_cache_indices
 
@@ -375,13 +438,89 @@ class KDAAttnBackend(MambaAttnBackendBase):
             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()
-            ),
+            # draft_extend_v2 must stay rollback-able, so kernels that commit state
+            # in place (e.g. FlashKDA) must not run for it.
+            is_spec_decode=forward_batch.forward_mode.is_draft_extend_v2(),
         )
 
         return core_attn_out
+
+    def _forward_target_verify(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+    ):
+        """MTP / speculative-decode verify (topk=1), mirroring the GDN backend.
+
+        Conv1d runs per draft token with intermediate-window checkpointing; the
+        SSM verify kernel writes each draft token's post-state into the
+        speculative `intermediate_ssm` scratch so the central post-verify rollback
+        (update_mamba_state_after_mtp_verify) can commit the accepted-length state.
+        """
+        fm = self.forward_metadata
+        seq_len = mixed_qkv.shape[0]
+        query_start_loc = fm.query_start_loc
+        cache_indices = fm.mamba_cache_indices
+        retrieve_next_token = fm.retrieve_next_token
+        retrieve_next_sibling = fm.retrieve_next_sibling
+        retrieve_parent_token = fm.retrieve_parent_token
+
+        mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = mamba_cache_params.conv[0]
+        ssm_states = mamba_cache_params.temporal
+        intermediate_state_cache = getattr(mamba_cache_params, "intermediate_ssm", None)
+        if intermediate_state_cache is None:
+            raise RuntimeError(
+                "KDA target_verify requires a speculative mamba cache "
+                "(MambaPool.SpeculativeState); none found."
+            )
+        intermediate_conv_window_cache = mamba_cache_params.intermediate_conv_window[0]
+        intermediate_state_indices = self.verify_intermediate_state_indices
+
+        draft_token_num = forward_batch.spec_info.draft_token_num
+        batch_size = seq_len // draft_token_num
+
+        # causal_conv1d_update expects [.., dim, width]. KDA keeps dense conv-window
+        # scratch because the deduplicated overlapping layout cannot be transposed.
+        mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1).transpose(
+            1, 2
+        )
+        mixed_qkv_processed = causal_conv1d_update(
+            mixed_qkv_reshaped,
+            conv_states.transpose(-1, -2),
+            layer.conv_weights,
+            layer.bias,
+            activation="silu",
+            conv_state_indices=cache_indices[:batch_size],
+            intermediate_conv_window=intermediate_conv_window_cache.transpose(-1, -2),
+            intermediate_state_indices=intermediate_state_indices[:batch_size],
+            retrieve_next_token=retrieve_next_token,
+            retrieve_next_sibling=retrieve_next_sibling,
+            retrieve_parent_token=retrieve_parent_token,
+        )
+        mixed_qkv = mixed_qkv_processed.transpose(1, 2).reshape(seq_len, -1)
+
+        q, k, v = mixed_qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
+        q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
+        k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)
+        v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)
+
+        return self.kernel_dispatcher.target_verify(
+            A_log=layer.A_log,
+            dt_bias=layer.dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            intermediate_states_buffer=intermediate_state_cache,
+            intermediate_state_indices=intermediate_state_indices,
+            cache_steps=draft_token_num,
+            retrieve_parent_token=retrieve_parent_token,
+        )

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

@@ -0,0 +1,288 @@
+"""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.
+
+Contract with the Triton KDA reference:
+  - raw per-K gate ``a`` is activated in-kernel as
+    ``-exp(A_log) * softplus(a + dt_bias)``;
+  - beta ``b`` is a logit, so this wrapper passes ``sigmoid(b)``;
+  - q/k are L2-normalized in-kernel;
+  - state layout is ``[N, HV, V, K]`` for committed and speculative state.
+"""
+
+import logging
+import os
+from typing import Optional
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+from sglang.srt.utils import is_cuda
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Lazy import for the FlashInfer KDA kernel
+# ---------------------------------------------------------------------------
+_flashinfer_kda_available: Optional[bool] = None
+_flashinfer_recurrent_kda = None
+
+
+def _get_flashinfer_kda_kernel():
+    """Lazy import for FlashInfer ``recurrent_kda`` (decode + MTP).
+
+    Returns (available, recurrent_kda_fn).
+    """
+    global _flashinfer_kda_available, _flashinfer_recurrent_kda
+    if _flashinfer_kda_available is None:
+        try:
+            os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
+
+            from flashinfer.kda_decode import recurrent_kda
+
+            _flashinfer_recurrent_kda = recurrent_kda
+            # recurrent_kda is SM100-only (CuTe DSL, Blackwell).
+            _flashinfer_kda_available = (
+                is_cuda() and torch.cuda.get_device_capability()[0] >= 10
+            )
+            if _flashinfer_kda_available:
+                logger.info("FlashInfer KDA kernel (recurrent_kda) loaded successfully")
+        except (ImportError, RuntimeError) as e:
+            logger.warning(f"FlashInfer KDA kernel not available: {e}")
+            _flashinfer_kda_available = False
+            _flashinfer_recurrent_kda = None
+    return _flashinfer_kda_available, _flashinfer_recurrent_kda
+
+
+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.
+    """
+
+    def __init__(self):
+        available, self._recurrent_kda = _get_flashinfer_kda_kernel()
+        if not available or self._recurrent_kda is None:
+            raise RuntimeError(
+                "FlashInfer KDA kernel (recurrent_kda) is not available. "
+                "Requires SM100 (Blackwell) and a FlashInfer build with KDA support."
+            )
+        # Cache the per-layer constant gate-param prep (A_log/dt_bias reshape+cast),
+        # keyed by tensor identity. Layer params are persistent weights so id() is
+        # stable; this removes the per-call reshape/float/contiguous work.
+        self._gate_cache: dict = {}
+        # Cache the constant per-(row-map, batch, T) verify scatter indices
+        # (ssm_state_indices), which never change across verify calls.
+        self._verify_idx_cache: dict = {}
+        logger.info("Using FlashInfer KDA kernel")
+
+    # ---- gate / beta normalization (shared by decode + verify) ----
+
+    def _prep_gate_params(self, A_log: torch.Tensor, dt_bias: torch.Tensor):
+        # A_log: [1, 1, H, 1] -> [H] fp32; dt_bias: [H*K] (1D) -> fp32. Cached per
+        # layer (constant weights) so this is a dict lookup on the hot path.
+        key = (id(A_log), id(dt_bias))
+        cached = self._gate_cache.get(key)
+        if cached is not None:
+            return cached
+        A_log_fi = A_log.reshape(-1).float().contiguous()
+        dt_bias_fi = (
+            dt_bias.reshape(-1).float().contiguous() if dt_bias is not None else None
+        )
+        self._gate_cache[key] = (A_log_fi, dt_bias_fi)
+        return A_log_fi, dt_bias_fi
+
+    @staticmethod
+    def _beta_logit_to_prob(b: torch.Tensor) -> torch.Tensor:
+        # Triton KDA does beta = sigmoid(b); recurrent_kda wants beta pre-sigmoided.
+        # torch.sigmoid computes in fp32 internally, so a single sigmoid on the bf16
+        # logit is enough (avoids an explicit fp32 upcast + downcast = 2 extra kernels).
+        return torch.sigmoid(b).to(torch.bfloat16)
+
+    # ---- decode ----
+
+    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:
+        batch_size = cache_indices.shape[0]
+        num_heads = q.shape[2]
+        head_k_dim = q.shape[3]
+        num_v_heads = v.shape[2]
+        head_v_dim = v.shape[3]
+
+        # Pack each request as a length-1 sequence ([1, B, ...] + cu_seqlens) so
+        # recurrent_kda indexes the committed pool IN-KERNEL via ssm_state_indices.
+        # The plain [B, 1, ...] path (no cu_seqlens) instead python-gathers
+        # initial_state[indices] and scatters it back with index_put around the
+        # kernel (~141us at B=64 in ncu); the cu_seqlens path skips both. q/k/v
+        # already arrive as [1, B, H, D] from forward_decode, so the reshape is a
+        # no-op view. recurrent_kda's cp.async + shared-mem staging are hardwired to
+        # bf16 (2-byte elements) for q/k/v/g/beta and the state, so every input is
+        # cast to bf16 -- a no-op for the common bf16 KDA model, a correct downcast
+        # otherwise (float16 bits would be reinterpreted as bf16 without the cast).
+        query_fi = q.reshape(1, batch_size, num_heads, head_k_dim).to(torch.bfloat16)
+        key_fi = k.reshape(1, batch_size, num_heads, head_k_dim).to(torch.bfloat16)
+        value_fi = v.reshape(1, batch_size, num_v_heads, head_v_dim).to(torch.bfloat16)
+        g_fi = a.reshape(1, batch_size, num_v_heads, head_k_dim).to(torch.bfloat16)
+        beta_fi = self._beta_logit_to_prob(b).reshape(1, batch_size, num_v_heads)
+
+        A_log_fi, dt_bias_fi = self._prep_gate_params(A_log, dt_bias)
+
+        # Softplus gate (lower_bound=None) to match the Triton KDA decode path;
+        # in-place state update into the committed pool (no rollback for decode).
+        # query_start_loc is the decode cu_seqlens (one token per request).
+        output_fi, _ = self._recurrent_kda(
+            q=query_fi,
+            k=key_fi,
+            v=value_fi,
+            g=g_fi,
+            beta=beta_fi,
+            A_log=A_log_fi,
+            dt_bias=dt_bias_fi,
+            scale=None,
+            initial_state=ssm_states,
+            output_final_state=False,
+            use_qk_l2norm_in_kernel=True,
+            use_gate_in_kernel=True,
+            lower_bound=None,
+            cu_seqlens=query_start_loc.to(torch.int32),
+            ssm_state_indices=cache_indices.to(torch.int32),
+        )
+
+        return output_fi.view(1, batch_size, num_v_heads, head_v_dim)
+
+    # ---- target_verify (MTP, topk=1) ----
+
+    def target_verify(
+        self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        intermediate_states_buffer: torch.Tensor,
+        intermediate_state_indices: torch.Tensor,
+        cache_steps: int,
+        retrieve_parent_token: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        if retrieve_parent_token is not None:
+            raise RuntimeError(
+                "FlashInfer KDA verify kernel only supports topk=1 "
+                "(retrieve_parent_token must be None)."
+            )
+
+        seq_len = q.shape[1]
+        batch_size = query_start_loc.shape[0] - 1
+        draft_token_num = cache_steps  # T = 1 + num_spec_tokens
+        num_spec_tokens = draft_token_num - 1
+        num_heads = q.shape[2]
+        head_k_dim = q.shape[3]
+        num_v_heads = v.shape[2]
+        head_v_dim = v.shape[3]
+
+        # Packed [1, N*T, ...] inputs, cu_seqlens = query_start_loc (draft stride).
+        # recurrent_kda is bf16-only (see decode), so cast every input to bf16.
+        q_fi = q.reshape(1, seq_len, num_heads, head_k_dim).to(torch.bfloat16)
+        k_fi = k.reshape(1, seq_len, num_heads, head_k_dim).to(torch.bfloat16)
+        v_fi = v.reshape(1, seq_len, num_v_heads, head_v_dim).to(torch.bfloat16)
+        g_fi = a.reshape(1, seq_len, num_v_heads, head_k_dim).to(torch.bfloat16)
+        beta_fi = self._beta_logit_to_prob(b).reshape(1, seq_len, num_v_heads)
+
+        A_log_fi, dt_bias_fi = self._prep_gate_params(A_log, dt_bias)
+
+        # recurrent_kda indexes a flat state pool. Map each request/step to the
+        # matching slot in SGLang's [scratch_row, allocated_step, HV, V, K] buffer.
+        scratch = intermediate_states_buffer  # [N_scratch, T, HV, V, K]
+        scratch_steps = scratch.shape[1]
+        if draft_token_num > scratch_steps:
+            raise RuntimeError(
+                f"KDA verify needs {draft_token_num} scratch steps, "
+                f"but intermediate_ssm only has {scratch_steps}."
+            )
+
+        base_rows = intermediate_state_indices[:batch_size]
+        cache_key = (
+            id(intermediate_state_indices),
+            batch_size,
+            draft_token_num,
+            scratch_steps,
+        )
+        ssm_state_indices = self._verify_idx_cache.get(cache_key)
+        if ssm_state_indices is None:
+            # The fast seed copy below assumes row n in scratch belongs to request n.
+            expected = torch.arange(
+                batch_size, device=base_rows.device, dtype=base_rows.dtype
+            )
+            if not torch.equal(base_rows, expected):
+                raise RuntimeError(
+                    "FlashInfer KDA verify requires an identity intermediate row-map "
+                    "(verify_intermediate_state_indices must be arange)."
+                )
+            step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32)
+            ssm_state_indices = (
+                base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :]
+            ).contiguous()  # [N, T]
+            self._verify_idx_cache[cache_key] = ssm_state_indices
+
+        # Seed step 0 from committed state, then recurrent_kda overwrites it with
+        # token-0 post-state. Padded graph rows clamp to slot 0; their output is ignored.
+        base_state = ssm_states.index_select(
+            0, cache_indices[:batch_size].clamp(min=0).to(torch.int64)
+        )
+        scratch[:batch_size, 0].copy_(base_state)
+
+        # Same storage as scratch, flattened over the allocated step stride.
+        state_pool = scratch.view(
+            scratch.shape[0] * scratch_steps, num_v_heads, head_v_dim, head_k_dim
+        )
+
+        output_fi, _ = self._recurrent_kda(
+            q=q_fi,
+            k=k_fi,
+            v=v_fi,
+            g=g_fi,
+            beta=beta_fi,
+            A_log=A_log_fi,
+            dt_bias=dt_bias_fi,
+            scale=None,
+            initial_state=state_pool,
+            output_final_state=False,
+            use_qk_l2norm_in_kernel=True,
+            use_gate_in_kernel=True,
+            lower_bound=None,
+            cu_seqlens=query_start_loc.to(torch.int32),
+            ssm_state_indices=ssm_state_indices,
+            num_spec_tokens=num_spec_tokens,
+        )
+
+        return output_fi.view(1, seq_len, num_v_heads, head_v_dim)
+
+    # ---- extend (prefill): not provided by FlashInfer ----
+
+    def extend(self, *args, **kwargs):
+        raise NotImplementedError(
+            "FlashInferKDAKernel has no prefill kernel; keep prefill on Triton / CuTe DSL."
+        )

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

@@ -141,6 +141,53 @@ class TritonKDAKernel(LinearAttnKernelBase):
             is_kda=True,
         )
 
+    def target_verify(
+        self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        intermediate_states_buffer: torch.Tensor,
+        intermediate_state_indices: torch.Tensor,
+        cache_steps: int,
+        retrieve_parent_token: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        # KDA MTP / speculative-decode verify via the fused KDA kernel (IS_KDA=True),
+        # mirroring the GDN triton verify path. Reads the committed state, writes
+        # per-draft-token intermediate states to the scratch buffer, does NOT mutate
+        # the committed pool (disable_state_update=True), and handles chain + tree
+        # (retrieve_parent_token). The verify kernel for the Triton / CuTe DSL KDA
+        # decode backends, and the reference the KDA correctness tests assert against.
+        return fused_sigmoid_gating_delta_rule_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            initial_state_source=ssm_states,
+            initial_state_indices=cache_indices,
+            cu_seqlens=query_start_loc,
+            use_qk_l2norm_in_kernel=True,
+            softplus_beta=1.0,
+            softplus_threshold=20.0,
+            is_kda=True,
+            disable_state_update=True,
+            intermediate_states_buffer=intermediate_states_buffer,
+            intermediate_state_indices=intermediate_state_indices,
+            cache_steps=cache_steps,
+            retrieve_parent_token=retrieve_parent_token,
+        )
+
     def extend(
         self,
         q: torch.Tensor,

python/sglang/srt/mem_cache/memory_pool.py

@@ -102,21 +102,19 @@ _use_aiter = bool(envs.SGLANG_USE_AITER.get()) and _is_hip
 
 
 def conv_window_dedup_enabled(
-    is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int]
+    is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int], is_kda: bool
 ) -> bool:
     """Whether the deduplicated sliding-window conv-intermediate layout is safe.
 
-    It is only correct for a *linear* draft chain (``speculative_eagle_topk <= 1``,
-    i.e. NEXTN / MTP): consecutive draft tokens then form a true sliding window, so
-    the overlapping physical columns hold identical values. Under EAGLE *tree*
-    verify (``topk > 1``) the conv kernel walks per-token tree ancestors, so aliased
-    columns can need different values from different parent chains -> fall back to
-    the dense layout. NPU/CPU also keep the dense layout (their kernels assume
-    contiguous per-step windows). See ``MambaPool.__init__``.
+    It is safe for CUDA linear draft chains whose kernels consume the window raw.
+    Tree verify, NPU/CPU, and KDA keep dense windows: tree ancestors need independent
+    windows, platform kernels expect contiguous steps, and KDA transposes the window
+    before conv so the overlapping ``as_strided`` layout would corrupt stores.
     """
     return (
         not is_npu
         and not is_cpu
+        and not is_kda
         and (speculative_eagle_topk is None or speculative_eagle_topk <= 1)
     )
 
@@ -576,7 +574,7 @@ class MambaPool:
                 # `fused_conv_window_scatter_with_mask` scatter is layout-agnostic,
                 # so the dense fallback reads correctly through the same code path.
                 dedup_conv_window = conv_window_dedup_enabled(
-                    _is_npu, _is_cpu, speculative_eagle_topk
+                    _is_npu, _is_cpu, speculative_eagle_topk, cache_params.is_kda
                 )
                 self._intermediate_conv_window_phys = []
                 if dedup_conv_window: