Six Green Lights: Validating a Custom CUDA MLA Attention Kernel for Speculative Decoding

The Message

The subject message is deceptively simple — a single bash command and its output:

cd /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine
ctest --test-dir build --output-on-failure -R verify_attn 2>&1 | tail -20
Test project /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/build
    Start 12: verify_attn_h16_p2048_b16
1/6 Test #12: verify_attn_h16_p2048_b16 ........   Passed    0.28 sec
    Start 13: verify_attn_h64_p128_b32
2/6 Test #13: verify_attn_h64_p128_b32 .........   Passed    0.23 sec
    Start 14: verify_attn_h64_p512_b8
3/6 Test #14: verify_attn_h64_p512_b8 ..........   Passed    0.32 sec
    Start 15: verify_attn_h8_p0_b8
4/6 Test #15: verify_attn_h8_p0_b8 .............   Passed    0.24 ...

Six tests, all passing. Each completes in roughly a quarter of a second. The output is terse, almost anticlimactic. But behind this quiet moment lies the culmination of an intense design and implementation effort: a custom CUDA kernel for Multi-Head Latent Attention (MLA) with tree-structured visibility masking, written from scratch, validated against a numpy reference, and now confirmed correct across a range of configurations that span from trivial to near-production scale.

This article examines this single message in depth — why it was written, what decisions led to it, what assumptions were baked into the kernel design, and what knowledge was produced by its successful execution.

The Broader Mission: A Native DDTree Inference Engine

To understand why this message matters, one must understand the project it belongs to. The assistant is building a complete native C/C++/CUDA inference engine for Draft-Tree-based speculative decoding (DDTree) — a technique where a lightweight "drafter" model proposes multiple candidate token sequences arranged as a tree, and the large "target" model verifies them in parallel using a single forward pass. The engine targets the Kimi K2.6 model running on 8× NVIDIA RTX PRO 6000 Blackwell GPUs.

The project is organized into phases. Phase 0 established the build infrastructure (CMake with CUDA 13 targeting sm_120 architecture), a binary container format called KDTR for sharing test data between Python and C++, and faithful numpy reference implementations of the DDTree algorithms. Phase 1 calls for three custom CUDA kernels that replace SGLang's existing CPU-based implementations:

  1. A GPU best-first tree builder — replaces SGLang's per-request CPU heapq for constructing the draft tree from top-k probabilities.
  2. A tree-verify MLA attention kernel — the subject of this message — which computes attention over the tree's candidate tokens using the visibility mask, with MLA's absorbed latent representation.
  3. A greedy tree-accept kernel — walks the verified tree to determine which tokens were accepted by the target model. The first kernel (tree builder) had already been committed with 11 passing tests, validated bit-exact against the numpy reference. The message analyzed here represents the validation of the second kernel.

Why This Message Was Written: The Verification Imperative

The assistant wrote this message to answer a single, critical question: Does the custom CUDA attention kernel produce numerically correct results? In the context of a production inference engine, this is not an academic question. A kernel that silently produces wrong logits would corrupt generations, waste GPU cycles, and undermine trust in the entire speculative decoding pipeline.

The assistant had just completed an intensive design-and-code session spanning several messages. In [msg 11874], the assistant worked through the MLA absorb attention mathematics, considering different kernel strategies — block-per-query-row with online softmax, atomicAdd-based accumulation, shared memory score storage — before settling on a three-phase design. In [msg 11875], the reference generator was written, producing KDTR bundles across diverse configurations. In [msg 11876], the kernel implementation was finalized with careful attention to shared memory budgets, thread mapping, and edge cases. In [msg 11877] and [msg 11878], the kernel source and test harness were written. In [msg 11879]–[msg 11880], the CMake build system was wired up. In [msg 11881], the reference data was generated and the project was rebuilt.

Now, in this message, the assistant runs the test suite. The -R verify_attn flag filters ctest to run only the verify-attention tests, and --output-on-failure ensures that any failure would print detailed diagnostics. The tail -20 truncation suggests the assistant expected clean output — a reasonable assumption given the careful construction — but the test run is the moment of truth where all assumptions meet reality.

The Six Configurations: What They Test

The six test configurations are not arbitrary. Each probes a different dimension of the kernel's correctness:

Assumptions Baked Into the Kernel

The kernel's correctness depends on several assumptions that the test suite implicitly validates:

Assumption 1: Shared memory is sufficient. The kernel stores attention scores in shared memory as float values, requiring kv_len * sizeof(float) bytes. For the largest test (kv_len=2065), this is ~8KB — well within the ~48KB shared memory per block on modern GPUs. The assumption holds for these test configurations, but the assistant explicitly noted in [msg 11876] that production use with very long contexts would require switching to a flash-attention-style tiling approach.

Assumption 2: The mask is well-formed. The kernel assumes every query position has at least one valid key to attend to (its own position). This prevents division-by-zero in the softmax normalization. The visibility mask generated by the tree builder guarantees this property, but the kernel includes a guard anyway.

Assumption 3: Float32 precision is sufficient. The reference implementation and the kernel both use float32. The assistant anticipated that numerical errors would be on the order of 1e-8 absolute and 1e-3 relative (confirmed in [msg 11883]), which is acceptable for attention computations where the final logit precision matters at roughly 1e-4 to 1e-5.

Assumption 4: The block-per-(batch, head, query) mapping is efficient. Each CUDA block handles one query position for one head in one batch item. This creates many small, independent blocks — potentially thousands for large batches — which is a natural fit for GPU parallelism. However, it also means each block processes the full key-value sequence independently, which could be inefficient for very long sequences. The assistant acknowledged this and planned to split attention into prefix (reusing FlashMLA) and tree tail (custom kernel) for production.

Assumption 5: The KDTR binary format is correct. The test harness reads reference data from KDTR files generated by the Python reference script. Any mismatch between the Python writer and the C++ reader would produce spurious failures. The assistant built both sides of this interface carefully, and the passing tests confirm the round-trip is correct.

What This Message Created: Output Knowledge

The six passing tests produce concrete knowledge:

  1. The kernel is numerically correct. The CUDA implementation matches the numpy reference across all tested configurations. This is not a proof of correctness for all inputs, but it is strong evidence that the algorithm is implemented faithfully.
  2. The kernel is fast enough for testing. Each test completes in 0.23–0.32 seconds, including data transfer to/from the GPU. This means the test suite can be run frequently during development without becoming a bottleneck.
  3. The shared-memory approach works at realistic scales. The h16_p2048_b16 test, with 2065 tokens, demonstrates that the kernel can handle prefix lengths in the thousands — relevant for production use where context windows of 32k–200k tokens are common (though the production split-attention strategy would handle the long prefix separately).
  4. The design decisions are validated. The three-phase approach (score computation → softmax normalization → weighted accumulation) works correctly. The choice to parallelize output accumulation over dimensions rather than key positions avoids atomicAdd contention. The decision to store scores in shared memory rather than recomputing them is validated.
  5. The KDTR test infrastructure is reliable. The round-trip from Python reference generation to C++ test execution produces consistent results, establishing a pattern that can be reused for future kernels (as it was for the tree builder and would be for the upcoming accept kernel).

The Thinking Process: A Window Into Kernel Design

The assistant's reasoning in the messages leading up to this test run reveals a sophisticated design process. In [msg 11874], the assistant works through the MLA absorb attention mathematics in detail, considering multiple implementation strategies:

"I'll go with a cleaner approach: one block per (b, h, i) where I first compute all scores into shared memory with threads strided over the key-value positions, then reduce to find the max and softmax denominator. Next, I'll normalize the probabilities in shared, then compute the output by parallelizing over the 512 output dimensions—each thread handles a few dimensions and loops through all key-value positions to accumulate the weighted sum. This avoids atomics entirely."

This passage reveals a key design tension: the trade-off between parallelism and memory contention. The initial approach — having each thread own a key position and accumulate its contribution via atomicAdd — would create 512 atomic operations per key position, a prohibitive cost. The chosen approach avoids atomics by having threads own output dimensions and loop over keys, which is naturally parallel and contention-free.

The assistant also thinks ahead to production integration:

"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. That's cleaner than trying to optimize everything in one kernel."

This is a mature architectural insight. Rather than building a monolithic kernel that handles all cases poorly, the assistant plans to compose existing optimized kernels (FlashMLA for the prefix) with the custom kernel (for the tree tail), using log-sum-exp merging — a standard technique for split attention. This shows that the current kernel is understood as a building block within a larger system, not an end in itself.

Mistakes and Incorrect Assumptions

The message itself reveals no mistakes — six tests pass cleanly. However, the broader context reveals some subtle issues:

The .kdtr file tracking. In [msg 11874], the assistant notes that KDTR files were committed to git despite the .gitignore excluding .npy and .bin files. The assistant decided this was acceptable since the files are small and make tests reproducible. This is a reasonable judgment, but it means the repository now contains binary artifacts that must be regenerated if the reference implementation changes — a potential maintenance burden.

The numerical tolerance. The test harness uses a floating-point comparison with tolerance. The assistant did not specify the exact tolerance in the visible reasoning, but the subsequent message ([msg 11883]) reveals max_abs errors of ~2e-8 and max_rel errors of ~4e-3. The relative error is inflated by near-zero denominators in the softmax output, but the assistant correctly interprets this as benign. A less experienced developer might have been alarmed by the 0.4% relative error without understanding its source.

The assumption that all configurations are equally representative. The six tests cover a range of parameters, but they do not test every combination. For example, there is no test with both high head count (64) and long prefix (2048) simultaneously — the h64_p512_b8 test has only 512 prefix tokens. The assistant likely judged that the kernel's behavior is separable across these dimensions (head count and prefix length don't interact in the kernel's control flow), but this is an unverified assumption.

Conclusion

The message analyzed here — a simple ctest invocation with six passing tests — represents a significant milestone in the construction of a custom CUDA inference engine. The tree-verify MLA attention kernel, designed through careful reasoning about memory layouts, thread mapping, and numerical stability, has been validated against a numpy reference across a strategic range of configurations. The six green lights are not just a passing grade; they are evidence that the assistant's design decisions — the three-phase computation, the dimension-parallel output accumulation, the shared memory score storage, the split-attention architecture for production — are sound.

This message also illustrates a broader principle of systems building: the moment of verification is always anticlimactic. The drama happens in the design phase, in the reasoning about trade-offs, in the careful construction of reference implementations and test harnesses. The test run itself is a formality — a check that reality matches the model the builder has constructed in their mind. When all six tests pass, it means the model was accurate. The kernel works as intended.

But the work is not done. The assistant immediately moved on to the third kernel — the greedy tree-accept kernel — and then to integration and deployment. Each step would require its own moment of verification, its own six green lights. This message is a snapshot of one such moment, frozen in time, showing what it looks like when careful design meets successful execution.