The Greedy Tree-Accept Kernel: Completing the DDTree Speculative Decode Trio

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When serving a 200-billion-parameter model like Kimi K2.6 across eight NVIDIA RTX PRO 6000 Blackwell GPUs, the difference between acceptable throughput and unusable latency often comes down to the quality of custom CUDA kernels. Message [msg 11886] captures a pivotal moment in the development of a native C/C++/CUDA speculative decoding engine: the creation of the greedy tree-accept kernel, the third and final piece of a custom GPU-side DDTree (Draft-and-Verify Tree) pipeline. This message, though brief in its surface appearance, represents the culmination of a carefully orchestrated effort to replace SGLang's CPU-bound tree verification logic with a fully GPU-accelerated alternative, and it embodies the kind of low-level systems thinking that separates production-grade inference infrastructure from research prototypes.

The Context: Building a Native DDTree Engine from Scratch

To understand why this message was written, one must first understand the broader project. The assistant had been working for many sessions on deploying the Kimi K2.6 model with speculative decoding using the DDTree algorithm — a technique where a small "drafter" model proposes multiple candidate token sequences arranged in a tree structure, and the large "target" model verifies them in parallel. The key insight of DDTree is that by batching the verification of multiple candidate paths through a single forward pass, the system can achieve significantly higher throughput than traditional autoregressive decoding.

The existing implementation in SGLang performed the tree-building step on the CPU using a Python heapq, which created a serial bottleneck — each request's tree had to be constructed on the host before being dispatched to the GPU. The assistant's vision was to replace this with a fully GPU-native pipeline consisting of three custom kernels: a GPU best-first tree builder (replacing the CPU heapq), a tree-verify MLA-absorb attention kernel (handling the masked attention computation for the tree's verification pass), and a greedy tree-accept kernel (determining which tokens from the verified tree to commit to the output). Messages [msg 11873] through [msg 11885] had already established the build infrastructure, created the numpy reference implementations, and validated the first two kernels. Message [msg 11886] completes the trio.

The Reasoning: Walking the Verified Tree on GPU

The assistant's reasoning in this message reveals a careful mental simulation of the kernel's control flow. The tree-accept kernel's job is deceptively simple: given the target model's predictions for each node in the draft tree, determine the longest prefix of the tree that the target agrees with. This is the "accept" step in speculative decoding — the system accepts tokens from the drafter's proposal up to the point where the target model disagrees, then uses the target's own prediction (the "bonus token") for the next position.

The assistant walks through the algorithm step by step. Each thread processes one request independently, starting at the root of that request's draft tree. At each node, it checks whether the target model's predicted token matches any of that node's children. The children are encoded using a next-token/next-sibling linked-list structure — a compact representation where each node stores its token ID, a pointer to its first child (next_token), and a pointer to the next sibling (next_sibling). This is the same encoding used in the tree builder kernel, ensuring interface compatibility.

The search for a matching child involves traversing the sibling chain: starting from the first child, the kernel follows next_sibling pointers until it finds a node whose token ID matches the target's prediction, or exhausts the list. If a match is found, traversal advances to that child and repeats. If no match is found, traversal stops, and the target's prediction at the stopping point becomes the bonus token.

The assistant carefully verifies the buffer layout against the reference implementation. The accepted_path buffer records the node indices along the accepted path (starting from the root), while the proposed buffer collects the token IDs of the accepted tokens (excluding the root) plus the bonus token. The accept_len counts how many nodes were visited including the root, and commit_len equals the number of proposed tokens. The assistant notes the need to pass q_len through the budget parameter for proper array access — a detail that reveals awareness of the memory layout constraints when operating on padded tensors.

Design Decisions and Their Rationale

Several implicit design decisions are visible in the assistant's reasoning. First, the choice of one thread per request rather than one block per request reflects the lightweight nature of the accept operation — it's essentially a sequential tree walk with minimal computation per node, so launching a full block per request would waste GPU resources. A grid-stride loop with one thread per request is appropriate when the work per thread is small and the number of requests (batch size) is modest.

Second, the decision to initialize output buffers with -1 padding is a standard CUDA convention for variable-length outputs. Since different requests may accept different numbers of tokens, the output arrays need sentinel values to distinguish valid data from padding. The value -1 is a natural choice since token IDs are non-negative.

Third, the assistant's verification 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" — demonstrates a disciplined approach to correctness. Rather than writing the kernel and testing blindly, the assistant mentally traces the execution path and checks for out-of-bounds access before writing a single line of code.

Assumptions and Their Implications

The message rests on several assumptions. The most fundamental is that the next-token/next-sibling encoding is sufficient for the accept walk. This encoding assumes that sibling nodes have distinct token IDs — an assumption that holds for DDTree because the tree builder selects top-k tokens at each depth, and the top-k set contains unique tokens by construction. If this assumption were violated (e.g., if the drafter proposed the same token via two different paths), the sibling chain traversal could match the wrong child, potentially accepting tokens that weren't actually verified.

The assistant also assumes that the target model's predictions are available in a contiguous array indexed by node position, and that the tree structure (node_token_ids, next_token, next_sibling) is laid out in the same padded format used by the tree builder. This is a reasonable assumption given that the kernels are designed as a matched set, but it creates a coupling between the kernels — any change to the tree builder's output format would require a corresponding change to the accept kernel.

Another assumption is that the accept operation is purely sequential and doesn't benefit from parallelism within a request. This is true for the greedy accept algorithm, but it's worth noting that the algorithm is inherently serial — you can't know whether to accept node N+1 until you've verified that node N was accepted. There's no opportunity for intra-request parallelism, which justifies the one-thread-per-request design.

Input Knowledge Required

To fully understand this message, the reader needs familiarity with several domains. First, speculative decoding: the concept of using a small drafter model to propose candidate tokens that a large target model then verifies in parallel. Second, the DDTree algorithm specifically: how draft tokens are organized into a tree structure where each path from root to leaf represents a different candidate continuation. Third, CUDA programming concepts: thread-block-grid hierarchy, shared memory, and the performance characteristics of different launch configurations. Fourth, the next-token/next-sibling encoding for trees in GPU memory — a compact representation that avoids storing explicit child arrays. Fifth, the broader context of the kdtree-engine project and its phased roadmap, which the assistant had established in earlier messages.

Output Knowledge Created

This message produces a CUDA kernel header file (tree_accept.cuh) that implements the greedy tree-accept algorithm on GPU. The kernel is the third and final piece of the GPU-side DDTree pipeline, completing the build→verify→accept cycle entirely on device without host round-trips. Together with the tree builder and verify-attention kernels, it enables the native engine to perform speculative decoding without CPU involvement in the critical path.

The kernel also serves as a reference implementation for the accept algorithm, documenting the exact semantics of the operation. Future developers can read this kernel to understand what "greedy tree acceptance" means in the context of DDTree — a concept that, while well-known in the speculative decoding literature, has specific implementation details when realized on GPU hardware.

The Thinking Process: A Window into Systems Programming

What makes this message particularly valuable is the window it provides into the assistant's thinking process. The reasoning section reads like a programmer talking through a design before writing code — a form of rubber-duck debugging where the act of verbalizing the algorithm reveals edge cases and correctness conditions.

The assistant starts with a high-level description ("each thread processes one request by traversing the path"), then progressively refines it with implementation details ("starting at the root, I predict the next token at each node, search for a matching child in the sibling list"). This top-down decomposition is characteristic of experienced systems programmers who think in terms of data flow and control flow before touching the keyboard.

The verification step — "Let me verify this against the reference implementation" — is particularly telling. The assistant doesn't just write the kernel and hope it works; it mentally simulates the execution and checks for consistency with the numpy reference. This is the kind of discipline that separates reliable code from buggy code, especially in CUDA where debugging tools are limited and runtime errors can manifest as subtle numerical discrepancies rather than crashes.

The final detail about passing q_len through the budget parameter reveals the assistant's awareness of the kernel's integration context. The tree builder and accept kernel share a common interface, and the budget parameter (which controls the maximum tree size) also encodes the query length for array bounds checking. This kind of parameter reuse is common in GPU kernels where register pressure and launch parameter limits constrain the number of explicit arguments.

Conclusion

Message [msg 11886] may appear at first glance to be a simple note about writing a CUDA kernel, but it represents much more. It is the completion of a carefully designed three-kernel pipeline that moves the entire DDTree speculative decoding loop onto the GPU, eliminating the CPU bottleneck that limited throughput in the SGLang implementation. The assistant's reasoning reveals a disciplined approach to systems programming: mental simulation before coding, verification against reference implementations, and careful attention to interface compatibility between kernels. For anyone interested in high-performance LLM inference, this message offers a concrete example of how speculative decoding algorithms are translated from mathematical descriptions into efficient GPU code — and the kind of thinking required to get it right.