diff --git a/docs/design/feature/skip_softmax.md b/docs/design/feature/skip_softmax.md index 0527c7eeb..632df5c54 100644 --- a/docs/design/feature/skip_softmax.md +++ b/docs/design/feature/skip_softmax.md @@ -1,90 +1,197 @@ # 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. +Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. It implements +[BLASST](https://arxiv.org/abs/2512.12087) (Dynamic BLocked Attention Sparsity via Softmax +Thresholding), which adds a per-tile skip test to the FlashAttention main loop. This page describes the algorithm, how the user configuration is resolved +into the value the kernel consumes, and how checkpoint calibration and timestep gating participate +in that resolution. Configuration keys and operating-point guidance are covered in +[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax); the +backend selection contract is covered in +[Diffusion Attention Backend Selection](attention_backend_selection.md). ## 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. +detects, per tile of keys, when a tile 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 +It is approximate: a skipped tile still carries a small non-zero contribution, so the mode is opt-in and off by default. -## The online-softmax pass +## Algorithm -Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the +![The BLASST algorithm: FlashAttention with a per-tile skip test on the local maximum](../figures/skip_softmax/blasst_algorithm.jpg) + +FlashAttention computes attention in a single streaming pass over KV tiles. 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. +and returns `O / l` at the end. For each KV tile it computes the scores `QK_j^T`, updates `m`, and +accumulates the tile's contribution into `l` and `O`. Rescaling `l` and `O` when `m` grows keeps the +dense pass numerically exact. -## The skip test - -Once a tile's scores are known, its largest score `tile_max` is compared against the running -maximum: +A tile is a block of query rows against a block of keys, and the kernel processes it as a unit. +BLASST inserts one test between computing a tile's scores and accumulating them. For each query +row `i` in the tile, let `tile_max[i]` be the row's largest score within the tile, `running_max[i]` +the row's running maximum, and `λ` the effective threshold: ```text -if exp(tile_max − running_max) < threshold: - skip this tile # do not compute its softmax weights or its P·V accumulation +if all(exp(tile_max[i] - running_max[i]) < λ for every row i in the tile): + skip this tile # do not compute its Softmax or PV contribution ``` -`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. +`exp(tile_max[i] − running_max[i])` is an upper bound on the softmax weight any key in the tile can +receive in row `i`: if even the row's best key in this tile is far below the row's current maximum, +the tile is unimportant for that row. Only when this holds for every row in the tile can the +Softmax and `PV` work for the tile be skipped, in which case its contribution to `l` and `O` is +simply omitted. A larger `λ` makes the test more aggressive. The figure states the same test in +log space, `m̃ − m < ln(λ)`, which is how the kernel evaluates it. -## What this bounds +The test needs the tile's scores, so `QK_j^T` always runs and only the Softmax and `PV` work can +be skipped. Since the two matmuls have the same FLOP count, skipping every eligible tile removes at +most roughly half of the attention arithmetic, and the kernel-level speedup stays well under 2×. -Two properties of the test shape the achievable speedup: +## From configuration to the kernel threshold -- **`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: +The FlashInfer kernel does not take `λ` directly. It takes a `threshold_scale_factor` and divides +it by the KV sequence length of the call: ```text -threshold = factor / seqlen +λ = threshold_scale_factor / max_kv_len ``` -`factor` comes from one of two sources: +vLLM-Omni resolves the factor from one of two mutually exclusive user controls: -- **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. +| Control | Factor passed to the kernel | Resulting `λ` | +| --- | --- | --- | +| `threshold` | `threshold * max_kv_len` | `threshold` | +| `target_sparsity=s` | `a * exp(b * s)` | `a * exp(b * s) / max_kv_len` | + +`threshold` is therefore `λ` itself and is independent of sequence length: the same value yields +the same per-tile test at any resolution or frame count. `threshold=0` skips no tiles. Because the +left-hand side of the skip test lies in `(0, 1]`, values in `(0, 1)` are the meaningful range; the +schema accepts any finite non-negative value. + +`target_sparsity` selects a point on a curve fitted per model by +[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity), +so that `s` lands near that fraction of skipped tiles on the calibration data. The coefficients +`a` and `b` come from the checkpoint; the achieved sparsity on another prompt, shape, or layer can +differ. + +### Skip-Softmax calibration config + +ModelOpt stores the calibration in each transformer component's `config.json` under +`sparse_attention_config`. The following is taken from the `transformer/config.json` of +[`nvidia/Wan2.2-T2V-A14B-Diffusers-FP8`](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8), +with the `ignore` list abridged: + +```json +{ + "sparse_attention_config": { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "targets": ["WanAttention"], + "ignore": ["blocks.0.attn1", "blocks.0.attn2", "blocks.1.attn1", "..."], + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "coefficients": {"a": 2822.6209952221557, "b": 4.388250623520793} + }, + "target_sparsity": 0.7, + "disabled_until_timestep": 0.93 + } + } + } +} +``` + +A checkpoint with several transformer components carries one such config per component, each +with its own coefficients; in this example, `transformer_2/config.json` holds the low-noise +expert's curve. + +vLLM-Omni consumes `formula`, `coefficients`, and `ignore`, and ignores `targets`, +`target_sparsity`, and `disabled_until_timestep`: + +- Only the formula `a * exp(b * target_sparsity)` is accepted; a checkpoint with any other formula + is rejected at startup. If a spec requests `target_sparsity` and the checkpoint carries no + calibration, startup fails and the error points to `threshold` as the calibration-free + alternative. +- Each component's curve is applied to that component's attention layers. If the config for a + secondary component such as `transformer_2` cannot be read, that component stays dense and a + warning is logged. +- `ignore` holds fnmatch patterns matched against both the full module name and the name relative + to the component, so `blocks.0.attn1` matches `transformer.blocks.0.attn1` as well as + `transformer_2.blocks.0.attn1`. Matching layers receive no coefficients and therefore stay dense + regardless of the user configuration. + +At each attention call, a layer with a `threshold` uses it directly; a layer with +`target_sparsity` and stamped coefficients evaluates the curve; a layer with `target_sparsity` but +no coefficients runs dense. The timestep gate below is applied to the result before it reaches the +kernel. ## 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. +The early, high-noise denoising steps set the global structure of the output, and their errors +propagate through every later step. Keeping just these steps dense markedly improves the quality +of the generated video at a small cost in skipped work. `disabled_until_timestep` implements this: +on each attention call the backend compares the current normalized timestep `t` against the +configured cutoff and passes no skip factor to the kernel while `t` is above it. Skip-Softmax +becomes active on the first step whose `t` is at or below the cutoff and stays active for the rest +of the run, since `t` only decreases. -`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. +The default `disabled_until_timestep = 0` is a sentinel rather than a cutoff: the gate is off, the +timestep is never read, and Skip-Softmax runs on every step. Any positive value enables the gate. +`disabled_until_timestep = 1.0` therefore also yields no dense steps on a pipeline that publishes +`t`, but on a pipeline that does not publish it the backend stays dense and logs a warning once +rather than guessing from the step index. + +### Where `t` comes from + +The pipeline publishes `t` once per denoising step before running the transformer. It is the +scheduler's own position in `[0, 1]`, decreasing from near `1.0` (pure noise) to `0.0` (clean +sample): scheduler-based pipelines publish the scheduler timestep divided by `num_train_timesteps`, +and rectified-flow pipelines publish the current sigma. `t` is deliberately not the step index +divided by the step count. Schedulers place their steps non-uniformly in `t`, so expressing the +cutoff in `t` ties it to the same noise level across step counts and schedulers. + +### Mapping the cutoff to denoising steps + +For a run whose published sequence is `t[0], ..., t[N-1]`, the number of dense steps is + +```text +dense_steps = count(t[i] > disabled_until_timestep) +skip_softmax_steps = N - dense_steps +``` + +This count **depends on the schedule**, and schedules differ from model to model: the +scheduler family, the number of steps, and any time shift all change where the steps land in `t`. +The same cutoff can therefore gate a very different fraction of the run on two models. Take +MiniMax-H3 as an example. Its video branch uses a flow-shifted rectified-flow schedule that applies +a shift `s` to `N + 1` uniform positions `u` from `1` to `0`, + +```text +t = s * u / (1 + (s - 1) * u) +``` + +and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and +`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence +starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving: + +| `disabled_until_timestep` | Dense steps | Skip-Softmax steps | +| :---: | ---: | ---: | +| `1.00` | 0 | 49 | +| `0.99` | 6 | 43 | +| `0.97` | 14 | 35 | +| `0.95` | 19 | 30 | +| `0.90` | 28 | 21 | +| `0.86` | 33 | 16 | + +These numbers hold only for that schedule. A model with a smaller shift, a different step count, +or a scheduler that is not rectified flow produces a different sequence and a different table; a +request that overrides `flow_shift` or the step count changes it for MiniMax-H3 too. +Choose the cutoff by counting against the schedule actually served, not by reusing a value from +another model. diff --git a/docs/design/figures/skip_softmax/blasst_algorithm.jpg b/docs/design/figures/skip_softmax/blasst_algorithm.jpg new file mode 100644 index 000000000..6973d45d6 Binary files /dev/null and b/docs/design/figures/skip_softmax/blasst_algorithm.jpg differ diff --git a/docs/user_guide/diffusion/attention_backends/dense_backends.md b/docs/user_guide/diffusion/attention_backends/dense_backends.md index 030550dac..d7182eb44 100644 --- a/docs/user_guide/diffusion/attention_backends/dense_backends.md +++ b/docs/user_guide/diffusion/attention_backends/dense_backends.md @@ -1,8 +1,9 @@ # Dense Attention Backends -Use these backends when you want dense attention without backend-specific -sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark -the fastest compatible kernel for your model, shape, and hardware. +Use these backends to run dense attention. Optional quantized or sparse modes +remain inactive unless they are explicitly configured. Start with +`TORCH_SDPA` for correctness comparisons, then benchmark the fastest compatible +kernel for your model, shape, and hardware. For selection precedence and per-role configuration, see the [attention backend overview](../attention_backends.md). @@ -41,6 +42,21 @@ Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on Blackwell. If the CuTe path is unavailable, the backend falls back to the compatible FlashAttention 3 or 2 path. +## `TRTLLM_ATTN` + +`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels and is the platform +default on datacenter Blackwell for models that declare a compatible attention +path. Selected without a `quant` or `skip_softmax` block, it runs dense BF16 +attention at FA4-level performance. + +```bash +vllm serve --omni \ + --diffusion-attention-backend TRTLLM_ATTN +``` + +See [TRTLLM Attention](trtllm.md) for its requirements and for the optional +SAGE quantization and Skip-Softmax modes. + ## `CUDNN_ATTN` `CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful diff --git a/docs/user_guide/diffusion/attention_backends/trtllm.md b/docs/user_guide/diffusion/attention_backends/trtllm.md index dd36ba61e..cd7e55d3d 100644 --- a/docs/user_guide/diffusion/attention_backends/trtllm.md +++ b/docs/user_guide/diffusion/attention_backends/trtllm.md @@ -1,103 +1,274 @@ # TRTLLM Attention -`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense -BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization. +`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels on datacenter +Blackwell GPUs. Selected on its own it computes dense BF16 attention, and its +dense performance is on par with +[FlashAttention-4](dense_backends.md#flashattention-4-on-blackwell). It also +provides two opt-in, lossy acceleration modes that can be enabled independently +or together: -For common selection and per-role configuration, see the -[attention backend overview](../attention_backends.md). +| Mode | Config key | What it changes | +| --- | --- | --- | +| [Skip-Softmax](#skip-softmax) | `skip_softmax` | Skips the Softmax and `PV` work of KV tiles whose scores are too low to matter | +| [SAGE quantization](#sage-quantization) | `quant` | Runs `QK^T` in INT8 or FP8 and `PV` in FP8 instead of BF16 | + +Both keys are set through `--diffusion-attention-config`, as JSON or as +vLLM-style dotted flags; the +[attention backend overview](../attention_backends.md#configuration) covers +both syntaxes and per-role resolution. The examples on this page use JSON. + +Attention computes scores `S = QK^T`, probabilities `P = softmax(S)`, and the +output `O = PV`. SAGE lowers the precision of both matrix multiplications. +Skip-Softmax keeps `QK^T` dense and removes the Softmax and `PV` work for +tiles that `QK^T` shows to be unimportant. The two modes therefore compose: +SAGE makes every tile cheaper, and Skip-Softmax reduces the number of tiles +that reach the second half of the kernel. ## Requirements -Selecting `TRTLLM_ATTN` requires: +`TRTLLM_ATTN` requires all of the following: -- datacenter Blackwell, sm_100 or sm_103; +- an `sm100a` or `sm103a` GPU (B200, B300, GB200, GB300); workstation Blackwell + (`sm120`/`sm121`) is not supported; - `head_dim=128`; -- FlashInfer; and -- a model-declared packed or mask-free compatible path. +- a FlashInfer build that exposes the trtllm-gen kernels (0.6.16rc1 or newer + for SAGE); +- an attention path that is mask-free or provides packed-padding metadata. + Structural suffix padding is expressed through that metadata rather than an + `attn_mask` tensor. -It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible -explicit selection raises instead of silently changing the backend. +An explicit selection that violates these requirements raises at startup +instead of silently falling back to another backend. -## Dense mode +### Sequence parallelism -Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16: +`TRTLLM_ATTN` runs under no sequence parallelism or under pure Ulysses. +Ulysses redistributes the sequence and attention heads around the attention +call, but the local computation still goes through the configured backend, so +both optional modes work unchanged. Ring and AllGather-KV do not: + +- Ring runs its own distributed attention and bypasses the backend. Combining + Ring with a `skip_softmax` key raises; a `quant` key would be silently + ignored, so do not combine them either. +- AllGather-KV changes the Q/KV distribution and is rejected when + `TRTLLM_ATTN` is selected. + +Configure eight-way Ulysses alone as: ```bash -vllm-omni serve --diffusion-attention-backend TRTLLM_ATTN +--usp 8 --ring 1 --allgather-degree 1 +``` + +A degree of `1` disables that sequence-parallel mode. It does not limit the +server to one GPU or affect tensor, pipeline, or VAE parallelism. + +## Quick start + +On datacenter Blackwell the platform selects `TRTLLM_ATTN` by default when the +model declares a compatible path. To select it explicitly, for dense BF16: + +```bash +vllm serve --omni \ + --diffusion-attention-backend TRTLLM_ATTN +``` + +To enable both optimizations with a conservative Skip-Softmax setting and FP8 +SAGE: + +```bash +vllm serve --omni \ + --diffusion-attention-config '{ + "default": { + "backend": "TRTLLM_ATTN", + "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}, + "skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97} + } + }' +``` + +The startup log reports the selection; look for one of: + +```text +Defaulting to diffusion attention backend TRTLLM_ATTN (datacenter Blackwell ..., head_dim 128) +Resolved diffusion attention backend 'TRTLLM_ATTN' for role='self' via attention_config.default ``` ## Skip-Softmax -Skip-Softmax trades some fidelity for speed by skipping low-contribution -attention work. See the [feature design](../../../design/feature/skip_softmax.md) -for the algorithm and implementation contract. +Skip-Softmax (BLASST) skips the Softmax and `PV` work of KV tiles whose +scores fall far below the running row maximum. `QK^T` always runs, and how +many tiles qualify depends on the input. The +[feature design](../../../design/feature/skip_softmax.md) gives the skip test. -| Key | Valid values | Meaning | +### Configuration keys + +| Key | Range | Meaning | | --- | --- | --- | -| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve | -| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` | -| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` | +| `threshold` | `>= 0`; useful values in `(0, 1)` | Skip threshold, independent of sequence length. Calibration-free. | +| `target_sparsity` | `[0, 1]` | Requested operating point on the checkpoint's calibrated curve. Requires calibration metadata. | +| `disabled_until_timestep` | `[0, 1]`; default `0` | Keeps attention dense while the normalized timestep `t > D`; `0` disables the gate. | + +`threshold` and `target_sparsity` are two ways to set the same kernel +threshold; setting both is a configuration error. Exactly one of them enables +Skip-Softmax. + +### Direct threshold + +Set `threshold` when the checkpoint carries no calibration. `threshold=0` +skips nothing; larger values skip more tiles and lower output fidelity. Values +around `0.05` are a reasonable first try for video DiTs; tune against dense +output on the same prompt and seed. The value is independent of sequence +length; see the +[feature design](../../../design/feature/skip_softmax.md#from-configuration-to-the-kernel-threshold) +for how it reaches the kernel. ```bash -vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\ - "skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}' +vllm serve --omni \ + --diffusion-attention-config \ + '{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{ + "threshold":0.05,"disabled_until_timestep":0.97}}}' ``` -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec +### Calibrated target sparsity -config = AttentionConfig( - default=AttentionSpec( - backend="TRTLLM_ATTN", - skip_softmax=SkipSoftmaxSpec( - target_sparsity=0.65, - disabled_until_timestep=0.86, - ), - ), -) +A fixed `threshold` does not produce a fixed fraction of skipped tiles because +the score distribution changes with the model, prompt, and shape. +[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer) +can calibrate a per-model curve from sparsity to threshold and store it in the +checkpoint's transformer `config.json` under `sparse_attention_config`. +`target_sparsity` selects a point on that curve. The achieved sparsity still +varies per prompt, layer, and denoising step; the calibration makes the +requested value a meaningful target, not a guarantee. + +What vLLM-Omni takes from the checkpoint: + +- The curve coefficients. Only the `a * exp(b * target_sparsity)` form ModelOpt + writes is supported; another formula is rejected at startup. +- The `ignore` list: attention layers matching those patterns stay dense. +- For multi-expert Diffusers checkpoints, `transformer_2/config.json` is read + separately; if it is missing, `transformer_2` stays dense with a warning. + +Checkpoint-level `target_sparsity` and `disabled_until_timestep` defaults are +not consumed; set them in the vLLM-Omni configuration. The +[feature design](../../../design/feature/skip_softmax.md#skip-softmax-calibration-config) +documents the checkpoint format and how the coefficients reach the kernel. + +To calibrate a checkpoint yourself, follow the +[ModelOpt Skip-Softmax example](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity), +which exports a copy of the checkpoint with the coefficients embedded. The +fitted curve depends on the attention statistics of the calibration shape, so +calibrate at the resolution and frame count you will serve. + +Requesting `target_sparsity` for a checkpoint without calibration is a startup +error that names the `threshold` alternative. The +[ModelOpt Wan2.2 FP8 checkpoint](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8/blob/main/transformer/config.json) +is a calibrated example: + +```bash +vllm serve nvidia/Wan2.2-T2V-A14B-Diffusers-FP8 --omni \ + --diffusion-attention-config \ + '{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{ + "target_sparsity":0.75,"disabled_until_timestep":0.86}}}' ``` -Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase -sparsity only after comparing quality against dense output at the same seed. -The benefit grows with sequence length because only attention is accelerated. +### Timestep gating + +The early, high-noise denoising steps fix the global layout of the output, and +their errors propagate through every later step. `disabled_until_timestep` +keeps those steps dense and enables Skip-Softmax once the normalized timestep +`t` reaches the configured cutoff. The default `0` is a sentinel that turns the +gate off: Skip-Softmax runs on every step and no timestep needs to be +published. `1.0` also leaves no step dense, but goes through the gate and +therefore requires the pipeline to publish `t`. + +`t` is the scheduler's own timestep normalized to `[0, 1]`, published by the +pipeline for each denoising step; for rectified-flow models it is the current +sigma. It runs from near `1.0` down to `0.0`, but not linearly in the step +index: flow-shifted schedules spend many of their steps at high `t`. The number +of dense steps a given cutoff produces therefore depends on the model's +schedule and step count, and follows from the actual timestep sequence: + +```text +dense_steps = count(t[i] > disabled_until_timestep) +``` + +Derive the cutoff from the schedule of the model you are serving rather than +reusing another model's value; the +[feature design](../../../design/feature/skip_softmax.md#mapping-the-cutoff-to-denoising-steps) +works through an example. + +A pipeline that does not publish `t` stays dense whenever +`disabled_until_timestep > 0` is set, and logs a warning once. Pipelines +publish it through `DenoiseProgressMixin.record_denoise_step`. ## SAGE quantization -TRTLLM SAGE quantizes Q/K per block and V per channel so both attention -matrix multiplications use low precision. +SAGE quantization runs both attention matrix multiplications in low precision: +Q and K are quantized to INT8 or FP8 E4M3 for `QK^T`, and P and V use FP8 E4M3 +for `PV`. P is quantized inside the FMHA kernel; V is quantized per channel +before the kernel call. +vLLM-Omni exposes the Q/K dtype and the Q/K scale granularity. The P and V +formats are fixed by the kernel, so the `quant` key has no V dtype for this +backend. This mode is distinct from the standalone +[SageAttention backends](sage.md), which use their own kernels. -| Key | Valid values | Meaning | +FP8 Q/K kernels exist on `sm100a` and `sm103a`; INT8 Q/K kernels exist on +`sm100a` only. + +| Key | Values | Meaning | | --- | --- | --- | -| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense | -| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` | -| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` | +| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype. Setting it enables SAGE. | +| `q_block_size` | `1`, `4`, `16` | Consecutive query tokens sharing one Q scale; default `1` | +| `k_block_size` | `1`, `4`, `16` | Consecutive key tokens sharing one K scale; default `16` | -V is always quantized per channel to FP8 E4M3 and K-smoothing is internal. -Only block sizes with compiled kernels are accepted. +Start with the defaults, `q_block_size=1` and `k_block_size=16`, and try +smaller K blocks only if output quality does not meet expectations; smaller +blocks give finer scales at some cost in speed. When a KV sequence in a call is +shorter than `k_block_size`, that call falls back to dense attention and a +warning is logged once. ```bash -vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\ - "quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}' +vllm serve --omni \ + --diffusion-attention-config \ + '{"default":{"backend":"TRTLLM_ATTN","quant":{ + "dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}' ``` -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec +The `quant` key is shared with `FLASHINFER_ATTN`, but each backend validates +its own fields: `float16`/`bfloat16` Q/K dtypes and `dtype_vo` are +`FLASHINFER_ATTN` options and are rejected here. -config = AttentionConfig( - default=AttentionSpec( - backend="TRTLLM_ATTN", - quant=AttnQuantSpec( - dtype_qk="fp8_e4m3", - q_block_size=1, - k_block_size=16, - ), - ), -) +## Composing both modes + +`skip_softmax` and `quant` may appear in the same `AttentionSpec`. Their +quality effects compound, so establish a dense baseline, enable one mode at a +time, and then evaluate the combination on the same prompts and seeds. + +Modes configured in `default` apply to every attention role that has no +`per_role` entry. A `per_role` spec replaces the whole spec for that role and +does not inherit `quant` or `skip_softmax` from `default`, so +`{"backend":"TRTLLM_ATTN"}` is the way to keep a short or sensitive attention +site dense while the long DiT sequence uses both modes: + +```bash +vllm serve MiniMaxAI/MiniMax-H3 --omni \ + --diffusion-attention-config '{ + "default": { + "backend": "TRTLLM_ATTN", + "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}, + "skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97} + }, + "per_role": { + "minimax_h3.token_refiner": {"backend": "TRTLLM_ATTN"} + } + }' ``` -This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are -available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only. -The shared `AttnQuantSpec` is backend typed: values intended for -`FLASHINFER_ATTN` are rejected here and vice versa. +Role names are declared by each model; the +[attention backend overview](../attention_backends.md#configuration) covers +the resolution order and the equivalent Python API. + +End-to-end speedup depends on the share of step time spent in attention, the +sequence length, the chosen Q/K precision, and the tile sparsity the input +actually yields. Benchmark the exact workload rather than extrapolating from +another model. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 9f63ea7ad..5ef1e9ebc 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -278,7 +278,7 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain ### Attention Backends -On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16 +On supported datacenter Blackwell systems, MiniMax H3 defaults to dense BF16 `TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly, use: @@ -298,17 +298,33 @@ FA4 remains available by explicitly selecting the `FLASH_ATTN` backend: --diffusion-attention-backend FLASH_ATTN ``` -On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains -`Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements. +With the optional `[fa4]` dependency installed, `FLASH_ATTN` prefers FA4 on +Blackwell. Confirm the server log contains `Using CuTe FlashAttention-4 on +Blackwell` before recording FA4 measurements. -`TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main -DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse -Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses -`fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V -to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short -attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled -for it. The example enables the calibration-free Skip-Softmax path with -`threshold=0.05`, after the normalized timestep reaches `0.97`: +`TRTLLM_ATTN` additionally offers two **lossy** optimizations for the long main +DiT attention sequence: SAGE attention quantization and Skip-Softmax sparse +attention. Both work under the pure Ulysses parallelism of the profile above +(`--usp 4 --ring 1`). The example below enables both: + +- SAGE with `fp8_e4m3` Q/K; P and V are always FP8 in this kernel. B200 + additionally supports `int8` Q/K, which preserves accuracy better than FP8. +- Skip-Softmax with a direct `threshold=0.05` (the calibrated + `target_sparsity` control needs ModelOpt metadata that the official H3 + checkpoint does not include). Together with the cutoff below this is a + **conservative** setting: a low threshold skips only clearly negligible tiles, + and a cutoff close to `1.0` leaves a substantial dense prefix. Raise + `threshold` or lower `disabled_until_timestep` for more speedup once quality + is verified. +- `disabled_until_timestep=0.97` keeps the early high-noise steps dense. The + gate compares against the video sigma, which H3's default flow shift of 12 + keeps high for much of the run: at 50 steps, `0.99`, `0.97`, and `0.95` + leave the first 6, 14, and 19 of 49 denoiser forwards dense. See the + [Skip-Softmax design](https://github.com/vllm-project/vllm-omni/blob/main/docs/design/feature/skip_softmax.md#timestep-gating) + for how the cutoff maps to steps. +- A `per_role` entry that keeps the token refiner, a short attention path, + dense. A per-role spec does not inherit `quant` or `skip_softmax` from + `default`. ```bash --diffusion-attention-config '{ @@ -332,10 +348,13 @@ for it. The example enables the calibration-free Skip-Softmax path with }' ``` -For configuration details, see -[TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax) +Both optimizations trade fidelity for speed and their effects compound. +Compare against dense output on the same prompt and seed before adopting them. +For the full key reference, see +[Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#skip-softmax) and -[TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization). +[SAGE quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#sage-quantization) +in the TRTLLM attention guide. ### Text encoder tensor parallelism