The Moment of Truth: Validating CUDA Graph Capture for a Custom sm_120 Verify Kernel

Introduction

In the high-stakes world of production LLM inference, few optimizations are as consequential—or as brittle—as CUDA graph capture. A CUDA graph records a sequence of GPU operations (kernel launches, memory copies, synchronization primitives) into a reusable executable that can be replayed with minimal CPU overhead. For inference engines like SGLang, graphs are the difference between a system that spends half its time in kernel launch overhead and one that saturates the GPU with continuous computation. But graph capture imposes a draconian constraint: every operation recorded must be capture-safe, meaning it cannot depend on host-side synchronization, dynamic memory allocation, or any operation that would change the CUDA execution state between captures and replays.

This article examines a single message ([msg 12308]) from an opencode coding session in which an AI assistant, having just completed the arduous task of making a custom sm_120 verify attention kernel capture-safe, runs the first performance benchmark with CUDA graphs enabled. This message is the payoff moment—the culmination of dozens of prior messages involving kernel rewriting, dtype analysis, workspace allocation strategy, and NSPLIT tuning. It is a message of verification, measurement, and transition from "does it work?" to "how fast is it?"

The Message in Full

The assistant begins with a reasoning block that reveals its mental model of the system state:

The CUDA graph capture completed successfully without errors, and the kdtree installation timing confirms the graphs were captured properly. The OpenAI warning is unrelated and benign. Now I need to verify that the decode phase is actually using CUDA graphs—the "cuda graph: False" lines are from prefill which doesn't use graphs by design, so I should check whether decode is properly leveraging them.

>

Let me pull the dtype logs to validate my assumptions, then run the performance benchmark in ON+graph mode and compare it against the original triton+graph baseline from the context scaling results and my earlier eager A/B test. I'll run the context decode benchmark at a few different context lengths to get a proper comparison.

The assistant then executes two commands: one to check the dtype logs and confirm the kernel installation, and another to run the bench_context_decode.py benchmark across four context lengths (1024, 4096, 16384, 65536 tokens). The output shows:

May 31 15:34:37 dflash-train python[31848]: [kdtree] installed verify override (mode=on, NSPLIT=16)
May 31 15:34:43 dflash-train python[31864]: [kdtree] installed verify override (mode=on, NSPLIT=16)
May 31 15:34:43 dflash-train python[31866]: [kdtree] installed verify override (mode=on, NSPLIT=16)
=== ON + cuda-graph perf ===
 ctx_req prompt_tok   TTFT_s prefill_tok/s decode_tok/s gen_tok decode_wall
    1024       1428     0.47          3070        135.1      16        0.11
    4096       5699...

The output is truncated, but the first row tells a compelling story: at 1024 tokens context, the custom kernel with CUDA graphs achieves 135.1 decode tokens per second—a dramatic improvement over the ~35 tok/s observed in eager mode (from the previous message's reasoning).

Why This Message Matters

This message sits at a critical juncture in a long engineering arc. To understand its significance, we must trace the path that led here.

The Problem: Triton MLA on sm_120

The session's broader context (segment 66) involves deploying the Kimi K2.6 model with DFlash speculative decoding on RTX PRO 6000 Blackwell GPUs (sm_120 architecture). The verify attention kernel—the component that checks speculative tokens against the full model's predictions—was initially locked to Triton's MLA (Multi-head Latent Attention) implementation with page_size=1. This caused scattered KV cache accesses that achieved only ~14 GB/s effective bandwidth, roughly 130× below the GPU's 1.8 TB/s peak. The result was abysmal decode performance: 0.7 tok/s at 185k context length, with GPU tensor core utilization at just 3%.

The Solution: A Custom sm_120 Kernel

The assistant discovered that all optimized MLA kernels (FlashMLA, cutlass-MLA, flashinfer-MLA) were compiled only for sm_90a (Hopper), sm_100a, and sm_103a (Blackwell datacenter)—none supported sm_120 (Blackwell consumer). These kernels rely on instructions like wgmma, TMA, and tcgen05 that simply don't exist on sm_120's Ada-like ISA. The user approved building a custom kernel, and the assistant designed a KV-split flash-decode MLA verify kernel with a partial+reduce design.

The Capture-Safety Challenge

The custom kernel worked in eager mode but was initially incompatible with CUDA graphs. The assistant had to rewrite it to consume SGLang's native static buffers directly—eliminating host synchronizations (.item() calls), tensor copies (.to(), .contiguous(), .cat()), and dynamic allocations (cudaMalloc). The NSPLIT parameter (which controls how many parallel chunks the KV sequence is divided into) had to be fixed at compile time rather than computed dynamically per request. The workspace had to be pre-allocated via torch.empty (which uses PyTorch's graph-private memory pools during capture) rather than allocated fresh on each call.

After this rewrite, CUDA graph capture succeeded in ~1.5 seconds, and the server generated correct output ("Paris. This is a well-known fact") in ON mode with graphs enabled ([msg 12307]).

The Reasoning Process: What the Assistant Is Thinking

The assistant's reasoning in [msg 12308] reveals several layers of understanding:

1. Confirming Capture Integrity

The assistant first checks that capture completed "without errors" and that there were no "invalidated" errors—a critical concern because CUDA graphs can be silently invalidated if a captured operation changes its memory addresses between capture and replay. The 1.53-second capture time and the absence of invalidation messages confirm the graph is stable.

2. Understanding SGLang's Graph Architecture

The assistant demonstrates deep knowledge of SGLang's internals by noting that "cuda graph: False" lines in the logs come from prefill, which "doesn't use graphs by design." This is an important distinction: in SGLang's architecture, prefill (processing the prompt) and decode (generating tokens one by one) use different execution paths. Prefill is typically compute-bound and benefits less from graph capture, while decode is launch-latency-bound and benefits enormously. The assistant correctly interprets log noise as expected behavior rather than a problem.

3. Planning the Validation Strategy

The assistant plans a three-step validation: (a) check dtype logs to confirm the kernel's data type assumptions were correct, (b) run the performance benchmark with graphs enabled, and (c) compare against the Triton+graph baseline. This reveals an assumption worth examining: the assistant assumes that the dtype assumptions (e.g., kv_indptr as int32, qo_indptr as int64, custom_mask as uint8) are correct and merely needs log confirmation. In fact, these dtype choices were the subject of extensive analysis in prior messages ([msg 12300]–[msg 12303]), where the assistant traced through SGLang's buffer construction to determine the exact native types.

4. The NSPLIT=16 Tradeoff

The kernel was configured with NSPLIT=16, meaning the KV sequence is divided into 16 parallel chunks. This choice reflects a careful tradeoff: higher NSPLIT values increase parallelism (more blocks launched) but require more workspace memory. The assistant had considered NSPLIT=64 but settled on 16 to keep workspace memory manageable (~605 MB for batch size 8) within the ~4 GB of free GPU memory after the model and KV cache. The assistant's reasoning in [msg 12303] shows this deliberation: "NSPLIT=16 for the kernel grid—it's modest enough to avoid excessive overhead while still handling the workload well. For a batch size of 1 at 200k tokens, that gives me 1152 blocks, which is solid occupancy."

The Performance Data: First Results

The benchmark output shows the first performance measurement of the custom kernel under production conditions with CUDA graphs:

| Context Length | Decode tok/s | Decode Wall Time | |----------------|--------------|------------------| | 1,024 | 135.1 | 0.11s | | 4,096 | (truncated) | (truncated) |

At 1024 tokens, 135.1 tok/s represents a 3.9× improvement over the ~35 tok/s observed in eager mode (from the assistant's earlier analysis in [msg 12307]). This is precisely the benefit of CUDA graphs: by eliminating kernel launch overhead, the GPU spends more time computing and less time waiting for CPU orchestration.

However, as the assistant will discover in the next message ([msg 12309]), this advantage diminishes at longer contexts. At 65k tokens, the custom kernel with graphs achieves only ~3.3 tok/s, compared to ~1.5 tok/s for Triton+graphs—a more modest 2.2× improvement. The graphs eliminate launch overhead, but at long contexts the bottleneck shifts to memory bandwidth and compute, where the custom kernel's advantages are less dramatic.

Assumptions and Their Validation

This message reveals several assumptions that the assistant is implicitly making:

Assumption 1: The Dtype Mapping Is Correct

The assistant assumes that kv_indptr is int32, qo_indptr/mask_indptr/kv_indices/out_cache_loc are int64, and custom_mask is uint8. These assumptions were derived from tracing SGLang's buffer construction code and are critical because mismatched dtypes would cause silent memory corruption. The dtype log check is meant to validate this.

Assumption 2: NSPLIT=16 Is Sufficient

The assistant assumes that 16 splits provide enough parallelism even at long context lengths. For a batch size of 1 at 200k tokens, this yields 1152 blocks (200k / 16 ≈ 12,500 tokens per split, but the actual grid is ceil(ctx_len / chunk_size) * NSPLIT). The assistant's reasoning suggests this is adequate for occupancy, but the benchmark results will reveal whether this holds.

Assumption 3: The Workspace Allocation Strategy Is Correct

The assistant chose per-call torch.empty for workspace allocation, relying on PyTorch's graph-private memory pools to ensure capture safety. This assumption was validated by the successful capture in [msg 12307].

Assumption 4: Prefill "cuda graph: False" Is Benign

The assistant assumes that prefill not using graphs is by design and not indicative of a problem. This is correct for SGLang's architecture, but it's a subtle point that could confuse someone less familiar with the system.

Input Knowledge Required

To fully understand this message, a reader needs:

  1. CUDA Graph Capture: Understanding that CUDA graphs record GPU operations for replay, and that capture requires all operations to be deterministic and not depend on host synchronization.
  2. SGLang's Architecture: Knowledge that SGLang separates prefill and decode paths, and that decode benefits disproportionately from graph capture due to its many small kernel launches.
  3. Verify Attention: Understanding that the verify kernel checks speculative tokens against the full model, and that its performance is critical for speculative decoding throughput.
  4. The sm_120 Constraint: Knowledge that RTX PRO 6000 Blackwell (sm_120) lacks the tensor core instructions available on Hopper and datacenter Blackwell, forcing a custom kernel implementation.
  5. KV Cache and NSPLIT: Understanding how KV-split flash attention works, where the KV sequence is divided into parallel chunks that are computed independently and then reduced.

Output Knowledge Created

This message produces several valuable outputs:

  1. Verified Capture Success: Confirmation that the capture-safe kernel works correctly in production, with no invalidation errors.
  2. First Performance Baseline: The initial measurement of the custom kernel with CUDA graphs, showing 135.1 tok/s at 1024 context length.
  3. Validation of the Engineering Approach: Evidence that the capture-safety rewrite was successful and that the kernel integrates correctly with SGLang's graph infrastructure.
  4. Data for A/B Comparison: The benchmark results that will be compared against the Triton baseline in the next message, enabling a quantitative assessment of the custom kernel's value.

The Broader Narrative

This message is the calm before the storm. In the next message ([msg 12309]), the assistant will analyze the benchmark results, discover that the graph-enabled kernel's advantage over Triton shrinks at longer contexts, and initiate a clean A/B comparison by switching to Triton mode. This will reveal that Triton+graphs achieves 100.4 tok/s at 1024 context (vs 135.1 for the custom kernel) but only 1.5 tok/s at 65k (vs 3.3 for the custom kernel). The custom kernel delivers a 1.35× speedup at short context and 2.2× at long context—a solid but not revolutionary improvement.

The assistant will then pivot to further optimization: increasing NSPLIT to 64, adding 128-bit vectorized bf16 KV loads, and implementing KV defragmentation. These changes will ultimately deliver a 3–6× end-to-end decode speedup over Triton+graphs, as documented in chunk 1 of segment 66.

Conclusion

Message [msg 12308] captures a pivotal moment in a complex engineering effort: the transition from "does it work?" to "how fast is it?" The assistant's reasoning demonstrates deep understanding of CUDA graph capture, SGLang's architecture, and the performance characteristics of custom GPU kernels. The message is deceptively simple—a few commands and some truncated output—but it represents the successful culmination of dozens of prior messages involving kernel rewriting, dtype analysis, memory management, and system integration.

The 135.1 tok/s decode throughput at 1024 context is not just a number; it is the payoff for hours of careful engineering. It validates the capture-safety rewrite, the NSPLIT=16 tradeoff, the workspace allocation strategy, and the dtype mapping. It also sets the stage for further optimization, as the assistant will discover that the real bottleneck has shifted from kernel launch overhead to memory bandwidth and MoE expert imbalance.

In the end, this message is about measurement—the discipline of replacing assumptions with data. The assistant doesn't declare victory based on the capture success alone; it immediately runs a benchmark to quantify the improvement and identify remaining bottlenecks. This empirical approach, combined with deep technical knowledge, is what transforms a working prototype into a production-ready optimization.