The Moment of Truth: Benchmarking a Custom sm_120 Flash Attention Kernel

In the high-stakes world of custom CUDA kernel development for large language model inference, there is no more revealing moment than the first benchmark run. All the careful reasoning about tile sizes, shared memory budgets, thread occupancy, and memory coalescing — the hours of design and implementation — condense into a single table of numbers. Message [msg 12238] captures exactly this moment: the first performance measurement of a custom flash attention kernel for the MLA (Multi-head Latent Attention) architecture, running on an NVIDIA RTX PRO 6000 Blackwell GPU (sm_120). The results were sobering, and the analysis that followed reshaped the entire optimization trajectory of the project.

The Context: A GPU Architecture Without Library Support

To understand why this message matters, one must appreciate the predicament that led to it. The project was deploying the Kimi K2.6 model with DFlash speculative decoding on RTX PRO 6000 Blackwell GPUs — a consumer-class Blackwell variant with an sm_120 architecture. The problem was stark: every optimized MLA kernel library available — FlashMLA, cutlass-MLA, flashinfer-MLA — was compiled exclusively for sm_90a (Hopper), sm_100a, or sm_103a (Blackwell data center). None supported sm_120, which uses an Ada-like ISA lacking the wgmma, TMA, and tcgen05 instructions that those optimized kernels depend on. The assistant was therefore building an owned sm_120 kernel from scratch, with a detailed plan documented in plans/0002-sm120-verify-kernel-defrag.md.

The kernel being benchmarked was a flash attention implementation of the "verify" step in DFlash speculative decoding — the operation that computes attention scores between the draft tree's candidate tokens and the full KV cache to determine which tokens to accept. The naive (oracle) baseline simply loaded the entire KV cache, computed a full score matrix in shared memory, and produced the output. The flash version was designed to tile across the KV dimension, loading chunks into shared memory and using online softmax to handle arbitrarily long prefixes within the sm_120's ~100 KB shared memory limit.

The Message: Raw Performance Data

The subject message is compact but dense. It opens with the agent's reasoning:

Both kernels pass correctness on all test bundles with the flash implementation matching the naive version to machine precision. Now I'm moving to the microbenchmark phase to compare performance across a range of prefix lengths, including longer sequences that the naive kernel can't handle.

This establishes that the correctness gate has been passed — the flash kernel produces token-exact results (max absolute error ~2e-8) matching the naive oracle. The focus now shifts to performance. The assistant then runs the benchmark:

== verify_attn (MLA absorb, H=64 Dl=512 Dr=64) ==
streams  q_len   prefix    naive_us    flash_us  speedup
     1      9      256       199.3       346.1     0.6x
     1      9     1024       738.5      1313.4     0.6x
     1      9     4096      2886.2      5184.2     0.6x
     1      9    16384     19596.5     20665.5     0.9x
     1      9    65536         n/a     93958.5        -
     1     33      256       652.1       753.3     0.9x
     1     33     1024      2213.4      2680.4     0.8x

The data tells an uncomfortable story. At short prefixes (256–4096 tokens), the flash kernel is actually slower than the naive baseline — a 0.6x speedup, meaning the naive kernel is nearly 1.7x faster. Only at 16k tokens does the flash kernel pull slightly ahead (0.9x, still not a win). The one unambiguous advantage is at 65k tokens, where the naive kernel cannot run at all (shown as "n/a") because its shared memory requirement — a full score row of kv_len × 4 bytes — exceeds the ~100 KB limit. The flash kernel handles it gracefully at 93.9 ms, but that's a far cry from the 10–30× speedup the assistant had hoped for.## The Analysis: Occupancy Starvation

The agent's reasoning in this message is terse — it simply notes that the correctness gate is met and runs the benchmark. But the next message ([msg 12239]) reveals the full depth of the analysis that this data triggered. The flash kernel was suffering from a classic GPU optimization pitfall: occupancy starvation.

The design used one thread block per (batch, query_position, head_tile). With batch size 1, query length 9, and 8 head tiles (H=64 divided by HT=8), this produced only 72 thread blocks total. The RTX PRO 6000 Blackwell has approximately 188 SMs. With only 72 blocks, the GPU was barely 38% utilized — and the problem was worse than that. Each block consumed 70 KB of shared memory, which meant only one block could reside on an SM at a time. With just 128 threads per block, the SM's warp scheduler had only 4 warps to work with — roughly 6% occupancy on a machine designed to context-switch between 2048 threads. Every stall on a global memory load or shared memory read would stall the entire SM, since there were no other warps to schedule in the gap.

The naive kernel, by contrast, used a much smaller shared memory footprint (just the score row, which at short prefixes fits easily) and generated 576 blocks (H × q_len) — 8× more parallelism. Its blocks were simpler and faster to launch, and the higher occupancy compensated for its redundant memory reads (each block read the KV cache twice, once for scoring and once for the PV computation).

This is a counterintuitive result: the "optimized" flash kernel that reads the KV cache 16× fewer times (amortized across 8 head tiles) was slower than the "naive" kernel that reads everything redundantly. The reason is that on modern GPUs, decode-phase attention is latency-bound, not memory-bandwidth-bound. When the GPU is barely occupied, memory bandwidth is not the limiting factor — the GPU simply isn't issuing enough instructions to saturate the memory subsystem. Reducing memory reads doesn't help if the GPU is idle most of the time anyway.

The Assumptions That Were Challenged

Several assumptions embedded in the initial design were called into question by this benchmark:

"Amortizing KV reads across head tiles is always beneficial." The flash kernel's head-tile design (HT=8) reduced global memory traffic by 16×, but the shared memory cost (70 KB) limited occupancy so severely that the reduction in memory reads was irrelevant. The naive kernel's redundant reads were actually faster because they came with much higher occupancy.
"Flash attention's tiling approach will naturally outperform a naive implementation." The flash attention literature typically assumes large batch sizes and compute-bound regimes where memory bandwidth is the bottleneck. In the decode regime (batch size 1, short query lengths), the bottleneck shifts to latency and occupancy, and the standard flash attention tradeoffs break down.
"The sm_120 architecture behaves similarly to other Blackwell variants." The lack of tensor-core instructions (wgmma, TMA) on sm_120 meant the kernel had to rely on plain CUDA cores for matrix operations, making occupancy even more critical since each FMA was slower than a tensor-core equivalent would have been.
"The flash kernel's ability to handle arbitrary prefix lengths is the primary win." While the flash kernel could handle 65k tokens where the naive kernel could not, the 93.9 ms latency at that length was still far from acceptable for production use. The ability to run at extreme lengths was necessary but not sufficient — performance at those lengths needed to be dramatically better.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

MLA (Multi-head Latent Attention): The attention mechanism used by Kimi K2.6, where queries and keys are decomposed into "nope" (no positional encoding) and "rope" (rotary positional encoding) components, with dimensions Dl=512 and Dr=64.
DFlash speculative decoding: A technique where a draft model generates a tree of candidate tokens, and the target model's attention mechanism "verifies" which tokens to accept. The verify attention operation is the core computational bottleneck.
CUDA thread block occupancy: The relationship between shared memory per block, registers per thread, and the number of blocks that can simultaneously reside on an SM. On sm_120, each SM has ~100 KB of shared memory and can hold up to 2048 threads (64 warps), but practical occupancy is constrained by resource usage per block.
Online softmax: A numerically stable technique for computing softmax incrementally, allowing attention over arbitrarily long sequences without storing the full score matrix. This is what enables the flash kernel to handle 65k+ token prefixes.
sm_120 architecture: A Blackwell consumer GPU variant that lacks the tensor-core matrix instructions (wgmma, TMA, tcgen05) present in Hopper and Blackwell data-center parts, forcing reliance on CUDA cores for matrix multiply-accumulate operations.## Output Knowledge Created This message, combined with the analysis it triggered, produced several critical insights that shaped the subsequent development: The occupancy bottleneck diagnosis. The benchmark proved that the flash kernel's design was fundamentally occupancy-starved in the decode regime. This led directly to a complete redesign in the following messages ([msg 12239], [msg 12242]), where the assistant pivoted to a "per-head, small-smem" design with KV-splitting and warp-per-key dot products. The shared memory target dropped from 70 KB to ~23 KB, allowing 4 blocks per SM instead of 1, and the block count increased from 72 to thousands through KV partitioning. The confirmation that the naive kernel is a valid baseline. The naive kernel's performance (19.6 ms at 16k tokens) established a reference point that any optimized kernel must beat. The fact that the naive kernel could not run at 65k tokens also validated one of the flash kernel's key advantages — arbitrary-length support — even if the performance at that length was not yet competitive. The discovery that decode-phase attention is latency-bound. This is a general insight that applies beyond this specific kernel. For batch-size-1 decode with short query lengths, the GPU is fundamentally underutilized, and optimization strategies must prioritize occupancy and latency hiding over memory bandwidth reduction. This runs counter to the conventional wisdom in the attention optimization literature, which focuses on large-batch training scenarios. The realization that the production win might come from contiguous memory access, not algorithmic efficiency. The microbenchmark used contiguous KV memory, but the production SGLang service used page_size=1 KV cache, causing scattered gather at ~14 GB/s effective bandwidth. The naive kernel with contiguous access might already be 10× faster than Triton's scattered path in production, even if it appears unimpressive in microbenchmarks. This reframed the optimization goal: the custom kernel didn't need to be faster than the naive baseline in microbenchmarks; it needed to be faster than Triton's scattered path in production.

The Thinking Process: From Data to Diagnosis

While the subject message itself is brief — essentially just the benchmark command and its output — it sits at a critical juncture in a much longer reasoning chain. The messages leading up to it show the assistant's design process:

In [msg 12227], the assistant designed the initial flash kernel with careful shared memory budgeting: TK=16 (key tile size), HT=8 (head tile count), BLOCK=256 threads, and 70 KB shared memory. The reasoning was thorough: "TK=16 gives me 72KB of shared memory, which fits safely under the 100KB limit with good margin." The assistant considered alternatives — HT=16 (107 KB, over budget), dropping q_nope from shared memory (69 KB but redundant global reads) — and settled on HT=8 as the best compromise.

In [msg 12234], the assistant anticipated the naive kernel's limitation at long prefixes: "the naive kernel requires full-score-row in shared memory, and at prefix 65536 that would need 262KB, which exceeds the 100KB shared memory maximum." The guard was added to show "n/a" when the naive kernel couldn't run, which also served to demonstrate the flash kernel's advantage.

Then came the benchmark in [msg 12238]. The numbers were clear: 0.6x at short prefixes, 0.9x at 16k, and a lonely 93.9 ms at 65k. The flash kernel was correct but not fast.

The analysis in [msg 12239] was where the real insight emerged. The assistant worked through the math: "With only 9 query nodes and 8 head tiles, I'm generating just 72 thread blocks total, but the RTX 6000 has ~188 SMs — the GPU sits mostly idle." The fix was identified as KV-splitting — partitioning the prefix across many blocks to generate thousands of blocks instead of 72, keeping all SMs busy. The assistant also noted that "the score computation step only uses half its threads (128 of 256), and the per-head online update serializes across just 8 threads with synchronization barriers."

By [msg 12242], the assistant had iterated further, exploring per-head designs with smaller shared memory footprints and warp-per-key dot products to improve thread utilization. The shared memory target dropped from 70 KB to ~23 KB, and the design philosophy shifted from "minimize memory reads" to "maximize occupancy." This was the direct consequence of the benchmark data in [msg 12238].

The Broader Significance

This message exemplifies a pattern that recurs throughout high-performance GPU programming: the gap between theoretical analysis and empirical measurement. The assistant's initial design was mathematically sound — fewer global memory reads should mean less time waiting on memory. But the GPU's complex scheduling dynamics, where occupancy and latency hiding interact with memory bandwidth in non-obvious ways, meant that the theoretically superior design was practically slower.

The lesson is not that flash attention is flawed, but that optimization must be guided by measurement, not intuition. The benchmark in [msg 12238] provided the empirical reality check that redirected the entire optimization effort. Without it, the assistant might have continued refining the HT=8 design, chasing marginal improvements while the fundamental occupancy problem remained unaddressed.

The message also illustrates the value of having a correct baseline. The naive kernel, for all its simplicity, established a performance floor and revealed the flash kernel's weakness. In GPU programming, a well-written naive kernel is often faster than a poorly optimized "advanced" kernel, because simplicity enables the compiler and hardware to do their jobs efficiently.

Finally, this message captures the moment when a project transitions from "does it work?" to "does it work well?" — the shift from correctness to performance that defines the difference between a prototype and a production system. The answer at this moment was "not yet," but the data provided the roadmap for getting there.