[KernelFlow-ops]

Summary

The Ampere port this branch builds on (pzhao-eng/FlashMLA@a100_replace, 21 commits, included) does not launch on GeForce Ampere: its kernel needs 113,280 B of shared memory per block and sm_86 allows 101,376 B. The 27 commits added here run MLA decode on an RTX 3090 at the hardware ceiling, then optimize the decode layer around it.

What is added

• v2 decode kernel (csrc/sm86/): Q in registers, split-K QK inside the CTA, a multi-stage cp.async KV ring; balanced scheduler, deterministic combine, graph-safe API; bf16 / fp16, page sizes 16 / 32 / 64.
• 8-bit KV cache: FP8 e4m3 or INT8 with one scale per token (644 instead of 1152 bytes per token), dequantized inside the kernel.
• Fused decode layer (csrc/fused/, flash_mla/mla_decode_layer.py): one prologue kernel (RoPE, concat, cache write, quantize-on-write), a weight-streaming "skinny" GEMM for the four projections, and the whole layer as one CUDA graph: 7 kernels per layer instead of about 44.
• Benchmarks, profiler helpers, 760 tests, and a write-up with four figures (docs/optimization/).

Performance

One RTX 3090; measured peaks 873.8 GB/s (read) and 70.0 TFLOPS (bf16). DeepSeek-V2-Lite layer (h_q = 16), 4096 cached tokens per request, 27-layer CUDA graph, microseconds per layer (lower is better), batch size b:

Decode layer	b=1	b=8	b=32	b=64
Unfused PyTorch, eager	820.1	895.1	903.1	899.2
Unfused PyTorch, one CUDA graph	114.7	184.9	330.0	515.0
Fused, Triton attention	76.4	133.8	355.7	599.6
Fused, port kernel (16-row fix)	71.1	115.7	258.2	440.9
Fused, v2 kernel	65.8	117.9	259.3	440.0
+ skinny GEMM, bf16 cache	61.5	102.3	239.5	431.4
+ FP8 cache	61.6	86.3	172.0	291.2
bf16 vs unfused eager	13.3×	8.7×	3.8×	2.1×
bf16 vs unfused graph	1.87×	1.81×	1.38×	1.19×
FP8 vs unfused graph	1.86×	2.14×	1.92×	1.77×

Reading the table: rows 1–2 are the unfused baselines (attention by the port's kernel), rows 3–5 the fused layer with three attention kernels, the bold rows the current build; the last three divide a baseline by a bold row. The eager layer is launch-bound: over 40 launches cost 820–900 µs at any batch size. The graph removes that cost; fusion removes the small kernels between projections and attention; the skinny GEMM gains 2–15%, most at b=8, where projections are half of the layer. At large batch attention dominates and the 8-bit cache is the lever (b=64: 431.4 → 291.2 µs).

The decode kernel alone:

Regime	Measured	Reference
h_q=16, b≥64, s_k≥4096 (memory-bound), graph	821–864 GB/s	94–99% of the read peak
h_q=128, b≥16, s_k≥4096 (compute-bound), eager	50.8–52.6 TFLOPS	73–75% of the bf16 peak
FP8 / INT8 cache vs bf16, b≥32, s_k≥4096	1.55–1.70× / 1.54–1.60×	ideal 1.79×
Time of other kernels relative to v2, h_q=16 / 128	port 1.07 / 1.17, Triton 1.47 / 2.23, FlashInfer 3.43 / 1.77	geometric mean, graph

Correctness

760 tests pass: fp32 reference (cosine distance below 8e-5), NaN padding, variable lengths, graph replay, bit-identical reruns; compute-sanitizer is clean. 8-bit accuracy on synthetic caches: worst cosine distance 2.4e-3 (FP8), 2.6e-4 (INT8).

Notes

• The branch predates the restructuring of main (merge base bcb90f2): 11 files conflict, so it cannot be merged as is and is offered as a reference Ampere branch.
• Phase 3 (prototypes: b=1 61.1 → 56.2 µs) is planned in docs/optimization/04-phase3-plan.md, harnesses in research/; not implemented.
• Clocks cannot be locked on this machine; medians are reported.

🤖 Generated with Claude Code