Cleanup code in get_weights_scaling_factor_from_quantizer of MXFP8QTensor

Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>

Author: danisereb

modelopt/torch/quantization/qtensor/mxfp8_tensor.py

@@ -132,15 +132,6 @@ def get_weights_scaling_factor_from_quantizer(
             assert scale.dtype == cls.SCALE_DTYPE, (
                 f"MXFP8 scale must be {cls.SCALE_DTYPE} (E8M0 format), got {scale.dtype}"
             )
-
-            # Reshape if needed (same number of elements but wrong shape)
-            if scale.shape != expected_shape:
-                expected_numel = 1
-                for dim in expected_shape:
-                    expected_numel *= dim
-                if scale.numel() == expected_numel:
-                    scale = scale.reshape(expected_shape)
-
             assert scale.shape == expected_shape, (
                 f"Scale shape {scale.shape} does not match expected shape {expected_shape}"
             )
@@ -179,12 +170,6 @@ def quantize_with_scale(
             f"Weight inner dimension ({in_dim}) must be divisible by MXFP8 block size ({cls.BLOCK_SIZE})"
         )
 
-        # Reshape scale if needed (same number of elements but wrong shape)
-        expected_shape = (*weight.shape[:-1], num_blocks)
-        if e8m0_scale.shape != expected_shape:
-            if e8m0_scale.numel() == weight.numel() // cls.BLOCK_SIZE:
-                e8m0_scale = e8m0_scale.reshape(expected_shape)
-
         # Convert E8M0 biased exponent to scale factor: scale = 2^(127 - exponent)
         scale_factor = torch.exp2(127 - e8m0_scale.float())
 
@@ -258,13 +243,6 @@ def dequantize(self, dtype: torch.dtype = None, **kwargs) -> torch.Tensor:
         # Convert E8M0 biased exponent back to scale factor: descale = 2^(exponent - 127)
         descale = torch.exp2(e8m0_scale.float() - 127)
 
-        # Reshape descale to match blocked tensor for broadcasting
-        expected_scale_shape = (*quantized_data.shape[:-1], num_blocks)
-        if descale.shape != expected_scale_shape and descale.numel() == num_blocks * (
-            quantized_data.numel() // quantized_data.shape[-1]
-        ):
-            descale = descale.view(expected_scale_shape)
-
         dequantized = quantized_blocked * descale.unsqueeze(-1)
 
         # Reshape and crop back to original shape

tests/gpu/torch/quantization/test_qtensor_cuda.py

@@ -897,3 +897,15 @@ class MockQuantizer:
 
         with pytest.raises(AssertionError, match="Scale shape .* does not match expected shape"):
             MXFP8QTensor.get_weights_scaling_factor_from_quantizer(weight, quantizer)
+
+    @pytest.mark.parametrize("device", ["cuda"])
+    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_mxfp8_dequantize_default_dtype(self, device, input_dtype):
+        """Test dequantize uses original dtype when dtype=None."""
+        input_tensor = torch.randn(64, 64, dtype=input_dtype, device=device)
+        qtensor, e8m0_scale = MXFP8QTensor.quantize(input_tensor)
+
+        # Dequantize without specifying dtype
+        dequant = qtensor.dequantize(scale=e8m0_scale)
+
+        assert dequant.dtype == input_dtype