diff --git a/docs/source/blogs/tech_blog/blog28_Accelerating_Video_Generation_with_GEMM_Quantization_Attention_Quantization_and_Skip_Softmax_Attention_in_TensorRT-LLM.md b/docs/source/blogs/tech_blog/blog28_Accelerating_Video_Generation_with_GEMM_Quantization_Attention_Quantization_and_Skip_Softmax_Attention_in_TensorRT-LLM.md index f8487a9fe7a0..9012a38179a9 100644 --- a/docs/source/blogs/tech_blog/blog28_Accelerating_Video_Generation_with_GEMM_Quantization_Attention_Quantization_and_Skip_Softmax_Attention_in_TensorRT-LLM.md +++ b/docs/source/blogs/tech_blog/blog28_Accelerating_Video_Generation_with_GEMM_Quantization_Attention_Quantization_and_Skip_Softmax_Attention_in_TensorRT-LLM.md @@ -75,9 +75,9 @@ Because the score distribution depends on the input, a fixed threshold $\lambda$ - `target_sparsity` expresses how aggressively to skip as an intuitive target. - `disabled_until_timestep` keeps the early denoising steps dense before enabling Skip Softmax. -The mapping from `disabled_until_timestep` to actual denoising steps depends on the scheduler and is not linear. In the 40-step UniPC schedule used here, `disabled_until_timestep=0.86` keeps the 14/40 steps dense and enables Skip Softmax for the remaining 26. See the [VisualGen Skip Softmax Attention documentation](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/visual-gen/features/sparse-attention.md#mapping-disabled_until_timestep-to-actual-denoising-steps) for the scheduler-dependent mapping. +The mapping from `disabled_until_timestep` to actual denoising steps depends on the scheduler and is not linear. In the 40-step UniPC schedule used here, `disabled_until_timestep=0.86` keeps the 14/40 steps dense and enables Skip Softmax for the remaining 26. See the [VisualGen Skip Softmax Attention documentation](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/visualgen-sparse-attention.md#mapping-disabled_until_timestep-to-actual-denoising-steps) for the scheduler-dependent mapping. -The ModelOpt checkpoints used in this experiment already include this calibration metadata. Without calibration, Skip Softmax can still be enabled by setting the threshold directly. See the [VisualGen Skip Softmax Attention documentation](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/visual-gen/features/sparse-attention.md#skip-softmax-attention) for direct-threshold configuration. +The ModelOpt checkpoints used in this experiment already include this calibration metadata. Without calibration, Skip Softmax can still be enabled by setting the threshold directly. See the [VisualGen Skip Softmax Attention documentation](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/visualgen-sparse-attention.md#skip-softmax-attention) for direct-threshold configuration. ## Results diff --git a/docs/source/visual-gen/features/cuda-graph.md b/docs/source/features/visualgen-cuda-graph.md similarity index 91% rename from docs/source/visual-gen/features/cuda-graph.md rename to docs/source/features/visualgen-cuda-graph.md index e199fe0d84b8..f16a4c288b72 100644 --- a/docs/source/visual-gen/features/cuda-graph.md +++ b/docs/source/features/visualgen-cuda-graph.md @@ -1,7 +1,7 @@ -# VisualGen CUDA Graphs +# VisualGen CUDA Graphs (Beta) ```{note} -This page is an unindexed draft until the VisualGen documentation hub is introduced. +This feature is in **beta** stage. APIs, supported models, and optimization options are actively evolving and may change in future releases. ``` - [Overview](#overview) @@ -62,7 +62,7 @@ Models add non-shape graph-key contributors through `BaseDiffusionModel.register During key construction, the runner calls each registered callback with the same `*args` and `**kwargs` passed to the wrapped `model.forward`. If the callback returns a non-`None` hashable value, the runner appends `(name, value)` to the graph key. If it returns `None`, that key part is omitted for the current call. -Use an extra key when a forward input affects captured kernels or control flow but is not already represented by tensor shapes. For example, [Skip Softmax Attention](sparse-attention.md) can require separate CUDA graphs across its dense and sparse phases even when tensor shapes are identical. The base `BaseDiffusionModel` implementation registers a callback through `runner.register_extra_key_fn(...)` for this case. +Use an extra key when a forward input affects captured kernels or control flow but is not already represented by tensor shapes. For example, [Skip Softmax Attention](visualgen-sparse-attention.md) can require separate CUDA graphs across its dense and sparse phases even when tensor shapes are identical. The base `BaseDiffusionModel` implementation registers a callback through `runner.register_extra_key_fn(...)` for this case. Subclasses can override `register_cuda_graph_extra_key_fns()` to add model-specific contributors. They should call `super()` unless they intentionally replace the shared registrations. This changes graph capture partitioning without exposing internal CUDA graph keys as public model-forward arguments. diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md new file mode 100644 index 000000000000..3f23649e99e8 --- /dev/null +++ b/docs/source/features/visualgen-quantized-attention.md @@ -0,0 +1,224 @@ +# VisualGen Quantized Attention (Beta) + +```{note} +This feature is in **beta** stage. APIs, supported models, and optimization options are actively evolving and may change in future releases. +``` + +- [Overview](#overview) + - [Recipes](#recipes) + - [Choosing and Tuning a Recipe](#choosing-and-tuning-a-recipe) + - [Configuration Surface](#configuration-surface) +- [QK16PV8 Attention Kernels in the CUTEDSL Backend](#qk16pv8-attention-kernels-in-the-cutedsl-backend) +- [SageAttention (TRTLLM)](#sageattention-trtllm) +- [MXFP8 / NVFP4 (CUTEDSL / FlashInfer)](#mxfp8--nvfp4-cutedsl--flashinfer) +- [Interaction With Other Features](#interaction-with-other-features) + +## Overview + +Visual generation models spend a large fraction of each denoising step inside attention, and every step is a full-context pass rather than an autoregressive decode. Quantized attention lowers the precision of the tensors the attention kernel itself consumes (Q, K, V), so that BMM1 (`Q·Kᵀ`) and/or BMM2 (`P·V`) run on narrower Tensor Core instructions. This is orthogonal to `VisualGenArgs.quant_config`, which quantizes the linear layers' *weights*: quantized attention quantizes *activations* inside the attention op and leaves the checkpoint untouched, so it needs no calibrated checkpoint and can be switched on for any supported model. + +Quantized attention is configured through `VisualGenArgs.attention_config.quant_attention_config` (`QuantAttentionConfig`). + +### Recipes + +A recipe is the tuple `(qk_dtype, v_dtype, (q_block_size, k_block_size, v_block_size))`. Only the combinations below are accepted; `AttentionConfig` validates the recipe against the selected backend at construction time and raises `ValueError` otherwise (`tensorrt_llm/visual_gen/args.py`, `_validate_quant_attention_config`). + +| Backend | `qk_dtype` | `v_dtype` | `(q, k, v)` block sizes | Common name | +|---|---|---|---|---| +| `TRTLLM` | `int8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)`, `(1, 16, 1)` | SageAttention (INT8 QK) | +| `TRTLLM` | `fp8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)` | SageAttention (FP8 QK) | +| `CUTEDSL` | `bf16` | `fp8` | `(0, 0, 0)` | QK16PV8 | +| `CUTEDSL` | `mxfp8` | `fp8` | `(0, 0, 0)`, `(0, 0, 1)` | MXFP8 Q/K | +| `CUTEDSL` | `nvfp4` | `fp8` | `(0, 0, 0)`, `(0, 0, 1)` | NVFP4 Q/K | +| `FLASHINFER` | `mxfp8` | `fp8` | `(0, 0, 0)` | MXFP8 Q/K | +| `FLASHINFER` | `nvfp4` | `fp8` | `(0, 0, 0)` | NVFP4 Q/K | +| `FLASHINFER` | `nvfp4` | `nvfp4` | `(0, 0, 0)` | NVFP4 Attention | + +### Choosing and Tuning a Recipe + +Choose a quantized-attention recipe by output quality first. Establish an unquantized quality baseline, then evaluate each compatible recipe with representative prompts, input media, resolutions, and fixed seeds. Quantization sensitivity varies by model, so do not assume that a recipe validated for one model will preserve quality for another. + +Video quality is generally more sensitive to BMM1 accuracy than BMM2 accuracy, so preserving Q/K precision is the most conservative starting point: + +- QK16PV8 keeps Q/K in BF16 and only quantizes V, making it the most conservative quantized-attention recipe. +- On B200/GB200, SageAttention with INT8 Q/K typically matches QK16PV8 quality while delivering higher end-to-end throughput. +- On B300/GB300, start with MXFP8 when optimizing the quality-throughput balance. SageAttention with FP8 Q/K remains an alternative when the `TRTLLM` backend is preferred for the surrounding workload. +- For SageAttention with INT8 Q/K, the default `(1, 16, 1)` block-size recipe works well for most cases. Use `(1, 4, 1)` when video quality is not satisfactory. +- For `CUTEDSL` MXFP8 or NVFP4 recipes, `v_block_size: 1` uses a separate V scale per head and channel, while `v_block_size: 0` uses one tensor-wide V scale. Try the per-channel variant when the tensor-wide scale loses quality. + +After a recipe meets the quality target, benchmark its end-to-end throughput with the production workload. + +### Configuration Surface + +| Field | Type | Default | Meaning | +|---|---|---|---| +| `qk_dtype` | `"bf16" \| "int8" \| "fp8" \| "mxfp8" \| "nvfp4"` | `"bf16"` | Q/K element format for BMM1. `bf16` leaves Q/K unquantized. | +| `v_dtype` | `"fp8" \| "nvfp4"` | `"fp8"` | V element format for BMM2. | +| `q_block_size` | int ≥ 0 | `0` | Q tokens per SageAttention quantization block. `0` outside SageAttention. | +| `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` outside SageAttention. | +| `v_block_size` | int ≥ 0 | `0` | V block size on the hidden dimension. `0` = one tensor-wide V scale; `1` = one scale per channel. | + +Routing (`tensorrt_llm/_torch/visual_gen/attention_backend/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. + +## QK16PV8 Attention Kernels in the CUTEDSL Backend + +**What it does.** Q and K stay in BF16, so BMM1 runs at full input precision. Only V is quantized to FP8 e4m3, so BMM2 runs on FP8 Tensor Cores. + +**Configuration.** + +```python +from tensorrt_llm import VisualGenArgs +from tensorrt_llm.visual_gen import AttentionConfig, QuantAttentionConfig + +args = VisualGenArgs( + model="", + attention_config=AttentionConfig( + backend="CUTEDSL", + quant_attention_config=QuantAttentionConfig( + qk_dtype="bf16", + v_dtype="fp8", + q_block_size=0, + k_block_size=0, + v_block_size=0, + ), + ), +) +``` + +```yaml +attention_config: + backend: CUTEDSL + quant_attention_config: + qk_dtype: bf16 + v_dtype: fp8 + q_block_size: 0 + k_block_size: 0 + v_block_size: 0 +``` + +## SageAttention in the `TRTLLM` backend + +**What it does.** SageAttention quantizes all three tensors with fine-grained scales, so both BMM1 and BMM2 run in low precision: + +- **Q and K** are quantized to INT8 or FP8 e4m3 with one scale per *token block* per head. The block size is `q_block_size` for Q and `k_block_size` for K, measured in tokens along the sequence axis; a larger K block amortizes more scales but is coarser. +- **V** is quantized to FP8 e4m3 with `v_block_size` elements per scale along the hidden dimension. All supported recipes use `v_block_size = 1`, i.e. one scale per head per channel. + +**Requirements and behavior.** + +- SageAttention is supported on B200/GB200 and B300/GB300 GPUs. +- On B200/GB200, use the recommended `qk_dtype: "int8"` recipe. +- On B300/GB300, use `qk_dtype: "fp8"` and evaluate output quality because it can be less accurate than the INT8 Q/K recipe on B200/GB200. + +**Configuration.** + +```python +from tensorrt_llm import VisualGenArgs +from tensorrt_llm.visual_gen import AttentionConfig, QuantAttentionConfig + +args = VisualGenArgs( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + attention_config=AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype="int8", + v_dtype="fp8", + q_block_size=1, + k_block_size=16, + v_block_size=1, + ), + ), +) +``` + +```yaml +attention_config: + backend: TRTLLM + quant_attention_config: + qk_dtype: int8 + v_dtype: fp8 + q_block_size: 1 + k_block_size: 16 + v_block_size: 1 +``` + +## MXFP8 and NVFP4 in the `CUTEDSL` and `FlashInfer` backends + +**What it does.** MXFP8 and NVFP4 quantize Q/K with fixed blocks of 32 and 16 elements, respectively. `q_block_size` and `k_block_size` remain `0` because MXFP8 and NVFP4 follows specialized block-scaling schema which divide into both token dimensions and channel dimensions. V uses FP8 on `CUTEDSL` and FlashInfer SM10X, or NVFP4 on FlashInfer SM12X. + +**Requirements and behavior.** + +- `CUTEDSL` requires a head dimension of 128 and supports `v_block_size` of `0` (tensor-wide scale) or `1` (per-head, per-channel scale). +- FlashInfer SM100/SM103 requires a head dimension of 128 and supports MXFP8 or NVFP4 Q/K with FP8 V. +- FlashInfer SM120/SM121 supports NVFP4 Q/K/V and requires self-attention, equal Q/K/V shapes and head counts, a head dimension of 64 or 128, and a sequence length divisible by 128. + +**Configuration.** + +CUTEDSL with MXFP8 Q/K and FP8 V: + +```python +from tensorrt_llm import VisualGenArgs +from tensorrt_llm.visual_gen import AttentionConfig, QuantAttentionConfig + +args = VisualGenArgs( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + attention_config=AttentionConfig( + backend="CUTEDSL", + quant_attention_config=QuantAttentionConfig( + qk_dtype="mxfp8", + v_dtype="fp8", + q_block_size=0, + k_block_size=0, + v_block_size=1, + ), + ), +) +``` + +```yaml +attention_config: + backend: CUTEDSL + quant_attention_config: + qk_dtype: mxfp8 + v_dtype: fp8 + q_block_size: 0 + k_block_size: 0 + v_block_size: 1 +``` + +FlashInfer with NVFP4 Q/K/V: + +```python +from tensorrt_llm import VisualGenArgs +from tensorrt_llm.visual_gen import AttentionConfig, QuantAttentionConfig + +args = VisualGenArgs( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + attention_config=AttentionConfig( + backend="FLASHINFER", + quant_attention_config=QuantAttentionConfig( + qk_dtype="nvfp4", + v_dtype="nvfp4", + q_block_size=0, + k_block_size=0, + v_block_size=0, + ), + ), +) +``` + +```yaml +attention_config: + backend: FLASHINFER + quant_attention_config: + qk_dtype: nvfp4 + v_dtype: nvfp4 + q_block_size: 0 + k_block_size: 0 + v_block_size: 0 +``` + +## Interaction With Other Features + +- **Linear-layer quantization** (`VisualGenArgs.quant_config`, e.g. FP8 block scales or NVFP4) is independent and can be combined with any attention recipe. +- **Sparse attention.** On `CUTEDSL`, quantized attention and Video Sparse Attention (VSA) are mutually exclusive and rejected by the validator. On `TRTLLM`, Skip Softmax uses the same backend and the SageAttention unit tests exercise the two together. +- **Parallelism.** SageAttention is covered by a multi-GPU Ulysses test (`tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py`). The CuTe DSL dense backend produces LSE, so it also composes with Attention2D / Ring context parallelism; the TRTLLM Sage path does not expose LSE through this wrapper. diff --git a/docs/source/visual-gen/features/sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md similarity index 97% rename from docs/source/visual-gen/features/sparse-attention.md rename to docs/source/features/visualgen-sparse-attention.md index 7edeff8fb63c..5ead2416d717 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -1,7 +1,7 @@ -# VisualGen Sparse Attention +# VisualGen Sparse Attention (Beta) ```{note} -This page is an unindexed draft until the VisualGen documentation hub is introduced. +This feature is in **beta** stage. APIs, supported models, and optimization options are actively evolving and may change in future releases. ``` - [Overview](#overview) @@ -210,7 +210,7 @@ attention_config: ### CUDA Graphs -`disabled_until_timestep` creates two sparse-attention phases when it is set: the high-timestep disabled phase and the enabled phase after the cutoff. VisualGen includes that phase in CUDA graph keys so graph capture does not reuse a graph across different Skip Softmax Attention settings. See [VisualGen CUDA Graphs](cuda-graph.md) for the general capture and replay design. +`disabled_until_timestep` creates two sparse-attention phases when it is set: the high-timestep disabled phase and the enabled phase after the cutoff. VisualGen includes that phase in CUDA graph keys so graph capture does not reuse a graph across different Skip Softmax Attention settings. See [VisualGen CUDA Graphs](visualgen-cuda-graph.md) for the general capture and replay design. Graphs are captured lazily. The first denoising step seen for a given tensor shape and sparse-attention phase captures a graph; later steps with the same shape and phase replay that graph. When denoising crosses the cutoff, the phase key changes, so VisualGen captures a second graph for the enabled phase instead of replaying the graph from the disabled phase. diff --git a/docs/source/index.rst b/docs/source/index.rst index aafc01a71b7d..de78a2d67dfe 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -88,6 +88,9 @@ Welcome to TensorRT LLM's Documentation! features/helix.md features/kv-cache-connector.md features/sparse-attention.md + features/visualgen-cuda-graph.md + features/visualgen-quantized-attention.md + features/visualgen-sparse-attention.md .. toctree:: diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index e3bf2e99bc2a..18356a2ac20e 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -14,8 +14,8 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode - A shared pipeline abstraction covering the denoising loop, guidance strategies, and component loading. - Pluggable attention backends: PyTorch SDPA (`VANILLA`), TRT-LLM kernels (`TRTLLM`), FlashInfer FP16/BF16 dense prefill (`FLASHINFER`), TRT-LLM CuTe DSL kernels (`CUTEDSL`, Blackwell-class GPUs), and Flash Attention 4 (`FA4`). - Quantization support (dynamic and static) using the [ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer) configuration format. -- Quantized attention support: `QK16PV8` to quantize Bmm2 on `CUTEDSL`, `SAGE` to run SageAttention on `TRTLLM` (requires Blackwell SM100), and FlashInfer block-scaled MXFP8/NVFP4 on `FLASHINFER`. -- Sparse attention support: see [VisualGen Sparse Attention](../visual-gen/features/sparse-attention.md). +- Quantized attention support: see [VisualGen Quantized Attention](../features/visualgen-quantized-attention.md). +- Sparse attention support: see [VisualGen Sparse Attention](../features/visualgen-sparse-attention.md). - Multi-GPU parallelism (CFG parallel, Ulysses sequence parallel, Tensor parallelism). - **Step caching** — two runtime caching backends (**TeaCache** and **Cache-DiT**) that skip transformer computation on steps where the step-to-step change is small. - CPU offloading to reduce peak GPU memory usage. @@ -213,71 +213,11 @@ By default, `strict=True` raises when adapter tensors cannot be matched, have un ### Quantized Attention -In addition to linear-layer quantization, VisualGen supports several **attention-level** quantization recipes, which define the Q/K/V data types and scaling granularity used inside the attention kernel. They are configured through `AttentionConfig.quant_attention_config` and are mutually exclusive with each other. - -- **QK16PV8** (`CUTEDSL` backend): Keeps Q & K in BF16 and quantizes only V to FP8 (E4M3, per-tensor), thus Bmm1 will be carried out in BF16 with Bmm2 in FP8. Targets Blackwell-class GPUs (`sm_100a` / `sm_103a`) with `head_dim = 128`. -- **SAGE** (`TRTLLM` backend): Quantizes Q, K, and V with per-block scaling factors. Q/K are stored as INT8 or FP8 (e4m3) and V as FP8 (e4m3); block sizes are tunable per axis (typically `(q, k, v) = (1, 4, 1)` for Wan-1.3B and `(1, 16, 1)` for larger Wan / FLUX checkpoints). Supported recipes are validated at runtime. - -- **FlashInfer block-scaled attention** (`FLASHINFER` backend): Uses FlashInfer's architecture-specific FMHA. On SM100/SM103 (B200/B300), set `qk_dtype` to `mxfp8` or `nvfp4` with `v_dtype: fp8`. On SM120/SM121 (RTX PRO 6000), only dense Q/K/V NVFP4 self-attention is supported, using `qk_dtype: nvfp4` and `v_dtype: nvfp4`; this path requires a sequence length divisible by 128. - -Python API for SageAttention: - -```python -from tensorrt_llm import VisualGenArgs - -args = VisualGenArgs( - model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - attention_config={ - "backend": "TRTLLM", - "quant_attention_config": { - "qk_dtype": "int8", - "q_block_size": 1, - "k_block_size": 16, - "v_block_size": 1, - }, - }, -) -``` - -Python API for QK16PV8: - -```python -from tensorrt_llm import VisualGenArgs - -args = VisualGenArgs( - model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - attention_config={ - "backend": "CUTEDSL", - "quant_attention_config": { - "qk_dtype": "bf16", - "q_block_size": 0, - "k_block_size": 0, - "v_block_size": 0, - }, - }, -) -``` - -Python API for FlashInfer MXFP8 on B200/B300: - -```python -from tensorrt_llm import VisualGenArgs - -args = VisualGenArgs( - model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - attention_config={ - "backend": "FLASHINFER", - "quant_attention_config": { - "qk_dtype": "mxfp8", - "v_dtype": "fp8", - }, - }, -) -``` +In addition to linear-layer quantization, VisualGen exposes multiple backend-specific **quantized-attention** recipes that operate inside the attention kernel. They are configured through `AttentionConfig.quant_attention_config` and can be enabled independently with any linear layer configuration. See [VisualGen Quantized Attention](../features/visualgen-quantized-attention.md) for the full recipe table, the V scale-granularity trade-off, and the block-scaled MXFP8 / NVFP4 recipes. ### CUDA Graphs -VisualGen CUDA graphs capture transformer forward calls during denoising and replay them for later steps with compatible inputs. See [VisualGen CUDA Graphs](../visual-gen/features/cuda-graph.md) for capture scope, graph keys, and sparse-attention phase behavior. +VisualGen CUDA graphs capture transformer forward calls during denoising and replay them for later steps with compatible inputs. See [VisualGen CUDA Graphs](../features/visualgen-cuda-graph.md) for capture scope, graph keys, and sparse-attention phase behavior. ### Step Caching