vllm

TileSparse

Arithmetic-Intensity-Aware Sparse Attention for Compute-Bound LLM Decoding

Chao Wang, Pengfei Zuo, Zhangyu Chen, Qihui Zhou, Tsung-Yi Ho, Ming-Chang Yang · ICML 2026

TileSparse is the reference implementation for the ICML 2026 paper of the same name. It is a research fork of vLLM (based on v0.13.0) that adds an arithmetic-intensity-aware sparse-attention backend for Multi-head Latent Attention (MLA) models such as DeepSeek-V2/V3.

Motivation

Conventional sparse attention assumes decode attention is memory-bound and therefore only prunes along the key dimension (K-only / "token-retrieval" sparsity), treating the per-query-head compute as free. MLA and Multi-Token Prediction (MTP) break that assumption: they raise the attention arithmetic intensity above the hardware ridge point, pushing decoding into a compute-bound regime where computing every query head for each retrieved key wastes significant compute.

TileSparse targets this regime with three ideas:

• QK 2D Sparsity — prune unnecessary (query-head, key-page) pairs, not just keys, and spend the freed compute budget on retrieving more semantically important pages.
• Tiered QK 2D Sparsity — a Critical Tier (α pages computed with all query heads) protects the most important tokens, while an Ambient Tier (β pages computed with only h heads) covers broad context cheaply.
• AutoTuner — an online search that picks the best tier/head configuration per layer by maximizing Estimated Attention Mass Coverage.

Under a 256-token budget the paper reports +36.4% (RULER) / +7.7% (LongBench) accuracy over the state-of-the-art K-only method, and preserves 99% of full-attention accuracy while cutting attention compute by 40.8% on LongBench.

Two implementations, by design

The paper's accuracy results and latency results come from two different kernels, and this repository keeps both:

Path	Where	Sparsity	Used for
Accuracy	vllm/attention/ops/tilesparse/ (the TILESPARSE_MLA serving backend)	mask-based (computes full QK, masks unselected entries — does not save compute)	RULER / LongBench accuracy (Fig. 5–8, 10)
Performance	evaluations/attention_latency/utils/attention_ops/	page-table-based (a true per-head page-list Flash-decode kernel)	latency (Fig. 4, 9)

The serving backend reproduces the paper's accuracy numbers exactly; it does not itself accelerate serving. Integrating the efficient kernel into the serving path is on the roadmap.

Quick start

Environment: A100/H200-class GPU, CUDA 12.x, Python 3.12. The build uses vLLM's precompiled extension (nvcc is not required), so the interpreter's torch must match the precompiled .so (torch 2.9.x).

# uv-based venv + precompiled vLLM + evaluation deps
bash scripts/environment/init_venv.sh
source .venv/bin/activate

The backend is selected with VLLM_ATTENTION_BACKEND=TILESPARSE_MLA plus enforce_eager=True, and is configured through a small JSON file that the backend hot-reloads (SPARSE_CONFIG_PATH, default .local/sparse_config.json). The easiest way to drive it is the shared harness in evaluations/llm_generator/; a minimal sanity check on DeepSeek-V3-0324 is:

bash scripts/evaluations/longbench_dsv3.sh

The backend currently supports DeepSeek V2/V3 only: it reads a per-layer attention context that deepseek_v2.py injects. Selecting it for another MLA model raises a clear error.

Reproducing the paper

See evaluations/README.md for the full directory-to-figure map, reproduction commands, and hardware requirements.

Configuration

The sparse method maps onto the paper's symbols as follows (details in evaluations/llm_generator/generation_config.py and vllm/attention/ops/tilesparse/sparse_config.py):

Field	Paper	Meaning
token_budget	compute budget B	total compute-entry budget
key_top_count	page budget N_p	pages kept overall
critical_top_count	Critical Tier (α)	pages using all query heads
ambient_head_count	Ambient Tier heads h	heads used for the remaining pages (h ≥ τ_hw·m_kv/2, e.g. 78 on A100 MLA)
dynamic_sparse_*	Tiered QK 2D + AutoTuner	online tier/head search

The decode kernel is chosen with TILESPARSE_MODE (default full): tilesparse (the paper's QK 2D sparse method), streaming (StreamingLLM baseline), torch (pure-PyTorch reference), plain (non-fused step-by-step), or full (fused Flash-decode). Other optional flags: TILESPARSE_TIMING=1 prints per-stage decode timing (off by default); ATTENTION_DUMP=1 dumps per-layer attention scores (and forces plain).

Relationship to vLLM

This is a fork of vLLM v0.13.0. The TileSparse additions are confined to the sparse-attention backend (vllm/attention/ops/tilesparse/, vllm/v1/attention/backends/mla/tilesparse_mla.py, vllm/attention/attention_dump_context.py) plus small registration hooks in vllm/platforms/cuda.py, vllm/attention/backends/registry.py, and vllm/model_executor/models/deepseek_v2.py. Upstream vLLM documentation applies to everything else.

Citation

@inproceedings{wang2026tilesparse,
  title     = {TileSparse: Arithmetic-Intensity-Aware Sparse Attention for Compute-Bound {LLM} Decoding},
  author    = {Chao Wang and Pengfei Zuo and Zhangyu Chen and Qihui Zhou and Tsung-Yi Ho and Ming-Chang Yang},
  booktitle = {Forty-third International Conference on Machine Learning},
  year      = {2026},
  url       = {https://openreview.net/forum?id=lK2AzLDJal},
}

License

Apache-2.0, inherited from vLLM. Vendored evaluation code retains its original license (RULER: Apache-2.0; LongBench: MIT).