vllm-omni/docs/design/feature/skip_softmax.md

# Skip-Softmax

Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
and how to pick an operating point — is in
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
The shared selector contract is documented in
[Diffusion Attention Backend Selection](attention_backend_selection.md). This
page explains the algorithm.

## Motivation

In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight
concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the
output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax
detects, per block of keys, when a block cannot matter and skips its softmax and its value
multiply.

It is approximate: a skipped block still carries a small non-zero contribution, so the mode is
opt-in and off by default.

## The online-softmax pass

Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the
kernel maintains three running values:

- `m` — the largest score seen so far,
- `l` — the running denominator `Σ exp(sⱼ − m)`,
- `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,

and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
`m` grows keeps the single pass numerically exact.

## The skip test

Once a tile's scores are known, its largest score `tile_max` is compared against the running
maximum:

```text
if exp(tile_max − running_max) < threshold:
    skip this tile          # do not compute its softmax weights or its P·V accumulation
```

`exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
receive: if even the tile's best key is far below the current maximum, every key in the tile is
negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
dropped. The tile's contribution to `l` and `O` is simply skipped.

## What this bounds

Two properties of the test shape the achievable speedup:

- **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
  score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
  and the value matmul are comparable in cost, so skipping every eligible tile removes at most
  roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.

- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
  collectively negligible; a single important key keeps the whole tile. How many tiles actually
  qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
  operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
  skipped.

## The threshold

The per-tile threshold is normalized by sequence length:

```text
threshold = factor / seqlen
```

`factor` comes from one of two sources:

- **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
  per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
  skipped tiles on the calibration data.
- **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
  the threshold themselves.

## Timestep gating

`disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
schedule). The early steps set the global structure of the output and their errors propagate through
every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.

`t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
`DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
stays dense when `disabled_until_timestep` is set.