The Occupancy Probe: A Methodical Approach to GPU Kernel Optimization

Introduction

In the high-stakes world of large language model inference, every millisecond counts. When serving the DeepSeek-V4-Flash-NVFP4 model on NVIDIA Blackwell GPUs (RTX PRO 6000), a single assistant message captured a pivotal moment: the transition from deep analysis to concrete action. This message, indexed as <msg id=13551> in the conversation, represents the moment when the assistant switched from theoretical reasoning about GPU kernel occupancy into build mode, executing a carefully planned A/B testing campaign to improve decode throughput by an estimated 15–25%.

The message is remarkable not for its length—it's relatively concise—but for the density of engineering judgment packed into its reasoning trace. It reveals how an experienced systems engineer navigates the treacherous waters between a promising hypothesis and safe execution, balancing ambition with the hard constraints of production GPU programming.

The Scene: What Came Before

To understand this message, we must first understand the problem it addresses. The conversation had been wrestling with a persistent performance bottleneck in the DeepSeek-V4-Flash model's decode attention kernel (_mma_sparse_decode_split_kernel). This kernel is the heart of the autoregressive generation process, responsible for attending to the cached key-value pairs for each new token. On Blackwell GPUs, the kernel was latency-bound—spending most of its time stalled on scattered fp8 KV cache gathers rather than doing useful computation.

The previous message (<msg id=13549>) contained a deep analysis of the kernel's resource usage. The assistant had identified the binding constraint: a per-program f32 accumulator of shape [BLOCK_H=32, 512] consuming approximately 144 registers per thread. This accumulator is intrinsic to the Multi-head Latent Attention (MLA) architecture that DeepSeek-V4 uses, where the latent head dimension is 512. Combined with K/V-tile shared memory staging of about 80 KB, the kernel was limited to exactly 1 CTA (cooperative thread array) per streaming multiprocessor (SM), with only 4 warps per SM—just 8% of the 48 warps available on Blackwell.

The key insight was that the attention kernel was latency-bound, not compute-bound or memory-bandwidth-bound. The SMs were 97% utilized but only 27% of memory controller bandwidth was in use, meaning the GPU was stalled on scattered memory accesses. In such a regime, increasing the number of concurrent warps per SM should help hide memory latency through what's known as Little's law: more in-flight memory requests means the memory subsystem stays busy while individual warps wait.

The assistant had proposed two paths forward:

The Message: Switching to Build Mode

The subject message (<msg id=13551>) is the assistant's response. It opens with a clear declaration of intent:

I'm switching into build mode now. My plan is to first write out a detailed A/B testing document that explains what we're probing (the num_warps and num_stages parameters for the attention kernel), why we're testing these specific values, and how the experiment will work. Once that's committed, I'll move into the research phase to gather any additional context needed for the probe.

This is a classic engineering workflow: document the plan before executing. The assistant recognizes that the A/B test involves modifying Triton kernel autotune configurations, which is a delicate operation that could break CUDA graph capture if done incorrectly. The document serves as both a specification and a safety checklist.

The reasoning trace then dives into the specific challenges:

The key challenge I'm seeing is understanding how the autotune mechanism selects configs at runtime. The autotune key is topk_rounded, which means it picks a single winning config per topk value and caches it. Since the split kernel dominates performance, I need to figure out whether I should remove the existing num_warps=4 configs to force num_warps=8, add an environment variable to control this, or introduce new num_stages=3 configs.

This reveals a critical insight: the Triton autotuner selects configurations based on a key that doesn't include batch size. The winning config is chosen during a warmup benchmark at whatever batch size the system encounters first. If the first warmup happens at a small batch (where 4 warps might be optimal due to better MMA tile subdivision), the autotuner will cache that config and never try 8 warps at high batch sizes where it might be superior.

The assistant considers three implementation approaches:

  1. Remove existing num_warps=4 configs to force num_warps=8 (simplest, but may hurt low-batch performance)
  2. Add an environment variable to control the choice
  3. Introduce new num_stages=3 configs for deeper gather pipelining But there's a constraint: num_stages=3 increases shared memory usage significantly, and with BLOCK_T=32 it might exceed the 100 KB limit and fail to compile. The assistant correctly identifies that a smaller BLOCK_T=16 might be needed.

The Thinking Process: Risk Assessment and Safety

The most valuable part of this message is the reasoning trace, which reveals how the assistant thinks about risk in GPU programming. Three concerns stand out:

1. The autotune-capture interaction. Triton's autotuning performs timed benchmark launches with synchronization. These cannot happen during CUDA graph capture, which requires deterministic, replayable operations. If the autotuner tries to benchmark a new configuration during graph capture, it will crash. The assistant flags this explicitly: "The tricky part is that the autotune benchmarks based on whatever shape and batch it encounters first, so I need to confirm whether the key captures enough information to ensure consistent behavior across different batch sizes."

2. The SMEM budget. Increasing num_stages from 2 to 3 adds another tile of K/V data to the shared memory pipeline. With BLOCK_T=32 and 512 dimensions in bf16, each tile is 32 × 512 × 2 bytes = 32 KB. Three stages would require ~96 KB just for K/V tiles, plus the accumulator and other temporaries. The assistant correctly identifies this might exceed the 100 KB shared memory limit per SM on Blackwell.

3. The batch-aware key problem. Adding batch size to the autotune key would multiply the number of (topk, bs) combinations, each requiring benchmarking during warmup. If warmup doesn't cover all batch buckets, the first runtime call at an uncovered batch size could trigger autotune during capture—breaking the graph. The assistant recognizes this as a future concern, not something to solve in the initial probe.

The Todo List: Structured Execution

The message includes a structured todo list with three items:

  1. Write ./DSV4_ATTN_OCCUPANCY_AB.md (what/why/how, baseline, result table) + commit
  2. Subagent research: Triton autotune selection mechanics, per-config SMEM/register feasibility, how to force/add configs safely
  3. A/B execute: Force num_warps=8 on split kernel, benchmark at C48–96, validate correctness This is a textbook example of staged execution: document first, research the risks, then execute. The assistant is building a safety net before touching any code.

Assumptions and Potential Pitfalls

Several assumptions underpin this message, some of which later prove to be partially incorrect:

Assumption 1: BLOCK_H=32. The assistant assumes 32 heads per GPU, but subsequent research (in the subagent task spawned in <msg id=13554>) reveals that with TP4 (tensor parallelism across 4 GPUs), the 64 attention heads are split into 16 per GPU, making block_h = min(32, 16) = 16. This changes the register pressure calculations significantly—with 16 heads instead of 32, the accumulator is half the size, and the occupancy analysis needs revision.

Assumption 2: The autotune key is topk_rounded. This is correct for the sparse indexer kernel, but the split-kernel attention may use a different key. The subagent research is designed to verify this.

Assumption 3: num_warps=8 will fit. The assistant believes that 8 warps (256 threads) with the 32×512 accumulator gives 64 f32 values per thread instead of 128, halving register pressure. This is correct in theory, but the actual fit depends on the compiler's register allocation, which can be unpredictable.

Assumption 4: The cheap probe has the same ceiling as a full rewrite. The assistant argues that forcing 8 warps in a single CTA gives the same latency-hiding benefit as 2 CTAs per SM. This is approximately true—both give 8 warps per SM—but the single-CTA approach may have different MMA efficiency characteristics.

Input Knowledge Required

To fully understand this message, one needs:

Output Knowledge Created

This message produces several concrete artifacts:

  1. DSV4_ATTN_OCCUPANCY_AB.md: A detailed A/B testing plan document covering what is being tested (num_warps=8, num_stages=3), why (occupancy-driven latency hiding), and how (benchmarking methodology, correctness gates, result tables)
  2. A git commit: docs(dsv4): #3a attention occupancy A/B plan — num_warps=8/num_stages probe (what/why/how, baseline, result table); raise warps/SM for gather-latency hiding, projected +15-25% high-C
  3. A subagent task: A research task investigating Triton autotune mechanics, per-config SMEM/register footprints, and the safest approach to force or add configurations
  4. A structured todo list: Three prioritized items that guide the subsequent work
  5. An engineering workflow: The pattern of "document → research → execute → validate" that serves as a template for future optimization efforts

The Broader Significance

This message exemplifies a mature approach to GPU kernel optimization that goes beyond mere trial and error. The assistant:

  1. Builds a quantitative model before touching any code (the throughput model with fixed overhead + marginal cost per request)
  2. Identifies the binding constraint (register pressure from the f32 accumulator, not shared memory or compute)
  3. Proposes a minimal intervention (change a config parameter, not rewrite the kernel)
  4. Documents the plan before executing (the A/B doc)
  5. Researches the risks before making changes (the subagent task)
  6. Stages the work to de-risk: probe first, rewrite only if needed This is the opposite of the "hail mary" approach to GPU optimization—randomly tweaking parameters and hoping for the best. It's systematic, evidence-driven, and safety-conscious. The message also reveals the unique challenges of optimizing for Blackwell GPUs (compute capability sm_120), which have different resource limits and scheduling characteristics than previous generations. The assistant is effectively doing frontier engineering—optimizing a model that was likely designed for Hopper (H100) GPUs on a newer architecture with different trade-offs.

Conclusion

Message <msg id=13551> is a masterclass in structured engineering execution. It captures the precise moment when analysis transitions to action, when the assistant commits to a plan and begins building. The reasoning trace shows a mind working through trade-offs: safety vs. performance, simplicity vs. flexibility, immediate gains vs. long-term correctness.

The document that results from this message—DSV4_ATTN_OCCUPANCY_AB.md—will guide the next phase of optimization. The subagent research will validate (or invalidate) the assumptions about autotune mechanics and SMEM budgets. And the actual A/B tests will determine whether the 15–25% throughput improvement is real or illusory.

But regardless of the outcome, the message itself stands as a testament to the value of methodical engineering. In a field where the temptation is always to "just try it and see," the assistant instead builds a framework for understanding, a plan for safe experimentation, and a commitment to documentation. That is the mark of an engineer who has learned—perhaps the hard way—that the fastest path to a correct solution is rarely the shortest one.