PaperlayerView exact revision on arXiv

Source-linked AI summary

SpecFold: Folding Multi-Branch Redundancy for Faster Speculative Decoding in Diffusion Language Models

Chung-En Ho, Weiyu Sun, Cheng-Jhih Shih, He Li, Yong Liu, Yingyan Celine Lin

arXiv:2610.04875v2cs.AIcs.CLcs.LG

TL;DR

DLLM decoding remains costly because repeated target-model evaluations refine denoising blocks, while dense multi-branch verification recomputes drafts that largely inherit their parents’ tokens. SPECFOLD folds this cross-branch redundancy into sparse execution. Across two DLLM families, five models, and five benchmarks, it reports up to 1.64× throughput over Spiffy and 1.99× over vanilla decoding while maintaining comparable task performance.

  • Problem

    Repeated target-model evaluations remain a dominant DLLM decoding cost, and dense multi-branch verification recomputes branches despite their substantial parent–child hidden-state similarity.

    Source: §1 · pp. 1–2

  • Method

    SPECFOLD uses token-level residual gating to reuse parent computation at similar positions, folds attention and FFN work, and maps recomputation to sparse execution with a Triton kernel.

    Source: §1 · p. 2, §3.2 · p. 6

  • Results

    Up to 1.64× throughput over Spiffy and 1.99× over Vanilla decoding across two DLLM families, five models, and five benchmarks, while maintaining comparable task performance.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    SPECFOLD removes redundant computation within speculative verification and is orthogonal to temporal caching and compatible with existing DLLM speculation strategies.

    Source: §1 · p. 2

    Reuse is premised on inherited-token layer-input hidden states being nearly identical across layers, so the stated reuse opportunity concerns those positions.

    Source: §3.1 · p. 5

Abstract

from arXiv · show

Diffusion large language models (DLLMs) generate text through iterative block denoising, and multi-branch speculative decoding accelerates this process by verifying a main branch together with multiple draft branches in a single forward pass. While prior DLLM acceleration methods primarily exploit temporal redundancy across denoising steps, we identify a complementary redundancy axis within each speculative verification step: multi-branch computational redundancy. During speculative verification, draft branches inherit most tokens from their parents while unmasking a small set of additional positions, causing large portions of hidden states to remain highly similar across branches. We propose SpecFold, an algorithm-system co-design that exploits this multi-branch redundancy to reduce the cost of multi-branch speculative verification. Algorithmically, SpecFold performs token-level residual gating and selectively reuses parent computation through folded attention and FFN while preserving residual hidden states. Systemically, a Triton kernel implementation translates this fine-grained reuse into end-to-end throughput gains through efficient sparse multi-branch execution. SpecFold is orthogonal to temporal caching and compatible with existing DLLM speculation strategies. Across two DLLM families, five models, and five standard benchmarks, SpecFold achieves up to 1.64x throughput over Spiffy and up to 1.99x over vanilla decoding, while maintaining comparable task performance.

1 INTRODUCTION

DLLM multi-branch speculation reduces denoising steps but densely recomputes every branch, despite substantial overlap between parent and child drafts. SPECFOLD folds that within-step redundancy into sparse execution, and the paper reports throughput gains over Spiffy and vanilla decoding across models and benchmarks.

1 INTRODUCTION

  • Dense verification lowers the number of denoising forward passes but raises per-pass cost by evaluating the main branch and every draft branch together.DLLM drafts are full block states with bidirectional attention, and a matching draft can supply logits that skip future denoising passes.Source: §1 · p. 2
  • More than 72% of measured draft–layer–position–step tuples have high parent–child hidden-state similarity, with divergence concentrated near recently unmasked positions.Drafts inherit almost all tokens from their parents, yet existing methods densely recompute every branch.Source: §1 · p. 2
  • SPECFOLD uses token-level parent–child residual gating to reuse similar-position computation and recompute divergent positions, with folded attention and FFN while preserving branch residual streams.Its Triton implementation compacts recomputed positions, resolves fold chains, and applies attention correction without materializing redundant branch tensors.Source: §1 · p. 2
  • SPECFOLD targets redundancy within each speculative verification step, complementing temporal caching and composing with existing speculation approaches such as Spiffy without changing draft construction policies.Temporal caching reuses computation across denoising steps along a single trajectory.Source: §1 · p. 2
  • Up to 1.64× throughput over Spiffy and 1.99× over Vanilla decoding across two DLLM families, five models, and five benchmarks, while maintaining comparable task performance.The proposed co-design combines token-level folding with a dedicated kernel implementation to reduce redundant computation within each speculative verification step.Source: §1 · p. 2

2 BACKGROUND

DLLMs iteratively denoise token blocks, while multi-branch speculation verifies related draft states together to reduce sequential denoising passes. The section distinguishes this approach from temporal reuse and identifies redundancy across branches within a verification step.

2.1 MASKED DIFFUSION LLMS AND TEMPORAL ACCELERATION

  • DLLMs generate block by block, repeatedly predicting distributions and unmasking selected high-confidence positions while refining the remaining masked positions.A denoising step can update multiple positions in parallel, unlike autoregressive decoding.Source: §2.1 · p. 3

2.2 MULTI-BRANCH SPECULATIVE DECODING IN DLLMS

  • Multi-branch DLLM speculation verifies a main active-block state and parent-linked draft states together, with each draft representing a full block configuration.Bidirectional attention distinguishes these drafts from autoregressive next-token continuations.Source: §2.2 · p. 3
  • Draft branches inherit parent states and add unmasked positions; dense verification computes every position of every branch, increasing per-pass cost as branch count grows.A matching cached draft can supply precomputed logits to skip a later denoising pass.Source: §2.2 · p. 3

2.3 REDUNDANCY EXPLOITATION IN TRANSFORMER INFERENCE

  • SPECFOLD targets a third redundancy axis: folding draft-branch computation within one DLLM verification step, rather than reusing activations across denoising steps or exploiting within-pass sparsity.Multi-branch methods batch draft graphs but otherwise recompute sibling branches’ suffix activations independently.Source: §2.3 · p. 4

3 THE PROPOSED SPECFOLD FRAMEWORK

SPECFOLD targets cross-branch redundancy within speculative verification, where inherited-token representations remain similar across drafts. It profiles this redundancy and uses gated parent-computation reuse with a Triton implementation.

3.1 SPECFOLD: PROFILING THE REDUNDANCY

  • Across all 28 layers, 72% of (draft, step, layer, position) cells satisfy ρ ≤ 0.1 and 88% satisfy ρ ≤ 0.3; high-divergence regions concentrate along unmasking positions and positions draft i unmasked.Source: §3.1 · p. 5
  • Mean per-forward latency increased with every added draft from B=2 to B=8 on Nemotron-Labs-Diffusion-3B on H100, despite weights streaming once for the batch.Dense verification computes every position of every draft, paying latency per forward pass.Source: §3.1 · p. 5, Figure 3 · p. 5

3.2 SPECFOLD: FOLDING REDUNDANT DRAFTS

  • At each layer, a parent–child residual gate folds similar positions: QKV projections and FFN reuse parent outputs, while residual streams and corrected attention preserve branch-specific diversity.Positions are recomputed when their divergence exceeds the threshold.Source: §3.2 · p. 6

3.3 TRITON IMPLEMENTATION

  • Triton kernels compact recomputed positions for packed QKV and FFN work, correct parent attention states for folded queries, and support both evaluated model families, including grouped-query attention.The implementation resolves fold chains so folded positions can access parent states.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

The experiments evaluate SPECFOLD across two DLLM families, five models, and five benchmarks, then examine threshold and draft-budget sensitivity and parent–child residual patterns. The results show throughput gains over speculative and vanilla decoding while accuracy remains comparable.

4.1 SETUP

  • Evaluation covers Fast-dLLM-v2 at 1.5B and 7B and Nemotron-Labs-Diffusion at 3B, 8B, and 14B, using five benchmarks and NFE, TPS, and accuracy.Experiments use a single NVIDIA H100 80GB GPU; SPECFOLD and Spiffy use the same draft graph for comparison.Source: §4.1 · p. 7

4.2 MAIN RESULTS

  • SPECFOLD reaches up to 1.64× throughput over Spiffy and 1.99× over Vanilla on Fast-dLLM-v2-1.5B, with comparable accuracy across five benchmarks.On Fast-dLLM-v2-7B, the corresponding maxima are 1.22× over Spiffy and 1.31× over Vanilla; on Nemotron-Labs-Diffusion 3B, 8B, and 14B, maxima are 1.34×, 1.18×, and 1.07× over Spiffy and 1.48×, 1.32×, and 1.24× over Vanilla.Source: §4.2 · pp. 7–8, Table 1 · p. 8

4.3 ABLATIONS

  • At δ = 0.2, accuracy stays at Vanilla’s level across five benchmarks while folding yields 1.3–1.4× throughput over Vanilla; accuracy declines beyond this plateau.Throughput continues to rise, but gains beyond δ = 0.3 come at an accuracy cost; as δ approaches 0, the accuracy–throughput point approaches Spiffy.Source: §4.3 · p. 8, Figure 4 · p. 8
  • The efficiency gains do not depend on a precise δ choice; the main evaluation fixes δ = 0.1 across models and benchmarks without per-task tuning.Source: §4.3 · p. 8
  • With draft budget B, acceptance grows and saturates around B = 6, where NFE reductions reach 30% relative to Vanilla; SPECFOLD reaches up to 1.73× throughput gain.Across B, Spiffy stays within 0.99–1.17× of Vanilla’s TPS, while SPECFOLD shares the same drafts and NFE as Spiffy; Figure 5 plots throughput gain against draft budget.Source: §4.3 · p. 9, Figure 5 · p. 9

4.4 ANALYSIS OF THE RESIDUAL RIPPLE

  • The parent–child residual stays low at committed positions, while its divergence ripple concentrates at masked and recently unmasked positions.The child draft adds one unmasked token to a parent with 17 of 32 positions committed; Figure 6 visualizes the per-position relative residual.Source: §4.4 · p. 9, Figure 6 · p. 9

5 CONCLUSION

SPECFOLD combines token-level folding with a dedicated kernel implementation to reduce redundant computation within each speculative verification step. Its folded-layer design reuses parent computation while retaining branch-specific hidden states and correcting attention for child-specific updates.

5 CONCLUSION

  • Across two DLLM families, five models, and five benchmarks, SPECFOLD achieves up to 1.64× throughput over Spiffy and 1.99× over Vanilla while maintaining comparable task performance.Source: §5 · p. 9

APPENDIX

  • The appendix covers the folded Transformer algorithm, Triton implementation, numerical discrepancy with Vanilla, and additional ablations.The additional ablations include draft-budget sweeps across model scales and families, drafting strategy, and hardware.Source: APPENDIX · p. 13

A FOLDED TRANSFORMER ALGORITHM

  • Each branch inherits its parent’s projection, attention-state, and FFN-output tables, overwriting only recomputed positions while preserving its own hidden states.The main branch recomputes every position; child hidden states retain cross-branch divergence.Source: A FOLDED TRANSFORMER ALGORITHM · p. 13

A.1 FOLDED SELF-ATTENTION

  • Folded attention corrects a parent’s softmax state by removing its contributions at recomputed positions and adding the child’s, yielding the child’s attention state exactly.Bidirectional attention couples positions, so folded positions require correction for the child’s recomputed positions.Source: §A.1 · pp. 13–14

A.2 FOLDED FFN

The folded FFN reuses parent outputs at folded positions but forms residual additions from each branch’s own hidden state. A folded layer composes this FFN reuse with folded attention and propagates corrections along branch chains.

A.2 FOLDED FFN

  • Because the FFN is position-wise, folded positions reuse the parent’s output table, while the residual add retains each branch’s own hidden state and preserves divergence.Source: §A.2 · p. 14

A.3 PUTTING IT TOGETHER: ONE FOLDED LAYER

  • A folded layer visits branches topologically so each inherits updated parent tables and applies attention correction against its immediate parent; corrections propagate along fold chains.The kernel resolves each folded row to its nearest recomputed ancestor and applies accumulated correction in one order-free launch across branches.Source: §A.3 · p. 14

B TRITON IMPLEMENTATION

The Triton implementation maps SPECFOLDLAYER onto six kernel families operating over flattened branch-position pairs. It gates and compacts positions, resolves folded sources, and computes attention and FFN updates while preserving each position’s residual stream.

B TRITON IMPLEMENTATION

  • The implementation runs six kernel families per layer over a flattened space of (B+1)N branch-position pairs.Source: B TRITON IMPLEMENTATION · p. 14
  • A gate marks each position for recomputation or reuse and compacts recomputed indices into global and per-branch lists.The main branch bypasses the gate and recomputes every position.Source: B TRITON IMPLEMENTATION · p. 14
  • A source-resolution kernel follows folded-position chains to the nearest ancestor where that position was recomputed.Source: B TRITON IMPLEMENTATION · p. 14
  • Each branch inherits its parent’s tables; recomputed positions receive QKV projections and attention, while folded positions receive an incremental attention correction.Branches are processed in topological order, with corrections applied against the immediate parent.Source: B TRITON IMPLEMENTATION · p. 15
  • Recomputed positions receive RMSNorm and packed QKV projections with rotary embeddings; attention kernels process recomputed queries and folded rows reuse attention state with corrections.The attention kernel retains per-head softmax maximum and denominator state for the correction kernel.Source: B TRITON IMPLEMENTATION · p. 15
  • The FFN runs on recomputed rows, then every position adds its own residual stream, which is never folded.The implementation uses Triton over HuggingFace checkpoints, with embeddings and the language-model head remaining in PyTorch or cuBLAS.Source: B TRITON IMPLEMENTATION · pp. 15–16

C ON THE NUMERICAL DISCREPANCY WITH VANILLA DECODING

The appendix explains why speculative decoding can show small metric discrepancies from vanilla decoding despite theoretical equivalence. It attributes the drift to batched matmul numerics and reports that serialized verification reproduces vanilla metrics.

C ON THE NUMERICAL DISCREPANCY WITH VANILLA DECODING

  • Batched bfloat16 matmuls can change accumulation order when speculative branches alter tensor shape, producing rounding differences that may flip tokens and shift exact-match or pass@1 scores.The appendix says short bfloat16 mantissas amplify these differences over denoising steps.Source: C ON THE NUMERICAL DISCREPANCY WITH… · p. 16
  • Serializing the B+1 branches preserves vanilla matmul shapes and reproduces Vanilla metrics on every tested benchmark, ruling out an algorithmic verification error.Source: C ON THE NUMERICAL DISCREPANCY WITH… · p. 16

D ADDITIONAL ABLATIONS

This section presents additional ablations of the SPECFOLD algorithm and implementation.

D ADDITIONAL ABLATIONS

D.1 DRAFT-BUDGET SWEEP ACROSS MODEL SCALES AND FAMILIES

The draft-budget sweep extends the earlier setup to the four remaining models. Across models and benchmarks, SPECFOLD retains a throughput advantage over Spiffy, although that gain diminishes as draft budgets grow.

D.1 DRAFT-BUDGET SWEEP ACROSS MODEL SCALES AND FAMILIES

  • The sweep extends the Sec. 4.3 draft-budget setup from Nemotron-Labs-Diffusion-3B to the four remaining models in Table 1.Source: §D.1 · p. 16
  • Across models and benchmarks, SPECFOLD sustains a consistent throughput advantage over Spiffy, while increasing draft budgets diminishes the end-to-end gain toward the dense baseline.Draft acceptance saturates as every additional branch adds verification cost.Source: §D.1 · p. 16

D.2 ABLATION ON DRAFTING STRATEGY

SPECFOLD retains throughput gains over Vanilla across draft-graph depths and benchmarks, from seven sibling drafts to deeper chains.

D.2 ABLATION ON DRAFTING STRATEGY

  • At fixed budget B=7, the depth sweep spans seven one-token sibling drafts at D=1 through deep parent-extending chains at D=4.Both methods are compared at every depth.Source: §D.2 · p. 16
  • 1.26–1.54× throughput gain over Vanilla at every depth on every benchmark, with fold ratios of 0.52–0.67 across all shapes.Drafting affects acceptance, but SPECFOLD’s throughput gains remain consistent.Source: §D.2 · pp. 16–17
  • Figure 7 sweeps draft budget on four remaining models across five benchmarks; Spiffy and SPECFOLD share drafts and NFE at each B, so curve gaps reflect folding.Source: Figure 7 · p. 17

D.3 ABLATION ON HARDWARE

On Nemotron-Labs-Diffusion-3B, SPECFOLD’s relative throughput gain is larger on A100 than H100, while task accuracy is preserved.

D.3 ABLATION ON HARDWARE

  • The comparison repeats the Nemotron-Labs-Diffusion-3B configuration on NVIDIA A100 40GB PCIe GPUs against H100 runs.The setup uses recompute threshold δ=0.1, B=8 draft branches, and identical calibration graphs.Source: §D.3 · p. 17
  • On A100, lower compute capacity makes the dense draft batch a larger share of each forward pass, so SPECFOLD recovers more throughput.Source: §D.3 · p. 17
  • Spiffy stays within 1.04–1.23× TPS gain on both GPUs.Source: §D.3 · p. 17
  • 1.43–1.67× relative gain on A100 versus 1.26–1.48× on H100, with task accuracy preserved.Source: §D.3 · p. 17, Table 3 · p. 18
Loading 2610.04875v2…