Skip to content

FlashInfer PR #4262: CAKE KDA Prefill Highlights

A short technical note for FlashInfer PR #4262: an optimized BF16 recurrent KDA prefill backend generated by CAKE for NVIDIA B200 / SM100a.

TL;DR

The main idea is to keep almost the entire KDA prefill dataflow inside one fused kernel:

  • prepare five future 32-token chunks in parallel;
  • consume those chunks in strict recurrent order;
  • retain the FP32 recurrent state in TMEM;
  • use Tensor Cores for both chunk-local and recurrent matrix products;
  • reuse SMEM aggressively enough to fit the five-stage pipeline in one CTA.

The PR reports a 2.0512× geometric-mean speedup over MoonshotAI/FlashKDA on its six B200 workloads.

1. Chunk 32 with a fixed exponent anchor

Ordinary SafeGate processing uses a 16-token chunk because directly materializing both positive and inverse gate factors over 32 tokens can exceed the useful BF16 exponent range.

The CAKE kernel instead centers the exponent around a fixed scalar:

\[ b = \mathrm{lower\_bound}\log_2(e), \qquad C = 16b. \]

For chunk prefix sums \(G_i\), the centered Q/K operands contain factors such as

\[ 2^{G_i-C} \quad\text{and}\quad 2^{C-G_j}. \]

The anchor cancels in Mqk:

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

Therefore:

  • the mathematical result is unchanged;
  • each individual BF16 operand sees only half of the original chunk-32 exponent radius;
  • chunk 32 has approximately the same operand exponent radius as an uncentered chunk 16;
  • \(C\) is derived from the configured lower bound, not sampled from the actual middle token.

Chunk-32 constant-anchor BF16 safety

2. Fully fused M128 recurrent pipeline

The numerical prefill path is fused rather than split into a FlashKDA-style preparation kernel plus recurrent kernel. Chunk-local intermediates stay in registers, SMEM, or TMEM instead of being written to a global workspace.

For each 32-token chunk, the kernel roughly performs:

  1. gate prefix accumulation and centered Q/K preparation;
  2. construction of the causal Mqk and triangular system;
  3. inversion of the 32×32 lower-triangular matrix;
  4. recurrent state decay and projection;
  5. residual/beta transformation;
  6. fused state and output update;
  7. final output transpose and store.

The final two updates share the same left operand \(U\):

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

They are concatenated into one logical GEMM:

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

For the M128 schedule:

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

Thus the logical tcgen05 GEMM is M=128, N=160, K=32. Hardware lowering uses two K=16 tcgen05.mma operations on the same accumulator, followed by one completion boundary.

Other important details:

  • the recurrent state remains FP32 in TMEM across chunks;
  • Mqk and other small 32×32 products use warp-level mma.sync.aligned.m16n8k16;
  • M64 and M128 are alternative physical schedules, trading per-CTA work against grid parallelism;
  • M128 uses 256 TMEM columns for the state, output, accumulators, and temporary operands.

M128 SMEM reuse and mathematical flow

3. Five-stage SMEM ring buffer

Chunk preparation is parallel, but recurrent state consumption must remain ordered. The M128 kernel connects them with five fixed SMEM slots:

\[ \mathrm{slot}(c)=c\bmod 5. \]

Five producer groups prepare chunk sequences

\[ p,\ p+5,\ p+10,\ldots, \qquad p\in\{0,1,2,3,4\}, \]

while the consumer processes chunks in order

\[ 0,1,2,3,4,5,\ldots \]

After chunk \(c\) is consumed, slot \(c\bmod5\) can be refilled with chunk \(c+5\). Mbarriers provide readiness notification and backpressure.

The five stages are a ring buffer for chunk-local operands, not five copies of the recurrent state. Their purpose is to:

  • expose parallelism in gate/Q/K/Mqk/inverse preparation;
  • hide preparation latency behind recurrent MMA work;
  • absorb short producer/consumer speed differences;
  • preserve strict recurrent ordering without global synchronization.

Five-stage ring-buffer scheduling

4. Why the memory footprint fits

The M128 kernel does not allocate independent SMEM buffers for every logical tensor. It aliases storage according to tensor lifetime, for example:

  • raw gate storage becomes centered/decayed Q storage;
  • raw K storage becomes transformed K storage;
  • raw Q storage becomes inverse-scaled and restored K storage;
  • gate-prefix and inverse workspaces later hold Mqk, restore factors, or V;
  • output ping-pong buffers are reused during the consumer pipeline.

Key sizes are:

Region Size
One preparation stage 41,984 bytes
Five-stage payload 209,920 bytes
Two output buffers 16,384 bytes
Barriers/control 1,024 bytes
Total M128 SMEM 227,328 bytes (222 KiB)

This lifetime-based reuse is what makes five-stage look-ahead possible within the B200 per-CTA SMEM budget.

5. Practical takeaways

  • Numerical representation unlocks a larger tile. The fixed anchor makes chunk 32 safe without changing Mqk math.
  • Separate dependency domains. Chunk-local preparation is parallel; recurrent state evolution remains serialized.
  • Use each Tensor Core path where it fits. Warp MMA handles small 32×32 matrices, while tcgen05 handles the large state/output GEMMs.
  • Keep the expensive state on chip. FP32 TMEM residency removes repeated state traffic between chunks.
  • Fuse common operands. Concatenating state and output along N turns two updates into one logical M128×N160×K32 operation.
  • Pipeline depth is a scheduling choice. Five stages provide enough look-ahead to overlap preparation and consumption, subject to the SMEM budget.
  • The optimization is specialized. It targets B200/SM100a, BF16, head dimension 128, and supported recurrent-prefill layouts; unsupported cases fall back to the existing backend.
  • Grid parallelism remains a limitation. The kernel is not a general persistent megakernel, so small batch×head counts can still underutilize the GPU.

References