The Custom sm_120 Verify Kernel: From 0.7 tok/s to 3–6× Decode Speedup on Blackwell Consumer GPUs

Introduction

In the high-stakes world of large language model inference, few metrics matter more than decode throughput. When a system is deployed with a 200,000-token context length, users expect usable generation speeds — not a crawl of minutes per response. This article chronicles a pivotal arc within a sprawling opencode coding session: the journey from a catastrophic 0.7 tokens-per-second decode collapse at 185k context to a 3–6× speedup achieved through custom CUDA kernel engineering on NVIDIA RTX PRO 6000 Blackwell GPUs (sm_120 architecture).

The story spans multiple phases: a hard-won deployment of 200k context-length on the CT200 server, a systematic benchmark that revealed a devastating performance cliff, a user's skeptical challenge that forced a deeper investigation, the identification of the Triton MLA verify attention kernel as the root cause, the discovery that no optimized MLA kernels support the sm_120 architecture, and ultimately the design, optimization, and deployment of a custom sm_120 verify attention kernel with CUDA graph support and KV defragmentation. This is a case study in disciplined performance debugging, GPU architecture understanding, and the art of knowing when off-the-shelf solutions must be replaced with custom engineering.

Part I: The 200k Context Deployment

The journey began with a seemingly straightforward goal: deploy the Kimi K2.6 model with DFlash DDTree speculative decoding on an 8-GPU RTX PRO 6000 Blackwell server (CT200) with 200,000 tokens of context length. The assistant quickly discovered that setting --context-length 200000 was not enough — SGLang sizes its KV cache pool based on memory budget (--mem-fraction-static), not the context-length parameter. The KV pool remained stubbornly pinned at 101,134 tokens.

The assistant methodically diagnosed the issue, calculating the per-token memory cost of MLA latent KV (~68 KB/token/GPU across 61 layers), raising mem-fraction-static from 0.85 to 0.92, then to 0.94, dropping max-running-requests from 64 to 8 to free the req_to_token pool, and restarting the service multiple times (each restart taking ~6 minutes). After an OOM crash at 32k context with the 0.94 setting, the assistant dialed back to 0.93, enabled chunked-prefill-size 4096 and expandable_segments:True, achieving a stable service with a KV pool of 195,011 tokens — 97.5% of the 200k target.

The deployment was technically complete. The server ran, handled requests, and produced correct tokens. But the question remained: how fast was it at long context?

Part II: The Benchmark That Exposed a Crisis

The user's directive was three words: "Benchmark long context perf." This simple request launched a comprehensive performance characterization that would reshape the entire project.

The assistant designed a streaming benchmark methodology to avoid prefix cache interference, running requests at context lengths from 1.4k to 185k tokens. The results were devastating:

| Actual Context | Prefill (tok/s) | Decode (tok/s) | |---|---|---| | 1,427 | 3,104 | 118.6 | | 11,402 | 2,659 | 19.9 | | 31,952 | 1,820 | 6.3 | | 63,903 | 1,255 | 2.8 | | 127,802 | 696 | 1.4 | | 184,728 | 479 | 0.7 |

Decode throughput collapsed from 118.6 tok/s at 1.4k context to just 0.7 tok/s at 185k — a 170× slowdown. Step time ballooned from 8ms to 1.43 seconds. The DDTree speculative drafter reported avg_commit_len=1.00, meaning it was accepting only a single token per step — zero speedup from speculation.

The assistant initially attributed this to the drafter being "undertrained for long contexts" and the synthetic benchmark prompt being "novel prose" the drafter couldn't predict. But the user immediately challenged this interpretation.

Part III: The User's Challenge That Changed Everything

The user raised four sharp concerns that fundamentally challenged the assistant's narrative:

  1. "decode step per token increase seems extremely high, no?" — The latency scaling was far too steep to be explained by physics alone.
  2. "Do we have sliding window attention working in the drafter?" — If the drafter was attending the full context instead of a bounded window, its cost would scale linearly with context.
  3. "commit_len=1 is probably a bug" — A perfectly flat 1.0 across all samples, not a statistical distribution, looked like a systematic verification failure.
  4. "GPU use appeared to be quite low pointing to compute being far from saturated" — If the GPU was nearly idle during decode, the bottleneck wasn't compute but something else. These questions forced the assistant to pivot from "reporting mode" to "debugging mode." The assistant launched a systematic investigation spanning code inspection, controlled experiments, and GPU profiling.

Part IV: The Investigation — Ruling Out Hypotheses

The assistant's investigation was methodical and hypothesis-driven. First, it inspected the drafter's configuration, confirming it was a 6-layer Qwen3 model with sliding window attention (5 of 6 layers, window 2048) and YaRN position scaling. But configuration alone didn't prove runtime behavior.

The decisive experiment was a controlled commit-length probe using trivially predictable text ("The cat sat on the mat" repeated) at multiple context lengths. The results were unambiguous:

| Context | Commit Len | Depth | Drafts | |---|---|---|---| | 1k | 7.88 | 7.00 | 7.00 | | 16k | 7.00 | 6.78 | 6.11 | | 64k | 7.00 | 6.11 | 6.11 | | 128k | 7.88 | 6.88 | 6.88 |

The drafter maintained commit lengths of 7–8 tokens at all context lengths. The sliding window attention, YaRN scaling, and KV caching were all working correctly. The commit_len=1 in the original benchmark was purely a text-difficulty artifact — the synthetic "continue this discussion" prompt produced genuinely unpredictable novel prose that the small drafter couldn't anticipate. Not a bug.

But this raised a deeper question: if the drafter was fine, why was decode so slow?

Part V: The GPU Utilization Revelation

The assistant deployed a GPU utilization probe during a 64k-context decode run. The results revealed a stunning paradox:

PREFILL: SM%= 99.4  memctrl%=  5.0  power_W= 301.7
DECODE : SM%= 99.8  memctrl%=  1.2  power_W= 133.6

During decode, nvidia-smi reported 99.8% SM utilization — suggesting the GPU was fully saturated. Yet the memory controller was barely active at 1.2%, and power draw had collapsed to 133.6W — roughly 22% of the GPU's thermal design power. The SMs were "occupied" but doing almost no useful work.

The user independently confirmed this with cufall screenshots showing tensor core utilization at ~3% during decode, versus 50–90% during prefill. This was the smoking gun: the decode kernel was latency-bound, not compute-bound or bandwidth-bound. The GPU was spending 97% of its cycles stalled, waiting on something.

The assistant calculated the effective bandwidth for KV cache reads: at 64k context, each decode step reads ~4.5 GB of compressed latents. With a step time of ~327ms, the effective bandwidth was ~14 GB/s — 130× below the GPU's 1.8 TB/s peak.

Part VI: The Root Cause — page_size=1 and the Triton Lock

Tracing the structural cause led to a hard constraint in SGLang's DDTree implementation. Reading the source code of dflash_worker.py revealed:

if self.is_ddtree and self.page_size != 1:
    raise RuntimeError(
        "DDTREE currently requires --page-size 1 because accepted tree nodes "
        "are non-contiguous in the verify block."
    )

DDTree requires page_size=1. The tree-structured speculative decoding means accepted tokens are scattered non-contiguously across the KV cache, forcing per-token page granularity. And the verify attention — the step where the target model evaluates the draft tree — is locked to the Triton MLA backend because no other backend (FlashInfer, FlashMLA, cutlass-MLA) supports custom tree attention masks for MLA on sm_120.

The combination was catastrophic: page_size=1 forces scattered, non-coalesced memory access; Triton MLA on sm_120 lacks the advanced tensor core instructions (wgmma, TMA, tcgen05) that make attention efficient on Hopper and datacenter Blackwell; and batch size 1 with 9 queries provides minimal parallelism. The result was 3% tensor core utilization and 0.7 tok/s.

The assistant also discovered that all optimized MLA kernels — FlashMLA, cutlass-MLA, flashinfer-MLA — are compiled only for sm_90a/sm_100a/sm_103a. None support sm_120 (the RTX PRO 6000 Blackwell consumer architecture). There was no off-the-shelf fix.

Part VII: Building the Custom sm_120 Verify Kernel

With the root cause confirmed, the assistant pivoted to building a custom solution. The strategy was outlined in plans/0002-sm120-verify-kernel-defrag.md and executed in phases.

The core insight was to separate the verify attention into two parts: the contiguous prefix (which doesn't need masking and can use coalesced reads at full bandwidth) and the small tree block (typically 9 nodes, which needs the custom mask and scattered access). By handling the prefix efficiently and applying the tree mask only to the small block, the native kernel could eliminate the 130× penalty.

The assistant wrote verify_attn_flash.cu, a KV-split flash-decode MLA verify kernel with a partial+reduce design to improve occupancy. The kernel was iteratively optimized through several key modifications.

Making the Kernel CUDA Graph Capture-Safe

One of the most challenging requirements was CUDA graph capture safety. CUDA graphs allow a sequence of GPU operations to be captured and replayed with minimal CPU overhead, dramatically reducing launch latency for short decode steps. However, graph capture imposes strict constraints: all operations within the captured region must use static buffers, no host-device synchronization is permitted, and no dynamic memory allocation (cudaMalloc) can occur during capture.

The assistant's initial kernel implementation used host-side synchronization and dynamic allocations that violated these constraints. The fix required a fundamental rewrite: instead of synchronizing between kernel launches, the kernel consumed SGLang's native static buffers directly. The NSPLIT parameter (which controls how many KV tiles are processed per thread block) was fixed at compile time rather than being dynamically determined. A torch-allocated workspace provided scratch space without cudaMalloc calls.

The assistant systematically identified every capture-breaking pattern:

The Optimization That Unlocked 3–6× Speedup

With the kernel working correctly, the assistant turned to performance tuning. The profiler had revealed that the verify attention was occupancy-starved in the TP8 regime — with only 8 heads per rank (64 heads ÷ 8 GPUs), there was insufficient parallelism to hide memory latency.

The insight was simple but profound. 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.

Two targeted fixes addressed this:

1. Increasing NSPLIT from 16 to 64. NSPLIT controls how many KV tiles each thread block processes before writing its partial result. A higher NSPLIT means each thread block does more work, improving occupancy by keeping more warps active. This multiplied the grid by 4× (to 4608 blocks), giving roughly 24 waves — enough to keep the SMs fed with work. The result, achieved with just an environment variable change and no kernel rebuild, was immediate: 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.

2. Adding 128-bit vectorized bf16 KV loads. The naive kernel loaded KV data using 32-bit scalar loads, achieving only a fraction of the memory bandwidth. By switching to 128-bit vectorized loads (using uint4 or float4 intrinsics), the kernel could issue four times fewer memory requests for the same data, reducing instruction overhead and improving bandwidth utilization. This was particularly impactful for the MLA latent, which has a contiguous layout in memory that vectorized loads can exploit.

The combined effect was dramatic. The cumulative speedup over the Triton+graph baseline was:

| Context | Triton+Graph | Custom Kernel | Speedup | |---------|-------------|---------------|---------| | 4k | 39 tok/s | 116.5 tok/s | 3.0× | | 16k | 6.8 tok/s | 29.1 tok/s | 4.3× | | 65k | 1.5 tok/s | 9.2 tok/s | 6.1× |

The 3–6× end-to-end decode speedup across all context lengths was a transformative result. The effective bandwidth rose from ~14 GB/s to over 100 GB/s — still below the 1.8 TB/s peak, but a massive improvement that brought the system into a usable regime.

Part VIII: Tier 0 KV Defragmentation

The third major workstream was KV cache defragmentation. The assistant had identified that DDTree's page_size=1 requirement scattered each token's KV cache across physical memory, and that server churn (requests being allocated and freed) exacerbated this fragmentation over time.

Tier 0 defragmentation was implemented by monkeypatching SGLang's TokenToKVPoolAllocator to force need_sort=True. This caused the allocator's merge_and_sort_free() function to return ascending/contiguous slot indices, improving allocation contiguity without any data movement. The change was confirmed active on all 8 TP workers via a runtime probe.

The assistant correctly assessed that Tier 0 was the low-risk, high-impact first step. Tier 1 (live relocation of fragmented KV slots under a copy budget) was deferred with clear reasoning: single-request KV was already contiguous after the Tier 0 fix, and the bottleneck had shifted to MoE expert imbalance — a structural ceiling beyond the verify kernel's scope.

Part IX: 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 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 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 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.

Part X: The Bottleneck Shift — MoE Expert Imbalance

With the verify attention bottleneck addressed, the assistant turned to understanding the remaining performance ceiling. The user had shared a GPU utilization heatmap showing decode-phase utilization that was "patchy/imbalanced" — some GPUs saturated on compute and PCIe bandwidth while others sat idle. The assistant correctly diagnosed this as MoE expert imbalance at batch size 1.

With tensor parallelism across 8 GPUs, each rank holds 48 of the 384 experts. When only ~9 tokens route to ~8 experts each, the activated experts are scattered across GPUs, and some GPUs may have zero activated experts for a given layer. Those idle GPUs must still participate in the all-reduce synchronization, creating a synchronization tax without contributing useful computation.

This is not a bug — it is a structural property of TP+MoE at low batch sizes. The assistant's analysis was precise: "the fundamental low-batch MoE behavior" cannot be fixed by kernel optimization. The real levers are batching (more concurrent requests → more tokens → more experts activated → better balance) or expert parallelism with load balancing.

The assistant documented this as the remaining bottleneck, noting that the MoE inefficiency is a separate concern from the attention bottleneck. The attention fix unlocked 3–6× speedup; the MoE fix would require a different engineering effort.

Part XI: Engineering Principles Emergent from the Journey

Several recurring principles emerge from this engineering journey that are worth articulating explicitly:

Constraint-driven design. The capture-safety requirement was not an optional optimization — it was a hard constraint that shaped every aspect of the kernel interface. The fixed NSPLIT, the in-kernel mask indexing, the pre-allocated workspace — none of these would have been necessary without the capture constraint. But the constraint forced a cleaner, more GPU-centric design that ultimately benefited performance even beyond graph replay.

Instrumentation before optimization. The assistant repeatedly used profiler data to guide decisions. When the profiler showed CPU orchestration (tree-build at 1.8ms) was negligible, the assistant abandoned the marshaling-optimization path and focused on kernel occupancy instead. When the MoE imbalance emerged as the new ceiling, the assistant chose the minimal defrag intervention (Tier 0) rather than building a complex Tier 1 system that would not improve current performance.

The art of deferral. One of the hardest skills in engineering is knowing what not to build. The assistant explicitly deferred Tier 1 defrag with clear reasoning, documented it as a future capability, and focused on delivering the minimal useful intervention. This is not procrastination — it is disciplined prioritization based on bottleneck analysis.

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.

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.

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: A Case Study in Performance Engineering

This segment of the opencode session is a masterclass in systematic performance engineering. It demonstrates several crucial lessons:

Measurement before optimization. The user's three-word directive — "Benchmark long context perf" — forced the assistant to characterize the system before attempting to optimize it. This revealed a 170× decode slowdown that no amount of configuration tuning could fix.

Trust but verify metrics. The nvidia-smi SM% metric reported 99.8% utilization, suggesting a fully saturated GPU. Cross-referencing with power draw, memory controller utilization, and tensor-core profiling revealed the truth: the GPU was 97% stalled.

Know when to build custom. The assistant explored every off-the-shelf option — FlashMLA, cutlass-MLA, flashinfer-MLA — and found that none supported sm_120. Building a custom kernel was not a luxury; it was the only viable path.

Systematic hypothesis elimination. The investigation ruled out drafter bugs, sliding window failures, KV caching issues, and compute-bound bottlenecks through controlled experiments before converging on the true root cause.

The result — a 3–6× decode speedup through custom kernel engineering — transformed a system that was technically functional but practically unusable into a viable long-context inference platform. The live service now runs 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 to MoE expert imbalance at batch size 1 — a structural ceiling beyond the verify kernel's scope. And that is precisely the mark of a successful optimization campaign: not the elimination of all bottlenecks, but the honest identification of the next one.

References

[1] "From 0.7 tok/s to 3-6× Speedup: The Custom Kernel Journey That Rescued Long-Context Decode on Blackwell GPUs" — Chunk 0 article covering the full arc from deployment to kernel diagnosis.

[2] "The Execution Phase: From Plan to 3–6× Decode Speedup on Blackwell Consumer GPUs" — Chunk 1 article covering the systematic execution of the optimization plan.

[3] "The Custom sm_120 Verify Kernel: A Case Study in GPU Kernel Engineering from Bottleneck Diagnosis to Production Deployment" — Chunk 2 article covering the kernel lifecycle.

[4] "From Capture Safety to MoE Ceiling: The Complete Arc of a Custom CUDA Verify Kernel Optimization" — Chunk 3 article covering the capture-safety and bottleneck-shift analysis.