metal: opt-in Metal 4.1 direct Q2 operand experiment

docs/development/metal-direct-q2.md

@@ -0,0 +1,31 @@
+# Experimental Metal 4.1 direct Q2 operands
+
+This path is disabled by default. Enable it before device initialization:
+
+```sh
+GGML_METAL_Q2_DIRECT_MAX=16 ./build/bin/llama-bench -m model.gguf -p 8,16 -n 0 -d 128 -r 3 -ngl 99 -fa 0
+```
+
+`GGML_METAL_Q2_DIRECT_MIN` defaults to 8. Unset `GGML_METAL_Q2_DIRECT_MAX`, or set it to zero, to retain existing dispatch. The path requires an embedded Metal source build, enabled TensorOps, an M5 device, macOS 27 or later, and a successful Metal 4.1 compiler probe. Unsupported configurations retain existing kernels. The experimental shaders are excluded when the gate is off.
+
+The 16 x 32 x 16 tile decodes Q2_0/PQ2_0 blocks directly into half cooperative operands, preserving the F32/F16 activation type with relaxed TensorOps precision. Unlike the existing PQ2_0 few-row implementation, this implementation also handles Q2_0 and uses the generic matrix batch strides. It has no split-K reduction. Explicit fragment layout is validated only on M5.
+
+## Current result
+
+Do not select this as the default performance path. Against the current branch's existing PQ2_0 few-row kernel, three alternating paired model runs on M5 Pro measured median block times of 59.18 ms vs 95.79 ms at width 8, and 72.57 ms vs 114.53 ms at width 16 (existing vs experimental). Each run used three repetitions, prefix depth 128 and flash attention disabled. Earlier improvements against the older staged kernel do not establish an improvement over the current branch.
+
+Focused Q2_0/PQ2_0 MUL_MAT backend checks pass with the gate off and on; unsupported reference combinations are skipped by the harness. Broader accuracy, current-branch full-model logits, other devices, and end-to-end speculative decoding remain release gates. This is an experimental comparison path, not a production recommendation.
+
+## Follow-up experiments
+
+Three changes were tested separately against matched current-default controls. Each passed the focused backend oracle before three alternating timing pairs; every timing run used three repetitions. Ratios below are candidate/default latency, so values above 1 are regressions.
+
+| Candidate | Width 8 ratio | Width 16 ratio |
+| --- | ---: | ---: |
+| Four-way K split, original activation type | 1.331 | 1.345 |
+| Four-way K split, half activation conversion | 1.531 | 1.555 |
+| Four-way K split, original activation type, packed-word sharing | 1.511 | 1.462 |
+
+K splitting reduced the candidate's latency, but none of these variants beat the current default. Half conversion and packed-word sharing were tested independently from the original-type four-way split. The follow-up variants were removed rather than added as more runtime knobs.
+
+A separate compiler-language control kept the existing kernel selected in both arms: `GGML_METAL_Q2_DIRECT_MIN=17` with MAX=0 versus MAX=16. Current-kernel latency was 58.22/58.12 ms at width 8 and 71.36/71.28 ms at width 16 (Metal 4.0/4.1). This difference is within measurement variation and does not explain the candidate's regression.

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -854,6 +854,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta
 
     const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor;
 
+    const char * direct_env     = getenv("GGML_METAL_Q2_DIRECT_MAX");
+    const char * direct_min_env = getenv("GGML_METAL_Q2_DIRECT_MIN");
+    const int    direct_min     = direct_min_env ? atoi(direct_min_env) : 8;
+    const bool   direct         = ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor_direct &&
+                                  direct_env && op->ne[1] >= direct_min && op->ne[1] <= atoi(direct_env) &&
+                                  (tsrc0 == GGML_TYPE_Q2_0 || tsrc0 == GGML_TYPE_PQ2_0) &&
+                                  (tsrc1 == GGML_TYPE_F32 || tsrc1 == GGML_TYPE_F16);
+
     const bool bc_out = has_tensor
         ? (op->ne[0] % NRA != 0 || op->ne[1] % NRB != 0)
         : (op->ne[0] % 64  != 0 || op->ne[1] % 32  != 0);
@@ -865,6 +873,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta
     const int16_t r3   = (int16_t) (ne13 / op->src[0]->ne[3]);
 
     snprintf(base, 256, "kernel_mul_mm_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
+    if (direct) {
+        snprintf(base, 256, "kernel_mul_mm_direct_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
+    }
     snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d",
              base, bc_inp, bc_out, ne12, ne13, r2, r3);
 
@@ -899,6 +910,13 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta
 
     res.nsg = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
 
+    if (direct) {
+        res.nr0  = 32;
+        res.nr1  = 16;
+        res.nsg  = 1;
+        res.smem = 32;
+    }
+
     return res;
 }
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -289,6 +289,7 @@ struct ggml_metal_device_props {
     bool has_unified_memory;
     bool has_bfloat;
     bool has_tensor;
+    bool has_tensor_direct;
     bool use_residency_sets;
     bool use_shared_buffers;
 

ggml/src/ggml-metal/ggml-metal-device.m

@@ -277,12 +277,12 @@ static bool ggml_metal_library_flatten_file(NSMutableString * dst, NSString * pa
 }
 
 // tensor API headers need Metal 4.0, and an unset language version follows the build SDK
-static void ggml_metal_compile_options_set_lang(MTLCompileOptions * options, bool has_tensor) {
+static void ggml_metal_compile_options_set_lang(MTLCompileOptions * options, bool has_tensor, bool has_tensor_direct) {
     if (!has_tensor) {
         return;
     }
 
-    options.languageVersion = (MTLLanguageVersion) MTLLanguageVersion4_0_GGML;
+    options.languageVersion = (MTLLanguageVersion) (has_tensor_direct ? ((4 << 16) | 1) : MTLLanguageVersion4_0_GGML);
 }
 
 // Compile all per-kind libraries in parallel. `source_for_kind` returns the MSL
@@ -324,7 +324,7 @@ static bool ggml_metal_library_compile_all(
             @autoreleasepool {
                 MTLCompileOptions * options = [MTLCompileOptions new];
                 options.preprocessorMacros = prep;
-                ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(res->dev)->has_tensor);
+                ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(res->dev)->has_tensor, ggml_metal_device_get_props(res->dev)->has_tensor_direct);
 
                 lib = [device newLibraryWithSource:src options:options error:&error];
 
@@ -397,6 +397,9 @@ ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
     if (ggml_metal_device_get_props(dev)->has_tensor) {
         [prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
     }
+    if (ggml_metal_device_get_props(dev)->has_tensor_direct) {
+        [prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR_DIRECT"];
+    }
 #if GGML_METAL_EMBED_LIBRARY
     [prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
 #endif
@@ -569,7 +572,7 @@ ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev
 
         MTLCompileOptions * options = [MTLCompileOptions new];
         options.preprocessorMacros = prep;
-        ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(dev)->has_tensor);
+        ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(dev)->has_tensor, ggml_metal_device_get_props(dev)->has_tensor_direct);
 
         library = [device newLibraryWithSource:src options:options error:&error];
         if (error) {
@@ -1146,6 +1149,29 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
                 }
             }
 
+            // The direct operand layout has only been validated on M5 with the OS 27 runtime.
+#if TARGET_OS_OSX
+            if (@available(macOS 27.0, *)) {
+                dev->props.has_tensor_direct = dev->props.has_tensor &&
+                    [[dev->mtl_device name] containsString:@"M5"] &&
+                    getenv("GGML_METAL_Q2_DIRECT_MAX") && atoi(getenv("GGML_METAL_Q2_DIRECT_MAX")) > 0;
+            }
+
+#endif
+            if (dev->props.has_tensor_direct) {
+                MTLCompileOptions * options = [MTLCompileOptions new];
+                ggml_metal_compile_options_set_lang(options, true, true);
+                NSError * error = nil;
+                id<MTLLibrary> probe = [dev->mtl_device newLibraryWithSource:@"#include <metal_stdlib>\nkernel void direct_version_probe() {}"
+                                                                  options:options error:&error];
+                if (!probe) {
+                    dev->props.has_tensor_direct = false;
+                    GGML_LOG_WARN("%s: Metal 4.1 unavailable, direct operand path disabled\n", __func__);
+                }
+                [probe release];
+                [options release];
+            }
+
             // try to compile a dummy kernel to determine if the tensor API is supported for bfloat
             if (dev->props.has_tensor && dev->props.has_bfloat) {
                 const char * src_tensor_bf16 = "\n"

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2795,7 +2795,13 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
     // 4-8 stay on mul_mv_ext. GGML_METAL_PQ2_0_NR1=1 restores mul_mv_ext for 2-3 as well.
     static const bool pq2_0_ext_enable = getenv("GGML_METAL_PQ2_0_NR1") && atoi(getenv("GGML_METAL_PQ2_0_NR1")) == 1;
 
-    const int ne11_mm_min = op->src[0]->type == GGML_TYPE_Q1_0 ? std::max(8, q1_0_mv_max) : 8;
+    const char * direct_env     = getenv("GGML_METAL_Q2_DIRECT_MAX");
+    const char * direct_min_env = getenv("GGML_METAL_Q2_DIRECT_MIN");
+    const int    direct_min     = direct_min_env ? atoi(direct_min_env) : 8;
+    const bool direct = props_dev->has_tensor_direct && direct_env && ne11 >= direct_min && ne11 <= atoi(direct_env) &&
+                        (op->src[0]->type == GGML_TYPE_Q2_0 || op->src[0]->type == GGML_TYPE_PQ2_0) &&
+                        (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16);
+    const int  ne11_mm_min = direct ? 0 : op->src[0]->type == GGML_TYPE_Q1_0 ? std::max(8, q1_0_mv_max) : 8;
 
     if (ggml_metal_op_mul_mat_q1_0_pc_supported(op)) {
         const int32_t nblk = ne00/128;
@@ -2887,7 +2893,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
         static const int  fewrow_max = getenv("GGML_METAL_PQ2_0_FEWROW_MAX") ? atoi(getenv("GGML_METAL_PQ2_0_FEWROW_MAX")) : 32;
         static const int  fewrow_cfg = getenv("GGML_METAL_PQ2_0_FEWROW_CFG") ? atoi(getenv("GGML_METAL_PQ2_0_FEWROW_CFG")) : 0;
 
-        if (fewrow_on && props_dev->has_tensor &&
+        if (!direct && fewrow_on && props_dev->has_tensor &&
             op->src[0]->type == GGML_TYPE_PQ2_0 && op->src[1]->type == GGML_TYPE_F32 &&
             ne11 >= fewrow_min && ne11 <= fewrow_max &&
             ne00 % 128 == 0 && // QK_PQ2_0
@@ -2934,7 +2940,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
 
     // first try to use small-batch mat-mv kernels
     // these should be efficient for BS [2, ~8]
-    if (op->src[1]->type == GGML_TYPE_F32 && (ne00%128 == 0) &&
+    if (!direct && op->src[1]->type == GGML_TYPE_F32 && (ne00%128 == 0) &&
         (
          (
           (

ggml/src/ggml-metal/kernels/mul_mm.metal

@@ -10,6 +10,79 @@ constant short FC_mul_mm_r3    [[function_constant(FC_MUL_MM + 5)]];
 
 // each block_q contains 16*nl weights
 #ifdef GGML_METAL_HAS_TENSOR
+#ifdef GGML_METAL_HAS_TENSOR_DIRECT
+// Experimental M5 direct-operand Q2/PQ2 path, enabled only by an explicit host switch.
+template <typename Block, int QK, typename T>
+kernel void kernel_mul_mm_direct(constant ggml_metal_kargs_mul_mm & args,
+                                 const device char *                srcA,
+                                 const device char *                srcB,
+                                 device char *                      dst,
+                                 threadgroup char *                 shmem [[threadgroup(0)]],
+                                 uint3                              tile [[threadgroup_position_in_grid]],
+                                 ushort                             tiitg [[thread_index_in_threadgroup]],
+                                 ushort                             sgitg [[simdgroup_index_in_threadgroup]]) {
+    (void) shmem;
+    (void) sgitg;
+    const int           i12 = tile.z % args.ne12, i13 = tile.z / args.ne12;
+    const uint64_t      offset0 = (i12 / args.r2) * args.nb02 + (i13 / args.r3) * args.nb03;
+    const device char * weights = srcA + offset0;
+    const device T *    x       = (const device T *) (srcB + args.nb12 * i12 + args.nb13 * i13);
+    const int           stride  = args.nb11 / sizeof(T);
+    const int           qid     = tiitg >> 2;
+    const int           fm = (qid & 4) | ((tiitg >> 1) & 3), fn = ((qid & 2) | (tiitg & 1)) * 4;
+    const int           m0 = int(tile.x) * 16 + fm, n0 = int(tile.y) * 32 + fm;
+    constexpr auto      desc = mpp::tensor_ops::matmul2d_descriptor(
+        16, 32, 16, false, true, true, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);
+    mpp::tensor_ops::matmul2d<desc, execution_simdgroup> mm;
+    auto a = mm.template get_left_input_cooperative_tensor<T, half, float>();
+    auto b = mm.template get_right_input_cooperative_tensor<T, half, float>();
+    auto c = mm.template get_destination_cooperative_tensor<metal::remove_addrspace_t<decltype(a)>,
+                                                            metal::remove_addrspace_t<decltype(b)>, float>();
+    for (uint i = 0; i < c.get_capacity(); i++) {
+        c[i] = 0;
+    }
+    for (int g = 0; g < args.ne00 / QK; g++) {
+        float                scales[4];
+        const device Block * rows[4];
+#pragma unroll
+        for (int j = 0; j < 4; j++) {
+            int n     = n0 + 8 * j;
+            rows[j]   = (const device Block *) (weights + min(n, args.ne0 - 1) * args.nb01) + g;
+            scales[j] = n < args.ne0 ? float(rows[j]->d) : 0;
+        }
+        for (int kk = 0; kk < QK; kk += 16) {
+            int k = g * QK + kk + fn;
+#pragma unroll
+            for (int j = 0; j < 4; j++) {
+                uchar packed = rows[j]->qs[(kk + fn) / 4];
+#pragma unroll
+                for (int i = 0; i < 4; i++) {
+                    b[j * 4 + i] = half((float((packed >> (2 * i)) & 3) - 1.0f) * scales[j]);
+                }
+            }
+#pragma unroll
+            for (int j = 0; j < 2; j++) {
+                int m = m0 + 8 * j;
+#pragma unroll
+                for (int i = 0; i < 4; i++) {
+                    a[j * 4 + i] = m < args.ne1 ? x[m * stride + k + i] : T(0);
+                }
+            }
+            mm.run(a, b, c);
+        }
+    }
+    auto output = tensor((device float *) dst + tile.z * args.ne0 * args.ne1, dextents<int, 2>{ args.ne0, args.ne1 },
+                         array<int, 2>{ 1, args.ne0 });
+    c.store(output.slice(int(tile.y) * 32, int(tile.x) * 16));
+}
+
+typedef decltype(kernel_mul_mm_direct<block_pq2_0, QK_PQ2_0, float>) direct_mm_t;
+template [[host_name("kernel_mul_mm_direct_pq2_0_f32")]] kernel direct_mm_t kernel_mul_mm_direct<block_pq2_0, QK_PQ2_0, float>;
+template [[host_name("kernel_mul_mm_direct_pq2_0_f16")]] kernel direct_mm_t kernel_mul_mm_direct<block_pq2_0, QK_PQ2_0, half>;
+template [[host_name("kernel_mul_mm_direct_q2_0_f32")]]  kernel direct_mm_t kernel_mul_mm_direct<block_q2_0, QK2_0, float>;
+template [[host_name("kernel_mul_mm_direct_q2_0_f16")]]  kernel direct_mm_t kernel_mul_mm_direct<block_q2_0, QK2_0, half>;
+#endif // GGML_METAL_HAS_TENSOR_DIRECT
+
 template<
     typename SA, typename SA_4x4, typename SA_8x8,
     typename SB, typename SB_2x4, typename SB_8x8,