Building a Native CUDA DDTree Inference Engine: From First Principles to Validated Kernels

Introduction

In the high-stakes world of large language model inference, few engineering challenges match the complexity of building a custom GPU inference engine from scratch. This article chronicles one such effort: the construction of a complete native C/C++/CUDA Draft-Driven Tree (DDTree) speculative decoding engine for the Kimi K2.6 model, targeting 8× NVIDIA RTX PRO 6000 Blackwell GPUs. Over the course of a single intensive coding session spanning nearly 75 messages, an AI assistant and its user partner progressed from architectural planning through infrastructure construction, custom CUDA kernel development, numerical validation, and finally to a working MVP engine that matched a numpy golden reference token-for-token.

The work is organized into the kdtree-engine/ repository, a self-contained codebase that implements the entire DDTree speculative decode pipeline—GPU tree building, tree-verify MLA attention, greedy tree acceptance, and a full FP32 transformer forward pass—without depending on SGLang's Python runtime. This article synthesizes the key phases, decisions, and outcomes of that effort, drawing on the detailed message articles that document each step.

Phase 0: Foundations

Every ambitious engineering project begins with infrastructure, and this one was no exception. Phase 0 established three critical foundations: a build system, a data interchange format, and reference implementations.

The Build System

The assistant configured a CMake-based build system targeting CUDA 13.2 with -arch=sm_120, the compute capability corresponding to NVIDIA's Blackwell GPU architecture (see [msg 11863], [msg 11864]). This was not a trivial exercise—the build system had to compile custom CUDA kernels alongside C++ test harnesses, link against the existing sgl-kernel library for Marlin INT4 MoE GEMMs, and support both local development on an RTX 5070 Ti and deployment on the target 8× PRO 6000 system. The CMake configuration used enable_language(CUDA) and set CMAKE_CUDA_ARCHITECTURES=120, ensuring that kernels compiled for one sm_120 device would run on any other [39][40].

The first build attempt revealed a critical lesson: CUDA 13.2 on Blackwell requires explicit __syncwarp calls rather than the deprecated __syncthreads_count and similar intrinsics (see [msg 11866]). The assistant fixed these issues, achieving a clean build that became the foundation for all subsequent kernel development [42][43].

The KDTR Binary Container

One of the most elegant infrastructure decisions was the creation of the KDTR binary container format (see [msg 11856]). This format serves as a bridge between Python (where reference implementations are written in numpy) and C++/CUDA (where the actual kernels run). A KDTR file stores typed, named tensors with shape information, allowing the Python reference generators to produce test inputs that the C++ test harness can load directly onto the GPU for comparison.

The KDTR format was designed with simplicity and debuggability in mind. It stores tensors in a flat binary layout with a header describing the number of tensors, their names, data types, and shapes. This makes it easy to inspect test data with a hex editor or write custom tools to manipulate it. The format proved essential for the entire validation pipeline: every kernel test in Phases 1 and 2 relied on KDTR files to carry reference data from Python to C++ [32].

Numpy Reference Implementations

Before writing a single line of CUDA, the assistant implemented faithful numpy reference versions of all DDTree algorithms (see [msg 11857], [msg 11858]). These references were direct ports of SGLang's build_ddtree_tree_from_topk function and the associated attention and acceptance logic. The references served as the "golden standard" against which all GPU kernels would be measured.

The reference implementations were tested across a range of configurations—default, mid, wide, underfull, chain, maxbudget—and their outputs were validated for structural invariants (e.g., parent indices strictly increasing, visibility matrices correctly encoding ancestor relationships). This validation of the references themselves was a critical meta-step: if the reference is wrong, the GPU kernel will faithfully reproduce the wrong behavior, and all tests will pass against an incorrect standard [33][34][35].

Phase 1: The Three CUDA Kernels

Phase 1 delivered the heart of the DDTree engine: three custom CUDA kernels implementing the complete greedy speculative decode pipeline on the GPU.

Kernel 1: GPU Best-First Tree Builder

The tree builder kernel (see [msg 11860], [msg 11861]) replaces SGLang's per-request CPU heapq with a parallel GPU implementation. The design uses one CUDA block per request, with thread 0 managing a best-first heap expansion in shared memory while all threads cooperatively populate output arrays.

The kernel's shared memory layout was calculated with meticulous attention to resource constraints: maximum queue length of 65 (budget 64 + 1), heap storage of ~576 bytes for keys, ~1152 bytes for metadata, ~1920 bytes for token/depth/logw arrays, ~260 bytes for parent indices, and ~1040 bytes for visibility masks—totaling a few kilobytes, well within the 48KB shared memory budget of modern GPUs [36].

One of the most elegant design choices was the use of bitmask-based ancestor sets for computing the visibility matrix. With a maximum tree size of 65 nodes, each node's ancestor set fits in two uint64 words (128 bits). Building these masks is O(n) sequential work: each node's mask is simply its parent's mask OR'd with its own bit. Expanding to the full visibility matrix is then O(n²) parallel work that all threads can share [36].

The assistant paid extraordinary attention to numerical determinism, ensuring that the GPU kernel produces bit-exact outputs matching the numpy reference. This required matching the reference's operation order exactly, using double-precision accumulation for log-probabilities, and providing deterministic tie-breaking via insertion sequence numbers [36][37].

Validation came through 11 test cases covering edge cases like chain topology (topk=1), underfull trees (fewer candidates than budget), maximum budget (64 nodes), maximum depth (15), and high batch counts (64 simultaneous requests). All 11 tests passed bit-exact against the numpy reference (see [msg 11871]) [47].

Kernel 2: Tree-Verify MLA Attention Kernel

The verify-attention kernel (see [msg 11874], [msg 11876], [msg 11877]) implements MLA-absorb attention with a tree visibility mask. This is the most technically novel component: it takes a shared prefix of latent KV (paged, page_size=1) plus q_len (≤64) tree-node queries and a [q_len × q_len] visibility mask, and produces the attention output for each tree node.

The kernel 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 assistant's 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 [50][52].

Validation was performed against random MLA-shaped tensors with realistic dimensions (head counts from 8 to 64, prefix lengths from 0 to 2048, budgets from 8 to 32). Six test configurations passed with a maximum absolute error of approximately 2e-8—essentially float32 precision limits (see [msg 11882], [msg 11883]) [57][58][59].

Kernel 3: Greedy Tree-Accept Kernel

The tree-accept kernel (see [msg 11886], [msg 11887]) implements the greedy follow-verified-tree algorithm on the GPU. One thread per request walks the verified tree using the next_token/next_sibling encoding, checking target predictions against children. If a child matches, the thread advances to that child; otherwise, it stops and records the bonus token.

The assistant designed synthetic test data that exercises accept lengths from 1 to 8 tokens, covering partial accepts, complete chains, and early terminations (see [msg 11885]). The test generation strategy used a 70% probability of following a child token to create long acceptance chains, with random tokens for early-stop and bonus cases [61][62][63].

All three kernels were validated together in an on-device composition test that chained build → accept without host round-trips, proving that the entire greedy DDTree pipeline could run entirely on the GPU (see [msg 11894], [msg 11895]) [70][71].

Phase 2: The MVP Native Engine

With the three kernels validated, the assistant built a working MVP native engine implementing a full DeepSeekV3/Kimi-style MLA+MoE transformer in FP32 (see [msg 11884]). The engine used cuBLAS GEMMs as placeholders for the INT4 Marlin kernels that would be integrated later, but included all the critical infrastructure:

Deployment Preparation and Diagnostics

After the engine was validated locally, the assistant prepared for deployment on the CT200 server (the 8× PRO 6000 Blackwell box). This involved:

  1. Adding nvcc-direct build scripts since CT200 lacks CMake
  2. Kernel microbenchmarks at K2.6-realistic shapes to measure performance characteristics
  3. A Python head-to-head benchmark comparing the GPU tree builder against SGLang's actual CPU implementation Upon deploying to CT200, a crash in the tree_accept kernel from cyclic random test data was fixed with a safety bound. The verify_attn benchmarks were collected, confirming that the documented production architecture—reuse FlashMLA for long prefix attention, custom kernel only for the small tree tail—is the correct path forward [70].

The Engineering Methodology

Throughout this effort, the assistant demonstrated a disciplined engineering methodology that is worth examining in its own right:

Build infrastructure first. Before writing any kernel code, the assistant established the CMake build system, the KDTR binary format, and the numpy reference implementations. This ensured that every kernel could be tested immediately upon completion, with no waiting for infrastructure.

Validate against references, not against the full model. By using synthetic test data with known expected outputs, the assistant could validate kernels on a local RTX 5070 Ti without needing the 1-trillion-parameter model or 8 GPUs. This dramatically accelerated the development cycle.

Commit milestones cleanly. Each validated kernel was committed with a descriptive message enumerating test counts, numerical tolerances, and documented interfaces (see [msg 11873], [msg 11899]). This created recoverable checkpoints and clean boundaries between phases [49][75].

Document interfaces before integration. The assistant wrote docs/kernels.md specifying the exact contracts (input tensors, output tensors, memory layouts) for each kernel before attempting to wire them together. This prevented integration surprises.

Accept pragmatism over perfection. The assistant accepted token-level rather than bitwise parity for the full engine, recognizing that floating-point non-determinism in the MoE reduction makes bitwise matching impossible for a 1-trillion-parameter model. The bar was set at "greedy output matches token-for-token," which was achieved and validated.

Conclusion

The construction of the kdtree-engine/ repository represents a remarkable engineering achievement: a complete native C/C++/CUDA DDTree inference engine for Kimi K2.6, built from scratch across a single intensive coding session. The engine delivers validated custom CUDA kernels for GPU tree building, tree-verify MLA attention, and greedy tree acceptance, along with an MVP FP32 transformer engine that matches a numpy golden reference token-for-token.

The work demonstrates that building efficient inference for trillion-parameter models requires not just knowledge of the models themselves, but deep understanding of GPU architecture, memory hierarchies, numerical analysis, and the subtle interactions between software components. The disciplined methodology—infrastructure first, reference implementations, incremental validation, clean milestones—provides a template for similar GPU kernel development projects.

The engine is now ready for the next phase: integration with INT4 Marlin kernels, multi-GPU tensor parallelism, and production deployment on the 8× RTX PRO 6000 Blackwell system. The foundation is solid, the kernels are correct, and the path forward is clear.## The Broader Context: Three Parallel Workstreams

The kernel engine development described above was not the only activity in this segment. In parallel, the assistant addressed two other critical workstreams that together paint a complete picture of the production inference stack.

Diagnosing the Throughput Regression

While building the native engine, the user reported a severe throughput regression on the live SGLang service: approximately 32 tokens/second at 5k context length, compared to a 138 t/s baseline. The assistant built a context-sweep diagnostic tool (bench_context_decode.py) and ran controlled measurements to isolate the cause.

The diagnosis revealed two compounding issues rather than a single bug. First, the undertrained drafter model's acceptance rate drops from ~5 tokens/step on predictable text to ~2.9 tokens/step on hard reasoning/analysis text. Second, KV cache fragmentation under page_size=1 with non-contiguous DDTree commits degrades step time over sustained generation. The measured step time at the user's context length combined with the low acceptance rate mathematically produces exactly the observed 32 t/s—confirming the system was behaving as expected for this workload, not suffering from a regression [2].

This diagnostic work was essential for setting realistic expectations. The native engine's performance gains would come from eliminating Python overhead and GPU-side tree building, but the fundamental throughput ceiling imposed by the drafter's acceptance rate would remain until a better-trained drafter became available.

Extending Context Length to 200k

The user also requested extending the maximum context length from 32k to 200k tokens. The assistant verified that the Kimi K2.6 model supports up to 262,144 tokens via YaRN scaling, then analyzed the memory feasibility. MLA's per-token KV overhead is exceptionally low (~8.6KB per token), meaning the current KV cache pool of ~101k tokens uses only a fraction of the available GPU memory (~10GB free per GPU). Concluding that the increase was memory-feasible, the assistant modified the systemd service file to set --context-length 200000 and issued a service restart [2].

This workstream demonstrated the practical realities of production inference: the model architecture (MLA) that makes the engine efficient also makes long-context deployment straightforward, as the KV cache footprint is dramatically smaller than traditional multi-head attention.

Key Architectural Insights

Several architectural insights emerged from this work that are worth highlighting:

The "verify is free" property. The central throughput insight is that the MoE GEMM is the throughput lever, not the attention kernel. The tree-verify attention kernel is primarily a correctness and robustness problem, not a throughput one—MLA's latent KV is tiny (576 elements per token), so attention is cheap. The real cost is in the INT4 MoE decode, which is HBM-bandwidth-bound [21].

Bounded shapes simplify the engine. With a fixed maximum of ~16 streams and a fixed budget, all tensor shapes are bounded and predictable. This means CUDA graph capture becomes a controlled tool rather than the liability it was in SGLang's variable-shape environment [21].

The three-kernel decomposition is the right abstraction boundary. The build → verify-attn → accept pipeline maps cleanly onto GPU execution, with each kernel producing outputs that the next consumes directly in device memory without host round-trips. This decomposition was validated by the on-device composition test [60].

Numerical determinism requires deliberate engineering. The assistant's meticulous attention to operation order, accumulation precision, and tie-breaking ensured that the GPU kernels matched the numpy references bit-exactly. This was not an accident—it was the result of careful analysis of every potential divergence point [36][37].

The Validation Philosophy

The validation strategy employed throughout this project is worth examining as a model for similar efforts. Rather than testing end-to-end on the full model (which would require 8 GPUs and hours per run), the assistant created a multi-layered validation approach:

  1. Kernel unit tests against synthetic data with known expected outputs
  2. Composition tests chaining multiple kernels together on-device
  3. Full engine validation against a numpy golden reference with model configurations
  4. Token-level parity as the acceptance criterion for the full engine This approach catches errors at the level where they are cheapest to fix. A bug in the tree builder is caught by the unit tests before it can corrupt the verify-attn kernel's inputs. A mismatch in the interface contract between kernels is caught by the composition test before it can cause a silent failure in the full engine. The final token-level parity check against the golden reference provides end-to-end confidence without requiring bitwise exactness (which is impossible for a 1-trillion-parameter MoE model due to floating-point non-determinism in the reduction).

Looking Forward

The kdtree-engine/ repository as committed at the end of Phase 1 represents a solid foundation for production deployment. The three custom CUDA kernels are validated and ready, the MVP FP32 engine demonstrates the architecture works end-to-end, and the diagnostic tools provide visibility into production behavior.

The next steps are clear: integrate the INT4 Marlin kernels for actual model weights, implement multi-GPU tensor parallelism via NCCL, optimize the verify-attn kernel for production shapes (using FlashMLA for long prefixes), and deploy on the CT200 hardware. The engine's architecture—bounded shapes, GPU-side tree construction, fused operations—is designed to exploit the Blackwell GPU's capabilities while avoiding the fragility that plagued the SGLang integration.

The work also highlights the importance of the drafter model's quality. Even the most optimized inference engine is limited by the drafter's acceptance rate, and the diagnostic work in this segment confirmed that the undertrained drafter is the dominant factor in the throughput regression. Improving the drafter—whether through further training, larger block sizes, or architectural changes—would provide more throughput improvement than any inference optimization alone.## References

[1] "The Architecture of Knowledge: How a Comprehensive Status Message Bridges Two Phases of DDTree Inference Optimization" — Analysis of message 11825, the comprehensive status document that synthesized weeks of work and laid out the blueprint for the C/CUDA stack.

[2] "The Pivot to Native: A User's Directive to Build a C/CUDA DDTree Inference Engine" — Analysis of message 11826, the user's three-line directive that pivoted the project from analysis to engine construction.

[21] "The Architecture of a Thousand-Parameter Decision: Designing a Native DDTree Inference Engine for Kimi K2.6" — Analysis of message 11845, the comprehensive engineering plan with quantitative performance model and phased execution strategy.

[32] "Bridging Languages: The KDTR Binary Container and the Architecture of Cross-Language Kernel Validation" — Analysis of the KDTR binary format design.

[33] "The Precision of Thought: Building a Golden Reference for GPU Tree Search" — Analysis of the numpy reference implementation methodology.

[36] "The Architecture of a CUDA Kernel: Designing a GPU Best-First Tree Builder for Speculative Decoding" — Analysis of the tree builder kernel design, including shared memory layout and numerical determinism strategy.

[39] "The Glue That Binds: Building the CMake Infrastructure for a Custom CUDA DDTree Engine" — Analysis of the CMake build system configuration.

[40] "The Build System That Binds: A CMakeLists.txt as the Keystone of a CUDA DDTree Engine" — Analysis of the CMakeLists.txt structure.

[42] "The First Compilation: When Custom CUDA Meets Reality" — Analysis of the first build attempt and the __syncwarp fix.

[47] "A Moment of Validation: The 11 Tests That Proved a CUDA Kernel Correct" — Analysis of the tree builder validation results.

[49] "The Milestone Commit: Cementing a Native DDTree Inference Engine" — Analysis of the tree builder milestone commit.

[50] "Designing the Tree-Verify MLA Attention Kernel: A Window into GPU Kernel Engineering" — Analysis of the verify-attn kernel design.

[52] "Designing the Verify-Attention Kernel: A Deep Dive into CUDA MLA Attention for DDTree Speculative Decoding" — Detailed analysis of the verify-attn kernel architecture.

[57] "The Moment of Truth: Validating a Custom CUDA Attention Kernel for Speculative Decoding" — Analysis of verify-attn kernel validation.

[58] "Six Green Lights: Validating a Custom CUDA MLA Attention Kernel for Speculative Decoding" — Analysis of the six verify-attn test configurations.

[60] "The Pivot Point: How a Single Decision Completed the DDTree CUDA Kernel Trio" — Analysis of the decision to implement the tree-accept kernel.

[61] "The Art of Synthetic Test Data: Building a Greedy Tree-Accept Kernel for Speculative Decoding" — Analysis of the accept kernel test data generation strategy.

[70] "The Moment of Truth: Validating Three Custom CUDA Kernels for Speculative Decoding on Blackwell GPUs" — Analysis of the full Phase 1 test suite (23 tests).

[71] "The Meta-Validation Pause: Ensuring Tests Mean Something in CUDA Kernel Development" — Analysis of the test coverage verification.

[75] "The Milestone Commit: Cementing Phase 1 of a Custom CUDA Speculative Decoding Engine" — Analysis of the Phase 1 completion commit.