Remove FbgemmConfig and remaining Fbgemm tensors

docs/source/torchao_vllm_integration.md

@@ -171,7 +171,7 @@ class MyNewQuantConfig(AOBaseConfig):
     VERSION: ClassVar[int] = 1
 
 class MyQuantizedTensor(TorchAOBaseTensor):
-    """Example based on FbgemmFp8Tensor - stores quantized data + scale"""
+    """Example based on Float8Tensor - stores quantized data + scale"""
 
     tensor_data_attrs = ["quantized_data", "scale"]
     tensor_attributes = ["dtype"]

test/core/test_config.py

@@ -24,7 +24,6 @@
     AWQStep,
 )
 from torchao.quantization.quant_api import (
-    FbgemmConfig,
     Float8DynamicActivationFloat8WeightConfig,
     Float8DynamicActivationInt4WeightConfig,
     Float8WeightOnlyConfig,
@@ -92,7 +91,6 @@
     ),
     AWQConfig(Int4WeightOnlyConfig(group_size=128), step=AWQStep.PREPARE_FOR_LOADING),
     AWQConfig(Int4WeightOnlyConfig(group_size=128), step="prepare_for_loading"),
-    FbgemmConfig(torch.bfloat16, torch.int4, torch.bfloat16, [1, 1, 256]),
 ]
 
 

test/dtypes/test_affine_quantized.py

@@ -24,9 +24,7 @@
     to_affine_quantized_intx,
     to_affine_quantized_intx_static,
 )
-from torchao.float8.config import e4m3_dtype
 from torchao.quantization import (
-    FbgemmConfig,
     Float8WeightOnlyConfig,
     GemliteUIntXWeightOnlyConfig,
     Int4DynamicActivationInt4WeightConfig,
@@ -44,7 +42,6 @@
     is_fbcode,
     is_ROCM,
     is_sm_at_least_89,
-    is_sm_at_least_90,
 )
 
 is_cusparselt_available = (
@@ -100,10 +97,6 @@ def get_quantization_functions(
     if is_sm_at_least_89():
         base_functions.append(Float8WeightOnlyConfig())
 
-    if is_sm_at_least_90():
-        base_functions.append(FbgemmConfig(torch.bfloat16, torch.int4, torch.bfloat16))
-        base_functions.append(FbgemmConfig(e4m3_dtype, e4m3_dtype, torch.bfloat16))
-
     return base_functions
 
 

torchao/_models/llama/generate.py

@@ -434,25 +434,6 @@ def ffn_or_attn_only(mod, fqn):
                 model,
                 Int4WeightOnlyConfig(group_size=group_size, use_hqq=use_hqq, version=1),
             )
-        elif "fbgemm" in quantization and "int4" in quantization:
-            from torchao.quantization import FbgemmConfig
-
-            _, precision, group_size = quantization.split("-")
-            group_size = int(group_size)
-            block_size = [1, group_size]
-            assert precision == "int4", f"FbegemmConfig({precision=}) not supported yet"
-            quantize_(
-                model,
-                FbgemmConfig(torch.bfloat16, torch.int4, torch.bfloat16, block_size),
-            )
-        elif "fbgemm" in quantization and "fp8" in quantization:
-            from torchao.float8.config import e4m3_dtype
-            from torchao.quantization import FbgemmConfig
-
-            quantize_(
-                model,
-                FbgemmConfig(e4m3_dtype, e4m3_dtype, torch.bfloat16),
-            )
         elif "int4dq-" in quantization:
             from torchao.dtypes import CutlassInt4PackedLayout
 

torchao/dtypes/__init__.py

@@ -8,7 +8,6 @@
     to_affine_quantized_intx,
     to_affine_quantized_intx_static,
 )
-from .fbgemm_fp8_tensor import FbgemmFp8Tensor, to_fbgemm_fp8
 from .floatx import (
     CutlassSemiSparseLayout,
     Float8Layout,

torchao/quantization/__init__.py

@@ -43,7 +43,6 @@
 )
 from .quant_api import (
     CutlassInt4PackedLayout,
-    FbgemmConfig,
     Float8DynamicActivationFloat8SemiSparseWeightConfig,
     Float8DynamicActivationFloat8WeightConfig,
     Float8DynamicActivationInt4WeightConfig,
@@ -161,7 +160,6 @@
     "GemliteUIntXWeightOnlyConfig",
     "AOPerModuleConfig",
     "ModuleFqnToConfig",
-    "FbgemmConfig",
     # tensor subclasses
     "Int4Tensor",
     "Int4PlainInt32Tensor",