feat(kda): add fused speculative decode kernel

docs/api/kda_decode.rst

@@ -10,6 +10,16 @@ The public ``recurrent_kda`` API supports standard decode with one token per
 sequence (``T=1``) and packed speculative decode with two or more tokens per
 sequence (``T>=2``).
 
+``fused_kda_decode`` combines width-four causal convolution, one KDA update,
+and gated RMSNorm for ``T=1``. ``packed_fused_kda_decode`` extends the same
+pipeline to packed ragged speculative decode, retaining one recurrent
+checkpoint per token and an extended rolling convolution window. Its optional
+``t1_state_indices`` argument is a T=1-only fast path containing caller-resolved
+per-row cache slots. T>1 always uses the ``state_indices[N, T]``,
+``query_start_loc``, and ``num_accepted_tokens`` metadata with the multitoken
+kernel; passing ``t1_state_indices`` for T>1 is rejected. No T>1 kernel
+behavior or performance is changed by this optimization.
+
 Pass ``backend="cake"`` to select the exported Cake backend. On SM100-family
 SM100a (B200/GB200) and SM103a (B300/GB300) devices, its D128 ``T=1..6``
 family with in-kernel QK normalization exports 25 frozen CUDA bodies:
@@ -125,5 +135,6 @@ small ``B * HV * T``. Requires SM90+ for the WY path and ``K = V = 128``;
     :toctree: ../generated
 
     fused_kda_decode
+    packed_fused_kda_decode
     packed_kda_decode
     recurrent_kda

flashinfer/__init__.py

@@ -218,6 +218,7 @@ def __dir__() -> list[str]:
 from .kda import recurrent_kda as recurrent_kda
 from .kda_decode import fused_kda_decode as fused_kda_decode
 from .kda_decode import packed_kda_decode as packed_kda_decode
+from .kda_decode import packed_fused_kda_decode
 from .cake_minimax_h3 import MiniMaxH3Mxfp8PreAttention as MiniMaxH3Mxfp8PreAttention
 from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper
 from . import mhc as mhc

flashinfer/kda_decode.py

@@ -29,6 +29,7 @@
 
 from .api_logging import flashinfer_api
 from .trace.templates.kda import (
+    packed_fused_kda_decode_trace,
     fused_kda_decode_trace,
     packed_kda_decode_trace,
     recurrent_kda_trace,
@@ -45,6 +46,16 @@
     _FUSED_KDA_DECODE_AVAILABLE = False
 
 from .kda_kernels import run_packed_kda_decode as _run_packed_kda_decode
+
+try:
+    from .kda_kernels.fused_kda_decode_multitoken import (
+        run_fused_kda_decode_multitoken as _run_packed_fused_kda_decode,
+    )
+
+    _FUSED_KDA_DECODE_PACKED_AVAILABLE = True
+except (ImportError, RuntimeError):
+    _run_packed_fused_kda_decode = None
+    _FUSED_KDA_DECODE_PACKED_AVAILABLE = False
 from .kda_kernels import run_recurrent_kda as _run_recurrent_kda
 
 try:
@@ -476,6 +487,227 @@ def fused_kda_decode(
     )
 
 
+def _validate_packed_t1_remap_inputs(x, query_start_loc, state_indices):
+    """Validate remap metadata structure without reading device values."""
+    if query_start_loc.dtype != torch.int32:
+        raise TypeError(
+            f"query_start_loc must have dtype torch.int32, got {query_start_loc.dtype}"
+        )
+    if state_indices.dtype != torch.int32:
+        raise TypeError(
+            f"state_indices must have dtype torch.int32, got {state_indices.dtype}"
+        )
+    if state_indices.ndim != 2 or state_indices.shape[1] != 1:
+        raise ValueError("state_indices must have shape [N, 1]")
+    if (
+        query_start_loc.ndim != 1
+        or query_start_loc.shape[0] != state_indices.shape[0] + 1
+    ):
+        raise ValueError("query_start_loc must have shape [state_indices.shape[0] + 1]")
+    if not query_start_loc.is_contiguous():
+        raise ValueError("query_start_loc must be contiguous")
+    if not state_indices.is_contiguous():
+        raise ValueError("state_indices must be contiguous")
+    if query_start_loc.device != x.device or state_indices.device != x.device:
+        raise ValueError("T=1 remap metadata must be on the same device as x")
+
+
+@flashinfer_api(trace=packed_fused_kda_decode_trace)
+def packed_fused_kda_decode(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    conv_state: torch.Tensor,
+    raw_gate: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    state_indices: torch.Tensor,
+    state: torch.Tensor,
+    output_gate: torch.Tensor,
+    norm_weight: torch.Tensor,
+    lower_bound: Optional[float] = -5.0,
+    norm_eps: float = 1e-5,
+    output: Optional[torch.Tensor] = None,
+    *,
+    query_start_loc: torch.Tensor,
+    num_accepted_tokens: torch.Tensor,
+    t1_state_indices: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    r"""Run packed T>=1 fused KDA with per-token cache checkpoints.
+
+    Tokens use a packed ragged layout described by ``query_start_loc``.
+    ``state_indices`` has shape ``[N, T]``, where T is the maximum verification
+    length. ``num_accepted_tokens[n] - 1`` selects the source checkpoint and
+    convolution-history offset for sequence ``n``; token ``t`` writes its
+    recurrent checkpoint to ``state_indices[n, t]``. The convolution cache is
+    one rolling window of length ``T + 2`` in ``state_indices[n, 0]``.
+
+    This facade is CuTe-only: T=1 uses the CuTe DSL implementation underlying
+    ``fused_kda_decode`` (without Cake or backend selection), including
+    bfloat16-state and softplus-gate support. The optional
+    ``t1_state_indices`` optimization is strictly T=1-only; passing it for
+    T>1 fails closed. T>1 uses the SM10x CuTe DSL packed backend, which
+    requires FP32
+    recurrent state and a finite negative ``lower_bound``. Both paths require
+    head dimension 128, convolution width four, and 12, 24, 32, 48, or 96
+    heads. Other tensors follow :func:`fused_kda_decode`, with the packed row
+    count supplied by ``x``.
+
+    ``query_start_loc`` must be contiguous int32 of shape ``[N+1]``, start at
+    zero, and contain nondecreasing offsets within ``x`` with active lengths no
+    larger than T.
+    ``num_accepted_tokens`` must be contiguous int32 of shape ``[N]`` with
+    values in ``[1, T]``. For T=1, ``t1_state_indices`` may optionally provide
+    the already-resolved contiguous per-row cache slots with shape
+    ``[num_rows]``. When supplied, it is used by the direct fused T=1
+    specialization and ``query_start_loc`` is retained only as structural
+    metadata; its values are not read by the host or kernel. This is intended
+    for CUDA-Graph-stable caller-owned buffers updated in-place.
+    Active recurrent destinations must be positive cache
+    slots and must not alias destinations from another active sequence;
+    non-positive or zero-length rows are null rows and do not mutate either
+    cache. In T=1, packed rows beyond ``query_start_loc[-1]`` are trailing
+    unused capacity: they receive zero output and do not mutate either cache.
+    These value constraints are caller-owned so CUDA Graph replay never performs
+    a device-to-host validation sync.
+
+    Args:
+        x:
+            Packed QKV projection with shape ``[num_rows, 3 * H * 128]`` and
+            dtype bfloat16.
+        weight:
+            Depthwise convolution weights with shape ``[3, 4, H * 128]`` and
+            dtype float32.
+        conv_state:
+            Paged bfloat16 convolution cache. T=1 uses history length three;
+            T>1 uses the extended rolling history length ``T + 2``.
+        raw_gate:
+            Raw recurrence gate with shape ``[1, num_rows, H, 128]`` and dtype
+            bfloat16.
+        raw_beta:
+            Raw delta-rule learning-rate logits with shape
+            ``[1, num_rows, H]`` and dtype bfloat16.
+        A_log:
+            Log decay parameter with ``H`` elements and dtype float32.
+        dt_bias:
+            Per-channel decay bias with ``H * 128`` elements and dtype float32.
+        state_indices:
+            Contiguous int32 cache slots with shape ``[N, T]``. For T>1,
+            column ``t`` is the recurrent checkpoint destination for token
+            ``t``. Column zero also selects the rolling convolution cache.
+        query_start_loc:
+            Contiguous int32 packed-row offsets with shape ``[N + 1]``.
+        num_accepted_tokens:
+            Contiguous int32 accepted-token counts with shape ``[N]``.
+        t1_state_indices:
+            Optional contiguous int32 per-packed-row cache slots with shape
+            ``[num_rows]``. T=1 only; null and trailing rows must be zero or
+            non-positive. The tensor is never allocated or populated by this
+            function.
+        state:
+            Paged recurrent state with shape ``[num_slots, H, 128, 128]``.
+            T=1 accepts float32 or bfloat16; T>1 requires float32.
+        output_gate:
+            Gated RMSNorm logits with shape ``[num_rows, H, 128]`` or
+            ``[1, num_rows, H, 128]`` and dtype bfloat16.
+        norm_weight:
+            RMSNorm weight with 128 elements and dtype float32.
+        lower_bound:
+            Negative recurrence-gate lower bound. T=1 also accepts ``None``
+            for the softplus gate; T>1 requires a finite negative value.
+        norm_eps:
+            Non-negative RMSNorm epsilon.
+        output:
+            Optional preallocated contiguous bfloat16 output with shape
+            ``[1, num_rows, H, 128]``.
+
+    Returns:
+        The packed bfloat16 output with shape ``[1, num_rows, H, 128]``.
+    """
+    if state_indices.ndim < 2:
+        raise ValueError("state_indices must have shape [N, T]")
+    if t1_state_indices is not None:
+        if state_indices.ndim != 2 or state_indices.shape[1] != 1:
+            raise ValueError(
+                "t1_state_indices is only supported for T=1; T>1 uses the "
+                "multitoken kernel unchanged"
+            )
+        if (
+            t1_state_indices.ndim != 1
+            or t1_state_indices.shape != (x.shape[0],)
+            or t1_state_indices.dtype != torch.int32
+            or not t1_state_indices.is_contiguous()
+            or t1_state_indices.device != x.device
+        ):
+            raise ValueError(
+                "t1_state_indices must be contiguous int32 [num_rows] on x.device"
+            )
+        if _run_fused_kda_decode is None:
+            raise NotImplementedError("fused KDA decode backend is unavailable")
+        _validate_packed_t1_remap_inputs(x, query_start_loc, state_indices)
+        return _run_fused_kda_decode(
+            x=x,
+            weight=weight,
+            conv_state=conv_state,
+            raw_gate=raw_gate,
+            raw_beta=raw_beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            state_indices=t1_state_indices,
+            state=state,
+            output_gate=output_gate,
+            norm_weight=norm_weight,
+            lower_bound=lower_bound,
+            norm_eps=norm_eps,
+            output=output,
+        )
+    if state_indices.ndim == 2 and state_indices.shape[1] == 1:
+        if _run_fused_kda_decode is None:
+            raise NotImplementedError("fused KDA decode backend is unavailable")
+        _validate_packed_t1_remap_inputs(x, query_start_loc, state_indices)
+        return _run_fused_kda_decode(
+            x=x,
+            weight=weight,
+            conv_state=conv_state,
+            raw_gate=raw_gate,
+            raw_beta=raw_beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            # T=1 does not need a remap buffer.  The direct ABI argument is
+            # unused by the packed constexpr specialization; pass the first
+            # packed destination as the compatible [N] view.
+            state_indices=state_indices[:, 0],
+            query_start_loc=query_start_loc,
+            packed_state_indices=state_indices,
+            state=state,
+            output_gate=output_gate,
+            norm_weight=norm_weight,
+            lower_bound=lower_bound,
+            norm_eps=norm_eps,
+            output=output,
+        )
+    if _run_packed_fused_kda_decode is None:
+        raise NotImplementedError("packed fused KDA decode backend is unavailable")
+    return _run_packed_fused_kda_decode(
+        x=x,
+        weight=weight,
+        conv_state=conv_state,
+        raw_gate=raw_gate,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        state_indices=state_indices,
+        query_start_loc=query_start_loc,
+        num_accepted_tokens=num_accepted_tokens,
+        state=state,
+        output_gate=output_gate,
+        norm_weight=norm_weight,
+        lower_bound=lower_bound,
+        norm_eps=norm_eps,
+        output=output,
+    )
+
+
 # Cached trivial metadata for the frozen mode's uniform batched form: the
 # arange cu_seqlens / slot indices are content-stable per (device, B, T), so
 # building them once keeps the hot path allocation- and launch-free (required

flashinfer/kda_kernels/__init__.py

@@ -26,6 +26,7 @@
 - run_recurrent_kda: Recurrent KDA standard decode and speculative decode backend
 - run_fused_kda_decode: Fused Kimi K3 conv, recurrent KDA, and RMSNorm backend
 - run_packed_kda_decode: Packed Kimi K3 T=1 recurrent decode backend
+- run_fused_kda_decode_multitoken: Packed T>=1 fused KDA backend
 
 The ``*_kda_prefill_sm120`` names below are the optional facade that
 ``flashinfer.kda_prefill`` dispatches through; they are not part of this
@@ -55,6 +56,14 @@
 # ``run_recurrent_kda`` in ``recurrent_kda.py``); import it by module path
 # only for tests and benchmarks.
 
+try:
+    from .fused_kda_decode_multitoken import run_fused_kda_decode_multitoken
+
+    fused_kda_decode_multitoken = run_fused_kda_decode_multitoken
+except (ImportError, RuntimeError):
+    run_fused_kda_decode_multitoken = None  # type: ignore
+    fused_kda_decode_multitoken = None  # type: ignore
+
 try:
     if _torch.cuda.is_available():
         from ..cute_dsl.availability import is_cute_dsl_arch_supported as _dsl_arch_ok
@@ -106,9 +115,11 @@
 
 __all__ = [
     "fused_kda_decode",
+    "fused_kda_decode_multitoken",
     "packed_kda_decode",
     "recurrent_kda",
     "run_fused_kda_decode",
+    "run_fused_kda_decode_multitoken",
     "run_packed_kda_decode",
     "run_recurrent_kda",
 ]