diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 304b74925..8e242e66a 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -1376,6 +1376,7 @@ Strata itself: [MIT](../LICENSE). The model files are not part of it; their lice - [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) (MIT): the i-quant formats, the GPU dot products and dequantizers transcribed in `src/kernels/cuda/iq_kernels.cu`, the CPU backend linked for the i-quant experts, the `mtmd` library behind the image encoder (`tools/vision/`), and `gguf-py` used by the tools. See `third_party/ggml/LICENSE`. +- [antirez/ds4](https://github.com/antirez/ds4) (MIT, commit 890898cf): the Metal vision-tower kernels of the Mac DeepSeek engine (`metal/src/engine/msl/ds4_vision.metal`, `ds4_vision_metal.cpp`) follow its `metal/deepseek4_vision.metal` and the tower's host code. See `third_party/antirez-ds4/NOTICE.md`. - Ideas from [Splash](https://github.com/incoai/splash), [ninfer](https://github.com/Neroued/ninfer) and [HyperQwen](https://github.com/syv-ai/HyperQwen); references in the paper. - The DFlash2 drafter's Q4_0 selector (174 MiB less graphics memory, same drafts and same text) is nazeshinjite's idea: diff --git a/metal/include/metal_compat/ds4_vision_metal.h b/metal/include/metal_compat/ds4_vision_metal.h new file mode 100644 index 000000000..5b09097ad --- /dev/null +++ b/metal/include/metal_compat/ds4_vision_metal.h @@ -0,0 +1,32 @@ +// metal/include/metal_compat/ds4_vision_metal.h - the Metal side of src/dsv41/ds4_vision.cu's encode(): one call that runs the +// whole vision tower on the Metal runtime (metal/src/engine/ds4_vision_metal.cpp, kernels in msl/ds4_vision.metal). +// +// The layout of a DeepSeek-V4 vision tower on Metal follows antirez/ds4 (MIT, commit 890898cf); see third_party/antirez-ds4. +#pragma once + +#include +#include +#include +#include + +namespace strata::ds4::vismetal { + +// every pointer is a device address (what strata::gpu's cudaMalloc / cudaHostGetDevicePointer hand out) +struct Args { + int width = 0, height = 0, gh = 0, gw = 0, n = 0, l = 0, nw = 0; // image, patch grid, patches, aligner cells, aligner width + int dim = 0, inter = 0, patch = 0, merge = 0, heads = 0, hd = 0, layers = 0, out = 0, rb = 0, qb = 0; + float eps = 0.0f, scale = 0.0f; + uint64_t rgb = 0, patches = 0, x = 0, xn = 0, qkv = 0, y = 0, o = 0, mid = 0, s = 0, p = 0, un = 0, a16 = 0, outb = 0, inv_freq = 0; + uint64_t weights = 0; // the streamed weights (BF16): every block's [wqkv, wo, w1, w2], then aligner w1, w2 + size_t block_n = 0, al1_off = 0, al2_off = 0; + uint64_t pe_w = 0, pe_b = 0, norm = 0, al1_b = 0, al2_b = 0; + std::vector n1, n2, bqkv, bo; +}; + +/// Runs patchify .. aligner into a.outb ([a.l][a.out] BF16) and returns once the GPU is done. +bool encode(const Args& a, std::string& err); + +/// The capability string the picture tower's startup line carries (the Control looks for it in the binary). +const char* marker(); + +} // namespace strata::ds4::vismetal diff --git a/metal/src/engine/ds4_vision_metal.cpp b/metal/src/engine/ds4_vision_metal.cpp new file mode 100644 index 000000000..29a99ddbe --- /dev/null +++ b/metal/src/engine/ds4_vision_metal.cpp @@ -0,0 +1,184 @@ +// metal/src/engine/ds4_vision_metal.cpp - DeepSeek-V4.1's vision tower on Metal: the host side. +// +// Runs src/dsv41/ds4_vision.cu's encode() as dispatches of msl/ds4_vision.metal on the dsv41 engine's one mtl::Context (the +// queue every other dsv41 kernel orders with). The weights are read in place (the tower's pinned host block is an MTLBuffer of +// the Metal runtime, so the GPU takes its device address): no slot copies. +// +// Credit: the structure of a DeepSeek-V4 vision tower on Metal (which operations are their own kernels, the BF16 rounding +// points) follows antirez/ds4 (MIT), commit 890898cf, metal/deepseek4_vision.metal and ds4_metal.m +// (ds4_gpu_deepseek4_vision_encode). Copyright (c) 2026 The ds4.c authors, (c) 2023-2026 The ggml authors, (c) 2023 DeepSeek. +// See third_party/antirez-ds4/LICENSE and NOTICE.md. +#include "ds4_vision_metal.h" + +#include "mtl_rt.h" + +#include +#include +#include +#include + +namespace strata::dsv41::metal_host { +mtl::Context& finish_cx(); +} + +namespace strata::ds4::vismetal { + +// the capability string: compiled in only here, where the Metal tower is built +static const char kMarker[] = "STRATA_DS4_METAL_VISION_V1"; +const char* marker() { return kMarker; } + +namespace { + +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 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 GemmArgs { uint64_t X, W, Y; uint32_t rows, in, out, ldx, ldw, ldy; float alpha; uint32_t flags; }; + +struct Runner { + mtl::Context& cx; + int lib = -1; + std::string err; + bool ok = true; + explicit Runner(mtl::Context& c) : cx(c) {} + + mtl::Pipeline pipe(const char* name) { + try { + return cx.pipeline(lib, name); + } catch (...) { + err = std::string("vision: no pipeline ") + name; + ok = false; + return mtl::Pipeline(); + } + } + template + void go(const mtl::Pipeline& p, const A& args, uint32_t gx, uint32_t gy, uint32_t tg) { + if (!ok || gx == 0 || gy == 0) return; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::V(args)}; + d.grid[0] = gx; + d.grid[1] = gy; + d.grid[2] = 1; + d.tg[0] = tg; + d.tg[1] = d.tg[2] = 1; + mtl::require_fit(p, tg, p.name.c_str()); + cx.run({d}); + } +}; + +} // namespace + +bool encode(const Args& a, std::string& err) { + mtl::Context& cx = strata::dsv41::metal_host::finish_cx(); + static int s_lib = -2; + if (s_lib == -2) s_lib = cx.compile({"../src/engine/msl/ds4_vision.metal"}, "", false); + if (s_lib < 0) { + err = "vision: compiling ds4_vision.metal failed"; + return false; + } + 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"); + if (!R.ok) { + err = R.err; + return false; + } + const int D = a.dim, I = a.inter, P = a.patch, M = a.merge, H3 = 3 * D, hd = a.hd, N = a.n, L = a.l; + const uint32_t smT = (uint32_t) std::min(512, (int) pSm.max_threads / 32 * 32); + if ((int) smT * 20 < N) { + err = "vision: this GPU's softmax group (" + std::to_string(smT) + " threads) is too small for " + std::to_string(N) + " patches"; + return false; + } + 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); + }; + auto add_rows = [&](uint64_t x, uint64_t y, int ldy, uint64_t bias, int rows) { + AddArgs g{x, y, bias, ldy, D}; + R.go(pAdd, g, (uint32_t) rows, 1, 256); + }; + auto rmsnorm = [&](uint64_t x, uint64_t w, uint64_t out) { + NormArgs g{x, w, out, D, a.eps}; + R.go(pNorm, g, (uint32_t) N, 1, 256); + }; + + // the patches + { + PatchArgs g{a.rgb, a.patches, a.width, a.gw, P}; + R.go(pPatch, g, (uint32_t) N, 1, 196); + } + gemm(a.patches, 3 * P * P, a.pe_w, 3 * P * P, false, a.x, D, false, N, 3 * P * P, D, 1.0f); + add_rows(a.x, 0, 0, a.pe_b, N); + + for (int l = 0; l < a.layers && R.ok; ++l) { + const uint64_t wqkv = at(a.weights, (size_t) l * a.block_n, 2), wo = at(wqkv, (size_t) H3 * D, 2), + w1 = at(wo, (size_t) D * D, 2), w2 = at(w1, (size_t) 2 * I * D, 2); + rmsnorm(a.x, a.n1[l], a.xn); + for (int r0 = 0; r0 < N; r0 += a.rb) { + const int nr = std::min(a.rb, N - r0); + gemm(at(a.xn, (size_t) r0 * D, 2), D, wqkv, D, false, a.y, H3, false, nr, D, H3, 1.0f); + RopeArgs g{a.y, a.bqkv[l], a.inv_freq, a.qkv, r0, a.gw, a.heads, hd}; + R.go(pRope, g, (uint32_t) nr, 1, 256); + } + for (int q0 = 0; q0 < N; q0 += a.qb) { + const int nq = std::min(a.qb, N - q0); + 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}; + 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, + 1.0f); + } + gemm(a.o, D, wo, D, false, a.y, D, false, nq, D, D, 1.0f); + add_rows(at(a.x, (size_t) q0 * D, 4), a.y, D, a.bo[l], nq); + } + rmsnorm(a.x, a.n2[l], a.xn); + for (int r0 = 0; r0 < N; r0 += a.rb) { + const int nr = std::min(a.rb, N - r0); + gemm(at(a.xn, (size_t) r0 * D, 2), D, w1, D, false, a.y, 2 * I, false, nr, D, 2 * I, 1.0f); + SiluArgs sg{a.y, a.mid, I}; + R.go(pSilu, sg, (uint32_t) ((I + 255) / 256), (uint32_t) nr, 256); + gemm(a.mid, I, w2, I, false, a.y, D, false, nr, I, D, 1.0f); + add_rows(at(a.x, (size_t) r0 * D, 4), a.y, D, 0, nr); + } + cx.split_runs(); // each block is its own command buffer: no long-running one for the system to cut short + } + rmsnorm(a.x, a.norm, a.xn); + + // the aligner + const int in1 = D * M * M; + { + UnfoldArgs g{a.xn, a.un, a.gh, a.gw, M, D, a.nw}; + R.go(pUnf, g, (uint32_t) L, 1, 256); + } + gemm(a.un, in1, at(a.weights, a.al1_off, 2), in1, false, a.y, a.out, false, L, in1, a.out, 1.0f); + { + BiasActArgs g{a.y, a.al1_b, a.a16, a.out, 1}; + R.go(pBias, g, (uint32_t) ((a.out + 255) / 256), (uint32_t) L, 256); + } + gemm(a.a16, a.out, at(a.weights, a.al2_off, 2), a.out, false, a.y, a.out, false, L, a.out, a.out, 1.0f); + { + BiasActArgs g{a.y, a.al2_b, a.outb, a.out, 0}; + R.go(pBias, g, (uint32_t) ((a.out + 255) / 256), (uint32_t) L, 256); + } + cx.flush_runs(); + if (!R.ok) { + err = R.err; + return false; + } + return true; +} + +} // namespace strata::ds4::vismetal diff --git a/metal/src/engine/msl/ds4_vision.metal b/metal/src/engine/msl/ds4_vision.metal new file mode 100644 index 000000000..3bd45f278 --- /dev/null +++ b/metal/src/engine/msl/ds4_vision.metal @@ -0,0 +1,281 @@ +// metal/src/engine/msl/ds4_vision.metal - DeepSeek-V4.1's vision tower on Metal. +// +// The kernels are the Metal form of src/dsv41/ds4_vision.cu's eight CUDA kernels (patchify, add_rows, rmsnorm, qkv_rope, +// softmax, silu_mul, unfold, bias_act), with the same arithmetic and the same BF16 rounding points, plus one tiled GEMM that +// takes the place of the cuBLAS calls (BF16 operands, FP32 accumulation, FP32 or BF16 result). +// +// Credit: the way a DeepSeek-V4 vision tower is laid out on Metal (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) follows antirez/ds4 (MIT), commit 890898cf, +// metal/deepseek4_vision.metal, metal/glm53_vision.metal and metal/glm53_bf16.metal. Copyright (c) 2026 The ds4.c authors, +// (c) 2023-2026 The ggml authors, (c) 2023 DeepSeek. See third_party/antirez-ds4/LICENSE and NOTICE.md. +// The tiled GEMM and the host side are Strata's own. +#include +#include + +using namespace metal; + +// every pointer is a device address (strata::gpu keeps all allocations in one residency set) +static inline float vis_bf(ushort b) { return as_type(uint(b) << 16); } +static inline ushort vis_bfr(float f) { // F32 -> BF16 bits, round to nearest even; a NaN stays a NaN + const uint u = as_type(f); + if ((u & 0x7fffffffu) > 0x7f800000u) return ushort((u >> 16) | 0x40u); + return ushort((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +// the block reduction of ds4_vision.cu: xor shuffles inside a warp, then the first warp over the warps' values +static inline float vis_block_reduce(float v, bool is_max, threadgroup float* red, ushort tid, ushort nthreads) { + for (ushort o = 16; o > 0; o >>= 1) { + const float u = simd_shuffle_xor(v, o); + v = is_max ? fmax(v, u) : v + u; + } + const ushort lane = tid & 31, w = tid >> 5, nw = (nthreads + 31) >> 5; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lane == 0) red[w] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float r = lane < nw ? red[lane] : (is_max ? -FLT_MAX : 0.0f); + for (ushort o = 16; o > 0; o >>= 1) { + const float u = simd_shuffle_xor(r, o); + r = is_max ? fmax(r, u) : r + u; + } + return r; // every thread of the group has it (the first warp's value is the same in all warps: each reduces red[]) +} + +// erf: MSL has none. erf(x) = 2/sqrt(pi) exp(-x^2) sum_n 2^n x^(2n+1) / (1 3 5 ... (2n+1)): positive terms only, so no cancellation. +static inline float vis_erf(float x) { + const float a = fabs(x); + if (a >= 4.0f) return x < 0.0f ? -1.0f : 1.0f; + const float a2 = a * a; + float term = a, sum = a; + for (int n = 1; n < 100; ++n) { + term *= 2.0f * a2 / float(2 * n + 1); + sum += term; + if (term < sum * 1.0e-8f) break; + } + const float r = 1.1283791670955126f * precise::exp(-a2) * sum; + return x < 0.0f ? -r : r; +} + +struct VisPatchArgs { + device const uchar* rgb; + device ushort* out; + int width, gw, p; +}; +// patch i of the gh x gw grid: its 3 x p x p values, channel-major. One group a patch. +kernel void vis_patchify(constant VisPatchArgs& a [[buffer(0)]], uint gi [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nt [[threads_per_threadgroup]]) { + const int pp = a.p * a.p; + for (int e = tid; e < 3 * pp; e += nt) { + const int c = e / pp, r = e % pp; + const int y = (int(gi) / a.gw) * a.p + r / a.p, x = (int(gi) % a.gw) * a.p + r % a.p; + const float u = float(a.rgb[(ulong(y) * ulong(a.width) + ulong(x)) * 3ul + ulong(c)]); + a.out[ulong(gi) * ulong(3 * pp) + ulong(e)] = vis_bfr(((u / 255.0f) - 0.5f) / 0.5f); + } +} + +struct VisAddArgs { + device float* x; + device const float* y; // 0: none + device const ushort* bias; // 0: none + int ldy, dim; +}; +// x[r] += y[r] (if any) + bias (if any). One group a row. +kernel void vis_add_rows(constant VisAddArgs& a [[buffer(0)]], uint r [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nt [[threads_per_threadgroup]]) { + for (int e = tid; e < a.dim; e += nt) { + float v = a.x[ulong(r) * ulong(a.dim) + ulong(e)]; + if (a.y != nullptr) v += a.y[ulong(r) * ulong(a.ldy) + ulong(e)]; + if (a.bias != nullptr) v += vis_bf(a.bias[e]); + a.x[ulong(r) * ulong(a.dim) + ulong(e)] = v; + } +} + +struct VisNormArgs { + device const float* x; + device const ushort* w; + device ushort* out; + int dim; + float eps; +}; +kernel void vis_rmsnorm(constant VisNormArgs& a [[buffer(0)]], uint r [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nt [[threads_per_threadgroup]]) { + threadgroup float red[32]; + device const float* xr = a.x + ulong(r) * ulong(a.dim); + float ss = 0.0f; + for (int e = tid; e < a.dim; e += nt) ss = fma(xr[e], xr[e], ss); + ss = vis_block_reduce(ss, false, red, tid, nt); + const float rs = rsqrt(ss / float(a.dim) + a.eps); + for (int e = tid; e < a.dim; e += nt) a.out[ulong(r) * ulong(a.dim) + ulong(e)] = vis_bfr(vis_bf(a.w[e]) * (xr[e] * rs)); +} + +struct VisRopeArgs { + device const float* y; + device const ushort* bias; + device const float* inv_freq; + device ushort* qkv; + int r0, gw, heads, hd; +}; +// a piece's rows of wqkv's output (+ bias, rounded to BF16) into qkv, q and k rotated. One group a row. +kernel void vis_qkv_rope(constant VisRopeArgs& a [[buffer(0)]], uint r [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nt [[threads_per_threadgroup]]) { + const int gr = a.r0 + int(r), dim = a.heads * a.hd, half_ = a.hd / 2, nf = a.hd / 4; + const int row = gr / a.gw, col = gr % a.gw; + device const float* yr = a.y + ulong(r) * ulong(3 * dim); + device ushort* o = a.qkv + ulong(gr) * ulong(3 * dim); + for (int i = tid; i < 2 * a.heads * half_; i += nt) { + const int base = (i / half_) * a.hd, d = i % half_; + const float x1 = vis_bf(vis_bfr(yr[base + d] + vis_bf(a.bias[base + d]))); + const float x2 = vis_bf(vis_bfr(yr[base + d + half_] + vis_bf(a.bias[base + d + half_]))); + const float ang = float(d < nf ? row : col) * a.inv_freq[d % nf]; + const float sn = precise::sin(ang), cs = precise::cos(ang); + o[base + d] = vis_bfr(x1 * cs - x2 * sn); + o[base + d + half_] = vis_bfr(x2 * cs + x1 * sn); + } + for (int e = 2 * dim + tid; e < 3 * dim; e += nt) o[e] = vis_bfr(yr[e] + vis_bf(a.bias[e])); +} + +struct VisSoftmaxArgs { + device const float* s; + device ushort* p; + int n; +}; +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]], + ushort tid [[thread_index_in_threadgroup]], ushort nt [[threads_per_threadgroup]]) { + threadgroup float red[32]; + device const float* sr = a.s + ulong(r) * ulong(a.n); + device ushort* pr = a.p + ulong(r) * ulong(a.n); + float v[kSmPer]; + 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; + m = fmax(m, v[k]); + } + m = vis_block_reduce(m, true, red, tid, nt); + float sum = 0.0f; + for (int k = 0; k < kSmPer; ++k) { + const int e = int(tid) + k * int(nt); + v[k] = e < a.n ? precise::exp(v[k] - m) : 0.0f; + sum += v[k]; + } + const float inv = 1.0f / vis_block_reduce(sum, false, red, tid, nt); + for (int k = 0; k < kSmPer; ++k) { + const int e = int(tid) + k * int(nt); + if (e < a.n) pr[e] = vis_bfr(v[k] * inv); + } +} + +struct VisSiluArgs { + device const float* y; + device ushort* mid; + int inter; +}; +// mid = silu(gate) * up, gate | up the halves of w1's row. Grid (ceil(inter / 256), rows). +kernel void vis_silu_mul(constant VisSiluArgs& a [[buffer(0)]], uint2 g [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + const int e = int(g.x) * 256 + int(tid); + if (e >= a.inter) return; + const ulong r = g.y; + const float gt = vis_bf(vis_bfr(a.y[r * ulong(2 * a.inter) + ulong(e)])); + const float up = vis_bf(vis_bfr(a.y[r * ulong(2 * a.inter) + ulong(a.inter + e)])); + a.mid[r * ulong(a.inter) + ulong(e)] = vis_bfr(vis_bf(vis_bfr(gt / (1.0f + precise::exp(-gt)))) * up); +} + +struct VisUnfoldArgs { + device const ushort* xn; + device ushort* un; + int gh, gw, m, dim, lw; +}; +// the aligner's input: block (bi, bj) of the m x m patches as one row, value c * m * m + kh * m + kw, zero past the grid. One group a row. +kernel void vis_unfold(constant VisUnfoldArgs& a [[buffer(0)]], uint l [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nt [[threads_per_threadgroup]]) { + const int bi = int(l) / a.lw, bj = int(l) % a.lw, mm = a.m * a.m; + for (int f = tid; f < a.dim * mm; f += nt) { + const int c = f / mm, k = f % mm, y = bi * a.m + k / a.m, x = bj * a.m + k % a.m; + a.un[ulong(l) * ulong(a.dim * mm) + ulong(f)] = + (y < a.gh && x < a.gw) ? a.xn[(ulong(y) * ulong(a.gw) + ulong(x)) * ulong(a.dim) + ulong(c)] : ushort(0); + } +} + +struct VisBiasActArgs { + device const float* y; + device const ushort* bias; + device ushort* out; + int n, gelu; +}; +// out = act(bf16(y + bias)): GELU (erf) or none. Grid (ceil(n / 256), rows). +kernel void vis_bias_act(constant VisBiasActArgs& a [[buffer(0)]], uint2 g [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + const int e = int(g.x) * 256 + int(tid); + if (e >= a.n) return; + const ulong r = g.y; + float v = vis_bf(vis_bfr(a.y[r * ulong(a.n) + ulong(e)] + vis_bf(a.bias[e]))); + if (a.gelu != 0) v = 0.5f * v * (1.0f + vis_erf(v * 0.70710678118654752f)); + a.out[r * ulong(a.n) + ulong(e)] = vis_bfr(v); +} + +// ---- the GEMM: Y[r][o] = alpha * sum_k X[r][k] * W(o, k), BF16 operands, FP32 accumulation. +// X: [rows][in] with row stride ldx. W: flags bit 0 clear: [out][in] with row stride ldw (k contiguous); +// bit 0 set: [in][out] with row stride ldw (o contiguous). Y: [rows][out], row stride ldy, FP32 (bit 1 clear) or BF16 (set). +// A 64 x 64 tile a group (4 simdgroups of 32 x 32), k in steps of 32, operands staged as F32 in threadgroup memory. +struct VisGemmArgs { + device const ushort* X; + device const ushort* W; + device uchar* Y; + uint rows, in, out, ldx, ldw, ldy; + float alpha; + uint flags; +}; +kernel void vis_gemm(constant VisGemmArgs& a [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float smem[4096]; // As [64][32] then Bs [32][64]; reused for the result tile [64][64] + threadgroup float* As = smem; + threadgroup float* Bs = smem + 2048; + const uint o0 = tgp.x * 64u, r0 = tgp.y * 64u; + const uint sr = (uint(sg) >> 1) * 32u, so = (uint(sg) & 1u) * 32u; + const bool wn = (a.flags & 1u) != 0u; + simdgroup_float8x8 acc[4][4]; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (uint k0 = 0; k0 < a.in; k0 += 32u) { + for (uint idx = tid; idx < 2048u; idx += 128u) { + const uint r = idx >> 5, kk = idx & 31u; + const uint gr = r0 + r, gk = k0 + kk; + As[idx] = (gr < a.rows && gk < a.in) ? vis_bf(a.X[ulong(gr) * ulong(a.ldx) + ulong(gk)]) : 0.0f; + } + if (!wn) { + for (uint idx = tid; idx < 2048u; idx += 128u) { + const uint o = idx >> 5, kk = idx & 31u; + const uint go = o0 + o, gk = k0 + kk; + Bs[kk * 64u + o] = (go < a.out && gk < a.in) ? vis_bf(a.W[ulong(go) * ulong(a.ldw) + ulong(gk)]) : 0.0f; + } + } else { + for (uint idx = tid; idx < 2048u; idx += 128u) { + const uint kk = idx >> 6, o = idx & 63u; + const uint go = o0 + o, gk = k0 + kk; + Bs[idx] = (go < a.out && gk < a.in) ? vis_bf(a.W[ulong(gk) * ulong(a.ldw) + ulong(go)]) : 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint k8 = 0; k8 < 4u; ++k8) { + simdgroup_float8x8 fa[4], fb[4]; + for (uint i = 0; i < 4u; ++i) simdgroup_load(fa[i], As + (sr + i * 8u) * 32u + k8 * 8u, 32); + for (uint j = 0; j < 4u; ++j) simdgroup_load(fb[j], Bs + k8 * 8u * 64u + so + j * 8u, 64); + for (uint i = 0; i < 4u; ++i) + for (uint j = 0; j < 4u; ++j) simdgroup_multiply_accumulate(acc[i][j], fa[i], fb[j], acc[i][j]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i = 0; i < 4u; ++i) + for (uint j = 0; j < 4u; ++j) simdgroup_store(acc[i][j], smem + (sr + i * 8u) * 64u + so + j * 8u, 64); + threadgroup_barrier(mem_flags::mem_threadgroup); + const bool bf = (a.flags & 2u) != 0u; + for (uint idx = tid; idx < 4096u; idx += 128u) { + const uint r = idx >> 6, o = idx & 63u; + const uint gr = r0 + r, go = o0 + o; + if (gr >= a.rows || go >= a.out) continue; + const float v = a.alpha * smem[idx]; + if (bf) ((device ushort*) a.Y)[ulong(gr) * ulong(a.ldy) + ulong(go)] = vis_bfr(v); + else ((device float*) a.Y)[ulong(gr) * ulong(a.ldy) + ulong(go)] = v; + } +} diff --git a/packaging/mac/build_dsv41_engine.sh b/packaging/mac/build_dsv41_engine.sh index 2ade4187d..0fb32281f 100644 --- a/packaging/mac/build_dsv41_engine.sh +++ b/packaging/mac/build_dsv41_engine.sh @@ -51,10 +51,12 @@ 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 mtp mtp_chain ple; do +for f in ds4_init ds4_iq ds4_mm ds4_quant ds4_vision 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/" +# antirez/ds4 (MIT): the Metal vision kernels follow it; its licence travels with the MSL +cp "$ROOT/third_party/antirez-ds4/LICENSE" "$STAGE/metal/src/engine/msl/LICENSE-antirez-ds4" python3 "$ROOT/packaging/package_engine.py" --src "$STAGE" --backend metal --version "$VERSION" \ --platform macos-arm64 --arch apple-gpu --entrypoint strata-dsv41 --build-id "$BUILD_ID" --no-picture-reader --include 'metal/**' --out "$OUT" || exit 1 diff --git a/src/dsv41/ds4_model.cu b/src/dsv41/ds4_model.cu index 1727eba04..91835edb2 100644 --- a/src/dsv41/ds4_model.cu +++ b/src/dsv41/ds4_model.cu @@ -31,6 +31,9 @@ #include "ds4_pool.hpp" #include "ds4_store.hpp" #include "ds4_vision.cuh" +#if defined(STRATA_GPU_METAL) +#include "ds4_vision_metal.h" // the Metal tower's capability marker +#endif #include "strata/dsv41/exl3_cpu.hpp" #include "strata/dsv41/exl3.hpp" #include "strata/kernels/f16_bits.hpp" @@ -7736,6 +7739,11 @@ struct Model::Impl { CK(cudaSetDevice(d.id)); CK(cudaMemGetInfo(&d.free0, &d.total)); // (the split sees what the tower left) d.free0 = uma_free(d.id, d.free0); +#if defined(STRATA_GPU_METAL) + std::fprintf(stderr, "ds4: vision tower (%d blocks) on Metal: %.0f MiB, up to %d patches, %.1f s [%s]\n", vg.layers, + (double) vt.device_bytes() / (1 << 20), vt.max_patches(), (now_ms() - t0) / 1000.0, + strata::ds4::vismetal::marker()); +#endif if (o.verbose) std::fprintf(stderr, "ds4: vision tower (%d blocks) on CUDA%d: %.0f MiB%s, up to %d patches, %.1f s\n", vg.layers, d.id, (double) vt.device_bytes() / (1 << 20), diff --git a/src/dsv41/ds4_vision.cu b/src/dsv41/ds4_vision.cu index ea0970b79..aa026813c 100644 --- a/src/dsv41/ds4_vision.cu +++ b/src/dsv41/ds4_vision.cu @@ -15,6 +15,9 @@ #include #include "ds4_numerics.cuh" +#if defined(STRATA_GPU_METAL) +#include "ds4_vision_metal.h" // the tower as Metal dispatches (metal/src/engine/ds4_vision_metal.cpp) +#endif #include #include @@ -545,6 +548,30 @@ bool VisionTower::encode(const uint8_t* rgb, int width, int height, std::vector< cudaSetDevice(dev_); cublasHandle_t h = (cublasHandle_t) blas_; cudaStream_t s = (cudaStream_t) st_, cs = (cudaStream_t) cs_; +#if defined(STRATA_GPU_METAL) + { // Metal: the whole tower as dispatches of msl/ds4_vision.metal; the weights are read in place from the pinned block + uint16_t* wdev = nullptr; + VK(cudaHostGetDevicePointer(&wdev, (void*) host_, 0)); + auto A = [](const void* p) { return (uint64_t) (uintptr_t) p; }; + strata::ds4::vismetal::Args a; + a.width = width, a.height = height, a.gh = gh, a.gw = gw, a.n = N, a.l = L, a.nw = n_w; + a.dim = D, a.inter = I, a.patch = P, a.merge = M, a.heads = g.heads, a.hd = hd, a.layers = g.layers, a.out = g.out; + a.rb = rb_, a.qb = (int) std::max(1, std::min((size_t) N, (size_t) qb_ * max_n_ / (size_t) N)); + a.eps = g.eps, a.scale = 1.0f / std::sqrt((float) hd); + a.rgb = A(rgb_), a.patches = A(patch_), a.x = A(x_), a.xn = A(xn_), a.qkv = A(qkv_), a.y = A(y_), a.o = A(o_), a.mid = A(mid_); + a.s = A(s_), a.p = A(p_), a.un = A(un_), a.a16 = A(a16_), a.outb = A(outb_), a.inv_freq = A(inv_freq_); + a.weights = A(wdev), a.block_n = block_n_, a.al1_off = al1_off_, a.al2_off = al2_off_; + a.pe_w = A(pe_w_), a.pe_b = A(pe_b_), a.norm = A(norm_), a.al1_b = A(al1_b_), a.al2_b = A(al2_b_); + for (int l = 0; l < g.layers; ++l) { + a.n1.push_back(A(n1_[l])), a.n2.push_back(A(n2_[l])), a.bqkv.push_back(A(bqkv_[l])), a.bo.push_back(A(bo_[l])); + } + VK(cudaMemcpy(rgb_, rgb, (size_t) N * 3 * P * P, cudaMemcpyHostToDevice)); + if (!strata::ds4::vismetal::encode(a, err)) return false; + rows.resize((size_t) L * g.out); + VK(cudaMemcpy(rows.data(), outb_, rows.size() * 2, cudaMemcpyDeviceToHost)); + return true; + } +#endif // the weights in the order the encode takes them: the blocks, then the aligner's in pieces of a slot std::vector pcs; for (int l = 0; l < g.layers; ++l) pcs.push_back({host_ + (size_t) l * block_n_, block_n_}); diff --git a/third_party/antirez-ds4/LICENSE b/third_party/antirez-ds4/LICENSE new file mode 100644 index 000000000..c5b10d432 --- /dev/null +++ b/third_party/antirez-ds4/LICENSE @@ -0,0 +1,23 @@ +MIT License + +Copyright (c) 2026 The ds4.c authors +Copyright (c) 2023-2026 The ggml authors +Copyright (c) 2023 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/third_party/antirez-ds4/NOTICE.md b/third_party/antirez-ds4/NOTICE.md new file mode 100644 index 000000000..d19089e90 --- /dev/null +++ b/third_party/antirez-ds4/NOTICE.md @@ -0,0 +1,14 @@ +# antirez/ds4, MIT + +The Metal form of DeepSeek-V4.1's vision tower follows antirez/ds4 (https://github.com/antirez/ds4), commit 890898cf, MIT License. +The licence text is in `LICENSE` beside this file; it is copied unchanged next to the Metal sources in the Mac engine package. + +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/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` + +The tiled BF16 GEMM, the rest of the host side and the wiring into `src/dsv41/ds4_vision.cu` are Strata's own. + +Copyright (c) 2026 The ds4.c authors; (c) 2023-2026 The ggml authors; (c) 2023 DeepSeek.