PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Keqi Deng, Ruchao Fan, Jinyu Li

arXiv:2610.31911v1cs.CL

TL;DR

Long-context attention is costly, while many sparse-attention methods focus on inference or do not accelerate training of standard attention. LazyFlash keeps the forward pass exact and sparsifies backward computation with gradient-aware block selection and unbiased sampling. It matches dense-training quality on text and speech while reducing attention-backward and end-to-end training time.

  • Problem

    Long-context attention is a major training cost, but most sparse-attention methods target inference rather than accelerating training of standard attention.

    Source: §1 · p. 1

  • Method

    LazyFlash keeps attention forward computation dense and exact, while a gradient-mass budget and importance sampling sparsify the backward pass over query blocks.

    Source: §1 · p. 2, §3.2 · p. 4, §3.3 · p. 5

  • Results

    LazyFlash matches dense-training quality on text and speech LLMs while achieving up to 5.7× faster attention backward and 1.7× faster end-to-end training.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The trained model remains a standard dense-attention model that existing inference frameworks can serve without deployment changes.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-context attention is costly because its computation grows quadratically, yet sparse-attention methods largely target inference or change standard attention. LazyFlash keeps the forward pass exact and sparsifies the backward pass, matching dense-training quality while accelerating training.

Motivation

  • Long-context attention remains a training bottleneck because its quadratic computation grows rapidly with sequence length despite memory-optimized kernels.Doubling context length quadruples leading attention computation; FlashAttention reduces memory traffic but not quadratic FLOPs.Source: §1 · p. 1

Research gap

  • Most sparse-attention methods target inference, while several trainable approaches require dense targets, alter standard attention, or add selection overhead.These constraints leave acceleration of standard-attention training an open challenge.Source: §1 · pp. 1–2

Contributions

  • Because the forward pass stays exact, the trained model remains a standard dense model requiring no inference changes.It can be served by existing inference frameworks.Source: §1 · p. 2
  • LazyFlash scores query blocks by output-gradient norm and uses a gradient-mass budget to adapt which blocks are skipped.The budget replaces a fixed skip threshold with one based on the gradient distribution.Source: §1 · p. 2
  • Importance sampling keeps blocks in proportion to their gradient contribution and rescales them to preserve an unbiased gradient estimate.The method is implemented in a FlashAttention backward kernel that supports dense and variable-length inputs.Source: §1 · p. 2
  • Up to 5.7× faster attention backward and 1.7× faster end-to-end training are reported for LazyFlash on text and speech LLMs.The forward pass remains exact, and the paper reports quality matching dense training.Source: §1 · p. 2

2 RELATED WORK

Efficient attention kernels reduce memory traffic but retain quadratic arithmetic, whereas sparse-attention approaches generally target inference or modify the forward pass. LazyFlash complements these lines of work by sparsifying the backward pass while keeping forward attention exact.

2.1 EFFICIENT ATTENTION KERNELS

  • FlashAttention-family kernels reduce memory traffic with tiled exact attention, but their arithmetic still grows quadratically with sequence length.LazyFlash complements them by removing backward computation rather than simply executing the same computation more efficiently.Source: §2.1 · pp. 2–3

2.2 SPARSE ATTENTION

  • Existing sparse-attention families target decoding or change the forward pass, while prior backward sparsification for linear layers does not directly transfer to tiled attention.The paper frames sparse attention backward computation with an exact forward pass as an unresolved challenge.Source: §2.2 · p. 3

3 METHODOLOGY

LazyFlash exploits the greater cost and uneven gradient contribution of attention’s backward pass by removing selected query-block computations while retaining exact forward attention. A gradient-mass budget controls sparsity, importance sampling keeps the gradient estimate unbiased, and the kernel skips zero-weight tiles before loading them.

3.1 THE ATTENTION BACKWARD PASS

  • At long contexts, the backward pass performs roughly twice the forward-pass work and dominates attention training.FlashAttention’s tiled computation enables removal at query-block granularity; blocks with small output gradients contribute little to dQ, dK, and dV.Source: §3.1 · pp. 3–4

3.2 GRADIENT-AWARE BLOCK SELECTION

  • The gradient-mass budget limits discarded gradient energy per head to a fraction δ, adapting the number of skipped blocks to each head’s gradient distribution.Because supervised fine-tuning scores are nonuniform, a small gradient-mass budget can permit a much larger reduction in computation.Source: §3.2 · pp. 4–5, Figure 1 · p. 4

3.3 UNBIASED IMPORTANCE SAMPLING

  • Blocks are sampled in proportion to their gradient contribution and reweighted, yielding an unbiased gradient estimate while concentrating randomness in the low-gradient tail.Large-gradient blocks are always kept; proportional sampling also limits the magnitude of kept contributions in the tail.Source: §3.3 · p. 5

3.4 KERNEL REALIZATION

  • A compact CUDA pipeline scores, budgets, and samples blocks, then the backward kernel skips zero-weight tiles before loading their data.The forward kernel is untouched, and packed variable-length inputs are supported by padding subsequences to tile boundaries and removing padding afterward.Source: §3.4 · p. 5

4 EXPERIMENTS

Across text and speech benchmarks, LazyFlash matches or exceeds dense attention while reducing attention-backward time; its end-to-end speedup grows with context length. Ablations show that gradient-based selection and a workload-adaptive budget are important to the reported quality, and the gains persist across model sizes.

4.2 LONG-CONTEXT TEXT UNDERSTANDING

  • 30.24 vs 29.64 on LongBench-Pro overall, LazyFlash’s score exceeds dense attention after fine-tuning Qwen3-8B-Base.It leads on five of seven columns, with 27.51 vs 25.40 on the Hard split.Source: Table 1 · p. 6, §4.2 · p. 6

4.1 SETUP

  • The text and speech experiments use a shared attention implementation, with Qwen3-8B backbones and a Conformer speech encoder for ASR.The setup compares against FlashAttention-2 dense attention and MoBA under shared backbones, data order, and training steps.Source: §4.1 · pp. 5–6

4.3 LONG-FORM SPEECH RECOGNITION

  • 14.90 vs 16.16 macro-average and 12.77 vs 14.80 EER average, LazyFlash outperforms dense attention across all six long-form ASR test-set columns.The largest reported gain is online-meeting EER, 23.04 vs 28.43.Source: Table 2 · p. 7, §4.3 · p. 7

4.4 TRAINING EFFICIENCY

  • 1.66× end-to-end speedup at 32K tokens; LazyFlash’s attention-backward speedup rises from 2.77× at 8K to 5.72× at 32K.The forward pass matches dense attention, and the computed-block share falls from 18% at 8K to 10% at 32K at the same budget.Source: Figure 2 · p. 7, §4.4 · p. 8

4.5 ABLATIONS

  • Gradient-norm ranking with an adaptive gradient-mass budget is the only tested selection rule that matches dense training quality.At similar backward time, fixed-fraction gradient-norm top-k scores 28.45 vs LazyFlash’s 30.24; random selection scores 21.18.Source: Table 3 · p. 8, §4.5 · p. 8

4.6 MODEL SIZE SCALING

  • LazyFlash remains ahead at Qwen3-4B, 8B, and 14B by +1.76, +0.60, and +2.43 overall, respectively.It wins 19 of 21 splits and all seven at 14B; attention-backward speedups are 2.85×, 2.82×, and 2.93×.Source: Table 5 · p. 9, §4.6 · p. 9

5 CONCLUSIONS

LazyFlash keeps attention forward computation exact and sparsifies only backward computation, matching dense attention on text and speech tasks while accelerating training. The paper also reports that its text experiments use public artifacts for reproducibility.

5 CONCLUSIONS

  • At 32K context, LazyFlash reaches up to 5.72× attention-backward and 1.66× end-to-end model training speedup.The conclusion also reports 2.82× and 2.45× attention-backward speedups on LongBench-Pro and long-form ASR, respectively.Source: §5 · p. 9
  • LazyFlash matches dense attention on LongBench-Pro across three Qwen3 backbones and on long-form ASR while cutting attention-backward time.Its exact forward pass leaves the trained model ordinary dense attention that standard inference stacks can serve.Source: §5 · p. 9
  • The authors report using generative AI tools for language polishing and grammar correction, with edits reviewed and verified by the authors.They state that AI tools were not used for research ideation, data collection, model design, or result analysis.Source: AI USE STATEMENT · p. 10
  • The text experiments use public Qwen3 checkpoints, LongAlign-10k, ShareGPT, and LongBench-Pro, supporting reproduction from public artifacts.Section 3 specifies selection and sampling, Appendix B provides pseudocode and kernel details, and Appendix A lists training configuration.Source: REPRODUCIBILITY STATEMENT · p. 10

A HYPER-PARAMETERS

LazyFlash uses 128-token query-row and key blocks, with the gradient-mass budget δ as its only method hyper-parameter. Text and speech experiments use different training configurations and budgets.

A HYPER-PARAMETERS

  • Across experiments, LazyFlash uses 128-token query-row blocks against 128-token key blocks, and δ is its only method hyper-parameter.The forward pass is bit-exact dense in every run.Source: A HYPER-PARAMETERS · p. 12
  • Text fine-tuning trains Qwen3-4B, 8B, and 14B models for 2,500 steps at maximum sequence length 65,536, using δ = 0.1.Training uses a 1:1 LongAlign-10k and ShareGPT mixture, eight H200 GPUs, and a global batch of 32 sequences.Source: A HYPER-PARAMETERS · p. 12
  • Speech fine-tuning trains the encoder, adapter, and LoRA components for 10,000 steps on English long-form data, using δ = 0.4.The Qwen3-8B backbone and its final norm, embeddings, and output head remain frozen; utterances span 60 to 1,800 seconds.Source: A HYPER-PARAMETERS · p. 12

B IMPLEMENTATION AND REPRODUCIBILITY DETAILS

LazyFlash selects query blocks per batch and head using output-gradient norms and a gradient-mass budget, then applies sampled block weights in the FlashAttention backward kernel. Timing measurements isolate attention backward on H200 GPUs.

B IMPLEMENTATION AND REPRODUCIBILITY DETAILS

  • For each batch and head, Algorithm 1 scores each of M query-row blocks by the Frobenius norm of its output gradient and determines a drop count under the gradient-mass budget.The selection runs once per backward call; the algorithm sorts squared scores in ascending order to choose the largest allowable number of blocks to drop.Source: B IMPLEMENTATION AND REPRODUCIBILITY… · p. 12
  • Kept blocks receive weight 1/q_m and skipped blocks weight 0, making each block's weight unbiased in expectation.The FlashAttention backward kernel reads the weights and iterates only over active query blocks; skipped blocks incur a weight read but no tile traffic.Source: B IMPLEMENTATION AND REPRODUCIBILITY… · p. 12
  • The CUDA selection pipeline reduces output gradients into block energies, applies the budget, and broadcasts weights; measurements were taken on H200 GPUs.CUDA-event timing brackets attention output through query gradient via autograd hooks, and the methods share the same container image and compiled kernel.Source: B IMPLEMENTATION AND REPRODUCIBILITY… · pp. 12–13

Related papers

Loading 2610.31911v1…