The Occupancy Epiphany: How a CUDA Kernel Designer Discovered That GPU Utilization Trumps Memory Bandwidth

Introduction

In the high-stakes world of custom CUDA kernel development for large language model inference, few moments are as decisive as the one captured in message 12242 of this coding session. The assistant, having just benchmarked its fourth iteration of a custom verify-attention kernel for the RTX PRO 6000 Blackwell GPU (sm_120), stares at a table of disappointing numbers. The KV-split optimization—a standard technique borrowed from flash-attention literature—has delivered only a 1.6–1.9× speedup over the naive baseline at long context lengths, and actually regressed performance at short prefixes. The kernel is taking 11.6 milliseconds for a computation that, by the assistant's own back-of-the-envelope calculation, should theoretically complete in 0.22 milliseconds. That is a 50× gap between theoretical peak and measured performance.

This message is the story of how the assistant diagnosed that gap, traced it to its root cause, and made a fundamental architectural pivot that would eventually unlock a 3–6× end-to-end speedup over the production Triton baseline. It is a masterclass in GPU kernel optimization methodology: the disciplined use of back-of-the-envelope math, the courage to abandon a design that is theoretically elegant but practically broken, and the willingness to let empirical data override architectural intuition.

Context: The Problem and the Stakes

To understand why this message matters, we must first understand what the assistant is trying to build. The broader project is deploying the Kimi K2.6 model with DFlash speculative decoding on NVIDIA RTX PRO 6000 Blackwell GPUs. The "verify attention" kernel is the computational heart of the speculative decoding process: it computes the attention scores between a set of draft tokens (the "queries") and the full key-value cache, determining which draft tokens to accept and which to reject. This kernel runs on every decode step, so its latency directly determines the throughput of the entire inference pipeline.

The production system was using a Triton-based MLA (Multi-head Latent Attention) kernel, but it suffered from a critical performance pathology: with page_size=1 in the vLLM paged attention implementation, the Triton kernel performed scattered gather memory accesses that achieved only ~14 GB/s effective bandwidth—roughly 130× below the GPU's 1.8 TB/s peak. The assistant's mission was to write a custom CUDA kernel that would eliminate this scatter penalty by using contiguous memory access patterns, and ideally deliver a dramatic speedup.

The assistant had already gone through several design iterations. The first flash-attention kernel (v1–v3) used a head-tile amortization strategy: each thread block processed 8 heads simultaneously, sharing the KV cache loads across heads to reduce memory traffic. This design was theoretically appealing—it minimized the number of times the KV cache had to be read—but it had a fatal flaw that the benchmarks now made undeniable.

The Moment of Truth: Confronting the Benchmark Data

The message opens with the assistant staring at the benchmark results from the KV-split optimization (v4):

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        -

The numbers tell a sobering story. At short prefixes (256 tokens), the flash kernel is half the speed of the naive baseline. At 4096 tokens, it's still 30% slower. Only at 16384 tokens does it eke out a 1.7× advantage, and even then, it's taking 11.6 milliseconds—far too slow for production use. The KV-split optimization helped at long prefixes (cutting latency roughly in half from the pre-split ~20ms), but it's still orders of magnitude away from the theoretical potential.

What makes this moment particularly compelling is the assistant's reaction. Rather than reaching for more sophisticated optimization techniques—tensor core instructions, warp-level matrix multiply, or double-buffered shared memory pipelines—the assistant does something more fundamental: it reaches for a calculator.

The Back-of-the-Envelope Calculation That Changed Everything

The assistant's reasoning in this message is a textbook example of how to diagnose GPU kernel performance problems. It starts with a simple question: what should the theoretical minimum runtime be?

The computation involves approximately 10 billion floating-point operations (10^10 FMAs). The RTX PRO 6000 Blackwell has roughly 90 TFLOP/s of fp32 compute throughput. At peak compute utilization, the kernel should complete in 0.22 milliseconds. The kernel is taking 11.6 milliseconds. That's a 50× gap between theory and practice.

The assistant then considers the memory bandwidth angle. Even accounting for reading the KV cache multiple times across head tiles, the memory traffic should only account for a fraction of a millisecond. The conclusion is inescapable: the bottleneck is neither compute nor memory bandwidth in the traditional sense. It is occupancy—the GPU's inability to keep enough warps in flight to hide the latency of individual instructions.

This diagnostic step is crucial because it reframes the entire optimization problem. Many GPU kernel developers default to optimizing for memory bandwidth, assuming that the bottleneck is always loading data from global memory. But the assistant correctly recognizes that when the gap between theoretical and actual performance is 50×, the problem is not about moving data faster—it's about keeping the GPU's execution units fed with work.

The Occupancy Analysis: A Deep Dive into GPU Architecture

The assistant's analysis of the occupancy problem is remarkably thorough. The flash kernel uses 70 KB of shared memory per thread block. On the RTX PRO 6000's Streaming Multiprocessors (SMs), which have a finite shared memory capacity, this means only one block can reside on each SM at a time. With 128 threads per block, and an SM designed to support up to 2048 threads (64 warps), the kernel achieves only 6% occupancy.

To understand why this is catastrophic, consider how a GPU SM works. When a warp issues an instruction that depends on a memory load (from global memory, shared memory, or registers), the warp stalls until the data arrives. The SM hides this latency by switching to another warp that is ready to execute. If there are many warps available, the SM can keep its execution units busy even when individual warps stall frequently. But with only 4 warps per SM (6% occupancy), every stall becomes a visible gap in utilization. The SM spends most of its time waiting, not computing.

The assistant's insight is that decode is fundamentally latency-bound, not bandwidth-bound. In a batch-size-1 decode scenario, there is only one query to process at a time. The GPU cannot amortize latency across many independent pieces of work. Every microsecond spent waiting for a memory load is a microsecond added directly to the end-to-end latency. In this regime, occupancy matters far more than minimizing the total number of memory loads.

This is a subtle but critical distinction. The head-tile amortization strategy (HT=8) was designed to reduce the total number of KV cache reads by sharing them across 8 heads. It succeeded at that goal—the flash kernel reads the KV cache 16× fewer times than the naive kernel. But the price of that reduction was 70 KB of shared memory per block, which crushed occupancy. The naive kernel, with its tiny shared memory footprint, could fit many more blocks per SM and thus achieved much higher occupancy. The naive kernel's "wasteful" memory reads were actually faster because the GPU never stalled waiting for them—there were always other warps ready to execute.

Three Paths Forward: The Design Space Exploration

Having identified occupancy as the root cause, the assistant systematically explores three possible solutions:

Path 1: Increase threads per block. By raising BLOCK from 128 to 256 or 512, more warps would be active even with only one block per SM. This is the simplest change but has limited upside—going from 4 warps to 8 or 16 warps still leaves the SM far below its 64-warp capacity.

Path 2: Reduce shared memory to fit multiple blocks per SM. By shrinking the shared memory footprint below 49 KB, the kernel could fit 2 blocks per SM, doubling occupancy. Below 33 KB, it could fit 3 blocks. This is more promising but requires sacrificing the head-tile amortization that was the original design's raison d'être.

Path 3: Restructure the dot product computation. Instead of having one thread compute a 576-element dot product serially, use a full warp cooperatively, with each thread handling a slice of the embedding dimension and then performing a warp-level reduction. This would increase thread utilization during the score computation phase.

The assistant's reasoning process is notably iterative and self-correcting. It starts by considering Path 1 (bumping BLOCK to 256), then realizes that the real win comes from Path 2 (shrinking shared memory), and finally recognizes that Path 3 (warp-per-key dot product) is necessary to keep all threads busy in the per-head design. The synthesis of these three paths produces the v5 architecture: per-head blocks (no head-tile amortization), small shared memory (~23 KB), warp-per-key dot product computation, and KV-splitting for long sequences.

The Per-Head Design: Trading Memory Traffic for Occupancy

The assistant's decision to abandon head-tile amortization is the most consequential choice in this message. It represents a fundamental rethinking of what matters in this specific workload.

The head-tile design (HT=8) was elegant: each block processes 8 heads simultaneously, loading the KV cache once and reusing it across heads. This minimizes the total memory traffic. But it requires storing 8 heads' worth of query vectors, score accumulators, and output accumulators in shared memory, bloating the footprint to 70 KB and limiting occupancy to one block per SM.

The per-head design (HT=1) is almost comically simple by comparison: each block processes one head. The KV cache is loaded once per head, meaning the total memory traffic increases by 8× (or more precisely, by a factor of H/HT where H is the total number of heads). But the shared memory footprint drops to ~23 KB, allowing 4 blocks per SM and achieving 67% occupancy—an 11× improvement.

The assistant's key insight is that in the latency-bound decode regime, the cost of additional memory traffic is dwarfed by the benefit of higher occupancy. A warp that stalls waiting for a memory load is only a problem if there are no other warps to run. With 4 blocks per SM and 256 threads per block, there are 32 warps available—enough to hide most memory latency. The additional KV cache reads are "free" in the sense that they happen while other warps are computing, and the GPU's memory subsystem can sustain the bandwidth if the pipeline is kept full.

This is a counterintuitive result that runs against the conventional wisdom of flash-attention design. The original flash-attention papers emphasize minimizing memory reads as the primary optimization goal. But those papers assume a high-occupancy regime where memory bandwidth is the bottleneck. In the low-occupancy regime of batch-1 decode on a consumer Blackwell GPU, the bottleneck is latency hiding, not bandwidth. The assistant correctly identifies that the optimization objective has changed.

The Warp-Per-Key Dot Product: Solving the Thread Utilization Problem

Having committed to the per-head design, the assistant immediately identifies a new problem: thread utilization during the score computation phase. With 256 threads per block and only 16 keys per tile, a naive implementation would activate only 16 threads (one per key) while the remaining 240 threads sit idle at a synchronization barrier. This would create a serial bottleneck that defeats the purpose of higher occupancy.

The assistant's solution is the warp-per-key dot product. Instead of assigning one thread per key, it assigns one warp (32 threads) per key. Each thread in the warp handles 18 elements of the 576-element embedding dimension, computing a partial dot product. The warp then performs a reduction (using __shfl_xor_sync or similar warp-level primitives) to sum the partial results into the final score. With 8 warps covering 8 keys per tile iteration, and two tile iterations to cover all 16 keys, all 256 threads remain active throughout the score computation phase.

This design is elegant because it solves two problems simultaneously. First, it keeps all threads busy, maximizing the utilization of each block's compute resources. Second, it distributes the dot product computation across many threads, reducing the serial work per thread from 576 FMAs to 18 FMAs plus a warp reduction. The warp reduction is essentially free (a few cycles) compared to the serial FMA loop.

The assistant's reasoning about this design choice reveals a deep understanding of GPU architecture. It recognizes that warp-level primitives are fundamentally faster than shared-memory-based reductions because they use the SM's register crossbar rather than the shared memory pipeline. It also recognizes that the optimal tile size (TK=8, not TK=16) is determined by the occupancy constraint: smaller tiles mean fewer shared memory bytes per tile, which means more blocks per SM, which means higher occupancy. The tradeoff is more tile iterations (more loop overhead), but the occupancy win more than compensates.

The Production Awareness: Thinking Beyond the Microbenchmark

One of the most sophisticated aspects of this message is the assistant's awareness of how the kernel will behave in production versus in the microbenchmark. The microbenchmark tests the kernel with H=64 heads and contiguous KV memory. But in production, the model uses tensor parallelism across 8 GPUs (TP8), meaning each rank processes only H/8 = 8 heads. And the KV cache in production uses paged memory with potential scatter.

The assistant makes two critical observations:

First, the per-head design is actually better suited to the TP8 production scenario. With only 8 heads per rank, the per-head design reads the KV cache 8 times per decode step. The head-tile design (HT=8) would read it once. The difference is a factor of 8 in memory traffic, but with 67% occupancy versus 6%, the per-head design is likely to be faster despite reading more data. In the microbenchmark with 64 heads, the per-head design reads the KV cache 64 times, which exaggerates the memory traffic cost. The production scenario is more favorable.

Second, the microbenchmark uses contiguous KV memory, but the production system uses paged memory with potential scatter. The Triton kernel's scattered gather achieved only 14 GB/s. Even if the custom kernel achieves only modest absolute performance, it may still be dramatically faster than Triton in production because it avoids the scatter penalty entirely. The microbenchmark may be understating the real-world win.

This production awareness prevents the assistant from over-optimizing for the microbenchmark at the expense of real-world performance. It also provides a healthy reality check: the goal is not to write the world's fastest MLA kernel in absolute terms, but to beat the Triton baseline in the specific production configuration.

The Honest Assessment: Knowing When to Declare Victory

Perhaps the most impressive aspect of this message is the assistant's willingness to be honest about what the kernel can and cannot achieve. It explicitly acknowledges that "writing a peak flash-MLA kernel is a multi-day project" and that the current kernel is "50× off peak" compute utilization. It recognizes that getting close to peak performance would require proper tensor-core tiling, which is a full flash-attention implementation beyond the scope of this effort.

This honesty is not defeatism—it's strategic prioritization. The assistant identifies the key metric that matters for production: beating Triton's scattered memory access pattern. The custom kernel with contiguous memory access and high occupancy is likely to achieve that goal even if it's far from theoretical peak. The remaining optimization (tensor-core tiling, float4 vectorization, etc.) can be pursued in subsequent phases if needed.

The assistant also correctly identifies the next bottleneck after occupancy is fixed: the per-head design reads the KV cache H times, and at the current ~400 GB/s achieved bandwidth, this accounts for a significant fraction of the runtime. The float4 vectorization optimization (added in the following message) is a straightforward way to improve achieved bandwidth by using 128-bit memory instructions instead of scalar 32-bit loads. This is a "cheap, high-impact win" that can be implemented quickly without changing the kernel architecture.

The Implementation: From Design to Code

The message concludes with the assistant writing the v6 kernel implementation. The key parameters are:

The Broader Lessons: What This Message Teaches About GPU Kernel Design

This message is valuable not just for its specific technical content, but for the methodology it demonstrates. Several lessons emerge that are applicable to any GPU kernel optimization effort:

1. Always compute the theoretical minimum. Before diving into optimization, calculate what the kernel should cost in a perfect world. If the gap between theory and practice is more than 2–3×, the bottleneck is likely occupancy or algorithm choice, not micro-optimization.

2. Occupancy is not just a number—it's a constraint that shapes the entire design. Many GPU developers treat occupancy as a secondary concern, focusing first on memory traffic and compute efficiency. But occupancy determines the GPU's ability to hide latency, and in latency-bound workloads (like batch-1 decode), it can be the dominant factor.

3. Be willing to abandon elegant designs for practical ones. The head-tile amortization design was theoretically elegant—it minimized memory reads, which is the textbook optimization goal. But it was practically broken because it destroyed occupancy. The per-head design is almost embarrassingly simple by comparison, but it works better because it keeps the GPU busy.

4. Think about production, not just benchmarks. The microbenchmark is a useful tool, but it can mislead if it doesn't reflect production conditions. The assistant's awareness of the TP8 scenario and the scatter penalty prevented it from over-optimizing for the microbenchmark at the expense of real-world performance.

5. Know when to stop optimizing. The assistant correctly identifies that the kernel is "good enough" for its primary goal (beating Triton in production) and that further optimization (tensor-core tiling) is a multi-day project with diminishing returns. This is a mature engineering judgment that balances effort against impact.

The Assumptions and Their Validity

The assistant makes several assumptions in this message, most of which are sound but worth examining:

Assumption: Decode is latency-bound, not bandwidth-bound. This is correct for batch-1 decode on a high-bandwidth GPU like the RTX PRO 6000. The working set (KV cache) is too large to fit in any cache, so every access goes to global memory. With low occupancy, the GPU cannot hide the latency of these accesses. The assumption is validated by the 50× gap between theoretical and actual performance.

Assumption: Higher occupancy will improve performance even with more memory traffic. This is the central bet of the per-head design. It is a reasonable bet given the occupancy analysis, but it is not guaranteed. If the additional memory traffic from reading the KV cache H times (instead of H/HT times) saturates the memory bus, the kernel could become bandwidth-bound despite higher occupancy. The assistant implicitly assumes that the memory subsystem can sustain the additional traffic, which is plausible given that the achieved bandwidth (~400 GB/s) is well below the theoretical peak (1.8 TB/s).

Assumption: The warp-per-key dot product is faster than the serial dot product. This is almost certainly correct for the specific dimensions involved (576 elements, 256 threads). Warp-level reductions are extremely fast (a few cycles), and distributing the 576 FMAs across 32 threads reduces the serial work per thread from 576 to 18. The synchronization overhead of the warp reduction is negligible compared to the FMA savings.

Assumption: The microbenchmark's contiguous memory access pattern is representative of production. This is partially true. The custom kernel uses contiguous access, which is the same in both microbenchmark and production. But the production KV cache may have additional fragmentation or alignment constraints that the microbenchmark doesn't capture. The assistant acknowledges this limitation and correctly notes that the custom kernel's advantage over Triton may be larger in production due to Triton's scatter penalty.

The Mistakes and Incorrect Assumptions

No analysis would be complete without examining what the assistant got wrong. Several points in the reasoning are worth questioning:

The initial assumption that KV-split alone would fix the problem. The assistant implemented KV-split in v4 expecting a 5–20× speedup at 65k tokens, but got only 1.7×. The mistake was underestimating the severity of the occupancy problem. KV-split increases the number of blocks, but if each block still has 70 KB of shared memory, only one block per SM can run at a time. The blocks are serialized, not parallelized. The assistant correctly diagnosed this after seeing the data, but the initial expectation was overly optimistic.

The assumption that the naive kernel's performance is a useful baseline. The naive kernel is itself far from optimal (it's also occupancy-limited, just less so). Using it as a baseline sets a low bar. A better baseline would be the theoretical minimum (0.22 ms) or the Triton kernel's performance in the same microbenchmark. The assistant partially addresses this by noting that Triton's scatter penalty makes direct comparison difficult, but the microbenchmark numbers are still somewhat misleading.

The assumption that 4 blocks per SM is achievable. The assistant estimates ~23 KB of shared memory per block, which would theoretically allow 4 blocks per SM (100 KB / 23 KB ≈ 4.3). But shared memory is not the only constraint on occupancy—register file size, thread scheduling hardware, and the maximum number of blocks per SM also matter. On sm_120, the maximum blocks per SM may be lower than 4, especially with 256 threads per block. The actual occupancy may be closer to 2 or 3 blocks per SM. The assistant's estimate is optimistic.

The assumption that float4 vectorization is a "cheap, high-impact win." This is true in general, but the assistant doesn't account for the alignment requirements of 128-bit vector loads. If the KV cache is not 16-byte aligned, float4 loads will either fail or be silently broken into scalar loads by the compiler. The assistant would need to ensure proper alignment in the production KV cache allocation, which may require changes to the memory allocator.

The Knowledge Required to Understand This Message

To fully appreciate the assistant's reasoning, the reader needs a solid understanding of several concepts:

GPU architecture: The distinction between SMs, warps, threads, and blocks; the concept of occupancy and how it affects latency hiding; the shared memory hierarchy and its capacity constraints.

CUDA programming model: Thread block scheduling, shared memory allocation, synchronization primitives (__syncthreads()), warp-level primitives (__shfl_xor_sync()), and the execution model of grid-stride loops.

Attention mechanisms: The difference between standard multi-head attention and multi-head latent attention (MLA); the role of the KV cache in autoregressive decoding; the online softmax algorithm and its numerical stability properties.

Flash attention: The tiling strategy for partitioning the attention computation across thread blocks; the concept of KV-splitting for long sequences; the tradeoff between memory traffic and recomputation.

Performance analysis: The distinction between compute-bound, memory-bound, and latency-bound kernels; the use of roofline analysis to identify bottlenecks; the relationship between theoretical peak performance and achieved performance.

The Knowledge Created by This Message

This message produces several important outputs:

A clear diagnosis of the occupancy bottleneck. The assistant demonstrates that the 50× gap between theoretical and actual performance is due to low occupancy (6%), not memory bandwidth or compute throughput. This diagnosis is supported by back-of-the-envelope calculations and architectural analysis.

A concrete design for the v6 kernel. The per-head, small-smem, warp-per-key design is fully specified, with specific parameters (BLOCK=256, HT=1, TK=8) and a detailed shared memory budget. This design is ready for implementation.

A methodology for GPU kernel optimization. The message demonstrates a disciplined approach to performance debugging: compute theoretical minimums, measure the gap, identify the bottleneck, generate hypotheses, test them, and iterate. This methodology is transferable to any GPU kernel optimization effort.

A strategic framework for prioritizing optimization work. The assistant distinguishes between optimizations that are "worth it" (occupancy fix, float4 vectorization) and those that are "multi-day projects" (tensor-core tiling). This framework helps avoid the trap of over-optimizing for marginal gains.

Production-aware performance analysis. The assistant considers how the kernel will behave in the TP8 production scenario and correctly identifies that the microbenchmark may understate the real-world advantage over Triton. This prevents premature optimization for the wrong workload.

The Thinking Process: A Window into Expert Reasoning

The assistant's reasoning in this message is notable for its structure and self-awareness. Several patterns are worth highlighting:

Iterative refinement. The assistant doesn't arrive at the v6 design in a single leap. It starts with the v4 design, identifies the occupancy problem, considers three paths forward, explores the tradeoffs of each, and gradually converges on the per-head + warp-per-key design. Each step builds on the previous one, and the assistant is willing to backtrack when an idea doesn't pan out (e.g., the HT=4, TK=16 config is considered and then abandoned in favor of HT=1, TK=8).

Quantitative reasoning. Every claim is supported by numbers. The assistant computes theoretical minimums (0.22 ms), measures actual performance (11.6 ms), calculates the gap (50×), and uses these numbers to drive decisions. When considering shared memory budgets, it breaks down the allocation component by component (4 KB for KV tile, 512 bytes for key embeddings, etc.) and sums them to verify feasibility.

Explicit tradeoff analysis. The assistant doesn't just list options—it analyzes the tradeoffs between them. The central tradeoff (occupancy vs. memory traffic) is explored in detail, with the assistant correctly concluding that occupancy matters more in the latency-bound decode regime. This tradeoff analysis is what separates expert reasoning from novice guesswork.

Self-correction. The assistant catches itself several times during the reasoning process. It starts to design a complex warp-level reduction, then realizes it's overcomplicating things and steps back. It considers the HT=4 config, then realizes HT=1 is better. It initially focuses on memory bandwidth, then corrects to focus on occupancy. This self-correction is a hallmark of expert problem-solving.

Awareness of limitations. The assistant is honest about what it doesn't know. It acknowledges that "writing a peak flash-MLA kernel is a multi-day project" and that the current kernel is "50× off peak." It recognizes that the microbenchmark may not reflect production conditions. This honesty prevents overconfidence and keeps the optimization effort focused on what matters.

Conclusion: The Turning Point

Message 12242 is the turning point of the entire kernel development effort. Before this message, the assistant was pursuing a design that was theoretically elegant but practically broken. After this message, the assistant has a clear diagnosis of the problem, a concrete plan for fixing it, and a strategic framework for prioritizing further optimization.

The v6 kernel that emerges from this reasoning process would go on to achieve a 3–6× end-to-end decode speedup over the Triton baseline, with CUDA graph support, KV defragmentation, and production deployment on 8 GPUs. But those achievements are downstream of the fundamental insight captured in this message: that in the latency-bound regime of batch-1 decode on a consumer Blackwell GPU, occupancy trumps memory bandwidth as the primary optimization objective.

This insight is counterintuitive. It runs against the conventional wisdom of flash-attention design, which emphasizes minimizing memory reads. It requires a deep understanding of GPU architecture and a willingness to let empirical data override theoretical elegance. And it demonstrates that the most important optimization tool is not a profiler or a microbenchmark, but a clear understanding of what the hardware is actually doing and why.

The assistant's reasoning in this message is a masterclass in GPU kernel optimization methodology. It shows that the path to peak performance is not a straight line of incremental improvements, but a winding road of diagnosis, hypothesis, experiment, and iteration. And it proves that sometimes the best optimization is not adding more complexity, but removing it—trading elegant amortization for brute-force occupancy, and letting the GPU's massive parallelism do the rest.