support Per Group in Float8DynamicActivationFloat8WeightConfig (#4182)

Summary:
Pull Request resolved: #4182

as title

Differential Revision: D97987011

Author: zlin888

torchao/float8/inference.py

@@ -8,6 +8,7 @@
 """
 
 import math
+import logging
 from typing import List, NamedTuple, Optional, Tuple, Union
 
 import torch
@@ -19,6 +20,8 @@
     is_sm_at_least_89,
 )
 
+logger: logging.Logger = logging.getLogger()
+
 Tensor = torch.Tensor
 
 
@@ -237,6 +240,7 @@ def _normalize_granularity(
     ],
 ) -> Tuple[FP8Granularity, FP8Granularity]:
     from torchao.quantization.granularity import (
+        PerGroup,
         PerRow,
         PerTensor,
     )
@@ -253,9 +257,12 @@ def _normalize_granularity(
         is_per_row = isinstance(granularity[0], PerRow) and isinstance(
             granularity[1], PerRow
         )
+        is_per_group = isinstance(granularity[0], PerGroup) and isinstance(
+            granularity[1], PerGroup
+        )
         is_a_1_128_w_128_128 = _granularity_is_a_1_128_w_128_128(granularity)
 
-        if not (is_per_tensor or is_per_row or is_a_1_128_w_128_128):
+        if not (is_per_tensor or is_per_row or is_per_group or is_a_1_128_w_128_128):
             raise ValueError(f"Unsupported granularity types: {granularity}.")
         if not isinstance(granularity[0], type(granularity[1])):
             raise ValueError(
@@ -281,6 +288,7 @@ def _check_hardware_support(
         ValueError: If invalid granularity type is provided
     """
     from torchao.quantization.granularity import (
+        PerGroup,
         PerRow,
         PerTensor,
     )
@@ -291,6 +299,9 @@ def _check_hardware_support(
     is_per_row = isinstance(granularities[0], PerRow) and isinstance(
         granularities[1], PerRow
     )
+    is_per_group = isinstance(granularities[0], PerGroup) and isinstance(
+        granularities[1], PerGroup
+    )
     is_a_1_128_w_128_128 = _granularity_is_a_1_128_w_128_128(granularities)
 
     if is_per_tensor or is_per_row:
@@ -304,5 +315,10 @@ def _check_hardware_support(
         assert is_sm_at_least_89(), (
             "Float8 1x128 activation and 128x128 weight scaling requires CUDA compute capability ≥8.9."
         )
+    elif is_per_group:
+        logger.warning(
+            "PerGroup blockwise FP8 quantization: no hardware check performed. "
+            "Ensure the target device supports blockwise FP8 operations."
+        )
     else:
         raise ValueError(f"Invalid granularities {granularities}.")