Fix per-tensor FP8 weight dequantization

modelopt/torch/quantization/plugins/huggingface.py

@@ -1330,13 +1330,16 @@ class _QuantFP8Linear(QuantModule):
     def _setup(self):
         self.input_quantizer = TensorQuantizer()
         self.weight_quantizer = TensorQuantizer()
-        assert self.weight_scale_inv.ndim == 2, "Weight scale inverse must be 2D"
         assert self.weight.ndim == 2, "Weight must be 2D"
-        self.block_size = max(
-            self.weight.shape[0] // self.weight_scale_inv.shape[0],
-            self.weight.shape[1] // self.weight_scale_inv.shape[1],
-        )
-        assert self.block_size == 128, "Block size must be 128"
+        if self.weight_scale_inv.ndim == 0:
+            self.block_size = None
+        else:
+            assert self.weight_scale_inv.ndim == 2, "Weight scale inverse must be 0D or 2D"
+            self.block_size = max(
+                self.weight.shape[0] // self.weight_scale_inv.shape[0],
+                self.weight.shape[1] // self.weight_scale_inv.shape[1],
+            )
+            assert self.block_size == 128, "Block size must be 128"
 
     def _get_weight_and_scale_inv(self):
         if isinstance(self.weight, torch.distributed.tensor.DTensor):
@@ -1347,12 +1350,17 @@ def _get_weight_and_scale_inv(self):
             scale_inv = self.weight_scale_inv.contiguous()
         return weight, scale_inv
 
-    def forward(self, input: Tensor) -> Tensor:
+    def _dequantize_weight(self, dtype: torch.dtype) -> Tensor:
+        weight, scale_inv = self._get_weight_and_scale_inv()
+        if self.block_size is None:
+            return weight.to(dtype) * scale_inv.to(dtype)
         assert weight_dequant is not None, "Triton is not available"
+        return weight_dequant(weight, scale_inv, self.block_size, dtype=dtype)
+
+    def forward(self, input: Tensor) -> Tensor:
         if self.weight.element_size() == 1:
             with torch.cuda.device(self.weight.device):
-                weight, scale_inv = self._get_weight_and_scale_inv()
-                weight = weight_dequant(weight, scale_inv, self.block_size, dtype=input.dtype)
+                weight = self._dequantize_weight(input.dtype)
         else:
             weight = self.weight
         return linear(
@@ -1362,11 +1370,9 @@ def forward(self, input: Tensor) -> Tensor:
         )
 
     def unpack_weight(self):
-        assert weight_dequant is not None, "Triton is not available"
         with torch.cuda.device(self.weight.device):
-            weight, scale_inv = self._get_weight_and_scale_inv()
             self.weight = nn.Parameter(
-                weight_dequant(weight, scale_inv, self.block_size, dtype=torch.get_default_dtype()),
+                self._dequantize_weight(torch.get_default_dtype()),
                 requires_grad=False,
             )
         if hasattr(self, "weight_scale_inv"):

modelopt_recipes/huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib.yaml

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  nvfp4: configs/numerics/nvfp4
+  fp8: configs/numerics/fp8
+  kv_fp8: configs/ptq/units/kv_fp8
+
+metadata:
+  recipe_type: ptq
+  description: >-
+    NVFP4 W4A4 on interior MLP layers, FP8 W8A8 on edge MLP and attention
+    layers, and an FP8 KV cache; uses max calibration.
+quantize:
+  algorithm:
+    method: max
+    layerwise: false
+  quant_cfg:
+    - $import: base_disable_all
+    # NVFP4 on MLP and expert weight matrices.
+    - quantizer_name: '*mlp.experts*weight_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*mlp.experts*input_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*block_sparse_moe*weight_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*block_sparse_moe*input_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*mlp*weight_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*mlp*input_quantizer'
+      cfg:
+        $import: nvfp4
+    # FP8 on language-model attention projections.
+    - quantizer_name: '*self_attn.q_proj*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.q_proj*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.k_proj*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.k_proj*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.v_proj*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.v_proj*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.o_proj*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*self_attn.o_proj*input_quantizer'
+      cfg:
+        $import: fp8
+    # Keep the first four and final decoder MLP layers in FP8.
+    - quantizer_name: '*layers.0.mlp*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.0.mlp*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.1.mlp*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.1.mlp*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.2.mlp*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.2.mlp*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.3.mlp*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.3.mlp*input_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.87.mlp*weight_quantizer'
+      cfg:
+        $import: fp8
+    - quantizer_name: '*layers.87.mlp*input_quantizer'
+      cfg:
+        $import: fp8
+    - $import: kv_fp8
+    - $import: default_disabled_quantizers

tests/unit/recipe/test_loader.py

@@ -166,6 +166,7 @@ def test_load_recipe_builtin_description():
     "general/ptq/nvfp4_experts_only-kv_fp8",
     "general/ptq/nvfp4_experts_only-kv_fp8_cast",
     "general/ptq/nvfp4_experts_only-kv_fp8_layerwise",
+    "huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib",
     "general/ptq/nvfp4_mlp_only-kv_fp8",
     "general/ptq/nvfp4_mlp_only-novit-kv_fp8",
     "general/ptq/nvfp4_mlp_only-kv_fp8_cast",

tests/unit/torch/quantization/plugins/test_huggingface.py

@@ -45,6 +45,7 @@
 
 import transformers
 from transformers import AutoModelForCausalLM, LlamaForCausalLM
+from transformers.integrations.finegrained_fp8 import FP8Linear
 from transformers.models.dbrx.configuration_dbrx import DbrxConfig, DbrxFFNConfig
 from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN
 
@@ -109,6 +110,21 @@ def test_convert_conv1d():
     assert torch.allclose(out_1, out_2)
 
 
+def test_fp8_linear_per_tensor_dequant(monkeypatch):
+    module = FP8Linear(2, 2, block_size=(128, 128))
+    module.weight_scale_inv = nn.Parameter(torch.tensor(2.0))
+    with torch.no_grad():
+        module.weight.copy_(torch.tensor([[-2.0, 1.0], [0.5, 4.0]], dtype=torch.float8_e4m3fn))
+
+    mtq.replace_quant_module(module)
+    monkeypatch.setattr("modelopt.torch.quantization.plugins.huggingface.weight_dequant", None)
+
+    assert module.block_size is None
+    torch.testing.assert_close(
+        module._dequantize_weight(torch.float32), module.weight.float() * 2.0
+    )
+
+
 @pytest.mark.skipif(
     Version(transformers.__version__) < Version("5.0"),
     reason="test_dbrx is not supported for transformers<5.0",