The Moment of Truth: Validating a Custom CUDA Attention Kernel in Production

In the high-stakes world of custom CUDA kernel development for large language model inference, there is a moment that separates confident progress from wishful thinking: the moment you flip the switch from "compare against the reference" to "trust your own output." Message [msg 12289] captures exactly this moment for a custom sm_120 verify attention kernel being integrated into a live SGLang serving stack for the Kimi K2.6 model with DDTree speculative decoding. It is a message that appears, on its surface, to be a routine service restart and smoke test. But beneath the bash loops and curl commands lies a carefully orchestrated validation ritual that represents the culmination of hours of debugging, architectural reasoning, and kernel optimization.

The Road to This Moment

To understand why this message matters, 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 across Blackwell GPUs (RTX PRO 6000, sm_120 architecture). A critical bottleneck had been identified: the DDTree verify attention step — the mechanism that checks whether draft tokens generated by the speculative drafter are accepted by the target model — was locked to a Triton-based MLA (Multi-head Latent Attention) kernel that achieved only ~14 GB/s effective bandwidth, roughly 130× below the hardware's 1.8 TB/s peak. The root cause was that all optimized MLA kernels (FlashMLA, cutlass-MLA, flashinfer-MLA) were compiled only for sm_90a, sm_100a, and sm_103a architectures — none supported sm_120, the consumer Blackwell ISA that lacks Hopper/Blackwell-DC instructions like wgmma, TMA, and tcgen05.

The assistant's response was to build a custom sm_120 verify attention kernel from scratch, documented in a detailed plan (plans/0002-sm120-verify-kernel-defrag.md). The initial implementation (verify_attn_flash.cu) used a KV-split flash-decode design with partial+reduce, but early microbenchmarks showed it was slower than the naive Triton kernel at short prefixes. Through systematic profiling and tuning — increasing NSPLIT from 16 to 64, adding 128-bit vectorized bf16 KV loads — the assistant achieved a dramatic 3–6× end-to-end decode speedup over Triton with CUDA graphs enabled.

But speed is worthless without correctness. The kernel had been producing wildly incorrect output (relative error ~1.0, meaning results were essentially uncorrelated with Triton's) due to a marshaling mismatch: the assistant had assumed that draft tokens' key-value data was included in the kv_indices passed to the attention kernel, when in fact SGLang's extend attention structure places prefix tokens in kv_indices and draft tokens as separate k/v tensors written to out_cache_loc. The mask row-stride was also miscalculated. Fixing these marshaling issues in [msg 12283]–[msg 12284] brought the numerical error down to bf16 rounding precision — a maximum absolute difference of 3.9e-3 to 7.8e-3, or approximately one unit in the last place (ULP) for bfloat16.

What the Message Actually Does

Message [msg 12289] is the validation step that follows numerical parity. The assistant has already confirmed, in [msg 12287], that the kernel's output matches Triton's to bf16 precision on real K2.6 serving tensors (TP8, 8 heads per rank, paged KV, real tree mask). But numerical parity on intermediate tensors does not guarantee that the generated text will be identical — subtle accumulation of bf16 rounding errors across multiple decode steps could theoretically diverge, especially in the speculative decoding context where a single wrong token could cascade.

The message executes three phases:

Phase 1: Service restart with KDTREE_VERIFY=on. In the previous message ([msg 12288]), the assistant had already edited the systemd service file to change the environment variable from KDTREE_VERIFY=validate to KDTREE_VERIFY=on. The validate mode runs both the custom kernel and the Triton reference, computing diffs but returning Triton's output. The on mode runs only the custom kernel and returns its output directly — this is the mode that will be used in production. The service restart triggers a full SGLang initialization, model loading, and warmup.

Phase 2: Wait loop with health check. The assistant uses a for loop with 30-second intervals (up to 26 iterations, ~13 minutes max) to wait for the service to become ready. The health check sends a minimal completion request ("prompt":"hi","max_tokens":3) and checks for the presence of choices in the JSON response. This is a pragmatic approach: rather than parsing systemd status or log files for readiness signals, the assistant directly probes the HTTP endpoint — the only definitive indicator that the model is loaded, the kernel is compiled, and the inference pipeline is operational. The service becomes ready at approximately 330 seconds (5.5 minutes), which is consistent with loading a large MoE model like K2.6 across 8 GPUs.

Phase 3: Generation comparison. The assistant runs the same three prompts that were used to capture the Triton baseline in [msg 12287]:

  1. "The capital of France is" → generates 24 tokens at temperature 0
  2. "Write one sentence about the ocean:" → generates 24 tokens at temperature 0
  3. "List three primes:" → generates 24 tokens at temperature 0 These prompts are chosen strategically. The first is a factual knowledge question with a deterministic answer (Paris). The second tests creative language generation. The third tests mathematical reasoning. Together, they cover different modes of model behavior. Temperature 0 ensures deterministic sampling — the same logits should always produce the same tokens. The output for the first prompt is: ' Paris. Paris is located in the north-central part of the country, along the Seine River. It is the most pop' This is identical to the Triton baseline captured in [msg 12287]. The second prompt also matches. The third prompt's output is truncated in the message but the assistant's reasoning (visible in the agent reasoning block) confirms all three match.

The Significance of Identical Output

The fact that the generations match exactly is not trivial. The kernel is operating at bf16 precision with a fundamentally different computation structure — a KV-split flash-decode with partial sums and reduction, rather than Triton's monolithic attention. The numerical paths diverge in how intermediate sums are accumulated, how floating-point rounding occurs, and how the softmax is computed. Yet the final token sequences are identical.

This validates several assumptions:

  1. The marshaling is correct. The combined kv_indices (prefix slots concatenated with draft token slots) and the visibility mask extraction with the correct stride produce the same attention pattern as SGLang's native extend mechanism.
  2. The kernel's numerical precision is sufficient. Bfloat16 rounding errors, while present at the individual element level (as shown by the 3.9e-3 max_abs_diff), do not accumulate to change the argmax decision in the final linear layers. This is expected for bf16 — the 7-bit mantissa provides enough precision that small perturbations in attention output are absorbed by subsequent layers without changing the most likely token.
  3. The CUDA graph capture works correctly. The kernel had been made capture-safe in the previous chunk by rewriting it to consume SGLang's native static buffers directly, with no host syncs, copies, or cudaMalloc calls. The fact that the service serves correct generations confirms that the CUDA graph replays the kernel correctly on every invocation.
  4. The TP8 (tensor parallelism across 8 GPUs) communication is intact. The kernel operates on a per-rank basis (8 heads per rank out of 64 total), and the all-reduce across ranks must produce identical results to the Triton path. The matching generations confirm that the distributed attention computation is correct.

What This Message Does Not Say

The message is notably sparse in its celebration. There is no "success!" banner, no summary of what was achieved. The assistant simply reports the generations and checks for errors. This is characteristic of the engineering mindset at this stage — the validation is a checkpoint, not a finish line. The assistant knows, from the analysis in the previous chunk, that the remaining bottleneck is MoE expert imbalance at batch size 1 — a structural limitation of tensor parallelism with small batch sizes, where the attention cost has been fixed but the MoE feed-forward layers are still underutilized.

The message also does not show the third prompt's output in full (it's truncated in the display), but the assistant's reasoning confirms it matches. The error check returns no entries, confirming clean operation.

The Broader Context: Why This Validation Matters

This message sits at the intersection of several themes that define the entire session:

Custom kernel development for novel hardware. The RTX PRO 6000 Blackwell (sm_120) is a consumer-grade GPU that lacks the specialized tensor core instructions of its data-center siblings (sm_100a/sm_103a). The assistant had to build a kernel from scratch because no existing optimized MLA implementation supported this architecture. The success of this validation proves that a custom approach can match — and in throughput, exceed — the Triton baseline on unsupported hardware.

The tension between speed and correctness. The kernel optimization journey — increasing NSPLIT, adding vectorized loads, making CUDA graph capture-safe — could have introduced subtle bugs at any step. The validation methodology (numerical parity first, then generation comparison) provides a rigorous gate before declaring success.

Production deployment discipline. The assistant does not validate with a single prompt but with three diverse prompts. The service is restarted cleanly (not hot-swapped). The health check uses the actual HTTP endpoint rather than process-level signals. These are small but important practices that increase confidence in the deployment.

Conclusion

Message [msg 12289] is the quiet validation that follows a storm of debugging. After hours of tracing tensor shapes, fixing marshaling bugs, tuning kernel parameters, and making CUDA graphs capture-safe, the assistant finally sees the custom sm_120 verify attention kernel produce the same text as the Triton reference — word for word, token for token. The generations match exactly. The service runs cleanly. The bottleneck has moved on.

This is not the end of the story — MoE imbalance awaits, and the Tier 1 defrag (live KV relocation) remains unimplemented. But it is a critical inflection point. The custom kernel is no longer an experimental prototype running in validate mode; it is the active compute engine for one of the most challenging operations in speculative decoding. The 3–6× throughput improvement over Triton is now real, validated, and deployed.