The Mask That Made Trees Work: Debugging DDTree Speculative Decoding on Blackwell GPUs

Introduction

In the high-stakes world of large language model inference, every token per second counts. When deploying a 1-trillion-parameter Mixture-of-Experts model like Kimi K2.6 across eight RTX PRO 6000 Blackwell GPUs, the difference between a working optimization and a silently corrupted one can be the difference between a 1.5× speedup and a regression that gets written off as "trees don't help." Message 11689 is the capstone of a multi-day debugging odyssey—a concise summary that distills three critical bug fixes, their performance implications, and the corrected understanding of what makes tree-based speculative decoding actually work. This message is not merely a status update; it is a postmortem that rewrites the project's earlier conclusions and sets a clear path forward.

The assistant wrote this message after completing a systematic sweep of DDTree (Draft Tree) configurations, verifying temperature sampling, and committing all fixes to git. It represents the moment when a long sequence of "trees hurt" observations was finally explained by a single, subtle bug: the attention mask was corrupted whenever the tree budget didn't match the block size. The message's tone is one of quiet triumph—the bugs are fixed, the results are in, and the earlier pessimism about tree-based speculative decoding has been overturned.

Context: The Long Road to Working Trees

To understand why this message matters, one must appreciate the journey that preceded it. The project had deployed Kimi K2.6 with DFlash speculative decoding—a technique where a smaller "drafter" model proposes tokens that the larger target model verifies in parallel. The standard DFlash approach uses a linear chain of draft tokens. DDTree extends this by proposing a tree of candidates at each step, increasing the chance that the target model accepts a longer prefix. In theory, trees should outperform chains. In practice, earlier benchmarks showed trees performing worse—sometimes producing garbled output, sometimes achieving lower throughput than the simpler chain.

The assistant had spent days debugging CUDA graph crashes, buffer sizing mismatches, and temperature sampling issues. Each fix brought incremental progress, but the core question remained: why were trees not delivering their theoretical advantage? The answer, as message 11689 reveals, was a mask corruption bug that caused tree nodes to attend to padded garbage in the KV cache. Every benchmark run with budget ≠ 7 was silently producing corrupted attention patterns. The assistant had been comparing a broken tree against a working chain—an unfair fight that the tree could never win.

The Subject Message

The message that caps this debugging effort reads as follows:

All mission items complete. Summary of what was accomplished:

## TP8 + DDTree + cuda graphs + temperature — done

### Three bugs fixed (all committed + code backed up in git)

1. cuda graph crash (cuda_graph_runner.py): the DDTree target-verify graph was sized for block_size (8) but verify runs over budget+1 tree nodes (e.g. 33) → token-count mismatch on replay. Split the sizing: draft worker keeps block_size, DDTree target uses budget+1.

2. DDTree custom-mask corruption (triton_backend.py) — the big one: num_draft_tokens was hardcoded to block_size, so whenever budget+1 ≠ block_size the attention mask offsets misaligned and real tree nodes attended padded/garbage KV → garbled, non-exact output. It only worked by coincidence at budget=7 (8==8). This is why wide trees seemed to "hurt" earlier — they were silently corrupted, not just slow. After the fix, trees genuinely help: commit_len ~4-5 vs the chain's ~2.

3. DDTree temperature (ddtree_utils.py + dflash_info.py + dflash_worker.py): implemented tree-structured rejection sampling by converting the best-first tree into EAGLE's first-child/next-sibling encoding and reusing the existing tree_speculative_sampling_target_only sgl_kernel. Validated coherent + diverse output at temp 0/0.6/1.0. (Linear DFlash temperature already worked via the same kernel.)

### Key results (K2.6 1T MoE, 8× RTX PRO 6000 Blackwell, PCIe, C=1) | Config | C=1 tok/s | |---|---| | TP8 autoregressive, no cuda graph | 26 | | TP8 autoregressive + cuda graph | 98 (3.8× from graphs) | | DDTree budget=8 topk=4 + cuda graph | 113-150 (accept ~4) |

- cuda graphs are the dominant single-stream win and now work with speculation. - DDTree wins at low concurrency; autoregressive EP4 wins raw aggregate at C≥64. - The "tmp-long" drafter is undertrained (accept ~4) — your better drafter should push acceptance higher, and now large budgets can pay off since the mask is correct.

The DDTree service on CT200 is running the best config (budget=8 topk=4, cuda graphs on). Two correctness notes: temp=0 greedy is exact but shows tiny run-to-run divergence late in long outputs — that's inherent 1T-MoE FP non-determinism (expert reduction order), not the spec path.

This compact message packs an extraordinary amount of information: three bug descriptions, a performance comparison, deployment status, and a forward-looking assessment—all in under 300 words.

The Three Bugs: A Taxonomy of Failure

The message enumerates three bugs, each with distinct root causes and symptoms. Understanding them reveals the complexity of deploying speculative decoding on modern hardware.

Bug 1: CUDA Graph Buffer Sizing

The first bug was a mismatch between the CUDA graph's expected token count and the actual number of tokens processed during verification. CUDA graphs are a performance optimization that captures a sequence of GPU kernel launches into a reusable "graph" that can be replayed with minimal CPU overhead. The assistant had previously achieved a 3.8× speedup from CUDA graphs alone (from 26 tok/s to 98 tok/s for the autoregressive baseline). However, the DDTree target-verify graph was sized for block_size (8 tokens), while the tree verification step actually runs over budget+1 tree nodes. When the budget was 16, the graph expected 8 tokens but received 17—a mismatch that caused illegal memory accesses and crashes.

The fix was straightforward once identified: split the sizing logic so the draft worker continues using block_size while the DDTree target verification path uses budget+1. This is a classic example of a latent assumption—that the number of tokens in the verify step equals the block size—that held true only for the linear DFlash case and broke under the tree formulation.

Bug 2: The Mask Corruption — The Big One

This is the bug that rewrites the project's history. The DDTree custom attention mask had num_draft_tokens hardcoded to block_size. In SGLang's triton attention backend, this mask controls which tokens in the sequence can attend to which other tokens. When budget+1 ≠ block_size, the mask offsets misaligned, causing real tree nodes to attend to padded or garbage KV cache entries. The result was garbled, non-exact output that looked like a failed optimization rather than a bug.

The critical insight is that this bug was invisible during normal operation. The tree would produce output that looked like plausible text but was semantically corrupted—the model was attending to random memory instead of the actual draft tokens. The assistant had earlier observed that "trees hurt" and concluded that the verification overhead outweighed the acceptance gains. In reality, the trees were producing garbage tokens that the target model was then forced to reject, artificially depressing the acceptance rate.

The mask only worked correctly by coincidence when budget=7 (since 7+1=8, matching block_size). This is why the assistant's earlier benchmarks with budget=8 showed reasonable behavior—it happened to hit the one configuration where the corruption was absent. Any wider tree was silently broken.

This bug is a cautionary tale about silent correctness failures in ML inference. Unlike a crash or an obvious error message, a corrupted attention mask produces output that seems reasonable but is statistically degraded. Without a reference implementation to compare against, such bugs can persist indefinitely, leading to incorrect conclusions about algorithm performance.

Bug 3: DDTree Temperature Sampling

The third bug was an implementation gap rather than a corruption: DDTree lacked temperature sampling support. The assistant implemented tree-structured rejection sampling by converting the best-first tree representation into EAGLE's first-child/next-sibling encoding and reusing the existing tree_speculative_sampling_target_only kernel from sgl_kernel. This approach avoids writing a custom sampling kernel by leveraging an existing, verified implementation with a different tree representation.

The implementation required changes across three files: ddtree_utils.py for the tree conversion, dflash_info.py for metadata propagation, and dflash_worker.py for the sampling orchestration. The assistant validated that temperature 0.0, 0.6, and 1.0 all produced coherent output with appropriate diversity, confirming that the sampling path preserved the target distribution correctly.## The Performance Landscape: What the Numbers Actually Say

The message presents a clean three-row comparison that tells a compelling story:

| Config | C=1 tok/s | |---|---| | TP8 autoregressive, no cuda graph | 26 | | TP8 autoregressive + cuda graph | 98 (3.8× from graphs) | | DDTree budget=8 topk=4 + cuda graph | 113-150 (accept ~4) |

The dominant finding is that CUDA graphs provide the largest single-stream win—a 3.8× improvement over the uncaptured baseline. DDTree adds another 15-53% on top of that, depending on the prompt. The range (113-150 tok/s) reflects prompt sensitivity: coding prompts with predictable structure accept more readily than open-ended explanations.

The assistant's sweep across budget values (8, 16, 24, 32, 48) revealed a clear pattern: smaller budgets win at single-stream concurrency. This makes intuitive sense: verification cost scales linearly with budget+1 tokens, while the acceptance depth is capped by the drafter's block_size-1=7. A budget of 8 creates a tree wide enough to capture useful diversity without paying for excessive verification. The budget=8 configuration achieved 150.2 tok/s, while budget=48 dropped to 101.4 tok/s—barely above the non-speculative baseline.

The concurrency sweep (C=1 through C=128) with budget=8 showed DDTree winning at low concurrency but losing to autoregressive EP4 at high concurrency (785.9 tok/s vs 1449 tok/s at C=128). This is a fundamental tradeoff: speculative decoding shines when the GPU is underutilized (low concurrency, latency-sensitive applications) but its overhead becomes a liability when batching already saturates compute. The assistant correctly notes that the "tmp-long" drafter is undertrained—a better drafter would push acceptance higher, making larger budgets viable and potentially shifting the crossover point.

Assumptions and Their Consequences

The message reveals several assumptions that shaped the debugging process:

Assumption 1: Trees were working correctly. The assistant had been operating under the belief that the DDTree implementation was functionally correct, just slow. The mask corruption bug meant that every benchmark with budget ≠ 7 was measuring a broken system. The assumption that "no crash means correct" is dangerously seductive in ML systems, where silent data corruption can produce plausible-looking outputs.

Assumption 2: The mask parameter was a constant. The hardcoded num_draft_tokens = block_size was likely a legacy from the linear DFlash implementation, where the number of draft tokens always equals the block size. When DDTree introduced a variable number of tree nodes, this constant became a bug. The assumption that a parameter doesn't need to change when the algorithm changes is a common source of subtle regressions.

Assumption 3: Budget=7 was a valid baseline. The assistant's earlier benchmarks used budget=7, which happened to work because 7+1=8=block_size. This created a false baseline: trees at budget=7 appeared to work, while trees at budget=16 appeared to fail. The natural conclusion was that wider trees are worse—when in fact, wider trees were simply broken.

Input Knowledge Required

To fully understand this message, one needs familiarity with several concepts:

Speculative decoding: A technique where a small "draft" model proposes tokens that a large "target" model verifies in parallel, achieving speedup when the draft is accepted.
DDTree (Draft Tree): An extension of speculative decoding where the draft model proposes a tree of candidate tokens rather than a linear chain, increasing the probability of a long accepted prefix.
CUDA graphs: A CUDA API that captures a sequence of kernel launches into a reusable graph, eliminating CPU launch overhead.
Block size: In DFlash, the number of draft tokens generated per step (typically 8). The drafter's architecture constrains the maximum tree depth to block_size - 1.
Budget and top-k: DDTree parameters controlling tree width. Budget limits total tree nodes; top-k limits candidates per position.
Attention masks: In transformer models, masks control which tokens can attend to which other tokens. Custom masks are needed for tree-structured speculation.
TP8/EP4: Tensor parallelism across 8 GPUs vs expert parallelism with 4 expert groups. These are parallelism strategies for distributing model computation.
1T MoE: A 1-trillion-parameter Mixture-of-Experts model, where only a subset of parameters is activated per token.
FP non-determinism: Floating-point operations in CUDA are not deterministic across runs due to instruction scheduling and reduction order, causing tiny output variations even at temperature=0.

Output Knowledge Created

This message produces several lasting contributions:

A corrected understanding of DDTree performance on Blackwell GPUs: Trees genuinely help when implemented correctly, achieving commit lengths of 4-5 tokens vs the chain's ~2. The earlier "trees hurt" conclusion was wrong.
A validated configuration: Budget=8, topk=4 with CUDA graphs enabled is the optimal single-stream configuration for K2.6 on PCIe Blackwell. This is immediately actionable for deployment.
A taxonomy of DDTree bugs: Three distinct failure modes (graph sizing, mask corruption, missing temperature) with their fixes documented. This serves as a reference for anyone implementing tree-based speculation.
A performance baseline: The three-row comparison table provides a clear hierarchy of optimizations. CUDA graphs are the dominant win (3.8×), with DDTree adding incremental value on top.
A diagnostic framework: The message implicitly establishes that correct speculative decoding must produce exact greedy output (modulo FP non-determinism). Any deviation from exactness indicates a bug, not a feature of the algorithm.

The Thinking Process

The assistant's reasoning in this message is notable for its clarity and structure. Rather than presenting a narrative of the debugging process, it organizes findings into a bug taxonomy, a results summary, and a forward-looking assessment. This structure reflects a mind that has moved from reactive debugging to proactive documentation.

The opening line—"All mission items complete"—signals closure. The assistant has a checklist (visible in the preceding messages as a todo list) and is ticking off the final items. The three bugs are enumerated with root cause, symptom, and fix, demonstrating a systematic approach to debugging.

The parenthetical emphasis on "the big one" for the mask corruption bug reveals the assistant's own surprise at the finding. The realization that "trees seemed to 'hurt' earlier—they were silently corrupted, not just slow" is a moment of retrospective insight that reframes all prior work. The assistant is not just reporting a fix; it is acknowledging that its own earlier conclusions were wrong.

The performance comparison table is carefully constructed to isolate the contribution of each optimization. By showing the no-graph baseline (26 tok/s), the graph-only baseline (98 tok/s), and the graph+DDTree result (113-150 tok/s), the assistant makes it clear that CUDA graphs are the dominant factor. This prevents readers from over-attributing the speedup to DDTree.

The closing note about FP non-determinism is a mark of intellectual honesty. The assistant anticipates a potential concern—"is the output truly exact?"—and addresses it preemptively. This shows an understanding that correctness in ML inference is not binary but exists on a spectrum, and that communicating the limits of determinism builds trust.

Broader Implications

This message has implications beyond the specific K2.6 + DDTree deployment. The mask corruption bug is a cautionary tale about the dangers of hardcoded constants in ML inference code. When an algorithm changes (linear chain → tree), every constant must be re-examined. The fact that the bug only manifested at budget ≠ 7, and that budget=7 was the most common test configuration, meant the bug could have persisted indefinitely.

The finding that CUDA graphs provide a 3.8× speedup—far more than the 1.5× from DDTree—suggests that for many deployments, the highest-impact optimization is not algorithmic but infrastructural. Eliminating Python overhead through graph capture dwarfs the gains from smarter speculation. This is a reminder that in systems engineering, the biggest wins often come from removing bottlenecks rather than adding cleverness.

The message also demonstrates the value of systematic benchmarking. The budget sweep (6 configurations) and concurrency sweep (5 concurrency levels) provided a comprehensive picture of the tradeoff space. Without this data, the assistant might have concluded that budget=48 was "better" because it has more tree nodes, when in fact it was worse for single-stream throughput.

Conclusion

Message 11689 is a masterclass in technical communication. It distills days of debugging into three clear bug descriptions, presents results with appropriate caveats, and reframes the project's understanding of tree-based speculative decoding. The mask corruption bug—the "big one"—is a particularly valuable lesson: silent correctness failures in ML inference can lead to incorrect conclusions about algorithm performance, and the only defense is rigorous validation against a reference implementation.

The message also reveals the assistant's growth from debugging to synthesis. The bugs are fixed, the results are documented, and the path forward is clear. The DDTree service on CT200 is running the best configuration, and the project can now move to the next phase: training a better drafter to push acceptance higher, where the now-correct tree implementation can deliver its full potential.