Numerical Parity at Last: Validating a Custom CUDA Verify Attention Kernel Against Triton

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When deploying models like Kimi K2.6—a massive MoE (Mixture of Experts) architecture—on specialized hardware like the RTX PRO 6000 Blackwell GPUs, the default attention kernels provided by frameworks like Triton may not be optimal for every use case. This is the story of a single message in a coding session where, after an arduous debugging journey spanning multiple rounds and dozens of tool calls, an AI assistant finally achieved numerical parity between a custom CUDA verify attention kernel and the Triton reference implementation. Message 12287 represents the triumphant culmination of a debugging arc that transformed a fundamentally broken kernel—one producing outputs essentially uncorrelated with the reference—into a drop-in replacement matching Triton to within a single bfloat16 unit in the last place (ULP).

Context: The Verify Attention Kernel Project

To understand the significance of this message, we must first understand what was being built. The assistant was developing a custom "verify attention" kernel for the DDTree (Draft-Tree) speculative decoding system within the SGLang inference framework. Speculative decoding accelerates autoregressive generation by having a smaller "drafter" model propose multiple candidate tokens in parallel, which the larger "target" model then verifies in a single forward pass. The "verify" step computes attention over both the prefix context and the draft tokens, applying a tree-structured visibility mask that determines which draft tokens can attend to which other draft tokens.

The challenge was architectural: the RTX PRO 6000 Blackwell GPUs (sm_120 architecture) have a different instruction set than the data-center Blackwell GPUs (sm_100a/sm_103a) for which optimized MLA (Multi-head Latent Attention) kernels like FlashMLA and cutlass-MLA are compiled. Those kernels use Hopper/Blackwell-DC instructions (wgmma, TMA, tcgen05) that simply don't exist on sm_120's Ada-like ISA. The assistant had to build a custom verify attention kernel from scratch, using only the primitives available on consumer Blackwell hardware.

The message immediately preceding this one (msg 12286) showed the assistant restarting the SGLang service with a fix to the marshaling code—specifically, correcting how the kernel's indices and visibility mask were constructed. The key insight was that SGLang's extend attention structure separates prefix tokens (stored in the KV pool and accessed via kv_indices) from draft tokens (passed as separate k/v tensors and written to out_cache_loc). The assistant's initial implementation had assumed all tokens lived in the pool, leading to a structural mismatch that produced completely wrong attention scores.

The Message: Quantitative and Qualitative Validation

Message 12287 opens with the assistant's reasoning, which immediately signals the breakthrough:

Parity achieved. max_abs_diff is 3.9e-3–7.8e-3 = ~1 bf16 ULP, mean ~1e-4 — my paged+bf16 kernel matches triton's verify output to bf16 rounding on real K2.6 serving tensors (TP8, 8 heads/rank, paged KV, real tree mask).

This is a moment of genuine technical achievement. The numbers tell a compelling story:

The Debugging Journey: From Catastrophic Failure to Parity

To fully appreciate message 12287, we must trace the debugging arc that led to it. The story begins in message 12282, where the assistant first deployed the custom kernel and observed catastrophic failure:

The kernel runs but output is essentially uncorrelated (rel ~1.0) — so a structural marshaling mismatch.

A relative error of ~1.0 means the kernel's output had no correlation with Triton's—it was producing essentially random results. The assistant correctly diagnosed this as a "structural marshaling mismatch" rather than a numerical precision issue or a kernel logic bug. This distinction is crucial: the CUDA kernel itself might be computing attention correctly, but if the input tensors are misaligned or the indexing is wrong, the output will be garbage regardless of how well the kernel is written.

The assistant's investigation revealed the root cause through careful analysis of the tensor shapes logged during execution:

shapes q=(9, 8, 576) kv_buf=(195012, 576) bs=1 q_len=9 Dl=512 Dr=64 kv_len_max=21 scale=0.14468 mask=(270,)

The key insight came from the mask dimensions: mask=(270,) where 270 = 9 × 30. The 9 is the query length (draft tokens), and 30 is the full current sequence length (21 prefix tokens + 9 draft tokens). But kv_len_max=21 only reflected the prefix length stored in kv_indices. The assistant realized:

SGLang's extend structure puts the prefix in kv_indices and the 9 draft tokens are the extend region (written to out_cache_loc), and the mask row-stride is 30, not 21.

This was the fundamental misunderstanding. The assistant had assumed that all KV tokens—both prefix and draft—were accessed through the KV pool via kv_indices. In reality, SGLang's extend attention architecture treats the prefix as a cached pool lookup and the draft tokens as a fresh computation, with the draft tokens' KV being written to new pool slots (out_cache_loc) and passed as separate tensors. The custom kernel was only attending to the 21 prefix tokens, completely missing the 9 draft tokens, and using the wrong mask stride.

The fix was to build combined indices: concatenate the prefix slot indices with the draft token slot locations, and extract the visibility mask using the correct stride of 30 (prefix + q_len) rather than 21 (prefix only). This was implemented across messages 12283 and 12284, followed by a service restart in 12285.

Assumptions and Mistakes

The debugging journey reveals several assumptions that proved incorrect:

1. KV pool containment assumption: The assistant initially assumed that all tokens attended to in a single forward pass would be accessible through the KV pool's kv_indices. This is a natural assumption—the KV cache is designed to store all previously computed key-value pairs. However, SGLang's extend attention architecture makes a distinction between "prefix" tokens (already in the pool) and "extend" tokens (being computed in the current forward pass). The draft tokens in DDTree verification are treated as extend tokens, not prefix tokens, even though they were just computed.

2. Mask stride assumption: The assistant assumed the visibility mask's row stride equaled kv_len_max (the prefix length). In reality, the mask spans the full sequence (prefix + extend), and the stride is the total sequence length. This is because the mask is a logical matrix over all tokens, with the prefix columns skipped during computation but still accounted for in the stride.

3. Ordering assumption about MLA buffer layout: Earlier in the debugging process (messages 12278-12281), the assistant investigated whether the query and key-value tensors used the same dimension ordering (latent-first vs rope-first). This turned out to be correct—both use [512 latent || 64 rope] ordering—but it was a legitimate concern that required verification.

4. The "kernel is correct" assumption: Perhaps the most important meta-assumption was that the CUDA kernel itself was computing attention correctly, and the bug was in the marshaling. This turned out to be right, but it was a hypothesis that needed to be tested. The assistant could have spent days optimizing a kernel that was fundamentally correct but receiving wrong inputs.

Input Knowledge Required

To fully understand this message, one needs:

Technical knowledge:

Output Knowledge Created

This message creates several important pieces of knowledge:

1. Quantitative validation result: The custom kernel matches Triton to within bfloat16 precision on real serving tensors. This is the primary output—a proven drop-in replacement for Triton's verify attention.

2. Baseline generation outputs: Three deterministic generations from the Triton baseline, serving as reference for subsequent validation of the custom kernel.

3. Validation methodology: The approach of first achieving numerical parity on tensors, then validating generation quality on simple prompts, provides a template for similar kernel development efforts.

4. Confidence in the marshaling fix: The dramatic improvement from rel~1.0 (uncorrelated) to ~1 ULP (bfloat16 rounding) confirms that the marshaling fix was correct and complete.

The Thinking Process

The assistant's reasoning in this message reveals a disciplined, methodical approach to validation:

First, the assistant quantifies the error: "max_abs_diff is 3.9e-3–7.8e-3 = ~1 bf16 ULP, mean ~1e-4." This is precise, using the language of floating-point analysis. The assistant doesn't just say "it matches"—it provides the specific error metrics and interprets them in terms of bfloat16 precision.

Second, the assistant contextualizes the validation: "my paged+bf16 kernel matches triton's verify output to bf16 rounding on real K2.6 serving tensors (TP8, 8 heads/rank, paged KV, real tree mask)." Every qualifier matters: "paged" (the kernel handles paged KV cache), "bf16" (the precision), "real K2.6 serving tensors" (not synthetic data), "TP8" (tensor parallelism), "8 heads/rank" (the head count per GPU), "real tree mask" (the actual DDTree visibility pattern).

Third, the assistant plans the next step: capture Triton baseline generations, then flip to the custom kernel. This shows an understanding that tensor-level parity is necessary but not sufficient—the generations must also match. A kernel could theoretically produce slightly different outputs that still lead to the same token selections, or it could produce the same attention outputs but with a different numerical layout that causes downstream operations to diverge.

The choice of test prompts is also thoughtful: simple, factual, and deterministic (temperature=0). These are designed to produce the same output every time with the same kernel, making them ideal for cross-kernel comparison.

Broader Significance

Message 12287 represents a turning point in the session. After this, the assistant can move from debugging to optimization—tuning the kernel for performance, implementing CUDA graph capture, and integrating it into the live serving path. The numerical parity validation is the gate that must be passed before any performance work can begin. Without it, the assistant would be optimizing a broken kernel, chasing throughput improvements on fundamentally wrong computations.

The message also illustrates a broader principle in systems engineering: when building a replacement for an existing component, the first goal must be functional equivalence, not performance. The assistant spent multiple rounds achieving parity before even beginning to optimize. This discipline—validate correctness first, then optimize—is what separates professional engineering from hacking.

Conclusion

Message 12287 is a quiet triumph. It doesn't announce a breakthrough in attention algorithm design or a record-breaking throughput number. It simply states, with precise numerical evidence, that a custom CUDA kernel now matches its reference implementation to within the limits of floating-point precision. The debugging journey that led to this moment—from catastrophic failure (rel~1.0) to bfloat16 ULP—is a testament to systematic debugging, careful hypothesis testing, and the value of understanding the framework's data flow before writing kernel code. For the reader who hasn't followed the entire session, this message captures the essence of what makes good systems engineering: methodical validation, precise measurement, and the discipline to verify correctness before declaring victory.