From Capture Safety to MoE Ceiling: The Complete Arc of a Custom CUDA Verify Kernel Optimization

Introduction

In the high-stakes world of large language model inference, the difference between a research prototype and a production system often comes down to a series of disciplined engineering decisions: what to optimize, when to stop optimizing, and how to know the difference. This article synthesizes the complete arc of a single chunk (Segment 66, Chunk 3) from an opencode coding session — a sustained engineering campaign to deploy the Kimi K2.6 model with speculative decoding on RTX PRO 6000 Blackwell GPUs (sm_120 architecture). Over the course of dozens of messages, the assistant built a custom CUDA verify attention kernel, made it CUDA-graph-capture-safe, tuned it for the tensor-parallelism-8 (TP8) regime to achieve a 3–6× end-to-end decode speedup over the Triton baseline, implemented Tier 0 KV defragmentation, and ultimately identified MoE expert imbalance as the remaining structural bottleneck.

This article traces the full narrative arc, examining the key decisions, the reasoning processes that drove them, and the engineering principles that emerged.

The Starting Point: A Broken Bottleneck

The chunk opens with the assistant in a familiar position for any performance engineer: a working optimization that cannot be deployed. The custom sm_120 verify attention kernel (verify_attn_flash_paged.cu) had been built from scratch because no existing optimized MLA kernels — FlashMLA, cutlass-MLA, flashinfer-MLA — supported the sm_120 architecture. These kernels targeted sm_90a (Hopper), sm_100a, and sm_103a (Blackwell DC), relying on instructions like wgmma, TMA, and tcgen05 that simply do not exist on sm_120's Ada-like ISA. The assistant's custom kernel had been validated for token-exact correctness against a naive oracle and achieved a promising 2× decode speedup over Triton in eager mode at long context lengths.

But there was a catch: the integration forced the entire SGLang service to run with --disable-cuda-graph, meaning CUDA graphs — which record a sequence of GPU operations into a reusable replay object — were globally disabled. Short-context throughput actually regressed to 0.93× of the original baseline because the Python orchestration overhead could no longer be amortized across graph replays. The user's directive was unambiguous: "Fix cudagraphs, then do defrag and optimise marshaling" ([msg 12298]).

Phase 2d: The Capture-Safe Kernel

The first major sub-chapter of this chunk is the transformation of the verify kernel from a graph-incompatible eager-mode prototype into a capture-safe kernel that could participate in CUDA graph replay. This work is documented across multiple messages ([msg 12299] through [msg 12308]), each revealing a different facet of the capture-safety problem.

The Architecture of Capture Safety

The assistant's reasoning in [msg 12299] is a masterclass in constraint-driven design. CUDA graph capture records GPU operations with fixed buffer addresses and grid dimensions. Any operation that violates these constraints — host-device synchronization (.item() calls), dynamic memory allocation (cudaMalloc), or tensor copies (.to(), .contiguous(), .cat()) — breaks capture. The assistant systematically identified every capture-breaking pattern in the existing kernel integration and designed solutions for each:

Fix NSPLIT: The kernel used a split-K flash-decode design where the attention computation was divided across thread blocks (splits). If NSPLIT was dynamic, the grid dimensions would change between capture and replay. The assistant fixed NSPLIT to a conservative maximum value (eventually settling on 16 after iterating through 64 and 32), with the kernel computing chunk boundaries dynamically from the actual KV length. Empty splits (where the chunk range exceeds the KV length) early-out cheaply.

Eliminate host-side index rebuilding: The Python code concatenated prefix KV indices with draft token cache locations into a combined buffer — an allocation that broke capture. The assistant redesigned the kernel to read from two separate arrays (kv_indices for the prefix, out_cache_loc for draft tokens) and select between them based on position relative to the prefix length.

Inline mask indexing: The visibility mask was pre-gathered into a dense buffer. Instead, the kernel now receives the raw custom_mask array and mask_indptr offsets, computing the mask index inline during execution.

Pre-allocate workspace: The ensure_ws function used raw cudaMalloc, which bypasses PyTorch's caching allocator and breaks graph capture. The assistant replaced this with a torch.empty tensor — PyTorch's allocator is capture-aware, so allocations through it are safe during graph capture.

The Leap of Faith

Message [msg 12305] captures the tense moment when the assistant must actually enable CUDA graphs and test the capture-safe kernel. The assistant's reasoning reveals a nuanced understanding of the testing dilemma: validate mode (which compares kernel outputs against a reference using .item() calls) is structurally incompatible with CUDA graph capture, because the very mechanism that makes validate mode useful — host-side comparison of GPU outputs — is a capture-breaking synchronization. The assistant must go straight to ON mode with graphs enabled, relying on earlier eager-mode validation that already confirmed numerical correctness.

The fallback plan is explicit: "if the server comes up successfully I'll test generation — otherwise I'll debug the capture error." This is not blind faith but calculated confidence, built on careful design work. The memory fraction is lowered from 0.94 to 0.88 to accommodate the workspace and graph overhead. The service restarts. And CUDA graph capture succeeds in approximately 1.5 seconds, with decode replay producing generations matching the Triton baseline.

Phase 2e: The Occupancy Epiphany and the 3–6× Speedup

With CUDA graphs working, the assistant turned to performance tuning. The initial benchmark with graphs enabled showed promising short-context gains (135 tok/s at 1k context vs 35 in eager mode) but diminishing returns at longer contexts. The assistant ran a clean A/B comparison ([msg 12309]) between the custom kernel+graphs and Triton+graphs under identical conditions — same memory fraction, same server, same benchmark script. The results were nuanced: 1.35× speedup at 1k context, parity at 4k–16k, and 2.2× at 65k.

But the real breakthrough came from an unexpected direction. The assistant had been investigating "idle gaps" in the GPU timeline, initially suspecting CPU orchestration overhead (tree-building at 1.8ms, mask-building at 0.18ms). The profiler revealed these were negligible — the real cost was the verify attention kernel itself, which was occupancy-starved in the TP8 regime.

The Occupancy Calculation

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.

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. 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 — a 2.1–3.3× speedup over the Triton+graph baseline.

Vectorized Loads: The Final Push

The assistant then faced a classic engineering dilemma ([msg 12328]): should it push further with vectorized bf16 loads (estimated 1.3× gain), try higher NSPLIT values (128?), or pivot to the user's requested defragmentation work? The reasoning process cycles through multiple prioritizations, each incorporating new information: the magnitude of the NSPLIT gain, the user's stated priorities, the cost of kernel rebuilds and service restarts, and the estimated ROI of vectorized loads.

The assistant ultimately implemented the vectorized loads ([msg 12333]), and the results validated the decision. The cumulative speedup over the Triton+graph baseline was now:

| 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× speedup across all context lengths was a transformative result. But with attention no longer the bottleneck, a new ceiling emerged.

The Bottleneck Shift: MoE Expert Imbalance

The user had shared a GPU utilization heatmap ([msg 12330]) 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 ([msg 12331]). 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.

Tier 0 Defrag: The Art of Knowing When to Stop

With the attention bottleneck resolved and MoE imbalance identified as the next frontier, the assistant turned to the user's remaining request: KV defragmentation. The investigation ([msg 12337]–[msg 12338]) traced the need_sort parameter in SGLang's TokenToKVPoolAllocator to a single line of code:

need_sort = self.server_args.disaggregation_mode in ("decode", "prefill")

This meant that for standard (non-disaggregated) services — which is what the Kimi K2.6 deployment used — the free list remained unsorted, allowing fragmentation to accumulate over time as KV cache pages were allocated and freed across multiple concurrent requests.

The assistant implemented Tier 0 defrag ([msg 12339]–[msg 12341]) by monkeypatching the allocator's __init__ to force need_sort=True, gated behind an environment variable (KDTREE_DEFRAG_SORT=1). This is a "defrag-lite" approach: it doesn't move existing KV blocks, but it prevents fragmentation from worsening by keeping the free list sorted.

The critical decision was to defer Tier 1 (live KV relocation). The assistant's reasoning was explicit: Tier 1 would require a complex live-memory-mover with scattered relocation and reindexing — high complexity, high risk, and low ROI given that the bottleneck had shifted to MoE imbalance. Tier 0, by contrast, was cheap, low-risk, and provided production robustness for multi-tenant scenarios.

The Closing Signal

The chunk culminates in a todowrite message ([msg 12345]) that marks four high-priority items as completed:

  1. Phase 2d: CUDA graph capture-safety — The kernel is capture-safe, capture succeeds in ~1.5 seconds, and generations match the Triton baseline.
  2. Phase 2e: Tune verify kernel for TP8 (NSPLIT=64 + vectorized loads) → 3–6× — The cumulative speedup is validated across all context lengths.
  3. Diagnose idle gaps: attention occupancy (fixed); bottleneck now MoE imbalance bs=1 — The root cause was found, fixed, and the bottleneck has shifted.
  4. Phase 4 Tier 0 defrag (free-list sort) — enabled, active — The defrag patch is deployed and verified on all 8 TP workers. This todowrite is not merely a status update — it is a carefully constructed signal that closes a major engineering chapter. It encodes decisions about what to optimize and what to defer, assumptions about shared context, and knowledge about the system's performance characteristics and remaining bottlenecks.

Engineering Principles Emergent from the Chunk

Several recurring principles emerge from this chunk 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.

Honest communication: The assistant repeatedly told the user when assumptions no longer held. "The fragmentation premise doesn't apply to this benchmark." "MoE imbalance at batch size 1 is a structural ceiling." This intellectual honesty — knowing what problems have been solved and what problems remain — is a hallmark of effective engineering.

Conclusion

This chunk of the opencode session represents a complete optimization cycle: from identifying a bottleneck (scattered KV access in Triton MLA), through building a custom solution (the sm_120 verify kernel), making it production-ready (capture-safe for CUDA graphs), tuning it for the deployment architecture (NSPLIT=64 + vectorized loads in TP8), implementing supporting infrastructure (Tier 0 defrag), and ultimately identifying the next bottleneck (MoE expert imbalance).

The 3–6× end-to-end decode speedup is a significant achievement, but the more valuable output is the understanding of where the bottleneck now lies. By shifting the bottleneck from attention to MoE, the assistant has both delivered a concrete speedup and identified the next frontier for optimization. The attention problem has been solved with a custom sm_120 kernel, NSPLIT tuning, and vectorized loads. The MoE problem — inherent to TP+low-batch inference — will require a different set of tools: batching, expert parallelism, or architectural changes to the model itself.

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.