[TRTLLM-11484][feat] VisualGen fmha.py support per-channel V scaling

docs/source/models/visual-generation.md

@@ -14,7 +14,7 @@ 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`), 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).
+- Quantized attention support: see [VisualGen Quantized Attention](../visual-gen/features/quantized-attention.md).
 - Sparse attention support: see [VisualGen Sparse Attention](../visual-gen/features/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.
@@ -169,49 +169,7 @@ By default, `strict=True` raises when adapter tensors cannot be matched, have un
 
 ### Quantized Attention
 
-In addition to linear-layer quantization, VisualGen exposes two **attention-level** quantization presets that operate 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.
-
-
-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,
-        },
-    },
-)
-```
+In addition to linear-layer quantization, VisualGen exposes two **attention-level** quantization presets 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](../visual-gen/features/quantized-attention.md) for the full recipe table, the V scale-granularity trade-off, and the block-scaled MXFP8 / NVFP4 recipes.
 
 ### CUDA Graphs
 

docs/source/visual-gen/features/quantized-attention.md

@@ -0,0 +1,142 @@
+# VisualGen Quantized Attention
+
+```{note}
+This page is an unindexed draft until the VisualGen documentation hub is introduced.
+```
+
+- [Overview](#overview)
+  - [Recipes](#recipes)
+  - [Configuration Surface](#configuration-surface)
+- [QK16PV8 (CUTEDSL)](#qk16pv8-cutedsl)
+- [SageAttention (TRTLLM)](#sageattention-trtllm)
+- [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, 1)`, `(0, 0, 0)` | QK16PV8 |
+| `CUTEDSL` | `mxfp8` | `fp8` | `(0, 0, 0)`, `(0, 0, 1)` | Block-scaled MXFP8 Q/K |
+| `CUTEDSL` | `nvfp4` | `fp8` | `(0, 0, 0)`, `(0, 0, 1)` | Block-scaled NVFP4 Q/K |
+
+Notes:
+
+- `v_dtype` only accepts `fp8`. Every quantized-attention kernel currently loads V as FP8 (e4m3), so BMM2 is always FP8; the recipes differ in how BMM1 is handled.
+- `qk_dtype: "bf16"` means Q/K are **not** quantized — BMM1 stays in BF16. It's recommended to set `v_block_size` to 1 for this kernel for both accuracy and performance.
+- `quant_attention_config` requires `backend` to be `TRTLLM` or `CUTEDSL`.
+
+### 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"` | `"fp8"` | V element format for BMM2 (FP8 e4m3). |
+| `q_block_size` | int ≥ 0 | `0` | Q tokens per SageAttention quantization block. `0` on the CuTe DSL paths. |
+| `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` on the CuTe DSL paths. |
+| `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`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`.
+
+## QK16PV8 (CUTEDSL)
+
+**What it does.** Q and K stay in BF16 (or FP16), so BMM1 runs at full input precision. Only V is quantized to FP8 e4m3, so BMM2 runs on FP8 Tensor Cores. `v_block_size` selects how V is scaled:
+
+- **`v_block_size: 1` (Recommended)** — one scale per KV head and channel.
+- **`v_block_size: 0`** — a single per-tensor scale, folded into the kernel's `scale_output` scalar.
+
+```{note}
+Prefer `v_block_size: 1`. With `v_block_size: 0` the per-tensor scale is a *device* scalar, so folding it into `scale_output` requires reading it back to the host (`.item()`) inside `cute_dsl_fmha_fwd`. That readback drains the pipeline once per attention call, adding a device-host synchronization overhead. `v_block_size: 1` avoids the readback entirely, and its amax is additionally cheaper because a reduction to `(H, D)` parallelizes better than a reduction to one scalar. Per-head-per-channel scaling is also finer, so it's recommended to set `v_block_size` to `1` for both accuracy and performance.
+```
+
+**Configuration.**
+
+```python
+from tensorrt_llm import VisualGenArgs
+from tensorrt_llm.visual_gen import AttentionConfig, QuantAttentionConfig
+
+args = VisualGenArgs(
+    model="<path_or_hf_id>",
+    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=1,
+        ),
+    ),
+)
+```
+
+```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: 1
+```
+
+## SageAttention (TRTLLM)
+
+**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.**
+
+- Blackwell GPU.
+- The recommended `qk_dtype: "int8"` is only supported on `sm_100a`.
+  - `sm_103a` can use `qk_dtype: "fp8"` but its accuracy could be worse than `qk_dtype: "int8"`.
+
+**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
+```
+
+## 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.

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py

@@ -298,8 +298,6 @@ def cute_dsl_fmha_fwd(
             raise ValueError("Block-scaled path (qk_sf_vec != 0) requires q_sf and k_sf tensors.")
         if not q_sf.is_contiguous() or not k_sf.is_contiguous():
             raise ValueError("q_sf and k_sf must be contiguous.")
-    elif scale_v_channels is not None:
-        raise ValueError("scale_v_channels is only supported by MXFP8 and NVFP4 kernels.")
 
     if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
         raise ValueError("Q, K, and V must be contiguous before the CuTe DSL launch boundary.")
@@ -399,15 +397,15 @@ def _scalar_float(t):
         sf_dtype = cutlass.Float8E8M0FNU if qk_sf_vec == 32 else cutlass.Float8E4M3FN
         q_sf_cute = _to_cute_tensor(q_sf, leading_dim=0, cutlass_element_type=sf_dtype)
         k_sf_cute = _to_cute_tensor(k_sf, leading_dim=0, cutlass_element_type=sf_dtype)
-        scale_v_channels_cute = (
-            _to_cute_tensor(scale_v_channels.view(-1), leading_dim=0)
-            if scale_v_channels is not None
-            else None
-        )
     else:
         q_sf_cute = None
         k_sf_cute = None
-        scale_v_channels_cute = None
+    # Per-channel V dequant scales. Supported by both the block-scaled and the dense kernel;
+    scale_v_channels_cute = (
+        _to_cute_tensor(scale_v_channels.view(-1), leading_dim=0)
+        if scale_v_channels is not None
+        else None
+    )
     # lse_4d is (B, S_q, h_kv, h_r) contiguous → h_r is the stride-1 inner dim (index 3).
     lse_cute = (
         from_dlpack(lse_4d, assumed_align=16).mark_layout_dynamic(leading_dim=3)
@@ -501,6 +499,7 @@ def _scalar_float(t):
             cute_typing.Float32(scale_softmax_log2),
             cute_typing.Float32(scale_softmax),
             cute_typing.Float32(scale_output),
+            scale_v_channels_cute,
             skip_threshold_log2,
             ws_left,
             ws_right,
@@ -528,11 +527,11 @@ def _quantize_fp8_v(
 ) -> Tuple[torch.Tensor, float | torch.Tensor, torch.Tensor | None]:
     """Quantize V to FP8 with either one tensor scale or an (H, D) scale tensor."""
     if per_head_channel:
-        v_qscale = _FP8_E4M3_MAX / v_bshd.float().abs().amax(dim=(0, 1)).clamp(min=1e-3)
+        v_qscale = _FP8_E4M3_MAX / v_bshd.abs().amax(dim=(0, 1)).float().clamp(min=1e-3)
         v_quantized = (v_bshd * v_qscale).to(torch.float8_e4m3fn)
         return v_quantized, 1.0, v_qscale.reciprocal().contiguous()
 
-    v_qscale = _FP8_E4M3_MAX / v_bshd.abs().amax().clamp(min=1e-3)
+    v_qscale = _FP8_E4M3_MAX / v_bshd.abs().amax().float().clamp(min=1e-3)
     v_quantized = (v_bshd * v_qscale).to(torch.float8_e4m3fn)
     return v_quantized, v_qscale.reciprocal(), None
 
@@ -636,6 +635,12 @@ def __init__(
         self.num_kv_heads = num_kv_heads or num_heads
         self.dtype = dtype
         self.quant_attention_config = quant_attention_config
+        if quant_attention_config is not None and quant_attention_config.v_block_size not in (0, 1):
+            raise NotImplementedError(
+                "CuTeDSLAttention supports v_block_size == 0 for per-tensor quantization; "
+                "and v_block_size == 1 for per-channel quantization. "
+                f"Found unsupported value: {quant_attention_config.v_block_size}."
+            )
         if skip_softmax_threshold_scale is not None and sparse_params is not None:
             raise ValueError("Set either skip_softmax_threshold_scale or sparse_params, not both.")
         self.skip_softmax_threshold_scale = skip_softmax_threshold_scale
@@ -727,7 +732,8 @@ def _fwd(
             device=q.device,
         )
 
-        # V is tensor-scaled by default. MXFP8/NVFP4 with v_block_size=1 use an (H, D) scale.
+        # Set v_block_size=1 use an (H, D) scale for V (recommended);
+        # Set v_block_size=0 to per-tensor quantize V (introduces extra device-host sync);
         scale_v = kwargs.get("scale_v", 1.0)
         scale_q = kwargs.get("scale_q", 1.0)
         scale_k = kwargs.get("scale_k", 1.0)
@@ -744,7 +750,7 @@ def _fwd(
                 scale_k = scale_k * gs_k
                 qk_cutlass_dtype = cutlass.Float4E2M1FN if qk_sf_vec == 16 else cutlass.Float8E4M3FN
             v, v_dequant_scale, scale_v_channels = _quantize_fp8_v(
-                v, per_head_channel=qk_sf_vec != 0 and qac.v_block_size == 1
+                v, per_head_channel=qac.v_block_size == 1
             )
             scale_v = scale_v * v_dequant_scale