The Moment of Truth: Validating Three Custom CUDA Kernels for Speculative Decoding on Blackwell GPUs

Introduction

In the high-stakes world of large language model inference, few moments are as tense as the first full test suite run. After days of designing, implementing, and debugging custom CUDA kernels, the assistant in this opencode session reaches that critical juncture. Message [msg 11894] is deceptively simple — a single bash command followed by truncated test output — but it represents the culmination of an extraordinary engineering effort: building a complete native C/C++/CUDA inference engine for the Kimi K2.6 model, a 1-trillion-parameter Mixture-of-Experts (MoE) language model, running on 8× NVIDIA RTX PRO 6000 Blackwell GPUs.

The message itself reads:

[assistant] [bash] cd /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine bash scripts/build_and_test.sh 2>&1 | tail -30 Start 11: tree_build_wide 11/23 Test #11: tree_build_wide .................. Passed 0.24 sec Start 12: verify_attn_h16_p2048_b16 12/23 Test #12: verify_attn_h16_p2048_b16 ........ Passed 0.29 sec Start 13: verify_attn_h64_p128_b32 13/23 Test #13: verify_attn_h64_p128_b32 ......... Passed 0.25 sec Start 14: verify_attn_h64_p512_b8 14/23 Test #14: verify_attn_h64_p512_b8 .......... Passed 0.25 sec Start 15: verify_attn_h8_p0_b8 15/23 Test #15: veri...

The output is truncated, but the pattern is unmistakable: all tests are passing. This article examines this single message in depth — why it was written, what decisions led to this moment, the assumptions underpinning it, and the knowledge it creates.

Why This Message Was Written: The Culmination of a Multi-Phase Plan

This message was not an arbitrary status check. It was the validation gate for Phase 1 of a carefully planned engineering roadmap. To understand its significance, we must trace the reasoning that led here.

The session began with a user directive at [msg 11826]: "re-read relevant files, plan out C/CUDA high speed DDTree inference stack for Kimi K2.6 and 8x Nvidia PRO 6000 Blackwell." The assistant responded by reading dozens of files — the DDTree findings report, the SGLang patch files, the drafter configuration, the K2.6 model implementation, the existing CUDA kernels in sgl-kernel, and the hardware specifications. This research phase produced a comprehensive plan ([msg 11845]) organized into four phases:

How Decisions Were Made: Architecture and Design Choices

The path to this message was paved with deliberate engineering decisions. The most consequential was the architecture choice articulated in the plan at [msg 11845]: build a native C++ loop with new custom kernels for the fragile/missing pieces, but reuse the battle-tested Marlin INT4 MoE grouped GEMM kernels and NCCL from sgl-kernel. This avoided rewriting millions of lines of proven CUDA code while still eliminating the Python overhead and cuda-graph fragility that plagued the SGLang integration.

For the three kernels themselves, the design decisions reveal deep understanding of the hardware:

GPU Tree Builder ([msg 11860]): The assistant chose a block-per-request design where thread 0 manages a best-first heap expansion in shared memory, while all threads cooperatively populate output arrays. The heap uses double-precision accumulation to match the numpy reference bit-exactly, and ancestor relationships are encoded as bitmasks using two uint64 words per node — an elegant solution for the small tree sizes (budget ≤ 64 nodes).

Tree-Verify MLA Attention ([msg 11876]): This kernel implements MLA-absorb attention with a visibility mask. The design uses one block per (batch, head, query) position, stores attention scores in shared memory, and computes online softmax with block-wide reduction. The critical insight documented in the reasoning is that for production, the kernel should only handle the small tree tail — the long prefix should use the existing FlashMLA kernel, with results merged via log-sum-exp.

Greedy Tree-Accept ([msg 11886]): One thread per request walks the verified tree using the next_token/next_sibling encoding, checking target predictions against children. The design is simple by intent — this is the lightweight glue between the verify and commit stages.

The test strategy was equally deliberate. Rather than testing end-to-end on the 1-trillion-parameter model (which would require 8 GPUs and hours per run), the assistant created a binary container format (KDTR) that allows Python to generate reference test data with known inputs and expected outputs, which the C++ test harness loads and compares against kernel results. This enables rapid iteration on a local RTX 5070 Ti (same sm_120 architecture as the target PRO 6000) without needing the full model.

Assumptions Underpinning the Validation

Every test suite rests on assumptions, and this one is no exception:

  1. Architecture equivalence: The local RTX 5070 Ti (sm_120, CUDA 13.2) is assumed to be representative of the 8× RTX PRO 6000 Blackwell GPUs on the target machine (CT200). Both share the same compute capability, so kernels that pass locally should behave identically on the target.
  2. Synthetic data fidelity: The random test data generated by the Python reference scripts is assumed to exercise the same computational paths as real model activations. The tree builder tests use random top-k log probabilities and token IDs; the verify-attn tests use random MLA-shaped tensors with realistic dimensions.
  3. Bit-exact equivalence for tree builder and accept: The tree builder and tree-accept kernels are tested for bit-exact equality with the numpy reference. This assumes that the double-precision accumulation in the GPU kernel matches Python's float64 arithmetic exactly — a strong assumption that the assistant validated by carefully matching operation order.
  4. Float tolerance for verify-attn: The attention kernel is tested with a tolerance of ~2e-8 absolute error, acknowledging that floating-point non-determinism in CUDA block reductions can produce tiny variations.
  5. Isolation validity: Passing isolated kernel tests is assumed to predict correct behavior when the kernels are composed into the full DDTree pipeline. This is reasonable for the greedy path (where each kernel's output feeds directly into the next), but doesn't account for integration issues like memory layout mismatches or NCCL communication problems that would surface in Phase 2.

Mistakes and Incorrect Assumptions

The most visible issue in this message is the truncated output. The tail -30 command cuts off the test results, so we see tests 11–15 but not tests 1–10 or 16–23. The assistant likely intended to show the full suite but was limited by terminal output. This is a minor operational error — the full results are visible in the subsequent message ([msg 11895]), which confirms all 23 tests pass and verifies that the accept tests exercise meaningful accept lengths (1 through 8 tokens).

A more subtle assumption is that passing unit tests on synthetic data guarantees correctness on real model data. The tree builder tests use random log probabilities with distinct values, avoiding the edge case of exact ties. In production, log probabilities could theoretically produce ties, and the kernel's tie-breaking heuristic (insertion order) might diverge from Python's tuple comparison. The assistant acknowledged this risk in the reasoning at [msg 11855]: "exact float ties are measure-zero in practice and I can document this minor divergence."

Similarly, the verify-attn kernel is tested with random MLA-shaped tensors, but real model activations have specific statistical properties (e.g., the latent vectors lie on a low-dimensional manifold learned during training). The kernel might pass synthetic tests but encounter numerical issues with real activations — though the float32 tolerance of 2e-8 provides a comfortable safety margin.

Input Knowledge Required to Understand This Message

A reader needs substantial context to grasp the significance of these 23 passing tests:

Output Knowledge Created by This Message

This message produces several forms of knowledge:

  1. Validation evidence: All 23 tests pass, confirming that the three custom CUDA kernels produce correct results matching the numpy references. The test times (~0.2–0.3 seconds each) also provide a performance baseline — these kernels are fast enough for the target application.
  2. Milestone confirmation: Phase 1 of the four-phase plan is complete. The greedy DDTree step trio (build → verify-attn → accept) is fully implemented and validated in CUDA on sm_120 hardware. This de-risks the hardest algorithmic components before moving to Phase 2's full engine integration.
  3. Reproducible test infrastructure: The build_and_test.sh script (created at [msg 11870]) provides a repeatable validation pipeline that regenerates reference data, builds the project, and runs all tests. Anyone with the repository and compatible hardware can reproduce these results.
  4. Architecture portability confirmation: The tests run on a local RTX 5070 Ti (consumer GPU) but target 8× RTX PRO 6000 (workstation GPU). The passing tests confirm that sm_120 CUDA code developed on consumer hardware will work on the target system — a non-trivial validation given that Blackwell's sm_120 is a new architecture with potential compatibility surprises.

The Thinking Process Visible in Reasoning

While message [msg 11894] itself contains no reasoning (it is purely a command and its output), the surrounding messages reveal the thinking that led to this moment. The assistant's reasoning traces a careful path from design through implementation to validation:

In [msg 11884], after validating the verify-attn kernel, the assistant considers what to do next: "Before wrapping up, I should consider what else I can accomplish in this session that's self-contained: an end-to-end integration test chaining the tree builder output directly into the verify-attention kernel to validate their interface contract, or implementing the greedy accept logic as a kernel." This shows strategic thinking about maximizing the value of the current session — choosing the tree-accept kernel because it "completes the trio of build → verify-attn → accept entirely in CUDA."

In [msg 11855], the assistant wrestles with tie-breaking: "The main parity concern is handling tie-breaking in the heap when cumulative log-weights are equal. Python's heapq uses tuple comparison on the ranks, which is complex to replicate exactly." The solution — using insertion order as a secondary key and relying on the measure-zero probability of exact ties — shows pragmatic engineering judgment.

In [msg 11860], the assistant designs the shared memory layout with meticulous attention to resource constraints: "Bounding the heap size: I pop at most budget times, and each pop nets +1 element (remove 1, add up to 2). Starting from 1 element, the maximum size during execution is 1 + budget. With budget ≤ 64, I'll allocate capacity of 68 to be safe — that's about 1.6KB of shared memory for the heap arrays, which is fine."

In [msg 11874], the assistant considers the verify-attn kernel's 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." This forward-looking design acknowledges that the current correctness-first kernel will need optimization for production — a mature engineering perspective.

Conclusion

Message [msg 11894] appears mundane — a test suite run with passing results — but it represents a pivotal moment in a complex engineering project. The 23 passing tests validate that three custom CUDA kernels, designed and implemented over dozens of messages, correctly implement the core of the DDTree speculative decoding algorithm on Blackwell GPUs. The GPU tree builder replaces SGLang's per-request CPU heapq with a parallel block-per-request design. The verify-attn kernel handles MLA-absorb attention with tree visibility masks at sequence lengths up to 2065 tokens. The tree-accept kernel implements greedy follow-verified-tree on the GPU, completing the pipeline.

The message also demonstrates a disciplined engineering methodology: build reusable test infrastructure (KDTR format, numpy references, CMake/ctest integration) before writing kernels; validate against synthetic data before testing on the real model; and establish clear phase boundaries with explicit exit criteria. This approach de-risks the project by catching algorithmic errors early, when they are cheap to fix, rather than during expensive full-model integration.

The truncated output is a reminder that even well-executed plans have minor imperfections — but the underlying signal is clear: Phase 1 is complete, and the foundation for a native C/C++/CUDA DDTree inference engine is solid. The next phase — building the full transformer engine with weight loading, NCCL communication, and the complete decode loop — can proceed with confidence that the core kernels are correct.