LazyFlash: Forward-Dense, Backward-Sparse Attention for Accelerated Long-Context Fine-Tuning

Keqi Deng; Ruchao Fan; Jinyu Li

cs.CLOct 22, 2026

LazyFlash: Forward-Dense, Backward-Sparse Attention for Accelerated Long-Context Fine-Tuning

Authors: Keqi Deng, Ruchao Fan, Jinyu Li

Organizations: Microsoft, USA

Abstract

As large language models (LLMs) process increasingly long sequences, attention becomes a major computational bottleneck. Sparse attention mostly targets inference, yet many methods still train densely or cannot support sparse training, and their added complexity makes sparsity hard to turn into wall-clock speedups, leaving dense attention widely used in practice. This paper introduces LazyFlash, a forward-dense, backward-sparse attention method built on the observation that attention training is dominated by the backward pass. LazyFlash keeps the forward pass exactly dense, so the trained model remains a standard dense model that existing systems such as vLLM can serve without any change. For the backward pass, LazyFlash proposes a gradient-aware block sparsification design that ranks query blocks by their attention-output gradient norms and skips those contributing negligibly to parameter updates, turning the dense backward into a block-sparse one. A gradient-mass budget adapts the amount of skipped computation to each head and training step, while importance sampling preserves the dense gradient in expectation. These designs are implemented inside the tiled FlashAttention backward kernel, where skipped blocks are never loaded or computed, translating gradient sparsity into real wall-clock acceleration. Experiments on long-context text and speech LLMs show that LazyFlash maintains downstream performance while achieving up to 5.7x attention-backward speedup over FlashAttention and 1.7x end-to-end training speedup.

Figures & tables

Figure 1: Overview of LazyFlash. The forward pass computes every block, so the trained model remains an ordinary dense-attention model. In the backward pass, each query block is kept with a probability set by its output-gradient norm: skipped tiles are never loaded, and kept blocks are reweighted so that the gradient stays unbiased. As the bars illustrate, a small budget on gradient mass translates into a much larger saving in computation.

Method	Overall	Language		Difficulty				Training attn. (ms)
Method	Overall	En	Zh	Extreme	Hard	Moderate	Easy	fwd	bwd
Qwen3-8B-Base (no FT)	19.70	19.35	19.78	19.39	16.89	13.98	24.67	–	–
Dense attention	29.64	30.31	28.63	25.18	25.40	19.54	41.78	547.3	1429.8
MoBA								700.3	710.2
sparse inference	24.48	24.73	24.23	20.04	21.71	15.36	35.65
sparse prefill	23.64	24.99	22.30	20.27	19.21	14.46	34.89

Table 1: LongBench-Pro results after fine-tuning Qwen3-8B-Base (higher is better,

\uparrow

), with the attention time per microbatch measured during training. The three indented rows are one MoBA-trained model served in three ways: sparse attention at prefill and decode, sparse attention at prefill only, and fully dense attention. They share the training cost listed on the MoBA row.

Method	EER				WER				ALL AVG
Method	meeting	banking	medical	AVG	conversation	meeting	dictation	AVG	ALL AVG
Dense attention	28.43	8.72	7.25	14.80	21.53	19.61	11.42	17.52	16.16
MoBA (dense infer)	26.23	12.66	8.40	15.76	26.69	19.06	10.97	18.91	17.34
LazyFlash	23.04	8.66	6.61	12.77	21.35	19.01	10.71	17.02	14.90

Table 2: Long-form English ASR results on 10-minute utterances (lower is better,

\downarrow

). MoBA is shown with dense inference only, the configuration under which it performed best on text in Table 1 .

Figure 2: Attention time and end-to-end training step time against context length, measured on real LongAlign examples truncated to each length. Forward and backward are shown separately, and the arrows give the speedup over dense attention.

Selection rule	Overall	Language		Difficulty				Attn bwd
Selection rule	Overall	En	Zh	Extreme	Hard	Moderate	Easy	(ms)
Dense attention (no selection)	29.64	30.31	28.63	25.18	25.40	19.54	41.78	1429.8
Random	21.18	21.56	20.81	19.84	16.03	13.87	29.89	523.8
Top- $k$ , attention mass	27.07	27.96	26.19	23.02	24.62	17.50	37.98	548.0
Top- $k$ , gradient norm	28.45	29.84	27.07	23.87	25.20	19.16	40.16	508.0
LazyFlash	30.24	30.37	29.03	25.22	27.51	20.04	40.88	507.4

Table 3: Selection rules compared on LongBench-Pro. Random drops query blocks with no signal, the two top-

k

rules keep a fixed fraction ranked by forward attention mass and by output-gradient norm, and LazyFlash ranks by the gradient norm but sets the amount skipped with a budget on gradient mass. The three alternatives share the CUDA implementation of LazyFlash and are tuned to match its backward time; dense attention is a quality reference and is not matched in cost.

Setting	EER				WER				ALL AVG
Setting	meeting	banking	medical	AVG	conversation	meeting	dictation	AVG	ALL AVG
Dense attention	28.43	8.72	7.25	14.80	21.53	19.61	11.42	17.52	16.16
LazyFlash ( $\delta=0.1$ )	28.43	8.22	7.40	14.68	20.48	19.71	11.05	17.08	15.88
LazyFlash ( $\delta=0.3$ )	26.96	8.94	6.76	14.22	21.49	19.22	10.73	17.15	15.69
LazyFlash ( $\delta=0.4$ )	23.04	8.66	6.61	12.77	21.35	19.01	10.71	17.02	14.90

Table 4: Effect of the budget

\delta

on long-form ASR (lower is better,

\downarrow

). Training diverges at

\delta=0.5

Method	Overall	Language		Difficulty				Training attn. (ms)
Method	Overall	En	Zh	Extreme	Hard	Moderate	Easy	fwd	bwd
Qwen3-4B-Base
before fine-tuning	6.88	5.84	7.92	11.09	4.91	4.08	5.89	–	–
Dense attention	24.79	23.33	26.26	22.32	18.56	15.92	36.11	534.4	1427.0
LazyFlash	26.55	27.18	25.92	22.94	23.33	18.34	36.72	544.0	500.7
Qwen3-8B-Base

Table 5: LongBench-Pro accuracy and per-step attention time across backbone sizes. Qwen3-4B and Qwen3-8B have the same attention configuration, and their dense attention times are nearly identical. LazyFlash wins 19 of the 21 accuracy splits.

Sparse Fine-Tuning

Sparse fine-tuning adapts pretrained models by updating selected parameters, reducing training costs and limiting interference with existing capabilities.

See all 27 papers

Block Sparse Flash Attention
Dec 7, 2025Daniel Ohayon, Itay Lamprecht, Itay Hubara +3GPU Acceleration Self-Attention
Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel
FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel
Aug 25, 2025Ran Yan, Youhe Jiang, Zhuoming Chen +3GPU Kernel Optimization GPU Acceleration
1. The Hong Kong University of Science and Technology · 2. Carnegie Mellon University
Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
Dec 18, 2025Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1Self-Attention Long-Context Language Model Inference
Microsoft Research India.

Dec 7, 2025J/K move · Enter open · S save

Block Sparse Flash Attention

Daniel Ohayon, Itay Lamprecht, Itay Hubara, Israel Cohen+2

Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel

GPU Acceleration Self-Attention Long-Context Language Model Inference High-Performance Computing LLM Inference Acceleration Sparse Attention+3

Open paper AskPDF arXiv

Modern large language models increasingly require long contexts for reasoning and multi-document tasks, but attention's quadratic complexity creates a severe computational bottleneck. We present Block Sparse Flash Attention (BSFA), a drop-in replacement that accelerates long-context inference while preserving model quality. Unlike methods that predict importance before computing scores, BSFA computes exact query-key similarities to select the top-k most important value blocks for each query. By comparing per-block maximum scores against calibrated thresholds, we skip approximately 50% of the computation and memory transfers for pruned blocks. Our training-free approach requires only a one-time threshold calibration on a small dataset to learn the per-layer and per-head attention score distributions. We provide a CUDA kernel implementation that can be used as a drop-in replacement for FlashAttention. On Llama-3.1-8B, BSFA achieves up to 1.13x end-to-end speedup on LongBench with only a 1.1% accuracy drop, and up to 1.24x on Needle-in-a-Haystack retrieval at a 1% accuracy drop. The attention kernel itself accelerates by up to 1.38x. We compare BSFA against five recent sparse attention baselines (SpargeAttention, MInference, FlexPrefill, XAttention, and BLASST), and verify the method on Qwen2.5-7B and on A6000 and H100 GPUs. The implementation is available at https://github.com/Danielohayon/Block-Sparse-Flash-Attention.

Abstract

Sparse Fine-Tuning

Explore similar work

Block Sparse Flash Attention

FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel

Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference