The Decisive Microbench: When Theory Meets Silicon

Introduction

In the high-stakes world of custom CUDA kernel development, there is no moment more revealing than the first benchmark run after a major optimization. All the careful reasoning about occupancy, shared memory budgets, thread block counts, and memory bandwidth coalesces into a single table of numbers—and those numbers either validate weeks of work or send you back to the drawing board. Message [msg 12241] captures exactly such a moment in the development of a custom sm_120 verify attention kernel for the RTX PRO 6000 Blackwell GPU. It is a short message—barely a dozen lines of benchmark output—but it carries the weight of an entire optimization cycle, representing the culmination of a deep diagnostic effort and the implementation of a sophisticated flash-decode technique called KV-splitting.

The Context: An Occupancy Crisis

To understand the significance of this message, one must first understand the problem it was designed to solve. In the preceding messages ([msg 12238] and [msg 12239]), the assistant had implemented an initial "flash" verify attention kernel—a custom CUDA kernel for the Multi-Head Latent Attention (MLA) mechanism used in the Kimi K2.6 model. The kernel used a tiled approach where each thread block handled one (batch, query_position, head_tile) combination, loading KV latent vectors into shared memory once and reusing them across multiple heads to reduce global memory traffic by a factor of 2×HT (where HT is the head tile size, set to 8).

The initial benchmark results were sobering. At short prefix lengths (256–4096 tokens), the flash kernel was actually slower than the naive baseline—achieving only 0.6× speedup. At 16,384 tokens, it barely broke even at 0.9×. The naive kernel couldn't run at 65,536 tokens (its shared memory requirement exceeded the ~100 KB limit of sm_120), but the flash kernel managed 93.9 ms—far slower than the theoretical memory floor of ~8.6 ms.

The assistant diagnosed the root cause with precision: occupancy starvation. With batch size 1, 9 query positions, and 8 head tiles, the flash kernel launched only 72 thread blocks total. The RTX PRO 6000 Blackwell GPU has approximately 188 Streaming Multiprocessors (SMs). This meant that nearly two-thirds of the GPU's compute resources sat idle during the kernel's execution. The naive kernel, despite its redundant memory reads, launched 576 blocks (8× more) and kept the GPU better occupied. The flash kernel's sophisticated memory reuse was rendered moot by its inability to keep the hardware busy.

The KV-Split Solution

The assistant's response was to adopt a standard technique from the flash-attention literature: KV-splitting (also known as flash-decode). Instead of one block per (batch, query, head_tile), the KV dimension itself is partitioned across multiple blocks. Each block handles a contiguous slice of the key-value cache, computes a partial online softmax, and writes its unnormalized accumulator and statistics to a global scratch buffer. A second lightweight reduction kernel then merges these partial results using the online-softmax combine formula.

The effect on block count is dramatic. With NSPLIT partitions, the grid size becomes B × q_len × n_htiles × NSPLIT. For a 65,536-token prefix with NSPLIT=32 (the adaptive maximum), this yields 72 × 32 = 2,304 blocks—enough to saturate all 188 SMs with multiple waves of work. The assistant implemented this design in message [msg 12239], rewriting verify_attn_flash.cu from scratch, then synced and rebuilt on the CT200 server in [msg 12240], where all correctness tests passed with token-exact precision.

The Message: A Moment of Truth

Message [msg 12241] is the benchmark run that answers the central question: did KV-splitting work? The message is quoted in full below:

[assistant] ## Agent Reasoning Now I'm running the benchmark to see if the KV-split optimization actually resolved the occupancy issues and delivers the performance gains we're expecting for longer prefixes. Still token-exact with split-K. Now the decisive microbench: [bash] timeout 180 ssh -o StrictHostKeyChecking=no root@10.1.230.171 'cd /root/kdtree-engine && ./build/bench_kernels 2>&1 | sed -n "/verify_attn/,/tree_build/p" | head -26' == verify_attn (MLA absorb, H=64 Dl=512 Dr=64) == streams q_len prefix naive_us flash_us speedup 1 9 256 199.5 436.4 0.5x 1 9 1024 738.8 1665.6 0.4x 1 9 4096 2887.8 4342.1 0.7x 1 9 16384 19602.0 11600.7 1.7x 1 9 65536 n/a 48477.6 - 1 33 256 651.9 951.9 0.7x 1 33 1024 2231.5 3390.6 0.7x ...

The assistant's reasoning reveals the mindset: this is framed as a decisive test. The word "decisive" is carefully chosen—it acknowledges that the entire KV-split optimization stands or falls on these numbers. The assistant also notes "Still token-exact with split-K," confirming that correctness was preserved through the rewrite, a non-trivial achievement given the complexity of online softmax combine logic across partitioned KV ranges.

Interpreting the Numbers

The benchmark table tells a nuanced story. At short prefixes (256 tokens), the KV-split flash kernel achieves only 0.5× speedup relative to the naive baseline—actually worse than the pre-split version's 0.6×. At 1024 tokens, it's 0.4× (down from 0.6×). The KV-split overhead—launching a second reduction kernel, writing partial results to global memory, the extra synchronization—hurts in the regime where the naive kernel is already fast (199 μs at 256 tokens) and occupancy is less of an issue.

But as the prefix length grows, the trend reverses decisively. At 4096 tokens, the flash kernel reaches 0.7× (up from 0.6× pre-split). At 16,384 tokens, it achieves 1.7× speedup—a dramatic improvement from the pre-split 0.9×. And at 65,536 tokens, where the naive kernel cannot even run (shown as "n/a" because its shared memory requirement exceeds the hardware limit), the flash kernel completes in 48.5 ms, down from 93.9 ms in the pre-split version—a 1.94× improvement from the KV-split alone.

The 48.5 ms at 65k tokens is still far from the theoretical memory floor of ~8.6 ms (calculated in [msg 12239] as the time to read the KV cache once at ~1.8 TB/s bandwidth), but it represents real progress. The assistant has cut the long-context latency nearly in half with a single architectural change.

The second row group (q_len=33) shows a similar pattern: 0.7× at short prefixes, with the trend improving at longer contexts (the full output was truncated by head -26).

Assumptions and Their Validation

Several assumptions underpin this benchmark. The assistant assumed that occupancy starvation was the primary bottleneck at long prefixes—an assumption validated by the 1.7–1.94× speedup at 16k–65k tokens. It assumed that the KV-split overhead (extra kernel launch, scratch memory writes) would be amortized at long prefixes—validated by the crossover occurring between 4k and 16k tokens. It assumed that NSPLIT=32 would provide sufficient parallelism without excessive overhead—the results suggest this is reasonable, though further tuning might improve short-prefix performance.

One assumption that proved incorrect was that KV-splitting would not significantly regress short-prefix performance. The assistant had reasoned in [msg 12239] that "when NSPLIT=1, the partial kernel becomes a single split... so there's minimal overhead." Yet the benchmark shows short-prefix performance actually decreased from 0.6× to 0.5× at 256 tokens. This suggests that even with NSPLIT=1, the two-kernel launch architecture (partial + reduce) introduces overhead—perhaps from the scratch buffer allocation, the extra kernel launch latency, or the reduction kernel's synchronization cost—that wasn't present in the single-kernel design.

Input Knowledge Required

To fully understand this message, one needs knowledge of: CUDA thread block occupancy and its relationship to SM count; the flash attention algorithm and its tiling strategy; the online softmax technique for computing attention over arbitrarily long sequences without materializing the full score matrix; the MLA (Multi-Head Latent Attention) mechanism used in modern LLMs like Kimi K2.6, which decomposes the KV cache into latent and rope components; the sm_120 architecture's shared memory limits (~100 KB) and how they constrain kernel design; and the concept of KV-splitting (flash-decode) as a standard technique for long-context attention.

Output Knowledge Created

This message creates several pieces of critical knowledge. First, it establishes that KV-splitting is effective on sm_120 Blackwell GPUs for long-context verify attention, with the crossover point around 4k–16k tokens. Second, it quantifies the overhead of the two-kernel approach at short prefixes (roughly 2× slower than naive at 256 tokens). Third, it provides a baseline (48.5 ms at 65k tokens) against which future optimizations—such as cp.async double-buffering, vectorized memory loads, and increased head tile sizes—can be measured. Fourth, it confirms that the online softmax combine logic across KV splits is numerically stable (token-exact results).

The Thinking Process

The assistant's reasoning in this message is concise but revealing. The phrase "Now I'm running the benchmark to see if the KV-split optimization actually resolved the occupancy issues and delivers the performance gains we're expecting for longer prefixes" shows an expectation that the optimization should work, tempered by the humility of letting data decide. The word "actually" is telling—it acknowledges that prior assumptions (the initial flash kernel would be faster) were overturned by evidence, and the assistant is now cautious about declaring victory prematurely.

The comment "Still token-exact with split-K" is placed before the benchmark results, not after. This ordering is significant: the assistant is establishing correctness first, then performance. In kernel development, a fast but wrong kernel is useless; a correct but slow kernel can be optimized. By confirming correctness before even looking at the speedup numbers, the assistant demonstrates disciplined engineering practice.

Broader Significance

This message sits at a critical inflection point in the project. The assistant has been iterating on this verify attention kernel across multiple rounds (see [chunk 66.1]), and the KV-split optimization is the most significant architectural change yet. The 1.7–1.94× speedup at long contexts validates the approach and clears the way for further refinements: increasing NSPLIT, tuning head tile sizes, adding vectorized memory access patterns, and eventually integrating the kernel into the live SGLang service with CUDA graph capture.

The message also illustrates a fundamental tension in GPU kernel design: optimizations that improve asymptotic behavior at scale often regress small-problem performance. The KV-split kernel is 2× slower than naive at 256 tokens but nearly 2× faster at 65k tokens. For a production system serving long-context LLM inference, the long-context win is far more valuable—but the short-context regression must eventually be addressed, perhaps through a heuristic that selects between the naive and flash kernels based on prefix length.

Conclusion

Message [msg 12241] is a textbook example of evidence-driven kernel development. A hypothesis (occupancy starvation → KV-split fix) was formulated, implemented, and tested against a decisive benchmark. The results were mixed but directional: the KV-split technique works, delivering meaningful speedups at the long contexts that matter most, while revealing a short-prefix regression that will need future attention. The assistant's disciplined approach—confirming correctness before evaluating performance, framing the test as decisive, and interpreting the numbers without spin—models the scientific mindset that separates professional systems engineering from guesswork. The 48.5 ms at 65k tokens is not the final answer, but it is a clear signal that the architecture is on the right track.