The Third Kernel: Completing the DDTree Speculative Decode Pipeline on GPU

[assistant] [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/tree_accept.cu Wrote file successfully.

On its surface, message [msg 11887] appears to be the most mundane of events in a coding session: a single file write, acknowledged with a success message. But this message is the capstone of an intensive engineering push to build a complete native C/C++/CUDA DDTree (Drafting-Decoding Tree) inference engine for the Kimi K2.6 language model. The file tree_accept.cu is the implementation of the third and final custom CUDA kernel in the speculative decoding pipeline — the greedy tree-accept kernel. Its creation marks the moment when the entire "greedy DDTree step trio" — build, verify-attn, accept — became fully implemented on GPU, eliminating the need for host round-trips between speculative decoding phases.

The Motivation: Why This Kernel Exists

The tree-accept kernel solves a specific but critical problem in speculative decoding with tree-structured drafts. In the DDTree algorithm, the drafter proposes a tree of possible next tokens, the target model verifies which paths are valid through attention, and then the system must accept the verified tokens — walking the tree from the root, checking which children match the target's predictions, and committing the longest verified path to the output.

This accept step is the final bridge between the raw verification results and the actual output tokens the user sees. Without it, the GPU has computed all the necessary information but cannot produce the final answer. In the previous architecture, this walk was done on the CPU — the verified attention results would be transferred back to host memory, a Python or C++ routine would traverse the tree structure, and the accepted tokens would be sent back to the GPU for the next iteration. This round-trip introduced latency and synchronization overhead that limited throughput.

The assistant's reasoning in the preceding messages makes this motivation explicit. After validating the tree-builder and verify-attention kernels, the assistant considered what else could be accomplished in the session and concluded: "The accept step is the final piece of the greedy DDTree pipeline — it takes the target's per-node predictions and walks the verified tree to determine which tokens were accepted. This is small, testable against the Python reference, and completes the trio of build → verify-attn → accept entirely in CUDA. That's the highest-value next step."

Input Knowledge Required

To understand why this message matters, one needs knowledge of several domains. First, speculative decoding: the technique where a small "drafter" model proposes multiple candidate tokens, and a large "target" model verifies them in parallel, achieving speedup by processing multiple tokens in a single forward pass. Second, tree-structured speculative decoding (DDTree) : instead of a linear sequence of draft tokens, the drafter proposes a tree of possibilities, with branching at each position to cover more alternatives. Third, CUDA kernel design: the ability to write GPU code that processes each request independently in a single thread, using pointer-chasing through sibling-linked data structures. Fourth, the Kimi K2.6 model architecture: a DeepSeekV3-style model with Multi-head Latent Attention (MLA) and Mixture-of-Experts (MoE), which determines the specific tensor shapes and memory layouts the kernel must handle.

The assistant also needed knowledge of the KDTR binary container format (the custom format for sharing test data between Python and C++), the CMake build system configured for CUDA 13 with sm_120 architecture targeting Blackwell GPUs, and the numpy reference implementation of the follow_verified_tree function that was extended in the preceding message [msg 11884].

How Decisions Were Made

The kernel design choices are visible in the assistant's reasoning trace from the preceding message [msg 11886]. Several key decisions stand out:

Thread-per-request parallelism: The assistant chose to launch one thread per request, with each thread processing its assigned request sequentially. This is a natural fit for the tree-walk algorithm, which is inherently sequential — you cannot know which child to visit until you've checked the current node's prediction. The grid-block configuration is straightforward: enough threads to cover the batch size.

Sibling-linked tree traversal: The tree structure is encoded using next_token and next_sibling pointers, a compact representation that allows the kernel to search for a matching child by walking the sibling chain. The assistant verified that this is equivalent to a dictionary-based approach: "since siblings have distinct tokens (generated from top-k selections at each depth), scanning through next_sibling pointers will find the matching token regardless of traversal order, just like the dict lookup would."

Output buffer design: The kernel produces two output arrays: accepted_path (the token IDs along the verified path, excluding the root) and proposed (the full set of tokens to commit, including the bonus token from where traversal stopped). Both are initialized with -1 sentinel values. The accept_len tracks how many nodes were visited including the root, and commit_len equals the number of proposed tokens.

Safety bounds: The assistant noted the need to pass q_len through the budget parameter for proper array access, ensuring the kernel stays within allocated memory bounds.

The Thinking Process

The assistant's reasoning reveals a careful design process that balances correctness, simplicity, and performance. The kernel walk logic is described step by step: "starting at the root, I predict the next token at each node, search for a matching child in the sibling list, and stop if no match is found."

There is also a deliberate verification step against the reference implementation: "Let me verify this against the reference implementation — the accepted path starts at node 0, proposed contains the token IDs of each advanced node plus the final bonus, and since all accepted nodes stay within the target prediction bounds, the indexing is safe."

This cross-checking between the CUDA kernel design and the Python reference is a recurring pattern throughout the session. The assistant consistently validates that the GPU implementation matches the numpy reference before declaring it correct, and the tree-accept kernel is no exception — the reference follow_verified_tree function was extended in message [msg 11884] specifically to serve as the ground truth for testing.

Output Knowledge Created

The writing of tree_accept.cu creates several forms of knowledge. At the most concrete level, it is a CUDA source file containing the implementation of the greedy tree-accept kernel, ready to be compiled and linked into the kdtree-engine project. This file, combined with the header tree_accept.cuh written in [msg 11886], forms a complete, testable kernel.

More broadly, this message completes the third and final kernel of the DDTree speculative decode pipeline, establishing the architectural invariant that the entire build→verify→accept loop can run on GPU without host synchronization. This architectural decision has performance implications: eliminating CPU round-trips reduces latency and allows the pipeline to sustain higher throughput, especially under batch processing where synchronization overhead would compound.

The message also implicitly validates the design decisions made in the preceding messages — the sibling-linked tree encoding, the thread-per-request parallelism model, and the output buffer layout. By committing the implementation to disk, the assistant signals that these decisions have converged into a concrete artifact that can be built, tested, and integrated.

Assumptions and Potential Pitfalls

Several assumptions underpin this kernel. The most fundamental is that the sibling-linked tree encoding is sufficient for correct traversal — specifically, that scanning through next_sibling pointers will always find a matching child if one exists. This holds because the tree construction guarantees distinct tokens among siblings (they come from top-k selections at each depth), but it would break if duplicate tokens were allowed within a node's children.

Another assumption is that the target model's predictions arrive in a specific layout (target_predict array) that aligns with the tree's node ordering. If the verification step reorders nodes or produces predictions in a different sequence, the accept kernel would silently produce wrong results.

The kernel also assumes that the tree structure fits within the allocated buffer sizes, with q_len (the budget parameter) providing the correct bounds. The safety fix applied in a later chunk — adding a bound check after a crash from cyclic random test data — confirms that this assumption required explicit enforcement.

The Broader Context

Message [msg 11887] sits at the end of a sustained engineering effort documented across dozens of preceding messages. The assistant had already built and validated the GPU best-first tree builder kernel and the tree-verify MLA attention kernel, each with its own numpy reference implementation, test data generators, KDTR bundles, and CMake integration. The tree-accept kernel was the final piece needed to close the loop.

The file path — /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/tree_accept.cu — tells its own story. The glm-kimi-sm120-rtx6000bw prefix identifies the target hardware (Blackwell RTX PRO 6000 GPUs with sm_120 architecture), the kdtree-engine is the custom inference engine being built from scratch, and src/kernels/ houses the CUDA kernels. This is not a modification to an existing codebase like SGLang or vLLM; it is ground-up infrastructure for a novel speculative decoding algorithm.

The message also represents a deliberate scoping decision. The assistant could have moved on to integration work on the CT200 hardware, but instead chose to complete the kernel trio while still in a self-contained development environment. This decision maximized the value of the current session by producing testable, validated artifacts that would be immediately useful when deployment resumed.