Q4k/Q6k improvements (#20643)

Small quant improvements:

* Parallelizes writes in Q4k/Q6k kernels.
* Adds support for MLX 6-bit kernel

Author: metascroy

backends/mlx/builder/op_helpers.py

@@ -157,7 +157,8 @@ def emit_lifted_constant(P: "MLXProgramBuilder", value, dtype: torch.dtype) -> S
 
     if isinstance(value, (int, float, bool)):
         return P.make_or_get_constant(
-            f"_scalar_{value}", torch.tensor(value, dtype=dtype)  # 0-D
+            f"_scalar_{value}",
+            torch.tensor(value, dtype=dtype),  # 0-D
         )
 
     from executorch.backends.mlx.serialization.mlx_graph_schema import FullNode
@@ -533,11 +534,10 @@ def to_mlx_qparams(
     assert qdata.dtype == torch.int8
     offset = 2 ** (bits - 1)
 
-    # Pack data tightly into uint32
-    assert 32 % bits == 0
-    vals_per_uint32 = 32 // bits
-    assert qdata.shape[1] % vals_per_uint32 == 0
+    # Pack data into a contiguous uint32 bitstream. cols*bits must be a
+    # multiple of 32 (holds since in_features is a multiple of group_size>=32).
     rows, cols = qdata.shape
+    assert (cols * bits) % 32 == 0
 
     if bits == 4:
         # 4-bit: view(uint8) + wrapping add + pack 2 nibbles per byte → view as uint32
@@ -554,14 +554,18 @@ def to_mlx_qparams(
         q = qdata.view(torch.uint8) + offset
         Q = q.contiguous().view(torch.uint32).reshape(rows, -1)
     else:
-        # General fallback for other bit widths
-        Q = (qdata.to(torch.int32) + offset).reshape(-1, vals_per_uint32)
-        shifts = torch.arange(0, 32, bits, dtype=torch.int32)
-        shifted = Q << shifts
-        packed = shifted[:, 0]
-        for i in range(1, vals_per_uint32):
-            packed = packed | shifted[:, i]
-        Q = packed.view(torch.uint32).reshape(rows, -1)
+        # Contiguous bit-packing for widths that don't divide 32 (e.g. 6-bit),
+        # matching MLX's affine pack_and_quantize (ops.cpp): each value
+        # contributes `bits` LSB-first bits to a continuous stream that is
+        # chunked into uint32 words (also LSB-first), straddling word
+        # boundaries -> (rows, cols*bits//32) uint32.
+        q = qdata.to(torch.int32) + offset  # 0 .. 2**bits-1
+        bit_ids = torch.arange(bits, dtype=torch.int32)
+        stream = (q.unsqueeze(-1) >> bit_ids) & 1  # (rows, cols, bits) LSB-first
+        stream = stream.reshape(rows, -1, 32)  # (rows, n_words, 32)
+        word_shifts = torch.arange(32, dtype=torch.int64)
+        packed = (stream.to(torch.int64) << word_shifts).sum(-1)  # (rows, n_words)
+        Q = packed.to(torch.int32).contiguous().view(torch.uint32)
 
     if compute_biases:
         B = -scale * (zero_point.to(scale.dtype) + offset)
@@ -654,10 +658,11 @@ def parse_dequant_node(
     if group_size not in [16, 32, 64, 128]:
         return None
 
-    # TODO: MLX supports 3, 5, and 7, but we need to figure out the
-    # packing story in to_mlx_qparams to use them
+    # MLX supports 2,3,4,5,6,8-bit affine quantization. to_mlx_qparams packs
+    # 2/4/8 via fast paths and other widths (e.g. 6) via a general contiguous
+    # bit-packer, so enable 6 here too.
     bits = (qmax - qmin + 1).bit_length() - 1
-    if bits not in [2, 4, 8]:
+    if bits not in [2, 4, 6, 8]:
         return None
     return qdata, scale, zero_point, group_size, bits, out_dtype, quantized_dim
 

backends/mlx/custom_kernel_ops/gguf/q4k/linear.py

@@ -152,8 +152,8 @@ def _q4k_matmul_source(has_bias: bool) -> str:
     constexpr short NL0 = NK / 16;  // = 2 — dequant iterations per thread for weight
     constexpr short NL1 = NK / 8;   // = 4 — load iterations per thread for activation
 
-    threadgroup half sa[4096];  // NR0 * NK storage (strided by 64)
-    threadgroup half sb[4096];  // NR1 * NK storage (strided by 64)
+    threadgroup half sa[4096];  // NR0 * NK weight tile; reused as NR1*NR0 float output staging (8 KB)
+    threadgroup half sb[1024];  // NR1 * NK activation tile (strided by 64): 16 ib slots * 64
 
     const ushort tid   = thread_index_in_threadgroup;   // 0..127
     const ushort sgitg = simdgroup_index_in_threadgroup; // 0..3
@@ -252,7 +252,9 @@ def _q4k_matmul_source(has_bias: bool) -> str:
         }}
     }}
 
-    // --- Write results: always via threadgroup memory for float→OutT cast ---
+    // --- Write results: stage the output tile in threadgroup memory, then
+    // drain it to device. Staging is required for the float->OutT cast and the
+    // optional bias add.
     // Barrier needed: sa was used for weight tiles during the K-loop and is now
     // reused as float staging for the output. Without this barrier, a fast
     // simdgroup could start writing mc[] into sa while a slower one is still
@@ -268,14 +270,15 @@ def _q4k_matmul_source(has_bias: bool) -> str:
         }}
         threadgroup_barrier(mem_flags::mem_threadgroup);
 
-        if (sgitg == 0) {{
-            for (int j = tid; j < nr1; j += NR1) {{
-                device OutT * D = out + (uint)(r1 + j) * (uint)N + r0;
-                threadgroup float * Cp = ((threadgroup float *) sa) + j * NR0;
-                for (int i = 0; i < nr0; ++i) {{
-                    float v = Cp[i];
-                    D[i] = (OutT)(v {bias_add});
-                }}
+        // Drain all NR1*NR0 tile elements across the threadgroup's 128 threads.
+        // NR0/NR1 are compile-time constants, so idx / NR0 and idx % NR0 fold
+        // to a shift/mask.
+        for (int idx = tid; idx < NR1 * NR0; idx += 128) {{
+            const int j = idx / NR0;
+            const int i = idx % NR0;
+            if (j < nr1 && i < nr0) {{
+                const float v = ((threadgroup float *) sa)[j * NR0 + i];
+                out[(uint)(r1 + j) * (uint)N + (r0 + i)] = (OutT)(v {bias_add});
             }}
         }}
     }}

backends/mlx/custom_kernel_ops/gguf/q6k/linear.py

@@ -162,8 +162,8 @@ def _q6k_matmul_source(has_bias: bool) -> str:
     constexpr short NL0 = NK / 16;  // = 2 — dequant iterations per thread for weight
     constexpr short NL1 = NK / 8;   // = 4 — load iterations per thread for activation
 
-    threadgroup half sa[4096];  // NR0 * NK storage (strided by 64)
-    threadgroup half sb[4096];  // NR1 * NK storage (strided by 64)
+    threadgroup half sa[4096];  // NR0 * NK weight tile; reused as NR1*NR0 float output staging (8 KB)
+    threadgroup half sb[1024];  // NR1 * NK activation tile (strided by 64): 16 ib slots * 64
 
     const ushort tid   = thread_index_in_threadgroup;   // 0..127
     const ushort sgitg = simdgroup_index_in_threadgroup; // 0..3
@@ -262,7 +262,9 @@ def _q6k_matmul_source(has_bias: bool) -> str:
         }}
     }}
 
-    // --- Write results: always via threadgroup memory for float→InT cast ---
+    // --- Write results: stage the output tile in threadgroup memory, then
+    // drain it to device. Staging is required for the float->InT cast and the
+    // optional bias add.
     // Barrier needed: sa was used for weight tiles during the K-loop and is now
     // reused as float staging for the output. Without this barrier, a fast
     // simdgroup could start writing mc[] into sa while a slower one is still
@@ -278,22 +280,23 @@ def _q6k_matmul_source(has_bias: bool) -> str:
         }}
         threadgroup_barrier(mem_flags::mem_threadgroup);
 
-        if (sgitg == 0) {{
-            for (int j = tid; j < nr1; j += NR1) {{
-                device InT * D = out + (uint)(r1 + j) * (uint)N + r0;
-                threadgroup float * Cp = ((threadgroup float *) sa) + j * NR0;
-                for (int i = 0; i < nr0; ++i) {{
-                    float v = Cp[i];
-                    D[i] = (InT)(v {bias_add});
-                }}
+        // Drain all NR1*NR0 tile elements across the threadgroup's 128 threads.
+        // NR0/NR1 are compile-time constants, so idx / NR0 and idx % NR0 fold
+        // to a shift/mask.
+        for (int idx = tid; idx < NR1 * NR0; idx += 128) {{
+            const int j = idx / NR0;
+            const int i = idx % NR0;
+            if (j < nr1 && i < nr0) {{
+                const float v = ((threadgroup float *) sa)[j * NR0 + i];
+                out[(uint)(r1 + j) * (uint)N + (r0 + i)] = (InT)(v {bias_add});
             }}
         }}
     }}
 """
 
 
 # Number of simdgroups per threadgroup for the mat-vec kernel.
-_Q6K_MV_NSG = 4
+_Q6K_MV_NSG = 2
 # Tile sizes for the mat-mat kernel (from llama.cpp kernel_mul_mm).
 _Q6K_MM_NR0 = 64  # weight/output rows (N dim) per threadgroup
 _Q6K_MM_NR1 = 32  # activation rows (M dim) per threadgroup

backends/mlx/test/test_ops.py

@@ -6303,6 +6303,8 @@ def get_test_configs(cls) -> List["QuantizedLinearTest"]:
             cls(group_size=128),
             cls(qdtype=torch.int2),
             cls(qdtype=torch.int8),
+            cls(qdtype=torch.int6),
+            cls(qdtype=torch.int6, group_size=128),
             # group_size=16: exercises the non-fused dequantize+matmul path
             # (requires ET_MLX_ALLOW_NON_FUSED_QUANTIZED_OPS=1).
             cls(qdtype=torch.int8, group_size=16),