The DDTree CUDA Graph Fix: Unlocking Speculative Decoding at 109 tok/s on PCIe Blackwell

From Crash to Breakthrough

The journey from a crashing CUDA graph to a working 109 tok/s speculative decoding pipeline on PCIe-connected Blackwell GPUs is a story of methodical debugging, architectural understanding, and the surprising discovery that sometimes the simplest configuration wins. This chunk of the session represents a critical inflection point: the assistant moved from deploying speculative decoding on a new NVLink B300 machine back to the PCIe PRO6000 platform, fixing a fundamental CUDA graph sizing bug in the DDTree verification path, and ultimately characterizing the performance landscape of tree-based speculative decoding with a weak drafter.

The DDTree Architecture: Confirming Multiple-Path Evaluation

Before diving into debugging, the assistant first answered a critical architectural question from the user: does DDTree evaluate a single candidate path or multiple paths in parallel? The answer, confirmed by reading the source code in dflash_info.py and ddtree_utils.py, was definitive: DDTree genuinely evaluates all tree paths in a single target forward pass.

The flow works as follows. First, a single draft forward produces block_size=8 hidden states. From each depth position, the _topk_logprobs_from_vocab_parallel_head function computes top-k candidate tokens (k=topk_cap). Then build_ddtree_tree_from_topk runs a best-first heap expansion: it pops the highest cumulative-logprob node, adds it to the tree, pushes its best child and next sibling, and repeats until budget nodes are selected. This builds a tree of budget+1 nodes where high-probability prefixes branch wider.

The critical insight is that compile_ddtree_visibility builds a [q_len, q_len] attention mask where each node attends only to its ancestors plus the shared prefix. Then one target verify forward runs all budget+1 nodes with this custom tree-attention mask. Every path is evaluated simultaneously — the target produces an argmax prediction at every node in parallel. Finally, follow_verified_tree walks from the root following the target's argmax at each node, taking the matching child until no child matches, and returns the longest accepted chain plus one bonus token.

This means the verify cost is budget+1 tokens through the target model (as a prefill-style batch), not budget separate forwards. And crucially, context is loaded once and shared across all tree nodes — the mask sets allow[:, :prefix_len] = True for every tree node, and the target verify reads the prefix KV from cache. No per-path re-reading occurs.

The CUDA Graph Crash: A Tale of Two Token Counts

With the architecture confirmed, the assistant attempted to deploy TP8+DDTree with CUDA graphs on the PCIe PRO6000 machine. The service started but crashed on the first request with a cryptic error: RuntimeError: The size of tensor a (8) must match the size of tensor b (33) at non-singleton dimension 0.

This error was the key to a fundamental bug in the CUDA graph runner. The DDTree speculative decoding pipeline has two different token counts that matter for graph sizing:

  1. Draft forward: Uses block_size=8 tokens per request (the anchor plus 7 mask tokens)
  2. Target verify forward: Uses budget+1=33 tokens per request (the root plus 32 tree nodes) The CUDA graph runner was using speculative_num_draft_tokens (which equals block_size=8) to size the target verify graph. But DDTree's verify forward actually processes budget+1=33 tokens. The graph was captured with buffers sized for 8 tokens per batch element, but replay tried to copy 33 tokens into those buffers — hence the size mismatch. The fix was elegant and surgical. In cuda_graph_runner.py, the assistant added a conditional: when the speculative algorithm is DDTree and this is the target worker (not the draft worker), set num_tokens_per_bs = speculative_ddtree_budget + 1 instead of speculative_num_draft_tokens. The draft worker continues using block_size. This single 12-line change resolved the crash.

The Verification: 109 tok/s and the Acceptance Mystery

After applying the fix and restarting the service, the assistant ran benchmarks. The results were immediately encouraging:

The Budget Scaling Puzzle

The assistant also tested larger budgets to see if they could improve acceptance:

| Configuration | C=1 Throughput | Accept Length | |---|---|---| | budget=7, topk=1 (chain) | 109.7 tok/s | ~2.0 | | budget=32, topk=8 | ~50 tok/s | ~1.5-2.0 |

The larger tree actually halved throughput. The reason was straightforward: the verify forward cost jumped from 8 tokens (budget=7) to 33 tokens (budget=32) — a 4× increase — without any improvement in acceptance length. The drafter simply wasn't good enough to benefit from tree exploration.

This led to a clear recommendation: use budget=7 topk=1 (effectively a single chain, equivalent to linear DFlash) with CUDA graphs as the current best configuration at 109 tok/s. The tree infrastructure is in place and working correctly, but it needs a better drafter to realize its potential.

Temperature Support: Already Working for Linear, Planned for DDTree

The user had asked about temperature support for speculative decoding. The assistant discovered that linear DFlash already supports temperature sampling — the compute_dflash_sampling_accept_len_and_bonus function in dflash_utils.py uses the tree_speculative_sampling_target_only kernel from sgl_kernel. A quick test confirmed that temperature=0.0, 0.6, and 1.0 all worked correctly with CUDA graphs.

For DDTree, temperature support required more work. The current implementation raises "DDTREE currently supports greedy verification only". The assistant designed a plan to wire DDTree's tree structure into the same tree_speculative_sampling_target_only kernel that EAGLE and linear DFlash use. The key insight was that DDTree already has the tree structure encoded in parents_per_req and child_maps_per_req — these just need to be converted into the (retrieve_index, retrieve_next_token, retrieve_next_sibling) format the kernel expects. The kernel would then handle tree-structured rejection sampling with temperature, top-k, and top-p filtering, all on GPU with minimal performance impact.

The Git Baseline: Securing the Working State

Throughout this work, the assistant maintained disciplined version control. A git snapshot was created with tag k26-baseline-working, capturing the exact state of all modified SGLang files:

Implications and the Path Forward

This chunk represents a significant step forward in the inference optimization campaign. The key findings are:

  1. CUDA graphs are the dominant optimization: The 3.8× speedup from enabling CUDA graphs on TP8 (26→98 tok/s) dwarfs the ~1.1× from speculative decoding with the current drafter. The graph infrastructure is now working correctly with DDTree.
  2. Drafter quality is the bottleneck: The "tmp-long" drafter's ~2.0 token acceptance at C=1 limits speculative decoding gains. A better drafter (which the user is training) could push acceptance to 4-6 tokens, making DDTree's tree exploration valuable.
  3. The tree infrastructure is sound: DDTree's multi-path evaluation, context sharing, and verify logic are all working correctly. The architecture is ready for a better drafter.
  4. Temperature support is achievable: The path to DDTree temperature sampling is clear — convert the tree structure to the kernel's expected format and reuse the existing tree_speculative_sampling_target_only kernel. The assistant's work here — from fixing the CUDA graph sizing bug to characterizing the acceptance landscape to planning temperature support — has laid the groundwork for the next phase: deploying on the NVLink B300 machine where the real throughput gains await.## The Broader Context: NVLink and the Findings Report While this chunk focused on the PCIe PRO6000 platform, the broader segment context reveals that the same DDTree stack was also deployed on an 8× B300 SXM6 NVLink machine with dramatically different results. On NVLink, DDTree with budget=8, NVLS, and CUDA graphs achieved 303 tok/s at C=1 — a 2.15× speedup over the 133 tok/s autoregressive baseline — and scaled to 4723 tok/s at C=128. NVLS provided a free ~5% boost, and NVLink made TP8 the clear winner over EP for this workload. The user's hypothesis that larger budgets would better utilize the B300's spare compute was confirmed: budget=16 in eager mode lifted acceptance to 5.3–6.4 tokens per step (vs b8's 4.48) and passed all coding tests. However, a critical sm_103-specific CUDA graph bug blocked realizing this gain — any budget > 8 caused illegal memory accesses during tree-verify graph capture, while eager mode lost the 3.8× graph speedup. The workload was confirmed to be HBM-bandwidth-bound (100% GPU util but only 360–460 W / 1100 W), meaning the compute is genuinely idle and waiting on memory. This validated the potential for optimized kernels and informed the roadmap for a custom C/C++/CUDA inference stack documented in the comprehensive DDTREE_FINDINGS_REPORT.md.

The Methodological Thread

Throughout this chunk, a clear methodological pattern emerged. The assistant alternated between reading source code to understand system behavior, formulating hypotheses about what should happen, running experiments to test those hypotheses, and fixing bugs when reality diverged from expectations. This cycle — read, hypothesize, experiment, fix — was executed with remarkable discipline.

When the CUDA graph crashed with a tensor size mismatch, the assistant didn't guess at the fix. It traced the error through the stack, identified that DDTree has two distinct token counts (block_size for draft, budget+1 for verify), and applied a precise 12-line fix. When the acceptance metrics seemed contradictory (depth=7 but commit=2), the assistant cross-referenced against linear DFlash at the same concurrency level, discovering that the earlier 3.5-4.0 accept_len figures were a high-concurrency artifact.

This methodological rigor is what separates effective optimization from trial-and-error hacking. Every conclusion was grounded in evidence from logs, benchmarks, and source code analysis.

Looking Ahead

The work in this chunk sets the stage for several clear next steps:

  1. A better drafter is the single highest-impact improvement. The current "tmp-long" drafter's ~2.0 token acceptance at C=1 limits all speculative decoding gains. A drafter that can reliably predict 4-6 tokens ahead would make DDTree's tree exploration genuinely valuable.
  2. DDTree temperature support is designed and ready to implement. The path through the tree_speculative_sampling_target_only kernel is clear, requiring only a conversion of DDTree's parent-child structure into the kernel's expected format.
  3. The custom inference stack roadmap in DDTREE_FINDINGS_REPORT.md targets eliminating Python overhead, fusing dequantization with MoE GEMM, and building custom tree-attention MLA verify kernels to bypass the sm_103 graph instability. The foundation laid here — working CUDA graphs with speculative decoding, a clear understanding of the acceptance landscape, and a disciplined methodological approach — will serve as the springboard for these next phases.## References [1] "The PP8 Benchmark That Killed a Hypothesis: When Pipeline Parallelism Fails on PCIe-Bound MoE Inference" — [chunk 64.1] [2] "The Pipeline That Wouldn't Fill: Diagnosing a Parallelism Mystery in SGLang" — [chunk 64.1] [3] "The Expert Parallelism Pivot: Reasoning Through Communication Topologies for MoE Inference" — [chunk 64.1] [4] "The EP8 Breakthrough: Why Expert Parallelism Transformed Kimi K2.6 Inference on PCIe Blackwell GPUs" — [chunk 64.1] [5] "The Parallelism Reckoning: How One Summary Table Settled the Architecture for Kimi K2.6 Inference" — [chunk 64.1] [6] "The Pivot: From Parallelism Benchmarks to Speculative Decoding" — [chunk 64.1] [7] "The First DFlash Benchmark: Validating Speculative Decoding for Kimi K2.6 on Blackwell" — [chunk 64.1] [8] "The DFlash Verdict: When Speculative Decoding Meets PCIe Reality" — [chunk 64.1] [9] "Reading the Source: How One Bash Command Uncovered DDTree's Greedy-Only Constraint" — [chunk 64.1] [10] "Peering into the DFlash Worker: How One Bash Command Uncovered the Hidden-State Bridge in Speculative Decoding" — [chunk 64.1] [11] "Tracing the CUDA Graph Crash: A Deep Dive into DFlash Speculative Decoding Debugging" — [chunk 64.1] [12] "Peering into the CUDA Graph Runner: A Microscope on Memory Bandwidth Optimization" — [chunk 64.1]