Support DeepSeek V4 Flash Vision Experimental

metal/DS4-LICENSE.txt

@@ -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.

metal/deepseek4_vision.metal

@@ -0,0 +1,168 @@
+// Derived from antirez/ds4 0aaea5a238fb41a35106a551e73c8409dfb751ac (MIT).
+// Copyright (c) 2026 The ds4.c authors; (c) 2023 DeepSeek.
+// See DS4-LICENSE.txt for the complete license.
+// DeepSeek V4 Flash Vision-Exp operations not covered by shared BF16 matmuls.
+
+struct deepseek4_vision_rows_args {
+    uint width;
+    uint rows;
+};
+
+struct deepseek4_vision_qkv_args {
+    uint rows;
+    uint grid_width;
+};
+
+struct deepseek4_vision_align_args {
+    uint grid_height;
+    uint grid_width;
+    uint output_rows;
+};
+
+static inline float deepseek4_vision_round_bf16(float value) {
+    uint bits = as_type<uint>(value);
+    if ((bits & 0x7f800000u) == 0x7f800000u) return value;
+    bits += 0x00007fffu + ((bits >> 16u) & 1u);
+    return as_type<float>(bits & 0xffff0000u);
+}
+
+static inline float deepseek4_vision_erf(float x) {
+    const float sign = x < 0.0f ? -1.0f : 1.0f;
+    const float a = abs(x);
+    const float t = 1.0f / (1.0f + 0.3275911f * a);
+    const float p = (((((1.061405429f * t - 1.453152027f) * t) +
+                       1.421413741f) * t - 0.284496736f) * t +
+                       0.254829592f) * t;
+    return sign * (1.0f - p * exp(-a * a));
+}
+
+kernel void kernel_deepseek4_vision_qkv_rope(
+        constant deepseek4_vision_qkv_args &args,
+        device const float                 *qkv,
+        device const ushort                *bias,
+        device float                       *q,
+        device float                       *k,
+        device float                       *v,
+        uint2 group [[threadgroup_position_in_grid]],
+        ushort lane [[thread_index_in_simdgroup]]) {
+    const uint row = group.x;
+    const uint head = group.y;
+    if (row >= args.rows || head >= 16u) return;
+    const ulong qkv_base = (ulong)row * 3072u + (ulong)head * 64u;
+    const ulong out_base = (ulong)row * 1024u + (ulong)head * 64u;
+    const float q0 = deepseek4_vision_round_bf16(
+        qkv[qkv_base + lane] +
+        glm53_bf16_to_f32(bias[(ulong)head * 64u + lane]));
+    const float q1 = deepseek4_vision_round_bf16(
+        qkv[qkv_base + lane + 32u] +
+        glm53_bf16_to_f32(bias[(ulong)head * 64u + lane + 32u]));
+    const float k0 = deepseek4_vision_round_bf16(
+        qkv[qkv_base + 1024u + lane] +
+        glm53_bf16_to_f32(bias[1024u + (ulong)head * 64u + lane]));
+    const float k1 = deepseek4_vision_round_bf16(
+        qkv[qkv_base + 1024u + lane + 32u] +
+        glm53_bf16_to_f32(bias[1024u + (ulong)head * 64u + lane + 32u]));
+    const uint y = row / args.grid_width;
+    const uint x = row - y * args.grid_width;
+    const uint pos = lane < 16u ? y : x;
+    const uint freq_index = lane & 15u;
+    const float inv_freq = powr(10000.0f, -(float)freq_index / 16.0f);
+    const float angle = (float)pos * inv_freq;
+    const float cs = cos(angle);
+    const float sn = sin(angle);
+    q[out_base + lane] = deepseek4_vision_round_bf16(q0 * cs - q1 * sn);
+    q[out_base + lane + 32u] =
+        deepseek4_vision_round_bf16(q1 * cs + q0 * sn);
+    k[out_base + lane] = deepseek4_vision_round_bf16(k0 * cs - k1 * sn);
+    k[out_base + lane + 32u] =
+        deepseek4_vision_round_bf16(k1 * cs + k0 * sn);
+    v[out_base + lane] = deepseek4_vision_round_bf16(
+        qkv[qkv_base + 2048u + lane] +
+        glm53_bf16_to_f32(bias[2048u + (ulong)head * 64u + lane]));
+    v[out_base + lane + 32u] = deepseek4_vision_round_bf16(
+        qkv[qkv_base + 2048u + lane + 32u] +
+        glm53_bf16_to_f32(
+            bias[2048u + (ulong)head * 64u + lane + 32u]));
+}
+
+kernel void kernel_deepseek4_vision_round_bf16(
+        constant deepseek4_vision_rows_args &args,
+        device float                       *x,
+        uint2 gid [[thread_position_in_grid]]) {
+    if (gid.x >= args.width || gid.y >= args.rows) return;
+    const ulong off = (ulong)gid.y * args.width + gid.x;
+    x[off] = deepseek4_vision_round_bf16(x[off]);
+}
+
+kernel void kernel_deepseek4_vision_add_residual(
+        constant deepseek4_vision_rows_args &args,
+        device float                       *x,
+        device const float                 *residual,
+        uint2 gid [[thread_position_in_grid]]) {
+    if (gid.x >= args.width || gid.y >= args.rows) return;
+    const ulong off = (ulong)gid.y * args.width + gid.x;
+    x[off] = deepseek4_vision_round_bf16(x[off] + residual[off]);
+}
+
+static inline void deepseek4_vision_swiglu_split(
+        constant deepseek4_vision_rows_args &args,
+        device const float                 *gate_up,
+        device float                       *out,
+        uint2 gid,
+        bool round_silu) {
+    if (gid.x >= args.width || gid.y >= args.rows) return;
+    const ulong source = (ulong)gid.y * args.width * 2u + gid.x;
+    const float gate = gate_up[source];
+    const float up = gate_up[source + args.width];
+    float activated = gate / (1.0f + exp(-gate));
+    if (round_silu) activated = deepseek4_vision_round_bf16(activated);
+    out[(ulong)gid.y * args.width + gid.x] = deepseek4_vision_round_bf16(activated * up);
+}
+
+kernel void kernel_deepseek4_vision_swiglu_split(
+        constant deepseek4_vision_rows_args &args,
+        device const float                 *gate_up,
+        device float                       *out,
+        uint2 gid [[thread_position_in_grid]]) {
+    deepseek4_vision_swiglu_split(args, gate_up, out, gid, false);
+}
+
+/* Match F.unfold(x, 3, stride=3): channels are outermost, followed by the
+ * row-major 3x3 neighborhood. Values outside the ViT grid are zero. */
+kernel void kernel_deepseek4_vision_aligner_reorder(
+        constant deepseek4_vision_align_args &args,
+        device const float                  *x,
+        device float                        *out,
+        uint2 gid [[thread_position_in_grid]]) {
+    const uint d = gid.x;
+    const uint row = gid.y;
+    if (d >= 9216u || row >= args.output_rows) return;
+    const uint output_width = (args.grid_width + 2u) / 3u;
+    const uint block_y = row / output_width;
+    const uint block_x = row - block_y * output_width;
+    const uint channel = d / 9u;
+    const uint within = d - channel * 9u;
+    const uint source_y = block_y * 3u + within / 3u;
+    const uint source_x = block_x * 3u + within % 3u;
+    float value = 0.0f;
+    if (source_y < args.grid_height && source_x < args.grid_width) {
+        const ulong source_row = (ulong)source_y * args.grid_width + source_x;
+        value = x[source_row * 1024u + channel];
+    }
+    out[(ulong)row * 9216u + d] = value;
+}
+
+kernel void kernel_deepseek4_vision_gelu_bias(
+        constant deepseek4_vision_rows_args &args,
+        device const float                 *x,
+        device const ushort                *bias,
+        device float                       *out,
+        uint2 gid [[thread_position_in_grid]]) {
+    if (gid.x >= args.width || gid.y >= args.rows) return;
+    const ulong off = (ulong)gid.y * args.width + gid.x;
+    const float value = deepseek4_vision_round_bf16(
+        x[off] + glm53_bf16_to_f32(bias[gid.x]));
+    out[off] = deepseek4_vision_round_bf16(
+        0.5f * value *
+        (1.0f + deepseek4_vision_erf(value * 0.7071067811865475f)));
+}

metal/dsv4_misc.metal

@@ -6894,3 +6894,65 @@ kernel void kernel_dsv4_softmax_pool_ratio4_direct(
 
     dst[ic * args.head_dim + id] = acc/sum;
 }
+
+// Vision-Exp routing from DS4 0aaea5a238fb41a35106a551e73c8409dfb751ac.
+struct ds4_metal_args_dsv4_router_select_visual {
+    uint32_t hash_rows;
+    uint32_t vocab_size;
+    uint32_t n_tokens;
+    uint32_t has_bias;
+    uint32_t hash_mode;
+};
+/* Vision-Exp uses score-based visual routing even in the first three layers,
+ * where ordinary text rows use a token-id hash table. One threadgroup owns a
+ * row so the image/text decision is uniform across all barriers. */
+kernel void kernel_dsv4_router_select_visual_batch(
+        constant ds4_metal_args_dsv4_router_select_visual & args,
+        device const float *probs,
+        device const float *bias,
+        device const float *visual_bias,
+        device const int32_t *hash,
+        device const int32_t *tokens,
+        device int32_t *selected,
+        threadgroup float *scratch [[threadgroup(0)]],
+        uint tid [[thread_position_in_threadgroup]],
+        uint row [[threadgroup_position_in_grid]]) {
+    if (tid >= 256u || row >= args.n_tokens) return;
+
+    const int32_t token = tokens[row];
+    const bool image = token >= 0 && (uint32_t)token >= args.vocab_size;
+    device int32_t *out = selected + (uint64_t)row * 6u;
+    if (args.hash_mode && !image) {
+        const uint hash_row = token >= 0 && (uint32_t)token < args.hash_rows
+            ? (uint32_t)token : 0u;
+        if (tid < 6u) out[tid] = hash[(uint64_t)hash_row * 6u + tid];
+        return;
+    }
+
+    threadgroup float *scores = scratch;
+    threadgroup int32_t *indices = (threadgroup int32_t *)(scratch + 256u);
+    const float p = probs[(uint64_t)row * 256u + tid];
+    scores[tid] = p + (image ? visual_bias[tid]
+                             : (args.has_bias ? bias[tid] : 0.0f));
+    indices[tid] = (int32_t)tid;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    for (uint k = 2u; k <= 256u; k <<= 1u) {
+        for (uint j = k >> 1u; j > 0u; j >>= 1u) {
+            const uint other = tid ^ j;
+            if (other > tid) {
+                const bool descending = (tid & k) == 0u;
+                const int32_t a = indices[tid];
+                const int32_t b = indices[other];
+                const float sa = scores[(uint)a];
+                const float sb = scores[(uint)b];
+                if ((descending && sa < sb) || (!descending && sa > sb)) {
+                    indices[tid] = b;
+                    indices[other] = a;
+                }
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+    }
+    if (tid < 6u) out[tid] = indices[tid];
+}