SYNTOLOGY HomeExplorerAtlasCodeMethodologyAboutDevelopersFeedPricing
Paper · 2606.03819 · 2026

TreeFlash: Parallel AR-Approximation for Faster Speculative Decoding

Roger Wattenhofer, Frédéric Berdoz, Peer Rheinboldt

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 18 functions out of this paper's own repositories and ran 11 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

RepositoryRoleRan
ETH-DISCO/TreeFlash — 11 of 18
FunctionStatusWhere it lives
SwiGLUApproximation Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("5ea2b3d0bd393b88")
_ar_approximation_name Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("445625cc41768e52")
_build_candidate_tree_from_transition_tables Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("2d78b7ae74c67452")
_candidate_tokens Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("39e6760001044741")
_compile_tree_attention_mask Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("1965ea40568355c3")
_compute_tree_acceptance Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("bf25271b68149b27")
_cuda_time Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("dcd98c1c9250dc1f")
_treeflash_config Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("7f2870e4ac60f4c2")
build_target_layer_ids Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("0dc744c2b2fb89e2")
extract_context_feature Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("5d0bf8328d86b399")
sample Ran ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("12b15cc84384fe5e")
Qwen3DFlashAttention Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("6359f88b4c146796")
Qwen3DFlashBackbone Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("a74eed0c1d064f7d")
Qwen3DFlashDecoderLayer Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("1d52ce02d7ac58f2")
TreeFlashDraftModel Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("5a2e1317a43f4b39")
_compact_appended_window Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("20b7d3b6431eaccb")
_compact_dynamic_cache_tail Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("38d45e75bdd70512")
apply_rotary_pos_emb Not yet run ETH-DISCO/TreeFlash/tree_flash.py
pointer only (licence: NONE) · get_code("cce1f5188f2264e4")

Repositories linked to this paper

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

One-shot block drafters for speculative decoding generate the full draft in a single forward pass, achieving strong throughput by eliminating sequential token generation. However, they predict each draft token conditioned only on the prefix context, with no dependence on previously drafted tokens. This non-autoregressive conditioning causes the drafter's distribution to diverge from the verifier's true autoregressive distribution as draft depth grows. This problem becomes more severe in tree-based drafting, where distinct branches are forced to share the same marginal distribution for subsequent tokens. We propose TreeFlash, which addresses this by incorporating an MLP layer conditioned on the drafter's hidden state and the previous token to approximate an autoregressive distribution. TreeFlash retains the O(1) decoding time complexity of one-shot drafters by employing a two-stage approximation mechanism. TreeFlash achieves state-of-the-art performance across a variety of tasks and models, improving over marginal tree drafting by 12% higher block efficiency and 9% higher speedup.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2606.03819")
get_code_for_paper("2606.03819")
have("2606.03819")

Connect an agent — have() is free.