The Occupancy Epiphany: How One Environment Variable Unlocked 3× Throughput in a Custom CUDA Verify Kernel

Introduction

In the high-stakes world of speculative decoding on cutting-edge hardware, performance debugging often resembles detective work: you gather clues from profilers, formulate hypotheses, test them against reality, and iterate. Message <msg id=12328> captures a pivotal moment in such a investigation—a moment where a single environment variable change (NSPLIT 16→64) delivered a dramatic 2.1–3.3× speedup over a Triton baseline, and the assistant faced a classic engineering tradeoff: squeeze more performance from the current optimization, or pivot to the next bottleneck the user had explicitly requested.

This message is remarkable not for its code changes—the actual edit is a single line in a CUDA kernel file—but for the reasoning process it reveals. The assistant is standing at a crossroads, weighing the marginal return on investment (ROI) of vectorized bf16 loads against the user's explicit request for defragmentation and marshaling optimizations. The thinking is iterative, self-correcting, and deeply informed by the hardware architecture of the RTX PRO 6000 Blackwell GPU (sm_120). To understand this message, one must understand the months-long journey that led to it: a custom verify attention kernel, built from scratch because no existing MLA kernel supported sm_120, that had been diagnosed as occupancy-starved in the tensor-parallelism-8 (TP8) regime.

The Broader Context: A Custom Kernel for an Unsupported Architecture

The story begins with a fundamental hardware mismatch. The assistant was deploying the Kimi K2.6 model with DFlash speculative decoding on RTX PRO 6000 Blackwell GPUs—a consumer-grade Blackwell variant with an sm_120 architecture. Every optimized MLA (Multi-head Latent Attention) kernel in existence—FlashMLA, cutlass-MLA, flashinfer-MLA—was compiled exclusively for sm_90a (Hopper), sm_100a, or sm_103a (Blackwell DC). None supported sm_120, which uses an Ada-like instruction set lacking the Hopper/Blackwell-DC instructions (wgmma, TMA, tcgen05) that these kernels depend on.

The assistant had therefore built an owned verify attention kernel from scratch, implementing a KV-split flash-decode design in CUDA. The kernel passed correctness tests against a naive oracle, but initial microbenchmarks showed it was slower than the naive implementation at short prefixes (0.4–0.7×) and only modestly faster at longer ones (1.7–2.0×). The real shock came when the kernel was deployed in the live SGLang service: at 46k context length, decode throughput was a dismal 0.7 tokens per second, with GPU tensor core utilization at only ~3% during decode, despite 99.8% SM occupancy reported by nvidia-smi.

This paradox—high occupancy but near-zero utilization—was the key clue. The assistant initially suspected memory fragmentation (scattered KV slots from the paged allocator), then CPU orchestration overhead (tree building, mask construction), and even the draft model's forward pass. Each hypothesis was systematically tested and ruled out through instrumentation. The profiler revealed that the _prepare_for_speculative_decoding phase (draft forward + tree build) took only ~7ms, while the verify forward pass consumed ~185ms—and this was on a freshly restarted service with contiguous KV slots, ruling out fragmentation.

The Occupancy Revelation

The breakthrough came when the assistant recalculated the kernel's grid dimensions for the TP8 regime. In the single-GPU microbench, the kernel had 64 query heads (H=64), giving a grid of 1 × 9 × 64 × 16 = 9216 blocks—plenty to keep all 188 SMs busy. But in the TP8 deployment, each rank had only 8 heads (H=8), reducing the grid to 1 × 9 × 8 × 16 = 1152 blocks. With 188 SMs, that's only about 6 waves of blocks—far too few to hide memory latency. Each block still had to process ~2875 tokens per split, so the per-block work was large, but the overall number of blocks was insufficient to saturate the memory system.

The fix was deceptively simple: increase NSPLIT from 16 to 64. This multiplied the grid by 4× (to 4608 blocks), giving roughly 24 waves—enough to keep the SMs fed with work and hide memory latency. The assistant had initially lowered NSPLIT to 16 to fit within the CUDA graph capture workspace, but with H=8, the per-block workspace was 8× smaller, leaving ample room for more splits.

The result was immediate and dramatic. Without any kernel rebuild—just an environment variable change in the systemd service file—throughput jumped from 62 to 82 tok/s at 4k context, from 12.2 to 22.4 tok/s at 16k, and from 3.3 to 5.0 tok/s at 65k. The 2.1–3.3× speedup over the Triton+graph baseline validated the occupancy hypothesis conclusively.

The Central Dilemma: Vectorized Loads or Defrag?

Message <msg id=12328> opens with this victory fresh in mind. The assistant is "weighing whether to push further with vectorized bf16 loads (which should improve bandwidth efficiency) or try higher NSPLIT values for better occupancy at longer contexts." The reasoning immediately reveals a tension: the assistant knows that vectorized loads might yield only a modest 1.2–1.5× gain ("scalar bf16 loads already get decent hardware coalescing since consecutive threads read consecutive addresses"), while the user has explicitly requested defragmentation and marshaling optimizations.

What follows is a remarkable internal debate, spanning multiple paragraphs, as the assistant cycles through different prioritization schemes:

1. First impulse: "Let me start with vectorized bf16 loads since they should give better per-block efficiency."
2. Second thought: "Actually, let me step back. I've already hit 2.1-3.3x speedup with NSPLIT tuning, which was free and just an environment change. The user asked for defrag and marshaling, so let me focus."
3. Third reconsideration: "I'm wondering if NSPLIT=128 could push even better long-context performance... but each test takes 6 minutes to restart. Let me be pragmatic."
4. Fourth re-evaluation: "Now I'm weighing the effort against the gains—the NSPLIT improvement was substantial and free, but vectorized loads would only give maybe 1.3x more. Given the rebuild and restart overhead, I should focus on what the user actually needs."
5. Fifth pivot: "Actually, the ROI on vectorized loads is worth it—at NSPLIT=64 the verify is around 190ms, and a 1.3x speedup would push throughput from 5 to 6.5 tokens per second. Let me implement the vectorized load." This back-and-forth is not indecision; it is the hallmark of a disciplined engineer weighing multiple competing objectives against concrete data. Each cycle incorporates new information: the magnitude of the NSPLIT gain, the user's stated priorities, the cost of kernel rebuilds and service restarts, the estimated ROI of vectorized loads, and the difficulty of measuring defrag benefits on a fresh pool.

Assumptions Underlying the Reasoning

The assistant's reasoning rests on several key assumptions, some explicit and some implicit:

1. The verify kernel remains the dominant bottleneck. The assistant assumes that after the NSPLIT improvement, the verify forward pass (now ~190ms at 46k context) still dominates the 192ms decode step, and that further optimizing it will directly translate to end-to-end throughput gains. This is reasonable given the profiler data showing _prepare_for_speculative_decoding at only ~7ms.

2. Vectorized bf16 loads will yield a meaningful speedup. The assistant estimates 1.2–1.5×, later refining to 1.3×. This is an educated guess based on the difference between scalar and vectorized memory access patterns on NVIDIA GPUs. The actual gain could be smaller if the L2 cache or memory controller already coalesces the scalar loads effectively, or larger if the kernel is hitting a specific bandwidth wall that only vectorized loads can break through.

3. Defragmentation (Tier 0) is low-risk but hard to measure. The assistant notes that "the fresh-pool benchmark won't show the benefit since there's no fragmentation, but for a real multi-tenant server it matters." This is a correct assessment: defragmentation prevents future fragmentation rather than fixing current performance, making it invisible in single-request benchmarks.

4. The marshaling work is essentially done. The assistant states "the tree build is negligible at 1.8ms, so the marshaling work is essentially done." This assumes that tree-build time is the only component of marshaling, and that no other marshaling-related overhead exists. This is a reasonable but potentially incomplete view—there could be memory allocation or data transfer costs not captured by the tree-build profiler.

5. CUDA graph capture compatibility is maintained. The assistant had previously made the kernel capture-safe by consuming SGLang's native static buffers directly, with no host syncs, copies, or cudaMalloc. The NSPLIT increase from 16 to 64 was validated to fit within the capture workspace given the reduced H=8 head count. The vectorized load modification must also preserve this property.

Potential Mistakes and Incorrect Assumptions

While the assistant's reasoning is generally sound, several points warrant scrutiny:

The 1.3× ROI estimate for vectorized loads may be optimistic. The assistant acknowledges that "scalar bf16 loads already get decent hardware coalescing since consecutive threads read consecutive addresses." On modern NVIDIA GPUs, the memory controller can often coalesce scalar loads from consecutive threads into wider transactions, especially when the access pattern is regular (as it is in a flash-attention kernel where each thread reads contiguous elements). If the hardware is already achieving near-vectorized bandwidth with scalar loads, the actual gain could be closer to 1.0–1.1×—hardly worth a kernel rebuild and service restart.

The tradeoff between NSPLIT=64 and NSPLIT=128 was dismissed too quickly. The assistant considered NSPLIT=128 but rejected it because "each test takes 6 minutes to restart." However, NSPLIT=128 would be another environment-variable-only change (no rebuild needed), making it far cheaper to test than the vectorized load modification. If NSPLIT=128 could push throughput from 5.0 to, say, 6.0 tok/s at 65k context, that would be a 20% gain with zero code changes—arguably a better ROI than vectorized loads requiring a kernel rebuild. The assistant's dismissal of this option seems driven by impatience ("Let me be pragmatic") rather than a rigorous cost-benefit analysis.

The assumption that defrag benefits are unmeasurable may be overly conservative. While a fresh-pool benchmark won't show fragmentation, the assistant could simulate fragmentation by running multiple concurrent requests with varying context lengths before measuring decode throughput. This would quantify the defrag benefit and inform the prioritization. The assistant instead defers defrag to "a follow-up" without establishing a baseline.

The "marshaling work is essentially done" conclusion may be premature. The tree-build time of 1.8ms is indeed negligible, but marshaling could include other components: data transfer between draft and verify workers, synchronization costs, or memory allocation for tree structures. The assistant only profiled _build_ddtree_verify_input, not the full marshaling pipeline.

Input Knowledge Required

To fully understand this message, one needs knowledge spanning several domains:

GPU Architecture: Understanding sm_120 (Blackwell consumer) vs sm_90a (Hopper) vs sm_103a (Blackwell DC) is essential. The message assumes familiarity with NVIDIA's compute capability tiers and the fact that sm_120 lacks wgmma, TMA, and tcgen05 instructions. The concept of "occupancy" (how many blocks are active per SM) and its relationship to memory latency hiding is central.

CUDA Programming: The distinction between scalar bf16 loads (half) and vectorized loads (half2 or float4) and their impact on memory bandwidth utilization is a key technical detail. The concept of NSPLIT (splitting the KV sequence into chunks for parallel reduction) and its effect on grid dimensions requires understanding of flash-attention decomposition.

Speculative Decoding Architecture: The message references DFlash (a speculative decoding framework), DDTree (draft tree generation), verify attention (the kernel that scores draft tokens against the target model), and the distinction between draft forward and verify forward passes. The TP8 (tensor parallelism across 8 GPUs) regime and its impact on per-rank head count (H=8 instead of H=64) is crucial.

SGLang Deployment: The service runs under systemd, uses CUDA graphs for decode replay, and has a paged KV cache allocator that can fragment over time. The concept of "Tier 0 defrag" (enabling free-list sorting to keep per-request KV contiguous) is specific to this deployment.

The Prior Investigation: The message builds on weeks of debugging: the 200k context-length deployment, the 0.7 tok/s decode at 46k context, the 2% DRAM utilization during verify, the discovery that _prepare_for_speculative_decoding takes only 7ms while verify takes 185ms, and the occupancy hypothesis that led to the NSPLIT change.

Output Knowledge Created

This message contributes several important pieces of knowledge:

1. Confirmed occupancy hypothesis for TP8 regime. The 2.1–3.3× speedup from NSPLIT 16→64 provides strong evidence that the verify kernel was occupancy-starved in the multi-GPU deployment. This is a general lesson: when scaling a kernel from single-GPU (H=64) to TP8 (H=8), the grid dimensions shrink proportionally, potentially starving occupancy. The fix—increasing NSPLIT to compensate—is simple but not obvious without the grid-dimension analysis.

2. Established the ROI baseline for further optimization. The message explicitly estimates that vectorized loads would yield ~1.3×, while noting that NSPLIT=128 might give additional gains but requires a 6-minute restart to test. This creates a decision framework: any optimization that requires more than a 6-minute test cycle should deliver at least 1.3× to be worth pursuing over the vectorized-load option.

3. Documented the prioritization logic for the user. The assistant's internal debate reveals the tradeoffs: the user asked for defrag and marshaling, but the verify kernel was the real bottleneck. The assistant chose to optimize the verify kernel first (NSPLIT + vectorized loads) before addressing defrag, implicitly arguing that fixing the dominant bottleneck takes priority over the user's explicit request. This is a defensible engineering decision, but it's one the assistant should communicate clearly.

4. Identified the practical constraint of restart time. The 6-minute service restart (including model loading and CUDA graph capture) is a significant operational constraint that shapes the optimization strategy. Environment-variable-only changes (NSPLIT) are cheap; kernel rebuilds are expensive. This operational reality informs the prioritization.

The Thinking Process: A Window Into Engineering Judgment

The most valuable aspect of this message is the raw reasoning process it reveals. The assistant is not simply executing a plan; it is dynamically re-evaluating priorities based on new information. Let me trace the thinking in detail:

Phase 1: Assessment of the NSPLIT victory. The message opens with the assistant comparing the NSPLIT=64 results against the Triton baseline and noting the "solid improvement." There is a clear sense of satisfaction, but also a recognition that more work remains.

Phase 2: The fork in the road. The assistant identifies two paths: vectorized loads (bandwidth efficiency) or higher NSPLIT (occupancy). The initial inclination is toward vectorized loads, but the reasoning immediately reveals doubt: "scalar bf16 loads already get decent hardware coalescing... so the gain might only be 1.2-1.5x."

Phase 3: The user's priorities intrude. The assistant remembers that "the user asked for defrag and marshaling" and attempts to reorient: "lock in NSPLIT=64, add vectorized loads to speed up the verify bottleneck, defrag the free list for Tier 0, then commit." This is a plan, but it's immediately disrupted by a new thought.

Phase 4: The NSPLIT=128 temptation. The assistant briefly considers pushing NSPLIT further, then rejects it due to restart time. This is a pragmatic call, but as noted above, it may be overly hasty—NSPLIT=128 would be a zero-code-change experiment, cheaper than the vectorized load rebuild.

Phase 5: The ROI calculation. The assistant explicitly calculates: "at NSPLIT=64 the verify is around 190ms, and a 1.3x speedup would push throughput from 5 to 6.5 tokens per second." This is the moment of decision: the assistant commits to vectorized loads because the ROI, while modest, is quantifiable and positive.

Phase 6: The defrag deferral. The assistant decides to implement defrag "in a single build cycle" with the vectorized loads, then measures both together. This is a sensible bundling strategy—combine two changes that each require a rebuild into one cycle—but it also means the defrag benefit will be conflated with the vectorized load benefit in the measurement.

What makes this thinking process compelling is its honesty. The assistant does not present a clean, linear decision tree. It cycles, doubts, re-evaluates, and sometimes contradicts itself within a single message. This is not a flaw; it is the signature of genuine reasoning as opposed to post-hoc rationalization.

Conclusion

Message <msg id=12328> is a masterclass in applied engineering judgment under uncertainty. The assistant has just achieved a 2.1–3.3× speedup through a simple environment variable change, and now faces the classic question: take the win and move to the next task, or squeeze more from the same optimization? The reasoning process—cycling through vectorized loads, higher NSPLIT, defrag, and marshaling—reveals a disciplined mind weighing marginal gains against implementation cost, operational overhead, and user priorities.

The message also illuminates a deeper truth about GPU kernel optimization: that the most impactful changes are often not the most technically sophisticated. The NSPLIT fix was a one-line environment variable change, yet it delivered more speedup than weeks of kernel development. The vectorized loads that the assistant ultimately commits to are a modest optimization, estimated at 1.3×, but they represent the kind of incremental improvement that separates a good kernel from a great one.

In the end, this message is about the art of prioritization. The assistant must balance the user's explicit requests (defrag, marshaling) against the data-driven reality (verify kernel is the bottleneck). It must weigh the cost of a 6-minute restart against the potential gain of NSPLIT=128. It must decide whether to bundle changes into a single build cycle or iterate separately. These are not technical decisions in the narrow sense; they are engineering management decisions, made in real-time, under uncertainty, with incomplete information. And that is what makes this message worth studying.