metal: COOP_K4 (function constant 7) - per-element right-input decode for GPUs without the M5's layout (the M4)

- moe_mm.metal mm_gu3/mm_down3 and qsa_prompt.metal qsa_attn_mm/qsa_flash: with COOP_K4 = false every element is
  decoded on its own (the M4 holds runs of 2 n at one k); the experts are correct on the M4 that way
- tests/coop_probe.h: probes the layout (COOP_K4=0/1 overrides); moe_mm_parity, qsa_prompt_parity, prompt_chain
  set the constant; prompt_chain defaults to experts v2 and QSA v1 without it (faster / correct on the M4)
- qsa_prompt_parity: NaN-fills the partials before the matrix-unit run (a kernel that writes nothing shows),
  line-buffered output, QSA_NOFLASH

metal/kernels/moe_mm.metal

@@ -373,6 +373,13 @@ MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16
 // no threadgroup staging and no barriers in the K loop.  Each simdgroup owns TN output rows; a threadgroup SG * TN.
 // The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values
 // at a time (W4<F>).  Activations: the gathered fp16 rows from device memory.
+#ifndef STRATA_COOP_K4
+#define STRATA_COOP_K4
+// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's
+// layout; tests/coop_layout).  false (the M4: runs of 2 n at one k): every element decoded on its own.
+constant bool COOP_K4 [[function_constant(7)]];
+constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true;
+#endif
 template <int F> struct W4;
 // the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16
 static inline half4 w4_signed(uint g, uint s, float sc) {
@@ -570,6 +577,14 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[
         if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; }
     for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) {
         const uint16_t cap = rg.get_capacity(), half_cap = cap / 2;
+        if (!coop_k4) {
+            for (uint16_t i = 0; i < cap; ++i) {
+                const auto ix = rg.get_multidimensional_index(i);
+                const int k = k0 + ix[0], n = ix[1];
+                rg[i] = W4<F>::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3];
+                ru[i] = W4<F>::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3];
+            }
+        } else
 #pragma clang loop unroll(full)
         for (uint16_t i = 0; i < half_cap; i += 4) {
             const auto ix = rg.get_multidimensional_index(i);
@@ -633,6 +648,13 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS
     for (uint16_t i = 0; i < acc.get_capacity(); ++i)
         if (acc.is_valid_element(i)) acc[i] = 0.0f;
     for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) {
+        if (!coop_k4) {
+            for (uint16_t i = 0; i < rw.get_capacity(); ++i) {
+                const auto ix = rw.get_multidimensional_index(i);
+                const int k = k0 + ix[0];
+                rw[i] = W4<F>::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3];
+            }
+        } else
 #pragma clang loop unroll(full)
         for (uint16_t i = 0; i < rw.get_capacity(); i += 4) {
             const auto ix = rw.get_multidimensional_index(i);

metal/kernels/qsa_prompt.metal

@@ -8,6 +8,13 @@
 // The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32.
 #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
 using namespace mpp::tensor_ops;
+#ifndef STRATA_COOP_K4
+#define STRATA_COOP_K4
+// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's
+// layout; tests/coop_layout).  false (the M4: runs of 2 n at one k): every element decoded on its own.
+constant bool COOP_K4 [[function_constant(7)]];
+constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true;
+#endif
 
 // grid (chunks, QKV, n_tok), 32 threads.  qh: the q heads as fp16 [t][QH][QD] (qcur rounded).
 kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]],
@@ -47,6 +54,14 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char*
         for (uint16_t i = 0; i < acc.get_capacity(); ++i)
             if (acc.is_valid_element(i)) acc[i] = 0.0f;
         for (int k0 = 0; k0 < QD; k0 += 32) {
+            if (!coop_k4) {
+                for (uint16_t i = 0; i < rk.get_capacity(); ++i) {
+                    const auto ix = rk.get_multidimensional_index(i);
+                    const int k = k0 + ix[0], id = cid[ix[1]];
+                    const uint row = uint(id) * QKV + g;
+                    rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f);
+                }
+            } else
 #pragma clang loop unroll(full)
             for (uint16_t i = 0; i < rk.get_capacity(); i += 4) {
                 const auto ix = rk.get_multidimensional_index(i);
@@ -213,6 +228,14 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq
             for (uint16_t i = 0; i < s.get_capacity(); ++i)
                 if (s.is_valid_element(i)) s[i] = 0.0f;
             for (int k0 = d0; k0 < d0 + 64; k0 += 32) {
+                if (!coop_k4) {
+                    for (uint16_t i = 0; i < rk.get_capacity(); ++i) {
+                        const auto ix = rk.get_multidimensional_index(i);
+                        const int k = k0 + ix[0], id = cid[ix[1]];
+                        const uint row = uint(id) * QKV + g;
+                        rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f);
+                    }
+                } else
 #pragma clang loop unroll(full)
                 for (uint16_t i = 0; i < rk.get_capacity(); i += 4) {
                     const auto ix = rk.get_multidimensional_index(i);

metal/tests/coop_probe.h

@@ -0,0 +1,30 @@
+// metal/tests/coop_probe.h - which element layout this GPU gives matmul2d's single-simdgroup right-input cooperative
+// tensor: true = runs of 4 consecutive k at one n per lane (the M5), false = anything else (the M4: runs of 2 n at one
+// k).  The kernels that fill right inputs per lane (mm_gu3, mm_down3, qsa_attn_mm, qsa_flash) take it as function
+// constant 7.  COOP_K4=0/1 in the environment overrides.
+#pragma once
+#include "mtl_rt.h"
+
+#include <cstdlib>
+
+inline bool coop_k4(mtl::Context& cx) {
+    if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0;
+    const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n");
+    mtl::Pipeline p = cx.pipeline(lib, "coop_layout_16_64_32");
+    mtl::Buffer out = cx.alloc(4 * 3 * 32 * 128 * 4);
+    mtl::Dispatch d;
+    d.p = &p;
+    d.args = {mtl::Arg::B(out)};
+    d.tg[0] = 32;
+    cx.run({d});
+    // every lane: elements 4j .. 4j + 3 are (k, n), (k + 1, n), (k + 2, n), (k + 3, n)
+    const int* o = out.as<int>();
+    for (int l = 0; l < 32; ++l) {
+        const int* q = o + l * 128 * 4;
+        const int cap = q[2];
+        for (int i = 0; i + 3 < cap; i += 4)
+            for (int j = 1; j < 4; ++j)
+                if (q[(i + j) * 4] != q[i * 4] + j || q[(i + j) * 4 + 1] != q[i * 4 + 1]) return false;
+    }
+    return true;
+}

metal/tests/moe_mm_parity.cpp

@@ -4,6 +4,7 @@
 //
 //   moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4]          parity (64 experts)
 //   moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..]            per-layer time, 512 experts
+#include "coop_probe.h"
 #include "mtl_rt.h"
 #include "ref_quant.h"
 
@@ -90,9 +91,12 @@ int main(int argc, char** argv) {
     const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG);
     mtl::Constants probe;
     probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0);
-    mtl::Pipeline pgu = v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl)
+    const bool k4 = coop_k4(cx);
+    if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)");
+    mtl::Pipeline pgu = v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4))
                            : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe),
-                  pdn = cx.pipeline(lib, (v3 ? "mm_down3_" : v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)),
+                  pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4))
+                           : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)),
                   pgat = cx.pipeline(lib, "mm_gather");
     MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk};
     auto encode = [&](int layer, std::vector<mtl::Dispatch>& ds, bool plan, bool gemm) {

metal/tests/prompt_chain.cpp

@@ -17,6 +17,7 @@
 // Every dense matrix is dequantized to fp16 (dq_f16) once a chunk, before its GEMM.  After the last chunk: the final
 // HC read and the LM head on the LAST row only (the decode kernels: gr_norm / gr_down / gr_up, GEMV, argmax).
 // Left out: the PLE path and the MTP draft.
+#include "coop_probe.h"
 #include "gguf_lite.h"
 #include "mtl_rt.h"
 #include "ref_quant.h"
@@ -100,6 +101,8 @@ int main(int argc, char** argv) {
     const int lsmp = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n");
     std::printf("kernels compiled in %.1f s\n", mtl::now_s() - tc0);
     std::map<std::string, mtl::Pipeline> pcache;
+    const bool k4 = coop_k4(cx);
+    std::printf("matmul2d right-input layout: %s\n", k4 ? "runs of 4 k (per-lane run decoders)" : "other (per-element decode)");
     auto P = [&](int l, const std::string& fn, const mtl::Constants& c = {}, const std::string& key = "") -> const mtl::Pipeline* {
         const std::string k = key.empty() ? fn : key;
         auto it = pcache.find(k);
@@ -123,12 +126,14 @@ int main(int argc, char** argv) {
     const char* hc_env = std::getenv("PC_HC");
     const std::string hc_mode = hc_env ? hc_env : "fused";   // fused | plain
     const char* qsa_env = std::getenv("PC_QSA");
-    const std::string qsa_mode = qsa_env ? qsa_env : "mm";   // mm (matrix units) | v1 (qsa_attn, fp32)
+    // mm (matrix units; not on the M4's layout yet: qsa_attn_mm writes nothing there) | v1 (qsa_attn, fp32)
+    const std::string qsa_mode = qsa_env ? qsa_env : (k4 ? "mm" : "v1");
     // the expert GEMMs: "v2" (threadgroup-staged weights) or "v3:TM,TN,TK,SG" (register operands, TN rows a simdgroup)
     const char* gu_env = std::getenv("PC_GU");
     const char* dn_env = std::getenv("PC_DN");
     // defaults (M5 Max, moe_mm_parity --v3 --bench): gate/up 16,32,32,4 (every format), down 16,64,32,2
-    const std::string gu_mode = gu_env ? gu_env : "v3:auto", dn_mode = dn_env ? dn_env : "v3:16,64,32,2";
+    // without the M5's right-input layout (the M4) the per-element v3 decode is slower than v2: v2 there
+    const std::string gu_mode = gu_env ? gu_env : (k4 ? "v3:auto" : "v2"), dn_mode = dn_env ? dn_env : (k4 ? "v3:16,64,32,2" : "v2");
     auto v3_tile = [](const std::string& m, int* t) { std::sscanf(m.c_str() + 3, "%d,%d,%d,%d", &t[0], &t[1], &t[2], &t[3]); };
     int gu3[4] = {16, 64, 16, 2}, dn3[4] = {16, 64, 16, 2};
     if (gu_mode.rfind("v3:", 0) == 0 && gu_mode != "v3:auto") v3_tile(gu_mode, gu3);
@@ -364,7 +369,7 @@ int main(int argc, char** argv) {
                 if (qsa_mode == "mm") {
                     const uint32_t nq = (uint32_t)T * 24 * 256;
                     D(P(lqsa, "qsa_q_f16"), {mtl::Arg::B(qcur), mtl::Arg::B(qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256, "qsa attn");
-                    D(P(lqsa, "qsa_attn_mm"), {mtl::Arg::B(qcurh), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]),
+                    D(P(lqsa, "qsa_attn_mm", mtl::Constants().b(7, k4)), {mtl::Arg::B(qcurh), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]),
                                                mtl::Arg::B(vsc[qi]), mtl::Arg::B(qids), mtl::Arg::B(part), mtl::Arg::V(a)},
                       QCHUNKS, 2, 32, "qsa attn", T);
                 } else
@@ -415,7 +420,7 @@ int main(int argc, char** argv) {
             g.ind_off = 8 * 4;
             push(g, "experts");
             mtl::Dispatch gu;
-            gu.p = gu_v3 ? P(lib, "mm_gu3_" + std::to_string(eg.type) + "_" + gu3_for(eg.type))
+            gu.p = gu_v3 ? P(lib, "mm_gu3_" + std::to_string(eg.type) + "_" + gu3_for(eg.type), mtl::Constants().b(7, k4))
                          : P(lib, "mm_gu2_" + std::to_string(eg.type) + "_16_64_16_2", mtl::Constants().i(5, 0));
             gu.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
             gu.tg[0] = gu_v3 ? 32 * gu3[3] : 64;
@@ -423,7 +428,7 @@ int main(int argc, char** argv) {
             gu.ind_off = 2 * 4;
             push(gu, "experts gu");
             mtl::Dispatch dn;
-            dn.p = dn_v3 ? P(lib, "mm_down3_" + std::to_string(ed.type) + "_" + tl3(dn3))
+            dn.p = dn_v3 ? P(lib, "mm_down3_" + std::to_string(ed.type) + "_" + tl3(dn3), mtl::Constants().b(7, k4))
                          : P(lib, "mm_down2_" + std::to_string(ed.type) + (ed.type == ref::IQ4_NL ? "_16_64_32_2" : "_16_64_16_2"));
             dn.args = {mtl::Arg::B(W), mtl::Arg::B(HS), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
             dn.tg[0] = dn_v3 ? 32 * dn3[3] : 64;

metal/tests/qsa_prompt_parity.cpp

@@ -2,12 +2,14 @@
 // (fp32, itself checked by qsa_parity), both followed by qsa_merge: the attention output of T query tokens on random
 // int8 caches, within the budget (all cells, causal) and past it (2051 selected cells per token); then the speed.
 //   qsa_prompt_parity [--bench]
+#include "coop_probe.h"
 #include "mtl_rt.h"
 #include "ref_quant.h"
 
 #include <algorithm>
 #include <cmath>
 #include <cstdio>
+#include <cstdlib>
 #include <set>
 #include <string>
 #include <vector>
@@ -20,6 +22,8 @@ double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v
 
 int main(int argc, char** argv) {
     const bool bench = argc > 1 && std::string(argv[1]) == "--bench";
+    std::setvbuf(stdout, nullptr, _IOLBF, 0);
+    const bool noflash = std::getenv("QSA_NOFLASH") != nullptr;
     mtl::Context cx;
     const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
     std::printf("device %s\n", cx.device_name().c_str());
@@ -34,8 +38,8 @@ int main(int argc, char** argv) {
     for (size_t i = 0; i < ksc.bytes / 2; ++i) { ksc.as<uint16_t>()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); vsc.as<uint16_t>()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); }
     for (size_t i = 0; i < (size_t)T * QH * QD; ++i) qcur.as<float>()[i] = r.sym() * 3.0f;
     for (size_t i = 0; i < (size_t)T * QH * 2 * QD; ++i) qfull.as<float>()[i] = r.sym();
-    mtl::Pipeline pa = cx.pipeline(lib, "qsa_attn"), pm = cx.pipeline(lib, "qsa_merge"), pmm = cx.pipeline(lib, "qsa_attn_mm"),
-                  pq = cx.pipeline(lib, "qsa_q_f16"), pf = cx.pipeline(lib, "qsa_flash");
+    mtl::Pipeline pa = cx.pipeline(lib, "qsa_attn"), pm = cx.pipeline(lib, "qsa_merge"), pmm = cx.pipeline(lib, "qsa_attn_mm", mtl::Constants().b(7, coop_k4(cx))),
+                  pq = cx.pipeline(lib, "qsa_q_f16"), pf = cx.pipeline(lib, "qsa_flash", mtl::Constants().b(7, coop_k4(cx)));
     mtl::Buffer xf = cx.alloc(2 * (size_t)T * QH * QD);
     int fails = 0;
     for (uint32_t pos0 : {0u, 30000u}) {
@@ -75,6 +79,7 @@ int main(int argc, char** argv) {
             continue;
         }
         cx.run(ref_ds);
+        std::fill(part.as<float>(), part.as<float>() + part.bytes / 4, std::nanf(""));   // a kernel that writes nothing shows
         cx.run(mm_ds);
         double se = 0, sr = 0;
         for (size_t i = 0; i < (size_t)T * QH * QD; ++i) {
@@ -86,6 +91,7 @@ int main(int argc, char** argv) {
         const bool ok = rel < 3e-3;
         fails += !ok;
         std::printf("pos0 %5u, T %d: qsa_attn_mm vs qsa_attn (after merge) rel L2 %.2e  %s\n", pos0, T, rel, ok ? "PASS" : "FAIL");
+        if (noflash) continue;
         cx.run({fl});
         double sf = 0;
         for (size_t i = 0; i < (size_t)T * QH * QD; ++i) {