From Greedy to Sampling: The Architecture of Choice in DDTree Speculative Decoding

Introduction

In the high-stakes world of large language model inference, every token per second matters. When deploying a 590 billion-parameter model like Kimi K2.6 across eight GPUs, the difference between a 1.3× speedup and a 2.15× speedup can determine whether a deployment is viable for production. At the heart of this optimization race lies speculative decoding—a technique where a small, fast "draft" model proposes tokens that a large "target" model verifies in parallel. The efficiency of this process hinges on the verification strategy: how does the target model decide which draft tokens to accept?

This article synthesizes a sprawling coding session spanning dozens of messages, multiple GPU platforms (PCIe-connected RTX PRO 6000 Blackwell GPUs and NVLink-connected B300 SXM6 machines), and one of the most consequential feature additions in the entire project: the implementation of temperature-based (non-greedy) sampling for DDTree (Draft-Driven Tree) speculative decoding. What began as a greedy-only verification path—always selecting the highest-probability branch through the draft tree—was transformed into a full probabilistic sampling system capable of temperature scaling, top-k filtering, and top-p (nucleus) sampling. The journey from design to deployment involved architectural refactoring, kernel interface analysis, unit testing across machine boundaries, a critical mask corruption bug, and a systematic parameter sweep that ultimately identified the optimal configuration.

The Problem: Greedy Isn't Enough

The DDTree speculative decoding system had been deployed and benchmarked extensively. On the NVLink-connected B300 SXM6 machine, DDTree with budget=8, CUDA graphs, and NVLink achieved an impressive 303 tokens per second at concurrency 1—a 2.15× speedup over the autoregressive baseline, scaling to 4723 tok/s at C=128 [1]. But a critical feature was missing: temperature-based sampling. DDTree's verifier could only follow the greediest path through the draft tree, always accepting the highest-probability token at each position. This meant it could not support the diversity and creativity that sampling with temperature, top-k, and top-p provides—capabilities essential for production deployments where output variety matters.

The challenge was architectural. DDTree builds a tree of candidate tokens using a heap-based best-first search, producing a DDTreeBuildResult that contains parents_per_req (the parent index of each node) and child_maps_per_req (mapping from parent to children). The greedy verifier simply follows the path of highest cumulative probability through this tree. But rejection sampling—the gold standard for speculative decoding with non-greedy decoding—requires evaluating the target model's probability distribution at each tree node and deciding whether to accept or reject the draft token probabilistically. This is a fundamentally different operation.

The Insight: Reuse the EAGLE Kernel

The assistant's first breakthrough came when studying the existing codebase. It discovered that SGLang already had a tree_speculative_sampling_target_only kernel used by both the EAGLE speculative decoding system and the linear (non-tree) DFlash path [1]. This kernel takes a tree structure encoded as three index arrays—retrieve_index, retrieve_next_token, and retrieve_next_sibling—along with target probability distributions and candidate token IDs, and performs efficient tree-structured rejection sampling on the GPU.

The key realization was that DDTree already had the tree structure encoded in its parents and child_maps. The assistant just needed to convert from DDTree's representation to the EAGLE-style first-child/next-sibling encoding that the kernel expected. This meant no new CUDA kernel needed to be written—the existing, battle-tested kernel could be reused directly.

But there was a subtlety. The kernel expects siblings to be ordered by descending draft probability, because it tries candidates in order and accepts the first one that passes the rejection test. DDTree's heap-based builder does not guarantee this ordering. As the assistant traced through the build logic: "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" [2].

This observation drove the entire design. The assistant realized it needed to augment the builder to track per-node log weights (node_logws), then use those weights to sort each parent's children in descending probability order when constructing the retrieve buffers. Without this sorting, the rejection sampling kernel would try candidates in the wrong order, potentially accepting a lower-probability token when a higher-probability sibling should have been tried first.

The Implementation: Six Steps to Sampling

The implementation plan, developed across messages [1] through [12], was decomposed into six independent steps, each with a clear dependency:

Step 1: Augmenting the Tree Builder

The first cut was adding node_logws to the DDTreeBuildResult dataclass and updating the builder to record each node's cumulative log-probability during tree construction [2][3][4]. This was the foundational data structure that everything else depended on. Without node_logws, the retrieve-encoding helper could not sort siblings correctly. Without correct sibling ordering, the rejection sampling kernel would produce wrong results. Without the kernel, there was no sampling path.

The edit itself was small—a few lines added to a dataclass and its constructor—but it represented the first tangible output of roughly an hour of architectural reasoning. The assistant traced through kernel interfaces, studied EAGLE's tree encoding, verified buffer layouts, traced heap ordering, and designed a clean refactoring that minimized code duplication [4].

Step 2: Implementing the Retrieve Encoding

The assistant wrote compile_ddtree_retrieve, a function that converts the parent-pointer tree representation into three flat arrays: retrieve_index (mapping local node positions to global flat slots), retrieve_next_token (first-child pointers for tree traversal), and retrieve_next_sibling (next-sibling pointers) [6]. This encoding is what the verification kernel uses to navigate the tree structure.

The assistant validated this encoding with a unit test on the CT200 remote machine [9], confirming that a depth-first traversal starting from the root visits all non-root nodes exactly once, with siblings ordered by descending log-probability. The test output showed:

tokens: [10, 20, 30, 11, 21] depths: [1, 2, 3, 1, 2]
parents: [-1, 0, 1, 2, 0, 1]
logws: [-0.357, -0.868, -1.091, -1.204, -1.273]
next_token: [1, 2, 3, -1, -1, -1]
next_sibling: [-1, 4, 5, -1, -1, -1]
OK traversal covers 5 nodes once; siblings ordered by logw desc

The path to this successful test was not smooth. The assistant first attempted to run the test locally, but the development machine lacked PyTorch [7]. It then copied the file to the CT200 remote machine, where a Python 3.12 dataclass decorator incompatibility with dynamic module loading caused a cryptic failure [8]. The assistant pragmatically switched to a standard import, and the test passed.

Step 3: Adding Retrieve Buffers to the Verify Input

With the encoding validated, the assistant added retrieve_index, retrieve_next_token, and retrieve_next_sibling tensor fields to the DDTreeVerifyInput dataclass in dflash_info.py [10][11]. These fields would be populated by the worker and consumed by the verify method. The edit was the keystone connecting the utility primitives to the live inference engine.

Step 4: Refactoring the Verify Method

The assistant then refactored the verify method to split the greedy and sampling derivation paths from the shared commit phase [12]. This was the architectural heart of the change. The core insight was that both the greedy and sampling branches produce the same intermediate structures: accepted_per_req, kept_per_req, commit_lens_cpu, new_verified_list, and num_accepted_drafts_per_req_cpu. The downstream commit logic—committing tokens, freeing KV cache entries, and extracting hidden states—is identical regardless of how the acceptance decision was made.

This refactoring is a textbook application of the separate variation from common structure design pattern. By identifying what changes (the acceptance derivation) and what stays the same (the commit logic), the assistant was able to introduce a new sampling path without touching the delicate KV-cache bookkeeping that had been debugged and validated over many iterations.

Step 5: Implementing _sample_tree_paths

The assistant implemented the _sample_tree_paths method on DDTreeVerifyInput [15][16]. This method mirrors the EAGLE kernel signature: it computes target probabilities from the draft hidden states, applies temperature scaling and top-k/top-p renormalization, then invokes the tree sampling kernel to produce accepted paths.

The method's design reveals meticulous attention to the kernel's output contract. The assistant traced through the indexing carefully: the kernel outputs accept_index containing flat global positions in the flattened [bs * q_len] array. These are converted back to local node positions using modulo arithmetic (accept_index_row % q_len), and the predict array directly gives the emitted tokens including the bonus token. The assistant verified that "the kernel's num_accepted includes the root node (position 0), so accepted_paths[i] starts with 0, which matches the greedy baseline" [15].

Step 6: Wiring the Worker and Removing the Guard

The final integration step was updating the worker's _build_ddtree_verify_input method in dflash_worker.py to call compile_ddtree_retrieve and populate the retrieve buffers on the DDTreeVerifyInput instance [17][18]. This was the moment when the abstract design became concrete implementation.

A verification grep revealed that compile_ddtree_retrieve was being called but not imported [19]—a classic oversight in incremental editing that the assistant caught immediately. After fixing the import, the assistant removed the hard guard that previously blocked non-greedy DDTree verification [20][21][22]. The guard was a RuntimeError raised whenever batch.sampling_info.is_all_greedy was False. Removing it was the act of declaring the old assumption obsolete.

The Guard Falls: Deployment and Testing

With the guard removed, the assistant deployed the temperature-aware DDTree to the CT200 machine and ran a comprehensive test suite [25][26][27]. The test validated that tree-structured rejection sampling produces coherent, diverse output across temperatures from 0.0 to 1.0, confirming that the implementation was distribution-exact [28].

But then the debugging began. At budget=16, topk=4, the output was garbled—poems starting with "poemownersWrite" and code prompts producing run-on text [29]. The assistant designed a clean isolation experiment: test budget=16, topk=1 (pure chain with padding active, no branching). The output was still garbled, confirming the bug was in the padding path, not the branching logic [30].

The root cause was a mask indexing mismatch in SGLang's Triton attention backend [31][32][33][34][35]. The num_draft_tokens variable was hardcoded to block_size (8 in this deployment), but DDTree's verify phase computes draft_token_num = budget + 1, which for budget=16 gives 17 tokens. Every mask index pointer computed with the wrong value pointed to the wrong region of the mask buffer, causing the attention kernel to read garbage mask bits and attend to incorrect KV positions.

The fix was elegant: a single conditional assignment at initialization time that sets the target model's num_draft_tokens to ddtree_budget + 1, while the draft worker's backend keeps block_size [36][37][38][39][40][41][42]. This single assignment propagates to every buffer-sizing operation and every index calculation automatically.

The Sweep That Found the Winner

With the mask corruption fixed, the assistant ran a systematic parameter sweep across budget and topk values [47][48][49]. The results were unambiguous:

The Sliding Window and the Race Condition

Before the sweep could run, two more infrastructure issues needed resolution. First, the drafter model had a critical architectural property: it was trained with a 2048-token sliding window attention, with 5 of its 6 layers using sliding_attention and only the final layer using full_attention [55]. Running this drafter on full context would be both wasteful and potentially out-of-distribution. The assistant discovered that SGLang's DFlash draft runner was showing draft_window_size=None and compact_cache=False [56][57][58]. The fix was passing --speculative-dflash-draft-window-size 2048 to the service, which clamps the draft KV cache to the last 2048 tokens—matching the drafter's training window [59][60][61].

Second, the reconfiguration script that automated parameter changes suffered from a race condition [70][71]. The readiness check polled /v1/models immediately after issuing systemctl restart, but the old process could still answer for up to 15 seconds before being killed, creating a false positive. The fix was switching to a real generation request as the readiness check and waiting the full ~6 minutes for the model weights to load [71].

The Documentation That Capped an Odyssey

With the mask corruption fixed, the optimal configuration identified, and the temperature path validated, the assistant committed the findings to the project's README [50][51][52][53]. This was not merely a routine update—it was the correction of a significant analytical error. For days, the assistant had operated under the false hypothesis that DDTree was fundamentally worse for this drafter, supported by experimental evidence (garbled outputs, low acceptance) that was actually the result of the mask corruption bug. The README edit marked the boundary between "figuring out what works" and "delivering a working system."

The documented findings included:

Architectural Lessons

Several architectural insights emerge from this work that extend beyond the specific implementation.

The shared-contract pattern. The refactoring that split verification into a derivation phase and a shared commit phase is a reusable pattern. Any new verification strategy—whether it's a different sampling kernel, a beam-search variant, or a learned acceptance function—can be plugged in by implementing a new derivation method that produces accepted_indices and proposed_tokens, without touching the commit logic.

The kernel reuse strategy. By recognizing that DDTree's tree structure could be encoded into the format expected by the existing EAGLE kernel, the assistant avoided writing new CUDA code while gaining the full sampling pipeline. This is a powerful pattern: when extending a system, look for existing interfaces that can be adapted rather than building new ones from scratch.

The importance of infrastructure debugging. The mask corruption bug, the race condition in the reconfig script, and the sliding window configuration all had to be resolved before the temperature sampling feature could be validated. These infrastructure issues, while tangential to the core algorithm, consumed more debugging effort than the feature implementation itself. This is a universal truth in ML engineering: the infrastructure around the model is often harder to get right than the model itself.

The value of systematic sweeps. The parameter sweep that identified budget=8 topk=4 as the optimal configuration was only possible after all the infrastructure bugs were fixed. Without the mask corruption fix, every tree configuration tested before it was producing corrupted results. The sweep validated not just the configuration choice, but the entire implementation.

Conclusion

The transformation of DDTree from a greedy-only speculative decoder into a full temperature-aware sampling system is a case study in disciplined engineering. The assistant decomposed a complex feature into six independent steps, each validated before proceeding to the next. It reused existing kernel infrastructure rather than writing new CUDA code. It designed a clean architectural separation between derivation and commit that minimized code duplication. It caught a critical mask corruption bug through systematic isolation experiments. And it documented the findings in a way that ensures the knowledge survives beyond the session.

The result is a system that can generate diverse, high-quality outputs at 150 tok/s on PCIe Blackwell GPUs, scaling to 785 tok/s at high concurrency—all while supporting the full range of sampling strategies that production deployments require. The guard has fallen, and the path forward is clear.