Mac DeepSeek vision tower: GEMM on the Metal 4 tensor ops (BF16 operands read in place, FP32 accumulation; 798x448 in 93 ms, 1344x1260 in 1.2 s, was 2.6 s and 29 s), DS4_VISION_NO_MPP=1 for the plain kernel, sequential-FMA test GEMM (DS4_VISION_REF_GEMM)

metal/src/engine/ds4_vision_metal.cpp

@@ -33,10 +33,11 @@ struct PatchArgs { uint64_t rgb, out; int32_t width, gw, p; };
 struct AddArgs { uint64_t x, y, bias; int32_t ldy, dim; };
 struct NormArgs { uint64_t x, w, out; int32_t dim; float eps; };
 struct RopeArgs { uint64_t y, bias, inv_freq, qkv; int32_t r0, gw, heads, hd; };
-struct SoftmaxArgs { uint64_t s, p; int32_t n; };
+struct SoftmaxArgs { uint64_t s, p; int32_t n; float scale; };
 struct SiluArgs { uint64_t y, mid; int32_t inter; };
 struct UnfoldArgs { uint64_t xn, un; int32_t gh, gw, m, dim, lw; };
 struct BiasActArgs { uint64_t y, bias, out; int32_t n, gelu; };
+struct MppArgs { uint64_t X, W, Y; int32_t M, N, K, ldx, ldw, ldy; };
 struct GemmArgs { uint64_t X, W, Y; uint32_t rows, in, out, ldx, ldw, ldy; float alpha; uint32_t flags; };
 
 struct Runner {
@@ -81,11 +82,32 @@ bool encode(const Args& a, std::string& err) {
         err = "vision: compiling ds4_vision.metal failed";
         return false;
     }
+    // the tensor-op GEMM (ds4_vision_mpp.metal) when it compiles here; DS4_VISION_NO_MPP=1: the plain kernel
+    static const bool want_mpp = std::getenv("DS4_VISION_NO_MPP") == nullptr && std::getenv("DS4_VISION_REF_GEMM") == nullptr;
+    static int s_lib_mpp = -2;
+    if (want_mpp && s_lib_mpp == -2) s_lib_mpp = cx.compile({"../src/engine/msl/ds4_vision_mpp.metal"}, "", false);
+    const bool use_mpp = want_mpp && s_lib_mpp >= 0;
     Runner R(cx);
     R.lib = s_lib;
     const mtl::Pipeline pPatch = R.pipe("vis_patchify"), pAdd = R.pipe("vis_add_rows"), pNorm = R.pipe("vis_rmsnorm"),
                         pRope = R.pipe("vis_qkv_rope"), pSm = R.pipe("vis_softmax"), pSilu = R.pipe("vis_silu_mul"),
-                        pUnf = R.pipe("vis_unfold"), pBias = R.pipe("vis_bias_act"), pGemm = R.pipe("vis_gemm");
+                        pUnf = R.pipe("vis_unfold"), pBias = R.pipe("vis_bias_act"), pGemm0 = R.pipe("vis_gemm"),
+                        pGemmRef = R.pipe("vis_gemm_ref");
+    static const char* const ref_env = std::getenv("DS4_VISION_REF_GEMM");   // (tests) 1: the sequential-FMA GEMM, 2: summed in chunks of 64
+    static const bool ref_gemm = ref_env != nullptr;
+    static const uint32_t ref_flags = ref_env != nullptr && ref_env[0] == '2' ? 4u : 0u;
+    const mtl::Pipeline& pGemm = ref_gemm ? pGemmRef : pGemm0;
+    mtl::Pipeline pMppK128, pMppK64, pMppNbf;
+    if (use_mpp) {
+        try {
+            pMppK128 = cx.pipeline(s_lib_mpp, "vis_mpp_k_f32_128_64");
+            pMppK64 = cx.pipeline(s_lib_mpp, "vis_mpp_k_f32_64_64");
+            pMppNbf = cx.pipeline(s_lib_mpp, "vis_mpp_n_bf_64_64");
+        } catch (...) {
+            err = "vision: no tensor-op pipeline";
+            return false;
+        }
+    }
     if (!R.ok) {
         err = R.err;
         return false;
@@ -99,8 +121,20 @@ bool encode(const Args& a, std::string& err) {
     auto at = [](uint64_t base, size_t elems, size_t esz) { return base + (uint64_t) (elems * esz); };
     auto gemm = [&](uint64_t X, uint32_t ldx, uint64_t W, uint32_t ldw, bool wn, uint64_t Y, uint32_t ldy, bool ybf, int rows,
                     int in, int out, float alpha) {
-        GemmArgs g{X, W, Y, (uint32_t) rows, (uint32_t) in, (uint32_t) out, ldx, ldw, ldy, alpha, (wn ? 1u : 0u) | (ybf ? 2u : 0u)};
-        R.go(pGemm, g, (uint32_t) ((out + 63) / 64), (uint32_t) ((rows + 63) / 64), 128);
+        if (use_mpp && alpha == 1.0f && (wn ? ybf : !ybf)) {   // k-major W to FP32, or n-major W (V) to BF16
+            MppArgs m{X, W, Y, rows, out, in, (int32_t) ldx, (int32_t) ldw, (int32_t) ldy};
+            if (wn) {
+                R.go(pMppNbf, m, (uint32_t) ((out + 63) / 64), (uint32_t) ((rows + 63) / 64), 128);
+            } else if (rows >= 128) {
+                R.go(pMppK128, m, (uint32_t) ((out + 63) / 64), (uint32_t) ((rows + 127) / 128), 128);
+            } else {
+                R.go(pMppK64, m, (uint32_t) ((out + 63) / 64), (uint32_t) ((rows + 63) / 64), 128);
+            }
+            return;
+        }
+        GemmArgs g{X, W, Y, (uint32_t) rows, (uint32_t) in, (uint32_t) out, ldx, ldw, ldy, alpha, (wn ? 1u : 0u) | (ybf ? 2u : 0u) | (ref_gemm ? ref_flags : 0u)};
+        if (ref_gemm) R.go(pGemm, g, (uint32_t) ((out + 31) / 32), (uint32_t) rows, 32);
+        else R.go(pGemm, g, (uint32_t) ((out + 63) / 64), (uint32_t) ((rows + 63) / 64), 128);
     };
     auto add_rows = [&](uint64_t x, uint64_t y, int ldy, uint64_t bias, int rows) {
         AddArgs g{x, y, bias, ldy, D};
@@ -134,8 +168,8 @@ bool encode(const Args& a, std::string& err) {
             for (int hh = 0; hh < a.heads; ++hh) {
                 // S [query][key] = scale * Q K^T
                 gemm(at(a.qkv, (size_t) q0 * H3 + (size_t) hh * hd, 2), H3, at(a.qkv, (size_t) D + (size_t) hh * hd, 2), H3, false,
-                     a.s, N, false, nq, hd, N, a.scale);
-                SoftmaxArgs sg{a.s, a.p, N};
+                     a.s, N, false, nq, hd, N, use_mpp ? 1.0f : a.scale);
+                SoftmaxArgs sg{a.s, a.p, N, use_mpp ? a.scale : 1.0f};
                 R.go(pSm, sg, (uint32_t) nq, 1, smT);
                 // O[:, head] = P V
                 gemm(a.p, N, at(a.qkv, (size_t) 2 * D + (size_t) hh * hd, 2), H3, true, at(a.o, (size_t) hh * hd, 2), D, true, nq, N, hd,

metal/src/engine/msl/ds4_vision.metal

@@ -137,6 +137,7 @@ struct VisSoftmaxArgs {
     device const float* s;
     device ushort* p;
     int n;
+    float scale;   // applied to the scores first (the attention scale when the GEMM did not fold it in; exact: a power of two)
 };
 constant int kSmPer = 20;   // a row of up to nt * 20 keys in registers (512 x 20 = 10240)
 kernel void vis_softmax(constant VisSoftmaxArgs& a [[buffer(0)]], uint r [[threadgroup_position_in_grid]],
@@ -148,7 +149,7 @@ kernel void vis_softmax(constant VisSoftmaxArgs& a [[buffer(0)]], uint r [[threa
     float m = -FLT_MAX;
     for (int k = 0; k < kSmPer; ++k) {
         const int e = int(tid) + k * int(nt);
-        v[k] = e < a.n ? sr[e] : -FLT_MAX;
+        v[k] = e < a.n ? sr[e] * a.scale : -FLT_MAX;
         m = fmax(m, v[k]);
     }
     m = vis_block_reduce(m, true, red, tid, nt);
@@ -279,3 +280,34 @@ kernel void vis_gemm(constant VisGemmArgs& a [[buffer(0)]], uint2 tgp [[threadgr
         else ((device float*) a.Y)[ulong(gr) * ulong(a.ldy) + ulong(go)] = v;
     }
 }
+
+// (reference, off by default) the same GEMM one output a thread with a sequential FP32 FMA chain: another accumulation order to measure
+// how much of a difference to another implementation is the BF16 roundings' sensitivity to summation order
+kernel void vis_gemm_ref(constant VisGemmArgs& a [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]],
+                         ushort tid [[thread_index_in_threadgroup]]) {
+    const uint o = tgp.x * 32u + uint(tid), r = tgp.y;
+    if (o >= a.out || r >= a.rows) return;
+    const bool wn = (a.flags & 1u) != 0u;
+    float acc = 0.0f;
+    if ((a.flags & 4u) == 0u) {
+        for (uint k = 0; k < a.in; ++k) {
+            const float x = vis_bf(a.X[ulong(r) * ulong(a.ldx) + ulong(k)]);
+            const float w = wn ? vis_bf(a.W[ulong(k) * ulong(a.ldw) + ulong(o)]) : vis_bf(a.W[ulong(o) * ulong(a.ldw) + ulong(k)]);
+            acc = fma(x, w, acc);
+        }
+    } else {   // chunks of 64 summed apart, then added (a blocked order, as a library GEMM has)
+        for (uint k0 = 0; k0 < a.in; k0 += 64u) {
+            float part = 0.0f;
+            const uint k1 = min(k0 + 64u, a.in);
+            for (uint k = k0; k < k1; ++k) {
+                const float x = vis_bf(a.X[ulong(r) * ulong(a.ldx) + ulong(k)]);
+                const float w = wn ? vis_bf(a.W[ulong(k) * ulong(a.ldw) + ulong(o)]) : vis_bf(a.W[ulong(o) * ulong(a.ldw) + ulong(k)]);
+                part = fma(x, w, part);
+            }
+            acc += part;
+        }
+    }
+    const float v = a.alpha * acc;
+    if ((a.flags & 2u) != 0u) ((device ushort*) a.Y)[ulong(r) * ulong(a.ldy) + ulong(o)] = vis_bfr(v);
+    else ((device float*) a.Y)[ulong(r) * ulong(a.ldy) + ulong(o)] = v;
+}

metal/src/engine/msl/ds4_vision_mpp.metal

@@ -0,0 +1,59 @@
+// metal/src/engine/msl/ds4_vision_mpp.metal - the vision tower's GEMM on the Metal 4 tensor ops (MetalPerformancePrimitives matmul2d),
+// BF16 operands read in place with their row strides, FP32 accumulation.  Same contract as vis_gemm in ds4_vision.metal
+// (Y[r][o] = sum_k X[r][k] W(o, k)) without alpha: the host folds the attention scale into the softmax.  A separate file, so a Mac
+// whose compiler lacks bfloat tensors keeps the plain kernel (DS4_VISION_NO_MPP=1 asks for that too).
+//
+// Strata's own; the idea of BF16 weights read in place by a tensor-op matmul follows metal/dense_mm.metal of this tree.
+#include <metal_stdlib>
+#include <metal_tensor>
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+
+using namespace metal;
+using namespace mpp::tensor_ops;
+
+struct VisMppArgs {
+    device const bfloat* X;
+    device const bfloat* W;
+    device uchar* Y;
+    int M, N, K, ldx, ldw, ldy;
+};
+
+// WN: W is [K][N] (o contiguous), else [N][K].  OB: Y is BF16, else FP32.
+template <int TM, int TN, int SG, bool WN, bool OB>
+kernel void vis_mpp(constant VisMppArgs& a [[buffer(0)]], uint3 tg [[threadgroup_position_in_grid]]) {
+    using E2 = dextents<int32_t, 2>;
+    tensor<device bfloat, E2, tensor_inline> tA((device bfloat*) a.X, E2(a.K, a.M), array<int32_t, 2>{1, a.ldx});
+    constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>, false, !WN, false);
+    matmul2d<desc, execution_simdgroups<SG>> op;
+    auto mA = tA.slice(0, int(tg.y) * TM);
+    if (WN) {
+        tensor<device bfloat, E2, tensor_inline> tB((device bfloat*) a.W, E2(a.N, a.K), array<int32_t, 2>{1, a.ldw});
+        auto mB = tB.slice(int(tg.x) * TN, 0);
+        if (OB) {
+            tensor<device bfloat, E2, tensor_inline> tY((device bfloat*) a.Y, E2(a.N, a.M), array<int32_t, 2>{1, a.ldy});
+            auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM);
+            op.run(mA, mB, mY);
+        } else {
+            tensor<device float, E2, tensor_inline> tY((device float*) a.Y, E2(a.N, a.M), array<int32_t, 2>{1, a.ldy});
+            auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM);
+            op.run(mA, mB, mY);
+        }
+    } else {
+        tensor<device bfloat, E2, tensor_inline> tB((device bfloat*) a.W, E2(a.K, a.N), array<int32_t, 2>{1, a.ldw});
+        auto mB = tB.slice(0, int(tg.x) * TN);
+        if (OB) {
+            tensor<device bfloat, E2, tensor_inline> tY((device bfloat*) a.Y, E2(a.N, a.M), array<int32_t, 2>{1, a.ldy});
+            auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM);
+            op.run(mA, mB, mY);
+        } else {
+            tensor<device float, E2, tensor_inline> tY((device float*) a.Y, E2(a.N, a.M), array<int32_t, 2>{1, a.ldy});
+            auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM);
+            op.run(mA, mB, mY);
+        }
+    }
+}
+
+template [[host_name("vis_mpp_k_f32_128_64")]] kernel decltype(vis_mpp<128, 64, 4, false, false>) vis_mpp<128, 64, 4, false, false>;
+template [[host_name("vis_mpp_k_f32_64_64")]] kernel decltype(vis_mpp<64, 64, 4, false, false>) vis_mpp<64, 64, 4, false, false>;
+template [[host_name("vis_mpp_n_f32_64_64")]] kernel decltype(vis_mpp<64, 64, 4, true, false>) vis_mpp<64, 64, 4, true, false>;
+template [[host_name("vis_mpp_n_bf_64_64")]] kernel decltype(vis_mpp<64, 64, 4, true, true>) vis_mpp<64, 64, 4, true, true>;

packaging/mac/build_dsv41_engine.sh

@@ -51,7 +51,7 @@ rm -rf "$STAGE"
 mkdir -p "$STAGE/metal/kernels" "$STAGE/metal/src/engine/msl"
 cp "$EXE" "$STAGE/strata-dsv41"; chmod +x "$STAGE/strata-dsv41"
 cp "$ROOT"/metal/kernels/*.metal "$STAGE/metal/kernels/"
-for f in ds4_init ds4_iq ds4_mm ds4_quant ds4_vision mtp mtp_chain ple; do
+for f in ds4_init ds4_iq ds4_mm ds4_quant ds4_vision ds4_vision_mpp mtp mtp_chain ple; do
   cp "$ROOT/metal/src/engine/msl/$f.metal" "$STAGE/metal/src/engine/msl/" || exit 1
 done
 cp "$ROOT/third_party/ggml/ggml-common.h" "$STAGE/metal/"

third_party/antirez-ds4/NOTICE.md

@@ -5,7 +5,7 @@ The licence text is in `LICENSE` beside this file; it is copied unchanged next t
 
 Derived files (the way a DeepSeek-V4 vision tower is split into Metal kernels, which steps round to BF16, the 2-D rotary on q and k,
 the aligner's unfold order, a polynomial erf because MSL has none):
-- `metal/src/engine/msl/ds4_vision.metal`, after `metal/deepseek4_vision.metal`, `metal/glm53_vision.metal`, `metal/glm53_bf16.metal`
+- `metal/src/engine/msl/ds4_vision.metal` and `ds4_vision_mpp.metal`, after `metal/deepseek4_vision.metal`, `metal/glm53_vision.metal`, `metal/glm53_bf16.metal`
 - `metal/src/engine/ds4_vision_metal.cpp` and `metal/include/metal_compat/ds4_vision_metal.h`, after `ds4_gpu_deepseek4_vision_encode`
   and its helpers in `ds4_metal.m`