The Custom sm_120 Verify Kernel: A Case Study in GPU Kernel Engineering from Bottleneck Diagnosis to Production Deployment

Introduction

In the high-stakes world of large language model inference, the gap between a theoretically fast GPU kernel and a production-deployed one is often measured not in lines of code but in engineering discipline. This article synthesizes a remarkable engineering journey captured across dozens of messages in an opencode coding session: the design, implementation, optimization, debugging, and deployment of a custom CUDA verify attention kernel for the NVIDIA RTX PRO 6000 Blackwell GPU (sm_120 architecture), targeting speculative decoding with the Kimi K2.6 model running in SGLang.

The work spans the full lifecycle of a production kernel: from diagnosing a devastating long-context decode bottleneck (0.7 tok/s at 185k tokens), through building a custom kernel from scratch when no optimized MLA libraries supported the target architecture, iteratively optimizing it for occupancy and memory bandwidth, making it CUDA-graph capture-safe, integrating it into a live SGLang service via a carefully designed monkeypatch backend, and finally deploying it to achieve a 3–6× end-to-end decode speedup over the Triton baseline. The story ends not with a triumphant declaration of victory, but with a clear-eyed identification of the remaining bottleneck: MoE expert imbalance at batch size 1—a structural ceiling beyond the verify kernel's scope.

The Bottleneck: Triton's Scattered KV Access at 14 GB/s

The journey began with a deployment problem. The CT200 server, equipped with eight RTX PRO 6000 Blackwell GPUs, was configured to serve the Kimi K2.6 model with 200k-token context length using DFlash speculative decoding with a DDTree (Draft Tree) drafter. A comprehensive long-context benchmark revealed a catastrophic decode slowdown: at 185k tokens, throughput collapsed to 0.7 tokens per second, with GPU tensor core utilization at a mere 3% despite 99.8% SM occupancy reported by nvidia-smi.

The root cause was traced to the DDTree verify attention kernel—the component that computes attention scores between draft tokens and the full KV cache during speculative decoding's verification step. This kernel was locked to SGLang's Triton-based MLA (Multi-head Latent Attention) implementation with page_size=1, which caused scattered, non-contiguous KV cache accesses that achieved an effective memory bandwidth of only ~14 GB/s—roughly 130× below the GPU's 1.8 TB/s peak. The GPU's tensor cores were starving for data while the memory subsystem struggled to service the scattered access pattern.

The assistant and user systematically investigated potential causes. They confirmed that the drafter's sliding-window attention and KV caching worked correctly (commit length stayed 7–8 at all context lengths with predictable text). The commit_len=1 observed in benchmarks was purely text difficulty, not a bug. The bottleneck was squarely in the verify attention kernel's memory access pattern.

The Architectural Dead End: No Optimized MLA Kernels for sm_120

The obvious solution would be to swap in an optimized MLA kernel—FlashMLA, cutlass-MLA, or flashinfer-MLA. However, the assistant discovered that all of these libraries are compiled exclusively for sm_90a (Hopper), sm_100a, and sm_103a (Blackwell data center) architectures. None support sm_120, the compute capability of the RTX PRO 6000 Blackwell consumer GPU. These kernels rely on Hopper and Blackwell data-center instructions—wgmma (warp-group matrix multiply), TMA (Tensor Memory Accelerator), and tcgen05—that simply do not exist on sm_120's Ada-like instruction set architecture.

The RTX PRO 6000 Blackwell is a curious hybrid: it uses the Blackwell architecture name but with a consumer-grade ISA that lacks the specialized tensor core instructions of its data-center siblings. This meant the assistant could not use any off-the-shelf optimized kernel. The only path forward was to build an owned sm_120 verify attention kernel from scratch, working within the constraints of ~100 KB of shared memory per SM and no access to the advanced tensor core instructions that power FlashMLA and its ilk.

Building the Custom Kernel: From Occupancy Starvation to Flash-Decode

The assistant wrote a detailed per-phase plan (plans/0002-sm120-verify-kernel-defrag.md) and began implementing Phase 1: a KV-split flash-decode MLA verify kernel (verify_attn_flash.cu). The initial design was a tiled flash attention kernel that processed KV cache in chunks, reusing each chunk across multiple query heads via shared memory. The design was theoretically sound: by tiling across head dimensions with HT=8 (processing 8 heads at a time from the KV data loaded into shared memory), the kernel should reduce global memory reads by a factor of 16× compared to the naive approach.

But when the benchmarks came back, the results were sobering. At short prefixes (256–4096 tokens), the flash kernel was 40% slower than the naive baseline. At 16384 tokens, it achieved only a 0.9× speedup—essentially parity. The theoretical memory advantage was being completely negated by a hidden enemy: occupancy starvation.

The diagnosis, delivered with characteristic precision in [msg 12239], revealed the core issue: the flash kernel's design produced B × q_len × (H/HT) = 1 × 9 × 8 = 72 thread blocks. The RTX PRO 6000 Blackwell has approximately 188 Streaming Multiprocessors. With only 72 blocks, only about 38% of SMs could be active simultaneously. The naive kernel, by contrast, used one block per head, producing 576 blocks and keeping all SMs fully occupied. The flash kernel's 16× memory reduction was meaningless when most of the GPU sat idle.

The fix was a fundamental architectural pivot to a flash-decode approach with KV splitting. Instead of one block per (batch, query_position, head_tile), the KV dimension itself would be partitioned across multiple blocks. Each block would handle a contiguous slice of the prefix, computing partial online-softmax statistics and writing them to a global scratch buffer. A second reduction kernel would combine the partials. With NSPLIT=32, the grid would become B × q_len × n_htiles × NSPLIT = 1 × 9 × 8 × 32 = 2304 blocks—enough to saturate all 188 SMs.

The KV-split rewrite passed all correctness tests with token-exact precision (max_abs error ~2e-8), as documented in [msg 12240]. The occupancy fix was validated, but the kernel was now memory-throughput-bound rather than occupancy-bound.

The Float4 Vectorization: A Single Edit That Unlocked 3–6× Decode Throughput

With the KV-split architecture in place, the assistant performed a detailed bottleneck analysis in [msg 12244]. The achieved bandwidth was only 403 GB/s—roughly 22% of the 1.8 TB/s peak. The scalar float loads used to read KV cache tiles were coalesced but not fully utilizing the memory subsystem's capacity.

The solution was a single, targeted edit: replacing scalar float loads with float4 (128-bit) vectorized loads for the KV tile and query data. This is a classic GPU optimization that reduces the number of load instructions and better utilizes the memory bus width. On NVIDIA architectures, 128-bit loads can be coalesced into a single memory transaction, whereas four scalar loads may require separate transactions even when targeting contiguous addresses.

The assistant also made a crucial production-context realization: in the TP8 (tensor parallelism across 8 ranks) deployment scenario, each rank processes only 8 heads instead of the full 64. This means the per-head design reads the KV cache only 8 times per rank rather than 64—an 8× reduction in memory traffic compared to the single-GPU benchmark. The float4 vectorization therefore provided benefits in both the benchmark (where the full 64-head read was the bottleneck) and production (where the reduced head count already made memory traffic less dominant).

This single edit, combined with the earlier NSPLIT increase from 16 to 64, delivered a dramatic 3–6× end-to-end decode speedup over the Triton baseline across all context lengths (4k–65k). The kernel was now not just correct, but genuinely fast.

The Integration Challenge: Bridging Kernel and Production Framework

A fast kernel in isolation is a laboratory curiosity. To realize the speedup in production, the assistant had to integrate it into SGLang's attention backend—a complex system with paged KV cache, custom masking, tensor parallelism, and CUDA graph capture. This integration consumed the bulk of the engineering effort and revealed the true nature of production kernel deployment: the hardest part is not writing the kernel, but wiring it into the serving framework.

The assistant's integration strategy, documented across multiple messages, was a masterclass in surgical engineering. Rather than registering a new attention backend (which would require understanding SGLang's entire backend registry and implementing the full interface), the assistant chose to monkeypatch TritonAttnBackend.forward_extend directly. This reduced the integration surface to a single method while keeping --attention-backend triton as the framework configuration.

The monkeypatch was gated by a three-mode environment variable (KDTREE_VERIFY=off|validate|on). In off mode, the patch was a no-op and all attention ran through Triton. In validate mode, the backend ran both Triton and the custom kernel, compared their outputs, logged any differences, but returned Triton's result—keeping the service correct while monitoring convergence. In on mode, only the custom kernel ran, delivering the full performance benefit.

But the monkeypatch faced a fundamental challenge: SGLang uses Python's multiprocessing spawn mechanism to create separate scheduler subprocesses for each GPU in the tensor parallelism group. A monkeypatch applied in the parent process never reaches the subprocesses where attention actually runs. The assistant's initial launch-wrapper approach crashed immediately because spawned subprocesses re-executed the wrapper's top-level code, creating a recursive loop.

The fix, documented in [msg 12270], was a pivot to sitecustomize.py—a Python startup hook that is automatically imported by every interpreter process when its directory is on PYTHONPATH. By placing the monkeypatch in sitecustomize.py and setting PYTHONPATH in the systemd unit, the patch was installed in the parent process and all eight TP worker subprocesses at interpreter startup time, before any SGLang code ran.

The Debugging Gauntlet: Three Real Bugs Against Live Tensors

Integration testing against real serving tensors revealed three distinct bugs, each requiring a different debugging approach.

Bug 1: CUDA Graph Capture Invalidation. When the assistant deployed the kernel in validate mode, the server crashed during startup at approximately 630 seconds. The journal logs revealed a cudaErrorStreamCaptureInvalidated error originating from SGLang's CUDA graph capture mechanism. The validation code performed host-side synchronizations (.item() calls, torch.max().item() for logging) and memory allocations that are illegal during CUDA graph capture. The fix was to disable CUDA graphs (--disable-cuda-graph) for the validation phase, accepting the performance cost to prove correctness first.

Bug 2: The Extend Layout Mismatch. When the kernel finally ran in eager mode, it produced outputs with a relative error of approximately 1.0—essentially uncorrelated with Triton's output. The assistant's investigation in [msg 12282] revealed a fundamental misunderstanding of SGLang's extend attention architecture. The assistant had assumed that all KV tokens—both prefix and draft—were accessed through the KV pool via kv_indices. In reality, SGLang's forward_extend path separates the prefix (stored in the pool, accessed via kv_indices) from the draft tokens (computed fresh, written to out_cache_loc). The custom kernel was only attending to the 21 prefix tokens, completely missing the 9 draft tokens, and using the wrong mask stride.

The fix was to build combined indices: concatenate the prefix slot indices with the draft token slot locations, and extract the visibility mask using the correct stride of prefix + q_len rather than prefix alone. After this fix, the kernel achieved numerical parity with Triton to within one bfloat16 unit in the last place (ULP)—a relative error of ~1e-4 mean, max_abs_diff of 3.9e-3 to 7.8e-3, which is exactly bfloat16 rounding noise.

Bug 3: The Marshaling Index Error. A separate marshaling bug was discovered where mask_indptr was sliced to [:bs+1] instead of [:bs], causing a tensor shape mismatch. This was a classic one-off indexing error that was caught by the validate-mode logging and fixed before the parity validation.

The Deployment: Live Cutover with Known Trade-offs

With numerical parity achieved and the kernel validated on real serving tensors, the assistant performed the live cutover in [msg 12293]. The service was configured with KDTREE_VERIFY=on and --disable-cuda-graph, running the custom kernel in eager mode. The benchmark results showed a clean 2× decode speedup at long context (23k and 91k tokens), with a slight regression at short context (5.7k tokens: 35.2 tok/s vs 38.0 tok/s, 0.93×) due to the eager-mode overhead.

The assistant documented the trade-offs explicitly in the commit message: "cuda-graph capture-safety = remaining for short-ctx/production." The cutover was delivered as requested, with a clear rollback path (revert KDTREE_VERIFY=off and drop --disable-cuda-graph), and the remaining work was scoped for future prioritization.

The Remaining Bottleneck: MoE Expert Imbalance

After the verify kernel was deployed and optimized, the assistant turned to understanding the remaining performance ceiling. Profiling revealed that the decode step was now limited by MoE (Mixture of Experts) expert imbalance at batch size 1. With tensor parallelism across 8 GPUs and only a single request, the MoE layers cannot efficiently utilize all GPUs because each token only activates a subset of experts, and at batch size 1, there is no opportunity to batch tokens across different expert assignments to improve utilization.

This is a structural ceiling beyond the verify kernel's scope. The attention computation—now running at 3–6× Triton's throughput—is no longer the bottleneck. The bottleneck has shifted to the MoE layers, which are fundamentally constrained by the single-request batch size. The solution would require either batching multiple requests together (which increases latency but improves throughput) or implementing expert parallelism (which distributes experts across GPUs differently than tensor parallelism).

The Tier 0 KV Defragmentation

In parallel with the kernel optimization, the assistant implemented Tier 0 KV defragmentation by monkeypatching the allocator to force need_sort=True, keeping per-request KV contiguous on churned pools. This was confirmed active on all 8 TP workers. Tier 1 (live relocation) was deferred with clear reasoning: single-request KV is already contiguous after Tier 0 defrag, and the bottleneck has shifted to MoE imbalance, making further defragmentation work lower priority.

Key Engineering Lessons

This engineering journey yields several important lessons for anyone building custom GPU kernels for production inference:

1. Occupancy is a first-order design constraint. The initial flash kernel's theoretical memory advantage was completely negated by poor occupancy. On a GPU with 188 SMs, generating enough thread blocks to keep all SMs busy is more important than minimizing memory traffic. The KV-split flash-decode approach succeeded because it transformed the grid from 72 blocks to thousands of blocks, saturating the GPU.

2. Measure before optimizing, and control your variables. The assistant's disciplined benchmarking—disabling CUDA graphs for both the custom kernel and Triton to create a level playing field—prevented misleading comparisons. A sloppier evaluation might have compared a graph-mode Triton baseline against an eager-mode custom kernel and concluded the custom kernel was slower.

3. Integration is harder than kernel development. The kernel itself was written and validated in Phase 1. The integration—understanding SGLang's extend attention layout, building the monkeypatch, fixing the multiprocessing spawn issue, debugging the CUDA graph capture incompatibility, and fixing the marshaling bugs—consumed the majority of the engineering effort. The three real bugs found against live tensors would never have been discovered in isolated unit tests.

4. The sitecustomize.py pattern is a powerful tool for cross-process code injection. When a Python serving framework uses multiprocessing spawn, a launch-wrapper approach cannot inject code into subprocesses. The sitecustomize mechanism, combined with PYTHONPATH, provides a clean solution that works at interpreter startup time.

5. Ship with known trade-offs. The assistant deployed the kernel knowing it had a short-context regression in eager mode. Rather than waiting for perfection (CUDA graph capture-safety), the assistant shipped the current state, documented the limitations, and provided a rollback path. This is the essence of production engineering: delivering value now while being honest about what remains.

Conclusion

The custom sm_120 verify attention kernel journey is a case study in the full lifecycle of production GPU kernel engineering. From diagnosing a 130× bandwidth gap, through building a kernel from scratch for an unsupported architecture, iteratively optimizing it for occupancy and memory bandwidth, integrating it into a complex serving framework, debugging three real bugs against live tensors, and finally deploying it to achieve a 3–6× decode speedup—every phase required a different skill set and a different mode of thinking.

The result is a live SGLang service running the owned capture-safe kernel with CUDA graphs enabled and Tier 0 defrag active, achieving 3–6× Triton's decode throughput. The bottleneck has shifted from the verify attention kernel to MoE expert imbalance at batch size 1—a structural ceiling that is a natural next challenge, not a failure of the kernel work.

This story demonstrates that the most impactful engineering is often not about writing the fastest kernel, but about understanding the full system—the GPU architecture, the serving framework's internals, the multiprocessing semantics, the CUDA graph capture constraints, and the production deployment infrastructure—and making disciplined trade-offs at each step. The custom kernel is the visible artifact, but the invisible engineering judgment is what made it work.