.. _apikda_prefill:

flashinfer.kda_prefill
======================

Optimized recurrent Kimi Delta Attention (KDA) prefill support. The
:func:`flashinfer.kda.recurrent_kda` facade dispatches a strict ordinary
multi-token prefill subset to frozen FlashKDA-compatible SM100-family kernels.

.. currentmodule:: flashinfer.kda_prefill

.. autosummary::
    :toctree: ../generated

    RecurrentKDAPrefillWorkspace

Optimized SM100-family prefill subset
-------------------------------------

``flashinfer.kda.recurrent_kda`` uses the frozen prefill backend only when
every condition below holds:

* the device has compute capability 10.0 (SM100a; B200/GB200) or 10.3
  (SM103a; B300/GB300);
* input is ordinary multi-token prefill: fixed ``T > 1``, or packed input
  whose total token count is greater than its number of sequences;
* Q, K, V, and G are contiguous BF16 ``[B,T,H,128]`` tensors with one shared
  head count, and beta is contiguous BF16 ``[B,T,H]``;
* ``A_log`` is contiguous FP32 ``[H]`` and ``dt_bias`` is contiguous FP32
  ``[H,128]`` or flattened ``[H*128]``;
* ``use_qk_l2norm_in_kernel=True``, ``use_gate_in_kernel=True``,
  ``beta_is_logit=True``, and ``lower_bound`` is a finite negative value;
* speculative decode, GQA, state indices, committed-state sources, and
  accepted-token/checkpoint features are not enabled.

Calls outside that subset retain the existing CuTe-DSL path. In particular,
T=1 decode and speculative decode are not rerouted.

With CUDA 12.9 or newer, JIT and AOT use one ``sm_100f`` URI, JIT
specification, and cubin target per M64/M128 variant for both CC 10.0 and CC
10.3. FlashInfer can still place a built copy in a compilation-context-specific
workspace; the shared logical target does not promise one physical cache path
across device contexts. The binding accepts only the two validated family
members. CUDA 12.8 predates ``sm_100f``, so B200 retains one exact ``sm_100a``
module per variant. CC 10.3 requires CUDA 12.9 or newer.

Fixed input omits ``cu_seqlens``. Packed input has ``B=1`` and accepts a
contiguous CUDA int32 or int64 ``cu_seqlens``. The frozen binding consumes
int64 offsets; pass int64 directly for CUDA graph capture to avoid an
in-capture conversion allocation. Offset values are a caller contract:
``cu_seqlens[0] == 0``, entries are strictly increasing (every sequence is
non-empty), and ``cu_seqlens[-1] == total_tokens``. FlashInfer does not
synchronize the device to inspect these values; invalid offsets may cause
out-of-bounds device access.

Packed scheduling
-----------------

Packed prefill optionally accepts ``seq_order``, a contiguous CUDA int32
tensor with one entry per sequence. It is a caller contract that this tensor
is a permutation of ``[0, N)``. Ordering sequences by decreasing length
reduces the final partial wave. FlashInfer validates dtype, device, rank, and
size without synchronizing the device to inspect permutation values.

When ``seq_order=None``, a cached identity order is used. Fixed ``B=1,H=64``
selects the two-CTA M64 value-split kernel; every packed input and every other
head count selects M128.

State and graph semantics
-------------------------

The BF16 state layout remains ``[N,H,V,K]`` and an explicitly supplied
``initial_state`` is still updated in place, even when
``output_final_state=False``. The frozen kernels load each CTA's disjoint
state rows before writing the final rows back to the same storage, so no
separate state scratch or copy-back is required. If no initial state is
supplied, a final state is allocated only when ``output_final_state=True``.

The frozen kernel uses restricted output storage. A preallocated ``output``
must not overlap Q, K, V, G, beta, or ``initial_state``.

Eager calls without ``prefill_workspace`` use an internal serialized workspace
for the current CUDA stream. This default workspace is eager-only and cannot
be used during CUDA graph capture.

CUDA graph capture requires a caller-owned
``RecurrentKDAPrefillWorkspace(device)`` and a preallocated ``output``. The
workspace owns optional final-state scratch for calls without an initial
state, beta padding, and separate 768-byte M64 and M128 TMA descriptor blocks.
It binds to the device and CUDA stream of its first ``recurrent_kda`` call.
Warm it eagerly on the intended capture stream with the exact Q, K, V, G,
beta, and output tensors, then synchronize that stream before capture. Packed
graphs must also pass preallocated int64 ``cu_seqlens`` and int32
``seq_order``. The warm call prepares descriptors; capture accepts only the
exact warmed pointer, shape, stride, and dtype signature and performs no
descriptor preparation.

The workspace must outlive its graph and every replay. Use one distinct
workspace for each captured ``recurrent_kda`` invocation, including two KDA
invocations in the same graph. Once a workspace participates in capture, any
later Python use through ``recurrent_kda``—eager or another capture—is
rejected. ``graph.replay()`` does not re-enter Python and remains valid.
Sequential replay launches may be issued while a different PyTorch stream is
current, but the caller must provide normal stream ordering. The Python stream
binding applies to eager warmup and capture calls, which must use the same
stream.

When an explicit workspace is used with ``initial_state=None`` and
``output_final_state=True``, the returned final state is workspace-owned
stable scratch. Otherwise an explicitly supplied ``initial_state`` is updated
directly in place by the frozen kernel. Head counts that are not divisible by
eight capture the beta copy into workspace-owned storage padded to the next
eight-head boundary before the frozen launch. The public beta and state shapes
keep the caller's original head count.
