diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index ee5185ccaa42..46ce34197c7e 100644 --- a/docs/source/models/visual-generation.md +++ b/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 diff --git a/docs/source/visual-gen/features/quantized-attention.md b/docs/source/visual-gen/features/quantized-attention.md new file mode 100644 index 000000000000..b39a8a613ec7 --- /dev/null +++ b/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="", + 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. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index 93aa23fe6809..611e716da62d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/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 diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py index bca82e086559..2f5cf6cc17e1 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py @@ -368,6 +368,7 @@ def __call__( scale_softmax_log2: Float32, scale_softmax: Float32, scale_output: Float32, + scale_v_channels: Optional[cute.Tensor], skip_softmax_threshold_log2: Optional[Float32], window_size_left: Optional[Int32], window_size_right: Optional[Int32], @@ -487,6 +488,17 @@ def __call__( else: sink = None + if cutlass.const_expr(scale_v_channels is not None): + # scale_v_channels is per (h_k, dv): shape (h_k * dv,) row-major. + # Expose as (dv, ((h_r, h_k), b)) where h_r and b are 0-stride broadcasts. + scale_v_channels_layout = cute.make_layout( + (dv, ((h_r, h_k), b)), + stride=(1, ((0, dv), 0)), + ) + mScaleV_channels = cute.make_tensor(scale_v_channels.iterator, scale_v_channels_layout) + else: + mScaleV_channels = None + self.tile_sched_params, grid = fmha_utils.compute_grid( cute.shape((s_q_max, d, ((h_r, h_k), b))), self.cta_tiler, @@ -686,6 +698,7 @@ class SharedStorage: scale_softmax_log2, scale_softmax, scale_output, + mScaleV_channels, skip_softmax_threshold_log2, window_size_left, window_size_right, @@ -728,6 +741,7 @@ def kernel( scale_softmax_log2: Float32, scale_softmax: Float32, scale_output: Float32, + mScaleV_channels: Optional[cute.Tensor], skip_softmax_threshold_log2: Optional[Float32], window_size_left: Optional[Int32], window_size_right: Optional[Int32], @@ -1625,6 +1639,7 @@ def make_tmem_tensors( sO[None, None, 0], mLSE, mSink, + mScaleV_channels, ), ( s0_corr_consumer, @@ -1647,6 +1662,7 @@ def make_tmem_tensors( sO[None, None, 1], mLSE, mSink, + mScaleV_channels, ), ( s1_corr_consumer, @@ -1688,6 +1704,7 @@ def make_tmem_tensors( gO0, mLSE, mSink, + mScaleV_channels, ), (s0_corr_consumer, mma_corr_consumer), (row_idx, *value_args), @@ -1704,6 +1721,7 @@ def make_tmem_tensors( gO1, mLSE, mSink, + mScaleV_channels, ), (s1_corr_consumer, mma_corr_consumer), (row_idx, *value_args), @@ -3002,7 +3020,7 @@ def correction_epilog( :type thr_mma: cute.ThrMma :param tiled_tmem_load_vec: Tiled memory load operation for the vectorized row-wise max :type tiled_tmem_load_vec: cute.TiledCopy - :param tensor_args: Tuple containing (tOtO, tTMEM_LOAD_VECtSi, tTMEM_LOAD_VECcS, sO_or_gO, mLSE, mSink) + :param tensor_args: Tuple containing (tOtO, tTMEM_LOAD_VECtSi, tTMEM_LOAD_VECcS, sO_or_gO, mLSE, mSink, mScaleV_channels) :type tensor_args: Tuple :param pipeline_args: When use_tma_store: (si_corr_consumer, mma_corr_consumer, corr_epi_producer). When not use_tma_store: (si_corr_consumer, mma_corr_consumer). @@ -3010,7 +3028,15 @@ def correction_epilog( :param value_args: Tuple containing (row_idx, cuseqlen_q, seqlen_q, blk_coord, scale_softmax, scale_output) :type value_args: Tuple """ - tOtO, tTMEM_LOAD_VECtSi, tTMEM_LOAD_VECcS, dest_O, mLSE, mSink = tensor_args + ( + tOtO, + tTMEM_LOAD_VECtSi, + tTMEM_LOAD_VECcS, + dest_O, + mLSE, + mSink, + mScaleV_channels, + ) = tensor_args row_idx, cuseqlen_q, seqlen_q, blk_coord, scale_softmax, scale_output = value_args pv_tiled_mma_shape = ( @@ -3076,16 +3102,27 @@ def correction_epilog( row_sum = row_sum + sink_exp scale = scale_output / row_sum + if cutlass.const_expr(mScaleV_channels is not None): + scaleV_ch_h = mScaleV_channels[None, blk_coord[2]] for i in range(self.cta_tiler[2] // corr_tile_size): tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i] tTMEM_LOADdO_i = tTMEM_LOADdO[None, 0, 0, i] - tTMrO = cute.make_rmem_tensor(tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype) + tTMEM_LOADoO_i = tTMEM_LOADoO[None, 0, 0, i] + tTMrO = cute.make_rmem_tensor(tTMEM_LOADoO_i.shape, self.pv_acc_dtype) cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) for j in range(0, cute.size(tTMrO), 2): tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( (tTMrO[j], tTMrO[j + 1]), (scale, scale), ) + if cutlass.const_expr(mScaleV_channels is not None): + for j in range(0, cute.size(tTMrO), 2): + _, n0 = tTMEM_LOADoO_i[j] + _, n1 = tTMEM_LOADoO_i[j + 1] + tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scaleV_ch_h[n0], scaleV_ch_h[n1]), + ) tDMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) o_vec = tTMrO.load() tDMrO.store(o_vec.to(self.o_dtype)) diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha_blockscaled.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha_blockscaled.py index 01be74595ac9..8c7208e65d47 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha_blockscaled.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha_blockscaled.py @@ -621,11 +621,9 @@ def __call__( (dv, ((h_r, h_k), b)), stride=(1, ((0, dv), 0)), ) - m_scale_v_channels = cute.make_tensor( - scale_v_channels.iterator, scale_v_channels_layout - ) + mScaleV_channels = cute.make_tensor(scale_v_channels.iterator, scale_v_channels_layout) else: - m_scale_v_channels = None + mScaleV_channels = None self.tile_sched_params, grid = fmha_utils.compute_grid( cute.shape((s_q_max, d, ((h_r, h_k), b))), @@ -898,7 +896,7 @@ class SharedStorage: scale_softmax_log2, scale_softmax, scale_output, - m_scale_v_channels, + mScaleV_channels, skip_softmax_threshold_log2, window_size_left, window_size_right, @@ -947,7 +945,7 @@ def kernel( scale_softmax_log2: Float32, scale_softmax: Float32, scale_output: Float32, - m_scale_v_channels: Optional[cute.Tensor], + mScaleV_channels: Optional[cute.Tensor], skip_softmax_threshold_log2: Optional[Float32], window_size_left: Optional[Int32], window_size_right: Optional[Int32], @@ -2063,7 +2061,7 @@ def make_tmem_tensors( sO[None, None, 0], mLSE, mSink, - m_scale_v_channels, + mScaleV_channels, ), ( s0_corr_consumer, @@ -2086,7 +2084,7 @@ def make_tmem_tensors( sO[None, None, 1], mLSE, mSink, - m_scale_v_channels, + mScaleV_channels, ), ( s1_corr_consumer, @@ -2128,7 +2126,7 @@ def make_tmem_tensors( gO0, mLSE, mSink, - m_scale_v_channels, + mScaleV_channels, ), (s0_corr_consumer, mma_corr_consumer), (row_idx, *value_args), @@ -2145,7 +2143,7 @@ def make_tmem_tensors( gO1, mLSE, mSink, - m_scale_v_channels, + mScaleV_channels, ), (s1_corr_consumer, mma_corr_consumer), (row_idx, *value_args), @@ -3533,7 +3531,7 @@ def correction_epilog( dest_O, mLSE, mSink, - m_scale_v_channels, + mScaleV_channels, ) = tensor_args row_idx, cuseqlen_q, seqlen_q, blk_coord, scale_softmax, scale_output = value_args @@ -3600,8 +3598,8 @@ def correction_epilog( row_sum = row_sum + sink_exp scale = scale_output / row_sum - if cutlass.const_expr(m_scale_v_channels is not None): - scale_v_ch_h = m_scale_v_channels[None, blk_coord[2]] + if cutlass.const_expr(mScaleV_channels is not None): + scaleV_ch_h = mScaleV_channels[None, blk_coord[2]] for i in range(self.cta_tiler[2] // corr_tile_size): tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i] tTMEM_LOADdO_i = tTMEM_LOADdO[None, 0, 0, i] @@ -3613,13 +3611,13 @@ def correction_epilog( (tTMrO[j], tTMrO[j + 1]), (scale, scale), ) - if cutlass.const_expr(m_scale_v_channels is not None): + if cutlass.const_expr(mScaleV_channels is not None): for j in range(0, cute.size(tTMrO), 2): _, n0 = tTMEM_LOADoO_i[j] _, n1 = tTMEM_LOADoO_i[j + 1] tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( (tTMrO[j], tTMrO[j + 1]), - (scale_v_ch_h[n0], scale_v_ch_h[n1]), + (scaleV_ch_h[n0], scaleV_ch_h[n1]), ) tDMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) o_vec = tTMrO.load() diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index d314c92b24ae..da600c590514 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -82,7 +82,8 @@ class QuantAttentionConfig(StrictBaseModel): ge=0, status="prototype", description=( - "V quantization block size on the hidden dimension; 0 uses one tensor-wide V scale." + "V quantization block size on the hidden dimension; 1 for per-channel quantization; " + "0 for per-tensor quantization." ), ) @@ -132,6 +133,7 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": } CUTEDSL_RECIPES = { ("bf16", "fp8", (0, 0, 0)), + ("bf16", "fp8", (0, 0, 1)), ("mxfp8", "fp8", (0, 0, 0)), ("mxfp8", "fp8", (0, 0, 1)), ("nvfp4", "fp8", (0, 0, 0)),