diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 240c217a044e..9777ee4508ce 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/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]): diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index c636c618a0e2..a4fd4fe4c3bf 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/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, diff --git a/tensorrt_llm/_torch/visual_gen/quantization/__init__.py b/tensorrt_llm/_torch/visual_gen/quantization/__init__.py index 909629b1b361..ef4494e13573 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/quantization/__init__.py @@ -1,4 +1,4 @@ -# 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 """ @@ -6,10 +6,11 @@ """ 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", ] diff --git a/tensorrt_llm/_torch/visual_gen/quantization/loader.py b/tensorrt_llm/_torch/visual_gen/quantization/loader.py index 8b042ed25e38..b6bbcec94681 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/loader.py +++ b/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 { diff --git a/tensorrt_llm/_torch/visual_gen/quantization/ops.py b/tensorrt_llm/_torch/visual_gen/quantization/ops.py index ce376f01de0a..51c4de435e4e 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/ops.py +++ b/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]: diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 76e6c1a01a8c..09a8bdc89310 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/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(): diff --git a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py index 0455e25faff7..112b7096c7dd 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py +++ b/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. diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index d2bbf46d32a5..388478e805f6 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/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 ( diff --git a/tests/unittest/_torch/visual_gen/test_model_loader.py b/tests/unittest/_torch/visual_gen/test_model_loader.py index a7c969affee4..0026c186a39e 100644 --- a/tests/unittest/_torch/visual_gen/test_model_loader.py +++ b/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", diff --git a/tests/unittest/_torch/visual_gen/test_quant_ops.py b/tests/unittest/_torch/visual_gen/test_quant_ops.py index 5b141ee74db7..dfd938dabda2 100644 --- a/tests/unittest/_torch/visual_gen/test_quant_ops.py +++ b/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() diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py index 0a8ef1900479..043d9ba09c67 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py +++ b/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(