feat(kda): add fused speculative decode kernel

docs/api/kda_decode.rst

@@ -10,6 +10,11 @@ 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``. ``fused_kda_decode_packed`` extends the same
+pipeline to packed ragged speculative decode, retaining one recurrent
+checkpoint per token and an extended rolling convolution window.
+
 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 +130,6 @@ small ``B * HV * T``. Requires SM90+ for the WY path and ``K = V = 128``;
     :toctree: ../generated
 
     fused_kda_decode
+    fused_kda_decode_packed
     packed_kda_decode
     recurrent_kda

flashinfer/__init__.py

@@ -213,6 +213,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 fused_kda_decode_packed as fused_kda_decode_packed
 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 (
+    fused_kda_decode_packed_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_fused_kda_decode_packed,
+    )
+
+    _FUSED_KDA_DECODE_PACKED_AVAILABLE = True
+except (ImportError, RuntimeError):
+    _run_fused_kda_decode_packed = None
+    _FUSED_KDA_DECODE_PACKED_AVAILABLE = False
 from .kda_kernels import run_recurrent_kda as _run_recurrent_kda
 
 try:
@@ -476,6 +487,151 @@ def fused_kda_decode(
     )
 
 
+@flashinfer_api(trace=fused_kda_decode_packed_trace)
+def fused_kda_decode_packed(
+    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,
+    query_start_loc: torch.Tensor,
+    num_accepted_tokens: 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,
+) -> 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]``.
+
+    T=1 dispatches to :func:`fused_kda_decode`, including its bfloat16-state and
+    softplus-gate support. T>1 uses the experimental SM10x packed backend,
+    which currently 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]``. Active recurrent destinations must be positive cache
+    slots and must not alias destinations from another active sequence;
+    non-positive or zero-length rows are padding 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]``.
+        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 and state_indices.shape[1] == 1:
+        if _run_fused_kda_decode is None:
+            raise NotImplementedError("fused KDA decode backend is unavailable")
+        legacy_state_indices = state_indices[:, 0]
+        if x.shape[0] != legacy_state_indices.shape[0]:
+            packed_rows = torch.arange(
+                x.shape[0], dtype=query_start_loc.dtype, device=x.device
+            )
+            sequence_indices = torch.searchsorted(
+                query_start_loc[1:], packed_rows, right=True
+            )
+            legacy_state_indices = legacy_state_indices.index_select(
+                0, sequence_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=legacy_state_indices,
+            state=state,
+            output_gate=output_gate,
+            norm_weight=norm_weight,
+            lower_bound=lower_bound,
+            norm_eps=norm_eps,
+            output=output,
+        )
+    if _run_fused_kda_decode_packed is None:
+        raise NotImplementedError("packed fused KDA decode backend is unavailable")
+    return _run_fused_kda_decode_packed(
+        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",
 ]