Move marlin_qqq_tensor to prototype

Author: jainapurva

benchmarks/microbenchmarks/utils.py

@@ -218,7 +218,7 @@ def string_to_config(
         )
     if "marlin" in quantization:
         if "qqq" in quantization:
-            from torchao.dtypes import MarlinQQQLayout
+            from torchao.prototype.dtypes import MarlinQQQLayout
 
             return Int8DynamicActivationInt4WeightConfig(
                 group_size=128,

test/quantization/test_marlin_qqq.py

@@ -10,7 +10,7 @@
 from torch import nn
 from torch.testing._internal.common_utils import TestCase, run_tests
 
-from torchao.dtypes import MarlinQQQLayout
+from torchao.prototype.dtypes import MarlinQQQLayout
 from torchao.quantization.marlin_qqq import (
     pack_to_marlin_qqq,
     unpack_from_marlin_qqq,
@@ -132,5 +132,25 @@ def test_pack_unpack_equivalence(self):
             )
 
 
+def test_marlin_qqq_tensor_deprecation_warning():
+    """Test that importing from the old location raises a deprecation warning"""
+    import warnings
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        # Import from the old deprecated location
+        from torchao.dtypes.uintx.marlin_qqq_tensor import (  # noqa: F401
+            MarlinQQQLayout,
+        )
+
+        # Verify the deprecation warning was raised
+        assert len(w) == 1
+        assert issubclass(w[-1].category, DeprecationWarning)
+        assert "torchao.dtypes.uintx.marlin_qqq_tensor is deprecated" in str(
+            w[-1].message
+        )
+        assert "torchao.prototype.dtypes import" in str(w[-1].message)
+
+
 if __name__ == "__main__":
     run_tests()

torchao/_models/llama/generate.py

@@ -460,7 +460,7 @@ def ffn_or_attn_only(mod, fqn):
                 )
         if "marlin" in quantization:
             if "qqq" in quantization:
-                from torchao.dtypes import MarlinQQQLayout
+                from torchao.prototype.dtypes import MarlinQQQLayout
 
                 quantize_(
                     model,

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -39,10 +39,6 @@
     _linear_fp_act_uint4_weight_int8_zero_check,
     _linear_fp_act_uint4_weight_int8_zero_impl,
 )
-from torchao.dtypes.uintx.marlin_qqq_tensor import (
-    _linear_int8_act_int4_weight_marlin_qqq_check,
-    _linear_int8_act_int4_weight_marlin_qqq_impl,
-)
 from torchao.dtypes.uintx.marlin_sparse_layout import (
     _linear_fp_act_int4_weight_sparse_marlin_check,
     _linear_fp_act_int4_weight_sparse_marlin_impl,
@@ -94,6 +90,10 @@
     _linear_int8_act_int4_weight_cpu_check,
     _linear_int8_act_int4_weight_cpu_impl,
 )
+from torchao.prototype.dtypes.uintx.marlin_qqq_tensor import (
+    _linear_int8_act_int4_weight_marlin_qqq_check,
+    _linear_int8_act_int4_weight_marlin_qqq_impl,
+)
 from torchao.quantization.quant_primitives import (
     ZeroPointDomain,
     _dequantize_affine_no_zero_point,

torchao/dtypes/uintx/__init__.py

@@ -1,3 +1,9 @@
+from torchao.prototype.dtypes.uintx.marlin_qqq_tensor import (
+    MarlinQQQLayout,
+    MarlinQQQTensor,
+    to_marlinqqq_quantized_intx,
+)
+
 from .dyn_int8_act_int4_wei_cpu_layout import (
     Int8DynamicActInt4WeightCPULayout,
 )
@@ -7,11 +13,6 @@
 from .int4_xpu_layout import (
     Int4XPULayout,
 )
-from .marlin_qqq_tensor import (
-    MarlinQQQLayout,
-    MarlinQQQTensor,
-    to_marlinqqq_quantized_intx,
-)
 from .marlin_sparse_layout import (
     MarlinSparseLayout,
 )