Attention gets cheaper by matching precision to each region
TriAttention: Unified Sparse Mixed-Precision Attention for Long-Context LLM Inference · arXiv:2611.07256
In plain words
A language model uses attention to work out which parts of its input matter as it produces an answer. Some parts can be skipped or handled with less precise arithmetic, while others need more care, so one uniform rule can waste computation or lose useful information.
TriAttention chooses how to handle each region and reuses a compressed prompt cache during generation. In one Llama request on an H100 NVL, with an 8K prompt, 128 output tokens, and batch size four, it generated tokens faster than BF16 attention over the full request, including prompt processing.
The research
Long prompts make attention costly, yet removing or simplifying the wrong regions can hurt a model’s output.
The result changes the relevant question from how much attention to compress overall to where computation can be reduced safely: TriAttention’s measurements suggest region-specific choices can improve speed while keeping tested prediction quality close to the BF16 reference.
TriAttention is a sparse mixed-precision attention method that assigns each head’s regions to skipping, 4-bit, or 8-bit computation and reuses those choices across prompt processing and generation.
On Llama-3.2-3B-Instruct, TriAttention processed an 8K prompt and generated 128 tokens at batch size four on an H100 NVL, reaching 318.36 tokens per second for the complete request, including prompt processing; BF16 reached 223.61. On WikiText-2, a text-prediction test, word perplexity was 12.50 versus 12.49 for BF16.
The request comparison shows an inference speed gain in that workload, while near-matching perplexity provides a separate quality check. A fixed-cost placement comparison also found stronger long-context question answering when boundary choices followed head sensitivity.
The speed result is for one Llama request shape. The evaluation covers three model checkpoints, selected general and long-context question-answering tasks, and A100, H100 NVL, and RTX 4090 GPUs.
What the evidence shows
- Regions need different treatment
TriAttention is a sparse mixed-precision attention method for language-model inference. It addresses a mismatch: some attention regions can be skipped or computed at lower precision with little output error, while other regions are more sensitive.
Applying one retention range or precision everywhere can therefore spend computation unnecessarily or damage useful contributions. The method’s central claim is that choosing both retention and precision by region can use computation more effectively.
- One budget guides head choices
TriAttention calibrates candidate choices for each layer and query head, measuring error after the attention output projection and comparing it with the full computation cost of retained tiles. A model-wide allocator selects boundary pairs under a shared computation budget, then separate prompts check the assembled model against quality thresholds.
During inference, those boundaries drive both prompt processing and a compressed cache reused during generation.
- A faster complete request
On this request, TriAttention completes generation faster than either reference while its WikiText-2 prediction score remains close to BF16. Together, the comparisons show that the measured speed gain did not require a large loss on this text-prediction test.
- Prediction quality stays close
On WikiText-2, a test of text prediction, word perplexity was 12.50 with TriAttention and 12.49 with BF16; lower perplexity means better prediction. The near-match supports the claim that the measured speed gain can coexist with near-reference quality on this test.
It does not by itself establish unchanged quality across other tasks, which the paper evaluates separately.
- Placement matters at fixed cost
Because the Random control keeps the same boundary choices and modeled cost but assigns them to different heads, its lower score indicates that placement mattered in this Llama evaluation. This supports sensitivity-guided allocation for long-context question answering in this tested setting.
Deep
- Original paper authors
- Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang
- Submitted to arXiv
- License
- CC BY 4.0