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

In the high-stakes world of speculative decoding for large language models, attention kernels are the engine room. Every microsecond shaved off the verify-attention step translates directly into faster generation, lower latency, and a better user experience. Message 12291 in this opencode session represents a critical inflection point: the moment when a handcrafted, architecture-specific CUDA kernel—built from scratch for NVIDIA's sm_120 (RTX PRO 6000 Blackwell) instruction set—faces its first honest, apples-to-apples benchmark against the established incumbent, Triton. The results, as we shall see, are both surprising and illuminating.

The Context: A Long Road to a Custom Kernel

To understand the weight of message 12291, one must appreciate the journey that preceded it. The assistant had been working for dozens of rounds to deploy the Kimi K2.6 model with DFlash speculative decoding on a cluster of RTX PRO 6000 Blackwell GPUs (sm_120 architecture). A persistent bottleneck had been identified: the verify attention kernel—the step that checks draft tokens against the full model's distribution—was using Triton's MLA (Multi-head Latent Attention) implementation, which was compiled only for sm_90a, sm_100a, and sm_103a architectures. None of the optimized MLA kernels (FlashMLA, cutlass-MLA, flashinfer-MLA) supported sm_120, and they relied on Hopper/Blackwell-DC instructions (wgmma, TMA, tcgen05) that simply do not exist on sm_120's Ada-like ISA.

The assistant's response was ambitious: build an owned sm_120 verify attention kernel from scratch. This involved writing a KV-split flash-decode MLA kernel (verify_attn_flash.cu) with a partial+reduce design, validating it against a naive oracle, making it CUDA-graph capture-safe by eliminating host synchronizations and dynamic allocations, and then integrating it into the live SGLang service via a Python backend module. By message 12290, the kernel was serving real K2.6 DDTree attention in "ON" mode—the custom kernel's output was being used for actual generation, and it matched Triton's output to within bfloat16 precision.

But correctness is only half the battle. The kernel had to be faster.

The Question That Drove Message 12291

Message 12291 opens with the assistant's reasoning, which reveals a subtle but important concern: the benchmark data available was contaminated by confounding variables. The earlier measurements of the custom kernel's performance came from a different code revision and mixed CUDA-graph mode with eager mode. The assistant writes:

"The issue is that my earlier benchmark used different context windows entirely—64k versus the 91k I'm testing now—and mixed graph mode with eager mode, so they're not directly comparable."

This is a crucial methodological insight. In ML inference benchmarking, the difference between CUDA-graph mode (where the entire GPU execution graph is captured and replayed without CPU intervention) and eager mode (where the Python runtime orchestrates each kernel launch) can be enormous—often a factor of 2–3× in decode throughput. Comparing a graph-mode Triton baseline against an eager-mode custom kernel would be fundamentally unfair and would produce misleading conclusions.

The assistant's solution is elegant and disciplined: disable CUDA graphs for both configurations, creating a level playing field where the only variable is the attention kernel itself. By toggling the KDTREE_VERIFY environment variable between on (custom kernel) and off (Triton), and ensuring --disable-cuda-graph is active in both cases, the assistant isolates the kernel's intrinsic performance from the orchestration overhead.

The Benchmark Protocol

The assistant uses a dedicated benchmark script (bench_context_decode.py) that measures decode throughput at three context lengths: 4,096 tokens, 16,384 tokens, and 65,536 tokens. Each measurement generates 16 tokens and reports the decode tokens per second. The prompt lengths vary slightly between runs (5,702 vs 5,705 tokens at the 4k setting, for example) but are essentially identical, providing a fair comparison.

The protocol proceeds in two phases:

Phase 1 — ON mode (custom kernel, eager): The assistant restarts the service with KDTREE_VERIFY=on and runs the benchmark. The results show decode throughput degrading with context length: 35.2 tok/s at 5.7k context, 13.6 tok/s at 23k context, and 3.6 tok/s at 91k context. The degradation is expected—attention is O(n²) in compute and O(n) in memory bandwidth, so longer contexts always slow down—but the question is whether the custom kernel degrades more gracefully than Triton.

Phase 2 — OFF mode (Triton, eager): The assistant edits the systemd service file to set KDTREE_VERIFY=off, reloads the daemon, restarts the service, waits for it to become healthy (verified by a curl health check), and runs the identical benchmark. The results: 38.0 tok/s at 5.7k context, 6.7 tok/s at 23k context, and 1.8 tok/s at 91k context.

The Results: A Tale of Two Regimes

The comparison reveals a fascinating bifurcation:

| Context Length | Custom Kernel (ON) | Triton (OFF) | Ratio | |---|---|---|---| | ~5.7k tokens | 35.2 tok/s | 38.0 tok/s | 0.93× | | ~23k tokens | 13.6 tok/s | 6.7 tok/s | 2.03× | | ~91k tokens | 3.6 tok/s | 1.8 tok/s | 2.00× |

At short context (4k), Triton holds a slight edge—about 8% faster. This is unsurprising. Triton's JIT-compiled kernels benefit from years of optimization for short-sequence attention, and the overhead of the custom kernel's partial+reduce design (which splits KV into chunks and reduces) may not be amortized over a small number of tokens.

But at medium and long context (16k and 65k), the custom kernel achieves a consistent 2× speedup over Triton. This is a remarkable result for a first-generation hand-written kernel targeting an architecture that no existing MLA library supports. The speedup is almost exactly 2× at both 23k and 91k tokens, suggesting the kernel's advantage is structural—likely related to better memory coalescing, higher occupancy, or more efficient use of sm_120's shared memory and register file—rather than a one-time algorithmic improvement that saturates.

Why the Custom Kernel Wins at Long Context

The assistant's earlier analysis provides clues about the root cause of this speedup. The Triton MLA implementation, when running on sm_120, was observed to be "occupancy-starved in the TP8 regime" (only 8 heads per rank). The custom kernel's KV-split flash-decode design, combined with 128-bit vectorized bf16 KV loads and an increased NSPLIT parameter, was specifically engineered to improve occupancy and memory bandwidth utilization.

At short context, the attention computation is small enough that Triton's highly tuned tile sizes and scheduling can overcome any architectural mismatch. But as context grows, the memory bandwidth bottleneck becomes dominant. The custom kernel's ability to stream KV data more efficiently—achieving higher effective bandwidth from HBM to the compute units—manifests as a consistent 2× advantage that persists across the long-context regime.

This is precisely the kind of result that validates the enormous effort of writing a custom kernel: when the off-the-shelf solution is architecturally mismatched (Triton's MLA targets sm_90/sm_100 features that don't exist on sm_120), a purpose-built kernel can recover much of the lost performance, even in its first iteration.

Assumptions and Methodological Rigor

The assistant makes several assumptions in this message, all of which are reasonable and defensible:

Eager-mode comparison is fair. By disabling CUDA graphs for both configurations, the assistant ensures that any graph-related overhead or optimization applies equally. This is the cleanest way to isolate kernel performance.
The benchmark is representative. The bench_context_decode.py script generates 16 tokens per request at each context length. This is a reasonable proxy for the decode phase of speculative decoding, where each step verifies a small number of draft tokens against a long prefix.
The context lengths span the relevant regime. Testing at 4k, 16k, and 65k covers short conversations, medium documents, and long-form contexts, providing a clear picture of scaling behavior.
The service is in a steady state. The assistant waits 330 seconds for the service to become healthy after each restart, ensuring that any warm-up effects (CUDA context initialization, kernel compilation caching) have settled. One subtle assumption worth examining is that the KDTREE_VERIFY=off configuration truly represents "pure Triton" with no residual custom-kernel influence. The assistant's backend module monkeypatches SGLang's attention path; when verification is off, the patch is still present but returns Triton's output instead of the custom kernel's. The monkeypatch itself could theoretically introduce overhead (e.g., Python function call overhead, conditional branching). However, the near-identical prefill throughput between the two configurations (2,908 vs 2,909 tok/s at 4k context) suggests that any such overhead is negligible.

What This Message Creates: Output Knowledge

Message 12291 produces several pieces of valuable knowledge:

Quantitative benchmark data: The first clean, apples-to-apples comparison of the custom sm_120 verify kernel against Triton on real K2.6 serving workloads. The 2× speedup at long context is a concrete, reproducible result that justifies continued investment in the custom kernel approach.

Scaling law characterization: The data reveals that the custom kernel's advantage grows with context length, from parity at 4k to 2× at 16k+. This suggests the kernel's design is particularly well-suited to the memory-bandwidth-bound regime that dominates long-context inference.

Methodological precedent: The assistant's decision to control for CUDA graphs by disabling them in both configurations establishes a rigorous benchmark methodology. Future performance comparisons in this project will benefit from this clean separation of variables.

Confidence in production readiness: The custom kernel has now passed two critical tests: numerical correctness (bf16 parity with Triton on real tensors) and performance competitiveness (2× faster at long context). While further optimization is possible, the kernel is demonstrably ready for production deployment in the long-context regime where it excels.

The Thinking Process: A Window into Engineering Judgment

The assistant's reasoning in message 12291 reveals a mature engineering mindset. The key insight—that earlier benchmarks were confounded by mixing graph and eager modes—demonstrates an understanding that benchmarking is not just about collecting numbers but about constructing valid comparisons. The assistant recognizes that the question "Is my kernel faster than Triton?" can only be answered when all other variables are controlled.

The reasoning also shows awareness of the benchmark script's limitations. The assistant notes that the earlier measurements used "different context windows entirely" and that the current test at 65k context actually has 91k prompt tokens (due to the model's internal processing). This attention to the precise conditions of measurement is essential for reproducible science.

Perhaps most telling is the assistant's willingness to run the experiment even though the outcome was uncertain. At short context, the custom kernel was slower than Triton. An engineer who had invested days in writing a custom kernel might be tempted to cherry-pick favorable benchmarks or explain away unfavorable results. Instead, the assistant runs the comparison openly, lets the data speak, and discovers that the kernel's advantage emerges at the context lengths that matter most for the deployment scenario (long-context speculative decoding).

The Broader Significance

Message 12291 sits at the intersection of several important themes in modern ML systems engineering:

The death of the universal kernel: For years, the ML community relied on a small set of highly optimized kernels (cuBLAS, FlashAttention) that worked well across architectures. As GPU architectures diversify—with sm_90 (Hopper), sm_100 (Blackwell DC), and sm_120 (Blackwell consumer) diverging in ISA features, shared memory sizes, and tensor core capabilities—the era of universal kernels is ending. Architecture-specific optimization is becoming necessary for peak performance.

The value of owning your stack: By building a custom kernel rather than waiting for library support, the assistant gains performance that would otherwise be unavailable on sm_120 hardware. This is a powerful demonstration of the "own your inference stack" philosophy that companies like Apple, Tesla, and OpenAI have adopted.

The importance of clean measurement: The assistant's methodological care—controlling for graph mode, using identical benchmarks, verifying service health—is a model for how to evaluate performance improvements in complex inference systems. A sloppier evaluation might have concluded that the custom kernel was slower (by comparing against a graph-mode Triton baseline) and abandoned a 2× improvement.

Conclusion

Message 12291 is a quiet triumph. In a few lines of benchmark output, it validates weeks of kernel engineering and provides a clear, data-driven answer to the central question: is the custom sm_120 verify attention kernel worth the effort? The answer is a definitive yes—at the context lengths that matter for production speculative decoding, the kernel delivers a consistent 2× speedup over Triton. The message also demonstrates the kind of rigorous engineering thinking—controlling variables, questioning assumptions, letting data drive decisions—that separates professional systems work from hobbyist tinkering. For anyone building inference infrastructure on non-standard GPU architectures, this message offers both inspiration and a methodological template.