[Doc] Refine TRTLLM attention guidance (#6724)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

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.

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 <model> --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

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 <model> --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 <model> --omni \
+  --diffusion-attention-backend TRTLLM_ATTN
+```
+
+To enable both optimizations with a conservative Skip-Softmax setting and FP8
+SAGE:
+
+```bash
+vllm serve <model> --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 <model> --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 <model> --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.

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