FlashInfer PR #4262: Core Optimization Summary¶

Overview¶

FlashInfer PR #4262, “feat(cake_kda): add optimized B200 recurrent prefill backend,” adds a CAKE-generated BF16 recurrent Kimi Delta Attention (KDA) prefill backend specialized for NVIDIA B200 / SM100a.

The optimized path is deliberately narrow: it handles ordinary multi-token prefill with contiguous BF16 tensors, head dimension 128, equal Q/K/V head counts, in-kernel gate and Q/K normalization, and no speculative-decode or state-pool features. Calls outside this contract continue to use the existing CuTe-DSL backend.

The main performance result reported by the PR is a 2.0512× geometric-mean speedup over MoonshotAI/FlashKDA across six fixed and packed B200 workloads, while output and recurrent-state comparisons pass the stated BF16 tolerances.

Executive Summary¶

The speedup comes from combining several optimizations rather than from one instruction:

Fuse chunk preparation, recurrent state updates, and output production inside one M64/M128 CUDA kernel.
Increase the intra-chunk size from 16 to 32 tokens while keeping BF16 exponent scaling safe through a fixed midpoint anchor.
Keep the recurrent state resident as FP32 in Tensor Memory (TMEM) across chunks.
Use tcgen05.mma for the large recurrent GEMMs, while using conventional warp-level mma.sync for the small 32×32 intra-chunk matrices.
Concatenate the state and output update along the GEMM N dimension so one logical M128×N160×K32 update produces both results.
Run five independent chunk-preparation groups ahead of a strictly ordered recurrent consumer through a five-slot shared-memory ring buffer.
Aggressively alias and reuse shared-memory regions so the complete five-stage pipeline fits in 227,328 bytes (222 KiB) per M128 CTA.
Select M64 or M128 physical schedules to balance CTA-level parallelism and per-CTA work.
Use TMA, stable tensor-map storage, exact in-place state aliasing, and graph-safe workspaces to remove avoidable transfers and preserve CUDA Graph compatibility.

1. A Fused Recurrent-Prefill Kernel¶

The numerical prefill pipeline is fused into the generated M64 or M128 kernel. The kernel performs, for every 32-token chunk:

gate transformation and prefix accumulation;
Q/K normalization and exponent scaling;
construction of the causal Mqk matrix;
construction and inversion of the 32×32 lower-triangular system;
recurrent state/output GEMMs;
residual and beta application;
final state and output updates;
output transposition and store.

This avoids the classic two-kernel organization in which a preparation kernel writes Mqk, inverse matrices, and transformed Q/K tensors to a global workspace before a second recurrent kernel consumes them. Intermediate data remains in registers, SMEM, or TMEM instead of making a global-memory round trip.

The descriptor-publication step can still occur outside the numerical kernel when tensor maps must be prepared, but the recurrent prefill computation itself is fused.

Relevant source:

2. Chunk-32 Through Fixed-Anchor Exponent Centering¶

The gate is accumulated in the base-2 logarithmic domain. For one feature dimension, define

\[ b = \text{lower bound}\log_2(e), \qquad G_i = \sum_{t=0}^{i} \lambda_t. \]

For a 32-token chunk, the prefix range is approximately

\[ G_i \in [32b, 0]. \]

Directly storing factors such as \(2^{G_i}\) and \(2^{-G_i}\) would expose a BF16 operand to an exponent magnitude of \(32|b|\), which is too large for the standard lower_bound = -5 configuration.

The kernel introduces the fixed scalar anchor

\[ C = 16b = 16\cdot\text{lower bound}\log_2(e). \]

It prepares centered GEMM operands

\[ Q_c = \text{scale}\,\hat Q\,2^{G-C}, \qquad K_c = \hat K\,2^{G-C}, \qquad K_i = \hat K\,2^{C-G}. \]

The anchor cancels inside the matrix product:

\[ 2^{G_i-C}\,2^{C-G_j}=2^{G_i-G_j}. \]

Therefore, centering changes only the representation of the operands, not the mathematical Mqk result.

For lower_bound = -5,

\[ b\approx-7.2135, \qquad C\approx-115.42. \]

Both centered exponent ranges become approximately [-115.42, +115.42], which is inside the normal BF16 exponent window. This is the same exponent radius as an uncentered 16-token chunk:

\[ \frac{32}{2}|b|=16|b|. \]

The anchor is a scalar derived from lower_bound; it is not the measured prefix value of token 15. The implementation can be seen in the restore-factor and centered-operand code and centered Q/K preparation.

3. Intra-Chunk and Inter-Chunk Decomposition¶

The 32-token preparation path constructs the chunk-local matrices independently of the incoming recurrent state.

One useful high-level form is

\[ L = \mathrm{strictLower}\!\left(\mathrm{diag}(\beta)K_cK_i^T\right), \]

\[ M_{qk}=\mathrm{lowerTri}\!\left(Q_cK_i^T\right), \qquad INV=(I+L)^{-1}. \]

The recurrent-domain operands restore the chunk endpoint \(G_T\):

\[ Q_d=\text{scale}\,\hat Q\,2^G, \qquad K_d=\hat K\,2^G. \]

\[ K_r=\hat K\,2^{G_T-G}, \qquad a=2^{G_T}. \]

The inter-chunk computation can then be written as

\[ S_d=S_{old}\mathrm{diag}(a), \]

\[ O_0=S_{old}Q_d^T, \qquad P=S_{old}K_d^T, \]

\[ R=(V^T-P)\mathrm{diag}(\beta), \qquad U=R\,INV^T. \]

The final update is

\[ S_{new}=S_d+UK_r, \qquad O=O_0+UM_{qk}^T. \]

The important scheduling property is that L, Mqk, INV, and the transformed Q/K data are chunk-local and can be prepared ahead of time. The state update is recurrent and must consume chunks strictly in order.

4. Five-Stage SMEM Ring Buffer¶

The M128 kernel defines five chunk-pipeline stages. Each stage occupies 41,984 bytes, so the ring payload is

\[ 5\times41{,}984=209{,}920\ \text{bytes}=205\ \text{KiB}. \]

Including barriers/control storage and two 8 KiB output buffers, total M128 shared memory is

\[ 1{,}024+5\times41{,}984+2\times8{,}192 =227{,}328\ \text{bytes}=222\ \text{KiB}. \]

The five preparation groups use warps 12–31, four warps per group. Their mapping is

\[ \text{prep instance}=\frac{\text{warp}-12}{4}, \]

\[ \text{chunk index}=5\cdot\text{prep iteration}+\text{prep instance}. \]

Consequently:

Producer	Warps	Chunks	Fixed SMEM slot
P0	12–15	0, 5, 10, …	0
P1	16–19	1, 6, 11, …	1
P2	20–23	2, 7, 12, …	2
P3	24–27	3, 8, 13, …	3
P4	28–31	4, 9, 14, …	4

The recurrent consumer processes chunks in order 0, 1, 2, 3, 4, 5, …. After consuming chunk c, it releases slot c mod 5, allowing the corresponding producer to refill it with chunk c + 5.

Five stages are not five copies of the recurrent state. They buffer only future chunk data. The ring hides the relatively long preparation latency and absorbs producer/consumer speed variation, while mbarriers provide backpressure in both directions.

The constants and role mapping are visible in the M128 memory definitions and five-group preparation loop.

5. Aggressive SMEM Aliasing and Lifetime Reuse¶

Each 41 KiB stage is reused across multiple logical tensors whose lifetimes do not overlap. Major aliases include:

raw G storage → decayed/centered Q storage;
raw K storage → decayed/centered K storage;
raw Q storage → inverse-scaled K → restored K;
gate-prefix workspace → Mqk^T, compact inverse data, restore factors, and V storage;
inverse workspace → V after the inverse is no longer needed.

The 12 KiB final_trans region contains

\[ [K_r\;|\;M_{qk}^T\;|\;padding], \]

with physical BF16 storage equivalent to 32×192 elements:

K_r: 32×128, 8 KiB;
Mqk^T: 32×32, 2 KiB;
padding/layout space: 32×32, 2 KiB.

The final MMA consumes only the logical first 160 columns. The extra 32 columns are reserved physical padding/layout space and are not part of the mathematical GEMM.

This lifetime-driven reuse is what makes a five-stage pipeline possible within the B200 per-CTA SMEM budget.

6. TMEM-Resident FP32 State¶

The M128 CTA allocates 256 TMEM columns. The key logical regions are:

TMEM columns	Main use
0–63	Packed BF16 state input; later reused by the `U` accumulator
64–191	FP32 recurrent state, persistent across chunks
192–223	FP32 chunk output `[128,32]`
224–255	P/R/U operand workspace

The recurrent state is loaded, converted to FP32, and retained in TMEM across the entire sequence. It does not return to global memory between chunks. Only the final BF16 state is stored when requested.

Keeping the state resident removes the dominant repeated state-memory traffic that would otherwise occur at every recurrent step.

7. Tensor-Core Work Division¶

The kernel uses two tensor-core paths for different matrix sizes:

Mqk and the small dense products used to construct the chunk-local triangular system use warp-level mma.sync.aligned.m16n8k16 instructions; the inverse also uses warp/register arithmetic.
Large state/output GEMMs use asynchronous SM100a tcgen05.mma.cta_group::1.kind::f16 instructions with BF16 inputs and FP32 TMEM accumulation.

This distinction matters: the token-parallel preparation path is not purely CUDA-core code, and Mqk is not implemented by the final wide tcgen05 operation.

8. Fused State-and-Output GEMM¶

The two final equations share the same left operand U:

\[ S_{new}=S_d+UK_r, \qquad O=O_0+UM_{qk}^T. \]

The kernel concatenates the right operands and accumulators along N:

\[ [S_{new}\;|\;O] = [S_d\;|\;O_0] + U[K_r\;|\;M_{qk}^T]. \]

For M128, the logical shapes are

\[ U:[128,32], \quad [K_r|M_{qk}^T]:[32,160], \quad [S|O]:[128,160]. \]

The instruction descriptor 0x08290490 encodes M=128, N=160, BF16 A/B, and FP32 D. K=32 is lowered to two K=16 tcgen05.mma instructions targeting the same M128×N160 accumulator, followed by one commit. Thus this is one logical fused GEMM and one completion boundary, not literally one K32 hardware instruction.

The implementation is in the final concatenated tcgen05 update.

9. M64 and M128 Physical Schedules¶

M is the number of value/state rows owned by one CTA, not the chunk length.

M128¶

One CTA owns all 128 state rows for one (sequence, head) pair.
It is used for every eligible packed workload and for most eligible fixed-layout workloads.
It maximizes reuse and minimizes redundant per-head work.

M64¶

Two CTAs split the 128 value rows into disjoint 64-row partitions.
It is selected only for fixed-layout B=1, H=64.
The split doubles CTA-level parallelism for the low-grid-parallelism fixed case.

The dispatch is implemented in _select_flash_kda_prefill_variant.

For packed input, optional seq_order changes CTA sequence order to improve tail utilization without changing the mathematical result.

10. True In-Place State Update¶

When an initial state is supplied, the binding passes the same allocation as both the initial-state and final-state pointer. There is no temporary state buffer followed by copy_ or memcpy.

This alias is safe because:

M128 loads all 128 rows owned by its (sequence, head) CTA before storing those rows back.
M64 CTAs own disjoint 64-row partitions and follow the same load-before-store rule.
the binding allows only disjoint buffers or an exact full-range alias and rejects partial overlap.

The Python path can be seen in the initial/final state pointer selection, while the C++ binding performs exact-alias validation.

11. TMA and CUDA Graph Safety¶

The backend uses TMA for the major global-to-SMEM transfers and output stores. Six tensor maps cover Q, K, V, G, beta, and output.

Tensor-map descriptors are published to stable device storage outside graph capture. At kernel entry, one thread executes fence.proxy.tensormap::generic.acquire.gpu for each map, followed by a CTA barrier before any TMA operation. This explicitly handles tensor-map proxy ordering rather than relying on ordinary kernel-start ordering.

RecurrentKDAPrefillWorkspace owns:

optional final-state scratch when no initial state exists;
beta padding;
separate M64 and M128 tensor-map descriptor blocks;
stream binding and warmed tensor signatures required for graph capture.

The relevant synchronization is in the kernel tensor-map acquire prologue, and workspace rules are documented in kda_prefill.py.

12. Strict Dispatch and Safe Fallback¶

The optimized backend is selected only when the call exactly matches its frozen-kernel contract, including:

compute capability exactly 10.0;
ordinary multi-token prefill;
contiguous BF16 Q/K/V/G and beta;
head dimension and state dimensions equal to 128;
shared Q/K/V head count;
FP32 A_log and dt_bias in the supported layouts;
in-kernel gate and Q/K L2 normalization;
logit beta and finite negative lower_bound;
no speculative decode, GQA, indexed state pool, or checkpoint features.

This keeps the fast path simple and specialized without regressing the broader API. Unsupported cases retain the existing implementation.

See the eligibility checks.

13. Reported Performance¶

The PR reports B200 CUPTI measurements against a pinned and binary-verified MoonshotAI/FlashKDA peer:

Case	Variant	This PR (ms)	FlashKDA (ms)	Speedup
H=96 fixed, 8192 tokens	M128	0.506095	1.085277	2.1444×
H=96 mixed packed	M128	0.391135	0.894702	2.2875×
H=96 uniform packed	M128	0.434992	0.735535	1.6909×
H=64 fixed, 8192 tokens	M64	0.469711	0.992638	2.1133×
H=64 mixed packed	M128	0.269903	0.676574	2.5067×
H=64 uniform packed	M128	0.292567	0.495871	1.6949×
Geometric mean				2.0512×

The benchmark used cold-L2 flushing, CUPTI activity tracing, deterministic inputs, rotating preinitialized state buffers, and no CUDA Graph in the timed region. The state update happened in place inside the candidate kernel.

The PR reports:

28 targeted recurrent-KDA prefill tests passed;
M64/M128 memcheck and synccheck passed with zero reported errors;
pre-commit hooks passed;
the full repository test suite was not run.

These results apply to the strict B200 contract and should not be generalized to unsupported architectures, layouts, or modes.

14. Important Caveats¶

B200-only specialization. The fast path requires exact compute capability 10.0.
Narrow tensor contract. It is not a general replacement for every recurrent-KDA path.
Large per-CTA resource use. M128 uses 1,024 threads, 256 TMEM columns, and 222 KiB SMEM. The schedule trades occupancy for state residency and reuse.
Five stages are a physical scheduling choice. They are not part of the KDA mathematics.
Chunk-32 safety depends on the gate bound. The centered radius is approximately 16 * abs(lower_bound) * log2(e). The usual lower_bound = -5 is safe, but increasingly negative values reduce BF16 exponent margin even though the dispatch contract accepts any finite negative value.
The final fused update is not one instruction. It is one logical M128×N160×K32 GEMM lowered to two K16 tcgen05 instructions and one commit.

Final Takeaway¶

PR #4262 succeeds by restructuring recurrent KDA around the B200 memory hierarchy:

independent chunk-local work is prepared in parallel;
recurrent work remains strictly ordered;
the state stays in FP32 TMEM;
SMEM is treated as a lifetime-aliased five-entry ring;
Tensor Cores handle both small intra-chunk and large recurrent matrix products through different instruction paths;
the state and output are updated together through N-dimension concatenation;
in-place state aliasing and graph-safe TMA integration remove surrounding overhead.

The reported approximately 2× speedup is therefore the combined result of algorithmic reformulation, memory-traffic elimination, pipeline parallelism, B200-specific tensor-core scheduling, and careful integration—not a single isolated micro-optimization.