Bridging Trees and Temperatures: Engineering Rejection Sampling for DDTree Speculative Decoding

In the high-stakes world of large language model inference, every token per second counts. When deploying a 590-billion-parameter model like Kimi K2.6 across cutting-edge Blackwell GPUs, the difference between a 1.3× speedup and a 2.15× speedup can hinge on a single missing feature: temperature sampling support in tree-structured speculative decoding. Message [msg 11635] captures the precise moment when an AI assistant, having already deployed the K2.6+DDTree stack on an 8× B300 SXM6 NVLink machine and benchmarked it to impressive effect, confronts the architectural challenge of wiring temperature-controlled random sampling into a verification path that was designed exclusively for greedy (deterministic) decoding.

This message is a design document in miniature—a window into the assistant's reasoning as it maps the terrain between two existing code paths: DDTree's tree-structured draft verification and EAGLE's GPU-accelerated rejection sampling kernel. The assistant must answer a deceptively complex question: how do you take a tree built by one algorithm, with its own internal representation of parent-child relationships and cumulative log-probabilities, and feed it into a kernel designed for a different tree representation, without breaking either the correctness guarantees of rejection sampling or the performance characteristics that make speculative decoding worthwhile?

The Context: A Multi-Platform Deployment Odyssey

To understand the significance of this message, one must appreciate the journey that preceded it. The assistant had been working for hours—across multiple segments and dozens of messages—to deploy Kimi K2.6 with DFlash speculative decoding across two radically different hardware platforms. The first was an 8× RTX PRO 6000 system connected via PCIe, where bandwidth bottlenecks made expert parallelism (EP) the clear winner over tensor parallelism (TP). The second was an 8× B300 SXM6 system with NVLink interconnects, where the opposite was true: NVLink's high-bandwidth fabric made TP8 dramatically outperform EP, achieving 303 tok/s at concurrency 1—a 2.15× speedup over the autoregressive baseline.

The assistant had already resolved a cascade of infrastructure issues: CUDA toolkit conflicts, FlashInfer SM120 compatibility problems, Triton JIT compilation failures, missing Python.h headers, and a vision tower warmup that required unavailable flash_attn.cute kernels. It had benchmarked parallelism strategies (TP8, PP8, EP8, EP4), validated coding correctness (5/5 on the B300), and confirmed the user's hypothesis that larger DDTree budgets would better utilize the B300's spare compute—budget 16 in eager mode achieved 5.3–6.4 tokens accepted per step versus budget 8's 4.48.

But a critical sm_103-specific CUDA graph bug blocked realizing this gain: any budget greater than 8 caused illegal memory accesses or garbage output during tree-verify graph capture. Eager mode lost the 3.8× graph speedup, resulting in lower net throughput. The workload was confirmed to be HBM-bandwidth-bound (100% GPU utilization but only 360–460 W out of 1100 W), meaning the compute was genuinely idle, waiting on memory.

In the messages immediately preceding [msg 11635], the assistant had been investigating temperature support. Linear DFlash (the non-tree variant) already supported temperature sampling—the assistant tested it at temperatures 0.0, 0.6, and 1.0, confirming correct behavior with CUDA graphs and acceptance lengths around 2.0–2.7 tokens per step. But DDTree, the tree-structured variant, hard-blocked temperature. The assistant's mission was clear: implement tree-structured temperature sampling for DDTree.

The Core Architectural Challenge

The message opens with the assistant tracing how the kernel output maps to the existing DDTree commit logic. This is the central puzzle: the tree_speculative_sampling_target_only kernel—already used by both EAGLE and linear DFlash—returns accept_index and predict tensors that indicate which tree nodes were accepted and what tokens should be emitted. DDTree already has a commit loop that takes accepted nodes and handles KV cache management. The question is how to connect these pieces.

The assistant identifies two approaches. The first is to build the retrieve buffers (the index structures that tell the kernel how to navigate the tree) eagerly in the DDTree builder, storing them on the DDTreeVerifyInput object so the verify function can use them directly. The second is to build them lazily only when sampling is needed, avoiding unnecessary work for the greedy path. The assistant chooses the eager approach, reasoning that "building them always in the builder is simpler and the cost is minimal." This is a pragmatic engineering decision: it prioritizes code simplicity and correctness over marginal performance optimization, recognizing that the tree-building step is already CPU-bound and the additional index construction adds negligible overhead.

The Sibling Ordering Problem

Then comes a critical realization. The assistant is about to write a helper function that constructs EAGLE-style retrieve indices from the DDTree parent structure when it stops to reconsider how siblings are ordered.

DDTree builds its tree using a max-heap. Nodes are popped in order of cumulative log-probability (logw), so the most promising path through the tree is explored first. When a node is popped, it pushes its next sibling (rank+1) and its first child (rank 0) onto the heap. The assistant initially assumed that siblings of the same parent would naturally appear in descending log-probability order in the flattened node list, since each sibling's cumulative logw equals the parent's logw plus the edge log-probability for that rank, and log-probabilities are descending in rank.

But then the assistant catches itself: "Wait, I need to reconsider how siblings are ordered. The builder pops nodes in heap order (best-first by cumulative logprob), but that doesn't mean siblings of the same parent are popped in descending logprob order—they get interleaved with nodes from other parents."

This is a subtle but crucial insight. The heap doesn't process siblings of one parent contiguously. It processes the globally best node next, regardless of which parent it belongs to. So siblings of parent A might be interleaved with children of parent B, grandchildren of parent C, and so on. The flattened node list preserves heap order, not sibling order.

The rejection sampling kernel, however, needs to try siblings in descending draft-probability order. This is because rejection sampling works by attempting each candidate in order: for each child, it computes an acceptance probability min(1, q[token]/p[token]) where q is the target probability and p is the draft probability. If accepted, the walk stops. If rejected, it renormalizes the target distribution and tries the next child. Trying higher-probability candidates first maximizes the chance of early acceptance, which is both correct (the math requires ordered attempts) and efficient (fewer kernel launches).

The assistant's solution is to augment the builder to return node_logws alongside the existing outputs, then use those to sort the children of each parent. This guarantees that siblings appear in the retrieve buffers in descending draft-probability order, matching the kernel's expectations.

Draft Probabilities: Zero as a Signal

Another design decision concerns draft probabilities. The assistant notes that EAGLE passes draft_probs=zeros to the kernel, and plans to do the same for DDTree. This is worth unpacking: why would you pass zero probabilities for draft tokens?

The answer lies in how the kernel uses draft probabilities. In EAGLE's formulation, the draft model produces a probability distribution over the vocabulary at each position. The kernel uses these probabilities in the rejection sampling formula: accept with probability min(1, q[token] / p[token]). If p[token] is zero, the acceptance probability is effectively infinite (or undefined), which would break the algorithm.

But EAGLE passes zeros, and the kernel handles it. How? The key is that when draft_probs is all zeros, the kernel interprets this as a signal that the draft is deterministic—each position has exactly one candidate token with probability 1.0, and all others have probability 0. The rejection sampling then becomes a simple comparison: accept if the target probability of the draft token is high enough relative to some threshold, or equivalently, sample a bonus token from the residual distribution if the draft token is rejected.

For DDTree, this same convention applies. Even though DDTree has multiple candidates per position (the tree branches), the tree structure itself encodes the draft probabilities implicitly through the ordering of siblings and the cumulative log-probabilities. The kernel doesn't need explicit draft probabilities because the tree structure (encoded in retrieve_index, retrieve_next_token, and retrieve_next_sibling) tells it which tokens to try and in what order. The zeros signal "use the tree structure for ordering, not explicit probabilities."

This is a elegant reuse of an existing interface. The assistant doesn't need to modify the kernel or create a new one. It just needs to convert DDTree's representation into the format the kernel already understands.

The Four-Step Plan

By the end of the message, the assistant has synthesized a clear, four-step implementation plan:

Convert DDTree's parents/child_maps to EAGLE-style (retrieve_index, retrieve_next_token, retrieve_next_sibling), with siblings ordered by draft log-probability. This is the core data structure transformation.
Compute target_probs by applying softmax with temperature scaling and top-k/top-p filtering to the target model's logits, then pass draft_probs=zeros as EAGLE does.
Call the same tree_speculative_sampling_target_only kernel that EAGLE and linear DFlash already use, passing the converted tree structure and target probabilities.
Adapt the existing commit loop to use the kernel's accept_index and predict outputs, which indicate which tree nodes were accepted and what tokens to emit. This plan is notable for its economy of effort. The assistant is not writing a new kernel, not implementing CPU-side rejection sampling, not redesigning the commit machinery. It is finding the minimal set of changes—essentially a data structure conversion and a few lines of plumbing—that unlocks temperature sampling for DDTree.

What This Message Reveals About Engineering Practice

Message [msg 11635] is a case study in how experienced AI engineers think about system integration. Several patterns stand out:

Reading before writing. The assistant spends significant effort reading existing code: the kernel interface, EAGLE's tree setup, DDTree's builder logic, the chain verify buffers. It reads the _get_or_create_chain_verify_buffers function signature at the end of the message to verify dtype and shape compatibility. This is not procrastination; it's the essential work of understanding the existing interfaces before designing the integration.

Tracing execution paths. The assistant mentally simulates the tree-building algorithm to discover the sibling ordering issue. It traces through the heap operations: "when a node is popped, it pushes its next sibling (rank+1) and its first child (rank 0). So siblings of the same parent share the same depth and parent, but their relative order in the tree depends on when they were pushed and popped." This kind of algorithmic reasoning is critical for correctness.

Identifying the minimal change. The assistant repeatedly asks: what is the simplest thing that works? Eager buffer construction instead of lazy? Yes, because it's simpler. Zero draft probabilities instead of computing them? Yes, because the kernel already handles that convention. Reusing the existing kernel instead of writing a new one? Yes, because it's already correct and optimized.

Correcting assumptions. The initial assumption about sibling ordering was wrong, and the assistant catches it before writing code. This self-correction is a hallmark of careful engineering. The assistant doesn't just accept its first mental model; it stress-tests it by tracing through the algorithm.

The Broader Significance

This message sits at a pivotal moment in the larger deployment story. The assistant has already proven that DDTree works on NVLink hardware, achieving a 2.15× speedup. It has identified the CUDA graph bug that limits budget to 8. It has confirmed the workload is HBM-bandwidth-bound, meaning there's headroom for compute-intensive improvements. The temperature sampling implementation, once complete, will unlock a critical capability: the ability to serve the model with controlled randomness, which is essential for creative applications, diverse outputs, and proper integration into larger systems that expect sampling behavior.

The assistant's plan also sets the stage for the next phase of optimization. The DDTREE_FINDINGS_REPORT.md (written later in the segment) will outline a roadmap for a custom C/C++/CUDA inference stack that eliminates Python overhead, fuses dequantization with MoE GEMM, builds a custom tree-attention MLA verify kernel, and moves tree construction to the GPU. The temperature sampling work in this message is a necessary prerequisite—it ensures that the DDTree algorithm is fully featured before the optimization phase begins.

Conclusion

Message [msg 11635] captures the moment when abstract design crystallizes into concrete implementation. The assistant has identified the core data structure transformation, resolved the sibling ordering subtlety, chosen the eager-over-lazy tradeoff, and formulated a four-step plan that reuses existing infrastructure to the maximum extent possible. What makes this message compelling is not just the technical content—the tree representations, the kernel interfaces, the heap ordering analysis—but the engineering mindset it reveals: methodical, self-correcting, and relentlessly focused on the simplest path to a correct solution.

The temperature sampling for DDTree is not yet implemented at the end of this message. But the blueprint is complete. The assistant has done the hardest part of the work: understanding the problem deeply enough to see the solution clearly. The actual code—the data structure conversion, the kernel call, the commit loop adaptation—will follow in subsequent messages. But the intellectual heavy lifting is done here, in this single message of reasoning about trees, temperatures, and the art of connecting one system's output to another system's input.