Kimi Delta Attention (KDA)

Recurrent Kimi Delta Attention, a per-K-dimension gating variant of GDN. flashinfer.recurrent_kda() is the canonical entry point for every phase: it classifies each call and serves prefill or decode accordingly. The two modules below hold the phase-specific surface it dispatches into, and are documented here rather than on pages of their own so that one call’s whole route reads in one place.

flashinfer.kda

Phase-neutral recurrent Kimi Delta Attention (KDA) facade, and the canonical entry point for every phase. The public recurrent_kda classifies the call and serves decode and speculative decode through the same shared decode dispatcher that flashinfer.kda_decode uses, while dispatching eligible ordinary multi-token prefill to the optimized backend described in flashinfer.kda_prefill. The automatic prefill selection depends on the device and workload: eligible SM100a and SM103a calls with small logical batch-times-head count use the bundled small-BH CuTe-DSL kernel; other calls follow the standard SM100-family CuTe-DSL/Cake policy. SM120a uses a CuTe-DSL backend of its own. These prefill routes leave decode and speculative-decode dispatch unchanged.

recurrent_kda(q, k, v, g, beta[, A_log, ...])

Recurrent KDA (Kimi Delta Attention) decode and prefill kernel.

Static PTX prefill on B300

flashinfer.recurrent_kda(..., backend="ptx") explicitly selects four retained PTX programs from NVlabs/kda (260927-kda-for-kda). The programs target SM103a (B300) and use persistent recurrence scheduling with exact FP32 state handoff between pieces. The default backend selection is unchanged.

Install nvidia-cuda-nvcc>=13.2 and apache-tvm-ffi>=0.1.12 in addition to FlashInfer’s regular dependencies. The retained PTX uses ISA 9.2. It is assembled once with ptxas >= 13.2 and embedded as a cubin, allowing execution with a CUDA 13 driver that cannot JIT PTX 9.2 directly. A C++17 host compiler and CUDA driver headers (provided by Triton) are needed for the FFI shims. FLASHINFER_KDA_PTXAS selects the assembler; FLASHINFER_KDA_PTX_CACHE_DIR overrides the default FlashInfer JIT cubin cache directory. No build output is written into the installed package. This backend is JIT-only and is unavailable with FLASHINFER_DISABLE_JIT=1 in a fresh process.

Supported inputs are contiguous CUDA BF16 q/k/v/g [1, T, H, 128] with T >= 32 and H equal to 64 or 96, BF16 beta logits [1, T, H], FP32 A_log [H] and dt_bias [H, 128] (or flattened). All tensors must share a device. Q/K L2 normalization, the gate -5 * sigmoid(exp(A_log) * (g + dt_bias)), and beta sigmoid are fused. The normalization denominator is sqrt(sum(x*x) + 1e-6).

Use either one fixed sequence or strictly increasing int32/int64 packed cu_seqlens [N+1] starting at zero and ending at T. Each sequence must have 1 through 16384 tokens. State is FP32 V-first [N, H, 128, 128]; when provided, it is updated in place even if output_final_state=False. Omitted initial state means zero initialization. Output is BF16. The retained MMA kernels require finite initial state with absolute values <= 4096; callers must preserve this value constraint when changing graph inputs. Unsupported hardware, shapes, gate modes, strided inputs, indexed state pools, speculative decode and checkpoints are rejected explicitly.

Planning reads packed offsets and validates the state range on the host. For CUDA Graph capture, warm an explicit RecurrentKDAPrefillWorkspace with the exact tensors, preallocate output, and capture on the same stream. Retain the workspace for the graph lifetime. Packed offsets must remain unchanged during capture/replay; activations and state values may change within the documented contract. A captured workspace cannot be replanned or used for another capture. Separate workspaces isolate streams.

workspace = flashinfer.RecurrentKDAPrefillWorkspace(device=q.device)
out = torch.empty_like(v)
# Warm this call before capturing it on the same CUDA stream.
out, state = flashinfer.recurrent_kda(
    q, k, v, g, beta, A_log=A_log, dt_bias=dt_bias,
    initial_state=state, output_final_state=True,
    use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True,
    beta_is_logit=True, lower_bound=-5.0, cu_seqlens=cu_seqlens,
    output=out, prefill_workspace=workspace, backend="ptx",
)

The reproducible INT21 benchmark compares the complete captured public call with raw MoonshotAI/FlashKDA, including final-state writeback. Both use evolving FP32 state, reset before each trial. Timing is cold-L2 CUPTI GPU activity; planning, allocation, compilation and capture are excluded. The JSON report includes per-shape latencies, correctness, trial medians and baseline identity.

python benchmarks/bench_recurrent_kda_ptx.py \
    --flash-kda-source-dir /path/to/FlashKDA --output ptx-results.json
pytest tests/kda/test_recurrent_kda_ptx.py

flashinfer.kda_decode

Key-Driven Attention (KDA) decode API. The CuTe-DSL kernel lives under flashinfer.kda_kernels.

Deprecated since version 0.8: flashinfer.kda_decode.recurrent_kda is deprecated and now a shim over the decode dispatcher shared with flashinfer.recurrent_kda(), which is the canonical entry point: it serves this same decode contract, accepts a superset of what this one accepts, and defaults to backend="auto". Calling it emits a DeprecationWarning. fused_kda_decode and packed_kda_decode are unaffected and remain here.

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:

  • T=3 with raw gates, use_gate_in_kernel=True, a negative lower_bound, float32 A_log and dt_bias, H=HV=16, and N in {1, 2, 4, 8, 16};

  • four value-row splits for each T in {1, 2, 4, 5, 6} with precomputed gates, use_gate_in_kernel=False, and no A_log, dt_bias, or lower_bound;

  • two additional one-warp direct-state T=1 schedules with value-row splits 16 and 8. T=1 keeps the standard decode API and is normalized to the packed frozen ABI with zero-copy views and cached identity metadata; explicit T=1 cu_seqlens metadata is outside the Cake contract.

  • two Kimi-Linear T=1 equal-head direct-state schedules for lower_bound=None. They accept any positive runtime head count, including production per-rank H=HV=32/16/8/4 for TP1/2/4/8, and evaluate -exp(A_log) * softplus(g + dt_bias) and beta sigmoid in-kernel. Q, K, and V may be zero-copy views into one padded packed-projection row; raw G and beta may have independent positive token-row strides. The single frozen kernel therefore consumes SGLang’s production views without staging copies.

Let W=N*HV be the active sequence/value-head work and S the device SM count. SM100a retains the B200-measured policy: direct split 16 for T1 when W<=2S and direct split 8 otherwise; split 4 for T2; split 2 for T4; and the T5/T6 CTA-wave policy of split 8 for W<=3S/8, split 2 for 3S/8<W<=S/2, split 4 for S/2<W<=3S/4, split 2 for 3S/4<W<=3S/2, and split 1 above that range.

SM103a uses its separately measured GB300 policy. T1 selects direct split 16 through a conservative W<=32S extrapolation guard (measured through W/S=26.95), and direct split 8 beyond it. T2 selects split 8 through W<=S/2 and split 4 above it. T4 selects split 8 through W<=S/2, split 4 through W<=S, split 2 through W<=3S/2, split 1 through W<=2S, and split 2 above it. T5 keeps the SM100a CTA-wave policy except for a measured split-1 island at 3S/4<W<=S. T6 selects split 8 through W<=3S/8, split 2 through W<=S/2, and split 1 above it. T3 uses its sole exact lower-bound split-4 specialization on both architectures.

With CUDA 12.9 or newer, JIT and AOT compile all 25 checked-in bodies once for the sm_100f family target. The family module URI and cubin artifact can run on both CC 10.0 and CC 10.3; build workspaces may still materialize separate cache directories for their local architecture context. Runtime split selection remains device-specific. A cold-L2 CUPTI A/B against exact-target cubins measured no aggregate change on B200 (1.0000x exact/family) and 0.9987x on GB300. The GB300 direct-T1 path was the repeatable exception (0.9790x), so its two public direct variants retain exact sm_103a cubins while every other GB300 route uses sm_100f.

CUDA 12.8 cannot compile sm_100f. On B200 it therefore retains exact sm_100a modules for all 25 bodies. SM103a requires CUDA 12.9 or newer. Every binding validates its family or exact-device contract before launch, and the frozen generated body bytes are identical across all physical targets.

Once backend="cake" is selected, every supported call launches exactly one exported Cake kernel. An unsupported architecture, shape, gate mode, layout, aliasing pattern, or optional feature raises an error; it never falls back to CuTe-DSL. The default backend="cute-dsl" preserves the existing FlashInfer implementation. backend="auto" selects Cake only for the equal-head D128 T1 unbounded-softplus contract and preserves CuTe-DSL for other decode modes.

Serving-native packed Kimi K3 decode

packed_kda_decode is a separate serving adapter for the Kimi K3 T=1, H=12, K=V=128 contract. It consumes post-convolution packed QKV plus raw gate and beta tensors, then fuses Q/K extraction and L2 normalization, the fixed lower_bound=-5 gate transform, beta sigmoid, and the recurrent update into one exported Cake kernel. It is distinct from recurrent_kda (whose inputs are already split into Q, K, V, gate, and beta tensors) and fused_kda_decode (which also performs the convolution and gated RMSNorm).

The operator updates a caller-owned bfloat16 state pool in place. A contiguous int32 state_indices tensor selects one unique active slot per batch row; -1 is an inactive CUDA-graph padding row that emits zero without touching state. By default the operator allocates a contiguous bfloat16 output of shape [B, 1, 12, 128]. Supplying a caller-owned output with that exact layout makes replay allocation-free. All work runs on the caller’s current PyTorch CUDA stream.

Two frozen schedules are selected from host-visible batch size only: the eight-row value tile for B < 32 and the sixteen-row value tile for B >= 32. CUDA 12.8 uses a legacy exact sm_100a module on CC 10.0; CUDA 12.9 or newer uses one sm_100f module on CC 10.0 and CC 10.3. Unsupported devices or contracts raise an error without falling back to another KDA implementation.

Frozen / speculative-verify mode

recurrent_kda(disable_state_update=True) computes the KDA outputs for up to 16 tokens per sequence from a read-only committed-state pool and never writes state back — the speculative-decode verify path, mirroring GDN’s gated_delta_rule_mtp(disable_state_update=True). Optional slot-indexed correction_cache (float32 per-token delta-rule corrections, [num_slots, HV, T_max, V]) and kg_cache (raw key | raw gate, [num_slots, HV, T_max, 2K]) out-params feed a downstream commit/recovery kernel, paralleling GDN’s slot-indexed intermediate_states_buffer. The mode accepts the batched [B, T, ...] form and the packed cu_seqlens form with ragged per-sequence lengths and null slots.

Two implementations are dispatched internally by problem size: a WY-parallel tensor-core kernel that replaces the serial recurrence with T x T GEMMs, a log-depth triangular inverse, and TMA-loaded state GEMMs (per-K-channel decay folded into k * exp(±cumsum g) / q * exp(cumsum g) operand tiles), and a grouped register recurrence (no per-token state-checkpoint writes) for small B * HV * T. Requires SM90+ for the WY path and K = V = 128; backend="cake" raises in this mode.

fused_kda_decode(x, weight, conv_state, ...)

Run the fused Kimi KDA decode pipeline.

packed_fused_kda_decode(x, weight, ...[, ...])

Run packed T>=1 fused KDA with per-token cache checkpoints.

packed_kda_decode(mixed_qkv, raw_gate, ...)

Run serving-native packed Kimi K3 recurrent decode.

recurrent_kda(q, k, v, g, beta[, A_log, ...])

Recurrent KDA (Kimi Delta Attention) decode kernel.

flashinfer.kda_prefill

Optimized recurrent Kimi Delta Attention (KDA) prefill support. The flashinfer.kda.recurrent_kda() facade exposes frozen Cake and source-level CuTe DSL implementations for a strict ordinary multi-token prefill subset.

RecurrentKDAPrefillWorkspace(device)

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

Note

flashinfer.RecurrentKDAPrefillWrapper is experimental, so it has no generated reference page here until it graduates. Calling plan or run is itself the opt-in and needs no environment variable; each warns once per process. It is limited to compute capability 10.0 and 10.3, is specific to the CuTe DSL backend, and fixes the sequence count and packed token extent after the first warmup run. Its planning implementation lives in flashinfer.experimental.kda_prefill_wrapper; it contains no kernels of its own, and the kernels it dispatches to are the stable AOT-registered ones, alongside a package README. examples/experimental/kda_prefill_wrapper.py is a runnable plan-and-run example. See #5069 for the graduation plan.

Backend selection

backend="auto" selects the source-level CuTe DSL backend for eligible ordinary multi-token prefill and falls back to the frozen Cake backend for unsupported contracts. Decode retains the existing KDA decode routing. backend="cake" and backend="cute-dsl" select a backend strictly and raise when its contract is unsupported.

For multi-token prefill, backend="cute-dsl" selects a BT=16 CuTe DSL kernel. It supports contiguous BF16 Q, K, V, G, and beta with one shared head count and head dimension 128, the in-kernel lower-bound gate, fixed or packed-varlen layout, BF16 recurrent state, explicit seq_order, and the same checkpoint contract as Cake. checkpoint_cu_starts must always be int64. Packed cu_seqlens must be int64 during CUDA graph capture. The CuTe DSL schedule is non-persistent.

Optimized Blackwell prefill subset

The strict Cake backend is available 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; beta is BF16 [B,T,H] with unit head stride and non-overlapping token rows (an aligned slice of a fused projection is accepted directly);

  • 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, committed-state sources, and accepted-token features are not enabled. Plain int32 ssm_state_indices and native prefill checkpoints are supported by direct M128.

T=1 decode and speculative decode are not handled by either prefill backend.

CUDA 12.8 predates the family target, so CC 10.0 uses legacy exact sm_100a modules. With CUDA 12.9 or newer, JIT and AOT compile one sm_100f module per schedule for both CC 10.0 and CC 10.3. Cache keys also include the frozen module identity so an older schedule cannot satisfy a refreshed request. Runtime routing remains device-specific. The legacy head-grouped and mixed-length persistent M128 routes remain restricted to measured 148/152-SM CC 10.0 devices. On validated 148/152-SM CC 10.0 or CC 10.3 devices, eligible uniform eager calls with a caller-owned in-place initial and final state may instead split only the recurrence chains that create a partial final device wave. An occupancy and roofline model uses the live SM count plus the frozen schedule’s thread, shared-memory, tensor-memory, BF16, HBM, state-transfer, and refill costs; it selects recurrence pieces only when their dependency-DAG critical path is shorter than direct M128. Intermediate BF16 states use device-scope release/acquire handoffs, and each consumer CTA resets its ready counter before it completes so subsequent eager launches start from the same state. This route is not selected when the caller supplies an explicit workspace or seq_order; CUDA Graph capture therefore continues through the existing non-piece routes.

CC 10.3 also has a tensor-core state-decay specialization for uniform, complete N32 work when there are at least 64 heads, at least 96 sequence/head tasks, and the maximum sequence length is a multiple of 32 and at least 256. Mixed or partial N32 tails retain scalar state decay. On either capability, fixed-layout calls with at most eight total sequence/head tasks, at most eight heads, and at least 2,048 tokens per sequence use the small-BH owner/helper schedule when all eight CTAs per task can reside concurrently. Calls outside that measured region continue through the existing direct or fallback route; it is not a public-input allowlist.

The dense beta-TMA BT16 one-wave route uses the S9 chain schedule when both value-split CTAs for every task fit in one device wave.

At maximum sequence length 16 or below, generic head counts use a one-stage N16 retrace with one four-warp prepare owner. It preserves the variable-shape N16 arithmetic while reducing the CTA from 32 to 16 warps. H12 keeps its dedicated scalar-beta N16 schedule.

The N16 route describes aligned beta storage directly once it contains a full 16-token tile. Calls requiring token or head padding refresh the stable beta-TMA workspace inside the binding before launching the recurrence kernel.

The frozen H12 N16 schedule’s residual recurrence rounds four intermediates through BF16: the state/K prediction, the V-minus-prediction delta, sigmoid beta, and the post-beta update carrier. The final-state contraction starts from a zero accumulator; the old BF16 state is multiplied by the total decay and explicitly added to that product in FP32 before the 16-token chunk boundary rounds back to BF16. The N16 prepare carrier also matches the source-visible BF16 arithmetic graph: normalized Q/K, positive and inverse prefix decay, and every chained Qd/Kd/Ki/Kr multiplication round through BF16 at their respective boundaries. The N32 schedule comes from the same generated export.

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 non-decreasing, and cu_seqlens[-1] == total_tokens. CuTe DSL accepts equal adjacent offsets for zero-length sequences; Cake requires every sequence to be non-empty. 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.

For Cake, omitting seq_order uses its cached eager scheduling metadata. H12 selects the dedicated M128 schedule with a 16-token recurrence chunk for both fixed and packed layouts. Fixed B=1,H=64 selects the two-CTA M64 value-split kernel; the fixed small-BH region described above selects its eight-CTA owner/helper schedule. Eligible medium and long shapes instead use a BT16 prepare/chain route: dense fixed B=1,H=60..64 inputs qualify from 4,096 tokens when two value-split CTAs per head fit on the device; general M128 shapes qualify from 65,536 tokens for one to eight sequence/head tasks, or from 4,096 tokens for nine to 32 tasks when two CTAs per task fit. N16 alternatives additionally depend on SM count, chain waves, and sequence length. Uniform work may use the recurrence-piece route described above when its modeled critical path wins. Supplying seq_order disables persistent host task-bin planning but does not suppress BT16 or otherwise force direct M128. Remaining eligible inputs select the shape-appropriate non-persistent or general M128 schedule. The scalar-prepare/S8 BT16 pair is submitted by one native Cake binding, which performs both launch plans before enqueueing either kernel so the dependent launches do not expose a Python/FFI inter-kernel gap. CUDA Graph capture still records the same two kernels and preserves the workspace contract below.

For packed CuTe DSL engine calls, omitting seq_order generates a stable decreasing-length order on the device. CuTe DSL decomp retains the original sequence order in the eager path because its CTA grid fits in one wave. flashinfer.RecurrentKDAPrefillWrapper (experimental) provides fixed-address metadata for CUDA Graph capture: plan only copies packed offsets into its device buffer, while a captured GPU prepass generates the order and decomp cu_chunks prefix before the recurrent kernels run. The decomp prep kernel binary-searches this compact prefix instead of carrying a dense chunk-to-sequence tensor. Its workspace and launch use a graph-static chunk capacity derived from tensor shapes, while the GPU prefix supplies the actual chunk count on each replay.

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.

With ssm_state_indices, initial_state is a caller-owned pool [N_pool,H,V,K]. Sequence i loads and updates the named pool row in place. The pool may have padding between first-dimension slots, but each [H,V,K] slot must be contiguous and both the pool base and slot pitch must be 16-byte aligned. Slot ids must be unique and in range; this is a caller contract so the launch path does not synchronize to inspect them.

Native checkpoints use a preallocated BF16 state_checkpoints[C,H,V,K], int64 checkpoint_cu_starts[N+1], and a positive checkpoint_every_n_tokens divisible by 16. KDA checkpoints are states before each interval: every non-empty sequence contributes its initial state as row zero, followed by states after one, two, … intervals that strictly precede its end. Consequently each sequence contributes ceil(seq_len / interval) rows. The call returns (output, final_state, state_checkpoints) when enabled. Intervals and cumulative counts are caller-provided device metadata and are not value-scanned at launch.

An aligned beta base address and a token pitch divisible by 16 bytes are sent directly to TMA. Other eligible row-strided beta views remain valid API inputs, but the binding refreshes reusable padded workspace internally; callers never need to materialize beta.contiguous().

The frozen kernel uses restricted output and auxiliary storage. A preallocated output and checkpoint buffers must not overlap Q, K, V, G, beta, state, metadata, or descriptor storage.

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, separate TMA descriptor blocks, and the small-BH compact packet ring with its generation counters. BT16 schedules additionally own cu_chunks and chunk-to-sequence metadata, BF16 Qd/Kd/W/QK factors, FP32 diagonal factors, and independent prepare/chain descriptor storage. The workspace 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. Warm the largest intended small-BH shape before capturing so its packet-ring storage is already allocated.

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.

SM120a prefill subset

Compute capability 12.0 devices select their own prefill backend, flashinfer.kda_kernels.sm120_prefill. It is CuTe DSL like the BT=16 backend above and shares no code and no device with it, nor with Cake: those two are CC 10.0 and 10.3 and this one is CC 12.0, so at most one of the three can be eligible for any call, and adding this one cannot change which kernel an SM100 or SM103 call receives. Nothing about the public API changes — the entry point is still flashinfer.kda.recurrent_kda(), and no argument names the architecture.

backend selects an implementation family, not an architecture name. On a CC 12.0 device both "auto" and "cute-dsl" may reach this backend; the dispatcher tries it before the SM100-family CuTe DSL prefill path. An explicit "cake" request never probes or runs SM120. If the Cake prefill predicate does not support that ordinary multi-token prefill call, the request is refused rather than silently executed by another backend.

recurrent_kda uses it only when every condition below holds:

  • the device has compute capability 12.0, and the installed CuTe DSL and CUDA toolkit can natively target sm_120a. A family-conditional fallback target is refused rather than accepted, because the kernels are written against architecture-specific instructions;

  • input is ordinary multi-token prefill: fixed T > 1, or packed input whose total token count exceeds its number of sequences;

  • Q, K, V and G are contiguous BF16 [B,T,H,128] tensors sharing one head count, and beta is contiguous BF16 [B,T,H]. GQA and V != K are not supported;

  • the output fits an INT32 extent: T_total * H * 128 <= 2**31 - 1, which is 16383 tokens at H=1024 and no constraint at ordinary head counts. Larger is refused with the backend’s own error, because the two things that stop there — a device index built in INT32, and the DSL packing a memref extent as one — otherwise fail as a silent negative offset and as a compile-time overflow naming no tensor;

  • 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 in [-5.0, 0.0). The bound exists because the safe gate’s worst-case chunk prefix reaches a reciprocal approximation’s cliff at about -5.4585;

  • initial_state, if given, is a contiguous BF16 [N,H,128,128] tensor. A state pool with ssm_state_indices is not supported;

  • output, if given, is contiguous BF16 with V’s shape and does not overlap any input in GMEM;

  • speculative decode, seq_order, prefill checkpoints, committed-state sources and FP32 gate or state are not enabled.

Under backend="auto", calls outside that subset continue through the existing dispatcher. An explicit backend="cute-dsl" ordinary multi-token prefill request is refused when neither CuTe DSL prefill implementation is eligible. T=1 decode and speculative decode are not rerouted by this backend.

Two variants, chosen per shape

The backend has two implementations of the same contract. decomp runs a chunk-parallel prepare and a serial recurrence, issued through one compiled host entry; fused does both in a single kernel. They agree numerically, so the choice between them is a performance one and is made from a measured table keyed on the device’s SM count — not on its name, which is not a stable unique selector.

Thresholds exist for the 110-SM, 156-SM and 188-SM parts, measured on each. They do not agree: the 110-SM part switches to the fused kernel at CTA 128 and the other two at 144, because at CTA 128 with short sequences the larger parts still prefer the decomposed kernel and the smallest one does not. Any other CC 12.0 device uses the 156-SM thresholds as a labelled fallback. A benchmark run reports which case applies, so a number taken on an unprofiled card cannot be read as tuned.

State and graph semantics

State semantics match the SM100-family path: a supplied initial_state is updated in place whether or not output_final_state is set, and the second return value is None when it is not set. Without an initial state, a BF16 final state is allocated only when output_final_state=True.

CUDA graph capture requires a caller-owned RecurrentKDAPrefillWorkspace and a preallocated output, warmed eagerly on the capture stream with the exact tensors and then synchronized before capture. The warm call is where every compile, descriptor build, metadata table and allocation happens; capture performs none of them and a cold capture is refused rather than silently degraded. Both offsets dtypes are accepted for packed capture; eager warmup populates a workspace-owned canonical int32 buffer. The offset values must stay fixed for the graph’s lifetime. Changing them requires a fresh eager warmup and capture. Q, K, V, G, beta and state contents may change freely at unchanged addresses.

The offsets contract is not only a capture one. Validating cu_seqlens needs a device-to-host read, so what is derived from it — the sequence lengths, the canonical int32 copy, and the decomposed variant’s chunk tables — is cached against the tensor’s address and version counter rather than read again on every call. Under torch.inference_mode a tensor has no version counter, so refilling an offsets buffer in place with a different segmentation is not detectable and the stale tables are reused, silently computing against the previous sequence boundaries. Use a different tensor for a different segmentation, or call flashinfer.kda_kernels.sm120_prefill.clear_kda_prefill_sm120_caches() after refilling one in place. This applies to eager calls as much as to captured ones.

A workspace binds to one variant, one stream and one call signature on first use, and once it has participated in a capture it cannot be used again.

What the caches hold

A warm call is a memo lookup, and the memo addresses the caller’s buffers: the descriptors carry their base addresses and the flat views wrap them. Those buffers therefore stay allocated for as long as the entry lives, which is what makes reusing the entry safe — an allocator that had recycled the address would otherwise hand the kernel someone else’s memory.

The retention scales with the number of distinct buffer sets a process rotates through, not with the number of calls. On a 110-SM part at [1, 1024, 8, 128] one set holds about 14.5 MiB, and eight rotating sets about 73 MiB. Reuse one set and it stays at one set’s worth forever.

The entry ceilings are not a memory budget and lowering them does not trade speed for memory: below the ceiling the retention is the same whatever the ceiling is, and above it every call rebuilds its plan — about 7.3 ms against a 100 microsecond hit on that part. A deployment that needs the memory back should rotate fewer buffer sets, or call flashinfer.kda_kernels.sm120_prefill.clear_kda_prefill_sm120_caches(), which releases all of it.

Persistent CuTe DSL prefill with FP32 state

recurrent_kda(..., backend="cute-dsl-persistent") selects the persistent M128/BT32 kernel adapted from humanfia/kda-for-kda-release. It fuses Q/K L2 normalization, the bounded K3 gate, beta sigmoid, preparation, triangular solves and recurrence. A continuous five-stage preparation pipeline is shared across each CTA’s sequence/head chains. Recurrent state and inter-CTA state handoffs use FP32; tensor-core operands and output use BF16. Values, states and outputs up to about 2**69 in magnitude stay finite: the BF16 state operand is stored scaled by 2**-57 so that the anchored decay factors cannot overflow the FP32 accumulators.

The work schedule is derived per call from a cost model, independent of specific sequence lengths or SM counts. Whole (sequence, head) chains are packed longest-processing-time first onto at most one CTA per SM. When cutting chains lowers the modelled makespan by at least 5%, the scheduler either cuts short heads off the longest chains or wraps chains across CTAs. Every cut piece hands FP32 state from a producer that leads its CTA to a consumer that ends another one, so no producer waits on a consumer. A sequence’s partial last chunk is loaded as a full tile and masked. The gate-prefix table remains FP32 on all routes. FP32 state storage does not make the arithmetic an FP32-only recurrence: tensor-core operands, residual updates and beta carriers are rounded to BF16. In particular, near-unit decay can retain operand-rounding error across many chunks. The INT21 benchmark uses the source contract: relative L2 error at most 3%, with each element’s error bounded by the larger of 0.5 times reference RMS and 5% of its magnitude. Short-sequence tests additionally compare with an independent FP64 recurrence.

This backend is explicitly selected and requires SM100 or SM103 with nvidia-cutlass-dsl>=4.7.0. Its supported contract is:

  • contiguous, 16-byte-aligned BF16 Q/K/V/G [B,T,H,128] and beta [B,T,H], with T > 1, positive H divisible by eight and B*T*H*128 < 2**31;

  • contiguous FP32 A_log[H] and dt_bias[H*128] or dt_bias[H,128];

  • use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, beta_is_logit=True and -5 <= lower_bound < 0;

  • optional contiguous, 32-byte-aligned FP32 state [N,H,128,128] in value-first order, with N*H*128*128 < 2**31, updated in place even when output_final_state=False;

  • fixed batches or packed B=1 sequences with contiguous CUDA int32/int64 cu_seqlens[N+1]. Offsets start at zero, end at the total token count, and are strictly increasing; empty sequences are rejected;

  • an optional contiguous BF16 output buffer, disjoint from every input.

State pools/indices, checkpoints, external sequence ordering, unbounded gates, GQA, speculative decode and frozen-state mode are outside this backend’s contract. Calls requesting them raise an error. Other backend choices retain their existing behavior.

For CUDA Graphs, supply a RecurrentKDAPrefillWorkspace and preallocated output. Warm the exact tensors on the capture stream, then synchronize before capture. One workspace belongs to one captured call; it cannot subsequently be used by Python eager calls or another capture. The workspace and buffers must outlive the graph. Q/K/V/G, beta, gate parameters and state contents may change between replays, but sequence offsets and scalar arguments must stay fixed. Handoff flags are cleared on the launch stream before every invocation, so captured replays can be interleaved with independent eager calls. Launches use no thread-block clusters. Split schedules rely on the full persistent grid being resident. Serialize independent split-grid invocations; overlapping them on separate CUDA streams is not supported.

In eager mode, offset values are read on the host to prepare the schedule; this also detects edits made under torch.inference_mode. Only the most recent schedule is retained per stream or explicit workspace. A returned final state with no supplied initial state is caller-owned in implicit eager mode and workspace-owned when using an explicit workspace.

The six INT21 workloads use 8192 total tokens, H96/H64 and D128: one sequence, eight 1024-token sequences, or lengths [1300,547,2048,963,271,3063]. The benchmark checks both output and state against FLA and compares the captured public call with a source-identified MoonshotAI/FlashKDA build using cold-L2 CUPTI GPU timings:

python benchmarks/bench_recurrent_kda_persistent.py \
    --flash-kda-source-dir /path/to/FlashKDA --output results.json

Install fla-core==0.5.2 and cupti-python to run the benchmark. Both arms update state in place, resetting it outside each timing trial. JSON results include absolute latencies, every trial median, the baseline revision and extension digest, software versions and the six-shape geometric mean.

TIRx prefill

recurrent_kda(..., backend="tirx") selects the optional TIRx KDA kernels adapted from NVlabs/kda (260927-kda-for-kda). A fixed single sequence uses a BT32 front end and recurrent chain, with concurrent execution when the device has sufficient SMs; when its length is not divisible by 32, the final partial chunk continues on the fused kernel. Other inputs use a fused persistent BT64 kernel with a host-built work list. The work list is chosen per call from a cost model, independent of specific sequence lengths or SM counts: whole (sequence, head) chains packed by longest-processing-time, or equal-cost ranges that hand FP32 state between CTAs when that lowers the modelled makespan by at least 5%. TIRx is imported only when explicitly selecting this backend.

The fused kernel accumulates each chunk’s state update before applying the chunk decay (up to 2**120.5), so its TMEM state is kept at 2**-60 of the true value. For unit-order inputs, scaling V and the initial state by any power of two from 2**-40 to 2**60 scales outputs and final states exactly; the BT32 chain alone covers 2**-60 to 2**100.

Install CUDA-enabled TVM, the TIRx Lite frontend and a CUDA 13 toolkit with nvcc available. The validated compiler packages are:

pip install 'apache-tvm==0.27.0' 'apache-tvm-ffi==0.1.14.post1'
pip install 'tirx-kernels @ git+https://github.com/mlc-ai/tirx-kernels.git@c4b700e7e8c390f069b369b588ecfe20215e5850'

The pinned commit is the v0.1.0 release; tirx_kernels.tirx_lite must be present. CUDA_PATH can select the toolkit root. Compiled modules are stored in FlashInfer’s JIT cache and keyed by architecture, specialization, source and compiler versions. The shared JIT lifecycle provides cross-process locking and honors FLASHINFER_DISABLE_JIT.

The backend supports SM100/SM103 and requires:

  • contiguous, 16-byte-aligned BF16 Q/K/V/G [B,T,H,128] and beta [B,T,H];

  • B >= 1, T > 1, H divisible by eight and H no larger than the device’s SM count, with B*T < 2**21 and B*T*H*128 < 2**31;

  • contiguous FP32 A_log[H] and dt_bias[H*128] or dt_bias[H,128];

  • fused Q/K normalization, gate activation and beta sigmoid, with use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, beta_is_logit=True and lower_bound in [-5.0, 0.0); each lower bound compiles its own gate constant;

  • optional contiguous, 32-byte-aligned FP32 value-first state [N,H,128,128], with N < 65536 and N*H*128*128 < 2**31;

  • fixed batches or packed B=1 sequences with CUDA int32/int64 cu_seqlens[N+1]. Offsets start at zero, end at the total token count and strictly increase;

  • an optional contiguous output buffer disjoint from every input.

A supplied initial state is updated in place even if output_final_state=False. Omitting it starts from zero. A returned final state without a supplied initial state is caller-owned in implicit eager mode and workspace-owned with an explicit workspace. State pools, checkpoints, frozen-state mode, GQA, speculative decode and strided activations are rejected.

Initial and final state, and every state handoff between CTAs, are FP32. Tensor core operands and residuals round to BF16. Numerical validation uses the source INT21 contract: relative L2 error at most 3%, with each error bounded by max(0.5 * RMS(reference), 0.05 * abs(reference)). Short cases also compare against an independent FP64 token recurrence.

Supply RecurrentKDAPrefillWorkspace and a preallocated output for repeated calls with the same buffers. The first call prepares tensor maps, work lists, scratch and kernels. Packed offsets are read on the host in eager mode, including when they were changed under torch.inference_mode. One plan is retained per workspace; changing buffers or offsets replaces it. The fixed route executes its first invocation sequentially to finish CUDA module loading before enabling producer/consumer overlap; state is updated once per invocation throughout.

For CUDA Graph capture, warm the exact tensors and scalar arguments on the capture stream and synchronize first. Each workspace belongs to one stream and one captured call, and must outlive the graph. It cannot be reused eagerly or in another capture afterward. Activation, gate-parameter and state contents may change between replays; offsets and scalar arguments remain fixed. Flags are reset before each invocation. The fixed route joins its private chain stream back to the caller’s stream, so the caller observes completion of both kernels. Independent persistent launches must be serialized: their producer/ consumer schedules rely on the full grid being resident.