From Tree-Based Speculative Decoding to Drafter Training: A Research Pivot in the DDTree Investigation

Introduction

The pursuit of faster large language model inference often leads researchers down unexpected paths. In the course of a single opencode coding session spanning multiple days, a team working with the Qwen3.6-27B model on an 8-GPU Blackwell node embarked on a journey that began with a deep investigation into DDTree—a tree-based speculative decoding algorithm—and ended with the construction of a high-throughput hidden state extraction pipeline for training a better draft model. This article chronicles that journey, examining how a critical architectural discovery forced a strategic pivot, how the team methodically investigated the DDTree codebase, and how they ultimately recognized that the drafter's quality—not the tree construction—was the primary bottleneck.

The Context: Speculative Decoding at Scale

The root session had been wrestling with speculative decoding integration for the Qwen3.6-27B model. Speculative decoding accelerates inference by using a small, fast "draft" model to propose candidate tokens, which a large "target" model then verifies in parallel. The team had already achieved solid results with MTP (Multi-Token Prediction) speculation, reaching 73.5 tokens per second on a single request. They had then pivoted to DFlash, a block-diffusion draft model that promised further speedups by generating probability distributions for multiple future positions simultaneously.

However, deploying DFlash had been painful. The acceptance rate was catastrophically low—around 1.1%—and the assistant had diagnosed three root causes: a layer-ID offset bug in vLLM's hidden state extraction, ignored sliding window attention (SWA) layers in the drafter, and possible eagle cache drop issues. All required unmerged vLLM pull requests to fix ([msg 0]). Against this backdrop, DDTree represented a potential breakthrough. DFlash produces per-position distributions for all future positions in a single forward pass; DDTree exploits this by constructing a tree of candidate continuations from those distributions and verifying multiple paths in a single target model forward pass. If DDTree could be integrated into the existing DFlash pipeline, it might dramatically improve acceptance rates.

The DDTree Investigation: A Methodical Code Deep-Dive

The DDTree investigation was launched as a subagent session with a focused mandate: research the repository at https://github.com/liranringel/ddtree thoroughly and answer seven specific questions about its implementation [3]. The user's questions were not casual—they formed a systematic probe of the integration surface area between DDTree and production serving frameworks like SGLang and vLLM.

The assistant's approach was textbook top-down code comprehension. In the first message, it cloned the repository and fetched the GitHub page in parallel, demonstrating a sophisticated understanding of the opencode tool paradigm where independent operations can be batched to minimize latency [2]. The second message read all the core source files: ddtree.py (the core algorithm), dflash.py (the baseline DFlash implementation), benchmark.py (the evaluation harness), distributed.py (multi-GPU utilities), and requirements.txt (dependencies) [4]. The third message pivoted to the model/ subdirectory, reading the DFlash draft model implementation [5].

This progression reveals a deliberate research methodology. The assistant began with breadth-first understanding—the overall structure and call graph—then drilled into depth-first exploration of the dependency chain. The model/ directory contained the DFlashDraftModel class, the sample function, and the extract_context_feature function that the core algorithm files imported but did not define [7]. Without understanding these, the assistant could not trace the algorithm's data flow or answer the user's questions about the DFlash interface.

The Architecture of DDTree: Tree Construction and Verification

The comprehensive technical analysis that emerged in message 6 [6] is a masterclass in technical synthesis. The assistant organized its findings into seven sections, each corresponding to one of the user's questions, and produced what amounts to a complete algorithmic specification of DDTree.

Tree construction is the heart of the algorithm. DFlash, as a block-diffusion drafter, produces logits for all block_size - 1 future positions in a single forward pass. Unlike autoregressive drafters that produce one token at a time, DFlash gives a full probability distribution at each position simultaneously. DDTree exploits this by using a best-first heap expansion algorithm (build_ddtree_tree(), ddtree.py:84-166). The algorithm extracts top-K token IDs and their log-probabilities at each position, then seeds a max-heap with the single most-probable token at depth 1. It pops the globally highest-probability node and performs two expansions: sibling expansion (adding the next-ranked token at the same depth/parent) and child expansion (adding the best token at the next depth as a child of the current node). This continues until a budget of nodes is allocated.

Tree attention verification is the second key innovation. The tree is flattened into tensors suitable for a single target model forward pass: verify_input_ids (root + all tree nodes), verify_position_ids (nodes at the same depth share positions), and an attention_mask that enforces ancestor-only visibility. The mask ensures each tree node attends only to its own ancestor path plus the full KV prefix. After the forward pass, a greedy tree walk (follow_verified_tree) compares the target model's sampled tokens against the tree structure, walking down the tree as long as matches are found. The KV-cache is then compacted to keep only the accepted path's entries.

The analysis revealed that DDTree is almost entirely pure PyTorch—zero custom CUDA kernels. The one optional C++ extension for KV-cache compaction has a pure Python fallback. This portability is excellent news for integration [6].

The Critical Discovery: vLLM's Architectural Limitation

Despite the thorough analysis, the assistant soon discovered a fundamental problem when attempting to integrate DDTree into vLLM. The chunk summary reveals the critical finding: vLLM's verification pipeline uses a linear-chain rejection sampler, not a tree-walk sampler, even in its EAGLE tree mode. This means EAGLE's tree attention is only used during the drafting phase, not for verifying multiple candidate paths. Implementing true DDTree verification would require writing a new tree-walk rejection kernel from scratch.

This is a profound architectural constraint. The entire DDTree approach depends on the ability to verify multiple candidate paths in a single forward pass with a custom attention mask. If the serving framework's verification pipeline is hardcoded for linear chains, DDTree's tree attention mechanism cannot be directly plugged in. The assistant would need to either modify vLLM's core verification logic—a significant engineering undertaking—or find another path forward.

Faced with this complexity, the assistant pivoted to running the DDTree authors' standalone code directly, successfully patching it for the Qwen3.6-27B GDN hybrid model. The benchmark confirmed that DDTree works correctly, but the acceptance rate improvement over DFlash was marginal: 1.67 vs 1.59. The reason was sobering: the underlying DFlash drafter is "still under training," making the drafter's quality the primary bottleneck. Even with perfect tree construction and verification, a weak drafter produces poor candidates.

The Pivot to Training: Recognizing the Real Bottleneck

This realization triggered a fundamental strategic shift. The team had been pursuing better verification algorithms—first DFlash, then DDTree—but the data showed that the drafter itself was the limiting factor. No amount of clever tree construction could compensate for a draft model that simply didn't predict the target model's outputs well. The path to better speculative decoding lay not in better verification, but in training a better drafter.

The assistant pivoted to building a comprehensive training infrastructure. This was not a small task. Training a DFlash drafter requires:

  1. A curated dataset aligned with the target model's use case
  2. Hidden states extracted from the target model to serve as conditioning signals
  3. A training pipeline compatible with the vllm-project/speculators framework
  4. A hardware environment capable of running both the target model (for hidden state extraction) and the training process

Dataset Curation: 913K Samples for Drafter Alignment

The dataset curation effort was substantial. The assistant assembled a 913K-sample dataset mixing several sources:

Training Infrastructure: 8× Blackwell GPUs

The assistant orchestrated the setup across three remote machines before landing on a stable 8× RTX PRO 6000 Blackwell node (96GB each, 1.9TB disk). The environment was provisioned with uv, speculators, and vLLM. The 55GB Qwen3.6-27B model was downloaded from HuggingFace in approximately 10 seconds—a testament to the node's network bandwidth.

The tokenized data and DFlash drafter checkpoint were transferred, and a test training run (100 samples, 1 epoch) was successfully launched. The architecture split the GPUs: the vLLM server served hidden states on GPUs 0-3, while the DFlash training ran on GPUs 4-7. This dual-use of the GPU cluster demonstrates the resource constraints of training large draft models—the target model must be running to provide hidden states, consuming half the available GPU memory.

The Hidden State Extraction Pipeline

The training pipeline required hidden states from the target model to condition the drafter's predictions. The initial approach used the speculators online vLLM pipeline, but this proved fundamentally incompatible with Qwen3.6-27B's GDN hybrid KV cache. The assistant pivoted to a custom offline extraction using HuggingFace Transformers.

The initial extraction ran at only 7–11 samples per second per GPU with high CPU system overhead due to per-sample safetensors writes and individual GPU-to-CPU copies. The key breakthrough was batching the hidden state capture entirely on GPU—concatenating all samples in a batch before a single .cpu() transfer. This eliminated 2,725 individual copies per batch and boosted throughput to 140–155 samples per second per GPU (aggregate approximately 600/s). The assistant also integrated async S3 uploads via subprocess to offload storage, installed flash-linear-attention (FLA) to reduce kernel overhead from GDN layers, and added backpressure to pause extraction when /dev/shm exceeded 80% capacity.

The pipeline was made robust for unattended operation with marker-based resume (skipping already-uploaded batches), a Flask monitoring UI showing per-shard progress and GPU stats, and cleanup of stale state after a node migration. When the original instance was killed, the assistant re-provisioned everything on a new 4× RTX PRO 6000 Blackwell node, installed dependencies (including building causal-conv1d for CUDA 13.0), pre-warmed Triton kernels, and restarted extraction.

The Broader Narrative: From Integration to Training

This chunk represents a complete arc from deployment to training. The team began by trying to deploy existing speculative decoding methods (DFlash, DDTree) on their target model. They encountered integration bugs, architectural limitations, and ultimately a quality bottleneck in the drafter itself. Each failure taught them something: the layer-ID offset bug revealed the tight coupling between vLLM's hidden state extraction and the draft model's architecture; the SWA layer issue showed that attention type compatibility matters; the DDTree architectural limitation exposed the fundamental difference between tree-based drafting and tree-based verification; and the marginal acceptance rate improvement confirmed that the drafter's predictive quality is the dominant factor.

The pivot to training represents a maturation of the project. Instead of trying to squeeze more performance out of existing components, the team recognized that they needed to build better components. This is a common pattern in ML engineering: the first deployment reveals the bottlenecks, and the second phase addresses them at the source.

Conclusion

The DDTree investigation and subsequent pivot to drafter training illustrate a fundamental truth about speculative decoding research: the algorithm is only as good as the draft model that powers it. DDTree's tree construction and tree attention verification are elegant innovations, but they cannot compensate for a drafter that hasn't been trained to predict the target model's outputs. The team's journey—from deep code analysis to architectural discovery to dataset curation to training infrastructure—demonstrates the full stack of skills required to advance speculative decoding in production.

The 913K-sample dataset, the optimized hidden state extraction pipeline running at 600 samples per second, and the training environment on 8× Blackwell GPUs represent the tangible output of this pivot. With these assets in place, the team is positioned to train a DFlash drafter that is genuinely aligned with the Qwen3.6-27B target model and its agentic use case. The DDTree investigation was not a detour—it was the diagnostic that revealed the true bottleneck, and the training infrastructure is the treatment.