The SM120 Decode Kernel Autotune Investigation: When Shared Memory Constraints Overturn Assumptions

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When serving models like DeepSeek-V4 across multiple GPUs, the attention kernel—the computational heart of the transformer—becomes the primary bottleneck. Optimizing it requires navigating a labyrinth of hardware constraints, compiler behavior, and autotuning logic. This article examines a single message from an AI-assisted coding session that delivered a comprehensive research report on Triton autotuning for the SM120 (Blackwell GPU) sparse decode split-kernel. The message, written by an AI agent acting as a research engineer, is remarkable not only for its technical depth but for the way it systematically overturns the core premise of the optimization plan it was asked to investigate.

The subject message (global index 36) is the culmination of a multi-round investigation into why a particular kernel configuration was being selected by Triton's autotuner, whether the autotuning process was safe under CUDA graph capture, and how to improve occupancy for better latency hiding. What the agent discovered—by reading actual compiled artifacts rather than relying on estimates—fundamentally changed the understanding of the problem space. The candidate configuration everyone had assumed would work (BLOCK_T=32 with 8 warps) turned out to be unlaunchable due to shared memory constraints, and the real optimization opportunity lived at a different tile size entirely.

This article dissects that message in detail, exploring the reasoning process, the assumptions that were corrected, the knowledge required to understand the analysis, and the concrete deliverables produced. It serves as a case study in how rigorous empirical investigation can overturn plausible but incorrect mental models in GPU kernel optimization.

Context: The Investigation Leading Up to the Message

To understand the subject message, one must first understand what preceded it. The broader coding session (a subagent invocation within a larger root session) was tasked with researching Triton autotuning behavior, CUDA graph capture interactions, and per-config shared memory (SMEM) usage for the DeepSeek-V4 MLA (Multi-head Latent Attention) decode kernel on NVIDIA Blackwell (SM120) GPUs.

The root session had been working on deploying the GLM-5-NVFP4 model using SGLang, a high-performance inference engine. Along the way, the team had encountered and resolved numerous issues: bf16 corruption under CUDA graph capture (fixed by disabling multi-stream-overlap), decode throughput scaling investigations, production incidents from degraded PD (prefill-decode) bootstrap, and a client-side HTTP deadlock. The subagent was spawned specifically to investigate the Triton autotuner's behavior with CUDA graph capture and per-config shared memory usage—a narrow but critical piece of the larger puzzle.

In the messages leading up to the subject message (messages 28 through 35), the agent had been systematically gathering evidence. It had:

  1. Read the Triton autotuner source code to understand how OutOfResources errors are handled (they result in inf timing values, effectively pruning those configurations).
  2. Checked for persisted autotune cache files (*.autotune.json) and found none—confirming that autotuning results are computed fresh in-memory per process startup, not persisted to disk.
  3. Examined the backend call site in deepseek_v4_backend.py to understand how the kernel is invoked, how batch size flows through, and what the grid dimensions look like.
  4. Read the actual compiled kernel artifacts from /root/.triton/cache, using JSON metadata for shared memory sizes and cuobjdump -res-usage for register counts.
  5. Verified the autotuner's short-circuit behavior—when only a single config is provided, the benchmarking block is skipped entirely, making single-config autotune decorators unconditionally safe under CUDA graph capture. By message 36, the agent had assembled enough evidence to produce a definitive report. But the report's conclusions were not what the original plan had anticipated.

The Core Finding: BLOCK_T=32 Does Not Fit

The most striking revelation in the subject message is that the candidate configuration everyone had been working toward—BLOCK_T=32 with 8 warps and 2 stages—cannot be launched on SM120. The shared memory requirement of 104,704 bytes (102.25 KB) exceeds the 99 KB (101,376 bytes) opt-in cap for Blackwell GPUs. Triton compiles the kernel (a cubin exists), but when the runtime attempts to set cudaFuncSetAttribute(maxDynamicSharedMemorySize=104704), it fails with OutOfResources, and the autotuner assigns an infinite timing value, effectively discarding the configuration.

This is a classic case of a plausible-sounding optimization being ruled out by a hard hardware constraint. The reasoning that led to considering BLOCK_T=32 was straightforward: larger tiles mean more work per thread block, which can improve arithmetic intensity and reduce the number of thread blocks needed. But on SM120, the shared memory budget is tight—only 101,376 bytes per thread block when using the opt-in maximum. The agent's derived SMEM formula revealed the exact breakdown:

SMEM = 1152 · BLOCK_H + 1348 · num_stages · BLOCK_T   bytes

For BLOCK_H=16, BLOCK_T=32, and num_stages=2: 1152·16 + 1348·2·32 = 18,432 + 86,272 = 104,704 bytes—exceeding the cap by 3,328 bytes. The formula was validated against all four measured data points, giving high confidence in its correctness.

The practical implication is stark: any optimization plan that assumes BLOCK_T=32 is viable on SM120 for this kernel is fundamentally wrong. The agent's measured data showed that both BLOCK_T=32 configurations (4 warps and 8 warps) produce the same 104,704-byte SMEM footprint, confirming that shared memory usage is independent of warp count for this kernel. The occupancy win that was expected from 8 warps at BLOCK_T=32 is moot—the kernel cannot run at all.

Where the Optimization Actually Lives: BLOCK_T=16 with 8 Warps

With BLOCK_T=32 ruled out, the agent turned to BLOCK_T=16, which fits comfortably at 61,568 bytes (60.1 KB). The current production configuration is BLOCK_T=16 with 4 warps (w4), yielding 1 CTA per SM and approximately 8.3% occupancy (4 warps out of 48 available per SM). The agent's key insight was that switching to 8 warps (w8) at BLOCK_T=16 keeps the same shared memory footprint (61,568 bytes—unchanged because SMEM is independent of warp count) but doubles occupancy to approximately 16.7% (8 warps per SM).

This is the real optimization opportunity. The register pressure analysis confirmed that w8 is feasible: with 256 threads instead of 128, the per-thread accumulator register footprint drops from roughly 72 registers per thread to 36, freeing substantial register budget. The measured BT32/w8 configuration showed 212 registers per thread (compared to 253 for BT32/w4), confirming the trend. At BT16/w8, the agent estimated approximately 200 registers per thread, well under the 255-register-per-thread limit, with no spilling expected.

The occupancy improvement from 8.3% to 16.7% is significant for latency hiding. During decode, the kernel is primarily memory-bound (gathering sparse KV cache entries from HBM), and having more warps available to overlap memory access with computation can substantially improve throughput. The agent's recommendation (V1) was to force BLOCK_T=16 with 8 warps as a single configuration, which has the additional advantage of being unconditionally capture-safe (the autotuner skips benchmarking entirely when only one config is provided).

Capture Safety: A Detailed Analysis

A major concern motivating the investigation was whether Triton's autotuning process could interfere with CUDA graph capture. CUDA graphs capture a sequence of GPU operations for replay, but they cannot capture operations that synchronize with the host (like cudaStreamSynchronize). Triton's autotuner uses do_bench, which performs multiple kernel launches with synchronization—exactly the kind of operation that would break graph capture if it occurred inside the capture block.

The agent's analysis of the SGLang codebase revealed a robust safety mechanism. The FullCudaGraphBackend.capture_one method runs the forward function twice in eager mode (with synchronize() and barrier()) before opening the capture block:

for _ in range(2): synchronize(); barrier(); forward_fn(); post_warmup_hook()
with graph_ctx(...): out = forward_fn()   # capture

Triton's autotuner benchmarks only on the first call for an unseen key. Those benchmark launches land in the eager warmup phase, never inside the capture block. By the time forward_fn() runs under graph_ctx, the configuration for that key is already cached, so the in-capture launch is a pure replay of the cached choice. No benchmarking ever occurs inside capture.

This finding is important for two reasons. First, it confirms that the existing system is safe—no changes are needed to prevent autotuner-capture conflicts. Second, it opens the door for the agent's V3 proposal (batch-aware keying), which adds the batch size as an autotune key parameter. Since each CUDA graph capture bucket is eagerly warmed before its own capture, adding a new key dimension does not introduce any capture-unsafe behavior. The only cost is additional warmup time, as the autotuner must benchmark each unique (topk_rounded, batch_bucket) pair once.

The agent also verified that there is no environment variable or flag that disables autotuning under capture (only TRITON_PTXAS_BLACKWELL_PATH is set; maxnreg is null). Safety comes purely from the warmup-before-capture design, not from disabling the tuner. This is a robust architectural property that holds regardless of how the autotune decorator is configured.

The SMEM Model: Derivation and Validation

One of the most impressive aspects of the subject message is the agent's derivation of the exact shared memory formula for the split kernel. Rather than relying on estimates or rules of thumb, the agent read the actual JSON metadata from the Triton cache for four compiled configurations and reverse-engineered the formula:

SMEM = 1152 · BLOCK_H + 1348 · num_stages · BLOCK_T   bytes

The breakdown is:

Why w4 Is Currently Selected: Not Performance, But Pruning

A key conceptual contribution of the subject message is correcting the mental model of why the 4-warp configuration is currently chosen. The natural assumption is that the autotuner benchmarks all configurations and selects the fastest one. But the reality is more interesting: the autotune key is topk_rounded only, which is constant at 512 in production. This means there is exactly one autotune key, benchmarked once. Of the three configured options:

  1. BT16/w4 (61,568 B) — fits
  2. BT32/w4 (104,704 B) — exceeds 99 KB → OutOfResources → inf
  3. BT32/w8 (104,704 B) — exceeds 99 KB → OutOfResources → inf Only BT16/w4 is launchable, so it wins by default. The choice is forced by the SM120 99 KB shared-memory cap, not by batch-shape tuning or performance comparison. The capture order (largest batch first) is irrelevant to the outcome because the BLOCK_T=32 configurations are pruned at every batch size. This is a subtle but important distinction. The agent explicitly addresses the user's likely mental model: "Your mental model 'first/high-batch shape determines the cached choice' isn't what's driving it here—SMEM pruning is." Understanding this distinction is crucial because it changes the optimization strategy. If the selection were driven by batch-shape tuning, the solution might involve better benchmarking or more representative warmup data. But since it's driven by a hard hardware constraint, the only solution is to either (a) work within the constraint (use BLOCK_T=16) or (b) change the kernel design to reduce SMEM usage (e.g., drop to num_stages=1, which would fit BLOCK_T=32 at 61,568 bytes but would kill load-compute pipelining).

The Three Proposed Variants

The subject message concludes with three concrete, line-number-precise diff proposals for modifying the kernel's autotune decorator. Each represents a different trade-off between simplicity, flexibility, and performance.

V1: Force More Warps (Recommended, Safest)

This variant replaces the three-config decorator with a single configuration:

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_T": 16}, num_warps=8, num_stages=2),
    ],
    key=["topk_rounded"],
)

The advantages are compelling. First, with only one config, the autotuner's run method short-circuits at if len(self.configs) > 1 (line 214 of autotuner.py), skipping benchmarking entirely. This makes V1 unconditionally capture-safe with zero autotune overhead. Second, it achieves the occupancy goal: 61,568 bytes SMEM, 1 CTA per SM, 8 warps per SM (≈17% occupancy). Third, it is the simplest change—one line modified in the decorator, no new kernel parameters, no caller changes.

The downside is losing the w4 fallback. But as the agent demonstrates, the BLOCK_T=32 entries were dead anyway (unlaunchable), so the only thing lost is the freedom to choose between w4 and w8 at BLOCK_T=16. Given that w8 likely wins at all batch sizes (more warps for latency hiding, no SMEM penalty), this is a minimal risk.

V2: Add num_stages=3

This variant adds a third pipeline stage while keeping BLOCK_T=16 pinned:

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_T": 16}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK_T": 16}, num_warps=8, num_stages=2),
        triton.Config({"BLOCK_T": 16}, num_warps=8, num_stages=3),
    ],
    key=["topk_rounded"],
)

The three-stage configuration uses 83,136 bytes (81.2 KB)—fits, still 1 CTA per SM. The extra pipeline stage can improve memory latency hiding by allowing more prefetch buffers, but it comes with a register pressure risk. The agent notes that the extra prefetch buffer can push register usage up and potentially cause spilling, which would need to be verified after warmup by reading the compiled cubin's register usage.

V2 is multi-config, so the autotuner will benchmark all three options. This is still capture-safe (warmup-before-capture handles it), but it adds warmup time. The agent recommends verifying register usage after deployment.

V3: Batch-Aware Key

The most sophisticated variant adds the batch size as an autotune key, allowing the tuner to select different configurations for different batch buckets:

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_T": 16}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK_T": 16}, num_warps=8, num_stages=2),
    ],
    key=["topk_rounded", "bs_bucket"],
)

This requires three changes: adding a bs_bucket: tl.int32 parameter to the kernel signature, adding "bs_bucket" to the autotune key list, and computing a coarse bucket value in the caller. The agent proposes a simple bucketing scheme: bs_bucket = 1 if B <= 8 else (32 if B <= 48 else 96), which keeps the number of distinct keys (hence warmup benchmarks) small.

V3 is the most flexible option, potentially allowing w4 for small batches (where fewer warps might suffice) and w8 for large batches (where more warps help hide latency). However, the agent notes that with only BLOCK_T=16 tiles fitting, V3's only freedom is w4↔w8. Given that w8 likely helps both the throughput-bound high-bs case (B=96 → 576 CTAs) and the latency-bound low-bs case (B=1 → ~16 CTAs on 188 SMs), V1 (force w8) is likely as good and far simpler. The agent recommends reaching for V3 only if A/B testing shows w8 regressing at some bucket.

Assumptions Corrected

The subject message corrects several assumptions that were likely held by the human user and possibly by the agent itself earlier in the investigation:

  1. "BLOCK_T=32 with 8 warps will improve occupancy." This was the core premise of the optimization plan. The agent's measurement showed it's unlaunchable—SMEM exceeds the cap by 3,328 bytes. The occupancy win exists but at BLOCK_T=16, not 32.
  2. "The autotuner selects w4 because it's faster for high batch sizes." The agent showed that w4 is selected not because it's faster, but because it's the only launchable option. The BLOCK_T=32 configs are pruned by SMEM constraints, not outperformed in benchmarking.
  3. "D_TOT is approximately 576." The agent corrected this: 576 is the byte stride per token (_TOKEN_DATA_STRIDE), not the logical dimension. The actual D_TOT is 512.
  4. "The current production config uses 80 KB SMEM with 255 registers." The agent showed this is a stale build from June 17 with BLOCK_H=32. The fresh production config uses 61,568 bytes (60.1 KB) with 252 registers and no spilling.
  5. "Autotuning might fire inside CUDA graph capture." The agent confirmed this cannot happen due to the warmup-before-capture design. Every bucket is eagerly warmed twice before its capture block opens.
  6. "Adding batch to the autotune key might be unsafe." The agent showed it is safe, as long as the batch values seen during replay were all seen during warmup—which the warmup-before-capture design guarantees.

Input Knowledge Required

To fully understand the subject message, a reader needs knowledge spanning several domains:

GPU Architecture (Blackwell SM120): Understanding of shared memory hierarchies, the 99 KB opt-in cap, warp scheduling (48 warps per SM), register file size (65,536 registers per SM), and the distinction between per-thread-block and per-SM resource limits.

Triton Compiler and Runtime: Familiarity with the @triton.autotune decorator, Config objects, key parameters, the OutOfResources exception handling in autotuner.py, the do_bench benchmarking mechanism, and the single-config short-circuit behavior.

CUDA Graph Capture: Understanding of how CUDA graphs work (capture and replay), why host-side synchronization is illegal inside capture, and how the warmup-before-capture pattern prevents autotuner conflicts.

SGLang Architecture: Knowledge of the FullCudaGraphBackend, the capture_one method, the decode CUDA graph runner, bucket-based capture (largest-first iteration), and the replay padding mechanism.

DeepSeek-V4 MLA Kernel: Understanding of the split-kernel design (nsplit calculation), the sparse decode path, the role of topk_rounded, the BLOCK_H and BLOCK_T tile dimensions, and the data layout (q_nope, q_rope, K/V tiles).

Quantitative Reasoning: Ability to interpret shared memory formulas, register counts, occupancy calculations, and trade-offs between pipeline stages, warp counts, and tile sizes.

Output Knowledge Created

The subject message produces several valuable knowledge artifacts:

  1. A validated SMEM formula for the _mma_sparse_decode_split_kernel that predicts shared memory usage from BLOCK_H, BLOCK_T, and num_stages. This formula can be used to evaluate any proposed configuration without compiling and running it.
  2. A definitive resource table for all compiled configurations, showing SMEM, register count, spill status, and launchability. This replaces estimates with measured ground truth.
  3. A capture-safety proof that autotuning cannot fire inside CUDA graph capture, backed by code analysis of both the Triton autotuner and the SGLang backend.
  4. Three concrete, deployable diffs with exact line numbers, ready for implementation. Each diff is accompanied by a risk assessment and a recommendation.
  5. A corrected mental model of why w4 is currently selected (SMEM pruning, not batch-shape tuning), which changes how future optimization efforts should be directed.
  6. A register pressure estimate for the proposed BT16/w8 configuration (~200 registers per thread, no spill expected), with a method for verification (read the compiled cubin after warmup).

The Thinking Process: A Window into Rigorous Investigation

The subject message is notable for the quality of its reasoning. Several patterns stand out:

Empirical grounding: The agent repeatedly prioritizes measured data over estimates. When the user's plan assumed BLOCK_T=32 was viable, the agent went to the actual compiled artifacts and read the exact SMEM values. This empirical approach is what overturned the core premise of the plan.

Systematic verification: Every claim is cross-checked against multiple sources. The SMEM formula is validated against four data points. The capture-safety analysis checks both the autotuner source and the backend capture code. The register count trends are confirmed across both BT32 and BT16 configurations.

Explicit mental-model correction: The agent doesn't just present findings; it explicitly addresses likely misconceptions. The "Your mental model..." construction in section 3 is a deliberate pedagogical move, helping the reader update their understanding.

Conservative recommendations: Despite discovering that BLOCK_T=32 is unlaunchable, the agent doesn't over-recommend. V1 (force w8) is presented as the safest and most likely best option, but V2 and V3 are offered with clear trade-off analyses. The agent even advises against V3 unless A/B testing shows a regression, showing restraint in the face of complexity.

Precision in language: The agent carefully distinguishes between "does not fit" (exceeds SMEM cap and cannot launch) and "not selected" (loses in benchmarking). It corrects the D_TOT vs. stride confusion. It notes that the stale H=32 builds are from June 17, not fresh. This precision prevents downstream confusion.

Awareness of limitations: The agent acknowledges what it cannot know without testing. The register estimate for BT16/w8 is labeled as such, with a verification procedure provided. The V3 caveat about potential w8 regression at some bucket is explicitly stated. This intellectual honesty strengthens the credibility of the recommendations.

Broader Implications

The subject message has implications beyond the immediate DeepSeek-V4 optimization task. It demonstrates a methodology for GPU kernel optimization that is broadly applicable:

  1. Always measure, never assume. The BLOCK_T=32 case is a perfect example of a plausible configuration being ruled out by a hard constraint that was not obvious from first principles. The SMEM formula could have been derived theoretically, but the agent's willingness to read the actual compiled artifacts provided immediate, incontrovertible evidence.
  2. Understand the autotuner's behavior. The agent's deep dive into the Triton autotuner source code revealed the single-config short-circuit, the OutOfResources → inf mapping, and the key-based caching mechanism. This understanding enabled recommendations that work with the autotuner rather than against it.
  3. Capture safety is a system property, not a kernel property. The agent showed that autotuner-capture safety depends on the warmup-before-capture pattern in the SGLang backend, not on any property of the kernel or the autotune decorator. This means the same kernel could be capture-safe or capture-unsafe depending on how it's invoked—a crucial insight for anyone working with CUDA graphs.
  4. The best optimization may not be the one you were looking for. The investigation started with a focus on BLOCK_T=32 and ended with a recommendation for BLOCK_T=16/w8. The real win (doubling occupancy from 8.3% to 16.7%) was hiding in a configuration that was already in the autotune list but as the "conservative" option. The agent's willingness to follow the evidence where it led, rather than forcing the data to fit the original plan, is a model of scientific integrity.

Conclusion

The subject message (global index 36) is a masterclass in GPU kernel optimization research. It combines deep knowledge of GPU architecture, Triton compiler internals, CUDA graph capture semantics, and SGLang runtime behavior to produce a definitive analysis of the DeepSeek-V4 SM120 decode split-kernel autotuning landscape. The agent's empirical approach—reading actual compiled artifacts rather than relying on estimates—overturned the core premise of the optimization plan and revealed a simpler, safer, and more effective path forward.

The message's three proposed variants (V1, V2, V3) provide a menu of options with clear trade-offs, and the agent's recommendation (V1: force BLOCK_T=16 with 8 warps) is well-supported by the evidence. The capture-safety analysis is thorough and conclusive, the SMEM formula is validated against multiple data points, and the corrected mental model of why w4 is currently selected prevents future misdirection.

For anyone working on GPU kernel optimization, Triton autotuning, or LLM inference serving, this message offers both practical deliverables (the diffs, the SMEM formula, the verification procedures) and methodological lessons (measure over assume, understand the full system, follow the evidence). It is a testament to the power of rigorous, systematic investigation in a domain where the constraints are complex and the stakes are high.