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

Dec 7, 2025

Block Sparse Flash Attention

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.

Daniel Ohayon, Itay Lamprecht, Itay Hubara +3

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

Ask PDF

Aug 25, 2025

FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel

Recent advances in sparse attention mechanisms have demonstrated strong potential for reducing the computational cost of long-context training and inference in large language models (LLMs). Native Sparse Attention (NSA), one state-of-the-art approach, introduces natively trainable, hardware-aligned sparse attention that delivers substantial system-level performance boosts while maintaining accuracy comparable to full attention. However, the kernel implementation of NSA forces a loop order that is only efficient with a relatively large number of query heads in each Grouped Query Attention (GQA) group, whereas existing LLMs widely adopt a much smaller number of query heads in each GQA group -- such an inconsistency significantly limits the applicability of this sparse algorithmic advance. In this work, we propose Flash Sparse Attention (FSA), an alternative kernel implementation that enables efficient NSA computation across a wide range of popular LLMs with a varied, smaller number of heads in each GQA group on modern GPUs. Compared to vanilla NSA kernel implementation, our empirical evaluation demonstrates that FSA achieves (i) up to 3.5x and on average 1.6x kernel-level latency reduction, (ii) up to 1.25x and 1.09x on average end-to-end training speedup on state-of-the-art LLMs, and (iii) up to 1.36x and 1.11x on average for prefill-phase speedup in LLM generative inference. The source code is open-sourced and publicly available at https://github.com/Relaxed-System-Lab/Flash-Sparse-Attention.

Ran Yan, Youhe Jiang, Zhuoming Chen +3

1. The Hong Kong University of Science and Technology · 2. Carnegie Mellon University

GPU Kernel Optimization GPU Acceleration Self-Attention LLM Inference Acceleration Sparse Attention Attention Mechanisms

Ask PDF

Dec 18, 2025

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

Attention is the dominant source of latency during long-context LLM inference, an increasingly popular workload with reasoning models and RAG. We propose Kascade, a training-free sparse attention method that leverages known observations such as 1) post-softmax attention is intrinsically sparse, and 2) the identity of high-weight keys is stable across nearby layers. Kascade computes exact Top-k indices in a small set of anchor layers, then reuses those indices in intermediate reuse layers. The anchor layers are selected algorithmically, via a dynamic-programming objective that maximizes cross-layer similarity over a development set, allowing easy deployment across models. The method incorporates efficient implementation constraints (e.g. tile-level operations), across both prefill and decode attention. The Top-k selection and reuse in Kascade is head-aware and we show in our experiments that this is critical for high accuracy. Kascade achieves up to 4.1x speedup in decode attention and 2.2x speedup in prefill attention over FlashAttention-3 baseline on H100 GPUs while closely matching dense attention accuracy on long-context benchmarks such as LongBench and AIME-24.

Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1

Microsoft Research India.

Self-Attention Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Attention Mechanisms

Ask PDF

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