The Moment of Truth: Validating a Custom CUDA Attention Kernel for Speculative Decoding

The Message

cd /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine
export PATH=/opt/cuda/bin:$PATH
python3 python/gen_verify_attn_refs.py 2>&1
echo "=== reconfigure + build ==="
cmake -S . -B build -DCMAKE_CUDA_COMPILER=/opt/cuda/bin/nvcc -DKDTREE_CUDA_ARCH=120 >/dev/null 2>&1
cmake --build build -j 2>&1 | tail -6
wrote /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/../tests/refs/verify_attn_h8_p0_b8.kdtr  B=2 H=8 prefix=0 q_len=9 kv_len=9
wrote /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/../tests/refs/verify_attn_h8_p1_b4.kdtr  B=3 H=8 prefix=1 q_len=5 kv_len=6
wrote /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/../tests/refs/verify_attn_h8_p37_b8.kdtr  B=2 H=8 prefix=37 q_len=9 kv_len=46
wrote /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/../tests...

This unassuming bash command, issued at message index 11881 in a sprawling coding session spanning dozens of messages, represents a critical juncture in the development of a custom CUDA inference engine for speculative decoding. It is the moment where design meets reality—where carefully reasoned algorithms, painstakingly crafted in numpy reference implementations and CUDA kernel code, are finally compiled and tested against each other. The message is a single tool call—a bash invocation—but it encapsulates the culmination of an intense design and implementation phase spanning the preceding seven messages ([msg 11874] through [msg 11880]).

Context: Building a Native DDTree Inference Engine

To understand why this message matters, one must understand what the assistant is building. The project is a native C/C++/CUDA inference engine for the Kimi K2.6 language model, specifically targeting the DDTree (Draft-Driven Tree) speculative decoding algorithm. Speculative decoding accelerates autoregressive generation by using a smaller, cheaper "drafter" model to propose multiple candidate tokens in parallel, which a larger "target" model then verifies. DDTree is an advanced variant where the drafter proposes a tree of candidate sequences rather than a single linear chain, allowing more tokens to be verified per forward pass.

The assistant has been constructing this engine from scratch in a new repository called kdtree-engine/. The architecture is organized in phases, with Phase 1 dedicated to implementing three custom CUDA kernels that replace critical components of the SGLang inference framework's CPU-based implementations:

  1. Tree builder kernel — a GPU best-first tree builder that constructs the draft tree on-device, replacing SGLang's per-request CPU heapq approach.
  2. Tree-verify attention kernel — a custom MLA (Multi-head Latent Attention) attention kernel that computes attention scores only over the tree's visible positions, using the visibility mask generated by the tree builder.
  3. Tree-accept kernel — a greedy walk that determines which tokens from the verified tree are accepted into the output sequence. The subject message is specifically about the second kernel: the tree-verify MLA attention kernel. The preceding messages show the assistant working through the design of this kernel in exhaustive detail.

The Design Journey: From Theory to Silicon

The reasoning visible in the messages leading up to [msg 11881] reveals a remarkable depth of engineering thinking. In [msg 11874], the assistant works through the MLA attention mathematics from first principles. MLA is a form of attention used by DeepSeek and Kimi models where the key-value cache stores a compressed latent representation shared across all attention heads, rather than storing separate keys and values per head. This dramatically reduces memory consumption—a critical advantage for long-context inference.

The assistant's reasoning traces through the attention computation: "The attention score combines the dot product of the absorbed query with the cached latent plus the rope components, then I apply softmax over the allowed positions based on the mask." It then considers the kernel's output contract: "I'll keep the value projection separate as a standard GEMM since that's cleaner and matches how SGLang structures the operation. So the kernel outputs the latent directly, and a follow-up matrix multiply applies the per-head value weights."

This design decision—to have the kernel output the attention-weighted latent rather than the final per-head output—is a deliberate architectural choice. It separates concerns: the custom kernel handles the masked attention computation, while a standard cuBLAS GEMM handles the value projection. This mirrors SGLang's own structure, making eventual integration cleaner.

The assistant then works through the kernel's computational strategy in meticulous detail. It considers and rejects an approach using atomicAdd for the weighted sum accumulation, recognizing that "512 atomics per j" would be prohibitive. Instead, it devises a cleaner three-phase approach: first compute all attention scores into shared memory, then find the softmax maximum and denominator via block-wide reduction, then compute the output by parallelizing over the output dimensions rather than over key positions. This avoids atomics entirely.

Crucially, the assistant also anticipates the production deployment architecture: "For the production system, I can split attention into two parts: reuse the existing flash MLA decode kernel for the long prefix, then apply our custom masked kernel just to the tree tail, and merge the results using log-sum-exp." This is a sophisticated insight—the custom kernel is not meant to replace the entire attention computation, only the portion that requires the tree visibility mask. The long prefix (which can be thousands of tokens) can be handled by the highly optimized FlashMLA kernel, while the custom kernel handles only the small tree tail (typically fewer than 100 tokens). This hybrid approach maximizes performance while maintaining correctness.

The Reference Test Harness: A Bridge Between Python and CUDA

A key architectural decision visible in this message is the KDTR (Kernel Data Transfer and Reference) binary container format. The assistant designed this format to serve as a bridge between Python and C++/CUDA. The Python reference generator (gen_verify_attn_refs.py) produces .kdtr files containing the input tensors and expected output for each test configuration. The C++ test harness reads these files, runs the CUDA kernel, and compares the results.

This design choice embodies a critical engineering principle: test data should be generated once and version-controlled, not recomputed on every test run. By writing the reference data to binary files that are committed to the repository, the assistant ensures that tests are deterministic and reproducible. A developer can run the tests without having Python or numpy installed—they only need the compiled C++ test executable and the reference files.

The output of the message shows four reference files being generated, each with different configurations:

Assumptions and Their Validity

The message and its surrounding reasoning rest on several assumptions, most of which proved sound:

Assumption 1: The numpy reference implementation is correct. The entire testing strategy depends on the Python reference producing the "ground truth" output. If the reference has a bug, the kernel will be validated against incorrect behavior. The assistant mitigates this by keeping the reference simple and mathematically transparent—it's a direct implementation of the attention formula without any optimizations that could introduce subtle errors.

Assumption 2: Float32 precision is sufficient for validation. The kernel computes in float32, and the reference also uses float32. The assistant later confirms (in [msg 11883]) that the maximum absolute error is ~2e-8, which is at the limit of float32 precision. The relative error of 0.1–0.4% is noted as being "inflated by near-zero denominators rather than actual computational issues." This is a correct diagnosis—when the attention probability for a position is near zero, tiny absolute differences can produce large relative differences.

Assumption 3: The shared memory budget is sufficient for the test configurations. The assistant calculated that a key-value length of ~2000 tokens requires ~8KB of shared memory for the score array, which fits within the ~10.5KB budget per block. This assumption holds for the test configurations but the assistant correctly notes that for production with longer sequences, a flash-attention-style tiling approach would be needed.

Assumption 4: The block-per-query-head parallelism strategy is appropriate. The kernel launches one block per (batch, head, query_position) triple. This is a straightforward mapping that works well when the number of query positions is small (as it is in speculative decoding, where the tree typically has 10–100 nodes). For the test configurations with q_len up to 17, this is fine.

The Outcome: Validation and Its Implications

The message itself only shows the reference generation and compilation succeeding. The actual test results come in the next message ([msg 11882]), where all 6 verify-attention tests pass, including the challenging configurations with 64 heads and 2048-token prefixes. The subsequent message ([msg 11883]) reveals the numerical error margins: maximum absolute error of 2.98e-8, maximum relative error of 0.4%. These are excellent results for float32 computation.

The success of this validation has far-reaching implications for the project. With the tree builder (validated in <msg id=11868–11871>) and the verify-attention kernel (validated here), the assistant has two of the three Phase 1 kernels working correctly. The third—the tree-accept kernel—is implemented shortly after ([msg 11884]), completing the trio.

More importantly, the validation confirms that the architectural decisions were correct. The MLA absorb-form attention computation, the separation of value projection into a separate GEMM, the shared-memory-based softmax strategy—all of these design choices are validated by the bit-exact (within float32 precision) agreement between the CUDA kernel and the numpy reference.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

  1. Multi-head Latent Attention (MLA): The attention mechanism used by DeepSeek V3 and Kimi K2.6 models, where keys and values are compressed into a shared latent space.
  2. Speculative decoding with draft trees: The DDTree algorithm, where a drafter proposes a tree of candidate tokens and a target model verifies them in parallel.
  3. CUDA kernel design patterns: Block-per-element parallelism, shared memory management, online softmax, and memory coalescing considerations.
  4. The KDTR binary format: The custom container format designed by the assistant for sharing test data between Python and C++.
  5. The SGLang inference framework: The existing system that the custom kernels are designed to eventually replace or augment.

Output Knowledge Created

This message creates several forms of knowledge:

  1. Validated reference data: Four .kdtr files containing input tensors and expected outputs for various attention configurations. These serve as the ground truth for all future testing of the verify-attention kernel.
  2. A compiled test executable: The test_verify_attn binary that can validate the kernel against any KDTR file.
  3. Confidence in the kernel design: The successful compilation (with the CMake reconfigure and build) confirms that the CUDA code is syntactically correct and that the kernel interface matches the test harness expectations.
  4. A reproducible test pipeline: The combination of the Python reference generator, the CMake build system, and the C++ test harness creates a pipeline that can be rerun at any time to validate the kernel.

The Thinking Process: A Window into Engineering Excellence

What makes this message remarkable is not the command itself, but the thinking that produced it. The assistant's reasoning, visible in the preceding messages, shows a systematic approach to complex CUDA kernel development:

  1. Start with the mathematics: Before writing any code, the assistant works through the MLA attention formula in detail, understanding exactly what computation needs to happen.
  2. Build a reference implementation in numpy: This provides a ground truth that is easy to debug and modify, independent of CUDA complexity.
  3. Design the kernel architecture on paper: The assistant considers multiple strategies (atomicAdd vs. shared memory, per-key vs. per-dimension parallelism) and evaluates their trade-offs before committing to code.
  4. Anticipate production constraints: Even while building a "correctness-first" kernel, the assistant plans for the production deployment where the kernel will be used only for the tree tail, with the long prefix handled by FlashMLA.
  5. Create a robust test harness: The KDTR format, the Python reference generator, and the C++ test executable form a complete validation pipeline that ensures correctness is maintained across code changes.
  6. Validate incrementally: Rather than building all three kernels and testing them together, the assistant validates each kernel independently, catching issues early. This message, then, is the hinge point in this process—the moment where design, implementation, and validation converge. It is the sound of a carefully constructed intellectual edifice meeting the unforgiving reality of silicon and floating-point arithmetic. And it passes the test.