Yggdrasil: Bridging Dynamic Speculation and Static Runtime for Latency-Optimal Tree-Based LLM Decoding

AI summary

Overview

Research area: Efficient large language model (LLM) inference — specifically speculative decoding, tree-based drafting, and compiler/runtime co-design for GPU execution.

Technical level: Advanced. The paper assumes familiarity with speculative decoding, tree attention masks, KV caches, CUDA Graphs, and deep learning compilers such as TorchInductor.

Scope: The paper presents Yggdrasil, a co-designed speculative decoding system that combines a context-aware tree drafting algorithm with a compiler-friendly runtime to minimize end-to-end per-token decoding latency across several GPUs and LLM pairs.

What This Paper Is About

Speculative decoding speeds up LLM generation by drafting several candidate tokens and verifying them in parallel, but the paper argues that existing systems are stuck between two bad options: dynamic tree drafting that raises average accepted length (AAL) but breaks the static graphs modern compilers rely on, and static/compiled runtimes that cut per-step latency but cannot adapt to context. The goal of Yggdrasil is to get both at once — a tree structure that changes with the decoding context while still presenting fixed operator shapes that a compiler can fuse, tune, and capture into a graph, plus a runtime schedule that hides CPU-side control logic behind GPU work.

Key Contributions

1. A latency-aware optimization objective. Instead of approximating speedup with AAL, the authors formulate speedup in terms of the actual execution latency of the drafter and verifier, exposing the tight coupling between draft width, draft depth, and verification width (Equation 3).
2. The Equal-Growth Tree (EGT) drafting algorithm. A context-adaptive tree that predicts draft depth, greedily selects draft width, and prunes to a verification width via a bottom-up dynamic programming maximum-value subtree search — while keeping static operator shapes for graph compilation compatibility.
3. A stage-based scheduling runtime. Ahead-of-time tail drafting, ahead-of-time head drafting, and an offline profile-guided execution plan search that break conditional dependencies in the speculative decoding pipeline and overlap CPU management with GPU inference.
4. A full implementation and evaluation. Built on PyTorch and the TorchInductor compiler backend, with unmodified models, evaluated on real models and datasets, reporting up to 3.98x speedup over baselines including SpecInfer, Sequoia, and vLLM-Spec.

Main Findings

• End-to-end latency: Yggdrasil achieves average decoding latency improvements of 3.98x on A100 GPUs and 2.76x on A40 GPUs over the baselines, and up to 3.98x speedup overall across multiple hardware setups.
• Baseline comparisons: SpecInfer performs the worst. Sequoia and vLLM-Spec, which use TorchInductor, achieve average speedups of 2.45x and 2.66x over SpecInfer respectively; Yggdrasil improves this to 3.37x through context-aware speculation and CUDA-graph optimizations. Gains are more pronounced on A100 GPUs, which the authors attribute to greater hardware under-utilization.
• Graph compilation is the single biggest win: In the optimization breakdown, graph compilation with TorchInductor (O2) accounts for an average 2.775x speedup, the largest of the optimizations — but the authors note this graph can only be achieved together with the EGT design.
• Adaptive verification helps: Verification width pruning (O3) adds an average 1.07x over O2 by aligning the verification number with the saturation region of the inference wall-time curve.
• Scheduling matters: Graph-based scheduling (O4) delivers an average 1.21x improvement over previous optimizations.
• Depth prediction helps: The draft depth predictor (O5) brings a 1.10x speedup by explicitly accounting for context difficulty, compared against a fixed baseline of depth 16 and width 8.
• Speedup objective beats AAL objective: Incorporating the latency-aware speedup objective yields an extra 8% performance improvement over optimizing AAL directly across all drafter-verifier settings.
• Tree structure sensitivity: On the Wikitest dataset, sequence speculative decoding shows limited AAL that saturates quickly as verification number grows. Sequoia achieves good AAL with a single static structure for all inputs, while EGT dynamically adjusts width and consistently outperforms it. In static parameter exploration, the combination D_draft = 8, W_draft = 8, W_verify = 64 yielded the best performance.
• Temperature effects: Both Sequoia and Yggdrasil perform better at temperature 0, because the draft model aligns better with the target model when temperature is low, raising AAL. Across temperatures, Yggdrasil achieves an average 1.49x speedup over Sequoia.
• Motivating measurements: Capturing Llama-2-7B with CUDA Graphs gives a 2.32x speedup, but dynamic drafting control flow cannot be frozen into a single graph; kernel auto-tuning gives an extra 1.23x only when shapes stay constant.

Methodology in Plain English

The authors start from the observation that the metric everyone optimizes — average accepted length — is only a proxy. Accepting more tokens is not the same as finishing faster, because verifying more tokens costs more time per step, and because dynamic tree shapes prevent the compiler from generating fast static kernels.

Their approach has three parts. First, they replace the AAL objective with a direct latency model that accounts for the drafter and verifier runtimes, then decompose the resulting high-dimensional search into three coordinated greedy decisions: a lightweight multi-head predictor (a two-layer MLP encoder with several depth-prediction heads, trained offline per dataset and drafter/verifier pair from one-time profiling on an in-domain corpus) predicts how many draft steps to take; a width is chosen to maximize the speedup objective under that depth; and after drafting, the tree is pruned with dynamic programming to the verification width that maximizes the same objective, subject to the verification budget. The tree grows a fixed number of leaves per step — hence "equal growth" — so operator shapes stay constant and the compiler can still fuse and tune.

Second, the runtime rearranges the pipeline. Because speculative decoding has data dependencies that seem to forbid overlap, the authors speculatively break them: instead of conditionally drafting the tail token, they draft the whole candidate sequence ahead of time and reuse whichever leaf is accepted; instead of decoding the head token from the confirmed root, they issue the head draft right after the previous bonus draft. A small offline grid search over profiled stage latencies picks the best overlap plan at compile time.

Third, everything is wrapped in a TokenTree abstraction that manages the token sequences, KV cache tensors, tree structure array, attention mask, and leaf positions, interoperating with unmodified drafter and verifier models through PyTorch and TorchInductor.

Why This Matters

Impact on research: The paper reframes speculative decoding as a joint algorithm–runtime problem rather than a drafting-algorithm problem, and it supplies evidence that AAL is a misleading optimization target once verification cost grows. It also demonstrates that dynamic, context-aware speculation and static-graph compilation need not be mutually exclusive — a result that generalizes beyond this specific system.

Real-world applications:

• Latency-critical interactive assistants and chat deployments where per-token latency (TPOT) dominates user experience.
• On-premises or edge inference with a dedicated accelerator serving a single user, the deployment mode the authors explicitly target.
• Code completion and other IDE-integrated generation, where latency is directly on the developer's critical path.
• Latency-sensitive pipelines such as interactive agents or real-time document generation, where faster decoding compounds over long outputs.

Industry relevance: The system requires no source-level modification to model architectures, and it is implemented on widely used infrastructure (PyTorch, TorchInductor). That makes its techniques — EGT-compatible graph compilation, verification-width pruning, stage overlap — directly applicable to existing serving stacks rather than requiring a bespoke inference engine.

Future Directions

• Batched, throughput-oriented serving. The authors state plainly that their latency-optimal setting assumes one interactive request monopolizes all GPU memory and compute, and that this is not applicable to production stacks that batch tens to hundreds of prompts. A unified policy that jointly decides when to speculate and how to pack requests is called an open line of research.
• Joint latency–throughput optimization. Extending the framework into the space where speculative decoding and batch scheduling are solved together.
• Broader hardware and model coverage. The motivated claim that gains are larger on hardware with greater under-utilization (A100 versus A40) invites study of how EGT parameters should be searched for other accelerators and larger target models (the evaluation uses Llama-2-7B and Llama-2-13B with Llama-68M and Llama-160M drafters).
• Blending with model-invasive speculation. The authors note that model-invasive methods which modify the target model are outside their supporting scope, leaving room for combining EGT with those approaches.

Target Audience

Systems and machine learning researchers working on LLM inference acceleration, deep learning compiler engineers interested in dynamic workloads over static graphs, and performance engineers building production or on-premises LLM serving stacks who need to reduce per-token latency without retraining or modifying models.

Authors’ abstract

Speculative decoding improves LLM inference by generating and verifying multiple tokens in parallel, but existing systems suffer from suboptimal performance due to a mismatch between dynamic speculation and static runtime assumptions. We present Yggdrasil, a co-designed system that enables latency-optimal speculative decoding through context-aware tree drafting and compiler-friendly execution. Yggdrasil introduces an equal-growth tree structure for static graph compatibility, a latency-aware optimization objective for draft selection, and stage-based scheduling to reduce overhead. Yggdrasil supports unmodified LLMs and achieves up to $3.98\times$ speedup over state-of-the-art baselines across multiple hardware setups.