Add default filtering to remove mis-alinged weights

test/dtypes/test_affine_quantized_float.py

@@ -236,6 +236,39 @@ def test_serialization(self, mode: str):
                     original_layer.weight.scale, new_layer.weight.scale
                 ), f"Scales do not match for {layer_name}"
 
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    @unittest.skipIf(not is_cuda_8_9, "Requires GPU with compute capability >= 8.9")
+    def test_fp8_weight_dimension_warning(self):
+        # Create model with incompatible dimensions (not multiples of 16)
+        model = ToyLinearModel(10, 25).cuda()  # 10x25 and 25x10 weights
+
+        # Set up logging capture
+        with self.assertLogs(
+            "torchao.quantization.quant_api", level="INFO"
+        ) as log_context:
+            quantize_(
+                model, float8_dynamic_activation_float8_weight(granularity=PerTensor())
+            )
+            print(model)
+
+        # Verify warning messages for both layers
+        expected_messages = [
+            "Skipping float8 quantization: weight shape torch.Size([25, 10])",
+            "Skipping float8 quantization: weight shape torch.Size([10, 25])",
+        ]
+        # Check that we got warnings for both incompatible layers
+        warning_count = sum(
+            1 for msg in log_context.output if "Skipping float8 quantization" in msg
+        )
+        self.assertEqual(warning_count, 2, "Expected warnings for both linear layers")
+
+        # Check warning message content
+        for expected in expected_messages:
+            self.assertTrue(
+                any(expected in msg for msg in log_context.output),
+                f"Expected warning message containing: {expected}",
+            )
+
 
 common_utils.instantiate_parametrized_tests(TestAffineQuantizedFloat8Compile)
 

torchao/quantization/quant_api.py

@@ -733,6 +733,32 @@ def _input_activation_quant_func_fp8(
     return activation
 
 
+def _fp8_mm_compat(weight: torch.Tensor) -> bool:
+    """
+    Check if a weight tensor meets float8 quantization requirements.
+
+    Args:
+        weight (torch.Tensor): The weight tensor to check
+
+    Returns:
+        bool: True if the tensor can be quantized to float8, False otherwise
+    """
+    assert (
+        weight.dim() == 2
+    ), f"float8 quantization only works for 2-D tensors, got {weight.dim()}D tensor"
+
+    out_dim, in_dim = weight.shape
+    is_compatible = (in_dim % 16 == 0) and (out_dim % 16 == 0)
+
+    if not is_compatible:
+        logger.info(
+            f"Skipping float8 quantization: weight shape {weight.shape} is not compatible with _scaled_mm. "
+            f"Both input dimension ({in_dim}) and output dimension ({out_dim}) must be multiples of 16. "
+        )
+
+    return is_compatible
+
+
 def float8_dynamic_activation_float8_weight(
     activation_dtype: torch.dtype = torch.float8_e4m3fn,
     weight_dtype: torch.dtype = torch.float8_e4m3fn,
@@ -761,6 +787,8 @@ def float8_dynamic_activation_float8_weight(
     activation_granularity, weight_granularity = _normalize_granularity(granularity)
 
     def apply_float8_dynamic_activation_quant(weight: torch.Tensor):
+        if not _fp8_mm_compat(weight):
+            return weight
         if isinstance(weight_granularity, PerRow):
             assert (
                 weight.dtype == torch.bfloat16
@@ -818,6 +846,8 @@ def float8_static_activation_float8_weight(
     ), "Static quantization only supports PerTensor granularity"
 
     def apply_float8_static_activation_quant(weight: torch.Tensor):
+        if not _fp8_mm_compat(weight):
+            return weight
         block_size = get_block_size(weight.shape, weight_granularity)
         quantized_weight = to_affine_quantized_floatx(
             input_float=weight,