The Temperature Test: Validating Tree-Structured Rejection Sampling for DDTree Speculative Decoding

Introduction

In the high-stakes world of large language model inference, every token per second counts. When deploying a 590 GB model like Kimi K2.6 across eight GPUs, the difference between a working speculative decoding pipeline and a broken one can mean hours of wasted compute, corrupted outputs, or silently incorrect behavior. Message [msg 11663] represents a pivotal moment in a long engineering session: the first live validation of a newly implemented temperature (non-greedy) sampling path for the DDTree (Diverse Depth Tree) speculative decoding algorithm within SGLang. This message is not merely a test—it is the culmination of a multi-hour effort to extend DDTree from a greedy-only verification mode to a full stochastic sampling engine, reusing the same tree-structured rejection sampling kernel that powers EAGLE and linear DFlash.

The Context: Why This Message Exists

To understand the significance of this message, one must appreciate the architecture it validates. DDTree is a best-first tree-building algorithm for speculative decoding: rather than generating a single linear sequence of draft tokens, it constructs a tree of candidates ordered by cumulative log-probability, then verifies them in parallel through a single target-model forward pass. The original implementation, deployed in earlier segments of this conversation, was restricted to greedy decoding (temperature=0). This limitation was enforced by a hard guard in the worker code that raised a RuntimeError if any request had a non-greedy sampling configuration.

The restriction was pragmatic. Greedy verification is simpler: the tree path with the highest cumulative probability is deterministically accepted. Non-greedy (temperature > 0) verification requires stochastic rejection sampling across the tree structure, which is fundamentally more complex. The assistant had to:

  1. Convert the DDTree's best-first tree into EAGLE's first-child/next-sibling encoding via the compile_ddtree_retrieve function
  2. Compute target probabilities with temperature scaling, top-k filtering, and top-p renormalization
  3. Invoke the existing tree_speculative_sampling_target_only sgl_kernel (the same kernel EAGLE and linear DFlash use)
  4. Map the kernel's output (accepted indices in a flat global space) back to local tree positions for hidden-state selection and KV cache management
  5. Ensure the commit path—the code that actually writes accepted tokens into the generation buffer—works identically for both greedy and sampled paths The assistant's reasoning in [msg 11651] reveals the careful indexing logic required: "accept_index returns global flat indices—I can convert those back to local positions by subtracting i*q_len." This arithmetic is the bridge between the kernel's flat view of the batch and the tree's local node structure. A single off-by-one error here would corrupt the entire generation.

The Deployment Pipeline

Before the test in [msg 11663] could run, the assistant executed a carefully orchestrated deployment sequence. In [msg 11660], the changes were committed to git with a descriptive message: "Add DDTree temperature support via tree-structured rejection sampling." The commit message itself reveals the design philosophy: "Reuses the existing tree_speculative_sampling_target_only sgl_kernel (same one EAGLE and linear DFlash use)." This is a key architectural decision—rather than writing a custom sampling kernel for DDTree's tree structure, the assistant converted DDTree's representation into the format the existing kernel already expects.

In [msg 11661], the assistant deployed the modified Python files to CT200's site-packages, backed up the originals, and updated the systemd service configuration to use --speculative-ddtree-budget 16 --speculative-ddtree-topk-cap 4. This configuration change is significant: earlier tests used budget=7 and topk-cap=1, which produced very shallow trees. The new settings create substantially larger, deeper trees (budget=16 allows up to 16 nodes, topk-cap=4 means each parent considers up to 4 children), which should increase the acceptance length but also stress the sampling logic more severely.

The service restart in [msg 11662] took 600 seconds—a full ten minutes—because the Kimi K2.6 model is 590 GB and must be loaded into GPU memory and have its Triton kernels compiled. The assistant's polling loop checked every 15 seconds, alternating between systemd status and HTTP health checks. This wait is itself a testament to the scale of the deployment: a 590 GB model on 8 GPUs requires substantial weight-loading and JIT compilation before it can serve a single request.

The Test Design: What the Message Actually Does

The message [msg 11663] executes a Python script via a bash heredoc that performs three categories of tests:

1. Greedy Regression Check (temp=0.0)

The first test sends a prompt requesting a Python quicksort implementation with temperature=0.0. This is the regression check: does the greedy path still work after the code changes? The assistant's reasoning in earlier messages established that the verify method now branches: greedy uses follow_verified_tree (the original deterministic path), while non-greedy computes target probabilities and runs the sampling kernel. The greedy path must remain untouched.

The result: 61.2 tok/s | 150 tok. The output text begins "The user wants python quicksort algorithm with comments..." — this is the model's characteristic thinking-style response, where it first restates the user's request before generating code. The throughput of 61.2 tok/s is reasonable for a 590 GB model with speculative decoding on PCIe-connected GPUs.

2. Temperature Sampling Tests (temp=0.6, temp=1.0)

The second category tests actual stochastic sampling. At temp=0.6, the throughput jumps to 66.9 tok/s—slightly higher than greedy. This is counterintuitive: one might expect sampling to be slower due to the additional target-probability computation. However, the assistant's design reuses the same kernel and the target probability computation is relatively cheap (softmax with temperature scaling, top-k/top-p filtering). The throughput variation may also reflect differences in acceptance length: stochastic sampling can sometimes accept more tokens than the greedy path because it explores tree branches that the greedy path would reject.

At temp=1.0, throughput drops to 38.3 tok/s. This is a significant degradation. The most likely explanation is that high temperature flattens the probability distribution, causing the rejection sampling kernel to reject more draft tokens. When the distribution is uniform (or near-uniform), the draft model's predictions are less likely to match the target model's samples, so fewer tokens are accepted per verification step, and the effective speculation depth decreases.

3. Diversity Check

The third test sends the same prompt ("Name a random animal and a color.") twice at temp=1.0 and checks whether the outputs differ. This is a critical correctness test: if temperature>0 produces identical outputs for the same prompt, something is wrong—either the random seed is fixed, the sampling kernel isn't actually sampling, or the tree path selection is deterministic despite the temperature setting.

The outputs differ dramatically:

Assumptions and Their Validity

The message rests on several key assumptions:

Assumption 1: The tree_speculative_sampling_target_only kernel works identically for DDTree trees as for EAGLE trees. The assistant explicitly designed the retrieve encoding to match EAGLE's first-child/next-sibling format, but the kernel may have implicit assumptions about tree shape (e.g., maximum depth, branching factor) that DDTree violates. DDTree trees are built by a best-first algorithm that can produce arbitrary structures, while EAGLE trees follow a fixed pattern. The successful test results suggest this assumption holds, but the throughput degradation at temp=1.0 warrants further investigation.

Assumption 2: Temperature=0.0 is equivalent to greedy. The assistant uses temperature=0.0 as the regression test, assuming it activates the same deterministic path as the original greedy-only implementation. This is correct in practice: temperature=0.0 collapses the softmax to a delta distribution on the most likely token, and the sampling kernel should produce the same result as the greedy path. The 61.2 tok/s result is consistent with earlier greedy-only benchmarks.

Assumption 3: The service is healthy after the 600-second wait. The assistant's polling loop checked /v1/models and confirmed the endpoint was responding. However, the readiness check may be insufficient: the model might respond to the models endpoint but still be compiling Triton kernels in the background, leading to JIT compilation delays on the first request. The test script doesn't include a warmup request, so the first API call includes any one-time compilation overhead.

Assumption 4: The diversity check is a sufficient test of stochastic correctness. Two differing outputs confirm that the sampling is non-deterministic, but they don't confirm that the sampling follows the correct distribution. A proper statistical test would require many samples and a goodness-of-fit test against the expected target distribution. The assistant is implicitly trusting the kernel implementation and the correctness of the target probability computation.

Input Knowledge Required

To fully understand this message, one needs:

  1. Speculative decoding fundamentals: The concept of draft models generating candidate tokens that a target model verifies in parallel. The distinction between greedy verification (accept the best path) and stochastic verification (sample from the target distribution).
  2. Tree-structured speculative decoding: How DDTree builds a best-first tree from draft log-probabilities, how EAGLE uses a fixed tree structure, and the first-child/next-sibling encoding that bridges the two representations.
  3. SGLang's architecture: The relationship between the speculative worker (dflash_worker.py), the verify input dataclass (DDTreeVerifyInput), and the sampling kernel (tree_speculative_sampling_target_only). The service deployment model (systemd, site-packages, model loading).
  4. The Kimi K2.6 model characteristics: It is a 590 GB MoE (Mixture of Experts) model that uses a thinking/reasoning format, producing outputs that begin with "The user wants..." restatements. It is deployed on 8× RTX PRO 6000 GPUs with PCIe interconnect.
  5. CUDA graph limitations: The earlier chunk analysis mentions that CUDA graphs are incompatible with the DFlash verify path on certain architectures, which affects throughput comparisons.

Output Knowledge Created

This message produces several valuable pieces of knowledge:

  1. Empirical validation: DDTree with temperature sampling works correctly on a real 590 GB model. The greedy path is preserved (61.2 tok/s), and stochastic sampling produces diverse outputs.
  2. Throughput characteristics: Temperature has a non-monotonic effect on throughput. Low temperature (0.6) slightly outperforms greedy (66.9 vs 61.2 tok/s), while high temperature (1.0) significantly degrades performance (38.3 tok/s). This suggests an optimal temperature sweet spot for throughput.
  3. Acceptance behavior: The diversity check confirms that the rejection sampling kernel is functioning and producing genuinely different outputs. The garbled nature of the short-context responses hints at potential issues with the thinking model's output format under sampling.
  4. System health: The service can handle the budget=16 topk=4 configuration without crashing, despite the increased tree complexity. No errors or tracebacks appear in the journalctl output.
  5. A baseline for future optimization: The 61.2 tok/s greedy result and the 66.9 tok/s temp=0.6 result provide reference points for future improvements to the DDTree pipeline.

The Thinking Process Visible in the Message

The message itself is a test script, but its structure reveals the assistant's thought process:

First, establish the baseline. The greedy test (temp=0.0) runs first because it's the regression check. If this fails, everything else is moot—the code changes broke the existing functionality.

Then, probe the new functionality. The temperature tests (0.6, 1.0) test the new sampling path at two points: a moderate temperature that should produce reasonable outputs, and a high temperature that stresses the system.

Finally, verify correctness. The diversity check is the most important test: it confirms that the stochastic path is actually stochastic. The assistant prints both outputs and explicitly checks outs[0]!=outs[1], making the result unambiguous.

The follow-up command checks server status and acceptance metrics, showing the assistant is thinking about both correctness and performance. The journalctl grep for "accept len" would reveal the average number of tokens accepted per verification step, which is the key metric for speculative decoding efficiency.

Conclusion

Message [msg 11663] is a masterclass in practical ML engineering validation. It is not a unit test, not a formal benchmark, but a live integration test on a production-scale model that costs hundreds of thousands of dollars to train. The assistant's careful progression—regression check, parameter sweep, diversity verification—reflects a deep understanding of what can go wrong when modifying a complex inference pipeline. The results are encouraging: DDTree with temperature sampling works, achieves competitive throughput, and produces genuinely diverse outputs. But the throughput degradation at high temperature and the garbled short-context responses point to areas for further investigation. This message marks the moment when a theoretical extension to DDTree became a working feature, validated against the unforgiving reality of a 590 GB model serving real requests.