BackLite

BackLite is a lightweight wrapper around Flash Attention 3 which identifies and exploits the sparsity of attention matrices in order to speed up the attention backward pass, while mathematically approximating the attention gradients. See here for more technical details.

📖 Overview

Attention backward passes are a dominant cost in transformer training. BackLite reduces this cost by making the matrix multiplications for the backward pass block-sparse. The algorithm computes the forward pass exactly (similar to FA) while additionally recording the weights of the attention matrix tiles. During the backward pass, the algorithm skips the computations related to the smallest attention matrix tiles, which cumulatively weigh less than a configurable probability weight threshold (negl_prob).

Key properties:

• Accurate: Smallest weighted tiles with cumulative weight < negl_prob are skipped.
• Tunable: One can tune negl_prob online to target specific gradient similarity metrics (cosine similarity, relative L2 difference) between the gradients computed by FA and BackLite. (Calibration tool: coming soon!)
• Zero forward overhead when negl_prob = 0 (falls back to standard FA3).
• Adaptive: Sparsity is derived from actual attention statistics per sample, per head.
• Composable: Supports LSE output for combining partial attention results.

🔍 How It Works

Sparse Backward Masking

BackLite uses a two-phase mechanism:

Phase 1 — Forward pass with tile statistics: When negl_prob > 0, the FA3 forward kernel records the per-tile log-sum-exp (LSE) for every (query tile, key tile) pair alongside the standard output and full-row LSE.

Phase 2 — Backward pass with block-sparse mask: Before the backward kernel runs, a fused Triton kernel (mask_from_stats_fused) converts the tile LSEs into a block-sparse boolean mask. For each backward tile pair, it computes:

$$p_\text{tile} = \exp!\left(\mathrm{LSE}_\text{tile} - \mathrm{LSE}_\text{row}\right)$$

mask_from_stats_fused selects the smallest weighted tiles with cumulative weight summing less than negl_prob and marks them as skippable. The backward kernel bypasses it entirely — no memory reads, no FLOPs.

This approach:

• Introduces no approximation in the forward pass
• Introduces negligible gradient error for properly chosen negl_prob
• Generates the sparsity mask in a single Triton pass over the stored tile statistics

📊 Backward Sparsity

TBD

🔧 Installation

Requirements

• H100 / H200 GPU
• CUDA >= 12.8
• CUDA toolkit
• C++ 20
• PyTorch 2.2 and above
• packaging Python package (pip install packaging)
• ninja Python package (pip install ninja) *
• Linux

* Make sure that ninja is installed and that it works correctly (e.g. ninja --version then echo $? should return exit code 0). If not (sometimes ninja --version then echo $? returns a nonzero exit code), uninstall then reinstall ninja (pip uninstall -y ninja && pip install ninja). Without ninja, compiling can take a very long time (2h) since it does not use multiple CPU cores. With ninja compiling takes 3-5 minutes on a 64-core machine using CUDA toolkit.

Build from Source

Clone this repo and build from source:

git clone https://github.com/moonmath-ai/BackLite.git
cd BackLite/hopper
pip install .

If your machine has less than 96GB of RAM and lots of CPU cores, ninja might run too many parallel compilation jobs that could exhaust the amount of RAM. To limit the number of parallel compilation jobs, you can set the environment variable MAX_JOBS:

MAX_JOBS=4 pip install .

🚀 Usage

Basic Usage (Single GPU)

BackLite(negl_prob: float = 0.05)

Parameters:

• negl_prob (float): Negligible probability mass threshold for backward sparsity. Tiles whose cumulative attention mass is below this value are skipped during gradient computation. Set to 0.0 to use standard FA3 with no sparsity. Typical values: 0.01–0.1.

from back_lite import BackLite

# Standard FA3 (no backward sparsity)
attn = BackLite()
output = attn(query, key, value)

# Sparse backward pass — skip smallest tiles cumulatively contributing < 5% probability mass
attn = BackLite(negl_prob=0.05)
output = attn(query, key, value)

# With explicit softmax scale
output = attn(query, key, value, scale=1.0 / math.sqrt(head_dim))

# Forward + backward example
attn = BackLite(negl_prob=0.05)
output = attn(query, key, value)
loss.backward()  # backward uses the sparse mask generated in the forward pass

Forward signature:

attn(
    query,               # (batch, seq_len, heads, head_dim)
    key,                 # (batch, seq_len, heads_k, head_dim)
    value,               # (batch, seq_len, heads_k, head_dim)
    scale=None,          # softmax scale, default 1/sqrt(head_dim)
    return_softmax_lse=False,  # return (output, lse) instead of output
    tile_stats=None,     # pre-allocated tile-stats buffer (optional)
)

Note

When negl_prob > 0, BackLite automatically allocates a tile_stats buffer during the forward pass to store per-tile LSE statistics. This buffer is passed to the autograd backward function, where a Triton kernel converts it into a block-sparse mask before the backward kernel runs. No changes to your training loop are required.

📝 Integration Example

TBD

🐛 Debugging

Set BACK_LITE_VERBOSE=1 to enable additional logging (tile statistics, mask density, backward sparsity ratio) during training.

🙏 Acknowledgements

BackLite is built on top of FlashAttention3 by Tri Dao and contributors. We thank the FlashAttention team for their foundational work on efficient attention mechanisms.

We also thank the teams behind SparseVideoGen, RadialAttention, SageAttention, Wan2.1, and LTX-Video for their insights and benchmarking support.

License

BackLite is build on top of FA3 which has a BSD 3-Clause license. As such the original code maintains that license and any new code for BackLite is distributed under an MIT license.

See LICENSE-BSD and LICENSE-MIT for further details.