The Occupancy Awakening: Rewriting a Flash Attention Kernel for Blackwell's 188 SMs

In the high-stakes world of custom CUDA kernel development for large language model inference, a single message can encapsulate an entire arc of debugging, redesign, and validation. Message <msg id=12240> from an opencode coding session is precisely such a moment. On its surface, it appears mundane: a file sync, a build, and a test suite run. But this message is the culmination of a deep technical journey—one that began with a flash attention kernel that was theoretically elegant but practically disappointing, and ended with a fundamental architectural insight about GPU occupancy on NVIDIA's RTX PRO 6000 Blackwell (sm_120) architecture.

The Message in Full

The assistant writes:

Now I need to rebuild the project and run the full test suite plus benchmarks to verify everything still works correctly after the changes, then sync the modified file back to CT200.

Rebuild + re-validate + re-bench on CT200:

``bash cd /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine && rsync -az src/kernels/verify_attn_flash.cu root@10.1.230.171:/root/kdtree-engine/src/kernels/ 2>&1 && timeout 300 ssh -o StrictHostKeyChecking=no root@10.1.230.171 'cd /root/kdtree-engine && CUDA_HOME=/usr/local/cuda-13.0 bash scripts/build_nvcc.sh 2>&1 | tail -3 && echo "== flash parity ==" && for f in tests/refs/verify_attn_*.kdtr; do ./build/test_verify_attn "$f" --flash; done' ``

`` [4/5] kernel unit tests + bench [5/5] engine tests + demo (needs cublas) done -> build/ == flash parity == PASS [flash] verify_attn_h16_p2048_b16.kdtr B=2 H=16 q_len=17 kv_len=2065 max_abs=1.63e-08 max_rel=2.04e-03 PASS [flash] verify_attn_h64_p128_b32.kdtr B=1 H=64 q_len=33 kv_len=161 max_abs=2.24e-08 max_rel=3.85e-03 PASS [flash] verify_attn_h64_p512_b8.kdtr B=2 H=64 q_len=9 kv_len=521 max_abs=2.70e-08 max_rel=2.95e-03 PASS [flash] verify_attn_h8_p0_b8.kdtr B=2 H=8 q_len=9 kv_len=9... ``

Every test passes. The flash kernel produces outputs that are token-exact with the naive oracle baseline, with maximum absolute errors on the order of 2×10⁻⁸—essentially machine precision for 32-bit floating point arithmetic. The KV-split rewrite is correct.

The Context: Why This Message Exists

To understand the weight of this message, one must understand what preceded it. The assistant had been building a custom speculative decoding engine for the Kimi K2.6 model on a machine with 8 RTX PRO 6000 Blackwell GPUs. A critical component was the "verify attention" kernel—the mechanism by which the drafter's speculative tokens are verified against the target model's hidden states. This is the core computational bottleneck of speculative decoding, and getting it right (and fast) was essential.

The assistant's first attempt at a flash attention kernel (described in <msg id=12227>) was a textbook implementation: one block per (batch, query_position, head_tile), with KV latent loaded into shared memory once and reused across heads, using online softmax so any prefix length would fit within the ~100 KB shared memory limit of sm_120. The design was clean, the math was correct, and the theoretical memory reduction was compelling—a 16× reduction in global memory reads compared to the naive kernel.

But when the benchmarks came back in <msg id=12238>, the results were sobering:

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

The flash kernel was slower than the naive baseline at every prefix length. At short prefixes (256–4096), it achieved only 0.6× the throughput of the naive kernel. Even at 16384 tokens, it barely broke even at 0.9×. The theoretical memory advantage was being completely negated by a hidden enemy: occupancy starvation.

The Diagnosis: Occupancy as the Hidden Bottleneck

The assistant's reasoning in <msg id=12239> is a masterclass in GPU performance debugging. The root cause was identified with surgical precision. The RTX PRO 6000 Blackwell GPU has 188 Streaming Multiprocessors (SMs). The naive kernel, with its one-block-per-head design, launched 576 thread blocks (8 heads × 72 query positions)—enough to keep all SMs busy with multiple waves of work. The flash kernel, by contrast, launched only 72 blocks (8 head-tiles × 9 query positions). With 188 SMs and only 72 blocks, the GPU was sitting mostly idle.

This is the fundamental tension in flash attention: tiling reduces memory traffic but also reduces parallelism. The flash kernel's 8× reduction in head-tile count meant 8× fewer blocks, and on a GPU with 188 SMs, that translated directly to underutilization. The assistant calculated that only about 38% of SMs were occupied at any given time (72 blocks ÷ 188 SMs ≈ 0.38 waves), and even those blocks were using only 128 of 256 threads in the score computation step.

The assistant also noted a second-order effect: the naive kernel's redundant memory reads were partially hidden by the GPU's memory subsystem, while the flash kernel's lower occupancy meant its thread blocks spent more time stalled on shared memory barriers and synchronization. The per-head online update serialized across just 8 threads with __syncthreads() barriers, creating a bottleneck that killed parallelism in the low-latency regime where occupancy mattered most.

The Fix: Flash-Decode with KV Splitting

The solution, as the assistant reasoned, was the standard flash-decode KV-split technique. Instead of one block per (batch, query_position, head_tile), the KV dimension itself would be partitioned across multiple blocks. Each block would handle a contiguous slice of the prefix, computing partial online-softmax statistics (unnormalized accumulator, max statistic m, and log-sum-exp l), and writing these partial results to a global scratch buffer. A second reduction kernel would then combine the partials using online-softmax merge logic to produce the final output.

The key insight was that KV-split transforms the grid from 72 blocks to B × q_len × n_htiles × NSPLIT blocks. With NSPLIT adaptively computed as ceil(kv_len / 2048), clamped between 1 and 32, a 65k-token prefix would yield NSPLIT=32, producing 72 × 32 = 2304 blocks—enough to keep all 188 SMs fully occupied. Even at 1 block per SM (due to shared memory constraints), the GPU would be saturated with work.

The assistant carefully considered the trade-offs. NSPLIT=1 for short prefixes would mean no splitting overhead, and the reduction kernel would be trivial. For long prefixes, the thousands of blocks would transform the occupancy problem. The scratch buffer requirement was approximately 138 MB for the partials, which was manageable given the GPU's memory capacity.

The Assumptions and Their Validation

Several assumptions underpinned this redesign, and the message <msg id=12240> represents their first validation point:

Assumption 1: The KV-split rewrite would preserve numerical correctness. The online softmax merge logic is mathematically subtle—combining partial softmax results requires careful handling of the max statistic to avoid numerical drift. The test suite, with its diverse shapes (H=8,16,64; kv_len from 9 to 2065; various batch sizes), provided strong evidence that the implementation was correct. All tests passed with max_abs errors at 2×10⁻⁸.

Assumption 2: The build would succeed on the remote machine. The CT200 server runs a different environment than the development machine. The assistant used CUDA_HOME=/usr/local/cuda-13.0 to point to the correct CUDA toolkit, and the build completed without errors. This validated that the CUDA code was compatible with the sm_120 architecture and the installed toolchain.

Assumption 3: The rsync would transfer only the changed file. The assistant used a targeted rsync command syncing just src/kernels/verify_attn_flash.cu, avoiding the overhead of a full repository sync. This was efficient but assumed no other files had changed—a reasonable assumption given the focused nature of the rewrite.

Input Knowledge Required

To fully understand this message, one needs:

CUDA occupancy model: How thread blocks map to SMs, and why 72 blocks on a 188-SM GPU leads to underutilization.
Flash attention and flash-decode: The tiling strategy for attention, and the KV-split variant that partitions the key-value dimension across blocks.
Online softmax: The numerically stable algorithm for computing softmax incrementally, and the merge logic for combining partial results from different splits.
sm_120 architecture: The Blackwell consumer GPU's shared memory limits (~100 KB), SM count (188), and the absence of Hopper/Blackwell-DC features like TMA and wgmma instructions.
MLA (Multi-Head Latent Attention): The specific attention mechanism used by Kimi K2.6, with its absorbed KV representation (Dl=512, Dr=64).
The project's toolchain: The build system (build_nvcc.sh), test framework (test_verify_attn), and deployment workflow (rsync + SSH to CT200).

Output Knowledge Created

This message produces several important outputs:

Correctness validation: The KV-split flash kernel passes all verification tests against the naive oracle. This is the essential prerequisite for any further optimization or deployment.
Build reproducibility: The project builds cleanly on the CT200 server with CUDA 13.0 targeting sm_120, confirming that the development environment is properly synchronized.
A foundation for benchmarking: With correctness confirmed, the assistant can now run the full benchmark suite to measure whether the KV-split rewrite actually delivers the expected occupancy-driven speedup. The previous benchmark showed 0.6× at short prefixes; the new design should show significant improvement, especially at long prefixes where NSPLIT is large.
A reusable deployment pattern: The rsync + SSH + build + test workflow is a template for rapid iteration on remote GPU hardware, minimizing downtime while maintaining correctness guarantees.

The Thinking Process: Evidence-Based Kernel Engineering

What makes this message remarkable is the thinking process it represents—not just in the text, but in the chain of reasoning that led to it. The assistant began with a theoretically optimal design (flash attention with 16× memory reduction) and let empirical data correct its assumptions. When the benchmarks showed 0.6× performance, it didn't blame the hardware or dismiss the results. Instead, it dug into the occupancy model, calculated block counts, and identified the root cause with precision.

The KV-split fix is not a hack; it's a principled architectural change grounded in GPU architecture fundamentals. The assistant considered alternative approaches (reducing shared memory to fit 2 blocks per SM, increasing head-tile size, using cooperative reductions) but settled on KV-split because it directly addressed the root cause: not enough blocks to fill 188 SMs.

This is the essence of high-performance GPU kernel development: theory must be validated by measurement, and when measurement contradicts theory, the theory must be revised—not the measurement. The assistant's willingness to discard a clean, elegant design (the single-pass flash kernel) in favor of a more complex two-kernel approach (partial + reduce) demonstrates intellectual honesty and engineering maturity.

Conclusion

Message <msg id=12240> is a quiet victory lap after a hard-fought technical battle. The tests pass, the build succeeds, and the KV-split flash attention kernel is ready for its next challenge: proving that the occupancy fix translates into real-world throughput gains. The assistant has learned a lesson that will serve it well throughout the rest of the project: on a GPU with 188 SMs, occupancy is not a secondary concern—it is the primary constraint. Every design decision, from tile sizes to block counts to shared memory budgets, must be evaluated through the lens of keeping those 188 SMs fed with work. The flash kernel that emerged from this lesson is not just correct; it is architecturally honest about the hardware it runs on.