[TRTLLM-14616][feat] add VisualGen fp8 row-wise quant

tensorrt_llm/_torch/modules/linear.py

@@ -983,6 +983,12 @@ def create_weights(self, module: Linear, in_features: int,
 
     def apply(self, module: Linear, input: torch.Tensor,
               bias: Optional[torch.Tensor]):
+        # Handle multi-dimensional inputs (e.g., 3D: batch, seq, hidden).
+        # fp8_rowwise_gemm requires a 2D mat1; flatten here and unflatten the output below.
+        orig_shape = input.shape
+        if input.dim() > 2:
+            input = input.reshape(-1, input.shape[-1])
+
         # FP8 tensor inputs are from attention. Directly use ones as scale.
         if input.dtype == torch.float8_e4m3fn:
             qinput = input
@@ -1016,6 +1022,8 @@ def apply(self, module: Linear, input: torch.Tensor,
         )
         if bias is not None:
             output = output + bias
+        if len(orig_shape) > 2:
+            output = output.reshape(*orig_shape[:-1], output.shape[-1])
         return output
 
     def _get_scale_name(self, weights: List[Dict]):

tensorrt_llm/_torch/visual_gen/config.py

@@ -272,6 +272,7 @@ def load_diffusion_quant_config(
             algo_map = {
                 "FP8": QuantAlgo.FP8,
                 "FP8_BLOCK_SCALES": QuantAlgo.FP8_BLOCK_SCALES,
+                "FP8_PER_CHANNEL_PER_TOKEN": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN,
                 "NVFP4": QuantAlgo.NVFP4,
                 "W4A16_AWQ": QuantAlgo.W4A16_AWQ,
                 "W4A8_AWQ": QuantAlgo.W4A8_AWQ,

tensorrt_llm/_torch/visual_gen/quantization/__init__.py

@@ -1,15 +1,16 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 """
 Quantization support for diffusion models.
 """
 
 from .loader import DynamicLinearWeightLoader
-from .ops import quantize_fp8_blockwise, quantize_fp8_per_tensor
+from .ops import quantize_fp8_blockwise, quantize_fp8_per_tensor, quantize_fp8_rowwise
 
 __all__ = [
     "DynamicLinearWeightLoader",
     "quantize_fp8_per_tensor",
     "quantize_fp8_blockwise",
+    "quantize_fp8_rowwise",
 ]

tensorrt_llm/_torch/visual_gen/quantization/loader.py

@@ -13,6 +13,7 @@
 from tensorrt_llm._torch.visual_gen.quantization.ops import (
     quantize_fp8_blockwise,
     quantize_fp8_per_tensor,
+    quantize_fp8_rowwise,
     quantize_nvfp4,
 )
 from tensorrt_llm.quantization.mode import QuantAlgo
@@ -152,7 +153,11 @@ def _should_dynamic_quantize(
             return False
 
         # For FP8 algorithms: quantize if weight is high precision
-        if quant_algo in (QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES):
+        if quant_algo in (
+            QuantAlgo.FP8,
+            QuantAlgo.FP8_BLOCK_SCALES,
+            QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN,
+        ):
             if weight.dtype == torch.float8_e4m3fn and "weight_scale" in weight_dict:
                 return False  # Already quantized
             return weight.dtype in (torch.bfloat16, torch.float16, torch.float32)
@@ -185,6 +190,9 @@ def _maybe_dynamic_quantize(
             block_size = self.quant_config.group_size if self.quant_config else 128
             qweight, scale = quantize_fp8_blockwise(weight, block_size=block_size)
             return {**weight_dict, "weight": qweight, "weight_scale": scale}
+        elif quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN:
+            qweight, scale = quantize_fp8_rowwise(weight)
+            return {**weight_dict, "weight": qweight, "weight_scale": scale}
         elif quant_algo == QuantAlgo.NVFP4:
             qweight, weight_scale, weight_scale_2 = quantize_nvfp4(weight)
             return {

tensorrt_llm/_torch/visual_gen/quantization/ops.py

@@ -99,6 +99,24 @@ def quantize_fp8_blockwise(
     return qweight, block_scales
 
 
+def quantize_fp8_rowwise(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Quantize weight to FP8 E4M3 with per-channel (per-output-row) scale.
+
+    Uses torch.ops.tensorrt_llm.quantize_e4m3_activation CUDA kernel.
+
+    Args:
+        weight: Input weight tensor (BF16/FP16/FP32), shape (out_features, in_features)
+
+    Returns:
+        Tuple of:
+            - qweight: Quantized weight (FP8 E4M3), shape (out_features, in_features)
+            - weight_scale: Per-channel dequantization scales (FP32), shape (out_features,)
+    """
+    qweight, scale = torch.ops.tensorrt_llm.quantize_e4m3_activation(weight.contiguous())
+    return qweight, scale.squeeze(1).to(torch.float32)
+
+
 def quantize_nvfp4(
     weight: torch.Tensor, block_size: int = 16
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

@@ -435,7 +435,7 @@ def test_load_weights_and_forward(self, cosmos3_transformer):
             )
         _assert_finite_output(out.video, hs.shape)
 
-    @pytest.mark.parametrize("quant_algo", ["FP8"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_load_fp8_quantization(self, quant_algo: str):
         checkpoint_dir = _require_checkpoint()
         if not torch.cuda.is_available():

tests/unittest/_torch/visual_gen/test_flux_pipeline.py

@@ -226,7 +226,7 @@ class TestFluxQuantization:
     """Test FLUX quantization loading and FP8 weight verification."""
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo: str):
         """Test loading FLUX.1 with FP8 quantization and verify FP8 weights."""
         args = VisualGenArgs(
@@ -253,6 +253,14 @@ def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo:
                             assert hasattr(module, "weight_scale"), (
                                 f"Linear {name} missing weight_scale"
                             )
+                            if quant_algo == "FP8_PER_CHANNEL_PER_TOKEN":
+                                assert module.weight_scale.dim() == 1, (
+                                    f"Linear {name} rowwise weight_scale should be 1-D, "
+                                    f"got shape {tuple(module.weight_scale.shape)}"
+                                )
+                                assert module.weight_scale.shape[0] == module.weight.shape[0], (
+                                    f"Linear {name} weight_scale length mismatch"
+                                )
                             found_fp8 = True
                             print(
                                 f"\n[{quant_algo}] FP8 layer {name}: weight {module.weight.shape}"
@@ -267,7 +275,7 @@ def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo:
         torch.cuda.empty_cache()
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_load_flux2_with_quantization(self, flux2_checkpoint_exists, quant_algo: str):
         """Test loading FLUX.2 with FP8 quantization and verify FP8 weights."""
         args = VisualGenArgs(
@@ -293,6 +301,14 @@ def test_load_flux2_with_quantization(self, flux2_checkpoint_exists, quant_algo:
                             assert hasattr(module, "weight_scale"), (
                                 f"Linear {name} missing weight_scale"
                             )
+                            if quant_algo == "FP8_PER_CHANNEL_PER_TOKEN":
+                                assert module.weight_scale.dim() == 1, (
+                                    f"Linear {name} rowwise weight_scale should be 1-D, "
+                                    f"got shape {tuple(module.weight_scale.shape)}"
+                                )
+                                assert module.weight_scale.shape[0] == module.weight.shape[0], (
+                                    f"Linear {name} weight_scale length mismatch"
+                                )
                             found_fp8 = True
                             print(
                                 f"\n[{quant_algo}] FP8 layer {name}: weight {module.weight.shape}"
@@ -316,7 +332,7 @@ class TestFluxFP8NumericalCorrectness:
     """Test FP8 vs BF16 numerical accuracy at single-layer and full-transformer levels."""
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str):
         """Test FP8 vs BF16 numerical accuracy on a single Linear layer.
 
@@ -399,7 +415,7 @@ def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str
         torch.cuda.empty_cache()
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_fp8_vs_bf16_full_transformer_e2e(self, flux1_checkpoint_exists, quant_algo: str):
         """End-to-end test: Compare full FLUX.1 transformer FP8 vs BF16 output.
 

tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py

@@ -200,7 +200,7 @@ class TestLTX2Quantization:
     """Test LTX2 quantization loading and FP8 weight verification."""
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_load_with_quantization(self, ltx2_bf16_checkpoint_exists, quant_algo: str):
         """Test loading LTX2 with FP8 quantization and verify FP8 weights."""
         args = VisualGenArgs(
@@ -250,7 +250,7 @@ class TestLTX2FP8NumericalCorrectness:
     """Test FP8 vs BF16 numerical accuracy at single-layer level."""
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_fp8_vs_bf16_single_layer(self, ltx2_bf16_checkpoint_exists, quant_algo: str):
         """Test FP8 vs BF16 numerical accuracy on a single Linear layer.
 
@@ -1599,7 +1599,7 @@ def test_two_stage_lora_deltas_match_transformer(self, ltx2_two_stage_assets_exi
             torch.cuda.empty_cache()
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"])
+    @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"])
     def test_two_stage_with_quantization(self, ltx2_two_stage_assets_exist, quant_algo: str):
         """Two-stage pipeline loads correctly with FP8 quantization."""
         from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import (

tests/unittest/_torch/visual_gen/test_model_loader.py

@@ -199,6 +199,54 @@ def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_exists):
     assert found_fp8_linear, "No FP8 Linear modules found in transformer"
 
 
+def test_load_wan_pipeline_with_fp8_rowwise(checkpoint_exists):
+    """Test loading with FP8 row-wise (per-channel-per-token) dynamic quantization.
+
+    Verifies:
+    1. Config parses FP8_PER_CHANNEL_PER_TOKEN and sets dynamic_weight_quant=True
+    2. Linear weights are FP8 after loading
+    3. weight_scale is 1-D [out_features] — one scale per output row, not a scalar
+    """
+    if not checkpoint_exists:
+        pytest.skip("Checkpoint not available")
+
+    from tensorrt_llm._torch.modules.linear import Linear
+    from tensorrt_llm._torch.visual_gen import PipelineLoader
+    from tensorrt_llm.visual_gen.args import VisualGenArgs
+
+    args = VisualGenArgs(
+        model=CHECKPOINT_PATH,
+        quant_config={"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN", "dynamic": True},
+    )
+    pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS)
+
+    assert pipeline.pipeline_config.dynamic_weight_quant is True
+
+    found_fp8_linear = False
+    for name, module in pipeline.transformer.named_modules():
+        if isinstance(module, Linear):
+            if hasattr(module, "weight") and module.weight is not None:
+                assert module.weight.dtype == torch.float8_e4m3fn, (
+                    f"Linear {name} weight dtype is {module.weight.dtype}, expected float8_e4m3fn"
+                )
+                assert hasattr(module, "weight_scale") and module.weight_scale is not None, (
+                    f"Linear {name} missing weight_scale"
+                )
+                # Per-channel: one scale per output neuron, not a scalar
+                assert module.weight_scale.dim() == 1, (
+                    f"Linear {name} weight_scale should be 1-D [out_features], "
+                    f"got shape {tuple(module.weight_scale.shape)}"
+                )
+                assert module.weight_scale.shape[0] == module.weight.shape[0], (
+                    f"Linear {name} weight_scale length {module.weight_scale.shape[0]} "
+                    f"!= out_features {module.weight.shape[0]}"
+                )
+                found_fp8_linear = True
+                break
+
+    assert found_fp8_linear, "No FP8 Linear modules found in transformer"
+
+
 def test_load_wan_pipeline_with_fp8_blockwise(checkpoint_exists):
     """Test loading with FP8 blockwise quantization using VisualGenArgs."""
     if not checkpoint_exists:
@@ -261,6 +309,15 @@ def test_visual_gen_args_to_quant_config():
     qc, _, _, _ = parse(args.quant_config)
     assert qc.quant_algo == QuantAlgo.FP8_BLOCK_SCALES
 
+    # FP8 rowwise (per-token activations, per-channel weights)
+    args = VisualGenArgs(
+        model="/fake/path",
+        quant_config={"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN", "dynamic": True},
+    )
+    qc, _, dwq, _ = parse(args.quant_config)
+    assert qc.quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN
+    assert dwq is True
+
     # NVFP4
     args = VisualGenArgs(
         model="/fake/path",

tests/unittest/_torch/visual_gen/test_quant_ops.py

@@ -7,6 +7,7 @@
 from tensorrt_llm._torch.visual_gen.quantization.ops import (
     quantize_fp8_blockwise,
     quantize_fp8_per_tensor,
+    quantize_fp8_rowwise,
 )
 
 
@@ -115,6 +116,43 @@ def test_fp8_blockwise_zero_weight(self):
         self.assertEqual(qweight.dtype, torch.float8_e4m3fn)
         self.assertTrue(torch.all(qweight.to(torch.float32) == 0))
 
+    def test_fp8_rowwise(self):
+        """Test FP8 rowwise (per-channel) quantization."""
+        out_features, in_features = 256, 512
+        weight = torch.randn(out_features, in_features, dtype=torch.bfloat16, device="cuda")
+        qweight, scale = quantize_fp8_rowwise(weight)
+
+        self.assertEqual(qweight.dtype, torch.float8_e4m3fn)
+        self.assertEqual(qweight.shape, weight.shape)
+        self.assertEqual(scale.dtype, torch.float32)
+        # One scale per output row
+        self.assertEqual(scale.shape, (out_features,))
+
+        # Dequantize and check round-trip error
+        dequant = qweight.to(torch.float32) * scale.unsqueeze(1)
+        error = (dequant - weight.to(torch.float32)).abs().mean()
+        self.assertLess(error, 0.15)
+
+    def test_fp8_rowwise_different_shapes(self):
+        """Test FP8 rowwise quantization with various shapes."""
+        shapes = [(128, 256), (256, 512), (512, 1024)]
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                weight = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+                qweight, scale = quantize_fp8_rowwise(weight)
+
+                self.assertEqual(qweight.dtype, torch.float8_e4m3fn)
+                self.assertEqual(qweight.shape, weight.shape)
+                self.assertEqual(scale.shape, (shape[0],))
+
+    def test_fp8_rowwise_zero_weight(self):
+        """Test FP8 rowwise quantization with zero weight."""
+        weight = torch.zeros(128, 256, dtype=torch.bfloat16, device="cuda")
+        qweight, scale = quantize_fp8_rowwise(weight)
+
+        self.assertEqual(qweight.dtype, torch.float8_e4m3fn)
+        self.assertTrue(torch.all(qweight.to(torch.float32) == 0))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py

@@ -170,6 +170,13 @@ def test_qwen_pipeline_feature_args(args, expected):
             True,
             id="dynamic-fp4",
         ),
+        pytest.param(
+            {"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN", "dynamic": True},
+            QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN,
+            None,
+            False,
+            id="dynamic-fp8-rowwise",
+        ),
     ],
 )
 def test_qwen_pipeline_quant_config_parses_from_args(