PaperlayerView exact revision on arXiv

Source-linked AI summary

LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz, Ying Hong Tham, Yuanyi Lin, Wenyi Fang, Fan Wu, Philipp Petersen

arXiv:2609.39361v1cs.LGcs.NAmath.NA

TL;DR

Practical FlashAttention kernels do not exploit the potential to compute most attention logits in low precision, and prior LAMP analyses do not account for these hardware constraints. LampAttention adapts look-ahead mixed precision to the kernel, selectively recomputing sensitive sub-blocks. In simulations, recovered metrics broadly approach 32-bit baselines with a minority of key tiles requiring disruptive recomputation.

  • Problem

    Prior LAMP analyses do not account for practical FlashAttention kernel constraints, and explicit storage of earlier softmax probabilities is inefficient for long contexts.

    Source: §1 · p. 1, §2.4 · p. 4

  • Method

    LampAttention adapts look-ahead precision selection to blockwise computation, recomputing identified sub-blocks, their exponentials, and the running normalization constant.

    Source: §2.4 · p. 4, §3.1 · p. 5

  • Results

    At τ = 0.5, recovered perplexity and downstream metrics broadly approach 32-bit vanilla baselines, with about 20-30% of key tiles disruptive in the compact layout.

    Source: §4.2 · p. 8

  • Takeaways & Limitations

    LampAttention exposes a controllable trade-off between inference efficiency and downstream performance without using 32-bit arithmetic for computationally intensive sections.

    Source: §4.2 · p. 9

    The analysis assumes rounding errors in logits do not change their maximizer and computes α̂ in high precision.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

While most attention logits can be computed in low precision without degrading numerical stability, current attention kernels fail to exploit this phenomenon. We introduce a novel hardware-algorithm co-design in the form of mixed-precision FlashAttention. Our method accumulates key-query products and evaluates their exponentials in 8-bit formats, then adaptively identifies sensitive sub-blocks and recomputes them in 16-bit formats. We propose the specifications for a dedicated accelerator capable of executing this pipeline efficiently. Simulated experiments with Qwen3 and Gemma 3 show that rerouting a selective minority of sub-blocks to high precision is sufficient to recover the baseline model performance.

1 INTRODUCTION

Mixed-precision methods can reduce arithmetic precision, but prior LAMP analyses do not account for practical FlashAttention kernel constraints. LampAttention adapts LAMP to FlashAttention and specifies a dedicated accelerator for executing it.

1 INTRODUCTION

  • Prior LAMP analyses overlook hardware constraints imposed by practical kernels such as FlashAttention.They also empirically validated LAMP on feedforward activations and transformer softmax functions, but their algorithms did not account for those kernel constraints.Source: §1 · p. 1

1.1 CONTRIBUTIONS

  • LampAttention adapts LAMP to FlashAttention, identifying destabilizing logits in two stages using running maxima and an analytical error bound.The method controls precision of intermediate attention-kernel calculations, not model weights or key-query-value operands.Source: §1.1 · p. 2
  • The proposed dedicated accelerator supports native 8-bit accumulation and LUT-based exponential evaluation for LampAttention.The evaluation simulates arithmetic because commercially available hardware lacks native 8-bit accumulation, so it measures high-precision recomputations against task metrics rather than wall-clock timings.Source: §1.1 · p. 2

1.2 RELATED WORK

  • FlashAttention reorganizes attention in an I/O-aware manner to achieve linear memory cost, while LampAttention targets precision within its computations.Other attention approaches described here include sparse patterns, softmax-free linear or log-linear variants, and KV-cache optimizations.Source: §1.2 · p. 2
  • Unlike operand quantization and homogeneous mixed-precision products, LampAttention assigns heterogeneous precision within intermediate attention calculations.It can operate alongside external optimizations such as operand quantization or rotary positional encodings.Source: §1.2 · pp. 2–3, §1.1 · p. 2

2 LOOK-AHEAD MIXED-PRECISION COMPUTATION OF ONLINE SOFTMAX

The section develops a look-ahead mixed-precision method for online softmax: it analyzes rounding-error propagation, then allocates precision across logits while retaining FlashAttention’s running-state computation.

2 LOOK-AHEAD MIXED-PRECISION COMPUTATION OF ONLINE SOFTMAX

  • Online softmax updates its intermediate result as new logits arrive, tracking running values to compute the attention distribution incrementally.The update is expressed as ℓ1 normalization of a vector containing the carried-forward state and the new exponentials.Source: §2 · p. 3, §2.1 · p. 3

2.1 FLOATING-POINT NORMALIZATION

  • The normalization error bound separates rounding in the normalization itself from discrepancies in its evaluated input, whose propagation depends on the exponentials and their arguments.The analysis represents the update as η(z) and bounds its evaluation using the precision of η.Source: §2.1 · p. 3

2.2 FLOATING-POINT INPUTS OF NORMALIZATION

  • The input-error analysis accounts for rounding in logits and the running maximum, assuming rounding does not change the logit maximizer.It also treats the carried-forward exponential term as high precision, a choice the text says adds negligible computational overhead.Source: §2.2 · pp. 3–4

2.3 PROPAGATION OF MIXED-PRECISION ROUNDING ERRORS

  • The probability-mass-shift bound is largest when the running maximum is set to a low-precision logit.The propagation analysis models different precisions across components and bounds the resulting softmax evaluation error.Source: §2.3 · p. 4

2.4 LOOK-AHEAD DISTRIBUTION OF PRECISIONS

  • The method allocates precisions by a greedy knapsack procedure, but avoids exact softmax probabilities and historical probability storage by using a two-step look-ahead based on running state.It selects the smallest weighted probability entries until the threshold is met, then recomputes selected logits and exponentials at higher precision and renormalizes.Source: §2.4 · pp. 4–5

3 LOOK-AHEAD MIXED-PRECISION FLASHATTENTION

LampAttention adapts LAMP to FlashAttention’s blockwise online softmax, identifying and selectively recomputing sensitive sub-blocks. Its accelerator co-design pairs low-precision arithmetic with register layouts intended to retain efficiency when recomputations are sparse.

3.1 BLOCK MODIFICATION OF LAMP COMPUTATION OF ONLINE SOFTMAX

  • The block-LAMP procedure sorts sub-blocks by threat, keeps those satisfying a running-maximum safety margin in low precision, and recomputes the others in high precision.It then selects a precision distribution by comparing feasible configurations of the small sub-block set, a multidimensional knapsack problem without a general greedy solution.Source: §3.1 · p. 5

3.2 LAMPATTENTION

  • LampAttention integrates the modified LAMP procedure into FlashAttention-2, with the remaining computation following vanilla FlashAttention-2.Its mixed-precision steps identify maximum-logit indices, sort sub-blocks, test the safety margin, and update the running normalization quantities.Source: §3.2 · pp. 5–6

3.3 HARDWARE-ALGORITHM CO-DESIGN

  • The proposed accelerator uses e4m3 for low-precision accumulation and ue5m3 for shifted logits and exponentials, enabling one-cycle LUT exponential evaluation.High-precision recomputations use e4m11 accumulation and ue5m11 exponentials, with exponentials routed through the SFU.Source: §3.3 · pp. 6–7
  • The e4m3 format caps absolute pre-softmax logits at 448, and the proposed bound requires headwise QK-Norm with moderate gain weights.For head dimension 128, the stated maximum gain of 6.29 prevents overflow during e4m3 accumulation.Source: §3.3 · pp. 6–7
  • The 8-bit exponential path avoids the multi-cycle SFU, while identifying and recomputing salient sub-blocks adds latency that warp-group specialization could mask.The paper omits that specialization in its implementation.Source: §3.2 · p. 5, §3.3 · p. 7
  • Packing four 8-bit accumulators into one 32-bit register reduces accumulator footprint by 75%, enabling larger tiles and higher warp occupancy to hide memory-transfer latency when disruptive recomputations of two or more sub-blocks remain infrequent.The compact layout reserves registers for the entire low-precision Y and one high-precision sub-block; two or more sub-blocks marked for recomputation trigger a sequential, fully 16-bit update of Ow.Source: §3.3 · pp. 7–8

4 NUMERICAL EXPERIMENTS

Across C4 perplexity and MMLU downstream tasks, adaptive 16-bit recomputations improve on the pure 8-bit baseline and broadly approach the 32-bit baseline, with model- and task-dependent closeness. The experiments use simulated arithmetic and examine the accuracy-efficiency trade-off as the recomputation threshold varies.

4.1 SETUP AND IMPLEMENTATION

  • The experiments evaluate Qwen3 and Gemma 3 using a Triton kernel with simulated 8-bit and 16-bit arithmetic and compare against 32-bit and pure 8-bit baselines.Task metrics use the EleutherAI lm-evaluation-harness; the compared bit widths apply to attention-logit accumulators and exponentials, not model weights.Source: §4.1 · p. 8
  • The experiments ran on a cloud-based GPU equipped with 80 GB of VRAM and required approximately 540 GPU-hours.Source: §4.1 · p. 8

4.2 NUMERICAL RESULTS

  • About 20–30% of key tiles are disruptive in the compact layout, requiring recomputation of two or more sub-blocks.This is the reported share associated with the recovery at τ = 0.5.Source: §4.2 · p. 8
  • Adaptive 16-bit recomputations substantially improve on the pure 8-bit baseline across C4 perplexity and MMLU tasks, broadly approaching the 32-bit baseline.The degree of closeness varies by model and task, even at τ = 0.5.Source: Table 1 · p. 8, §4.2 · p. 8, Table 2 · p. 9
  • As τ decreases, performance converges toward the 32-bit baseline non-monotonically, while the disruptive-recomputation rate increases monotonically.The τ = 1 and τ = 0.5 settings show only marginal differences in recomputation rate and accuracy.Source: §4.2 · pp. 8–9
  • The running-maximum stage moves metrics toward the 32-bit baseline with a moderate disruptive-recomputation budget; the second stage provides finer correction at additional recomputation cost.At τ = 1, the second-stage optimization omits recomputations.Source: §4.2 · p. 9

AI USE STATEMENT

The authors report using generative AI tools for software development and text refinement, while stating that the core research work was conducted by the authors.

AI USE STATEMENT

  • Generative AI tools assisted with software development.The authors state that they reviewed all AI-assisted work.Source: AI USE STATEMENT · p. 10
  • Generative AI tools helped refine the clarity and flow of the text.Source: AI USE STATEMENT · p. 10
  • The authors state that they conducted all core research ideation, theoretical development, experimental design, and data analysis.They assume full responsibility for the work’s accuracy, integrity, and originality.Source: AI USE STATEMENT · p. 10

A PROOF OF PROPOSITION 1

The proof passage states differentiability conditions for the ℓ1 normalization function and a nonnegativity property of a matrix’s diagonal entries.

A PROOF OF PROPOSITION 1

B ADDITIONAL BENCHMARKS

Additional evaluations cover Wikitext and ARC-Challenge; on ARC-Challenge, Qwen3-32B’s 8-bit accuracy baseline exceeds its 32-bit baseline. Figure 3 extends the comparison across all models and tasks.

B ADDITIONAL BENCHMARKS

C IMPACT OF SAFETY MARGIN

A larger first-stage safety margin initiates more disruptive recomputations when τ = 1, while the trade-off curves show similar convergence as τ tightens.

C IMPACT OF SAFETY MARGIN

D IMPACT OF FIRST-STAGE RECOMPUTATION

Disabling first-stage recomputation violates the assumptions used to derive the block-LAMP problem, yet the comparisons suggest either recomputation phase can identify important sub-blocks.

D IMPACT OF FIRST-STAGE RECOMPUTATION

E IMPACT OF IN-CONTEXT LEARNING

On ARC-Challenge, 25-shot evaluation increases accuracy and the share of tiles processed entirely in 8-bit relative to 0-shot evaluation; the recomputation distribution also differs between the two settings.

E IMPACT OF IN-CONTEXT LEARNING

Related papers

Loading 2609.39361v1…