Remove excessive assertions in MXFP8QTensor

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

Author: danisereb

modelopt/torch/quantization/qtensor/mxfp8_tensor.py

@@ -148,14 +148,14 @@ def quantize_with_scale(
         This method is useful for export paths where the scale has already been computed.
 
         Args:
-            weight: The weight tensor to quantize. Must be at least 2D.
+            weight: The weight tensor to quantize. Must be at least 1D.
             e8m0_scale: E8M0 scale as uint8 biased exponent (bias = 127).
-                Shape should be [..., out_dim, in_dim // 32].
+                Shape should be [..., out_dim, in_dim // 32] for 2D+ tensors,
+                or [in_dim // 32] for 1D tensors.
 
         Returns:
             torch.Tensor: Quantized weight as float8_e4m3fn with same shape as input.
         """
-        assert weight.dim() >= 2, f"Weight must be at least 2D, got {weight.dim()}D"
         assert e8m0_scale.dtype == cls.SCALE_DTYPE, (
             f"e8m0_scale must be {cls.SCALE_DTYPE} (E8M0 format), got {e8m0_scale.dtype}"
         )
@@ -201,10 +201,6 @@ def quantize(cls, input: torch.Tensor) -> tuple:
         Returns:
             tuple: (MXFP8QTensor, e8m0_scale) where e8m0_scale is uint8 biased exponent.
         """
-        assert input.numel() > 0, "Input tensor must not be empty"
-        assert input.dim() >= 1, f"Input must have at least 1 dimension, got {input.dim()}D"
-        assert input.is_floating_point(), f"Input must be floating point, got {input.dtype}"
-
         original_shape = input.shape
         original_dtype = input.dtype
 
@@ -234,9 +230,6 @@ def dequantize(self, dtype: torch.dtype = None, **kwargs) -> torch.Tensor:
         assert "scale" in kwargs, "dequantize requires 'scale' in kwargs"
 
         e8m0_scale = kwargs["scale"]
-        assert e8m0_scale.dtype == self.SCALE_DTYPE, (
-            f"e8m0_scale must be {self.SCALE_DTYPE} (E8M0 format), got {e8m0_scale.dtype}"
-        )
 
         if dtype is None:
             dtype = self.metadata["dtype"]

tests/gpu/torch/quantization/test_qtensor_cuda.py

@@ -839,12 +839,6 @@ def test_mxfp8_get_weights_scaling_factor_asserts_non_divisible(self):
     @pytest.mark.parametrize("device", ["cuda"])
     def test_mxfp8_quantize_with_scale_asserts(self, device):
         """Test quantize_with_scale raises assertions for invalid inputs."""
-        # Test 1D weight assertion
-        weight_1d = torch.randn(64, dtype=torch.float32, device=device)
-        scale = torch.randint(0, 255, (2,), dtype=torch.uint8, device=device)
-        with pytest.raises(AssertionError, match="Weight must be at least 2D"):
-            MXFP8QTensor.quantize_with_scale(weight_1d, scale)
-
         # Test wrong scale dtype assertion
         weight = torch.randn(64, 64, dtype=torch.float32, device=device)
         wrong_dtype_scale = torch.randn(64, 2, dtype=torch.float32, device=device)
@@ -856,34 +850,3 @@ def test_mxfp8_quantize_with_scale_asserts(self, device):
         scale = torch.randint(0, 255, (64, 1), dtype=torch.uint8, device=device)
         with pytest.raises(AssertionError, match="must be divisible by MXFP8 block size"):
             MXFP8QTensor.quantize_with_scale(weight_bad_dim, scale)
-
-    @pytest.mark.parametrize("device", ["cuda"])
-    def test_mxfp8_quantize_dequantize_asserts(self, device):
-        """Test quantize and dequantize raise assertions for invalid inputs."""
-        # Test empty tensor assertion
-        empty_tensor = torch.empty(0, dtype=torch.float32, device=device)
-        with pytest.raises(AssertionError, match="Input tensor must not be empty"):
-            MXFP8QTensor.quantize(empty_tensor)
-
-        # Test 0D tensor assertion
-        scalar_tensor = torch.tensor(1.0, dtype=torch.float32, device=device)
-        with pytest.raises(AssertionError, match="Input must have at least 1 dimension"):
-            MXFP8QTensor.quantize(scalar_tensor)
-
-        # Test non-floating point assertion
-        int_tensor = torch.randint(0, 10, (32, 32), dtype=torch.int32, device=device)
-        with pytest.raises(AssertionError, match="Input must be floating point"):
-            MXFP8QTensor.quantize(int_tensor)
-
-        # Create a valid quantized tensor for dequantize tests
-        input_tensor = torch.randn(64, 64, dtype=torch.float32, device=device)
-        qtensor, e8m0_scale = MXFP8QTensor.quantize(input_tensor)
-
-        # Test missing scale assertion
-        with pytest.raises(AssertionError, match="dequantize requires 'scale' in kwargs"):
-            qtensor.dequantize(dtype=torch.float32)
-
-        # Test wrong scale dtype assertion
-        wrong_dtype_scale = torch.randn(64, 2, dtype=torch.float32, device=device)
-        with pytest.raises(AssertionError, match="e8m0_scale must be"):
-            qtensor.dequantize(dtype=torch.float32, scale=wrong_dtype_scale)