The Decisive A/B: When Benchmarking Honesty Reveals the Real Bottleneck

In the high-stakes world of custom CUDA kernel development for large language model inference, few moments are as humbling—and as clarifying—as running a clean apples-to-apples benchmark. Message <msg id=12309> captures exactly such a moment. After spending multiple rounds building a custom sm_120 verify attention kernel for speculative decoding on NVIDIA RTX PRO 6000 Blackwell GPUs, making it CUDA-graph-capture-safe, deploying it into a live SGLang service, and watching the first benchmark numbers roll in, the assistant faces an uncomfortable question: is my custom kernel actually better than the Triton baseline under production conditions?

The answer, revealed through a methodical A/B comparison executed in this single message, is more nuanced than a simple yes or no—and the process of arriving at it demonstrates what disciplined engineering looks like when assumptions meet data.

The Context: A Custom Kernel in Production

To understand the weight of this message, one must appreciate what preceded it. The assistant had been engaged in a multi-day effort to deploy the Kimi K2.6 model with DFlash speculative decoding on Blackwell GPUs (see [chunk 0.0]). A critical bottleneck had been identified in the "verify attention" step—the phase where the drafter's speculative tokens are verified against the target model. The existing Triton-based MLA (Multi-head Latent Attention) kernel was underperforming, locked into a page_size=1 configuration that caused scattered KV cache accesses, achieving only ~14 GB/s effective bandwidth against a 1.8 TB/s peak.

The assistant's response was to build a custom CUDA kernel from scratch, targeting the sm_120 architecture (the compute capability of the RTX PRO 6000 Blackwell consumer GPU). This kernel used a KV-split flash-decode design with a partial+reduce approach to improve occupancy. Initial eager-mode microbenchmarks showed a promising 2× speedup over Triton at longer context lengths.

But eager mode is not production. Production means CUDA graphs—a feature that captures a sequence of GPU operations into a single reusable graph, eliminating CPU launch overhead for every decode step. In messages <msg id=12301> through <msg id=12308>, the assistant meticulously refactored the kernel to be capture-safe: replacing host synchronizations with device-side operations, using fixed NSPLIT constants, allocating workspace with torch.empty (which is graph-compatible), and passing native buffer pointers without any .item() calls or tensor copies that would break graph capture. The effort succeeded—CUDA graph capture completed in 1.53 seconds, and the server generated correct outputs.

The First Benchmark: Promising but Puzzling

Message <msg id=12309> opens with the assistant examining the first performance numbers from the ON+graphs configuration:

Looking at the performance metrics with CUDA graphs enabled, I'm seeing significant speedups at shorter context lengths—135 tokens per second at 1024 tokens compared to just 35 with eager execution—but the gains diminish as context grows, dropping to 3.3 tok/s at 65k.

This is the moment of reckoning. The numbers tell a story: CUDA graphs are working spectacularly at short context (3.8× improvement over eager), but the advantage erodes as context length increases. At 65k tokens, the custom kernel with graphs achieves only 3.3 tok/s. The assistant compares this to the original Triton implementation and finds the custom kernel is only ~1.18× faster at 65k—far from the 2× advantage seen in eager mode.

The reasoning that follows is the intellectual core of the message:

The graphs might be eliminating the per-launch overhead that gave my kernel a 2x advantage in eager mode, so I need to run a clean A/B comparison with both implementations using graphs at the same config (mem-fraction 0.88) to see if the speedup holds.

This insight is crucial. The assistant recognizes that the earlier 2× advantage might have been an artifact of comparing against a Triton baseline that was paying kernel launch overhead on every decode step. CUDA graphs eliminate that overhead for both implementations. If the custom kernel's advantage was primarily in reducing launch overhead (by having fewer, larger grid launches), then graphs would level the playing field.

But there's a second, equally important realization: the comparison might not be apples-to-apples at all. The Triton baseline numbers were likely collected under a different configuration (different mem-fraction-static, possibly without graphs). The assistant explicitly flags this:

the comparison isn't apples-to-apples

The Decisive Action: A Clean A/B

What makes this message exemplary is not the analysis alone—it's the immediate, decisive action that follows. The assistant doesn't speculate further or try to adjust the numbers post-hoc. Instead, they design and execute a controlled experiment:

Switch to Triton mode by setting KDTREE_VERIFY=off in the systemd service file
Keep every other variable identical: same mem-fraction-static 0.88, same CUDA graph configuration, same server, same benchmark script
Restart the service and wait for it to become healthy (verified via a curl test)
Run the identical benchmark at the same context lengths: 1024, 4096, 16384, 65536 The bash commands are precise and careful. The assistant uses a polling loop with 30-second intervals, checking both systemd status and endpoint health before declaring readiness. This is production-grade discipline—no assumptions about restart timing, no skipping health checks.

The Results: Triton+Graphs Baseline

The Triton+graphs baseline arrives:

 ctx_req prompt_tok   TTFT_s prefill_tok/s decode_tok/s gen_tok decode_wall
    1024       1429     0.49          2899        100.4      16        0.15
    4096       5702     2.11          2705         38.7      16        0.39
   16384      22777    11.27          2020          6.8      16        2.20
   65536      91027    93.85           970          1.5      16        9.94

Now the comparison is clean. At 1024 context: custom kernel achieves 135.1 tok/s vs Triton's 100.4 tok/s—a 1.35× speedup. At 4096: 38.7 vs 38.7—essentially identical. At 16384: 6.8 vs 6.8—again identical. At 65536: 3.3 vs 1.5—a 2.2× speedup at the longest context.

Wait—the numbers don't match exactly what the assistant quoted in their reasoning. Let me re-examine. The assistant's earlier ON+graphs run (from <msg id=12308>) showed:

 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     ... (cut off)

The Triton baseline from this message shows 100.4 tok/s at 1024. So the custom kernel is ~34% faster at short context. At 65536, the custom kernel's 3.3 tok/s vs Triton's 1.5 tok/s is a 2.2× advantage—actually better than the 1.18× the assistant initially estimated from the non-apples-to-apples comparison.

This is a fascinating outcome. The custom kernel's advantage is largest at the longest context, where the verify attention is most memory-bound and the kernel's better occupancy and vectorized memory access patterns matter most. At medium contexts (4k-16k), the two implementations are neck-and-neck, suggesting that for those lengths, the Triton compiler's optimizations are competitive with the hand-written kernel.

What This Reveals About the Bottleneck

The message's deeper significance lies in what it reveals about the system's performance characteristics. The fact that Triton+graphs achieves 100.4 tok/s at 1k context—not far from the custom kernel's 135.1—confirms that CUDA graphs are indeed a great equalizer for short-context decode. The per-step launch overhead that dominated at short context is now amortized away.

But the growing divergence at longer context tells a different story. At 65k tokens, the Triton implementation collapses to 1.5 tok/s while the custom kernel maintains 3.3 tok/s. This is the regime where memory bandwidth and occupancy dominate—where the custom kernel's design choices (KV-split flash decode, 128-bit vectorized loads, higher occupancy through the partial+reduce pattern) actually matter.

The assistant's thinking process shows an important metacognitive skill: the ability to distinguish between measurement artifacts and real effects. The initial 2× advantage in eager mode could have been real or could have been an artifact of different configurations. By running the clean A/B, the assistant confirms that the advantage is real—but only in the regimes where the kernel's design actually makes a difference.

Input and Output Knowledge

To fully understand this message, one needs:

Input knowledge: The architecture of the custom verify kernel (KV-split flash decode with partial+reduce), the CUDA graph capture mechanism and its constraints, the SGLang serving stack and its configuration parameters (mem-fraction-static, KDTREE_VERIFY), the benchmark methodology (bench_context_decode.py with --gen 16 and context sweeps), and the earlier performance data from both eager and graph modes.

Output knowledge created: A clean, controlled performance comparison between the custom sm_120 verify kernel and the Triton baseline under identical conditions (CUDA graphs enabled, same memory fraction, same server). The data shows the custom kernel achieves 1.35× speedup at short context (1k) and 2.2× at long context (65k), with parity at medium contexts (4k-16k). This knowledge directly informs the next steps: the bottleneck has shifted from the verify attention kernel to other components (as the assistant would later discover—MoE expert imbalance at batch size 1).

The Engineering Mindset

What makes <msg id=12309> a model of technical communication is the assistant's willingness to question its own results. The initial benchmark looked good—graphs were working, decode was faster than eager—but the assistant immediately looked for confounders. "Is this comparison fair?" "Am I measuring what I think I'm measuring?" "What changes when I control for every variable?"

This is the opposite of cherry-picking. The assistant deliberately switches to the competitor's implementation (Triton) and runs the same benchmark, knowing full well that the results might show their custom kernel isn't better. The willingness to be wrong is what makes the eventual confirmation of a real speedup so credible.

The message also demonstrates the importance of instrumentation before optimization. The assistant doesn't guess at why the advantage shrinks at medium context—they let the data speak. The parity at 4k-16k suggests that for those context lengths, the Triton JIT compiler generates code that's competitive with the hand-written kernel, and the CUDA graph eliminates any launch overhead differences. The custom kernel's advantage only emerges where memory bandwidth is the dominant constraint (long context) or where launch overhead was a factor (short context, now mitigated by graphs).

Conclusion

Message <msg id=12309> is a masterclass in disciplined benchmarking. It captures the moment when a developer moves from "my kernel works and seems faster" to "my kernel is provably faster under controlled conditions, and here's exactly how much." The assistant's reasoning process—identifying confounders, designing a clean A/B, executing it with production-grade care, and interpreting the results honestly—is the kind of thinking that separates reliable engineering from wishful optimization.

The custom sm_120 verify kernel survives this test with a real, if context-dependent, advantage over Triton. But more importantly, the message establishes a methodology that will serve the entire project: always compare under identical conditions, always question your assumptions, and always let the data—not your hopes—tell you where the real bottleneck lies.