PaperlayerView exact revision on arXiv

Source-linked AI summary

TreeWY: Speculative Verification for Gated DeltaNet Hybrids

Sneha Murthy Ghantasala

arXiv:2608.20961v1cs.AIcs.CLcs.DCcs.LGcs.PF

TL;DR

Speculative verification is memory-intensive for GDN hybrids because rejected draft positions require rollback of a recurrent state that cannot be partially truncated. TreeWY replaces per-node snapshots with a tree-structured WY transform, jointly solves for draft outputs, and reconstructs only the accepted state; benchmarks preserve acceptance while reducing memory pressure and improving serving performance where memory binds, although wider trees are not yet a throughput win.

  • Problem

    GDN recurrent states cannot be partially rolled back, so speculative verification requires full snapshots at draft positions that cannot be shared across tree branches.

  • Method

    TreeWY applies a tree-structured WY/UT transform to the gated delta rule, verifies every draft node with one triangular solve, and stores a pseudo-value matrix for commit-time reconstruction.

  • Results

    Acceptance length remains essentially identical to baseline while TreeWY reduces peak KV usage by 2–3×, reaches up to 1.49× throughput and ∼40× lower p99 TTFT where memory binds, and trails by 0.97–0.99× where it does not.

  • Takeaways & Limitations

    TreeWY makes wider, higher-acceptance draft trees affordable by keeping stored GDN state flat at one block regardless of width.

  • Takeaways & Limitations

    Wider trees are enabled and correct but are not yet a throughput win because they require more target tokens per step and a costlier non-capturable tree-verification kernel.

Abstract

from arXiv · show

Modern open models are hybrids: most layers are linear-attention (Gated DeltaNet, GDN) layers carrying a small fixed-size recurrent state instead of a growing key-value (KV) cache. This makes ordinary decoding memory-efficient, but hurts speculative decoding. To verify a batch of draft tokens and then roll back the rejected ones, today's systems snapshot the full recurrent state at every draft position for GDN layers, and those snapshots cannot be shared across branches of a draft tree, so a wide, high-acceptance tree becomes memory-infeasible. We remove the snapshots. Using a tree-structured WY transform of the gated delta rule, we compute every draft node's output with a single triangular solve and reconstruct only the one accepted state on commit, storing a small pseudo-value matrix instead of per-node states; the derivation depends only on the gated delta rule, not on any other architectural detail. In serving benchmarks on two scales of one hybrid model family (Qwen3.5 35B and 397B) this cuts speculative recurrent-state memory and KV-cache pressure at identical acceptance length, turning the freed HBM into higher throughput and much lower time-to-first-token (TTFT) wherever memory binds, and costing a few percent where it does not. For tree width the same memory buys affordability: a wider, higher-acceptance draft becomes possible, though not yet a throughput win.

1 Background

Speculative decoding verifies multiple draft tokens in one target-model pass, but hybrid models use GDN recurrent states that cannot be partially rolled back. This makes GDN state handling the source of the speculative-decoding challenge.

  • Speculative decoding: Speculative decoding has a target model verify multiple drafter-proposed tokens in one pass and accept the longest matching prefix.A tree drafter can propose multiple alternatives at each position.
  • Hybrid models and GDN layers: GDN layers maintain one fixed-size recurrent state matrix per head instead of a growing KV cache.The state summarizes the whole prefix while softmax-attention layers store key and value vectors for previous tokens.
  • Hybrid models and GDN layers: The gated delta rule updates the state using decay, keys, values, queries, and a write-strength gate.Its delta term erases what the state already predicts for the current key before writing the new value.
  • Hybrid models and GDN layers: A GDN state is a lossy summary of the prefix, so it cannot be truncated or partially rolled back like a KV cache.That property directly creates the speculative-decoding problem.

2 The memory problem

Speculative verification is cheap for softmax KV caches but expensive for GDN layers because parallel verification advances the recurrent state past rejected tokens. Existing snapshotting therefore consumes memory proportional to draft size and cannot share state across tree branches.

  • Speculative verification: During speculative verification, softmax layers append draft KV entries and roll back rejected tokens by moving a pointer.Tree branches can share common-prefix KV entries, leaving only a small marginal memory cost.
  • Speculative verification: GDN verification advances the recurrent state through the entire draft before acceptance is known, requiring rollback to the accepted node.The state must therefore be recoverable for every possible accepted position.
  • Existing approaches and trade-offs: Full-state snapshotting stores k+1 GDN state blocks for a chain of k draft tokens and N+1 blocks for a tree of N nodes.For k = 3, this reaches 120 MiB per sequence at 35B and 360 MiB at 397B.
  • Existing approaches and trade-offs: Snapshot memory scales with draft size and cannot be shared across tree branches, limiting concurrency and wide-tree speculation.This is the bottleneck addressed by TreeWY.
  • Existing approaches and trade-offs: ReplaySSM defers state materialization and solves the same triangular system inside its verify window, while STree targets tree verification only for Mamba2.The paper positions TreeWY as a GDN tree method using a different state-materialization strategy from ReplaySSM.
  • Existing approaches and trade-offs: TreeWY applies a tree-structured WY/UT transform of the gated delta rule and stores one small pseudo-value matrix written only on commit.It uses one triangular solve to verify every draft node.

3 Method: TreeWY

TreeWY rewrites the gated delta recurrence as decay-weighted additive attention with corrected pseudo-values, then solves for all tree nodes jointly. It reconstructs only the accepted continuation state, replacing per-node full-state snapshots with a small matrix.

  • Key rewrite: The gated delta rule can be rewritten as decay-weighted additive attention using a pseudo-value that subtracts the state’s current prediction.This removes the need to walk the recurrence to obtain intermediate states.
  • Verifying a draft: TreeWY lays out chain or tree draft nodes in depth-first order so every ancestor precedes its descendant.The same ordering supports both chains and branching trees.
  • Verifying a draft: Chaining the rewrite along root-to-node paths produces one linear system for all pseudo-values, with a strictly lower-triangular ancestor structure.Forward substitution replaces recurrence execution, and every node output reads from the same solve.
  • Reconstruct on commit: After verification accepts node a, TreeWY reconstructs the continuation state by summing pseudo-values over a’s ancestors.The reconstructed state becomes the next round’s committed state.
  • Reconstruct on commit: TreeWY stores O(Nd_v) pseudo-values instead of one full state per node, yielding a 128× smaller per-head object when d_k=d_v=128.This changes chain and tree storage from N+1 state blocks to one block per sequence per layer.

4 Implementation

TreeWY is implemented in vLLM with configuration options for reconstructed commits and draft-tree widths. Chain verification uses a fused, CUDA-graph-capturable kernel, whereas real trees require piecewise capture and atomic scheduling.

  • Implementation: The vLLM fork enables reconstruction with mamba_state_commit="reconstruct" and configures branching with draft_tree_widths.The default commit mode is "store_all".
  • Implementation: Chains verify and commit in one fused, CUDAgraph-capturable Triton kernel.An all-ones tree is treated as a chain.
  • Implementation: Real trees require a non-causal ancestor mask that prevents CUDA-graph replay and causes the model to use piecewise capture.This also evicts the GDN mixer from graphs, making the cost larger than the mask itself.
  • Implementation: Tree requests must be scheduled atomically because truncating a tree changes its topology.Requests whose tree does not fit the per-step token budget skip speculation rather than being truncated.
  • Implementation: The implementation checks closed-form and production-kernel correctness against recurrence references within stated fp64, fp32, and bf16 tolerances.The closed form matches recurrence to ∼10−15 in fp64 and ∼10−7 in fp32.

5 Evaluation

The evaluation compares TreeWY with store-all snapshotting across two Qwen3.5 scales, memory budgets, workloads, and concurrency settings. TreeWY preserves acceptance while reducing memory pressure and improving serving performance where memory binds; wide trees are affordable but not yet faster.

  • Setup: The setup uses Qwen3.5-35B-A3B at TP1 and Qwen3.5-397B-A17B at TP8 with depth-3 MTP drafts on B200 GPUs.The sweeps cover six workloads and vary gmu across the reported budgets.
  • Correctness: 175 matched evaluations show acceptance length essentially identical to storeall, with mean |∆| = 0.039 and maximum 0.33.Acceptance matches within 0.01 at every depth on real prompts.
  • Memory and throughput: TreeWY reduces peak KV usage by 2–3× at the same load and preempts fewer requests than storeall.Preemptions were 1365 for TreeWY versus 2531 for storeall.
  • Memory and throughput: Where memory binds, TreeWY reaches up to 1.49× throughput and approximately 40× lower p99 TTFT.Where memory does not bind, throughput is 0.97–0.99× while the KV reduction remains 2–3×.
  • Memory and throughput: The memory-pressure knee shifts with budget: it occurs at concurrency 128 for gmu=0.6, 256 for gmu=0.75, and never at gmu=0.9.At b=1.00, the freed memory remains unspent while other metrics stay within 3% of parity.

6 Conclusion

TreeWY applies a tree-structured WY transform to verify GDN draft trees with one triangular solve and reconstruct only the accepted state. It reduces per-node snapshots to one pseudo-value matrix, improving memory-bound chain serving while making wider trees affordable but not yet faster.

  • TreeWY replaces O(N) recurrent-state snapshots with one small pseudo-value matrix for chains and trees.The accepted state is reconstructed only at commit.
  • On chains, freed HBM yields higher throughput and much lower TTFT wherever memory binds, with a residual per-step implementation cost.The chain kernel is fused and graph-capturable.
  • On trees, stored state remains one block regardless of width, making wider, higher-acceptance drafts affordable but not yet a throughput win.The non-capturable tree verification kernel costs more than the extra acceptance returns.
  • The paper identifies graph-capturable tree fusion, deferred state writes, and evaluation on a second model family as next steps.

A Remaining memory budgets

At higher memory budgets, Table 3 uses the same axes and orientation as Table 1. The 397B knee appears only at gmu=0.75 and concurrency 256, while b=1.00 indicates freed memory is not yet used.

  • At 397B, the memory-pressure knee appears only at gmu=0.75 and concurrency 256.
  • Where b reads 1.00, freed memory is real but unspent.

B Wide draft trees

TreeWY keeps per-request state cost flat as draft-tree width grows, whereas store-all cost increases with the number of nodes. Acceptance length rises across the wider-tree range, but remains bounded by draft depth.

  • Store-all cost grows 10×, from 4 to 40 blocks, across the reported tree shapes, while TreeWY remains at one block.
  • Acceptance length rises from 1.883 to 2.786 across the same shape range, despite being bounded by draft depth rather than width.Acceptance is nearly flat from (2, 2, 2) to (3, 3, 3) while blocks increase from 15 to 40.
  • The acceptance gain comes from matching each level’s probability mass more often, not from extending the accepted path.
  • Table 4 evaluates Qwen3.5-35B-A3B with depth-3 trees and reports acceptance matching the storeall reference within sampling noise.

C Admission vs. per-step cost

TreeWY’s serving gains arise mainly when eliminating snapshots admits more requests; when admission is matched, differences reflect per-step cost rather than memory.

  • At the 31 memory-bound 35B points, TreeWY reaches 1.15× throughput, 2.94× lower p99 TTFT, and 1.17× lower mean end-to-end latency at 0.83× TPOT.The lower TPOT reflects larger admitted batches, while mean latency improves because requests the baseline could not admit stop waiting.
  • Across the other 74 points, TreeWY matches batch size within 1% and has 0.98 TPOT, identifying this as its per-step cost.The same decomposition gives ReplaySSM r = −0.79.
  • ReplaySSM turns comparable freed headroom into 1.12–1.20× throughput, versus TreeWY’s 0.99–1.08×, while preserving acceptance on real prompts.The comparison normalizes each method to its own snapshotting baseline; the store-all throughputs agree within 1%.
  • At 256 concurrency, both methods admit essentially the same work, and throughput ratios close to within 3% once both are saturated and decode-bound.The table attributes the remaining difference to per-token efficiency, not the memory mechanism.
  • ReplaySSM leads on throughput and per-token speed at every load, while the methods converge on memory-driven columns after the knee.At 256 concurrency, peak-KV reductions are 1.56 versus 1.57 and admitted-batch ratios are 1.58 versus 1.46.

E Additional serving curves

Serving curves show TreeWY preserving KV headroom as load rises, preventing saturation-driven queuing and sharply reducing TTFT when memory binds.

  • TreeWY keeps 2–3× more KV headroom at the same offered load than the snapshotting baseline, which reaches the 100% ceiling early.The freed capacity supports admission, throughput, and TTFT improvements.
  • At 256 concurrency, baseline p99 TTFT is up to ∼30× worse for 35B and 26× worse for 397B at utilization 0.75.The gap arises as the baseline’s KV pool saturates and queues requests.
  • The largest sweep gap is 40×, comparing 683 versus 27489 ms at 128 concurrency for 35B and utilization 0.6.TreeWY remains unsaturated at this point while the baseline saturates.
  • When memory is slack, the serving gap closes, whereas memory-bound points show TreeWY reaching 1.15× throughput and 2.94× lower p99 TTFT.At 35B, 31 of 105 points are memory-bound; TreeWY also provides 1.60× more KV headroom and 0.83× per-token cost there.

F Compute resources

The evaluation used approximately 85 GPU-hours across continuous serving sweeps on B200 GPUs, covering both model scales and a ReplaySSM comparison.

  • ≈85 GPU-hours covered the 35B main sweep, the 397B main sweep, and six additional 35B ReplaySSM comparison jobs.All runs used B200 GPUs with 178 GiB HBM per device.
Loading 2608.20961v1…