add int4tilepackedto4dtensor subclass to safetensors

test/prototype/safetensors/test_safetensors_support.py

@@ -45,6 +45,7 @@ class TestSafeTensors(TestCase):
             (Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()), False),
             (Int4WeightOnlyConfig(), False),
             (Int4WeightOnlyConfig(), True),
+            (Int4WeightOnlyConfig(int4_packing_format="tile_packed_to_4d"), False),
         ],
     )
     def test_safetensors(self, config, act_pre_scale=False):

torchao/prototype/safetensors/safetensors_utils.py

@@ -6,21 +6,30 @@
 import torch
 
 import torchao
-from torchao.quantization import Float8Tensor, Int4Tensor
+from torchao.quantization import (
+    Float8Tensor,
+    Int4Tensor,
+    Int4TilePackedTo4dTensor,
+)
 from torchao.quantization.quantize_.common import KernelPreference
 from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs
 
 ALLOWED_CLASSES = {
     "Float8Tensor": Float8Tensor,
     "Int4Tensor": Int4Tensor,
+    "Int4TilePackedTo4dTensor": Int4TilePackedTo4dTensor,
     "Float8MMConfig": torchao.float8.inference.Float8MMConfig,
     "QuantizeTensorToFloat8Kwargs": QuantizeTensorToFloat8Kwargs,
     "PerRow": torchao.quantization.PerRow,
     "PerTensor": torchao.quantization.PerTensor,
     "KernelPreference": KernelPreference,
 }
 
-ALLOWED_TENSORS_SUBCLASSES = ["Float8Tensor", "Int4Tensor"]
+ALLOWED_TENSORS_SUBCLASSES = [
+    "Float8Tensor",
+    "Int4Tensor",
+    "Int4TilePackedTo4dTensor",
+]
 
 __all__ = [
     "TensorSubclassAttributeJSONEncoder",