Research map / Architectures and efficiency / Attention and transformers
DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DeFT, a prefix-aware, load-balanced tree attention kernel (KV-guided grouping plus flattened tree KV splitting) that cuts KV-cache IO for tree-structured LLM decoding, implemented in Triton.
Problem
Attention for tree-structured LLM workloads (few-shot prompting, multi-step reasoning, speculative decoding) repeatedly loads shared-prefix KV cache and suffers poor GPU load balancing, causing redundant IO and low utilization.
Approach
Group queries by KV node (KV-guided grouping) so shared-prefix KV is loaded once, and flatten the tree KV into evenly sized blocks for balanced partitions, using bit causal masks and tree-topology-aware global reduction.
Key results
- DeFT-Flatten achieves up to 2.23x decoding speedup and 3.59x attention speedup over Radix Attention across the tree workloads.
- On the sorting reasoning task, DeFT-Flatten reduces decoding latency from 104.79 s (Radix Attention) to 94.67 s, and it reduces KV-cache IO by 73-99%.
- Ablation: DeFT-Flatten outperforms DeFT-Node and DeFT-Node-Chunk across all tree settings, e.g. 5.87 s vs 10.85 s for few-shot b=50 attention latency.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Radix Attention with DeFT-Flatten | Llama3-8B tree-based decoding · Few-shot prompting | decoding latency: 12.37 -> 9.98; 14.08 -> 10.99; 16.54 -> 12.48 | Helped |
| Replace Radix Attention with DeFT-Flatten | Llama3-8B tree-based decoding · Multi-step reasoning | decoding latency: 104.79 -> 94.67; 69.61 -> 66.95; 11.25 -> 10.90; 17.03 -> 16.10 | Helped |
| Replace Radix Attention with DeFT-Flatten | Llama3-8B tree-based decoding · Speculative decoding | decoding latency: 54.66 -> 42.23; 69.75 -> 46.60; 108.56 -> 56.96; 188.66 -> 84.27 | Helped |
| Add flattened tree KV splitting instead of node-wise KV splitting (ablation) | DeFT tree attention · Few-shot prompting b=20 | attention latency: 10.59 -> 3.47 | Helped |
| Add flattened tree KV splitting instead of node-chunk splitting (ablation) | DeFT tree attention · Multi-step reasoning Sorting | attention latency: 49.63 -> 28.41 | Helped |
| Add flattened tree KV splitting instead of node-wise KV splitting (ablation) | DeFT tree attention · Multi-step reasoning Keyword | attention latency: 6.16 -> 2.57 | Helped |
| Add flattened tree KV splitting instead of node-chunk splitting (ablation) | DeFT tree attention · Speculative decoding t=64 | attention latency: 20.28 -> 16.79 | Helped |
| Add flattened tree KV splitting instead of node-wise KV splitting (ablation) | DeFT tree attention · Speculative decoding t=256 | attention latency: 41.78 -> 40.56 | Helped |
| Use node-chunk splitting instead of node-wise splitting (ablation) | DeFT tree attention · Few-shot prompting b=20 | attention latency: 10.59 -> 8.52 | Helped |
| Use node-chunk splitting instead of node-wise splitting (ablation) | DeFT tree attention · Speculative decoding t=256 | attention latency: 41.78 -> 57.26 | Hurt |
| Increase prompt length while using DeFT-Flatten (ablation) | Llama3-8B multi-step reasoning Sorting · Sorting, prompt length L=1k to L=10k | decoding speedup over Radix Attention: 1.09× -> 1.67× | Helped |
| Use Codellama-34B instead of Codellama-7B (ablation) | DeFT-Flatten relative to Radix Attention · Few-shot prompting b=30 | decoding time speedup: 1.34× -> 1.23× | Hurt |
Surprising findings
- Multi-step reasoning sees only modest decoding gains (up to 1.10x) because narrow trees limit KV reuse and attention is only ~30% of latency.
- Attention becomes the bottleneck (51-58%) only after paged memory is used; with unpaged memory, KV materialization dominates (69-83%).
Limitations
- Multi-step reasoning speedups are modest for narrow tree widths.
- Speedups are measured on decoding latency excluding prefill and framework overheads, not full end-to-end serving.
Experimental setup
Datasets: APPS Graph-of-Thoughts multi-step reasoning tasks (Sorting, Document, Keyword, Set) Medusa speculative decoding trees
Benchmarks: few-shot prompting (APPS, 4k prompt) multi-step reasoning (Sorting, Document, Keyword, Set) speculative decoding (APPS with Medusa trees)
Metrics: decoding latency (s) attention latency (s) IO (TB) speedup
Baselines: Flash-Decoding Tree Attention-Medusa Radix Attention DeFT-Node DeFT-Node-Chunk
Builds on: FlashAttention Flash-Decoding Medusa SpecInfer tree attention Triton
Closest papers by meaning
- ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference (ICML 2025 · similarity 0.90)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.89)
- ArborKV: Structure-Aware KV Cache Management for Scaling Tree-based LLM Reasoning (ICML 2026 · similarity 0.89)
- Yggdrasil: Bridging Dynamic Speculation and Static Runtime for Latency-Optimal Tree-Based LLM Decoding (NeurIPS 2025 · similarity 0.89)
- MagicPIG: LSH Sampling for Efficient LLM Generation (ICLR 2025 · similarity 0.89)
- Break the Sequential Dependency of LLM Inference Using Lookahead Decoding (ICML 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.