up

Author: metascroy

backends/mlx/builder/op_helpers.py

@@ -321,11 +321,11 @@ def parse_dequant_node(
     quantized_dim, group_size = non_one[0]
     if group_size not in [32, 64, 128]:
         return None
-    if qmin == -8 and qmax == 7:
-        bits = 4
-    elif qmin == -128 and qmax == 127:
-        bits = 8
-    else:
+
+    # TODO: MLX supports 3, 5, and 7, but we need to figure out the
+    # packing story in to_mlx_qparams to use them
+    bits = (qmax - qmin + 1).bit_length() - 1
+    if bits not in [2, 4, 8]:
         return None
     return qdata, scale, zero_point, group_size, bits, out_dtype, quantized_dim
 

backends/mlx/runtime/MLXInterpreter.h

@@ -149,24 +149,6 @@ exec_addmm(const AddmmNode& n, ExecutionState& st, StreamOrDevice s) {
   st.set_tensor(n.out, std::move(Y));
 }
 
-inline void
-exec_linear(const LinearNode& n, ExecutionState& st, StreamOrDevice s) {
-  const auto& X = st.const_tensor_ref(n.x);
-  auto W = st.const_tensor_ref(n.weight);
-  W = transpose(W, {1, 0}, s);
-
-  array Y = n.bias ? addmm(
-                         st.const_tensor_ref(*n.bias),
-                         X,
-                         W,
-                         /*alpha=*/1.0f,
-                         /*beta=*/1.0f,
-                         s)
-                   : matmul(X, W, s);
-
-  st.set_tensor(n.out, std::move(Y));
-}
-
 inline void
 exec_item_int(const ItemIntNode& n, ExecutionState& st, StreamOrDevice) {
   // Intentional sync: item() requires a concrete scalar value for SymInt
@@ -1601,9 +1583,6 @@ class Interpreter {
       case OpCode::ADDMM:
         ops::exec_addmm(std::get<AddmmNode>(instr.node), st, s);
         break;
-      case OpCode::LINEAR:
-        ops::exec_linear(std::get<LinearNode>(instr.node), st, s);
-        break;
       case OpCode::ITEM_INT:
         ops::exec_item_int(std::get<ItemIntNode>(instr.node), st, s);
         break;

backends/mlx/serialization/schema.fbs

@@ -86,13 +86,6 @@ table AddmmNode {
     beta: float = 1.0;         // Scalar multiplier for bias
 }
 
-table LinearNode {
-    x: Tid (required);
-    weight: Tid (required);
-    out: Tid (required);
-    bias: Tid;  // optional
-}
-
 table ItemIntNode {
     x: Tid (required);
     out: Vid (required);
@@ -916,7 +909,6 @@ union OpNode {
     NoopNode,
     IdCopyNode,
     AddmmNode,
-    LinearNode,
     ItemIntNode,
     ExpandDimsNode,
     TileNode,

backends/mlx/test/test_ops.py

@@ -5385,19 +5385,6 @@ def get_dynamic_shapes(self) -> Optional[Dict[str, any]]:
         }
 
 
-class QuantizedLinearModel(nn.Module):
-    """Simple linear layer that will be quantized."""
-
-    def __init__(
-        self, in_features: int = 64, out_features: int = 128, bias: bool = True
-    ):
-        super().__init__()
-        self.linear = nn.Linear(in_features, out_features, bias=bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.linear(x)
-
-
 @register_test
 class QuantizedLinearTest(OpTestCase):
     """Test case for TorchAO int4 quantized nn.Linear."""
@@ -5408,13 +5395,14 @@ class QuantizedLinearTest(OpTestCase):
 
     def __init__(
         self,
-        in_features: int = 64,
+        in_features: int = 128,
         out_features: int = 128,
         batch_size: int = 2,
         seq_len: int = 16,
         bias: bool = True,
         group_size: int = 32,
         dtype: torch.dtype = torch.bfloat16,
+        qdtype: torch.dtype = torch.int4,
     ):
         self.in_features = in_features
         self.out_features = out_features
@@ -5423,8 +5411,9 @@ def __init__(
         self.bias = bias
         self.group_size = group_size
         self.dtype = dtype
+        self.qdtype = qdtype
 
-        parts = ["quantized_linear", f"g{group_size}"]
+        parts = ["quantized_linear", f"{qdtype}", f"g{group_size}"]
         if not bias:
             parts.append("no_bias")
         self.name = "_".join(parts)
@@ -5434,26 +5423,25 @@ def get_test_configs(cls) -> List["QuantizedLinearTest"]:
         return [
             cls(),
             cls(bias=False),
+            cls(group_size=64),
+            cls(group_size=128),
+            cls(qdtype=torch.int2),
+            cls(qdtype=torch.int8),
         ]
 
     def create_model(self) -> nn.Module:
-        model = QuantizedLinearModel(
-            self.in_features, self.out_features, bias=self.bias
-        )
+        model = LinearModel(self.in_features, self.out_features, bias=self.bias)
         model = model.to(self.dtype)
 
-        try:
-            from torchao.quantization.granularity import PerGroup
-            from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
+        from torchao.quantization.granularity import PerGroup
+        from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
 
-            quantize_(
-                model,
-                IntxWeightOnlyConfig(
-                    weight_dtype=torch.int4, granularity=PerGroup(self.group_size)
-                ),
-            )
-        except ImportError:
-            raise RuntimeError("TorchAO not installed. Run: pip install torchao")
+        quantize_(
+            model,
+            IntxWeightOnlyConfig(
+                weight_dtype=self.qdtype, granularity=PerGroup(self.group_size)
+            ),
+        )
 
         return model
 
@@ -5464,21 +5452,6 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (x,)
 
 
-class QuantizedEmbeddingModel(nn.Module):
-    """Simple embedding layer that will be quantized."""
-
-    def __init__(
-        self,
-        num_embeddings: int = 1000,
-        embedding_dim: int = 64,
-    ):
-        super().__init__()
-        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.embedding(x)
-
-
 @register_test
 class QuantizedEmbeddingTest(OpTestCase):
     """Test case for TorchAO int4 quantized nn.Embedding."""
@@ -5490,48 +5463,51 @@ class QuantizedEmbeddingTest(OpTestCase):
     def __init__(
         self,
         num_embeddings: int = 1000,
-        embedding_dim: int = 64,
+        embedding_dim: int = 128,
         batch_size: int = 2,
         seq_len: int = 16,
         group_size: int = 32,
         dtype: torch.dtype = torch.bfloat16,
+        qdtype: torch.dtype = torch.int4,
     ):
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.batch_size = batch_size
         self.seq_len = seq_len
         self.group_size = group_size
         self.dtype = dtype
+        self.qdtype = qdtype
 
-        parts = ["quantized_embedding", f"g{group_size}"]
+        parts = ["quantized_embedding", f"{qdtype}", f"g{group_size}"]
         self.name = "_".join(parts)
 
     @classmethod
     def get_test_configs(cls) -> List["QuantizedEmbeddingTest"]:
         return [
             cls(),
+            cls(group_size=64),
+            cls(group_size=128),
+            cls(qdtype=torch.int2),
+            cls(qdtype=torch.int8),
         ]
 
     def create_model(self) -> nn.Module:
-        model = QuantizedEmbeddingModel(self.num_embeddings, self.embedding_dim)
+        model = EmbeddingModel(self.num_embeddings, self.embedding_dim)
         model = model.to(self.dtype)
 
-        try:
-            from torchao.quantization.granularity import PerGroup
-            from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
+        from torchao.quantization.granularity import PerGroup
+        from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
 
-            def embedding_filter(module: nn.Module, fqn: str) -> bool:
-                return isinstance(module, nn.Embedding)
+        def embedding_filter(module: nn.Module, fqn: str) -> bool:
+            return isinstance(module, nn.Embedding)
 
-            quantize_(
-                model,
-                IntxWeightOnlyConfig(
-                    weight_dtype=torch.int4, granularity=PerGroup(self.group_size)
-                ),
-                embedding_filter,
-            )
-        except ImportError:
-            raise RuntimeError("TorchAO not installed. Run: pip install torchao")
+        quantize_(
+            model,
+            IntxWeightOnlyConfig(
+                weight_dtype=torch.int4, granularity=PerGroup(self.group_size)
+            ),
+            embedding_filter,
+        )
 
         return model