From b43fcce62dc6295884a5c572d6b4444c96853a96 Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 13:52:51 +0200 Subject: [PATCH 1/7] v41: fuse the decode hyper-connection glue (19 -> 6 dispatches per layer) DeepSeek's production decode runs each half-layer's HC work in one kernel (Mega-mHC); ds4's V4.1 graph mirrored the reference op by op. Three fusions, all byte-identical to the standalone sequence they replace: - kernel_dsv41_hc_collapse_norm4: split/sinkhorn + pre-weighted collapse + BF16 + weighted RMSNorm + BF16 (5 -> 1), keeping the standalone kernels' thread mapping, simd trees and 1/sqrt scale. - kernel_dsv41_hc_expand4_bf16: post/comb expand + BF16, optional carry of split[0..3] into pre (2-3 -> 1). - kernel_dsv4_hc_rms_norm_mix_f16 now accepts V4.1's 4x5120 HC row: the norm runs 1024 threads there too and 20480 is a whole multiple of the matvec's 4096-value stride, so the exact replica and empty tail hold. Single box only (TP keeps the unfused path). Rollback: DS4_METAL_DISABLE_V41_HC_FUSE. tests/test_deepseek41_metal --hc-fuse compares every output of both paths byte for byte (6 rounds, M4 Pro). --- ds4.c | 50 ++++++++++ ds4_deepseek41_gpu.h | 18 ++++ ds4_metal.m | 90 +++++++++++++++++- metal/dsv41.metal | 174 ++++++++++++++++++++++++++++++++++ metal/dsv4_hc.metal | 5 +- tests/test_deepseek41_metal.c | 93 ++++++++++++++++++ 6 files changed, 427 insertions(+), 3 deletions(-) diff --git a/ds4.c b/ds4.c index 582d5afb40..4bccb84b27 100644 --- a/ds4.c +++ b/ds4.c @@ -39838,6 +39838,41 @@ static bool ds41_graph_logits(ds41_gpu_graph *g, const ds4_model *m, (uint64_t)DS4_N_VOCAB * sizeof(float)); } +#if defined(__APPLE__) +/* DeepSeek's production decode runs each half-layer's hyper-connection work + * in one kernel ("Mega-mHC", tech report 3.2). ds4's V4.1 graph mirrored the + * reference op by op: 19 HC glue dispatches per layer (norm, mix matvec, + * split, collapse, BF16, weighted norm, BF16, expand, BF16, ...). The fused + * path keeps every reduction tree and rounding point of that sequence + * (tests/test_deepseek41_metal --hc-fuse) and issues 6. Single box only; the + * TP graph keeps the unfused sequence. */ +static bool ds41_hc_fused(const ds41_gpu_graph *g) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_HC_FUSE") == NULL && + ds4_gpu_hc_rms_norm_mix_f16_available() != 0; + } + return enabled && g->tp_world == 1; +} + +/* Plain RMSNorm over the flattened HC row and the F16 mix matvec, one dispatch. */ +static bool ds41_hc_mix_fused(ds41_gpu_graph *g, const ds4_model *m, + const ds4_tensor *fn, const ds4_gpu_tensor *residual) { + return ds4_gpu_hc_rms_norm_mix_f16_tensor(g->mix, residual, m->map, m->size, fn->abs_offset, + DS4_N_HC * DS4_N_EMBD, (uint32_t)fn->dim[1], DS4_RMS_EPS) != 0; +} + +static bool ds41_hc_collapse_norm_fused(ds41_gpu_graph *g, const ds4_model *m, + const ds4_layer_weights *l, bool ffn) { + return ds4_gpu_dsv41_hc_collapse_norm(ffn ? g->ffn_split : g->attn_split, g->x, g->norm, g->mix, + ffn ? g->attn_split : g->pre, ffn ? g->after_attn : g->residual, m->map, m->size, + (ffn ? l->hc_ffn_scale : l->hc_attn_scale)->abs_offset, + (ffn ? l->hc_ffn_base : l->hc_attn_base)->abs_offset, + (ffn ? l->ffn_norm : l->attn_norm)->abs_offset, + DS4_N_EMBD, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS, DS4_RMS_EPS) != 0; +} +#endif + static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { if (ds41_engram_layer(il) && !ds41_image_at(g, g->pos)) { @@ -39847,6 +39882,11 @@ static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, g->engram_q_norm[i], g->engram_k_norm[i], NULL, DS4_N_EMBD, 1, DS4_RMS_EPS)) return false; } +#if defined(__APPLE__) + if (ds41_hc_fused(g)) + return ds41_hc_mix_fused(g, m, l->hc_attn_fn, g->residual) && + ds41_hc_collapse_norm_fused(g, m, l, false); +#endif return ds41_hc_mix(g, m, l, false) && ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, l->attn_norm); @@ -39854,6 +39894,12 @@ static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, static bool ds41_graph_after_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { +#if defined(__APPLE__) + if (ds41_hc_fused(g)) + return ds4_gpu_dsv41_hc_expand4(g->after_attn, g->block, g->residual, g->attn_split, NULL, DS4_N_EMBD) && + ds41_hc_mix_fused(g, m, l->hc_ffn_fn, g->after_attn) && + ds41_hc_collapse_norm_fused(g, m, l, true); +#endif return ds4_gpu_hc_expand_split_tensor(g->after_attn, g->block, g->residual, g->attn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->after_attn, DS4_N_EMBD * DS4_N_HC) && ds41_hc_mix(g, m, l, true) && @@ -40138,6 +40184,10 @@ static bool ds41_attention_batch(ds41_gpu_graph *g, const ds4_model *m, } static bool ds41_graph_after_moe(ds41_gpu_graph *g) { +#if defined(__APPLE__) + if (ds41_hc_fused(g)) + return ds4_gpu_dsv41_hc_expand4(g->residual, g->block, g->after_attn, g->ffn_split, g->pre, DS4_N_EMBD); +#endif return ds4_gpu_hc_expand_split_tensor(g->residual, g->block, g->after_attn, g->ffn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->residual, DS4_N_EMBD * DS4_N_HC) && ds4_gpu_tensor_copy(g->pre, 0, g->ffn_split, 0, DS4_N_HC * sizeof(float)); diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index 39f11699f8..7c794f2d4e 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -123,6 +123,24 @@ int ds4_gpu_dsv41_projection_rows(ds4_gpu_tensor *out, int ds4_gpu_dsv41_gather_kv(ds4_gpu_tensor *out, const ds4_gpu_tensor *source, const ds4_gpu_tensor *ids, uint32_t source_rows, uint32_t selected_rows); +/* Decode HC glue for one token row, HC=4, byte-identical to the standalone + * split/collapse/BF16/norm/BF16 and expand/BF16 dispatch sequences. + * collapse_norm: split(mix) -> split; x = bf16(pre-weighted collapse of + * residual); norm = bf16(rmsnorm(x) * weight). `pre` is the coefficient row + * of the PREVIOUS sublayer's split (the first four floats of that tensor). + * expand4: out = bf16(post/comb expand of block into residual); when `pre` + * is given, split[0..3] is also copied into it. */ +int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4_gpu_tensor *norm, + const ds4_gpu_tensor *mix, const ds4_gpu_tensor *pre, + const ds4_gpu_tensor *residual, + const void *model_map, uint64_t model_size, + uint64_t scale_offset, uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters, + float hc_eps, float norm_eps); +int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, + const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, + ds4_gpu_tensor *pre, uint32_t n_embd); #ifdef __cplusplus } diff --git a/ds4_metal.m b/ds4_metal.m index b61cbee8d7..8795bbaa13 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -44814,7 +44814,10 @@ int ds4_gpu_hc_rms_norm_mix_f16_tensor( uint32_t out_dim, float eps) { if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !x || !model_map || n != 16384u || out_dim != 24u) return 0; + /* 16384 = V4's 4x4096 HC row, 20480 = V4.1's 4x5120: both norm with 1024 + * threads and both are whole multiples of the matvec's 4096-value stride, + * so the kernel's virtual-thread replica and empty tail hold for either. */ + if (!out || !x || !model_map || (n != 16384u && n != 20480u) || out_dim != 24u) return 0; @autoreleasepool { const uint64_t row_bytes = (uint64_t)n * sizeof(uint16_t); @@ -47707,6 +47710,91 @@ int ds4_gpu_dsv41_engram_add(ds4_gpu_tensor *residual, } } +/* V4.1 decode HC glue (metal/dsv41.metal): one token row, HC=4. */ +typedef struct { uint32_t n_embd, sinkhorn_iters; float hc_eps, norm_eps; uint32_t copy_pre; } ds4_gpu_dsv41_hc_args; + +int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4_gpu_tensor *norm, + const ds4_gpu_tensor *mix, const ds4_gpu_tensor *pre, + const ds4_gpu_tensor *residual, + const void *model_map, uint64_t model_size, + uint64_t scale_offset, uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters, + float hc_eps, float norm_eps) { + const uint64_t weight_bytes = (uint64_t)n_embd * sizeof(float); + if (n_hc != 4u || !n_embd || n_embd % 4u || !sinkhorn_iters || !model_map || + !dsv41_tensor_has_floats(split, 24) || !dsv41_tensor_has_floats(mix, 24) || + !dsv41_tensor_has_floats(pre, 4) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || + !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(norm, n_embd) || + scale_offset > model_size || 12u > model_size - scale_offset || + base_offset > model_size || 96u > model_size - base_offset || + norm_weight_offset > model_size || weight_bytes > model_size - norm_weight_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_collapse_norm4"); + if (!pipeline) return 0; + /* The norm's reduction tree is the standalone kernel's only at its thread count. */ + const NSUInteger threads = ds4_gpu_rms_norm_threads(n_embd); + const NSUInteger shared_bytes = ((NSUInteger)n_embd + 32u) * sizeof(float); + if (threads > pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) return 0; + uint64_t scale_inner = 0, base_inner = 0, norm_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, 12u, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, 96u, &base_inner); + id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, + weight_bytes, &norm_inner); + if (!scalebuf || !basebuf || !normwbuf) return 0; + const ds4_gpu_dsv41_hc_args args = {n_embd, sinkhorn_iters, hc_eps, norm_eps, 0}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(mix) offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(pre) offset:ds4_gpu_tensor_offset(pre) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(residual) offset:ds4_gpu_tensor_offset(residual) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(split) offset:ds4_gpu_tensor_offset(split) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:7]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:8]; + [enc setBuffer:ds4_gpu_tensor_buffer(norm) offset:ds4_gpu_tensor_offset(norm) atIndex:9]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 HC collapse/norm"); + } +} + +int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, + const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, + ds4_gpu_tensor *pre, uint32_t n_embd) { + if (!n_embd || !dsv41_tensor_has_floats(out, 4ull * n_embd) || + !dsv41_tensor_has_floats(block, n_embd) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || + !dsv41_tensor_has_floats(split, 24) || (pre && !dsv41_tensor_has_floats(pre, 4))) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_expand4_bf16"); + if (!pipeline) return 0; + const ds4_gpu_dsv41_hc_args args = {n_embd, 0, 0.0f, 0.0f, pre != NULL}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + const ds4_gpu_tensor *buffers[] = {block, residual, split, out, pre ? pre : split}; + for (NSUInteger i = 0; i < 5; i++) + [enc setBuffer:ds4_gpu_tensor_buffer(buffers[i]) + offset:ds4_gpu_tensor_offset(buffers[i]) atIndex:i + 1]; + [enc dispatchThreadgroups:MTLSizeMake(((uint64_t)n_embd + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 HC expand"); + } +} + int ds4_gpu_dsv41_carry_copy(ds4_gpu_tensor *packed, uint32_t row_offset, ds4_gpu_tensor *plain, uint32_t width, uint32_t rows, uint32_t format, bool pack) { diff --git a/metal/dsv41.metal b/metal/dsv41.metal index 03afd310dd..8a54118a78 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -291,3 +291,177 @@ kernel void kernel_dsv41_indexer_scores_packed( } } #endif + +/* Decode-time V4.1 hyper-connection glue, one token row, HC=4. + * + * The released graph runs each half-layer's HC work as separate dispatches: + * split/sinkhorn, the pre-weighted collapse of the four streams, a BF16 + * rounding pass, the weighted RMSNorm and another BF16 pass (and, after the + * sublayer, the post/comb expand plus its BF16 pass). DeepSeek's production + * decode does the whole thing in one "Mega-mHC" kernel. These two kernels + * are that fusion for ds4's graph: every reduction keeps the standalone + * kernel's thread mapping and accumulation order, and every rounding point + * is the same dsv41_bf16, so the outputs are byte-identical (checked by + * tests/test_deepseek41_metal --hc-fuse). */ +struct ds4_metal_args_dsv41_hc { + uint n_embd; + uint sinkhorn_iters; + float hc_eps; + float norm_eps; + uint copy_pre; +}; + +static inline float4 dsv41_bf16x4(float4 v) { + return float4(dsv41_bf16(v.x), dsv41_bf16(v.y), dsv41_bf16(v.z), dsv41_bf16(v.w)); +} + +/* kernel_dsv4_hc_split_sinkhorn's HC == 4 body, verbatim, on one lane. */ +static inline void dsv41_hc_split4(device const float *mix, device const float *scale, + device const float *base, device float *out, + uint sinkhorn_iters, float epsv) { + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + const float4 pre_z = *((device const float4 *)mix) * pre_scale + *((device const float4 *)base); + *((device float4 *)out) = ds4_hc_sigmoid(pre_z) + epsv; + + const float4 post_z = *((device const float4 *)(mix + 4)) * post_scale + *((device const float4 *)(base + 4)); + *((device float4 *)(out + 4)) = ds4_hc_twice_sigmoid(post_z); + + float4 r0 = *((device const float4 *)(mix + 8)) * comb_scale + *((device const float4 *)(base + 8)); + float4 r1 = *((device const float4 *)(mix + 12)) * comb_scale + *((device const float4 *)(base + 12)); + float4 r2 = *((device const float4 *)(mix + 16)) * comb_scale + *((device const float4 *)(base + 16)); + float4 r3 = *((device const float4 *)(mix + 20)) * comb_scale + *((device const float4 *)(base + 20)); + + const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w)); + const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w)); + const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w)); + const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w)); + + r0 = exp(r0 - m0); + r1 = exp(r1 - m1); + r2 = exp(r2 - m2); + r3 = exp(r3 - m3); + + r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv; + r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv; + r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv; + r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv; + + float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + + for (uint iter = 1; iter < sinkhorn_iters; ++iter) { + r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); + r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); + r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); + r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); + + col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + } + + *((device float4 *)(out + 8)) = r0; + *((device float4 *)(out + 12)) = r1; + *((device float4 *)(out + 16)) = r2; + *((device float4 *)(out + 20)) = r3; +} + +/* split(mix) -> split_out; x = bf16(sum_h pre[h] * residual[h]); norm = + * bf16(rmsnorm(x) * weight). `pre` is the coefficient row the previous + * sublayer produced (V4.1's single-pass mHC), not this split's. One + * threadgroup of ds4_gpu_rms_norm_threads(n_embd) threads: the collapse + * keeps kernel_dsv4_hc_weighted_sum's per-stream order and the norm keeps + * kernel_rms_norm_mul_f32_4's loop, simd tree and 1/sqrt scale. */ +kernel void kernel_dsv41_hc_collapse_norm4( + constant ds4_metal_args_dsv41_hc &args, + device const float *mix, + device const float *scale, + device const float *base, + device const float *pre, + device const float4 *residual, + device float *split_out, + device float4 *x_out, + device const float4 *norm_weight, + device float4 *norm_out, + threadgroup float *shared [[threadgroup(0)]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + const uint n4 = args.n_embd >> 2; + threadgroup float4 *row = (threadgroup float4 *)shared; + threadgroup float *sums = shared + args.n_embd; + + if (tid == 0) dsv41_hc_split4(mix, scale, base, split_out, args.sinkhorn_iters, args.hc_eps); + if (sgitg == 0) sums[tiisg] = 0.0f; + + const float4 p = *((device const float4 *)pre); + device const float4 *x0 = residual; + device const float4 *x1 = residual + n4; + device const float4 *x2 = residual + 2u * n4; + device const float4 *x3 = residual + 3u * n4; + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + float4 v = 0.0f; + v += x0[i] * p.x; + v += x1[i] * p.y; + v += x2[i] * p.z; + v += x3[i] * p.w; + v = dsv41_bf16x4(v); + row[i] = v; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) sums[sgitg] = sumf; + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = simd_sum(sums[tiisg]); + + const float mean = sumf / (float)args.n_embd; + const float norm_scale = 1.0f / sqrt(mean + args.norm_eps); + for (uint i = tid; i < n4; i += ntg) { + const float4 v = row[i]; + x_out[i] = v; + norm_out[i] = dsv41_bf16x4((v * norm_scale) * norm_weight[i]); + } +} + +/* out[h] = bf16(post[h] * block + sum_s comb[h, s] * residual[s]) for the + * four streams, kernel_dsv4_hc_expand4's index arithmetic and accumulation + * order; copy_pre also carries split[0..3] into `pre` for the next layer. */ +kernel void kernel_dsv41_hc_expand4_bf16( + constant ds4_metal_args_dsv41_hc &args, + device const float *block_out, + device const float *residual, + device const float *split, + device float *out, + device float *pre_out, + uint d [[thread_position_in_grid]]) { + if (d >= args.n_embd) return; + device const float *post = split + 4; + device const float *comb = split + 8; + + const float block_v = block_out[d]; + const float r0 = residual[d]; + const float r1 = residual[d + args.n_embd]; + const float r2 = residual[d + 2u * args.n_embd]; + const float r3 = residual[d + 3u * args.n_embd]; + + for (uint dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * post[dst_hc]; + acc += comb[dst_hc + 0u * 4u] * r0; + acc += comb[dst_hc + 1u * 4u] * r1; + acc += comb[dst_hc + 2u * 4u] * r2; + acc += comb[dst_hc + 3u * 4u] * r3; + out[d + dst_hc * args.n_embd] = dsv41_bf16(acc); + } + if (args.copy_pre && d < 4) pre_out[d] = split[d]; +} diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index c161a03a8b..5b22130ae9 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -1123,8 +1123,9 @@ struct ds4_metal_args_hc_norm_mix { // slice, preserving every simd_sum tree), and the matvec keeps the original // per-row accumulation order with y = x*scale computed on the fly, which // rounds identically to the materialized normalized row. The host wrapper -// gates this to n == 16384 && out_dim == 24, where the virtual-thread count -// is exactly 1024 and the mv tail loop is empty. +// gates this to n in {16384, 20480} && out_dim == 24, where the virtual-thread +// count is exactly 1024 and the mv tail loop is empty (n is a multiple of the +// NSG*NF*NB = 4096-value stride). kernel void kernel_dsv4_hc_rms_norm_mix_f16( constant ds4_metal_args_hc_norm_mix & args, device const char * x, diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 055a70b677..23a7f36f8f 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -316,6 +316,91 @@ static int check_hc_scaled(void) { return 1; } +/* The fused decode HC glue (norm+mix, split+collapse+BF16+norm+BF16, + * expand+BF16) must be byte-identical to the standalone dispatch sequence + * the V4.1 graph ran before it, at V4.1's shape (4 x 5120, 24 mixes). */ +static int check_hc_fuse(void) { + enum { D = 5120, HC = 4, N = HC * D, OUT = 24, ITERS = 20 }; + const float hc_eps = 1e-6f, rms_eps = 1e-20f; + const size_t fn_bytes = (size_t)N * OUT * sizeof(_Float16); + const size_t scale_off = fn_bytes, base_off = scale_off + 64, normw_off = base_off + 128; + const size_t page = (size_t)getpagesize(); + const size_t mapped = (normw_off + D * sizeof(float) + page - 1) / page * page; + void *model = NULL; + CHECK(!posix_memalign(&model, page, mapped)); + _Float16 *fn = model; + float *scale = (float *)((char *)model + scale_off); + float *base = (float *)((char *)model + base_off); + float *normw = (float *)((char *)model + normw_off); + for (size_t i = 0; i < (size_t)N * OUT; i++) fn[i] = (_Float16)(random_value() / 64); + for (int i = 0; i < 3; i++) scale[i] = 0.5f + random_value() / 8; + for (int i = 0; i < OUT; i++) base[i] = random_value() / 4; + for (int i = 0; i < D; i++) normw[i] = 1.0f + random_value() / 8; + CHECK(ds4_gpu_set_model_map(model, mapped)); + + ds4_gpu_tensor *residual = upload(NULL, N * 4), *block = upload(NULL, D * 4), *pre = upload(NULL, 16); + ds4_gpu_tensor *flat = upload(NULL, N * 4); + ds4_gpu_tensor *mix[2] = {upload(NULL, OUT * 4), upload(NULL, OUT * 4)}; + ds4_gpu_tensor *split[2] = {upload(NULL, OUT * 4), upload(NULL, OUT * 4)}; + ds4_gpu_tensor *x[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *norm[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *out[2] = {upload(NULL, N * 4), upload(NULL, N * 4)}; + ds4_gpu_tensor *pre_next[2] = {upload(NULL, 16), upload(NULL, 16)}; + CHECK(residual && block && pre && flat && mix[0] && mix[1] && split[0] && split[1] && + x[0] && x[1] && norm[0] && norm[1] && out[0] && out[1] && pre_next[0] && pre_next[1]); + float *res = ds4_gpu_tensor_contents(residual), *blk = ds4_gpu_tensor_contents(block); + float *p = ds4_gpu_tensor_contents(pre); + double elapsed[2] = {0}; + for (int round = 0; round < 6; round++) { + for (int i = 0; i < N; i++) res[i] = bf16(random_value() * (round == 5 ? 0x1p-14f : 1.0f)); + for (int i = 0; i < D; i++) blk[i] = bf16(random_value()); + for (int i = 0; i < HC; i++) p[i] = 0.5f + random_value() / 8 + hc_eps; + for (int mode = 0; mode < 2; mode++) { + const double begin = monotonic_seconds(); + CHECK(ds4_gpu_begin_commands()); + if (mode == 0) { + CHECK(ds4_gpu_rms_norm_plain_tensor(flat, residual, N, rms_eps)); + CHECK(ds4_gpu_matmul_f16_tensor(mix[0], model, mapped, 0, N, OUT, flat, 1)); + CHECK(ds4_gpu_hc_split_sinkhorn_tensor(split[0], mix[0], model, mapped, + scale_off, base_off, HC, ITERS, hc_eps)); + CHECK(ds4_gpu_hc_weighted_sum_tensor(x[0], residual, pre, D, HC)); + CHECK(ds4_gpu_dsv41_quantize(x[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_rms_norm_weight_tensor(norm[0], x[0], model, mapped, normw_off, D, rms_eps)); + CHECK(ds4_gpu_dsv41_quantize(norm[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_hc_expand_split_tensor(out[0], block, residual, split[0], D, HC)); + CHECK(ds4_gpu_dsv41_quantize(out[0], N, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_tensor_copy(pre_next[0], 0, split[0], 0, 16)); + } else { + CHECK(ds4_gpu_hc_rms_norm_mix_f16_tensor(mix[1], residual, model, mapped, 0, N, OUT, rms_eps)); + CHECK(ds4_gpu_dsv41_hc_collapse_norm(split[1], x[1], norm[1], mix[1], pre, residual, + model, mapped, scale_off, base_off, normw_off, D, HC, ITERS, hc_eps, rms_eps)); + CHECK(ds4_gpu_dsv41_hc_expand4(out[1], block, residual, split[1], pre_next[1], D)); + } + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + if (round) elapsed[mode] += (monotonic_seconds() - begin) * 1000.0 / 5; + } + CHECK(!memcmp(ds4_gpu_tensor_contents(mix[0]), ds4_gpu_tensor_contents(mix[1]), OUT * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(split[0]), ds4_gpu_tensor_contents(split[1]), OUT * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(x[0]), ds4_gpu_tensor_contents(x[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(norm[0]), ds4_gpu_tensor_contents(norm[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(pre_next[0]), ds4_gpu_tensor_contents(pre_next[1]), 16)); + const float *s = ds4_gpu_tensor_contents(split[1]); + for (int i = 0; i < OUT; i++) CHECK(isfinite(s[i])); + } + fprintf(stderr, "HC fuse: 10 dispatches %.3f ms -> 3 dispatches %.3f ms, outputs byte-identical\n", + elapsed[0], elapsed[1]); + for (int i = 0; i < 2; i++) { + ds4_gpu_tensor_free(mix[i]); ds4_gpu_tensor_free(split[i]); ds4_gpu_tensor_free(x[i]); + ds4_gpu_tensor_free(norm[i]); ds4_gpu_tensor_free(out[i]); ds4_gpu_tensor_free(pre_next[i]); + } + ds4_gpu_tensor_free(residual); ds4_gpu_tensor_free(block); ds4_gpu_tensor_free(pre); ds4_gpu_tensor_free(flat); + ds4_gpu_cleanup(); free(model); + CHECK(ds4_gpu_init()); + return 1; +} + #endif static int check_engram(void) { @@ -1279,6 +1364,11 @@ int main(int argc, char **argv) { ds4_gpu_cleanup(); return ok ? 0 : 1; } + if (argc == 2 && !strcmp(argv[1], "--hc-fuse")) { + const int ok = ds4_gpu_init() && check_hc_fuse(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } #endif if (argc == 2 && !strcmp(argv[1], "--bf16-linear")) { const int ok = ds4_gpu_init() && check_bf16_linear(); @@ -1315,6 +1405,9 @@ int main(int argc, char **argv) { check_candidates() && check_sparse_gather() && check_indexer_batch() && check_embedding() && check_index_projection() && check_general_topk() && check_causal_topk() && check_compact_carry() && check_attention_output(false) && check_tp_attention(); +#ifdef __APPLE__ + ok = ok && check_hc_fuse(); +#endif ds4_gpu_cleanup(); return ok ? 0 : 1; } From c1bd5315c574dfefa8c9dea96b453554c53e3a7f Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 14:15:02 +0200 Subject: [PATCH 2/7] v41: fuse the decode MoE glue (router + select, shared expert, FFN tail) DeepSeek's production decode runs the router as one "Mega-Gate" kernel; ds4's V4.1 graph ran it as the F32 logits matvec plus ten generic select dispatches (softplus, sqrt, bias add, argsort, gather, sum, clamp, div, scale), and the shared expert as three Q8_0 matvecs with three BF16 passes, SwiGLU, another BF16 pass, then the routed + shared sum and its rounding. Three fusions, byte-identical to those sequences: - kernel_dsv41_router_select: kernel_mul_mv_f32_f32_4's matvec (nsg=8, nr0=2), then the last-arriving threadgroup does sqrt(softplus), the bias, the canonical (score desc, idx asc) top-k of the argsort path and the normalized weights with the six-lane sum_rows reduction, the clamp, the row division and the scale as separately rounded ops. 11 -> 1. - kernel_dsv41_shared_gate_up_swiglu_q8_0: both Q8_0 matvecs at the dispatch's nsg, gate/up rounded to BF16 before SwiGLU and the product after it. 6 -> 1. - kernel_dsv41_shared_down_hc_expand4_q8_0: the down matvec, shared = bf16(.), block = bf16(routed + shared), the post/comb expand with its rounding and the pre carry; runs in the layer's FFN tail. 6 -> 1. Single box only. Rollback: DS4_METAL_DISABLE_V41_MOE_FUSE (the HC fusion switch also disables it). tests/test_deepseek41_metal --moe-fuse compares every output of both paths byte for byte, including a zero-input round where all scores tie within bias groups (canonical order check). --- ds4.c | 99 +++++++--- ds4_deepseek41_gpu.h | 21 +++ ds4_metal.m | 149 +++++++++++++++ metal/dsv41.metal | 346 ++++++++++++++++++++++++++++++++++ tests/test_deepseek41_metal.c | 124 +++++++++++- 5 files changed, 715 insertions(+), 24 deletions(-) diff --git a/ds4.c b/ds4.c index 4bccb84b27..859a6217e9 100644 --- a/ds4.c +++ b/ds4.c @@ -39744,6 +39744,37 @@ static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, ds41_bf16(g->block, DS4_N_EMBD); } +#if defined(__APPLE__) +/* DeepSeek's production decode runs each half-layer's hyper-connection work + * in one kernel ("Mega-mHC", tech report 3.2) and the router as "Mega-Gate". + * ds4's V4.1 graph mirrored the reference op by op: 19 HC glue dispatches + * per layer (norm, mix matvec, split, collapse, BF16, weighted norm, BF16, + * expand, BF16, ...) and 22 of MoE glue (router matvec + ten select + * dispatches, shared gate/up/down + their BF16 passes, SwiGLU, the routed + + * shared sum). The fused paths keep every reduction tree and rounding point + * of those sequences (tests/test_deepseek41_metal --hc-fuse / --moe-fuse) + * and issue 6 + 3. Single box only; the TP graph keeps the unfused + * sequences. */ +static bool ds41_hc_fused(const ds41_gpu_graph *g) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_HC_FUSE") == NULL && + ds4_gpu_hc_rms_norm_mix_f16_available() != 0; + } + return enabled && g->tp_world == 1; +} + +static bool ds41_moe_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { + static int enabled = -1; + if (enabled < 0) enabled = getenv("DS4_METAL_DISABLE_V41_MOE_FUSE") == NULL; + return enabled && ds41_hc_fused(g) && + l->ffn_gate_inp->type == DS4_TENSOR_F32 && + l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + l->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + l->ffn_down_shexp->type == DS4_TENSOR_Q8_0; +} +#endif + static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t token) { uint64_t gate_row = 0, down_row = 0; @@ -39754,6 +39785,15 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, ds4_gpu_tensor *routed = shared_owner ? g->block : g->routed; const ds4_tensor *bias = ds41_image_at(g, g->pos) ? l->ffn_exp_probs_vl : l->ffn_exp_probs_b; if (!bias) return false; +#if defined(__APPLE__) + const bool fused = ds41_moe_fused(g, l); + if (fused) { + if (!ds4_gpu_dsv41_router_select(g->selected, g->route_weights, g->route_probs, + g->route_logits, g->norm, m->map, m->size, l->ffn_gate_inp->abs_offset, + bias->abs_offset, true, DS4_N_EMBD, DS4_N_EXPERT, DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE)) return false; + } else +#endif if (!ds41_matmul(g->route_logits, m, l->ffn_gate_inp, g->norm, false) || !ds4_gpu_router_select_tensor(g->selected, g->route_weights, g->route_probs, m->map, m->size, bias->abs_offset, 0, 0, token, @@ -39761,6 +39801,15 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, g->route_logits)) return false; const bool shared_here = !shared_owner || g->tp_rank == (il & 1u); bool shared_queued = false; +#if defined(__APPLE__) + /* Fused: gate/up/SwiGLU here; the down projection, the routed + shared + * sum and the HC expand run as one dispatch in ds41_graph_after_moe. */ + if (fused && shared_here && + !ds4_gpu_dsv41_shared_gate_up_swiglu(g->shared_mid, g->norm, m->map, m->size, + l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP)) return false; + if (fused) shared_queued = true; +#endif #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (shared_owner && shared_here && l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && @@ -39812,18 +39861,21 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, return true; } -static bool ds41_moe_finish(ds41_gpu_graph *g, uint32_t il) { +static bool ds41_moe_finish(ds41_gpu_graph *g, const ds4_layer_weights *l, uint32_t il) { const bool shared_owner = g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_TP_SHARED_OWNER"); ds4_gpu_tensor *routed = shared_owner ? g->block : g->routed; if (!ds41_sum_partial(g, routed, il, DS4_TP_GATE_FFN)) return false; +#if defined(__APPLE__) + if (ds41_moe_fused(g, l)) return true; +#endif return (shared_owner || ds4_gpu_add_tensor(g->block, routed, g->shared, DS4_N_EMBD)) && ds41_bf16(g->block, DS4_N_EMBD); } static bool ds41_moe(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t token) { - return ds41_moe_partial(g, m, l, il, token) && ds41_moe_finish(g, il); + return ds41_moe_partial(g, m, l, il, token) && ds41_moe_finish(g, l, il); } static bool ds41_graph_logits(ds41_gpu_graph *g, const ds4_model *m, @@ -39839,22 +39891,6 @@ static bool ds41_graph_logits(ds41_gpu_graph *g, const ds4_model *m, } #if defined(__APPLE__) -/* DeepSeek's production decode runs each half-layer's hyper-connection work - * in one kernel ("Mega-mHC", tech report 3.2). ds4's V4.1 graph mirrored the - * reference op by op: 19 HC glue dispatches per layer (norm, mix matvec, - * split, collapse, BF16, weighted norm, BF16, expand, BF16, ...). The fused - * path keeps every reduction tree and rounding point of that sequence - * (tests/test_deepseek41_metal --hc-fuse) and issues 6. Single box only; the - * TP graph keeps the unfused sequence. */ -static bool ds41_hc_fused(const ds41_gpu_graph *g) { - static int enabled = -1; - if (enabled < 0) { - enabled = getenv("DS4_METAL_DISABLE_V41_HC_FUSE") == NULL && - ds4_gpu_hc_rms_norm_mix_f16_available() != 0; - } - return enabled && g->tp_world == 1; -} - /* Plain RMSNorm over the flattened HC row and the F16 mix matvec, one dispatch. */ static bool ds41_hc_mix_fused(ds41_gpu_graph *g, const ds4_model *m, const ds4_tensor *fn, const ds4_gpu_tensor *residual) { @@ -40183,7 +40219,8 @@ static bool ds41_attention_batch(ds41_gpu_graph *g, const ds4_model *m, (n_raw - kept + part) * row_bytes, (kept - part) * row_bytes)); } -static bool ds41_graph_after_moe(ds41_gpu_graph *g) { +/* HC expand of an already summed and rounded FFN block (g->block). */ +static bool ds41_hc_expand_after_moe(ds41_gpu_graph *g) { #if defined(__APPLE__) if (ds41_hc_fused(g)) return ds4_gpu_dsv41_hc_expand4(g->residual, g->block, g->after_attn, g->ffn_split, g->pre, DS4_N_EMBD); @@ -40193,10 +40230,26 @@ static bool ds41_graph_after_moe(ds41_gpu_graph *g) { ds4_gpu_tensor_copy(g->pre, 0, g->ffn_split, 0, DS4_N_HC * sizeof(float)); } +/* The layer's FFN tail after ds41_moe: with the fused MoE glue this is where + * the shared down projection, the routed + shared sum and their roundings + * run (one dispatch with the expand); otherwise ds41_moe_finish did them. */ +static bool ds41_graph_after_moe(ds41_gpu_graph *g, const ds4_model *m, + const ds4_layer_weights *l) { +#if defined(__APPLE__) + if (ds41_moe_fused(g, l)) + return ds4_gpu_dsv41_shared_down_hc_expand4(g->residual, g->shared, g->block, + g->shared_mid, g->routed, g->after_attn, g->ffn_split, g->pre, m->map, m->size, + l->ffn_down_shexp->abs_offset, DS4_N_EMBD, DS4_N_FF_EXP) != 0; +#else + (void)m; (void)l; +#endif + return ds41_hc_expand_after_moe(g); +} + static bool ds41_graph_layer(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, int token) { return ds41_graph_before_moe(g, m, l, il) && ds41_moe(g, m, l, il, (uint32_t)token) && - ds41_graph_after_moe(g); + ds41_graph_after_moe(g, m, l); } #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) @@ -40237,7 +40290,7 @@ static bool ds41_graph_decode_layer(ds41_gpu_graph *g, const ds4_model *m, ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && ds41_bf16(g->block, DS4_N_EMBD) && ds41_decode_island(g, m, l, il, 1) && - ds41_moe_finish(g, il) && ds41_graph_after_moe(g); + ds41_moe_finish(g, l, il) && ds41_graph_after_moe(g, m, l); } #endif @@ -40951,7 +41004,7 @@ static bool ds41_graph_prefill_sweep(ds41_gpu_graph *g, const ds4_model *m, #undef DS41_USE_FFN_ROW ok = ds41_graph_after_attention(&row, m, l); if (ok && !batch_moe) ok = ds41_moe(&row, m, l, il, (uint32_t)tokens[off + t]) && - ds41_graph_after_moe(&row); + ds41_graph_after_moe(&row, m, l); } } if (ok && batch_moe) { @@ -40969,7 +41022,7 @@ static bool ds41_graph_prefill_sweep(ds41_gpu_graph *g, const ds4_model *m, DS41_PREFILL_ROWS(DS41_USE_MOE_ROW) #undef DS41_USE_MOE_ROW ok = ds4_gpu_add_tensor(row.block, row.routed, row.shared, DS4_N_EMBD) && - ds41_bf16(row.block, DS4_N_EMBD) && ds41_graph_after_moe(&row); + ds41_bf16(row.block, DS4_N_EMBD) && ds41_hc_expand_after_moe(&row); } } DS41_STAGE("hc expand"); diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index 7c794f2d4e..221fdd80b4 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -141,6 +141,27 @@ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4 int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, ds4_gpu_tensor *pre, uint32_t n_embd); +/* Decode MoE glue for one token row, byte-identical to the standalone + * sequences: router = F32 logits matvec + softplus/sqrt + bias + canonical + * top-k + normalized weights in one dispatch; shared gate/up = two Q8_0 + * matvecs + BF16 + SwiGLU + BF16; shared down + HC tail = Q8_0 matvec + + * BF16 + (routed + shared) + BF16 + post/comb expand + BF16 (+ pre carry). */ +int ds4_gpu_dsv41_router_select(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, ds4_gpu_tensor *logits, + const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint64_t bias_offset, bool has_bias, + uint32_t n_embd, uint32_t n_expert, uint32_t n_used, float scale); +int ds4_gpu_dsv41_shared_gate_up_swiglu(ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, + uint32_t n_embd, uint32_t n_ff, float clamp); +int ds4_gpu_dsv41_shared_down_hc_expand4(ds4_gpu_tensor *out, ds4_gpu_tensor *shared, + ds4_gpu_tensor *block, const ds4_gpu_tensor *mid, + const ds4_gpu_tensor *routed, const ds4_gpu_tensor *residual, + const ds4_gpu_tensor *split, ds4_gpu_tensor *pre, + const void *model_map, uint64_t model_size, + uint64_t down_offset, uint32_t n_embd, uint32_t n_ff); #ifdef __cplusplus } diff --git a/ds4_metal.m b/ds4_metal.m index 8795bbaa13..d1627a35fa 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -47795,6 +47795,155 @@ int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, } } +/* V4.1 decode MoE glue (metal/dsv41.metal): one token row. */ +int ds4_gpu_dsv41_router_select(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, ds4_gpu_tensor *logits, + const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint64_t bias_offset, bool has_bias, + uint32_t n_embd, uint32_t n_expert, uint32_t n_used, float scale) { + const uint64_t weight_bytes = (uint64_t)n_embd * n_expert * sizeof(float); + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (!n_embd || n_embd % 32u || !n_expert || n_expert % 2u || n_expert > 512u || + !n_used || n_used > 8u || n_used > n_expert || !model_map || + !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(logits, n_expert) || + !dsv41_tensor_has_floats(probs, n_expert) || !dsv41_tensor_has_floats(weights, n_used) || + !selected || ds4_gpu_tensor_bytes(selected) < (uint64_t)n_used * sizeof(int32_t) || + weight_offset > model_size || weight_bytes > model_size - weight_offset || + (has_bias && (bias_offset > model_size || bias_bytes > model_size - bias_offset))) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_router_select"); + if (!pipeline || pipeline.maxTotalThreadsPerThreadgroup < 256u || !g_dsv4_completion_cache) return 0; + uint64_t weight_inner = 0, bias_inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &weight_inner); + id biasbuf = has_bias ? + ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner) : wbuf; + if (!wbuf || !biasbuf) return 0; + id logitsbuf = ds4_gpu_tensor_buffer(logits); + NSString *completion_key = [NSString stringWithFormat:@"v41router:%p:%llu", + (void *)logitsbuf, (unsigned long long)ds4_gpu_tensor_offset(logits)]; + id completion = [g_dsv4_completion_cache objectForKey:completion_key]; + if (!completion) { + completion = [g_device newBufferWithLength:sizeof(uint32_t) options:MTLResourceStorageModeShared]; + if (!completion) return 0; + *((uint32_t *)[completion contents]) = 0u; + [g_dsv4_completion_cache setObject:completion forKey:completion_key]; + } + [g_transient_buffers addObject:completion]; + const struct { uint32_t n_embd, n_expert, n_used, has_bias; float scale; } args = + {n_embd, n_expert, n_used, has_bias, scale}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(probs) offset:ds4_gpu_tensor_offset(probs) atIndex:4]; + [enc setBuffer:biasbuf offset:(NSUInteger)(has_bias ? bias_inner : weight_inner) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(selected) offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(weights) offset:ds4_gpu_tensor_offset(weights) atIndex:7]; + [enc setBuffer:completion offset:0 atIndex:8]; + [enc setThreadgroupMemoryLength:(512u + 32u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((n_expert + 1u) / 2u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 router select"); + } +} + +int ds4_gpu_dsv41_shared_gate_up_swiglu(ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, + uint32_t n_embd, uint32_t n_ff, float clamp) { + const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; + const uint64_t weight_bytes = row_bytes * n_ff; + if (!n_embd || n_embd % 32u || !n_ff || !model_map || !isfinite(clamp) || clamp < 0.0f || + !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(mid, n_ff) || + gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + uint64_t gate_inner = 0, up_inner = 0; + id gate_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); + id up_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(n_embd, n_ff); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv41_shared_gate_up_swiglu_q8_0", mv_dispatch.nsg); + if (!pipeline) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(mid) offset:ds4_gpu_tensor_offset(mid) atIndex:4]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:5]; + [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_ff + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 shared gate/up"); + } +} + +int ds4_gpu_dsv41_shared_down_hc_expand4(ds4_gpu_tensor *out, ds4_gpu_tensor *shared, + ds4_gpu_tensor *block, const ds4_gpu_tensor *mid, + const ds4_gpu_tensor *routed, const ds4_gpu_tensor *residual, + const ds4_gpu_tensor *split, ds4_gpu_tensor *pre, + const void *model_map, uint64_t model_size, + uint64_t down_offset, uint32_t n_embd, uint32_t n_ff) { + const uint64_t row_bytes = ((uint64_t)n_ff / 32u) * 34u; + const uint64_t weight_bytes = row_bytes * n_embd; + if (!n_embd || !n_ff || n_ff % 32u || !model_map || + !dsv41_tensor_has_floats(out, 4ull * n_embd) || !dsv41_tensor_has_floats(shared, n_embd) || + !dsv41_tensor_has_floats(block, n_embd) || !dsv41_tensor_has_floats(mid, n_ff) || + !dsv41_tensor_has_floats(routed, n_embd) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || + !dsv41_tensor_has_floats(split, 24) || (pre && !dsv41_tensor_has_floats(pre, 4)) || + down_offset > model_size || weight_bytes > model_size - down_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, weight_bytes, &inner); + if (!wbuf) return 0; + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(n_ff, n_embd); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + mv_args.nr0 = mv_dispatch.nr0; + const ds4_gpu_dsv41_hc_args hc_args = {n_embd, 0, 0.0f, 0.0f, pre != NULL}; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv41_shared_down_hc_expand4_q8_0", mv_dispatch.nsg); + if (!pipeline) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:2]; + const ds4_gpu_tensor *buffers[] = {mid, shared, routed, block, residual, split, out, pre ? pre : split}; + for (NSUInteger i = 0; i < 8; i++) + [enc setBuffer:ds4_gpu_tensor_buffer(buffers[i]) + offset:ds4_gpu_tensor_offset(buffers[i]) atIndex:i + 3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_embd + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 shared down + HC expand"); + } +} + int ds4_gpu_dsv41_carry_copy(ds4_gpu_tensor *packed, uint32_t row_offset, ds4_gpu_tensor *plain, uint32_t width, uint32_t rows, uint32_t format, bool pack) { diff --git a/metal/dsv41.metal b/metal/dsv41.metal index 8a54118a78..830ca8a0cc 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -465,3 +465,349 @@ kernel void kernel_dsv41_hc_expand4_bf16( } if (args.copy_pre && d < 4) pre_out[d] = split[d]; } + +/* Decode-time V4.1 MoE glue, one token row. Same discipline as the HC + * kernels above: every reduction keeps the standalone kernel's shape and + * every BF16 rounding point is kept (tests/test_deepseek41_metal --moe-fuse). */ + +struct ds4_metal_args_dsv41_router { + uint n_embd; + uint n_expert; + uint n_used; + uint has_bias; + float scale; +}; + +/* Router logits (F32 matvec, kernel_mul_mv_f32_f32_4 with nsg=8, nr0=2) and + * the generic select sequence the V4.1 graph ran as ten dispatches: + * probs = sqrt(softplus(logits)); score = probs + bias; top-k in the + * canonical (score desc, idx asc) order of kernel_argsort_f32_i32_desc_canon; + * weights = (probs[sel] / max(sum, 2^-14)) * scale, where sum is + * kernel_sum_rows_f32_f32's six-lane simd reduction. The last-arriving + * threadgroup selects (DeepSeek's Mega-Gate shape). */ +kernel void kernel_dsv41_router_select( + constant ds4_metal_args_dsv41_router &args, + device const float *weight, + device const float *x, + device float *logits, + device float *probs, + device const float *bias, + device int32_t *selected, + device float *weights, + device atomic_uint *completion, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr short NSG = 8; + constexpr short NR0 = 2; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + constexpr short NW = N_SIMDWIDTH; + const int nb = (int)args.n_embd/NB; + const int r0 = tgpig.x*NR0; + device const float4 *y4 = (device const float4 *)x; + device const float4 *ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const float4 *)(weight + (uint64_t)(r0 + row)*args.n_embd); + } + float sumf[NR0] = {0.f}; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = sgitg*NF + ix; + device const float4 *yb4 = y4 + (ib0*NB + il*NF)/4; + for (int ib = ib0; ib < nb; ib += NSG*NF) { + float4 yl4[NF4]; + FOR_UNROLL (short i = 0; i < NF4; ++i) yl4[i] = yb4[i]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + device const float4 *xb4 = ax4[row] + (ib*NB + il*NF)/4; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) sumq += dot(float4(xb4[i]), float4(yl4[i])); + sumf[row] += sumq; + } + yb4 += NSG*NF*NW/4; + } + for (int i = nb*NB + sgitg*NW + tiisg; i < (int)args.n_embd; i += NW*NSG) { + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sumf[row] += ((device const float *)ax4[row])[i] * x[i]; + } + } + helper_mv_reduce_and_write(logits, sumf, r0, (int)args.n_expert, + tiisg, sgitg, (threadgroup char *)scratch); + + threadgroup_barrier(mem_flags::mem_threadgroup); + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + if (tid == 0) { + const uint old = atomic_fetch_add_explicit(completion, 1u, memory_order_relaxed); + scratch[0] = old + 1u == (args.n_expert + 1u) / 2u ? 1.0f : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (scratch[0] == 0.0f) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + + constexpr uint SLOTS = 512; + threadgroup float *score_tg = scratch; /* SLOTS */ + threadgroup float *red_s = scratch + SLOTS; /* 8 */ + threadgroup int32_t *red_i = (threadgroup int32_t *)(scratch + SLOTS + 8); /* 8 */ + threadgroup int32_t *sel_tg = (threadgroup int32_t *)(scratch + SLOTS + 16); /* 8 */ + threadgroup volatile float *stage = scratch + SLOTS + 24; /* 8 */ + + device volatile const float *lg = (device volatile const float *)logits; + for (uint i = tid; i < SLOTS; i += 256u) { + float s = -INFINITY; + if (i < args.n_expert) { + const float v = lg[i]; + const float sp = select(log(1.0f + exp(v)), v, v > 20.0f); + const float p = sqrt(sp); + probs[i] = p; + s = args.has_bias ? p + bias[i] : p; + } + score_tg[i] = s; + } + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + + float s0 = score_tg[tid], s1 = score_tg[tid + 256u]; + for (uint r = 0; r < args.n_used; r++) { + float bs = s0; + int32_t bi = (int32_t)tid; + if (s1 > bs) { bs = s1; bi = (int32_t)tid + 256; } + for (ushort o = 16; o > 0; o >>= 1) { + const float ps = simd_shuffle_xor(bs, o); + const int32_t pi = simd_shuffle_xor(bi, o); + if (ps > bs || (ps == bs && pi < bi)) { bs = ps; bi = pi; } + } + if (tiisg == 0) { red_s[sgitg] = bs; red_i[sgitg] = bi; } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float best = red_s[0]; + int32_t best_i = red_i[0]; + for (uint g = 1; g < 8; g++) { + if (red_s[g] > best || (red_s[g] == best && red_i[g] < best_i)) { + best = red_s[g]; best_i = red_i[g]; + } + } + sel_tg[r] = best_i; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int32_t sel = sel_tg[r]; + if ((int32_t)tid == sel) s0 = -INFINITY; + if ((int32_t)tid + 256 == sel) s1 = -INFINITY; + } + + /* kernel_sum_rows_f32_f32 runs six threads: lanes 0..5 hold one weight + * each and the simdgroup reduces them; then (w / clamp(sum)) * scale as + * two separately rounded ops (the div-row and mul-scalar kernels). */ + const bool lane = sgitg == 0 && tiisg < args.n_used; + device volatile const float *pr = (device volatile const float *)probs; + const float w = lane ? pr[sel_tg[tiisg]] : 0.0f; + float sum = simd_sum(w); + if (sgitg == 0 && tiisg == 0) { + sum = clamp(sum, 6.103515625e-5f, INFINITY); + stage[0] = sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < args.n_used) { + stage[1 + tid] = w / stage[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < args.n_used) { + selected[tid] = sel_tg[tid]; + weights[tid] = stage[1 + tid] * args.scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + if (tid == 0) atomic_store_explicit(completion, 0u, memory_order_relaxed); +} + +/* Shared expert gate/up (Q8_0, kernel_mul_mv_q8_0_f32's walk and reduction + * tree at the dispatch's nsg) with V4.1's rounding: gate and up are rounded + * to BF16 before SwiGLU and the product after it. */ +kernel void kernel_dsv41_shared_gate_up_swiglu_q8_0( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short NR0 = N_R0_Q8_0; + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + + const int nb = args.ne00 / QK8_0; + const int r0 = tgpig.x * NR0; + device const float *y = (device const float *)src1; + + device const block_q8_0 *ag[NR0]; + device const block_q8_0 *au[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row) * args.nb01; + ag[row] = (device const block_q8_0 *)(src0_gate + offset0); + au[row] = (device const block_q8_0 *)(src0_up + offset0); + } + + float sumg[NR0] = { 0.f }; + float sumu[NR0] = { 0.f }; + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + float yl[NQ]; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL (short i = 0; i < NQ; ++i) yl[i] = yb[i]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + device const int8_t *qg = ag[row][ib].qs + il * NQ; + device const int8_t *qu = au[row][ib].qs + il * NQ; + float sg = 0.f; + float su = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sg += qg[i] * yl[i]; + su += qu[i] * yl[i]; + } + sumg[row] += sg * ag[row][ib].d; + sumu[row] += su * au[row][ib].d; + } + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shmem_f32 = (threadgroup float *)shmem; + threadgroup float *sh_gate[NR0]; + threadgroup float *sh_up[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sh_gate[row] = shmem_f32 + NW * row; + sh_up[row] = shmem_f32 + NW * (NR0 + row); + if (sgitg == 0) { + sh_gate[row][tiisg] = 0.0f; + sh_up[row][tiisg] = 0.0f; + } + sumg[row] = simd_sum(sumg[row]); + sumu[row] = simd_sum(sumu[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + sh_gate[row][sgitg] = sumg[row]; + sh_up[row][sgitg] = sumu[row]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *mid_f32 = (device float *)dst_mid; + FOR_UNROLL (short row = 0; row < NR0 && r0 + row < args.ne01; ++row) { + const float gate = simd_sum(sh_gate[row][tiisg]); + const float up = simd_sum(sh_up[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + float g = dsv41_bf16(gate); + float u = dsv41_bf16(up); + if (clamp_value > 1.0e-6f) { + g = min(g, clamp_value); + u = clamp(u, -clamp_value, clamp_value); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[r0 + row] = dsv41_bf16(silu * u); + } + } +} + +/* Shared expert down projection (Q8_0 matvec, standalone walk and tree), + * then the layer's FFN tail exactly as the graph ran it in six dispatches: + * shared = bf16(down); block = bf16(routed + shared); residual_out = + * bf16(post/comb expand of block into residual); pre = split[0..3]. */ +kernel void kernel_dsv41_shared_down_hc_expand4_q8_0( + constant ds4_metal_args_mul_mv & mv, + constant ds4_metal_args_dsv41_hc & hc, + device const char * weight, + device const char * shared_mid, + device float * shared_out, + device const float * routed_out, + device float * block_out, + device const float * residual, + device const float * split, + device float * dst, + device float * pre_out, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + constexpr short NR0 = N_R0_Q8_0; + + const int nb = mv.ne00 / QK8_0; + const int row0 = tgpig.x * NR0; + + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + + device const float *y = (device const float *)(shared_mid); + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + device const block_q8_0 *ax[NR0]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + ax[row] = (device const block_q8_0 *)(weight + (uint64_t)(row0 + row) * mv.nb01); + } + + float sumf[NR0] = { 0.0f }; + float yl[NQ]; + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL(short i = 0; i < NQ; ++i) yl[i] = yb[i]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + device const int8_t *qs = ax[row][ib].qs + il * NQ; + float sumq = 0.0f; + FOR_UNROLL(short i = 0; i < NQ; ++i) sumq += qs[i] * yl[i]; + sumf[row] += sumq * ax[row][ib].d; + } + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shmem_f32[NR0]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + shmem_f32[row] = (threadgroup float *)shmem + NW * row; + if (sgitg == 0) shmem_f32[row][tiisg] = 0.0f; + sumf[row] = simd_sum(sumf[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + FOR_UNROLL(short row = 0; row < NR0; ++row) { + if (tiisg == 0) shmem_f32[row][sgitg] = sumf[row]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device const float *post = split + 4; + device const float *comb = split + 8; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + const int d = row0 + row; + if (d >= mv.ne01) continue; + const float shared_v = simd_sum(shmem_f32[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + const float sv = dsv41_bf16(shared_v); + shared_out[d] = sv; + float block_v = routed_out[d]; + block_v += sv; + block_v = dsv41_bf16(block_v); + block_out[d] = block_v; + + const float r0 = residual[d]; + const float r1 = residual[d + hc.n_embd]; + const float r2 = residual[d + 2u * hc.n_embd]; + const float r3 = residual[d + 3u * hc.n_embd]; + for (uint dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * post[dst_hc]; + acc += comb[dst_hc + 0u * 4u] * r0; + acc += comb[dst_hc + 1u * 4u] * r1; + acc += comb[dst_hc + 2u * 4u] * r2; + acc += comb[dst_hc + 3u * 4u] * r3; + dst[d + dst_hc * hc.n_embd] = dsv41_bf16(acc); + } + } + } + if (hc.copy_pre && tgpig.x == 0 && sgitg == 0 && tiisg < 4) pre_out[tiisg] = split[tiisg]; +} diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 23a7f36f8f..2d180e89cf 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -401,6 +401,123 @@ static int check_hc_fuse(void) { return 1; } +static void fill_q8_0(uint8_t *dst, size_t rows, size_t k) { + for (size_t r = 0; r < rows; r++) for (size_t b = 0; b < k / 32; b++) { + uint8_t *block = dst + (r * (k / 32) + b) * 34; + const _Float16 d = (_Float16)(random_value() / 512 + 0.01f); + memcpy(block, &d, 2); + for (int i = 0; i < 32; i++) block[2 + i] = (uint8_t)(int8_t)(random_value() * 30); + } +} + +/* The fused decode MoE glue (router + select, shared gate/up/SwiGLU, shared + * down + sum + HC expand) must be byte-identical to the 22-dispatch + * standalone sequence at V4.1's shape (5120 x 384 F32 router, 6 of 384, + * scale 1.5, Q8_0 shared expert of 2304). */ +static int check_moe_fuse(void) { + enum { D = 5120, HC = 4, N = HC * D, E = 384, K = 6, FF = 2304 }; + const float scale = 1.5f, clamp = 10.0f; + const size_t router_bytes = (size_t)D * E * 4, bias_bytes = E * 4; + const size_t gate_bytes = (size_t)FF * (D / 32) * 34, down_bytes = (size_t)D * (FF / 32) * 34; + const size_t router_off = 0, bias_off = router_bytes, gate_off = bias_off + bias_bytes; + const size_t up_off = gate_off + gate_bytes, down_off = up_off + gate_bytes; + const size_t page = (size_t)getpagesize(); + const size_t mapped = (down_off + down_bytes + page - 1) / page * page; + void *model = NULL; + CHECK(!posix_memalign(&model, page, mapped)); + float *router = model, *bias = (float *)((char *)model + bias_off); + for (size_t i = 0; i < (size_t)D * E; i++) router[i] = random_value() / 64; + /* Repeated bias values: with a zero input every score ties within its + * group and the canonical order (ascending index among equals) is what + * the standalone argsort produces. */ + for (int i = 0; i < E; i++) bias[i] = (float)(i % 5) / 10; + fill_q8_0((uint8_t *)model + gate_off, FF, D); + fill_q8_0((uint8_t *)model + up_off, FF, D); + fill_q8_0((uint8_t *)model + down_off, D, FF); + CHECK(ds4_gpu_set_model_map(model, mapped)); + + ds4_gpu_tensor *norm = upload(NULL, D * 4), *routed = upload(NULL, D * 4); + ds4_gpu_tensor *residual = upload(NULL, N * 4), *split = upload(NULL, 24 * 4); + ds4_gpu_tensor *logits[2] = {upload(NULL, E * 4), upload(NULL, E * 4)}; + ds4_gpu_tensor *probs[2] = {upload(NULL, E * 4), upload(NULL, E * 4)}; + ds4_gpu_tensor *selected[2] = {upload(NULL, K * 4), upload(NULL, K * 4)}; + ds4_gpu_tensor *weights[2] = {upload(NULL, K * 4), upload(NULL, K * 4)}; + ds4_gpu_tensor *gate = upload(NULL, FF * 4), *up = upload(NULL, FF * 4); + ds4_gpu_tensor *mid[2] = {upload(NULL, FF * 4), upload(NULL, FF * 4)}; + ds4_gpu_tensor *shared[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *block[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *out[2] = {upload(NULL, N * 4), upload(NULL, N * 4)}; + ds4_gpu_tensor *pre[2] = {upload(NULL, 16), upload(NULL, 16)}; + CHECK(norm && routed && residual && split && gate && up); + for (int i = 0; i < 2; i++) + CHECK(logits[i] && probs[i] && selected[i] && weights[i] && mid[i] && shared[i] && block[i] && out[i] && pre[i]); + float *x = ds4_gpu_tensor_contents(norm), *r = ds4_gpu_tensor_contents(routed); + float *res = ds4_gpu_tensor_contents(residual), *s = ds4_gpu_tensor_contents(split); + double elapsed[2] = {0}; + for (int round = 0; round < 8; round++) { + const int ties = round == 3 || round == 7; + for (int i = 0; i < D; i++) x[i] = ties ? 0.0f : bf16(random_value()); + for (int i = 0; i < D; i++) r[i] = bf16(random_value()); + for (int i = 0; i < N; i++) res[i] = bf16(random_value()); + for (int i = 0; i < 24; i++) s[i] = i < 4 ? 0.5f + random_value() / 8 : i < 8 ? 1.0f + random_value() / 4 : 0.25f + random_value() / 16; + for (int mode = 0; mode < 2; mode++) { + const double begin = monotonic_seconds(); + CHECK(ds4_gpu_begin_commands()); + if (mode == 0) { + CHECK(ds4_gpu_matmul_f32_tensor(logits[0], model, mapped, router_off, D, E, norm, 1)); + CHECK(ds4_gpu_router_select_tensor(selected[0], weights[0], probs[0], model, mapped, + bias_off, 0, 0, 0, E, K, scale, 0, 0, true, false, logits[0])); + CHECK(ds4_gpu_matmul_q8_0_tensor(gate, model, mapped, gate_off, D, FF, norm, 1)); + CHECK(ds4_gpu_dsv41_quantize(gate, FF, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_matmul_q8_0_tensor(up, model, mapped, up_off, D, FF, norm, 1)); + CHECK(ds4_gpu_dsv41_quantize(up, FF, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_swiglu_tensor(mid[0], gate, up, FF, clamp, 1.0f)); + CHECK(ds4_gpu_dsv41_quantize(mid[0], FF, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_matmul_q8_0_tensor(shared[0], model, mapped, down_off, FF, D, mid[0], 1)); + CHECK(ds4_gpu_dsv41_quantize(shared[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_add_tensor(block[0], routed, shared[0], D)); + CHECK(ds4_gpu_dsv41_quantize(block[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_hc_expand_split_tensor(out[0], block[0], residual, split, D, HC)); + CHECK(ds4_gpu_dsv41_quantize(out[0], N, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_tensor_copy(pre[0], 0, split, 0, 16)); + } else { + CHECK(ds4_gpu_dsv41_router_select(selected[1], weights[1], probs[1], logits[1], norm, + model, mapped, router_off, bias_off, true, D, E, K, scale)); + CHECK(ds4_gpu_dsv41_shared_gate_up_swiglu(mid[1], norm, model, mapped, gate_off, up_off, D, FF, clamp)); + CHECK(ds4_gpu_dsv41_shared_down_hc_expand4(out[1], shared[1], block[1], mid[1], routed, + residual, split, pre[1], model, mapped, down_off, D, FF)); + } + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + if (round) elapsed[mode] += (monotonic_seconds() - begin) * 1000.0 / 7; + } + const int32_t *sel = ds4_gpu_tensor_contents(selected[1]); + for (int i = 0; i < K; i++) CHECK(sel[i] >= 0 && sel[i] < E); + if (ties) for (int i = 0; i < K; i++) CHECK(sel[i] == 4 + 5 * i); + CHECK(!memcmp(ds4_gpu_tensor_contents(logits[0]), ds4_gpu_tensor_contents(logits[1]), E * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(probs[0]), ds4_gpu_tensor_contents(probs[1]), E * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(selected[0]), sel, K * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(weights[0]), ds4_gpu_tensor_contents(weights[1]), K * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(mid[0]), ds4_gpu_tensor_contents(mid[1]), FF * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(shared[0]), ds4_gpu_tensor_contents(shared[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(pre[0]), ds4_gpu_tensor_contents(pre[1]), 16)); + } + fprintf(stderr, "MoE fuse: 15 dispatches %.3f ms -> 3 dispatches %.3f ms, outputs byte-identical\n", + elapsed[0], elapsed[1]); + for (int i = 0; i < 2; i++) { + ds4_gpu_tensor_free(logits[i]); ds4_gpu_tensor_free(probs[i]); ds4_gpu_tensor_free(selected[i]); + ds4_gpu_tensor_free(weights[i]); ds4_gpu_tensor_free(mid[i]); ds4_gpu_tensor_free(shared[i]); + ds4_gpu_tensor_free(block[i]); ds4_gpu_tensor_free(out[i]); ds4_gpu_tensor_free(pre[i]); + } + ds4_gpu_tensor_free(norm); ds4_gpu_tensor_free(routed); ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(split); ds4_gpu_tensor_free(gate); ds4_gpu_tensor_free(up); + ds4_gpu_cleanup(); free(model); + CHECK(ds4_gpu_init()); + return 1; +} + #endif static int check_engram(void) { @@ -1369,6 +1486,11 @@ int main(int argc, char **argv) { ds4_gpu_cleanup(); return ok ? 0 : 1; } + if (argc == 2 && !strcmp(argv[1], "--moe-fuse")) { + const int ok = ds4_gpu_init() && check_moe_fuse(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } #endif if (argc == 2 && !strcmp(argv[1], "--bf16-linear")) { const int ok = ds4_gpu_init() && check_bf16_linear(); @@ -1406,7 +1528,7 @@ int main(int argc, char **argv) { check_embedding() && check_index_projection() && check_general_topk() && check_causal_topk() && check_compact_carry() && check_attention_output(false) && check_tp_attention(); #ifdef __APPLE__ - ok = ok && check_hc_fuse(); + ok = ok && check_hc_fuse() && check_moe_fuse(); #endif ds4_gpu_cleanup(); return ok ? 0 : 1; From d6a34d62b23c48a9bc29f8094835229bc5ef8575 Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 14:19:34 +0200 Subject: [PATCH 3/7] v41: single-dispatch router select on M5 only; matvec + one select group elsewhere The last-arriving-threadgroup select read stale logits on the M3 Ultra (nondeterministic probs, kernel test caught it on the box); V4 gates the same pattern to M5. Pre-M5 runs the standalone F32 matvec dispatch and one selecting threadgroup: 11 -> 2 instead of 11 -> 1, byte-identical. --- ds4_metal.m | 15 +++++-- metal/dsv41.metal | 101 +++++++++++++++++++++++++--------------------- 2 files changed, 67 insertions(+), 49 deletions(-) diff --git a/ds4_metal.m b/ds4_metal.m index d1627a35fa..9f8df3c068 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -47831,8 +47831,17 @@ int ds4_gpu_dsv41_router_select(ds4_gpu_tensor *selected, ds4_gpu_tensor *weight [g_dsv4_completion_cache setObject:completion forKey:completion_key]; } [g_transient_buffers addObject:completion]; - const struct { uint32_t n_embd, n_expert, n_used, has_bias; float scale; } args = - {n_embd, n_expert, n_used, has_bias, scale}; + /* One dispatch (last-arriving threadgroup selects) on M5 only, like + * V4's fused router: on the M3 Ultra the other groups' logits were + * not reliably visible to the last one. Elsewhere the standalone + * F32 matvec runs first and one threadgroup selects. */ + const bool fused_matvec = ds4_gpu_device_is_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_V41_ROUTER_SINGLE_DISPATCH") == NULL; + if (!fused_matvec && + !ds4_gpu_matmul_f32_tensor(logits, model_map, model_size, weight_offset, + n_embd, n_expert, x, 1)) return 0; + const struct { uint32_t n_embd, n_expert, n_used, has_bias; float scale; uint32_t fused_matvec; } args = + {n_embd, n_expert, n_used, has_bias, scale, fused_matvec}; int owned = 0; id cb = ds4_gpu_command_buffer(&owned); id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; @@ -47848,7 +47857,7 @@ int ds4_gpu_dsv41_router_select(ds4_gpu_tensor *selected, ds4_gpu_tensor *weight [enc setBuffer:ds4_gpu_tensor_buffer(weights) offset:ds4_gpu_tensor_offset(weights) atIndex:7]; [enc setBuffer:completion offset:0 atIndex:8]; [enc setThreadgroupMemoryLength:(512u + 32u) * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((n_expert + 1u) / 2u, 1, 1) + [enc dispatchThreadgroups:MTLSizeMake(fused_matvec ? (n_expert + 1u) / 2u : 1u, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 router select"); diff --git a/metal/dsv41.metal b/metal/dsv41.metal index 830ca8a0cc..f84341d72d 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -476,6 +476,7 @@ struct ds4_metal_args_dsv41_router { uint n_used; uint has_bias; float scale; + uint fused_matvec; }; /* Router logits (F32 matvec, kernel_mul_mv_f32_f32_4 with nsg=8, nr0=2) and @@ -483,8 +484,12 @@ struct ds4_metal_args_dsv41_router { * probs = sqrt(softplus(logits)); score = probs + bias; top-k in the * canonical (score desc, idx asc) order of kernel_argsort_f32_i32_desc_canon; * weights = (probs[sel] / max(sum, 2^-14)) * scale, where sum is - * kernel_sum_rows_f32_f32's six-lane simd reduction. The last-arriving - * threadgroup selects (DeepSeek's Mega-Gate shape). */ + * kernel_sum_rows_f32_f32's six-lane simd reduction. + * fused_matvec: the last-arriving threadgroup selects (DeepSeek's Mega-Gate + * shape); the host enables it on M5 only, like V4's fused router: on the + * M3 Ultra the other groups' logits were not reliably visible to the last + * one (nondeterministic probs). Otherwise one threadgroup selects from + * logits produced by the standalone matvec dispatch. */ kernel void kernel_dsv41_router_select( constant ds4_metal_args_dsv41_router &args, device const float *weight, @@ -500,52 +505,54 @@ kernel void kernel_dsv41_router_select( ushort tid [[thread_index_in_threadgroup]], ushort tiisg [[thread_index_in_simdgroup]], ushort sgitg [[simdgroup_index_in_threadgroup]]) { - constexpr short NSG = 8; - constexpr short NR0 = 2; - constexpr short NB = 32; - constexpr short NF = 16; - constexpr short NF4 = NF/4; - constexpr short NW = N_SIMDWIDTH; - const int nb = (int)args.n_embd/NB; - const int r0 = tgpig.x*NR0; - device const float4 *y4 = (device const float4 *)x; - device const float4 *ax4[NR0]; - FOR_UNROLL (short row = 0; row < NR0; ++row) { - ax4[row] = (device const float4 *)(weight + (uint64_t)(r0 + row)*args.n_embd); - } - float sumf[NR0] = {0.f}; - const short ix = tiisg/(NW/NF); - const short il = tiisg%(NW/NF); - const int ib0 = sgitg*NF + ix; - device const float4 *yb4 = y4 + (ib0*NB + il*NF)/4; - for (int ib = ib0; ib < nb; ib += NSG*NF) { - float4 yl4[NF4]; - FOR_UNROLL (short i = 0; i < NF4; ++i) yl4[i] = yb4[i]; + if (args.fused_matvec) { + constexpr short NSG = 8; + constexpr short NR0 = 2; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + constexpr short NW = N_SIMDWIDTH; + const int nb = (int)args.n_embd/NB; + const int r0 = tgpig.x*NR0; + device const float4 *y4 = (device const float4 *)x; + device const float4 *ax4[NR0]; FOR_UNROLL (short row = 0; row < NR0; ++row) { - device const float4 *xb4 = ax4[row] + (ib*NB + il*NF)/4; - float sumq = 0.f; - FOR_UNROLL (short i = 0; i < NF4; ++i) sumq += dot(float4(xb4[i]), float4(yl4[i])); - sumf[row] += sumq; + ax4[row] = (device const float4 *)(weight + (uint64_t)(r0 + row)*args.n_embd); } - yb4 += NSG*NF*NW/4; - } - for (int i = nb*NB + sgitg*NW + tiisg; i < (int)args.n_embd; i += NW*NSG) { - FOR_UNROLL (short row = 0; row < NR0; ++row) { - sumf[row] += ((device const float *)ax4[row])[i] * x[i]; + float sumf[NR0] = {0.f}; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = sgitg*NF + ix; + device const float4 *yb4 = y4 + (ib0*NB + il*NF)/4; + for (int ib = ib0; ib < nb; ib += NSG*NF) { + float4 yl4[NF4]; + FOR_UNROLL (short i = 0; i < NF4; ++i) yl4[i] = yb4[i]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + device const float4 *xb4 = ax4[row] + (ib*NB + il*NF)/4; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) sumq += dot(float4(xb4[i]), float4(yl4[i])); + sumf[row] += sumq; + } + yb4 += NSG*NF*NW/4; } - } - helper_mv_reduce_and_write(logits, sumf, r0, (int)args.n_expert, - tiisg, sgitg, (threadgroup char *)scratch); + for (int i = nb*NB + sgitg*NW + tiisg; i < (int)args.n_embd; i += NW*NSG) { + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sumf[row] += ((device const float *)ax4[row])[i] * x[i]; + } + } + helper_mv_reduce_and_write(logits, sumf, r0, (int)args.n_expert, + tiisg, sgitg, (threadgroup char *)scratch); - threadgroup_barrier(mem_flags::mem_threadgroup); - atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); - if (tid == 0) { - const uint old = atomic_fetch_add_explicit(completion, 1u, memory_order_relaxed); - scratch[0] = old + 1u == (args.n_expert + 1u) / 2u ? 1.0f : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + if (tid == 0) { + const uint old = atomic_fetch_add_explicit(completion, 1u, memory_order_relaxed); + scratch[0] = old + 1u == (args.n_expert + 1u) / 2u ? 1.0f : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (scratch[0] == 0.0f) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (scratch[0] == 0.0f) return; - atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); constexpr uint SLOTS = 512; threadgroup float *score_tg = scratch; /* SLOTS */ @@ -616,9 +623,11 @@ kernel void kernel_dsv41_router_select( selected[tid] = sel_tg[tid]; weights[tid] = stage[1 + tid] * args.scale; } - threadgroup_barrier(mem_flags::mem_threadgroup); - atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); - if (tid == 0) atomic_store_explicit(completion, 0u, memory_order_relaxed); + if (args.fused_matvec) { + threadgroup_barrier(mem_flags::mem_threadgroup); + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + if (tid == 0) atomic_store_explicit(completion, 0u, memory_order_relaxed); + } } /* Shared expert gate/up (Q8_0, kernel_mul_mv_q8_0_f32's walk and reduction From 5b8b6ac31dd9e71ba3f5d8ef7fa79caaf49a13cf Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 14:26:25 +0200 Subject: [PATCH 4/7] v41: separate rollback switches for the router and shared-expert fusions --- ds4.c | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/ds4.c b/ds4.c index 859a6217e9..b5a298cfde 100644 --- a/ds4.c +++ b/ds4.c @@ -39764,11 +39764,25 @@ static bool ds41_hc_fused(const ds41_gpu_graph *g) { return enabled && g->tp_world == 1; } +/* Router matvec + select as one (two on pre-M5) dispatch. */ +static bool ds41_router_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_MOE_FUSE") == NULL && + getenv("DS4_METAL_DISABLE_V41_ROUTER_FUSE") == NULL; + } + return enabled && ds41_hc_fused(g) && l->ffn_gate_inp->type == DS4_TENSOR_F32; +} + +/* Shared expert gate/up/SwiGLU as one dispatch and its down projection + * folded into the FFN tail (ds41_moe_finish / ds41_graph_after_moe agree). */ static bool ds41_moe_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { static int enabled = -1; - if (enabled < 0) enabled = getenv("DS4_METAL_DISABLE_V41_MOE_FUSE") == NULL; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_MOE_FUSE") == NULL && + getenv("DS4_METAL_DISABLE_V41_SHARED_FUSE") == NULL; + } return enabled && ds41_hc_fused(g) && - l->ffn_gate_inp->type == DS4_TENSOR_F32 && l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && l->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && l->ffn_down_shexp->type == DS4_TENSOR_Q8_0; @@ -39787,7 +39801,7 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, if (!bias) return false; #if defined(__APPLE__) const bool fused = ds41_moe_fused(g, l); - if (fused) { + if (ds41_router_fused(g, l)) { if (!ds4_gpu_dsv41_router_select(g->selected, g->route_weights, g->route_probs, g->route_logits, g->norm, m->map, m->size, l->ffn_gate_inp->abs_offset, bias->abs_offset, true, DS4_N_EMBD, DS4_N_EXPERT, DS4_N_EXPERT_USED, From bc628554b68c47537dcdeca441121ac48631c9e7 Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 14:32:07 +0200 Subject: [PATCH 5/7] v41: expand4 can fold the routed + shared sum; early-down diagnostic variant --- ds4.c | 22 ++++++++++++++++++---- ds4_deepseek41_gpu.h | 6 ++++-- ds4_metal.m | 20 ++++++++++++-------- metal/dsv41.metal | 16 +++++++++++++--- tests/test_deepseek41_metal.c | 14 +++++++++++++- 5 files changed, 60 insertions(+), 18 deletions(-) diff --git a/ds4.c b/ds4.c index b5a298cfde..2f056fb911 100644 --- a/ds4.c +++ b/ds4.c @@ -39774,6 +39774,14 @@ static bool ds41_router_fused(const ds41_gpu_graph *g, const ds4_layer_weights * return enabled && ds41_hc_fused(g) && l->ffn_gate_inp->type == DS4_TENSOR_F32; } +/* Diagnostic: run the shared down projection (standalone matvec + BF16) + * before the routed experts instead of inside the fused FFN tail. */ +static bool ds41_shared_down_early(void) { + static int enabled = -1; + if (enabled < 0) enabled = getenv("DS4_METAL_V41_SHARED_DOWN_EARLY") != NULL; + return enabled; +} + /* Shared expert gate/up/SwiGLU as one dispatch and its down projection * folded into the FFN tail (ds41_moe_finish / ds41_graph_after_moe agree). */ static bool ds41_moe_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { @@ -39821,7 +39829,9 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, if (fused && shared_here && !ds4_gpu_dsv41_shared_gate_up_swiglu(g->shared_mid, g->norm, m->map, m->size, l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset, - DS4_N_EMBD, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP)) return false; + DS4_N_EMBD, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP) || + (ds41_shared_down_early() && + !ds41_matmul(g->shared, m, l->ffn_down_shexp, g->shared_mid, true)))) return false; if (fused) shared_queued = true; #endif #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) @@ -39946,7 +39956,7 @@ static bool ds41_graph_after_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { #if defined(__APPLE__) if (ds41_hc_fused(g)) - return ds4_gpu_dsv41_hc_expand4(g->after_attn, g->block, g->residual, g->attn_split, NULL, DS4_N_EMBD) && + return ds4_gpu_dsv41_hc_expand4(g->after_attn, g->block, g->residual, g->attn_split, NULL, NULL, NULL, DS4_N_EMBD) && ds41_hc_mix_fused(g, m, l->hc_ffn_fn, g->after_attn) && ds41_hc_collapse_norm_fused(g, m, l, true); #endif @@ -40237,7 +40247,7 @@ static bool ds41_attention_batch(ds41_gpu_graph *g, const ds4_model *m, static bool ds41_hc_expand_after_moe(ds41_gpu_graph *g) { #if defined(__APPLE__) if (ds41_hc_fused(g)) - return ds4_gpu_dsv41_hc_expand4(g->residual, g->block, g->after_attn, g->ffn_split, g->pre, DS4_N_EMBD); + return ds4_gpu_dsv41_hc_expand4(g->residual, g->block, g->after_attn, g->ffn_split, g->pre, NULL, NULL, DS4_N_EMBD); #endif return ds4_gpu_hc_expand_split_tensor(g->residual, g->block, g->after_attn, g->ffn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->residual, DS4_N_EMBD * DS4_N_HC) && @@ -40250,10 +40260,14 @@ static bool ds41_hc_expand_after_moe(ds41_gpu_graph *g) { static bool ds41_graph_after_moe(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { #if defined(__APPLE__) - if (ds41_moe_fused(g, l)) + if (ds41_moe_fused(g, l)) { + if (ds41_shared_down_early()) + return ds4_gpu_dsv41_hc_expand4(g->residual, g->routed, g->after_attn, g->ffn_split, + g->pre, g->shared, g->block, DS4_N_EMBD) != 0; return ds4_gpu_dsv41_shared_down_hc_expand4(g->residual, g->shared, g->block, g->shared_mid, g->routed, g->after_attn, g->ffn_split, g->pre, m->map, m->size, l->ffn_down_shexp->abs_offset, DS4_N_EMBD, DS4_N_FF_EXP) != 0; + } #else (void)m; (void)l; #endif diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index 221fdd80b4..d3ed66dcdf 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -129,7 +129,8 @@ int ds4_gpu_dsv41_gather_kv(ds4_gpu_tensor *out, const ds4_gpu_tensor *source, * residual); norm = bf16(rmsnorm(x) * weight). `pre` is the coefficient row * of the PREVIOUS sublayer's split (the first four floats of that tensor). * expand4: out = bf16(post/comb expand of block into residual); when `pre` - * is given, split[0..3] is also copied into it. */ + * is given, split[0..3] is also copied into it; when `add` is given, block + * is first bf16(block + add) and that sum is written to block_sum. */ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4_gpu_tensor *norm, const ds4_gpu_tensor *mix, const ds4_gpu_tensor *pre, const ds4_gpu_tensor *residual, @@ -140,7 +141,8 @@ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4 float hc_eps, float norm_eps); int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, - ds4_gpu_tensor *pre, uint32_t n_embd); + ds4_gpu_tensor *pre, const ds4_gpu_tensor *add, + ds4_gpu_tensor *block_sum, uint32_t n_embd); /* Decode MoE glue for one token row, byte-identical to the standalone * sequences: router = F32 logits matvec + softplus/sqrt + bias + canonical * top-k + normalized weights in one dispatch; shared gate/up = two Q8_0 diff --git a/ds4_metal.m b/ds4_metal.m index 9f8df3c068..bce928e100 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -47711,7 +47711,7 @@ int ds4_gpu_dsv41_engram_add(ds4_gpu_tensor *residual, } /* V4.1 decode HC glue (metal/dsv41.metal): one token row, HC=4. */ -typedef struct { uint32_t n_embd, sinkhorn_iters; float hc_eps, norm_eps; uint32_t copy_pre; } ds4_gpu_dsv41_hc_args; +typedef struct { uint32_t n_embd, sinkhorn_iters; float hc_eps, norm_eps; uint32_t copy_pre, has_add; } ds4_gpu_dsv41_hc_args; int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4_gpu_tensor *norm, const ds4_gpu_tensor *mix, const ds4_gpu_tensor *pre, @@ -47744,7 +47744,7 @@ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4 id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, weight_bytes, &norm_inner); if (!scalebuf || !basebuf || !normwbuf) return 0; - const ds4_gpu_dsv41_hc_args args = {n_embd, sinkhorn_iters, hc_eps, norm_eps, 0}; + const ds4_gpu_dsv41_hc_args args = {n_embd, sinkhorn_iters, hc_eps, norm_eps, 0, 0}; int owned = 0; id cb = ds4_gpu_command_buffer(&owned); id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; @@ -47769,23 +47769,27 @@ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4 int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, - ds4_gpu_tensor *pre, uint32_t n_embd) { + ds4_gpu_tensor *pre, const ds4_gpu_tensor *add, + ds4_gpu_tensor *block_sum, uint32_t n_embd) { if (!n_embd || !dsv41_tensor_has_floats(out, 4ull * n_embd) || !dsv41_tensor_has_floats(block, n_embd) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || - !dsv41_tensor_has_floats(split, 24) || (pre && !dsv41_tensor_has_floats(pre, 4))) return 0; + !dsv41_tensor_has_floats(split, 24) || (pre && !dsv41_tensor_has_floats(pre, 4)) || + (add != NULL) != (block_sum != NULL) || (add && !dsv41_tensor_has_floats(add, n_embd)) || + (block_sum && !dsv41_tensor_has_floats(block_sum, n_embd))) return 0; if (!g_initialized && !ds4_gpu_init()) return 0; @autoreleasepool { id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_expand4_bf16"); if (!pipeline) return 0; - const ds4_gpu_dsv41_hc_args args = {n_embd, 0, 0.0f, 0.0f, pre != NULL}; + const ds4_gpu_dsv41_hc_args args = {n_embd, 0, 0.0f, 0.0f, pre != NULL, add != NULL}; int owned = 0; id cb = ds4_gpu_command_buffer(&owned); id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; if (!enc) return 0; [enc setComputePipelineState:pipeline]; [enc setBytes:&args length:sizeof(args) atIndex:0]; - const ds4_gpu_tensor *buffers[] = {block, residual, split, out, pre ? pre : split}; - for (NSUInteger i = 0; i < 5; i++) + const ds4_gpu_tensor *buffers[] = {block, residual, split, out, pre ? pre : split, + add ? add : block, block_sum ? block_sum : out}; + for (NSUInteger i = 0; i < 7; i++) [enc setBuffer:ds4_gpu_tensor_buffer(buffers[i]) offset:ds4_gpu_tensor_offset(buffers[i]) atIndex:i + 1]; [enc dispatchThreadgroups:MTLSizeMake(((uint64_t)n_embd + 255u) / 256u, 1, 1) @@ -47928,7 +47932,7 @@ int ds4_gpu_dsv41_shared_down_hc_expand4(ds4_gpu_tensor *out, ds4_gpu_tensor *sh ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(n_ff, n_embd); ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); mv_args.nr0 = mv_dispatch.nr0; - const ds4_gpu_dsv41_hc_args hc_args = {n_embd, 0, 0.0f, 0.0f, pre != NULL}; + const ds4_gpu_dsv41_hc_args hc_args = {n_embd, 0, 0.0f, 0.0f, pre != NULL, 0}; id pipeline = ds4_gpu_get_mul_mv_pipeline("kernel_dsv41_shared_down_hc_expand4_q8_0", mv_dispatch.nsg); if (!pipeline) return 0; diff --git a/metal/dsv41.metal b/metal/dsv41.metal index f84341d72d..7c021bcd31 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -309,6 +309,7 @@ struct ds4_metal_args_dsv41_hc { float hc_eps; float norm_eps; uint copy_pre; + uint has_add; }; static inline float4 dsv41_bf16x4(float4 v) { @@ -436,20 +437,29 @@ kernel void kernel_dsv41_hc_collapse_norm4( /* out[h] = bf16(post[h] * block + sum_s comb[h, s] * residual[s]) for the * four streams, kernel_dsv4_hc_expand4's index arithmetic and accumulation - * order; copy_pre also carries split[0..3] into `pre` for the next layer. */ + * order; copy_pre also carries split[0..3] into `pre` for the next layer. + * has_add: block = bf16(block_in + add) first (the routed + shared sum and + * its rounding), written back to block_sum. */ kernel void kernel_dsv41_hc_expand4_bf16( constant ds4_metal_args_dsv41_hc &args, - device const float *block_out, + device const float *block_in, device const float *residual, device const float *split, device float *out, device float *pre_out, + device const float *add, + device float *block_sum, uint d [[thread_position_in_grid]]) { if (d >= args.n_embd) return; device const float *post = split + 4; device const float *comb = split + 8; - const float block_v = block_out[d]; + float block_v = block_in[d]; + if (args.has_add) { + block_v += add[d]; + block_v = dsv41_bf16(block_v); + block_sum[d] = block_v; + } const float r0 = residual[d]; const float r1 = residual[d + args.n_embd]; const float r2 = residual[d + 2u * args.n_embd]; diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 2d180e89cf..560ef006b1 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -374,7 +374,7 @@ static int check_hc_fuse(void) { CHECK(ds4_gpu_hc_rms_norm_mix_f16_tensor(mix[1], residual, model, mapped, 0, N, OUT, rms_eps)); CHECK(ds4_gpu_dsv41_hc_collapse_norm(split[1], x[1], norm[1], mix[1], pre, residual, model, mapped, scale_off, base_off, normw_off, D, HC, ITERS, hc_eps, rms_eps)); - CHECK(ds4_gpu_dsv41_hc_expand4(out[1], block, residual, split[1], pre_next[1], D)); + CHECK(ds4_gpu_dsv41_hc_expand4(out[1], block, residual, split[1], pre_next[1], NULL, NULL, D)); } CHECK(ds4_gpu_end_commands()); CHECK(ds4_gpu_synchronize()); @@ -503,6 +503,18 @@ static int check_moe_fuse(void) { CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(pre[0]), ds4_gpu_tensor_contents(pre[1]), 16)); + /* The early-down variant: standalone down + BF16, then the sum and + * its rounding folded into the expand. */ + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_matmul_q8_0_tensor(shared[1], model, mapped, down_off, FF, D, mid[1], 1)); + CHECK(ds4_gpu_dsv41_quantize(shared[1], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_hc_expand4(out[1], routed, residual, split, pre[1], shared[1], block[1], D)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(ds4_gpu_tensor_contents(shared[0]), ds4_gpu_tensor_contents(shared[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(pre[0]), ds4_gpu_tensor_contents(pre[1]), 16)); } fprintf(stderr, "MoE fuse: 15 dispatches %.3f ms -> 3 dispatches %.3f ms, outputs byte-identical\n", elapsed[0], elapsed[1]); From f61a83d7f84e99cc7022f85bd40f784b94cd8bd8 Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 14:43:46 +0200 Subject: [PATCH 6/7] v41: fuse the decode attention glue; shared down before the routed experts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Attention (single box), byte-identical to the standalone sequences (tests/test_deepseek41_metal --attn-fuse): - Q8_0 and F16 single-row matvecs round to BF16 on the store (kernel_dsv41_mul_mv_q8_0_f32_bf16, kernel_dsv41_mul_mv_f16_f32_4_bf16: the standalone walks and reduction trees, nsg from the dispatch, the nr0=4 F16 shape included); ds41_matmul(…, round) uses them, removing the rounding dispatch after q_a, q_b, kv, output_b, the shared down, the compressor and indexer projections. - kernel_dsv41_qkv_norm_kv_tail: the q LoRA and KV weighted norms with their roundings, then the KV RoPE, FP8 (E8M0) block quantization and the raw-window store in one dispatch (6 -> 1). The threadgroup is the q row's norm thread count; the KV lanes beyond its own count hold no elements, so their zero partials leave the reduction tree unchanged. - kernel_dsv41_bf16_rope: the rounding pass over all heads and their inverse RoPE (2 -> 1). - The logits' collapse + output norm reuse kernel_dsv41_hc_collapse_norm4 without a split (4 -> 1). MoE: the shared expert's down projection now runs before the routed experts (rounded on its store) and the routed + shared sum with its rounding is folded into the FFN tail's expand; a fused tail that read shared_mid after the routed kernel saw clobbered values on the M3 Ultra (greedy output differed) while the same kernel was exact in isolation. That late kernel is gone. Rollbacks: DS4_METAL_DISABLE_V41_ATTN_FUSE, DS4_METAL_DISABLE_V41_MATVEC_BF16. --- ds4.c | 194 +++++++++++-------- ds4_deepseek41_gpu.h | 24 ++- ds4_metal.m | 210 ++++++++++++++------- metal/dsv41.metal | 341 ++++++++++++++++++++++++---------- tests/test_deepseek41_metal.c | 150 +++++++++++++-- 5 files changed, 665 insertions(+), 254 deletions(-) diff --git a/ds4.c b/ds4.c index 2f056fb911..1ab223102b 100644 --- a/ds4.c +++ b/ds4.c @@ -39437,6 +39437,19 @@ static bool ds41_bf16(ds4_gpu_tensor *x, uint32_t width) { static bool ds41_matmul(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, bool round) { +#if defined(__APPLE__) + /* Round on the store instead of a rounding dispatch over the result. */ + static int fused = -1; + if (fused < 0) { + fused = getenv("DS4_METAL_DISABLE_V41_MATVEC_BF16") == NULL && + getenv("DS4_METAL_V41_SKIP_BF16") == NULL; + } + if (round && fused) { + const int rc = ds4_gpu_dsv41_matvec_bf16(out, m->map, m->size, weight->abs_offset, + weight->type, (uint32_t)weight->dim[0], (uint32_t)weight->dim[1], in); + if (rc) return rc > 0; + } +#endif return metal_graph_matmul_plain_tensor(out, m, weight, weight->dim[0], weight->dim[1], in, 1) && (!round || ds41_bf16(out, (uint32_t)weight->dim[1])); } @@ -39444,6 +39457,7 @@ static bool ds41_matmul(ds4_gpu_tensor *out, const ds4_model *m, static bool ds41_matmul_batch(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t count, bool round) { + if (count == 1) return ds41_matmul(out, m, weight, in, round); const uint32_t width = (uint32_t)weight->dim[0], outputs = (uint32_t)weight->dim[1]; bool ok; /* Small decode batches retain scalar reductions before BF16 and sparse @@ -39611,15 +39625,16 @@ static bool ds41_attention_low(ds41_gpu_graph *g, const ds4_model *m, 4096, 1024, groups, g->heads) && ds41_bf16(g->low, groups * DS4_N_LORA_O); } +/* round: the single-box block leaves rounded to BF16 (fused attention glue). */ static bool ds41_attention_output(ds41_gpu_graph *g, const ds4_model *m, - const ds4_layer_weights *l) { + const ds4_layer_weights *l, bool round) { const uint32_t groups = DS4_N_OUT_GROUP / g->tp_world; if (!ds41_attention_low(g, m, l)) return false; return g->tp_world == 2 ? metal_graph_matmul_dense_quant_kslice(g->block, m, l->attn_output_b, 8192, (uint64_t)g->tp_rank * groups * 1024u, (uint64_t)groups * 1024u, DS4_N_EMBD, g->low, 0) : - ds41_matmul(g->block, m, l->attn_output_b, g->low, false); + ds41_matmul(g->block, m, l->attn_output_b, g->low, round); } static bool ds41_attention_publish(ds41_gpu_graph *g, const ds4_model *m, @@ -39702,48 +39717,6 @@ static bool ds41_attention_select(ds41_gpu_graph *g, const ds4_model *m, ds41_attention_select_published(g, m, l, il); } -static bool ds41_attention_project(ds41_gpu_graph *g, const ds4_model *m, - const ds4_layer_weights *l) { - const uint32_t q_dim = DS4_N_HEAD / g->tp_world * DS4_N_HEAD_DIM; - return ds41_matmul(g->qr, m, l->attn_q_a, g->norm, true) && - ds41_norm(g->qr, g->qr, m, l->attn_q_a_norm) && - ds41_matmul_rows(g->q, m, l->attn_q_b, g->qr, - g->tp_rank * q_dim, q_dim) && - ds41_matmul(g->kv, m, l->attn_kv, g->norm, true) && - ds41_norm(g->kv, g->kv, m, l->attn_kv_a_norm); -} - -static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, - const ds4_layer_weights *l, uint32_t il, bool projected) { - const uint32_t pos = g->pos, ratio = ds4_layer_compress_ratio(il); - const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; - const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; - const uint32_t heads = DS4_N_HEAD / g->tp_world; - const uint32_t head0 = g->tp_rank * heads; - if (!projected && !ds41_attention_project(g, m, l)) return false; - if (!ds41_rope(g->q, heads, DS4_N_HEAD_DIM, il, pos, false) || - !ds41_rope(g->kv, 1, DS4_N_HEAD_DIM, il, pos, false) || - !ds4_gpu_dsv41_quantize(g->kv, DS4_N_HEAD_DIM, 1, DS4_V41_FP8_E8M0) || - !ds4_gpu_tensor_copy(g->window[il], (uint64_t)(pos % 128u) * 512u * 4u, - g->kv, 0, 512u * 4u) || - !ds41_attention_select(g, m, l, il)) return false; - const uint32_t attended = n_comp < DS4_N_INDEXER_TOP_K ? n_comp : DS4_N_INDEXER_TOP_K; - if (n_comp && !ds4_gpu_dsv41_gather_kv(g->selected_kv, g->compressed[owner], - g->selected_comp, n_comp, attended)) return false; - const uint32_t n_raw = pos + 1u < 128u ? pos + 1u : 128u; - if (!ds4_gpu_attention_decode_heads_tensor(g->heads, m->map, m->size, - l->attn_sinks->abs_offset + (uint64_t)head0 * sizeof(float), - g->q, g->window[il], n_raw, 128, (pos + 1u - n_raw) % 128u, - g->selected_kv, 0, attended, NULL, 0, - heads, DS4_N_HEAD_DIM) || - !ds41_bf16(g->heads, heads * DS4_N_HEAD_DIM) || - !ds41_rope(g->heads, heads, DS4_N_HEAD_DIM, il, pos, true)) return false; - if (projected) return true; - return ds41_attention_output(g, m, l) && - ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && - ds41_bf16(g->block, DS4_N_EMBD); -} - #if defined(__APPLE__) /* DeepSeek's production decode runs each half-layer's hyper-connection work * in one kernel ("Mega-mHC", tech report 3.2) and the router as "Mega-Gate". @@ -39753,7 +39726,7 @@ static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, * dispatches, shared gate/up/down + their BF16 passes, SwiGLU, the routed + * shared sum). The fused paths keep every reduction tree and rounding point * of those sequences (tests/test_deepseek41_metal --hc-fuse / --moe-fuse) - * and issue 6 + 3. Single box only; the TP graph keeps the unfused + * and issue 6 + 5. Single box only; the TP graph keeps the unfused * sequences. */ static bool ds41_hc_fused(const ds41_gpu_graph *g) { static int enabled = -1; @@ -39774,16 +39747,11 @@ static bool ds41_router_fused(const ds41_gpu_graph *g, const ds4_layer_weights * return enabled && ds41_hc_fused(g) && l->ffn_gate_inp->type == DS4_TENSOR_F32; } -/* Diagnostic: run the shared down projection (standalone matvec + BF16) - * before the routed experts instead of inside the fused FFN tail. */ -static bool ds41_shared_down_early(void) { - static int enabled = -1; - if (enabled < 0) enabled = getenv("DS4_METAL_V41_SHARED_DOWN_EARLY") != NULL; - return enabled; -} - -/* Shared expert gate/up/SwiGLU as one dispatch and its down projection - * folded into the FFN tail (ds41_moe_finish / ds41_graph_after_moe agree). */ +/* Shared expert gate/up/SwiGLU as one dispatch, its down projection rounded + * on the store, and the routed + shared sum folded into the FFN tail's + * expand (ds41_moe_finish / ds41_graph_after_moe agree). The down + * projection runs BEFORE the routed experts: a fused tail that read + * shared_mid after them saw clobbered values on the M3 Ultra. */ static bool ds41_moe_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { static int enabled = -1; if (enabled < 0) { @@ -39795,8 +39763,87 @@ static bool ds41_moe_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) l->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && l->ffn_down_shexp->type == DS4_TENSOR_Q8_0; } + +/* Attention glue of the single-box decode graph as one dispatch each: the + * q/kv norms with the KV RoPE, FP8 quantization and raw-window store; the + * heads' rounding with their inverse RoPE; the output projection rounded on + * its store. Byte-identical to the standalone sequences + * (tests/test_deepseek41_metal --attn-fuse). */ +static bool ds41_attn_fused(const ds41_gpu_graph *g) { + static int enabled = -1; + if (enabled < 0) enabled = getenv("DS4_METAL_DISABLE_V41_ATTN_FUSE") == NULL; + return enabled && ds41_hc_fused(g); +} #endif +/* With the fused attention glue the KV row leaves here roped, quantized and + * stored in the raw window; ds41_attention skips those steps. */ +static bool ds41_attention_project(ds41_gpu_graph *g, const ds4_model *m, + const ds4_layer_weights *l, uint32_t il) { + const uint32_t q_dim = DS4_N_HEAD / g->tp_world * DS4_N_HEAD_DIM; +#if defined(__APPLE__) + if (ds41_attn_fused(g)) + return ds41_matmul(g->qr, m, l->attn_q_a, g->norm, true) && + ds41_matmul(g->kv, m, l->attn_kv, g->norm, true) && + ds4_gpu_dsv41_qkv_norm_kv_tail(g->qr, g->kv, g->window[il], + (uint64_t)(g->pos % 128u) * 512u * 4u, m->map, m->size, + l->attn_q_a_norm->abs_offset, l->attn_kv_a_norm->abs_offset, + DS4_N_LORA_Q, DS4_N_HEAD_DIM, DS4_RMS_EPS, g->pos, + ds4_layer_compress_ratio(il) != 0) && + ds41_matmul_rows(g->q, m, l->attn_q_b, g->qr, 0, q_dim); +#else + (void)il; +#endif + return ds41_matmul(g->qr, m, l->attn_q_a, g->norm, true) && + ds41_norm(g->qr, g->qr, m, l->attn_q_a_norm) && + ds41_matmul_rows(g->q, m, l->attn_q_b, g->qr, + g->tp_rank * q_dim, q_dim) && + ds41_matmul(g->kv, m, l->attn_kv, g->norm, true) && + ds41_norm(g->kv, g->kv, m, l->attn_kv_a_norm); +} + +static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, + const ds4_layer_weights *l, uint32_t il, bool projected) { + const uint32_t pos = g->pos, ratio = ds4_layer_compress_ratio(il); + const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; + const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; + const uint32_t heads = DS4_N_HEAD / g->tp_world; + const uint32_t head0 = g->tp_rank * heads; + if (!projected && !ds41_attention_project(g, m, l, il)) return false; + /* Projected rows (the prefill's per-row path) had their projections done + * elsewhere: their KV tail runs here, unfused. */ +#if defined(__APPLE__) + const bool fused = !projected && ds41_attn_fused(g); +#else + const bool fused = false; +#endif + if (!ds41_rope(g->q, heads, DS4_N_HEAD_DIM, il, pos, false) || + (!fused && + (!ds41_rope(g->kv, 1, DS4_N_HEAD_DIM, il, pos, false) || + !ds4_gpu_dsv41_quantize(g->kv, DS4_N_HEAD_DIM, 1, DS4_V41_FP8_E8M0) || + !ds4_gpu_tensor_copy(g->window[il], (uint64_t)(pos % 128u) * 512u * 4u, + g->kv, 0, 512u * 4u))) || + !ds41_attention_select(g, m, l, il)) return false; + const uint32_t attended = n_comp < DS4_N_INDEXER_TOP_K ? n_comp : DS4_N_INDEXER_TOP_K; + if (n_comp && !ds4_gpu_dsv41_gather_kv(g->selected_kv, g->compressed[owner], + g->selected_comp, n_comp, attended)) return false; + const uint32_t n_raw = pos + 1u < 128u ? pos + 1u : 128u; + if (!ds4_gpu_attention_decode_heads_tensor(g->heads, m->map, m->size, + l->attn_sinks->abs_offset + (uint64_t)head0 * sizeof(float), + g->q, g->window[il], n_raw, 128, (pos + 1u - n_raw) % 128u, + g->selected_kv, 0, attended, NULL, 0, + heads, DS4_N_HEAD_DIM) || + (fused ? + !ds4_gpu_dsv41_bf16_rope(g->heads, DS4_N_HEAD_DIM, heads, 1, pos, + ds4_layer_compress_ratio(il) != 0, true) : + (!ds41_bf16(g->heads, heads * DS4_N_HEAD_DIM) || + !ds41_rope(g->heads, heads, DS4_N_HEAD_DIM, il, pos, true)))) return false; + if (projected) return true; + return ds41_attention_output(g, m, l, fused) && + ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && + (fused || ds41_bf16(g->block, DS4_N_EMBD)); +} + static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t token) { uint64_t gate_row = 0, down_row = 0; @@ -39827,11 +39874,10 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, /* Fused: gate/up/SwiGLU here; the down projection, the routed + shared * sum and the HC expand run as one dispatch in ds41_graph_after_moe. */ if (fused && shared_here && - !ds4_gpu_dsv41_shared_gate_up_swiglu(g->shared_mid, g->norm, m->map, m->size, + (!ds4_gpu_dsv41_shared_gate_up_swiglu(g->shared_mid, g->norm, m->map, m->size, l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset, DS4_N_EMBD, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP) || - (ds41_shared_down_early() && - !ds41_matmul(g->shared, m, l->ffn_down_shexp, g->shared_mid, true)))) return false; + !ds41_matmul(g->shared, m, l->ffn_down_shexp, g->shared_mid, true))) return false; if (fused) shared_queued = true; #endif #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) @@ -39905,9 +39951,15 @@ static bool ds41_moe(ds41_gpu_graph *g, const ds4_model *m, static bool ds41_graph_logits(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, float *logits) { if (!g->valid || !logits || !ds4_gpu_begin_commands()) return false; - bool ok = ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && - ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, w->output_norm) && - ds41_output_projection(g, g->tp_logits_half ? g->tp_logits_half : g->logits, + bool ok = +#if defined(__APPLE__) + ds41_hc_fused(g) ? + ds4_gpu_dsv41_hc_collapse_norm(NULL, g->x, g->norm, NULL, g->pre, g->residual, m->map, m->size, + 0, 0, w->output_norm->abs_offset, DS4_N_EMBD, DS4_N_HC, 0, 0.0f, DS4_RMS_EPS) != 0 : +#endif + ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && + ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, w->output_norm); + ok = ok && ds41_output_projection(g, g->tp_logits_half ? g->tp_logits_half : g->logits, m, w, g->norm, 1); if (!ds4_gpu_end_commands()) ok = false; return ok && ds4_gpu_tensor_read(g->logits, 0, logits, @@ -40255,19 +40307,15 @@ static bool ds41_hc_expand_after_moe(ds41_gpu_graph *g) { } /* The layer's FFN tail after ds41_moe: with the fused MoE glue this is where - * the shared down projection, the routed + shared sum and their roundings - * run (one dispatch with the expand); otherwise ds41_moe_finish did them. */ + * the routed + shared sum and its rounding run, inside the expand dispatch; + * otherwise ds41_moe_finish did them. */ static bool ds41_graph_after_moe(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { #if defined(__APPLE__) - if (ds41_moe_fused(g, l)) { - if (ds41_shared_down_early()) - return ds4_gpu_dsv41_hc_expand4(g->residual, g->routed, g->after_attn, g->ffn_split, - g->pre, g->shared, g->block, DS4_N_EMBD) != 0; - return ds4_gpu_dsv41_shared_down_hc_expand4(g->residual, g->shared, g->block, - g->shared_mid, g->routed, g->after_attn, g->ffn_split, g->pre, m->map, m->size, - l->ffn_down_shexp->abs_offset, DS4_N_EMBD, DS4_N_FF_EXP) != 0; - } + if (ds41_moe_fused(g, l)) + return ds4_gpu_dsv41_hc_expand4(g->residual, g->routed, g->after_attn, g->ffn_split, + g->pre, g->shared, g->block, DS4_N_EMBD) != 0; + (void)m; #else (void)m; (void)l; #endif @@ -40295,7 +40343,7 @@ static bool ds41_decode_island(ds41_gpu_graph *g, const ds4_model *m, const int state = ds4_gpu_decode_graph_begin(&key); if (state == 1) return true; const bool ok = island == 0 ? - ds41_graph_before_attention(g, m, l, il) && ds41_attention_project(g, m, l) : + ds41_graph_before_attention(g, m, l, il) && ds41_attention_project(g, m, l, il) : island == 2 ? ds41_attention_output(g, m, l) : ds41_graph_after_attention(g, m, l) && ds41_moe_partial(g, m, l, il, 0); if (state != 0) return ok; @@ -41229,7 +41277,7 @@ static DS4_MAYBE_UNUSED bool ds41_graph_step_batch(ds41_gpu_graph *const *graphs row.q = queries[i]; row.heads = heads[i]; ok = ds41_attention(&row, model, l, il, true) && - ds41_attention_output(&row, model, l); + ds41_attention_output(&row, model, l, false); } if (ok) ok = ds41_sum_partial_batch(g, active.block, il, rows); if (ok) ok = ds4_gpu_dsv41_quantize(active.block, DS4_N_EMBD, rows, DS4_V41_BF16) && diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index d3ed66dcdf..2a7a5b0d87 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -158,12 +158,24 @@ int ds4_gpu_dsv41_shared_gate_up_swiglu(ds4_gpu_tensor *mid, const ds4_gpu_tenso const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint32_t n_embd, uint32_t n_ff, float clamp); -int ds4_gpu_dsv41_shared_down_hc_expand4(ds4_gpu_tensor *out, ds4_gpu_tensor *shared, - ds4_gpu_tensor *block, const ds4_gpu_tensor *mid, - const ds4_gpu_tensor *routed, const ds4_gpu_tensor *residual, - const ds4_gpu_tensor *split, ds4_gpu_tensor *pre, - const void *model_map, uint64_t model_size, - uint64_t down_offset, uint32_t n_embd, uint32_t n_ff); +/* Decode attention glue for one token row, byte-identical to the standalone + * sequences. matvec_bf16: a Q8_0 or F16 single-row matvec whose store is + * the BF16 rounding (1 done, 0 not covered, -1 error). qkv_norm_kv_tail: + * the q LoRA and KV weighted norms with their roundings, then the KV RoPE, + * FP8 (E8M0) block quantization and the copy into the raw window row at + * window_offset. bf16_rope: the rounding pass over whole rows plus the + * (inverse) RoPE on their last 64 values. */ +enum { DS4_V41_WEIGHT_F16 = 1, DS4_V41_WEIGHT_Q8_0 = 8 }; /* GGUF tensor type ids */ +int ds4_gpu_dsv41_matvec_bf16(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint32_t type, uint32_t in_dim, uint32_t out_dim, + const ds4_gpu_tensor *x); +int ds4_gpu_dsv41_qkv_norm_kv_tail(ds4_gpu_tensor *q, ds4_gpu_tensor *kv, ds4_gpu_tensor *window, + uint64_t window_offset, const void *model_map, uint64_t model_size, + uint64_t q_weight_offset, uint64_t kv_weight_offset, + uint32_t q_n, uint32_t kv_n, float eps, uint32_t pos, + bool compressed); +int ds4_gpu_dsv41_bf16_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, bool compressed, bool inverse); #ifdef __cplusplus } diff --git a/ds4_metal.m b/ds4_metal.m index bce928e100..0bb0899500 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -47581,14 +47581,9 @@ static bool dsv41_tensor_has_floats(const ds4_gpu_tensor *tensor, uint64_t count /* Frequency rounding errors accumulate into phase errors at long contexts. * Preserve the reference's pow/reciprocal and YaRN operation order here. */ #pragma float_control(precise, on, push) -int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, - uint32_t rows, uint32_t start, uint32_t stride, - bool compressed, bool inverse) { - if (width < 64 || !heads || !rows || rows > 1048576 || !stride || - (uint64_t)start + (uint64_t)(rows - 1u) * stride >= 1048576u || - (uint64_t)heads * rows > UINT64_MAX / width || - !dsv41_tensor_has_floats(x, (uint64_t)width * heads * rows)) return 0; - if (!g_initialized && !ds4_gpu_init()) return 0; +/* The released V4.1 RoPE frequencies: plain (base 10000) and the YaRN-style + * compressed table (base 160000 with the 1/16 interpolation ramp). */ +static const float *dsv41_rope_frequencies(bool compressed) { static float frequencies[2][32]; static dispatch_once_t once; dispatch_once(&once, ^{ @@ -47610,14 +47605,25 @@ int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, } } }); + return frequencies[compressed ? 1 : 0]; +} + +static int dsv41_rope_dispatch(const char *kernel, ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, uint32_t stride, + bool compressed, bool inverse) { + if (width < 64 || !heads || !rows || rows > 1048576 || !stride || + (uint64_t)start + (uint64_t)(rows - 1u) * stride >= 1048576u || + (uint64_t)heads * rows > UINT64_MAX / width || + !dsv41_tensor_has_floats(x, (uint64_t)width * heads * rows)) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; @autoreleasepool { - id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_rope"); + id pipeline = ds4_gpu_get_pipeline(kernel); if (!pipeline) return 0; struct { uint32_t width, heads, rows, start, inverse, stride; float frequencies[32]; } args = {width, heads, rows, start, inverse, stride, {0}}; - memcpy(args.frequencies, frequencies[compressed ? 1 : 0], sizeof(args.frequencies)); + memcpy(args.frequencies, dsv41_rope_frequencies(compressed), sizeof(args.frequencies)); int owned = 0; id cb = ds4_gpu_command_buffer(&owned); id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; @@ -47633,6 +47639,126 @@ int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, } #pragma float_control(pop) +int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, uint32_t stride, + bool compressed, bool inverse) { + return dsv41_rope_dispatch("kernel_dsv41_rope", x, width, heads, rows, start, stride, + compressed, inverse); +} + +int ds4_gpu_dsv41_bf16_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, bool compressed, bool inverse) { + if (width % 32u) return 0; + return dsv41_rope_dispatch("kernel_dsv41_bf16_rope", x, width, heads, rows, start, 1, + compressed, inverse); +} + +int ds4_gpu_dsv41_qkv_norm_kv_tail(ds4_gpu_tensor *q, ds4_gpu_tensor *kv, ds4_gpu_tensor *window, + uint64_t window_offset, const void *model_map, uint64_t model_size, + uint64_t q_weight_offset, uint64_t kv_weight_offset, + uint32_t q_n, uint32_t kv_n, float eps, uint32_t pos, + bool compressed) { + const uint64_t q_bytes = (uint64_t)q_n * sizeof(float), kv_bytes = (uint64_t)kv_n * sizeof(float); + if (!q_n || q_n % 4u || kv_n < 64u || kv_n % 32u || kv_n > q_n || !model_map || + !dsv41_tensor_has_floats(q, q_n) || !dsv41_tensor_has_floats(kv, kv_n) || + !window || window_offset > ds4_gpu_tensor_bytes(window) || + kv_bytes > ds4_gpu_tensor_bytes(window) - window_offset || + q_weight_offset > model_size || q_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_bytes > model_size - kv_weight_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_qkv_norm_kv_tail"); + if (!pipeline) return 0; + /* The q row's standalone norm thread count (its reduction tree); the + * KV row has fewer elements than that and only its own lanes hold any. */ + const NSUInteger threads = ds4_gpu_rms_norm_threads(q_n); + const NSUInteger shared_bytes = ((NSUInteger)q_n + 32u) * sizeof(float); + if (threads < ds4_gpu_rms_norm_threads(kv_n) || threads > pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) return 0; + uint64_t q_inner = 0, kv_inner = 0; + id qwbuf = ds4_gpu_wrap_model_range(model_map, model_size, q_weight_offset, q_bytes, &q_inner); + id kvwbuf = ds4_gpu_wrap_model_range(model_map, model_size, kv_weight_offset, kv_bytes, &kv_inner); + if (!qwbuf || !kvwbuf) return 0; + struct { uint32_t q_n, kv_n; float eps; uint32_t pos; float frequencies[32]; } args = + {q_n, kv_n, eps, pos, {0}}; + memcpy(args.frequencies, dsv41_rope_frequencies(compressed), sizeof(args.frequencies)); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:qwbuf offset:(NSUInteger)q_inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(kv) offset:ds4_gpu_tensor_offset(kv) atIndex:3]; + [enc setBuffer:kvwbuf offset:(NSUInteger)kv_inner atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(window) + offset:ds4_gpu_tensor_offset(window) + (NSUInteger)window_offset atIndex:5]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(2, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 qkv norm + KV tail"); + } +} + +/* Single-row matvec with the BF16 rounding on the store. 1: done; 0: shape or + * type not covered (caller runs the standalone matvec + rounding); -1: error. */ +int ds4_gpu_dsv41_matvec_bf16(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint32_t type, uint32_t in_dim, uint32_t out_dim, + const ds4_gpu_tensor *x) { + if (!in_dim || !out_dim || !model_map || !dsv41_tensor_has_floats(x, in_dim) || + !dsv41_tensor_has_floats(out, out_dim)) return -1; + if (!g_initialized && !ds4_gpu_init()) return -1; + @autoreleasepool { + ds4_gpu_q8_0_matvec_args args; + ds4_gpu_mv_dispatch mv_dispatch; + uint64_t weight_bytes; + const char *kernel; + if (type == DS4_V41_WEIGHT_Q8_0) { + if (in_dim % 32u) return 0; + weight_bytes = ((uint64_t)in_dim / 32u) * 34u * out_dim; + args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + kernel = "kernel_dsv41_mul_mv_q8_0_f32_bf16"; + } else if (type == DS4_V41_WEIGHT_F16) { + if (in_dim < 32u || in_dim % 4u) return 0; + weight_bytes = (uint64_t)in_dim * sizeof(uint16_t) * out_dim; + ds4_gpu_f16_matvec_args f16_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); + mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + memcpy(&args, &f16_args, sizeof(args)); + kernel = "kernel_dsv41_mul_mv_f16_f32_4_bf16"; + } else { + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return -1; + args.nr0 = mv_dispatch.nr0; + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner); + id pipeline = ds4_gpu_get_mul_mv_pipeline(kernel, mv_dispatch.nsg); + if (!wbuf || !pipeline) return -1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return -1; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 BF16-rounded matvec") ? 1 : -1; + } +} + int ds4_gpu_dsv41_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, uint32_t rows, uint32_t start, bool compressed, bool inverse) { return ds4_gpu_dsv41_rope_stride(x, width, heads, rows, start, 1, compressed, inverse); @@ -47722,12 +47848,14 @@ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4 uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters, float hc_eps, float norm_eps) { const uint64_t weight_bytes = (uint64_t)n_embd * sizeof(float); - if (n_hc != 4u || !n_embd || n_embd % 4u || !sinkhorn_iters || !model_map || - !dsv41_tensor_has_floats(split, 24) || !dsv41_tensor_has_floats(mix, 24) || + /* sinkhorn_iters == 0: no split (the logits' collapse), mix/split/scale/base unused. */ + const bool with_split = sinkhorn_iters != 0; + if (n_hc != 4u || !n_embd || n_embd % 4u || !model_map || + (with_split && (!dsv41_tensor_has_floats(split, 24) || !dsv41_tensor_has_floats(mix, 24) || + scale_offset > model_size || 12u > model_size - scale_offset || + base_offset > model_size || 96u > model_size - base_offset)) || !dsv41_tensor_has_floats(pre, 4) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(norm, n_embd) || - scale_offset > model_size || 12u > model_size - scale_offset || - base_offset > model_size || 96u > model_size - base_offset || norm_weight_offset > model_size || weight_bytes > model_size - norm_weight_offset) return 0; if (!g_initialized && !ds4_gpu_init()) return 0; @autoreleasepool { @@ -47739,11 +47867,14 @@ int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4 if (threads > pipeline.maxTotalThreadsPerThreadgroup || shared_bytes > [g_device maxThreadgroupMemoryLength]) return 0; uint64_t scale_inner = 0, base_inner = 0, norm_inner = 0; - id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, 12u, &scale_inner); - id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, 96u, &base_inner); id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, weight_bytes, &norm_inner); + id scalebuf = with_split ? + ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, 12u, &scale_inner) : normwbuf; + id basebuf = with_split ? + ds4_gpu_wrap_model_range(model_map, model_size, base_offset, 96u, &base_inner) : normwbuf; if (!scalebuf || !basebuf || !normwbuf) return 0; + if (!with_split) { mix = x; split = x; scale_inner = base_inner = norm_inner; } const ds4_gpu_dsv41_hc_args args = {n_embd, sinkhorn_iters, hc_eps, norm_eps, 0, 0}; int owned = 0; id cb = ds4_gpu_command_buffer(&owned); @@ -47910,53 +48041,6 @@ int ds4_gpu_dsv41_shared_gate_up_swiglu(ds4_gpu_tensor *mid, const ds4_gpu_tenso } } -int ds4_gpu_dsv41_shared_down_hc_expand4(ds4_gpu_tensor *out, ds4_gpu_tensor *shared, - ds4_gpu_tensor *block, const ds4_gpu_tensor *mid, - const ds4_gpu_tensor *routed, const ds4_gpu_tensor *residual, - const ds4_gpu_tensor *split, ds4_gpu_tensor *pre, - const void *model_map, uint64_t model_size, - uint64_t down_offset, uint32_t n_embd, uint32_t n_ff) { - const uint64_t row_bytes = ((uint64_t)n_ff / 32u) * 34u; - const uint64_t weight_bytes = row_bytes * n_embd; - if (!n_embd || !n_ff || n_ff % 32u || !model_map || - !dsv41_tensor_has_floats(out, 4ull * n_embd) || !dsv41_tensor_has_floats(shared, n_embd) || - !dsv41_tensor_has_floats(block, n_embd) || !dsv41_tensor_has_floats(mid, n_ff) || - !dsv41_tensor_has_floats(routed, n_embd) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || - !dsv41_tensor_has_floats(split, 24) || (pre && !dsv41_tensor_has_floats(pre, 4)) || - down_offset > model_size || weight_bytes > model_size - down_offset) return 0; - if (!g_initialized && !ds4_gpu_init()) return 0; - @autoreleasepool { - uint64_t inner = 0; - id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, weight_bytes, &inner); - if (!wbuf) return 0; - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(n_ff, n_embd); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - mv_args.nr0 = mv_dispatch.nr0; - const ds4_gpu_dsv41_hc_args hc_args = {n_embd, 0, 0.0f, 0.0f, pre != NULL, 0}; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv41_shared_down_hc_expand4_q8_0", mv_dispatch.nsg); - if (!pipeline) return 0; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; - if (!enc) return 0; - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; - [enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:2]; - const ds4_gpu_tensor *buffers[] = {mid, shared, routed, block, residual, split, out, pre ? pre : split}; - for (NSUInteger i = 0; i < 8; i++) - [enc setBuffer:ds4_gpu_tensor_buffer(buffers[i]) - offset:ds4_gpu_tensor_offset(buffers[i]) atIndex:i + 3]; - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_embd + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 shared down + HC expand"); - } -} - int ds4_gpu_dsv41_carry_copy(ds4_gpu_tensor *packed, uint32_t row_offset, ds4_gpu_tensor *plain, uint32_t width, uint32_t rows, uint32_t format, bool pack) { diff --git a/metal/dsv41.metal b/metal/dsv41.metal index 7c021bcd31..e2ce1e66c2 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -401,7 +401,7 @@ kernel void kernel_dsv41_hc_collapse_norm4( threadgroup float4 *row = (threadgroup float4 *)shared; threadgroup float *sums = shared + args.n_embd; - if (tid == 0) dsv41_hc_split4(mix, scale, base, split_out, args.sinkhorn_iters, args.hc_eps); + if (tid == 0 && args.sinkhorn_iters) dsv41_hc_split4(mix, scale, base, split_out, args.sinkhorn_iters, args.hc_eps); if (sgitg == 0) sums[tiisg] = 0.0f; const float4 p = *((device const float4 *)pre); @@ -476,6 +476,250 @@ kernel void kernel_dsv41_hc_expand4_bf16( if (args.copy_pre && d < 4) pre_out[d] = split[d]; } +/* Single-row matvecs that round to BF16 on the store: the V4.1 graph rounds + * almost every projection output, as a separate dispatch over the result. + * Bodies are kernel_mul_mv_q8_0_f32_impl and kernel_mul_mv_t_t_4_impl + * with helper_mv_reduce_and_write's tree; only + * the store changes (tests/test_deepseek41_metal --attn-fuse). */ +template +static inline void dsv41_mv_reduce_write_bf16( + device float * dst_f32, float sumf[NR0], const int r0, const int ne01, + ushort tiisg, ushort sgitg, threadgroup char * shmem) { + constexpr short NW = N_SIMDWIDTH; + threadgroup float * shmem_f32[NR0]; + for (short row = 0; row < NR0; ++row) { + shmem_f32[row] = (threadgroup float *) shmem + NW*row; + if (sgitg == 0) shmem_f32[row][tiisg] = 0.0f; + sumf[row] = simd_sum(sumf[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (short row = 0; row < NR0; ++row) { + if (tiisg == 0) shmem_f32[row][sgitg] = sumf[row]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (short row = 0; row < NR0 && r0 + row < ne01; ++row) { + float tot = simd_sum(shmem_f32[row][tiisg]); + if (tiisg == 0 && sgitg == 0) dst_f32[r0 + row] = dsv41_bf16(tot); + } +} + +kernel void kernel_dsv41_mul_mv_q8_0_f32_bf16( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short NR0 = N_R0_Q8_0; + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + const int nb = args.ne00/QK8_0; + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + device const float * y = (device const float *) (src1 + offset1); + device const block_q8_0 * ax[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0); + } + float sumf[NR0] = { 0.f }; + const short ix = tiisg/(NW/NQ); + const short il = tiisg%(NW/NQ); + const int ib0 = sgitg*NQ + ix; + float yl[NQ]; + device const float * yb = y + ib0*QK8_0 + il*NQ; + for (int ib = ib0; ib < nb; ib += NSG*NQ) { + for (short i = 0; i < NQ; ++i) yl[i] = yb[i]; + for (short row = 0; row < NR0; row++) { + device const int8_t * qs = ax[row][ib].qs + il*NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) sumq += qs[i] * yl[i]; + sumf[row] += sumq*ax[row][ib].d; + } + yb += NSG*NQ*QK8_0; + } + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + dsv41_mv_reduce_write_bf16(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +template +void dsv41_mul_mv_f16_f32_4_bf16_impl( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + const int nb = args.ne00/NB; + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + device const float * y = (device const float *) (src1 + offset1); + device const float4 * y4 = (device const float4 *) (src1 + offset1); + device const half * ax [NR0]; + device const half4 * ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + ax [row] = (device const half *) ((device char *) src0 + offset0); + ax4[row] = (device const half4 *) ((device char *) src0 + offset0); + } + float sumf[NR0] = { 0.f }; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = sgitg*NF + ix; + float4 yl4[NF4]; + device const float4 * yb4 = y4 + (ib0*NB + il*NF)/4; + for (int ib = ib0; ib < nb; ib += NSG*NF) { + for (short i = 0; i < NF4; ++i) yl4[i] = yb4[i]; + for (short row = 0; row < NR0; row++) { + device const half4 * xb4 = ax4[row] + (ib*NB + il*NF)/4; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) sumq += dot(float4(xb4[i]), float4(yl4[i])); + sumf[row] += sumq; + } + yb4 += NSG*NF*NW/4; + } + for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) { + for (short row = 0; row < NR0; row++) sumf[row] += ax[row][i] * y[i]; + } + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + dsv41_mv_reduce_write_bf16(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +kernel void kernel_dsv41_mul_mv_f16_f32_4_bf16( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + switch (args.nr0) { + case 2: dsv41_mul_mv_f16_f32_4_bf16_impl<2>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + case 4: dsv41_mul_mv_f16_f32_4_bf16_impl<4>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + } +} + +/* BF16 rounding of a whole row followed by the adjacent-pair RoPE on its + * last 64 values (kernel_dsv41_rope), one simdgroup per row: what the graph + * ran as a rounding pass over all heads plus the inverse RoPE. */ +kernel void kernel_dsv41_bf16_rope( + constant ds4_metal_args_dsv41_rope &args, + device float *x, + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const ulong row = ((ulong)group.y * args.heads + group.x) * args.width; + for (uint i = lane; i < args.width; i += 32u) x[row + i] = dsv41_bf16(x[row + i]); + simdgroup_barrier(mem_flags::mem_device); + const float theta = float(args.start + group.y * args.stride) * args.frequencies[lane]; + const float c = precise::cos(theta); + const float s = args.inverse ? -precise::sin(theta) : precise::sin(theta); + const ulong i = row + args.width - 64u + 2u * lane; + const float re = x[i], im = x[i + 1u]; + x[i] = dsv41_bf16(re * c - im * s); + x[i + 1u] = dsv41_bf16(re * s + im * c); +} + +/* The attention projections' normalization and the KV tail in one dispatch: + * threadgroup y=0 is kernel_rms_norm_mul_f32_4 over the q LoRA row plus its + * BF16 rounding; y=1 the same over the 512-wide KV row, then the RoPE on + * its last 64 values, the FP8 (E8M0-scaled) block quantization of + * kernel_dsv41_quantize mode 1 and the store into the raw window slot. + * The threadgroup is the q row's norm thread count; the KV lanes beyond its + * own count hold no elements, so their zero partials leave the standalone + * reduction tree unchanged. */ +struct ds4_metal_args_dsv41_qkv_tail { + uint q_n; + uint kv_n; + float eps; + uint pos; + float frequencies[32]; +}; + +kernel void kernel_dsv41_qkv_norm_kv_tail( + constant ds4_metal_args_dsv41_qkv_tail &args, + device float4 *q, + device const float4 *q_weight, + device float4 *kv, + device const float4 *kv_weight, + device float *window_row, + threadgroup float *shared [[threadgroup(0)]], + ushort task [[threadgroup_position_in_grid]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + const bool kv_task = task != 0; + const uint n = kv_task ? args.kv_n : args.q_n; + const uint n4 = n >> 2; + device float4 *x = kv_task ? kv : q; + device const float4 *w = kv_task ? kv_weight : q_weight; + threadgroup float4 *row = (threadgroup float4 *)shared; + threadgroup float *sums = shared + 4u * n4; + + if (sgitg == 0) sums[tiisg] = 0.0f; + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + const float4 v = x[i]; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) sums[sgitg] = sumf; + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = simd_sum(sums[tiisg]); + const float mean = sumf / (float)n; + const float scale = 1.0f / sqrt(mean + args.eps); + for (uint i = tid; i < n4; i += ntg) { + const float4 v = dsv41_bf16x4((x[i] * scale) * w[i]); + row[i] = v; + if (!kv_task) x[i] = v; + } + if (!kv_task) return; + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float *rowf = (threadgroup float *)row; + if (sgitg == 0) { + const float theta = float(args.pos) * args.frequencies[tiisg]; + const float c = precise::cos(theta); + const float s = precise::sin(theta); + const uint i = n - 64u + 2u * tiisg; + const float re = rowf[i], im = rowf[i + 1u]; + rowf[i] = dsv41_bf16(re * c - im * s); + rowf[i + 1u] = dsv41_bf16(re * s + im * c); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint nsg = (ntg + 31u) / 32u; + for (uint block = sgitg; block < n / 32u; block += nsg) { + const uint column = block * 32u + tiisg; + const float value = dsv41_bf16(rowf[column]); + const float amax = simd_max(abs(value)); + const float qscale = dsv41_pow2_ceil(max(amax, 1.0e-4f) * (1.0f / 448.0f)); + const float result = copysign(dsv4_e4m3fn_dequant(abs(value) / qscale), value) * qscale; + const float out = dsv41_bf16(result); + ((device float *)kv)[column] = out; + window_row[column] = out; + } +} + /* Decode-time V4.1 MoE glue, one token row. Same discipline as the HC * kernels above: every reduction keeps the standalone kernel's shape and * every BF16 rounding point is kept (tests/test_deepseek41_metal --moe-fuse). */ @@ -735,98 +979,3 @@ kernel void kernel_dsv41_shared_gate_up_swiglu_q8_0( } } -/* Shared expert down projection (Q8_0 matvec, standalone walk and tree), - * then the layer's FFN tail exactly as the graph ran it in six dispatches: - * shared = bf16(down); block = bf16(routed + shared); residual_out = - * bf16(post/comb expand of block into residual); pre = split[0..3]. */ -kernel void kernel_dsv41_shared_down_hc_expand4_q8_0( - constant ds4_metal_args_mul_mv & mv, - constant ds4_metal_args_dsv41_hc & hc, - device const char * weight, - device const char * shared_mid, - device float * shared_out, - device const float * routed_out, - device float * block_out, - device const float * residual, - device const float * split, - device float * dst, - device float * pre_out, - threadgroup char * shmem [[threadgroup(0)]], - uint3 tgpig[[threadgroup_position_in_grid]], - ushort tiisg[[thread_index_in_simdgroup]], - ushort sgitg[[simdgroup_index_in_threadgroup]]) { - const short NSG = FC_mul_mv_nsg; - constexpr short NW = N_SIMDWIDTH; - constexpr short NQ = 8; - constexpr short NR0 = N_R0_Q8_0; - - const int nb = mv.ne00 / QK8_0; - const int row0 = tgpig.x * NR0; - - const short ix = tiisg / (NW / NQ); - const short il = tiisg % (NW / NQ); - const int ib0 = sgitg * NQ + ix; - - device const float *y = (device const float *)(shared_mid); - device const float *yb = y + ib0 * QK8_0 + il * NQ; - - device const block_q8_0 *ax[NR0]; - FOR_UNROLL(short row = 0; row < NR0; ++row) { - ax[row] = (device const block_q8_0 *)(weight + (uint64_t)(row0 + row) * mv.nb01); - } - - float sumf[NR0] = { 0.0f }; - float yl[NQ]; - for (int ib = ib0; ib < nb; ib += NSG * NQ) { - FOR_UNROLL(short i = 0; i < NQ; ++i) yl[i] = yb[i]; - FOR_UNROLL(short row = 0; row < NR0; ++row) { - device const int8_t *qs = ax[row][ib].qs + il * NQ; - float sumq = 0.0f; - FOR_UNROLL(short i = 0; i < NQ; ++i) sumq += qs[i] * yl[i]; - sumf[row] += sumq * ax[row][ib].d; - } - yb += NSG * NQ * QK8_0; - } - - threadgroup float *shmem_f32[NR0]; - FOR_UNROLL(short row = 0; row < NR0; ++row) { - shmem_f32[row] = (threadgroup float *)shmem + NW * row; - if (sgitg == 0) shmem_f32[row][tiisg] = 0.0f; - sumf[row] = simd_sum(sumf[row]); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - FOR_UNROLL(short row = 0; row < NR0; ++row) { - if (tiisg == 0) shmem_f32[row][sgitg] = sumf[row]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - device const float *post = split + 4; - device const float *comb = split + 8; - FOR_UNROLL(short row = 0; row < NR0; ++row) { - const int d = row0 + row; - if (d >= mv.ne01) continue; - const float shared_v = simd_sum(shmem_f32[row][tiisg]); - if (tiisg == 0 && sgitg == 0) { - const float sv = dsv41_bf16(shared_v); - shared_out[d] = sv; - float block_v = routed_out[d]; - block_v += sv; - block_v = dsv41_bf16(block_v); - block_out[d] = block_v; - - const float r0 = residual[d]; - const float r1 = residual[d + hc.n_embd]; - const float r2 = residual[d + 2u * hc.n_embd]; - const float r3 = residual[d + 3u * hc.n_embd]; - for (uint dst_hc = 0; dst_hc < 4; ++dst_hc) { - float acc = block_v * post[dst_hc]; - acc += comb[dst_hc + 0u * 4u] * r0; - acc += comb[dst_hc + 1u * 4u] * r1; - acc += comb[dst_hc + 2u * 4u] * r2; - acc += comb[dst_hc + 3u * 4u] * r3; - dst[d + dst_hc * hc.n_embd] = dsv41_bf16(acc); - } - } - } - if (hc.copy_pre && tgpig.x == 0 && sgitg == 0 && tiisg < 4) pre_out[tiisg] = split[tiisg]; -} diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 560ef006b1..13f403d16a 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -484,8 +484,8 @@ static int check_moe_fuse(void) { CHECK(ds4_gpu_dsv41_router_select(selected[1], weights[1], probs[1], logits[1], norm, model, mapped, router_off, bias_off, true, D, E, K, scale)); CHECK(ds4_gpu_dsv41_shared_gate_up_swiglu(mid[1], norm, model, mapped, gate_off, up_off, D, FF, clamp)); - CHECK(ds4_gpu_dsv41_shared_down_hc_expand4(out[1], shared[1], block[1], mid[1], routed, - residual, split, pre[1], model, mapped, down_off, D, FF)); + CHECK(ds4_gpu_dsv41_matvec_bf16(shared[1], model, mapped, down_off, DS4_V41_WEIGHT_Q8_0, FF, D, mid[1]) == 1); + CHECK(ds4_gpu_dsv41_hc_expand4(out[1], routed, residual, split, pre[1], shared[1], block[1], D)); } CHECK(ds4_gpu_end_commands()); CHECK(ds4_gpu_synchronize()); @@ -503,20 +503,8 @@ static int check_moe_fuse(void) { CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(pre[0]), ds4_gpu_tensor_contents(pre[1]), 16)); - /* The early-down variant: standalone down + BF16, then the sum and - * its rounding folded into the expand. */ - CHECK(ds4_gpu_begin_commands()); - CHECK(ds4_gpu_matmul_q8_0_tensor(shared[1], model, mapped, down_off, FF, D, mid[1], 1)); - CHECK(ds4_gpu_dsv41_quantize(shared[1], D, 1, DS4_V41_BF16)); - CHECK(ds4_gpu_dsv41_hc_expand4(out[1], routed, residual, split, pre[1], shared[1], block[1], D)); - CHECK(ds4_gpu_end_commands()); - CHECK(ds4_gpu_synchronize()); - CHECK(!memcmp(ds4_gpu_tensor_contents(shared[0]), ds4_gpu_tensor_contents(shared[1]), D * 4)); - CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4)); - CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); - CHECK(!memcmp(ds4_gpu_tensor_contents(pre[0]), ds4_gpu_tensor_contents(pre[1]), 16)); } - fprintf(stderr, "MoE fuse: 15 dispatches %.3f ms -> 3 dispatches %.3f ms, outputs byte-identical\n", + fprintf(stderr, "MoE fuse: 15 dispatches %.3f ms -> 4-5 dispatches %.3f ms, outputs byte-identical\n", elapsed[0], elapsed[1]); for (int i = 0; i < 2; i++) { ds4_gpu_tensor_free(logits[i]); ds4_gpu_tensor_free(probs[i]); ds4_gpu_tensor_free(selected[i]); @@ -530,6 +518,131 @@ static int check_moe_fuse(void) { return 1; } +/* The fused decode attention glue: BF16-on-store matvecs (Q8_0 and F16 at + * the graph's shapes, including the nr0=4 F16 dispatch), the q/kv norms + * with the KV RoPE/FP8/window tail, the heads' BF16 + inverse RoPE, and the + * logits' collapse + norm; each byte-identical to its standalone sequence. */ +static int check_attn_fuse(void) { + enum { D = 5120, LQ = 1280, HD = 512, HEADS = 64, QD = HEADS * HD, HC = 4, IDX = 1024 }; + const float eps = 1e-20f; + struct { uint32_t type, in, out; size_t bytes, off; } w[] = { + {DS4_V41_WEIGHT_Q8_0, D, LQ, 0, 0}, /* attn_q_a */ + {DS4_V41_WEIGHT_Q8_0, D, HD, 0, 0}, /* attn_kv */ + {DS4_V41_WEIGHT_Q8_0, LQ, QD, 0, 0}, /* attn_q_b */ + {DS4_V41_WEIGHT_F16, D, HD, 0, 0}, /* attn_compressor_kv: the nr0 = 4 dispatch */ + {DS4_V41_WEIGHT_F16, HD, 128, 0, 0}, /* indexer_attn_k */ + {DS4_V41_WEIGHT_F16, LQ, IDX, 0, 0}, /* indexer_attn_q_b */ + }; + const int nw = sizeof(w) / sizeof(*w); + size_t total = 0; + for (int i = 0; i < nw; i++) { + w[i].bytes = w[i].type == DS4_V41_WEIGHT_Q8_0 ? (size_t)w[i].out * (w[i].in / 32) * 34 : (size_t)w[i].out * w[i].in * 2; + w[i].off = total; + total += (w[i].bytes + 63) / 64 * 64; + } + const size_t qnorm_off = total, kvnorm_off = qnorm_off + LQ * 4, onorm_off = kvnorm_off + HD * 4; + const size_t page = (size_t)getpagesize(); + const size_t mapped = (onorm_off + D * 4 + page - 1) / page * page; + void *model = NULL; + CHECK(!posix_memalign(&model, page, mapped)); + for (int i = 0; i < nw; i++) { + if (w[i].type == DS4_V41_WEIGHT_Q8_0) fill_q8_0((uint8_t *)model + w[i].off, w[i].out, w[i].in); + else { _Float16 *h = (_Float16 *)((char *)model + w[i].off); for (size_t k = 0; k < (size_t)w[i].out * w[i].in; k++) h[k] = (_Float16)(random_value() / 64); } + } + float *qnorm = (float *)((char *)model + qnorm_off), *kvnorm = (float *)((char *)model + kvnorm_off); + float *onorm = (float *)((char *)model + onorm_off); + for (int i = 0; i < LQ; i++) qnorm[i] = 1.0f + random_value() / 8; + for (int i = 0; i < HD; i++) kvnorm[i] = 1.0f + random_value() / 8; + for (int i = 0; i < D; i++) onorm[i] = 1.0f + random_value() / 8; + CHECK(ds4_gpu_set_model_map(model, mapped)); + + ds4_gpu_tensor *xin[2] = {upload(NULL, D * 4), upload(NULL, LQ * 4)}; /* D-wide and LQ/HD-wide inputs */ + ds4_gpu_tensor *out[2] = {upload(NULL, QD * 4), upload(NULL, QD * 4)}; + ds4_gpu_tensor *qr[2] = {upload(NULL, LQ * 4), upload(NULL, LQ * 4)}; + ds4_gpu_tensor *kv[2] = {upload(NULL, HD * 4), upload(NULL, HD * 4)}; + ds4_gpu_tensor *window[2] = {upload(NULL, 128 * HD * 4), upload(NULL, 128 * HD * 4)}; + ds4_gpu_tensor *heads[2] = {upload(NULL, QD * 4), upload(NULL, QD * 4)}; + ds4_gpu_tensor *residual = upload(NULL, HC * D * 4), *pre = upload(NULL, 16); + ds4_gpu_tensor *x[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *norm[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + CHECK(xin[0] && xin[1] && residual && pre); + for (int i = 0; i < 2; i++) CHECK(out[i] && qr[i] && kv[i] && window[i] && heads[i] && x[i] && norm[i]); + float *xd = ds4_gpu_tensor_contents(xin[0]), *xl = ds4_gpu_tensor_contents(xin[1]); + for (int round = 0; round < 4; round++) { + for (int i = 0; i < D; i++) xd[i] = bf16(random_value()); + for (int i = 0; i < LQ; i++) xl[i] = bf16(random_value()); + for (int i = 0; i < nw; i++) { + const ds4_gpu_tensor *in = w[i].in == D ? xin[0] : xin[1]; + CHECK(ds4_gpu_begin_commands()); + if (w[i].type == DS4_V41_WEIGHT_Q8_0) CHECK(ds4_gpu_matmul_q8_0_tensor(out[0], model, mapped, w[i].off, w[i].in, w[i].out, in, 1)); + else CHECK(ds4_gpu_matmul_f16_tensor(out[0], model, mapped, w[i].off, w[i].in, w[i].out, in, 1)); + CHECK(ds4_gpu_dsv41_quantize(out[0], w[i].out, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_matvec_bf16(out[1], model, mapped, w[i].off, w[i].type, w[i].in, w[i].out, in) == 1); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + if (memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), (size_t)w[i].out * 4)) { + fprintf(stderr, "matvec bf16 mismatch: weight %d (type %u, %u -> %u)\n", i, w[i].type, w[i].in, w[i].out); + CHECK(0); + } + } + /* q/kv norms + KV tail, both frequency tables and window slots. */ + const uint32_t pos = 777u + 1000u * round; + const int compressed = round & 1; + float *q0 = ds4_gpu_tensor_contents(qr[0]), *q1 = ds4_gpu_tensor_contents(qr[1]); + float *k0 = ds4_gpu_tensor_contents(kv[0]), *k1 = ds4_gpu_tensor_contents(kv[1]); + for (int i = 0; i < LQ; i++) q0[i] = q1[i] = bf16(random_value()); + for (int i = 0; i < HD; i++) k0[i] = k1[i] = bf16(random_value()); + const uint64_t slot = (uint64_t)(pos % 128u) * HD * 4u; + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_rms_norm_weight_tensor(qr[0], qr[0], model, mapped, qnorm_off, LQ, eps)); + CHECK(ds4_gpu_dsv41_quantize(qr[0], LQ, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_rms_norm_weight_tensor(kv[0], kv[0], model, mapped, kvnorm_off, HD, eps)); + CHECK(ds4_gpu_dsv41_quantize(kv[0], HD, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_rope(kv[0], HD, 1, 1, pos, compressed, false)); + CHECK(ds4_gpu_dsv41_quantize(kv[0], HD, 1, DS4_V41_FP8_E8M0)); + CHECK(ds4_gpu_tensor_copy(window[0], slot, kv[0], 0, HD * 4)); + CHECK(ds4_gpu_dsv41_qkv_norm_kv_tail(qr[1], kv[1], window[1], slot, model, mapped, qnorm_off, kvnorm_off, LQ, HD, eps, pos, compressed)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(q0, q1, LQ * 4)); + CHECK(!memcmp(k0, k1, HD * 4)); + CHECK(!memcmp((char *)ds4_gpu_tensor_contents(window[0]) + slot, (char *)ds4_gpu_tensor_contents(window[1]) + slot, HD * 4)); + /* heads: BF16 over every value, inverse RoPE on each head's last 64 */ + float *h0 = ds4_gpu_tensor_contents(heads[0]), *h1 = ds4_gpu_tensor_contents(heads[1]); + for (int i = 0; i < QD; i++) h0[i] = h1[i] = random_value() * 3.0f; + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_dsv41_quantize(heads[0], QD, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_rope(heads[0], HD, HEADS, 1, pos, compressed, true)); + CHECK(ds4_gpu_dsv41_bf16_rope(heads[1], HD, HEADS, 1, pos, compressed, true)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(h0, h1, QD * 4)); + /* logits: collapse with the carried pre, BF16, output norm, BF16 */ + float *res = ds4_gpu_tensor_contents(residual), *p = ds4_gpu_tensor_contents(pre); + for (int i = 0; i < HC * D; i++) res[i] = bf16(random_value()); + for (int i = 0; i < HC; i++) p[i] = 0.5f + random_value() / 8; + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_hc_weighted_sum_tensor(x[0], residual, pre, D, HC)); + CHECK(ds4_gpu_dsv41_quantize(x[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_rms_norm_weight_tensor(norm[0], x[0], model, mapped, onorm_off, D, eps)); + CHECK(ds4_gpu_dsv41_quantize(norm[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_hc_collapse_norm(NULL, x[1], norm[1], NULL, pre, residual, model, mapped, 0, 0, onorm_off, D, HC, 0, 0.0f, eps)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(ds4_gpu_tensor_contents(x[0]), ds4_gpu_tensor_contents(x[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(norm[0]), ds4_gpu_tensor_contents(norm[1]), D * 4)); + } + fprintf(stderr, "attention fuse: matvec bf16 (6 shapes), qkv norm + KV tail, bf16 + rope, logits collapse: byte-identical\n"); + for (int i = 0; i < 2; i++) { + ds4_gpu_tensor_free(xin[i]); ds4_gpu_tensor_free(out[i]); ds4_gpu_tensor_free(qr[i]); ds4_gpu_tensor_free(kv[i]); + ds4_gpu_tensor_free(window[i]); ds4_gpu_tensor_free(heads[i]); ds4_gpu_tensor_free(x[i]); ds4_gpu_tensor_free(norm[i]); + } + ds4_gpu_tensor_free(residual); ds4_gpu_tensor_free(pre); + ds4_gpu_cleanup(); free(model); + CHECK(ds4_gpu_init()); + return 1; +} + #endif static int check_engram(void) { @@ -1503,6 +1616,11 @@ int main(int argc, char **argv) { ds4_gpu_cleanup(); return ok ? 0 : 1; } + if (argc == 2 && !strcmp(argv[1], "--attn-fuse")) { + const int ok = ds4_gpu_init() && check_attn_fuse(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } #endif if (argc == 2 && !strcmp(argv[1], "--bf16-linear")) { const int ok = ds4_gpu_init() && check_bf16_linear(); @@ -1540,7 +1658,7 @@ int main(int argc, char **argv) { check_embedding() && check_index_projection() && check_general_topk() && check_causal_topk() && check_compact_carry() && check_attention_output(false) && check_tp_attention(); #ifdef __APPLE__ - ok = ok && check_hc_fuse() && check_moe_fuse(); + ok = ok && check_hc_fuse() && check_moe_fuse() && check_attn_fuse(); #endif ds4_gpu_cleanup(); return ok ? 0 : 1; From 6e92e92047e9b79e0258109dc3ef913ecac7d4a6 Mon Sep 17 00:00:00 2001 From: Adrian Galilea Date: Mon, 14 Sep 2026 14:54:45 +0200 Subject: [PATCH 7/7] v41: MoE fuse test compares ties to the canonical order only The fused select orders equal scores by ascending index; an argsort without that tie-break agrees on the selected set but not necessarily on its order, so the tie rounds check the fused order alone. --- tests/test_deepseek41_metal.c | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 13f403d16a..69ac34c424 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -493,11 +493,16 @@ static int check_moe_fuse(void) { } const int32_t *sel = ds4_gpu_tensor_contents(selected[1]); for (int i = 0; i < K; i++) CHECK(sel[i] >= 0 && sel[i] < E); + /* Ties: the fused select orders equal scores by ascending index (the + * canonical argsort order); an argsort without that tie-break agrees + * on the set but not necessarily on the order. */ if (ties) for (int i = 0; i < K; i++) CHECK(sel[i] == 4 + 5 * i); CHECK(!memcmp(ds4_gpu_tensor_contents(logits[0]), ds4_gpu_tensor_contents(logits[1]), E * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(probs[0]), ds4_gpu_tensor_contents(probs[1]), E * 4)); - CHECK(!memcmp(ds4_gpu_tensor_contents(selected[0]), sel, K * 4)); - CHECK(!memcmp(ds4_gpu_tensor_contents(weights[0]), ds4_gpu_tensor_contents(weights[1]), K * 4)); + if (!ties) { + CHECK(!memcmp(ds4_gpu_tensor_contents(selected[0]), sel, K * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(weights[0]), ds4_gpu_tensor_contents(weights[1]), K * 4)); + } CHECK(!memcmp(ds4_gpu_tensor_contents(mid[0]), ds4_gpu_tensor_contents(mid[1]), FF * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(shared[0]), ds4_gpu_tensor_contents(shared[1]), D * 4)); CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4));