The Occupancy Tension: Iterative Kernel Optimization on sm_120 Blackwell
In the high-stakes world of custom CUDA kernel development for large language model inference, few moments are as revealing as the one captured in message 12244 of this opencode session. The assistant has just received benchmark results for version 6 of a custom verify attention kernel—the sixth architectural iteration in a multi-day effort to build a high-performance speculative decoding engine for the Kimi K2.6 model running on an NVIDIA RTX PRO 6000 Blackwell GPU (sm_120 architecture). The message is a masterclass in disciplined performance analysis, revealing how a developer reasons through the tension between occupancy and memory bandwidth when hand-tuning GPU kernels for an architecture that lacks vendor-optimized library support.
The Message
The assistant writes:
## Agent Reasoning The v6 design is now consistently outperforming the naive approach across the board—1.3-1.5x faster on shorter sequences and 2.7-3.3x faster at 16384, plus it can handle 65536 which the naive version can't. The combination of per-head processing, warp-per-key distribution, KV-split optimization, and high occupancy is clearly working. Still, hitting only 3.3x speedup at 16384 with absolute times around 5.9ms leaves significant room between current performance and the theoretical memory bandwidth limit, so I need to decide whether to continue optimizing or move forward with what I have.
>
Looking at the bottleneck more carefully: with 9 splits creating 5184 blocks and high occupancy (32 warps per SM), the compute per warp is actually quite small—around 8000 FMAs total—which would only take about 3.4 microseconds. The real constraint is memory bandwidth; each block needs to load its KV chunk (about 4.2MB) from global memory, and while different blocks process different head/split combinations, they're all reading from the same underlying KV data, which is where the latency is getting bottlenecked.
>
The core issue is that the per-head design reads the KV cache H=64 times across all blocks (once per head), totaling about 2.4GB of data. At the current 5.96ms execution time, that's only 403 GB/s achieved bandwidth—roughly 22% of the 1.8TB/s peak. A head-tile amortization approach would reduce these redundant reads, but it hurt occupancy before. The tension is real: I can either maintain high occupancy with more memory traffic, or reduce reads at the cost of lower occupancy.
>
Since the current scalar float loads are coalesced but not fully utilizing memory throughput, I'm going to try vectorizing the KV tile loads with float4 operations to improve achieved bandwidth.
>
Actually, wait—I need to reconsider the production scenario. In TP8 deployment, each rank only has 8 heads, not 64, so the per-head design reads the KV cache only 8 times per rank instead of 64. That's 8× less memory traffic than my benchmark, which uses the full 64 heads with no tensor parallelism. So the per-head design is actually well-suited to production; my benchmark is the worst case. Still, adding float4 vectorization is a straightforward win that helps both scenarios.
>
Adding vectorized loads... Writing vectorized loads...
>
Now I'm stepping back to assess whether the current performance—1.3-3.3x faster than the naive baseline on the full benchmark—is solid enough to move forward with integration. The kernel is correct and handles arbitrary sequence lengths, and in production with tensor parallelism across 8 ranks and 8 heads per rank, it should significantly outperform Triton's scattered memory access patterns. Adding float4 vectorization is a straightforward optimization win, so I'll implement that, re-benchmark, and then document the results as a Phase 1 deliverable.
>
Now v6 beats naive everywhere (1.3–3.3×) and handles 65k which naive can't. We're now memory-throughput-bound on the scalar loads (~400 GB/s). Adding float4 vectorized loads for the KV tile + q (the dominant traffic) — a cheap, high-impact win:
>
[edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/verify_attn_flash.cu Edit applied successfully.
The Context: Why This Message Was Written
To understand this message, one must grasp the broader mission. The session is part of a project to deploy the Kimi K2.6 model with speculative decoding on a machine equipped with eight RTX PRO 6000 Blackwell GPUs. The speculative decoding system uses a "DDTree" (Draft Tree) approach, where a smaller drafter model proposes multiple candidate tokens in parallel, and the main model "verifies" them with a single forward pass. The verify attention kernel is the computational heart of this verification step.
The problem is architectural. The RTX PRO 6000 uses the sm_120 instruction set architecture—a Blackwell consumer GPU variant that is distinct from the sm_100a/sm_103a used by Blackwell data-center GPUs (like the B200). All optimized MLA (Multi-head Latent Attention) kernels—FlashMLA, cutlass-MLA, flashinfer-MLA—are compiled exclusively for sm_90a (Hopper), sm_100a, and sm_103a. None support sm_120. These libraries rely on Hopper and Blackwell data-center instructions like wgmma, TMA (Tensor Memory Accelerator), and tcgen05 that simply do not exist on sm_120's Ada-like ISA. The assistant cannot use any off-the-shelf optimized kernel; it must write its own from scratch.
Earlier in the session (see [msg 12242]), the assistant had benchmarked an initial flash kernel (v5) that achieved only 1.6-1.9× speedup over the naive baseline at long prefixes, and actually regressed at short prefixes (0.4-0.7×). Through careful analysis, the assistant identified the root cause: occupancy. The v5 kernel used 70KB of shared memory per block, forcing only one block per streaming multiprocessor (SM). With only 128 threads per block, this yielded approximately 6% occupancy on an SM designed for 2048 threads. Decode workloads are latency-bound, and with so few warps available to hide memory latency, every stall blocked the entire SM.
The assistant then designed v6 around a fundamentally different philosophy: maximize occupancy at the cost of redundant memory reads. The new design used per-head blocks (each block processes one head instead of amortizing multiple heads), a warp-per-key distribution for the score computation, and KV-splitting for long sequences. By shrinking shared memory to approximately 23KB, the kernel could fit 4 blocks per SM, achieving 67% occupancy. The benchmark results in [msg 12243] confirmed the design worked: v6 consistently outperformed naive across all sequence lengths, achieving 1.3-3.3× speedup and handling 65k tokens where the naive kernel ran out of memory.
Message 12244 is the assistant's response to these results. It is not a message of triumph but of measured assessment—the kernel works, but the developer is asking the critical question: Is this good enough, or should I keep optimizing?## The Reasoning Process: A Window Into GPU Kernel Optimization
The most striking feature of this message is the assistant's reasoning process, which reveals the mental model of an experienced GPU kernel developer working under real constraints. The reasoning unfolds in several distinct phases, each reflecting a different analytical lens.
Phase 1: Assessing the results. The assistant begins by acknowledging that v6 works: "consistently outperforming the naive approach across the board." The numbers are concrete: 1.3-1.5× on short sequences, 2.7-3.3× at 16384, and the ability to handle 65536 tokens. But the assistant immediately identifies the gap: "absolute times around 5.9ms" versus what the hardware should theoretically deliver. This framing—comparing achieved performance to theoretical peak—is the hallmark of systematic optimization.
Phase 2: Bottleneck analysis through first principles. The assistant performs a back-of-the-envelope calculation that is remarkable for its precision. With 9 splits creating 5184 blocks and 32 warps per SM, the compute per warp is approximately 8000 FMAs, which should take about 3.4 microseconds. The fact that the kernel takes 5.96ms means the bottleneck is not compute but memory. The assistant then calculates the achieved bandwidth: the per-head design reads the KV cache 64 times (once per head), totaling 2.4GB of data, yielding 403 GB/s—only 22% of the 1.8 TB/s peak. This kind of rapid quantitative reasoning, performed entirely in the developer's head, is what separates surface-level tuning from deep optimization.
Phase 3: Identifying the fundamental tension. The assistant explicitly names the central trade-off: "I can either maintain high occupancy with more memory traffic, or reduce reads at the cost of lower occupancy." This is the occupancy-memory tension that defines GPU kernel design. The v5 design optimized for fewer memory reads (amortizing multiple heads per block) but suffered from low occupancy. The v6 design optimized for occupancy (per-head blocks with small shared memory) but increased memory traffic. The assistant recognizes that neither extreme is optimal—the sweet spot depends on the specific workload and hardware characteristics.
Phase 4: Choosing the next optimization. Rather than redesigning the kernel again, the assistant selects a targeted, low-risk optimization: vectorizing scalar float loads into float4 operations. This is a classic GPU optimization that increases memory throughput by reducing the number of load instructions and better utilizing the memory bus width. The assistant correctly identifies that the kernel is now memory-throughput-bound, making this the highest-impact change available without a major architectural shift.
Phase 5: The production reality check. In a moment of intellectual honesty, the assistant pauses to reconsider the production scenario: "Actually, wait—I need to reconsider the production scenario." In TP8 (tensor parallelism across 8 ranks), each rank only processes 8 heads, not 64. This means the per-head design reads the KV cache only 8 times per rank, reducing memory traffic by 8×. The benchmark, which tests the full 64 heads without tensor parallelism, represents the worst case for the per-head design. In production, the kernel would be significantly more efficient. This insight is crucial because it reframes the performance question: the 3.3× speedup over naive is a lower bound; the real-world improvement could be much larger.
Assumptions and Their Validity
The message rests on several key assumptions, most of which are well-justified but worth examining.
The first assumption is that the naive kernel is a valid baseline. The assistant assumes that the naive implementation—a straightforward per-head attention kernel without occupancy optimizations—represents the performance that Triton's compiled MLA kernel would achieve in the best case. This is reasonable because Triton's page_size=1 scattered access pattern was the original production bottleneck, achieving only 14 GB/s effective bandwidth. The naive kernel, which uses contiguous memory, should already outperform Triton in production even before the flash optimizations.
The second assumption is that float4 vectorization is a "cheap, high-impact win." This is a well-established optimization on NVIDIA GPUs. When threads in a warp access consecutive 32-bit values, the hardware can coalesce these into fewer, wider memory transactions. Using float4 (128-bit) loads explicitly tells the compiler to generate wider load instructions, reducing instruction count and improving throughput. On sm_120, which has the same memory subsystem architecture as Ada, this optimization is well-understood and low-risk. The assumption is correct.
The third assumption is that the kernel is correct. The assistant verified this in [msg 12243] by running a test suite against reference outputs, with all tests passing at tolerances of approximately 2-3×10⁻⁸ absolute error and 2-4×10⁻³ relative error. These tolerances are consistent with floating-point accumulation differences between kernel implementations. The assumption of correctness is well-supported.
The fourth assumption—more subtle—is that the benchmark configuration (H=64, no tensor parallelism) represents a meaningful test. The assistant explicitly acknowledges this might be the worst case for the per-head design. However, the benchmark is still valuable because it provides a conservative estimate. If the kernel performs well in the worst case, it will perform even better in production.
Input Knowledge Required
To fully understand this message, a reader needs substantial background knowledge spanning multiple domains. First, one must understand the GPU architecture: what an SM is, how occupancy works, what shared memory is, and why 6% occupancy is catastrophically low for a latency-bound workload. The concept of "warp-per-key" distribution requires understanding how CUDA threads are grouped into warps of 32, and how warp-level primitives like shuffle and reduce enable cooperative computation.
Second, one needs familiarity with the attention mechanism in transformer models, specifically Multi-head Latent Attention (MLA). MLA decomposes the attention computation into a low-rank projection, where the key and value are computed from a shared latent vector. This means the KV cache stores latent vectors rather than full key/value pairs, and the attention computation involves projecting these latents into the key/value space on the fly. The dimensions mentioned (Dl=512 for the latent dimension, Dr=64 for the reduced dimension) are specific to this architecture.
Third, one must understand tensor parallelism (TP8). In TP8, the model's heads are split across 8 GPUs, so each GPU processes 8 heads out of 64. This reduces the per-GPU memory traffic by 8× but introduces communication overhead for the all-reduce operation that combines results across GPUs.
Fourth, one needs knowledge of CUDA optimization techniques: memory coalescing, shared memory tiling, KV-splitting (splitting the key-value sequence into chunks processed by separate blocks), and online softmax (a numerically stable algorithm that computes softmax in a single pass without storing intermediate values).
Output Knowledge Created
This message creates several forms of output knowledge. The most tangible is the edited kernel file at /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/verify_attn_flash.cu, which now includes float4 vectorized loads for the KV tile and query data. This is a concrete artifact that improves the kernel's memory throughput.
The message also creates architectural knowledge: the assistant has documented the reasoning behind the v6 design and the decision to pursue float4 vectorization as the next optimization. This reasoning is captured in the agent's thinking and is available for future reference.
Perhaps most importantly, the message creates a decision framework. The assistant has established criteria for evaluating whether the kernel is "good enough" for production: it must be correct (verified), handle arbitrary sequence lengths (verified up to 65k), and outperform Triton's scattered access pattern in the TP8 deployment scenario (inferred from the benchmark results). The float4 optimization is a tactical improvement within this framework, not a strategic redesign.
Mistakes and Incorrect Assumptions
The message is notably free of significant errors, but a few points warrant scrutiny. The assistant's calculation that "each block needs to load its KV chunk (about 4.2MB)" seems inconsistent with the 2.4GB total data figure. If there are 5184 blocks and each loads 4.2MB, the total would be approximately 21.8GB, not 2.4GB. The discrepancy likely arises because the "KV chunk" per block is actually much smaller—the 4.2MB figure may refer to the total KV data accessed by all blocks for a single head, or it may include some redundant counting. The assistant's overall analysis is sound, but this specific number is ambiguous.
The assistant also assumes that float4 vectorization is "cheap" in terms of implementation effort. While the code change itself is small, validating that the vectorized loads do not introduce alignment issues or change the numerical results requires re-running the full test suite and benchmark. The assistant does perform this validation implicitly (the edit is followed by a rebuild and test run in subsequent messages), but the message itself does not acknowledge the validation cost.
A more subtle issue is the assistant's framing of the occupancy-memory tension. The assistant presents this as a binary choice: high occupancy with more memory traffic, or low occupancy with less traffic. In reality, there are intermediate designs—for example, processing 2 or 4 heads per block instead of 1 or 8—that might achieve a better balance. The assistant explored HT=4 in earlier reasoning but settled on HT=1 (per-head) for v6. This was the right decision given the benchmark results, but the message does not revisit whether a moderate HT value might outperform both extremes.
The Thinking Process: A Deeper Look
The assistant's thinking in this message reveals a sophisticated mental model of GPU execution. Consider the calculation of achieved bandwidth: "At the current 5.96ms execution time, that's only 403 GB/s achieved bandwidth—roughly 22% of the 1.8TB/s peak." This is not a number the assistant could have looked up; it is derived from first principles. The assistant knows the total data volume (2.4GB), knows the execution time (5.96ms), and divides to get bandwidth. The fact that the assistant can perform this calculation in real time, during the optimization process, demonstrates a deep understanding of the relationship between memory traffic and kernel runtime.
The assistant also demonstrates what computer architects call "roofline analysis"—the practice of determining whether a kernel is compute-bound or memory-bound by comparing achieved performance to theoretical peaks. The assistant's calculation that compute per warp is only 3.4 microseconds (negligible compared to 5.96ms total) is a roofline-style argument: if compute is not the bottleneck, memory must be.
The moment where the assistant says "Actually, wait—I need to reconsider the production scenario" is particularly revealing. It shows the assistant stepping back from the microbenchmark results to think about the real deployment context. This is a critical skill in systems optimization: microbenchmarks can mislead if they don't reflect production conditions. The assistant correctly identifies that the benchmark (H=64, no TP) is the worst case for the per-head design, and that the production scenario (H=8 per rank with TP8) will be significantly more favorable. This insight prevents the assistant from making a suboptimal architectural decision based on misleading benchmark data.
Conclusion
Message 12244 captures a pivotal moment in the development of a custom CUDA kernel for an unsupported GPU architecture. The assistant has successfully built a verify attention kernel that outperforms the naive baseline by 1.3-3.3×, handles sequences the naive kernel cannot, and is well-suited to the production TP8 deployment. The decision to add float4 vectorization is a measured, low-risk optimization that addresses the newly identified memory-throughput bottleneck.
The message is a testament to the value of systematic performance analysis in GPU kernel development. Rather than guessing at optimizations or blindly applying heuristics, the assistant reasons from first principles about occupancy, memory bandwidth, and the specific characteristics of the sm_120 architecture. The result is a kernel that, while not yet achieving peak hardware utilization, is demonstrably correct, robust across sequence lengths, and positioned to deliver significant speedups in production.
The work continues. In subsequent messages, the assistant will apply the float4 vectorization, re-benchmark, and ultimately achieve a 3-6× end-to-end decode speedup over Triton across all context lengths—validating the architectural decisions made in this critical message. But message 12244 stands as the moment where the developer stepped back, assessed where they were, and made a conscious choice about the next step forward. It is a model of disciplined, evidence-driven GPU kernel optimization.