The Occupancy Trap: When a Carefully Designed CUDA Kernel Fails Because of What You Didn't Measure
In the high-stakes world of custom CUDA kernel development for large language model inference, few moments are as instructive as the collision between a well-reasoned design and an unforgiving benchmark. Message 12239 of this opencode session captures precisely such a moment. The assistant, having spent multiple rounds designing, implementing, and validating a custom flash attention kernel for the NVIDIA RTX PRO 6000 Blackwell GPU (sm_120 architecture), confronts benchmark data that reveals its creation is slower than the naive baseline it was meant to replace. What follows is a masterclass in diagnostic reasoning, architectural pivoting, and the subtle art of GPU occupancy analysis.
The Scene: A Kernel That Should Have Won
To understand the weight of this message, we must reconstruct the context. The assistant has been engaged in a months-long effort to deploy and optimize speculative decoding for large language models on a cluster of RTX PRO 6000 Blackwell GPUs. The specific problem at hand is the "verify attention" kernel—the component that computes attention scores between a draft tree of candidate tokens and the full key-value cache during speculative decoding. This kernel is called on every decode step, making it a critical latency path.
The assistant had discovered earlier (in chunk 0 of segment 66) that the existing verify attention implementation was locked to a Triton-based MLA kernel using page_size=1, which caused scattered KV cache accesses and achieved an effective memory bandwidth of only ~14 GB/s—roughly 130× below the GPU's 1.8 TB/s peak. All optimized MLA kernels from the broader ecosystem (FlashMLA, cutlass-MLA, flashinfer-MLA) were compiled exclusively for sm_90a (Hopper), sm_100a, and sm_103a architectures—none supported sm_120, the Blackwell consumer GPU's instruction set. The RTX PRO 6000 Blackwell lacks the warp-group matrix multiply (wgmma), Tensor Memory Accelerator (TMA), and tcgen05 instructions that power those optimized kernels. The assistant was building from scratch, on an architecture with only ~100 KB of shared memory per SM—roughly half of what Hopper offers.
The assistant wrote a detailed plan (plans/0002-sm120-verify-kernel-defrag.md) and implemented Phase 1: a tiled flash attention kernel that processes KV cache in chunks, reusing each chunk across multiple query heads via shared memory. The design was sound on paper: by tiling across head dimensions with HT=8 (processing 8 heads at a time from the KV data loaded into shared memory), the kernel should reduce global memory reads by a factor of 2×HT = 16× compared to the naive approach. Each thread block would load a KV tile once and use it for scoring and then for the value aggregation phase, rather than reading the entire KV cache twice (once for scores, once for values) as the naive kernel does.
The kernel passed correctness tests with token-exact precision (max_abs error ~2e-8). Then came the microbenchmark.
The Data That Changed Everything
The benchmark results, visible in the preceding message (msg 12238), painted a sobering picture:
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 -At short prefixes (256–4096 tokens), the flash kernel was 40% slower than the naive baseline. At 16384 tokens, it achieved only a 0.9× speedup—essentially parity. The naive kernel couldn't even run at 65536 tokens (its shared memory requirement for the full score row exceeded the 100 KB limit), but the flash kernel's 94 ms latency was far from the promised improvement.
This is the moment captured in message 12239. The assistant has just received these numbers and must now diagnose why a theoretically superior design is failing in practice.
The Diagnosis: Occupancy Starvation
The assistant's reasoning begins with a clear-eyed analysis of the performance numbers. The core insight is delivered with characteristic precision:
"The core issue is that my flash design creates severe GPU underutilization. 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."
This is the occupancy trap. The flash kernel's design—one thread block per (batch, query_position, head_tile)—produces B × q_len × (H/HT) = 1 × 9 × (64/8) = 72 blocks. The naive kernel, by contrast, uses one block per (batch, query_position, head) = 1 × 9 × 64 = 576 blocks. The RTX PRO 6000 Blackwell has approximately 188 Streaming Multiprocessors. With 72 blocks, only about 38% of SMs can be active simultaneously. With 576 blocks, all SMs are fully occupied.
The flash kernel's theoretical advantage—reducing global memory reads by 16×—is completely negated by the fact that most of the GPU is idle while the few active blocks churn through their work. The naive kernel, despite reading the KV cache 16× more times, keeps all SMs busy and finishes faster through sheer parallelism.
The assistant identifies two additional micro-architectural issues compounding the problem. First, the score computation step uses only 128 of the 256 threads in each block, wasting half the available compute resources. Second, the per-head online update serializes across just 8 threads with synchronization barriers, creating a bottleneck that kills parallelism in the low-latency regime where occupancy matters most.
The Pivot: Flash-Decode with KV Splitting
The assistant's response to this diagnosis is not to tweak parameters but to fundamentally redesign the kernel's decomposition strategy. The fix is the standard "flash-decode" approach with KV splitting:
"The fix is to adopt a flash-decode approach with KV splitting: instead of one block per (batch, query_position, head_tile), split the KV dimension across multiple blocks so each handles a slice of the prefix, producing partial results that get reduced together."
This is a two-kernel design. The first kernel—the "partial" kernel—processes B × q_len × n_htiles × NSPLIT blocks, where each block handles a contiguous chunk of the KV sequence. Each block computes online softmax over its chunk, producing unnormalized accumulators and normalization statistics (the running maximum m and the log-sum-exp l). These partial results are written to a global scratch buffer. The second kernel—the "reduce" kernel—merges the NSPLIT partial results for each (batch, query, head) using the online-softmax combine logic: it finds the maximum m across splits, reweights and sums the accumulators, then normalizes.
The key insight is that KV splitting transforms the grid size from 72 blocks to thousands of blocks. With NSPLIT=32 (a reasonable value for long sequences), the grid becomes B × q_len × n_htiles × NSPLIT = 1 × 9 × 8 × 32 = 2304 blocks. Now all 188 SMs stay busy, and the kernel can approach the memory floor—the theoretical minimum number of global memory reads required to compute the attention.
The Trade-Off Analysis
The assistant's reasoning reveals a sophisticated understanding of the architectural trade-offs involved. The message contains an extended internal debate about block configuration:
"I'm reconsidering the block configuration with BLOCK=128 to improve thread utilization in the score step, where all 128 threads would be fully occupied computing score pairs. The PV and KV loading would still be efficient at 32 and 64 elements per thread respectively. However, this creates an occupancy bottleneck—the 70KB shared memory limits us to 1 block per SM, which hurts latency hiding despite the within-block streaming pipeline."
The tension is between two forms of occupancy. Grid-level occupancy—the number of thread blocks relative to the number of SMs—determines whether all SMs have work. SM-level occupancy—the number of active warps per SM—determines how well latency hiding works within each SM. The flash-decode approach fixes grid-level occupancy by generating enough blocks to cover all SMs. But SM-level occupancy remains a concern: with 70 KB of shared memory per block and ~100 KB available per SM, only one block can reside on each SM at a time, limiting the SM's ability to hide memory latency by switching between warps.
The assistant resists the temptation to over-optimize analytically:
"Rather than over-analyzing this analytically, I should make HT, TK, and BLOCK tunable parameters and empirically determine the best configuration, starting with BLOCK=128, HT=8, TK=16."
This is a pragmatic engineering decision. The assistant recognizes that the first-order problem—grid-level occupancy—will yield the largest gains, and that second-order tuning can follow once the basic architecture is proven.
The Memory Floor Calculation
A particularly impressive piece of reasoning is the assistant's calculation of the memory floor:
"at q=16384, naive takes 19.6ms while the memory floor (reading KV once) is ~8.6ms, so naive is roughly 2x off its memory pattern. Flash reads KV 16x less but measures 20ms due to occupancy starvation (only 72 blocks), not memory bandwidth—KV-split should fix occupancy and approach the memory floor."
This calculation reveals the assistant's understanding of what performance is possible given the hardware's memory bandwidth. The RTX PRO 6000 Blackwell has a peak memory bandwidth of approximately 1.8 TB/s. Reading the KV cache once for a 16384-token prefix with the MLA architecture's dimensions (Dl=512, Dr=64) requires moving about 16384 × (512 + 64) × 2 bytes (assuming bf16) ≈ 18.9 MB. At 1.8 TB/s, that's about 10.5 μs—but this is per query position, and the naive kernel reads the KV cache twice (once for scores, once for values) across all heads. The actual memory floor is more complex, but the assistant's estimate of ~8.6 ms for the full operation is a reasonable ballpark.
The flash kernel, even in its occupancy-starved state, reads the KV cache 16× fewer times than the naive kernel. If it could achieve the same effective bandwidth, it should be 16× faster. Instead, it's slower, because its effective bandwidth is catastrophically lower due to idle SMs. The KV-split fix aims to restore effective bandwidth by keeping all SMs busy, allowing the kernel to approach the memory floor.
Assumptions and Their Consequences
The message reveals several assumptions the assistant made that proved incorrect:
Assumption 1: Reduced memory traffic automatically translates to reduced latency. The flash kernel's design was predicated on the idea that reading the KV cache 16× fewer times would yield a proportional speedup. This ignored the occupancy dimension: the naive kernel's higher block count gave it better GPU utilization, compensating for its less efficient memory access pattern.
Assumption 2: The block count from head tiling would be sufficient. With HT=8, the assistant expected H/HT = 8 blocks per query position, times 9 query positions = 72 blocks. On a GPU with 188 SMs, this leaves most SMs idle. The assistant had not accounted for the fact that the RTX PRO 6000 Blackwell has more SMs than the flash kernel could generate blocks for.
Assumption 3: The thread utilization within each block was adequate. The score computation step used only 128 of 256 threads, and the per-head update serialized across 8 threads. These micro-architectural inefficiencies compounded the grid-level occupancy problem.
Assumption 4: Short-prefix performance was less important. The assistant notes that "short prefixes like 256 tokens are already fast enough at 199us, and production workloads care far more about long-context performance." While this is true for the specific deployment (which targets 200k context lengths), the short-prefix regime still matters for the initial tokens of any sequence and for workloads with mixed-length requests.
The Input Knowledge Required
To fully understand this message, the reader needs knowledge spanning several domains:
GPU Architecture: Understanding of Streaming Multiprocessors (SMs), thread blocks, warps, shared memory, global memory bandwidth, and the relationship between block count and SM occupancy. Familiarity with the Blackwell sm_120 architecture's constraints (100 KB shared memory, no wgmma/TMA/tcgen05 instructions) is essential.
Flash Attention Theory: Understanding of the tiled flash attention algorithm, online softmax, and the flash-decode variant with KV splitting. Knowledge of how attention is decomposed into score computation and value aggregation phases, and how KV splitting enables parallelization across the sequence dimension.
MLA (Multi-head Latent Attention): The specific attention mechanism used by the Kimi K2.6 and GLM models, which separates query/key/value into latent and rope components (Dl=512, Dr=64). Understanding how MLA's KV cache is structured and how verify attention operates on draft trees.
Speculative Decoding: Understanding of how draft models generate candidate token trees and how the verify attention kernel scores these candidates against the full model's KV cache. The distinction between draft tokens (tree nodes) and prefix tokens (KV cache entries).
CUDA Programming: Understanding of kernel launch parameters, grid dimensions, block dimensions, shared memory allocation, and the mechanics of CUDA graph capture. Familiarity with the constraints of writing kernels for architectures without access to cutting-edge tensor core instructions.
The Output Knowledge Created
This message creates several important contributions to the project's knowledge base:
1. The Occupancy Diagnosis: A clear explanation of why the tiled flash kernel underperforms, grounded in quantitative analysis of block counts, SM counts, and utilization. This diagnosis becomes the foundation for the KV-split redesign.
2. The KV-Split Architecture: A detailed design for a two-kernel flash-decode approach, including the partial kernel (which computes online softmax over KV chunks) and the reduce kernel (which merges partial results). The design includes adaptive NSPLIT selection based on KV length, workspace management for scratch buffers, and handling of edge cases like empty splits.
3. The Memory Floor Benchmark: A calculation of the theoretical minimum latency for the verify attention operation, providing a target for optimization. This establishes that the naive kernel is ~2× off the memory floor and that the flash kernel should be able to approach it with proper occupancy.
4. The Tuning Framework: A decision to make HT, TK, and BLOCK tunable parameters and determine optimal values empirically, rather than over-analyzing analytically. This sets up a systematic optimization process for subsequent rounds.
5. The Occupancy vs. Latency Trade-off Map: A nuanced understanding of the relationship between grid-level occupancy (blocks covering SMs) and SM-level occupancy (warps hiding latency), and how shared memory constraints create tension between the two.
The Thinking Process: A Window into Diagnostic Reasoning
The message's "Agent Reasoning" section provides a rare window into the assistant's cognitive process. Several patterns are worth highlighting:
Hypothesis testing: The assistant begins by stating the observed facts ("the flash kernel is underperforming") and immediately formulates a hypothesis ("severe GPU underutilization"). It then quantifies the hypothesis (72 blocks vs 188 SMs) and checks it against the data (0.6× at short prefixes, 1.5–1.8× at 16k).
Quantitative reasoning: Every claim is backed by numbers. The assistant calculates block counts (72 vs 576), SM counts (~188), utilization percentages (38%), memory floor estimates (8.6 ms), and shared memory budgets (70 KB vs 100 KB limit). This quantitative grounding prevents vague speculation.
Counterfactual thinking: The assistant considers alternative explanations and design choices. It asks whether a simpler occupancy fix would suffice, whether the complexity of KV splitting is justified, and whether the short-prefix regime deserves optimization attention. It rejects these alternatives with reasoned arguments.
Iterative refinement: The assistant does not treat the KV-split design as final. It explicitly notes that "I can add optimizations like cp.async double-buffering and further tuning as a follow-up once I measure the baseline improvement." This sets up a measurement-driven optimization loop.
Awareness of limitations: The assistant acknowledges the tension between analytical optimization and empirical tuning: "Rather than over-analyzing this analytically, I should make HT, TK, and BLOCK tunable parameters and empirically determine the best configuration." This reflects a mature understanding that GPU performance is often too complex to predict from first principles.
The Broader Significance
This message is significant beyond its immediate context because it illustrates a fundamental truth about GPU kernel optimization: theoretical advantages in memory traffic can be completely negated by poor occupancy. The flash attention literature often emphasizes the reduction in global memory reads, but this case demonstrates that the block count—and therefore the ability to keep all SMs busy—is equally important.
The message also showcases the value of benchmarking early and often. The assistant could have spent more rounds refining the tiled flash kernel, adding micro-optimizations, and chasing marginal gains. Instead, the benchmark data forced a rapid diagnosis and a fundamental architectural pivot. The 3–6× speedup ultimately achieved (as documented in chunk 1 of segment 66) vindicates this decision.
For anyone working on custom CUDA kernels for LLM inference, this message offers a cautionary tale: occupancy is not a detail to be optimized after correctness—it is a first-order design constraint that must shape the kernel's decomposition from the outset. The flash-decode KV-split approach that the assistant pivots to is now a standard technique precisely because it addresses this occupancy challenge by generating enough blocks to saturate modern GPUs with their hundreds of SMs.
Conclusion
Message 12239 captures a pivotal moment in the development of a custom sm_120 verify attention kernel. Confronted with benchmark data showing its carefully designed flash kernel underperforming the naive baseline, the assistant diagnoses the root cause as occupancy starvation—a mismatch between the kernel's block count and the GPU's SM count. The response is a fundamental architectural pivot to a KV-split flash-decode design that generates thousands of thread blocks instead of dozens, ensuring all SMs stay busy and the kernel can approach the memory bandwidth floor.
The message is a testament to the value of measurement-driven development, quantitative reasoning, and the willingness to abandon a design when the data demands it. It also serves as a detailed case study in GPU occupancy analysis, demonstrating how block count, shared memory constraints, and SM count interact to determine real-world performance. For anyone building custom CUDA kernels in the era of large language models, the lessons in this message are invaluable.