flashinfer.kda_prefill

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

RecurrentKDAPrefillWorkspace(device)

Caller-owned storage required for recurrent-KDA CUDA graph capture.

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. The small-head H < 8 path captures the beta copy into workspace-owned padded storage before the frozen launch.