The Moment of Truth: Benchmarking a Custom sm_120 Verify Kernel Against Triton

In the high-stakes world of large language model inference, few moments are as revealing as the first head-to-head benchmark between a hand-written CUDA kernel and the established baseline. Message [msg 12292] captures exactly such a moment: the assistant has just completed an A/B performance comparison of its custom sm_120 verify attention kernel against Triton's implementation, running both in eager mode (with CUDA graphs disabled) on a live Kimi K2.6 service. The results are in, and they tell a story of triumph, trade-offs, and unfinished business.

This message is the culmination of a long engineering arc that began when the assistant discovered that all existing optimized MLA (Multi-head Latent Attention) kernels—FlashMLA, cutlass-MLA, flashinfer-MLA—were compiled only for sm_90a, sm_100a, and sm_103a architectures, none of which support the sm_120 instruction set found on the RTX PRO 6000 Blackwell consumer GPU. With no off-the-shelf solution available, the assistant built an owned sm_120 verify attention kernel from scratch, iterated through correctness validation, CUDA graph capture safety, occupancy tuning, and vectorized memory loads, and finally integrated it into the live SGLang serving stack. Message [msg 12292] is where the assistant processes the benchmark data that validates (and qualifies) that entire effort.

The Raw Numbers: A 2× Win at Long Context

The core of the message is a clean three-row table that the assistant constructs from the benchmark runs in the preceding messages:

| Context Length | Triton (tok/s) | Custom Kernel (tok/s) | Speedup | |---|---|---|---| | 5.7k tokens | 38.0 | 35.2 | 0.93× | | 23k tokens | 6.7 | 13.6 | 2.0× | | 91k tokens | 1.8 | 3.6 | 2.0× |

These numbers are the product of careful experimental design. In message [msg 12290], the assistant measured the custom kernel's performance in ON mode (eager, no CUDA graphs). In message [msg 12291], it switched the KDTREE_VERIFY environment variable to off, restarted the service (a ~5.5-minute wait for the model to load), and ran the identical benchmark suite against Triton, also in eager mode. This apples-to-apples comparison removes CUDA graph acceleration as a confounding variable—both sides run with --disable-cuda-graph—so the difference is purely the attention kernel itself.

The results are striking. At the shortest context (5.7k tokens), the custom kernel is actually slower than Triton by about 7%. But at 23k tokens, it achieves exactly 2× the throughput, and at 91k tokens, it maintains that exact 2× advantage. The crossover point where the custom kernel pulls ahead lies somewhere between 5.7k and 23k tokens.

Interpreting the Crossover: The Marshaling Tax

The assistant immediately diagnoses the short-context regression: "The short-context regression comes from Python marshaling overhead per layer that becomes negligible once attention computation dominates at longer sequences." This is a classic pattern in GPU kernel optimization. When the sequence length is short, the attention computation itself is fast, so the per-layer overhead of marshaling data—building index arrays, copying them to the GPU, launching the kernel—dominates the wall time. As the sequence grows, the O(n²) attention computation (or in the case of MLA, the O(n) decode with O(n) memory access) eventually dwarfs the fixed marshaling cost, and the kernel's intrinsic advantages shine through.

The assistant identifies two possible remedies: "building combined indices once per step instead of per layer, or eliminate it entirely with CUDA graph capture." The first is a software optimization that would reduce Python-side overhead; the second is a more fundamental restructuring that would allow the entire kernel launch sequence to be captured as a CUDA graph and replayed without host involvement. The assistant correctly notes that CUDA graph capture is the more powerful solution, as it eliminates all per-step Python overhead, not just the index-building portion.

The CUDA Graph Dilemma

A central tension runs through this message: the custom kernel delivers a clean 2× speedup at long context in eager mode, but the production service was originally running with CUDA graphs enabled, which provided substantial acceleration for the Triton baseline. The assistant weighs two scenarios:

  1. Leave the service in ON+eager mode: The custom kernel runs, but CUDA graphs are disabled globally (because the kernel path isn't capture-safe). At long context, the 2× kernel win more than compensates for the loss of graph acceleration. At short context, the service would be slower than the original Triton+graph configuration.
  2. Revert to Triton+graphs: The known-good production baseline is restored. The custom kernel remains available via environment flags for experimentation and future graph-capture work. The assistant's reasoning is nuanced. It compares the custom kernel's 3.6 tok/s at 91k (eager) against an interpolated ~2.0 tok/s for the original Triton+graph baseline, concluding that "even with eager overhead I'm ahead at long context where attention dominates." But it also acknowledges that "at short context though, graphs matter more and eager would be slower than the original setup."

The Decision: Revert, but Keep the Kernel Alive

After weighing the options, the assistant makes a pragmatic decision: "Since the service isn't handling real traffic and the cuda-graph path is incomplete, I'll revert to the original production config as the known-good baseline, keeping the kdtree integration available via environment flags."

This decision reflects a mature engineering judgment. The service is a development/deployment testbed, not a production system serving user traffic. The priority is to maintain a known-good baseline while continuing to develop the graph-capture path. Reverting doesn't mean abandoning the custom kernel—it means treating the current state as a successful milestone (2× long-context speedup validated) and deferring the final integration until the graph-capture work is complete.

The assistant also demonstrates awareness of the user's priorities: "I'm weighing whether to leave the kdtree optimization enabled in production or revert to the safe baseline—the long-context win is real and what the user cared about." This shows the assistant is not just optimizing in a vacuum but is attuned to the project's goals.

The Thinking Process: A Model of Systematic Reasoning

The "Agent Reasoning" section of the message reveals the assistant's internal thought process in remarkable detail. It proceeds through several stages:

  1. Observation: The benchmark data shows a clear 2× speedup at long context with a slight regression at short context.
  2. Diagnosis: The short-context regression is attributed to Python marshaling overhead per layer, which becomes negligible once attention dominates.
  3. Forward path identification: CUDA graph capture is identified as the ideal solution, requiring "overriding the metadata replay logic and using pre-allocated static buffers for a capturable kernel call."
  4. Trade-off analysis: The assistant compares eager-mode custom kernel vs. Triton+graphs, acknowledging that the comparison is context-dependent.
  5. Decision: Revert to Triton+graphs as the production baseline, document the graph-capture work as the next phase.
  6. Action: Write the results file, commit the integration code, update the plan. This structured reasoning—observe, diagnose, identify solutions, analyze trade-offs, decide, act—is a template for disciplined engineering decision-making. The assistant doesn't just report numbers; it interprets them in the context of the system's architecture and the project's priorities.

Assumptions and Their Validity

The message contains several assumptions worth examining:

Assumption 1: The short-context regression is purely Python marshaling overhead. This is a reasonable hypothesis given the architecture (per-layer Python loops building index arrays), but it's not yet proven by profiling. The assistant acknowledges this implicitly by noting that CUDA graph capture would "eliminate it entirely," which would only be true if the overhead is indeed on the CPU side. If there's a kernel-level issue (e.g., the custom kernel has higher launch latency or worse occupancy at small problem sizes), graph capture alone wouldn't fully close the gap. The assistant's subsequent work in the chunk (increasing NSPLIT and adding vectorized loads) would address kernel-level issues, but at this point in the conversation, the diagnosis is still a hypothesis.

Assumption 2: The interpolated Triton+graph baseline at 91k is ~2.0 tok/s. The assistant doesn't have a direct measurement at 91k with graphs enabled, so it interpolates from earlier benchmarks at different context lengths. This is a reasonable approximation but introduces uncertainty. The actual graph-mode performance could be higher or lower depending on how CUDA graph acceleration interacts with the Triton kernel at very long contexts.

Assumption 3: CUDA graph capture is the correct path forward. This is almost certainly correct for production deployment, but it's worth noting that the assistant doesn't consider alternative approaches, such as batching multiple decode requests to amortize the marshaling overhead, or moving the marshaling logic into a CUDA kernel (a "marshaling kernel" that runs on the GPU). The focus on graph capture is consistent with SGLang's architecture and the project's direction, but it's not the only possible optimization.

Input Knowledge Required

To fully understand this message, the reader needs:

Output Knowledge Created

This message produces several forms of output knowledge:

  1. Benchmark results: The 2× speedup at long context is a validated result that can inform deployment decisions and be cited in reports. The assistant writes these to a file (verify_kernel_inservice.txt) for permanent record.
  2. A decision framework: The trade-off analysis between eager-mode custom kernel and Triton+graphs provides a template for similar decisions in the future. The assistant's reasoning about when to prioritize long-context performance over short-context performance is reusable knowledge.
  3. A roadmap for graph capture: The assistant identifies the specific technical challenges for making the kernel graph-capturable: "overriding the metadata replay logic and using pre-allocated static buffers for a capturable kernel call." This is actionable knowledge for the next engineering phase.
  4. A documented baseline: By reverting to Triton+graphs and recording the benchmark numbers, the assistant establishes a clear baseline against which future improvements can be measured.

The Broader Context: From Kernel to System

This message is remarkable not just for the benchmark numbers it reports, but for the system-level thinking it demonstrates. The assistant doesn't treat the kernel as an isolated component; it evaluates the kernel's performance in the context of the full serving stack, considering CUDA graph interactions, Python marshaling overhead, and the service's operational state (not handling real traffic). This holistic perspective is essential for deploying custom kernels in production systems, where a 2× kernel improvement can be negated by integration overhead or regressions in other parts of the system.

The message also illustrates the iterative nature of GPU kernel development. The custom kernel started as a correctness prototype, was validated against Triton's output, then optimized for occupancy and memory bandwidth, then made capture-safe, and finally benchmarked in the live service. Each iteration revealed new bottlenecks and opportunities. The 2× speedup at long context is a significant milestone, but the assistant correctly identifies that the journey isn't over—CUDA graph capture and KV defragmentation remain as next steps.

Conclusion: A Pivotal Moment of Validation

Message [msg 12292] is a pivotal moment in the engineering narrative of this coding session. It marks the transition from "does the kernel work?" to "is the kernel worth it?" The answer is a qualified yes: at long context, the custom sm_120 verify kernel delivers a clean 2× speedup over Triton, validating the investment in building an owned kernel from scratch. But the qualification matters—the short-context regression and the absence of CUDA graph support mean the kernel isn't ready for unconditional production deployment.

The assistant's handling of this moment is exemplary. It doesn't oversell the results or ignore the trade-offs. It presents the data clearly, reasons through the implications, makes a defensible decision, and produces a concrete artifact (the results file) that captures the state of knowledge. This is engineering decision-making at its best: grounded in measurement, aware of context, and oriented toward action.

The 2× speedup at long context is a genuine achievement, but perhaps the more important output of this message is the framework for thinking about kernel integration that the assistant demonstrates—a framework that will serve well in the subsequent phases of CUDA graph capture, KV defragmentation, and beyond.