diff --git a/metal/dense.metal b/metal/dense.metal index ffcdb36..b56f509 100644 --- a/metal/dense.metal +++ b/metal/dense.metal @@ -197,33 +197,6 @@ kernel void kernel_mul_mv_q8_0_f32( kernel_mul_mv_q8_0_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); } -[[host_name("kernel_mul_mv_q8_0_f32_r4")]] -kernel void kernel_mul_mv_q8_0_f32_r4( - 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]]) { - kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); -} - -// Output projection alias used by the optimized host dispatch. -[[host_name("kernel_mul_mv_q8_0_f32_nr4")]] -kernel void kernel_mul_mv_q8_0_f32_nr4( - 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]]) { - kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>( - args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); -} // Decode Q-A/KV pair. Both projections consume the same activation row but // have independent weight ranges and output extents. Keep the standalone Q8_0 @@ -497,25 +470,190 @@ kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0( clamp_value, shmem, tgpig, tiisg, sgitg); } -[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0_r4")]] -kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0_r4( +// Decode-only fusion of the router logits matvec (F16, embd -> n_expert) +// with the shared-expert gate/up SwiGLU (Q8_0, embd -> shared). Both read +// the same normalized FFN input back to back; one dispatch removes one +// launch per decode layer. Router threadgroups replicate +// kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2); shared threadgroups host two +// virtual 4-simdgroup cohorts replicating +// kernel_dsv4_shared_gate_up_swiglu_q8_0 (nsg=4, nr0=2), including its +// per-row simd/shmem reduction trees. Bit-exact by construction. +kernel void kernel_dsv4_router_shared_gate_up_q8_0( constant ds4_metal_args_mul_mv & args, + constant ds4_metal_args_mul_mv & sargs, + device const char * src0_router, device const char * src0_gate, device const char * src0_up, device const char * src1, + device char * dst_router, device char * dst_gate, device char * dst_up, 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]]) { - kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, true>( - args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, - clamp_value, shmem, tgpig, tiisg, sgitg); + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr short NW = N_SIMDWIDTH; + const uint router_tgs = ((uint)args.ne01 + 1u) / 2u; + + if (tgpig.x < router_tgs) { + // Exact replica of kernel_mul_mv_f16_f32_4 with NSG=8, NR0=2. + constexpr short NSG = 8; + constexpr short NR0 = 2; + 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; + + device const float4 * y4 = (device const float4 *) src1; + + device const half4 * ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const half4 *) + (src0_router + (uint64_t)(r0 + row)*args.nb01); + } + + 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 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]), yl4[i]); + } + + sumf[row] += sumq; + } + + yb4 += NSG*NF*NW/4; + } + + device float * dst_f32 = (device float *) dst_router; + helper_mv_reduce_and_write(dst_f32, sumf, r0, args.ne01, + tiisg, sgitg, shmem); + return; + } + + // Shared-expert part: two virtual nsg=4 cohorts per threadgroup, each an + // exact replica of kernel_dsv4_shared_gate_up_swiglu_q8_0 (NR0=2). + constexpr short NSG = 4; + constexpr short NR0 = 2; + constexpr short NQ = 8; + + const uint cohort = sgitg >> 2; + const ushort vsg = sgitg & 3u; + const uint vt = (tgpig.x - router_tgs) * 2u + cohort; + + const int nb = sargs.ne00 / QK8_0; + const int r0 = vt * 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 = (uint64_t)(r0 + row) * sargs.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 = vsg * 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 + cohort * (2*NR0*NW); + 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 (vsg == 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][vsg] = sumg[row]; + sh_up[row][vsg] = sumu[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *gate_f32 = (device float *)dst_gate; + device float *up_f32 = (device float *)dst_up; + device float *mid_f32 = (device float *)dst_mid; + + FOR_UNROLL (short row = 0; row < NR0 && r0 + row < sargs.ne01; ++row) { + const float gate = simd_sum(sh_gate[row][tiisg]); + const float up = simd_sum(sh_up[row][tiisg]); + if (tiisg == 0 && vsg == 0) { + const uint out_row = r0 + row; + gate_f32[out_row] = gate; + up_f32[out_row] = up; + float g = gate; + float u = 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[out_row] = silu * u; + } + } } + [[host_name("kernel_dsv4_shared_mid_swiglu_q8_0")]] kernel void kernel_dsv4_shared_mid_swiglu_q8_0( constant ds4_metal_args_mul_mv & args, @@ -535,24 +673,6 @@ kernel void kernel_dsv4_shared_mid_swiglu_q8_0( clamp_value, shmem, tgpig, tiisg, sgitg); } -[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_r4")]] -kernel void kernel_dsv4_shared_mid_swiglu_q8_0_r4( - constant ds4_metal_args_mul_mv & args, - device const char * src0_gate, - device const char * src0_up, - device const char * src1, - device char * dst_gate, - device char * dst_up, - 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]]) { - kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, false>( - args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, - clamp_value, shmem, tgpig, tiisg, sgitg); -} template void kernel_mul_mv_t_t_impl( @@ -974,6 +1094,311 @@ kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4( state_score[dst] = projected_score[col] + ape_v; } +// Decode compressor + indexer-compressor projection in one dispatch. Both +// pairs read the same normalized activation with the same F16 matvec shape, +// so one launch covers all four matrices: threadgroups below the first +// range boundary run the exact paired matvec + state store of +// kernel_mul_mv_f16_f32_pair_compressor_store_4 for the attention +// compressor, the rest for the indexer compressor. Per-row reduction trees +// and the per-threadgroup state stores are unchanged, keeping the fused +// result bit-identical to the two separate dispatches while removing one +// dispatch per decode layer. +kernel void kernel_mul_mv_f16_f32_quad_compressor_store_4( + constant ds4_metal_args_mul_mv & args, + constant ds4_metal_args_compressor_pair_store & store0, + constant ds4_metal_args_compressor_pair_store & store1, + device const char * src0_a0, + device const char * src0_b0, + device const char * src0_a1, + device const char * src0_b1, + device const char * src1, + device char * dst_a0, + device char * dst_b0, + device char * dst_a1, + device char * dst_b1, + device const char * ape0, + device const char * ape1, + device float * state0_kv, + device float * state0_score, + device float * state1_kv, + device float * state1_score, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr short NR0 = 2; + const uint tgs0 = ((uint)store0.width + NR0 - 1u) / NR0; + const bool second = tgpig.x >= tgs0; + + uint3 local_tgpig = tgpig; + if (second) local_tgpig.x = tgpig.x - tgs0; + + ds4_metal_args_mul_mv largs = args; + largs.nr0 = NR0; + largs.ne01 = second ? (int32_t)store1.width : (int32_t)store0.width; + + if (!second) { + kernel_mul_mv_f16_f32_pair_4_impl( + largs, src0_a0, src0_b0, src1, dst_a0, dst_b0, + shmem, local_tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_f16_f32_pair_4_impl( + largs, src0_a1, src0_b1, src1, dst_a1, dst_b1, + shmem, local_tgpig, tiisg, sgitg); + } + + threadgroup_barrier(mem_flags::mem_device); + + // State append: identical to the paired store kernel, scoped to the + // range this threadgroup just projected (its own outputs only). + constant ds4_metal_args_compressor_pair_store & store = second ? store1 : store0; + if (tiitg >= NR0 || store.width == 0u || store.ratio == 0u) { + return; + } + const uint col = local_tgpig.x * (uint)NR0 + tiitg; + if (col >= store.width) return; + + const uint pos_mod = store.pos % store.ratio; + const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod; + const uint dst = dst_row * store.width + col; + const uint ape_i = pos_mod * store.width + col; + + device volatile const float * projected_kv = second + ? (device volatile const float *)dst_a1 + : (device volatile const float *)dst_a0; + device volatile const float * projected_score = second + ? (device volatile const float *)dst_b1 + : (device volatile const float *)dst_b0; + device const char * ape = second ? ape1 : ape0; + device float * state_kv = second ? state1_kv : state0_kv; + device float * state_score = second ? state1_score : state0_score; + + float ape_v; + if (store.ape_type == 1u) { + ape_v = (float)(((device const half *)ape)[ape_i]); + } else { + ape_v = ((device const float *)ape)[ape_i]; + } + + state_kv[dst] = projected_kv[col]; + state_score[dst] = projected_score[col] + ape_v; +} + +/* Decode-only fusion: one dispatch covers the q_a/kv Q8 pair projection and + * the four F16 compressor projections (attention + indexer) with their + * state-store epilogue. Both stages read the same normalized attention + * input and write disjoint outputs. The q_a/kv range hosts two virtual + * NSG=4 cohorts per threadgroup, each an exact replica of + * kernel_mul_mv_q8_0_f32_pair (same per-lane K walk and reduction tree, cf. + * kernel_dsv4_router_shared_gate_up_q8_0); the compressor ranges run + * kernel_mul_mv_f16_f32_pair_4_impl<2> and the paired store epilogue + * verbatim, so every output bit matches the two separate dispatches. */ +kernel void kernel_dsv4_qkv_pair_quad_compressor_store_q8_0( + constant ds4_metal_args_mul_mv & args0, + constant ds4_metal_args_mul_mv & args1, + constant ds4_metal_args_mul_mv & cargs, + constant ds4_metal_args_compressor_pair_store & store0, + constant ds4_metal_args_compressor_pair_store & store1, + constant uint & pair_vtgs, + device const char * qw0, + device const char * qw1, + device const char * cw0a, + device const char * cw0b, + device const char * cw1a, + device const char * cw1b, + device const char * src1, + device char * dst0, + device char * dst1, + device char * cdst_a0, + device char * cdst_b0, + device char * cdst_a1, + device char * cdst_b1, + device const char * ape0, + device const char * ape1, + device float * state0_kv, + device float * state0_score, + device float * state1_kv, + device float * state1_score, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr short NW = N_SIMDWIDTH; + const uint pair_ctgs = (pair_vtgs + 1u) / 2u; + + if (tgpig.x < pair_ctgs) { + /* Q8 pair range: cohort c of threadgroup t runs virtual pair + * threadgroup 2t+c with the original NSG=4 mapping. */ + constexpr short NSG = 4; + constexpr short NQ = 8; + constexpr short NR0 = 2; + const uint cohort = sgitg >> 2; + const ushort vsg = sgitg & 3u; + const uint vt = tgpig.x * 2u + cohort; + const bool valid = vt < pair_vtgs; + + const int r0 = vt * NR0; + const bool active_a = valid && r0 < args0.ne01; + const bool active_b = valid && r0 < args1.ne01; + const int nb = args0.ne00 / QK8_0; + + device const float *y = (device const float *)src1; + device const block_q8_0 *ax_a[NR0]; + device const block_q8_0 *ax_b[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const int out_row = r0 + row; + ax_a[row] = active_a && out_row < args0.ne01 + ? (device const block_q8_0 *)(qw0 + (uint64_t)out_row * args0.nb01) + : (device const block_q8_0 *)qw0; + ax_b[row] = active_b && out_row < args1.ne01 + ? (device const block_q8_0 *)(qw1 + (uint64_t)out_row * args1.nb01) + : (device const block_q8_0 *)qw1; + } + + float suma[NR0] = { 0.f }; + float sumb[NR0] = { 0.f }; + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = vsg * NQ + ix; + float yl[NQ]; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + if (valid) { + 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) { + const int out_row = r0 + row; + if (active_a && out_row < args0.ne01) { + device const int8_t *qs = ax_a[row][ib].qs + il * NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + suma[row] += sumq * ax_a[row][ib].d; + } + if (active_b && out_row < args1.ne01) { + device const int8_t *qs = ax_b[row][ib].qs + il * NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + sumb[row] += sumq * ax_b[row][ib].d; + } + } + yb += NSG * NQ * QK8_0; + } + } + + threadgroup float *shared = + (threadgroup float *)shmem + cohort * (2 * NR0 * NW); + threadgroup float *sha[NR0]; + threadgroup float *shb[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sha[row] = shared + NW * row; + shb[row] = shared + NW * (NR0 + row); + if (vsg == 0) { + sha[row][tiisg] = 0.0f; + if (active_b) shb[row][tiisg] = 0.0f; + } + suma[row] = simd_sum(suma[row]); + if (active_b) sumb[row] = simd_sum(sumb[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + sha[row][vsg] = suma[row]; + if (active_b) shb[row][vsg] = sumb[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out_a = (device float *)dst0; + device float *out_b = (device float *)dst1; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const float total_a = simd_sum(sha[row][tiisg]); + if (tiisg == 0 && vsg == 0) { + const int out_row = r0 + row; + if (active_a && out_row < args0.ne01) out_a[out_row] = total_a; + } + if (active_b) { + const float total_b = simd_sum(shb[row][tiisg]); + if (tiisg == 0 && vsg == 0) { + const int out_row = r0 + row; + if (out_row < args1.ne01) out_b[out_row] = total_b; + } + } + } + return; + } + + /* Compressor quad range: verbatim body of + * kernel_mul_mv_f16_f32_quad_compressor_store_4 on the shifted grid. */ + constexpr short NR0 = 2; + const uint lx = tgpig.x - pair_ctgs; + const uint tgs0 = ((uint)store0.width + NR0 - 1u) / NR0; + const bool second = lx >= tgs0; + + uint3 local_tgpig = tgpig; + local_tgpig.x = second ? lx - tgs0 : lx; + + ds4_metal_args_mul_mv largs = cargs; + largs.nr0 = NR0; + largs.ne01 = second ? (int32_t)store1.width : (int32_t)store0.width; + + if (!second) { + kernel_mul_mv_f16_f32_pair_4_impl( + largs, cw0a, cw0b, src1, cdst_a0, cdst_b0, + shmem, local_tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_f16_f32_pair_4_impl( + largs, cw1a, cw1b, src1, cdst_a1, cdst_b1, + shmem, local_tgpig, tiisg, sgitg); + } + + threadgroup_barrier(mem_flags::mem_device); + + // State append: identical to the paired store kernel, scoped to the + // range this threadgroup just projected (its own outputs only). + constant ds4_metal_args_compressor_pair_store & store = second ? store1 : store0; + if (tiitg >= NR0 || store.width == 0u || store.ratio == 0u) { + return; + } + const uint col = local_tgpig.x * (uint)NR0 + tiitg; + if (col >= store.width) return; + + const uint pos_mod = store.pos % store.ratio; + const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod; + const uint dst = dst_row * store.width + col; + const uint ape_i = pos_mod * store.width + col; + + device volatile const float * projected_kv = second + ? (device volatile const float *)cdst_a1 + : (device volatile const float *)cdst_a0; + device volatile const float * projected_score = second + ? (device volatile const float *)cdst_b1 + : (device volatile const float *)cdst_b0; + device const char * ape = second ? ape1 : ape0; + device float * state_kv = second ? state1_kv : state0_kv; + device float * state_score = second ? state1_score : state0_score; + + float ape_v; + if (store.ape_type == 1u) { + ape_v = (float)(((device const half *)ape)[ape_i]); + } else { + ape_v = ((device const float *)ape)[ape_i]; + } + + state_kv[dst] = projected_kv[col]; + state_score[dst] = projected_score[col] + ape_v; +} + template void kernel_mul_mv_t_t_short_impl( args_t args, @@ -1476,125 +1901,6 @@ constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]]; constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]]; #ifdef DS4_METAL_HAS_TENSOR -template< - short NR0, short NR1, - typename SA, typename SA_4x4, typename block_q, short nl, - void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &), - typename T0, typename T0_4x4, typename T1> -kernel void kernel_mul_mm_mpp( - constant ds4_metal_args_mul_mm & args, - device const char * srcA, - device const char * srcB, - device char * dst, - threadgroup char * shmem [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiitg [[thread_index_in_threadgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - (void) sgitg; - - constexpr int NK = 32; - constexpr int NL = NK/16; - constexpr int NUM_THREADS = 128; - - const int K = args.ne00; - const int M = args.ne0; - const int N = args.ne1; - const int im = tgpig.z; - const int i12 = im%args.ne12; - const int i13 = im/args.ne12; - const int r0 = tgpig.y*NR0; - const int r1 = tgpig.x*NR1; - - const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; - - threadgroup SA *sa = (threadgroup SA *)shmem; - threadgroup SA *sb = sa + NR0*NK; - auto tA = tensor(sa, dextents(NK, NR0)); - auto tB = tensor(sb, dextents(NK, NR1)); - - device const T1 *ptrB = (device const T1 *)(srcB + args.nb12*i12 + args.nb13*i13); - const int strideB = args.nb11/sizeof(T1); - - matmul2d< - matmul2d_descriptor(NR1, NR0, NK, false, true, false, - matmul2d_descriptor::mode::multiply_accumulate), - execution_simdgroups<4>> mm; - - auto cT = mm.template get_destination_cooperative_tensor(); - - #pragma unroll - for (uint16_t i = 0; i < cT.get_capacity(); ++i) { - if (cT.is_valid_element(i)) { - cT[i] = 0.0f; - } - } - - for (int loop_k = 0; loop_k < K; loop_k += NK) { - for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { - const int row = work/NL; - const int k_chunk = work%NL; - const int k_pos = loop_k + k_chunk*16; - const short k_base = k_chunk*16; - - if (!FC_mul_mm_bc_out || r0 + row < M) { - if (is_same::value && FC_mul_mm_bc_inp) { - device const T0 *row_ptr = (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0); - FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (k_pos + i < K) ? (SA)row_ptr[k_pos + i] : (SA)0; - } - } else { - const int block_idx = k_pos/(16*nl); - const short il = (k_pos/16)%nl; - device const block_q *row_ptr = (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0); - - SA_4x4 temp_a; - dequantize_func(row_ptr + block_idx, il, temp_a); - FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0; - } - } - } else { - FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (SA)0; - } - } - } - for (int work = tiitg; work < NK*NR1; work += NUM_THREADS) { - const int col = work/NK; - const int k = work%NK; - if ((!FC_mul_mm_bc_out && !FC_mul_mm_bc_inp) || - (r1 + col < N && loop_k + k < K)) { - sb[col*NK + k] = (SA)ptrB[(uint64_t)(r1 + col)*strideB + loop_k + k]; - } else { - sb[col*NK + k] = (SA)0; - } - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - auto mA = tA.slice(0, 0); - auto mB = tB.slice(0, 0); - mm.run(mB, mA, cT); - - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - device float *dst_batch = (device float *)dst + im*N*M; - if (!FC_mul_mm_bc_out) { - device float *dst_tile = dst_batch + r0 + (uint64_t)r1*M; - auto tD = tensor(dst_tile, dextents(NR0, NR1), array({1, M})); - cT.store(tD); - } else { - auto tD = tensor(dst_batch, dextents(M, N), array({1, M})); - auto mD = tD.slice(r0, r1); - cT.store(mD); - } -} - -typedef decltype(kernel_mul_mm_mpp<64, 32, half, half4x4, float4x4, 1, dequantize_f32, float, float4x4, float>) mul_mm_mpp_t; - -template [[host_name("kernel_mul_mm_f16_f32_mpp")]] kernel mul_mm_mpp_t kernel_mul_mm_mpp<64, 32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; - // Retained Metal4/TensorOps dense prefill kernel. The legacy MPP prototype // staged both operands in threadgroup memory; this version stages only the // model weight tile and lets MPP read the dense RHS activation matrix directly @@ -2144,242 +2450,6 @@ kernel void kernel_mul_mm_f16_f32_scaled( } } -kernel void kernel_mul_mm_f16_f32_pair( - constant ds4_metal_args_mul_mm & args, - device const char * src0_a, - device const char * src0_b, - device const char * src1, - device char * dst_a, - device char * dst_b, - threadgroup char * shmem [[threadgroup(0)]], - uint3 tgpig[[threadgroup_position_in_grid]], - ushort tiitg[[thread_index_in_threadgroup]], - ushort sgitg[[simdgroup_index_in_threadgroup]]) { - threadgroup half * sa_a = (threadgroup half *)(shmem); - threadgroup half * sa_b = (threadgroup half *)(shmem + 4096); - threadgroup half * sb = (threadgroup half *)(shmem + 8192); - - constexpr int NR0 = 64; - constexpr int NR1 = 32; - constexpr int NK = 32; - constexpr int NL0 = NK/16; - constexpr int NL1 = NK/8; - - const int im = tgpig.z; - const int r0 = tgpig.y*NR0; - const int r1 = tgpig.x*NR1; - - const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; - const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1; - - const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; - const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; - - const short il0 = (tiitg % NL0); - short il = il0; - - const int i12 = im%args.ne12; - const int i13 = im/args.ne12; - - const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; - const short offset1 = il0; - - device const half4x4 * xa = (device const half4x4 *)(src0_a + args.nb01*(r0 + lr0) + offset0) + offset1; - device const half4x4 * xb = (device const half4x4 *)(src0_b + args.nb01*(r0 + lr0) + offset0) + offset1; - - const short iy = 8*(tiitg % NL1); - - device const float * y = (device const float *)(src1 - + args.nb13*i13 - + args.nb12*i12 - + args.nb11*(r1 + lr1) - + args.nb10*iy); - - simdgroup_half8x8 ma[4]; - simdgroup_half8x8 mb[2]; - - simdgroup_float8x8 mc_a[8]; - simdgroup_float8x8 mc_b[8]; - - for (short i = 0; i < 8; i++) { - mc_a[i] = make_filled_simdgroup_matrix(0.f); - mc_b[i] = make_filled_simdgroup_matrix(0.f); - } - - for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { - half4x4 temp_a; - half4x4 temp_b; - dequantize_f16(xa, il, temp_a); - dequantize_f16(xb, il, temp_b); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - FOR_UNROLL (short i = 0; i < 16; i++) { - const short sx = 2*il0 + i/8; - const short sy = (tiitg/NL0)/8; - - const short lx = (tiitg/NL0)%8; - const short ly = i%8; - - const short ib = 8*sx + sy; - - *(sa_a + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; - *(sa_b + 64*ib + 8*ly + lx) = temp_b[i/4][i%4]; - } - - if (FC_mul_mm_bc_inp) { - for (short i = 0; i < 8; ++i) { - const short sx = (tiitg%NL1); - const short sy = (tiitg/NL1)/8; - - const short lx = i; - const short ly = (tiitg/NL1)%8; - - const short ib = 4*sx + sy; - - *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (half) *((device float *) y + i) : 0; - } - } else { - const short sx = (tiitg%NL1); - const short sy = (tiitg/NL1)/8; - - const short ly = (tiitg/NL1)%8; - - const short ib = 4*sx + sy; - - *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)(*((device float2x4 *) y)); - } - - il = (il + 2 < 1) ? il + 2 : il % 2; - xa = (il < 2) ? xa + 2 : xa; - xb = (il < 2) ? xb + 2 : xb; - - y += NK; - - threadgroup_barrier(mem_flags::mem_threadgroup); - - threadgroup const half * lsma_a = (sa_a + 4*64*(sgitg%2)); - threadgroup const half * lsma_b = (sa_b + 4*64*(sgitg%2)); - threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); - - FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 2; i++) { - simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 4; i++) { - simdgroup_load(ma[i], lsma_a + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 8; i++) { - simdgroup_multiply_accumulate(mc_a[i], mb[i/4], ma[i%4], mc_a[i]); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 4; i++) { - simdgroup_load(ma[i], lsma_b + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 8; i++) { - simdgroup_multiply_accumulate(mc_b[i], mb[i/4], ma[i%4], mc_b[i]); - } - - lsma_a += 8*64; - lsma_b += 8*64; - lsmb += 4*64; - } - } - - if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) { - device float * C_a = (device float *) dst_a + - (r0 + 32*(sgitg & 1)) + - (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; - device float * C_b = (device float *) dst_b + - (r0 + 32*(sgitg & 1)) + - (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; - - for (short i = 0; i < 8; i++) { - simdgroup_store(mc_a[i], C_a + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); - simdgroup_store(mc_b[i], C_b + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); - } - } else { - threadgroup_barrier(mem_flags::mem_threadgroup); - - threadgroup float * temp_str = (threadgroup float *) shmem; - - for (short i = 0; i < 8; i++) { - simdgroup_store(mc_a[i], - temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 8*(i%4) + 8*NR0*(i/4), - NR0, - 0, - false); - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - if (sgitg == 0) { - for (int j = tiitg; j < nr1; j += NR1) { - device float * D = (device float *) dst_a + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; - device float4 * D4 = (device float4 *) D; - - threadgroup float * C = temp_str + (j*NR0); - threadgroup float4 * C4 = (threadgroup float4 *) C; - - int i = 0; - for (; i < nr0/4; i++) { - *(D4 + i) = *(C4 + i); - } - - i *= 4; - for (; i < nr0; i++) { - *(D + i) = *(C + i); - } - } - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (short i = 0; i < 8; i++) { - simdgroup_store(mc_b[i], - temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 8*(i%4) + 8*NR0*(i/4), - NR0, - 0, - false); - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - if (sgitg == 0) { - for (int j = tiitg; j < nr1; j += NR1) { - device float * D = (device float *) dst_b + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; - device float4 * D4 = (device float4 *) D; - - threadgroup float * C = temp_str + (j*NR0); - threadgroup float4 * C4 = (threadgroup float4 *) C; - - int i = 0; - for (; i < nr0/4; i++) { - *(D4 + i) = *(C4 + i); - } - - i *= 4; - for (; i < nr0; i++) { - *(D + i) = *(C + i); - } - } - } - } -} - typedef decltype(kernel_mul_mm) mul_mm_t; // Host-visible prefill matmul variants for F16 and Q8_0 weights. diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index af5c0e9..c467946 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -21,20 +21,6 @@ struct ds4_metal_args_dsv4_hc_weighted_sum { uint64_t nb1; }; -struct ds4_metal_args_dsv4_hc_weighted_sum_norm { - int64_t n_embd; - int64_t n_hc; - int64_t n_tokens; - uint64_t nb_x0; - uint64_t nb_x1; - uint64_t nb_x2; - uint64_t nb_w0; - uint64_t nb_w1; - uint64_t nb0; - uint64_t nb1; - uint64_t nb_norm1; - float norm_eps; -}; struct ds4_metal_args_dsv4_output_hc_weights4 { float post_scale; @@ -411,6 +397,68 @@ kernel void kernel_dsv4_hc_split_weighted_sum( // kernel_dsv4_hc_split_weighted_sum, stores the HC-pre row for diagnostics, and // reuses the just-collapsed values from threadgroup memory for the RMSNorm // reduction. +static __attribute__((always_inline)) inline void ds4_hc_comb_weights4_exact( + constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args, + device volatile const float *mix, + device const float *scale, + device const float *base, + device float *out) { + const float epsv = args.eps; + const float comb_scale = scale[2]; + + float4 r0 = + *((device volatile const float4 *)(mix + 8)) * comb_scale + + *((device const float4 *)(base + 8)); + float4 r1 = + *((device volatile const float4 *)(mix + 12)) * comb_scale + + *((device const float4 *)(base + 12)); + float4 r2 = + *((device volatile const float4 *)(mix + 16)) * comb_scale + + *((device const float4 *)(base + 16)); + float4 r3 = + *((device volatile 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 (int iter = 1; iter < args.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; +} + kernel void kernel_dsv4_hc_split_weighted_sum_norm4( constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args, device const char * mixes, @@ -519,7 +567,6 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4( } threadgroup_barrier(mem_flags::mem_threadgroup); - float sumf = 0.0f; for (uint i = tid; i < n4; i += ntg) { device const float4 *x0 = (device const float4 *)(x + 0 * args.nb_x1 + (uint64_t)row * args.nb_x2); @@ -884,6 +931,119 @@ kernel void kernel_dsv4_q8_hc_expand4_q8_0( } } +kernel void kernel_dsv4_q8_hc_expand4_q8_0_vec_hc( + constant ds4_metal_args_mul_mv & mv, + constant ds4_metal_args_dsv4_hc_expand & hc, + device const char * weight, + device const char * input, + device char * block_out, + device const char * residual, + device const char * post, + device const char * comb, + 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]]) { + if (hc.n_hc != 4 || hc.n_tokens != 1) { + return; + } + + 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 *)(input); + 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) { + const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01; + ax[row] = (device const block_q8_0 *)(weight + off0); + } + + 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); + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + const int d = row0 + row; + if (d >= mv.ne01) { + continue; + } + + const float block_v = simd_sum(shmem_f32[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + *((device float *)(block_out + (uint64_t)d * sizeof(float))) = block_v; + + const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1)); + const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1)); + const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1)); + const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1)); + + const float4 post4 = *((device const float4 *)post); + const float4 comb0 = *((device const float4 *)(comb + 0 * hc.nb_comb1)); + const float4 comb1 = *((device const float4 *)(comb + 1 * hc.nb_comb1)); + const float4 comb2 = *((device const float4 *)(comb + 2 * hc.nb_comb1)); + const float4 comb3 = *((device const float4 *)(comb + 3 * hc.nb_comb1)); + float4 acc = block_v * post4; + acc += comb0 * r0; + acc += comb1 * r1; + acc += comb2 * r2; + acc += comb3 * r3; + FOR_UNROLL (short dst_hc = 0; dst_hc < 4; ++dst_hc) { + *((device float *)(dst + (uint64_t)d * hc.nb0 + + (uint64_t)dst_hc * hc.nb1)) = acc[dst_hc]; + } + } + } +} + // Reduces HC channels to a normal embedding row with the learned pre weights. // This is the input adapter before the attention block and before the FFN block. kernel void kernel_dsv4_hc_weighted_sum( @@ -910,76 +1070,6 @@ kernel void kernel_dsv4_hc_weighted_sum( *((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc; } -// The one-row output head immediately applies a learned RMSNorm after reducing -// its four HC streams. Preserve the standalone scalar HC accumulation, write -// the collapsed row for diagnostics, then reload its F32 values from -// threadgroup memory using the standalone RMSNorm's float4 reduction mapping. -kernel void kernel_dsv4_hc_weighted_sum_norm4( - constant ds4_metal_args_dsv4_hc_weighted_sum_norm & args, - device const char * x, - device const char * weights, - device char * dst, - device const char * norm_weight, - device char * norm_dst, - 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]]) { - if (args.n_tokens != 1 || args.n_hc != 4 || - args.n_embd <= 0 || (args.n_embd & 3) != 0) { - return; - } - - const uint n_embd = uint(args.n_embd); - const uint n4 = n_embd >> 2; - threadgroup float *row_shmem = shared; - threadgroup float *sum_shmem = shared + n_embd; - - if (sgitg == 0) { - sum_shmem[tiisg] = 0.0f; - } - - for (uint d = tid; d < n_embd; d += ntg) { - float acc = 0.0f; - for (int64_t h = 0; h < args.n_hc; ++h) { - const float xv = *((device const float *)( - x + (uint64_t)d*args.nb_x0 + (uint64_t)h*args.nb_x1)); - const float wv = *((device const float *)( - weights + (uint64_t)h*args.nb_w0)); - acc += xv * wv; - } - row_shmem[d] = acc; - *((device float *)(dst + (uint64_t)d*args.nb0)) = acc; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - threadgroup const float4 *row4 = - (threadgroup const float4 *)row_shmem; - float sumf = 0.0f; - for (uint i = tid; i < n4; i += ntg) { - sumf += dot(row4[i], row4[i]); - } - sumf = simd_sum(sumf); - - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tiisg == 0) { - sum_shmem[sgitg] = sumf; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - sumf = sum_shmem[tiisg]; - sumf = simd_sum(sumf); - - const float mean = sumf/args.n_embd; - const float scale = 1.0f/sqrt(mean + args.norm_eps); - device const float4 *w4 = (device const float4 *)norm_weight; - device float4 *norm4 = (device float4 *)norm_dst; - for (uint i = tid; i < n4; i += ntg) { - norm4[i] = (row4[i]*scale)*w4[i]; - } -} - // The one-row HC=4 output head historically materializes four device-F32 // stages across separate launches. Collapse those launches into one tiny // two-thread group while preserving the scalar/vector lane mapping and every @@ -1015,3 +1105,440 @@ kernel void kernel_dsv4_output_hc_weights4( args.post_scale * x + args.eps; } } + + +struct ds4_metal_args_hc_norm_mix { + int32_t n; + int32_t out_dim; + float eps; +}; + +// Fused unweighted RMSNorm + F16 HC-mix projection for DS4 decode HC-pre. +// The standalone decode path runs kernel_rms_norm_f32_4 over the flattened +// 4*embd HC row (1024 threads, one threadgroup) and then +// kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2) over the normalized row. Both +// stages are reproduced bit-exactly in one dispatch: every threadgroup +// redundantly recomputes the norm partials with the original 1024-thread +// mapping (each real lane covers one virtual thread of each 256-thread +// 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. +kernel void kernel_dsv4_hc_rms_norm_mix_f16( + constant ds4_metal_args_hc_norm_mix & args, + device const char * x, + device const char * weight, + 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 NSG = 8; // ds4_gpu_make_plain_mv_dispatch(16384) + constexpr short NW = N_SIMDWIDTH; + constexpr short NR0 = 2; // plain mv nr0 + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + constexpr uint VTHREADS = 1024u; // rms norm threads at n == 16384 + constexpr short VSLICES = VTHREADS/(NSG*NW); // virtual 256-thread slices + + const uint n = (uint)args.n; + const uint n4 = n >> 2; + + device const float4 *x4 = (device const float4 *)x; + + threadgroup float *norm_shmem = (threadgroup float *)shmem; // NW slots + threadgroup float *mv_shmem = (threadgroup float *)shmem + NW; // NW*NR0 slots + + // Phase A: exact replica of kernel_rms_norm_f32_4's reduction tree with + // the 1024 virtual threads folded onto this threadgroup's 8 simdgroups. + for (short v = 0; v < VSLICES; ++v) { + const uint vt = (uint)(sgitg + NSG*v)*NW + tiisg; + float sumf = 0.0f; + for (uint i00 = vt; i00 < n4; i00 += VTHREADS) { + sumf += dot(x4[i00], x4[i00]); + } + sumf = simd_sum(sumf); + if (tiisg == 0) { + norm_shmem[sgitg + NSG*v] = sumf; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + float total = norm_shmem[tiisg]; + total = simd_sum(total); + const float mean = total/(float)args.n; + const float scale = 1.0f/sqrt(mean + args.eps); + + // Phase B: exact replica of kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2) with + // the normalized operand recomputed as x*scale instead of reloaded. + const int nb = args.n/NB; + const int r0 = tgpig.x*NR0; + + device const half4 * ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const half4 *) + (weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half)); + } + + float sumf_mv[NR0] = { 0.f }; + + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = sgitg*NF + ix; + + for (int ib = ib0; ib < nb; ib += NSG*NF) { + float4 yl4[NF4]; + FOR_UNROLL (short i = 0; i < NF4; ++i) { + yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale; + } + + FOR_UNROLL (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]), yl4[i]); + } + + sumf_mv[row] += sumq; + } + } + + // n == 16384 makes the scalar tail loop of the original empty. + device float * dst_f32 = (device float *) dst; + helper_mv_reduce_and_write(dst_f32, sumf_mv, r0, args.out_dim, + tiisg, sgitg, (threadgroup char *)mv_shmem); +} + +// M5 specialization: pack two exact NR0=2 HC-mix producer groups into one +// 512-thread group. Two independent eight-simdgroup clusters retain the +// matvec reductions while the exact RMS scale is redundantly formed six, +// rather than twelve, times. +kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2( + constant ds4_metal_args_hc_norm_mix & args, + device const char * x, + device const char * weight, + 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 NSG_CLUSTER = 8; + constexpr short NCLUSTER = 2; + constexpr short NSG_TOTAL = NSG_CLUSTER * NCLUSTER; + constexpr short NW = N_SIMDWIDTH; + constexpr short NR0 = 2; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + constexpr uint VTHREADS = 1024u; + constexpr short VSLICES = VTHREADS/(NSG_TOTAL*NW); + + const uint n = (uint)args.n; + const uint n4 = n >> 2; + device const float4 *x4 = (device const float4 *)x; + threadgroup float *norm_shmem = (threadgroup float *)shmem; + threadgroup float *mv_shmem = norm_shmem + NW; + + // Exact 1024-virtual-thread RMS reduction, now folded two ways over + // the 16 physical simdgroups instead of four ways over eight. + for (short v = 0; v < VSLICES; ++v) { + const uint vt = (uint)(sgitg + NSG_TOTAL*v)*NW + tiisg; + float sumf = 0.0f; + for (uint i00 = vt; i00 < n4; i00 += VTHREADS) { + sumf += dot(x4[i00], x4[i00]); + } + sumf = simd_sum(sumf); + if (tiisg == 0) { + norm_shmem[sgitg + NSG_TOTAL*v] = sumf; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float total = norm_shmem[tiisg]; + total = simd_sum(total); + const float mean = total/(float)args.n; + const float scale = 1.0f/sqrt(mean + args.eps); + + // Two independent eight-simdgroup clusters reproduce two original + // NR0=2 matvec threadgroups inside this 512-thread threadgroup. + const short cluster = sgitg / NSG_CLUSTER; + const short local_sg = sgitg - cluster*NSG_CLUSTER; + const int nb = args.n/NB; + const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0; + + device const half4 *ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const half4 *) + (weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half)); + } + + float sumf_mv[NR0] = { 0.f }; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = local_sg*NF + ix; + for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) { + float4 yl4[NF4]; + FOR_UNROLL (short i = 0; i < NF4; ++i) { + yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale; + } + FOR_UNROLL (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]), yl4[i]); + } + sumf_mv[row] += sumq; + } + } + + threadgroup float *cluster_shmem[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + cluster_shmem[row] = mv_shmem + + ((uint)cluster*NR0 + row)*NW; + if (local_sg == 0) { + cluster_shmem[row][tiisg] = 0.0f; + } + sumf_mv[row] = simd_sum(sumf_mv[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + cluster_shmem[row][local_sg] = sumf_mv[row]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *mixes_f32 = (device float *)dst; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const float tot = simd_sum(cluster_shmem[row][tiisg]); + if (tiisg == 0 && local_sg == 0 && r0 + row < args.out_dim) { + mixes_f32[r0 + row] = tot; + } + } +} + +kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm( + constant ds4_metal_args_hc_norm_mix & args, + constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & split_args, + device const char * x, + device const char * weight, + device char * dst, + device const float * hc_scale, + device const float * hc_base, + device char * split, + device char * collapse_dst, + device const char * norm_weight, + device char * norm_dst, + device atomic_uint * completion, + 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 NSG_CLUSTER = 8; + constexpr short NCLUSTER = 2; + constexpr short NSG_TOTAL = NSG_CLUSTER * NCLUSTER; + constexpr short NW = N_SIMDWIDTH; + constexpr short NR0 = 2; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + constexpr uint VTHREADS = 1024u; + constexpr short VSLICES = VTHREADS/(NSG_TOTAL*NW); + + const uint n = (uint)args.n; + const uint n4 = n >> 2; + device const float4 *x4 = (device const float4 *)x; + threadgroup float *norm_shmem = (threadgroup float *)shmem; + threadgroup float *mv_shmem = norm_shmem + NW; + + // Exact 1024-virtual-thread RMS reduction, now folded two ways over + // the 16 physical simdgroups instead of four ways over eight. + for (short v = 0; v < VSLICES; ++v) { + const uint vt = (uint)(sgitg + NSG_TOTAL*v)*NW + tiisg; + float sumf = 0.0f; + for (uint i00 = vt; i00 < n4; i00 += VTHREADS) { + sumf += dot(x4[i00], x4[i00]); + } + sumf = simd_sum(sumf); + if (tiisg == 0) { + norm_shmem[sgitg + NSG_TOTAL*v] = sumf; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float total = norm_shmem[tiisg]; + total = simd_sum(total); + const float mean = total/(float)args.n; + const float scale = 1.0f/sqrt(mean + args.eps); + + // Two independent eight-simdgroup clusters reproduce two original + // NR0=2 matvec threadgroups inside this 512-thread threadgroup. + const short cluster = sgitg / NSG_CLUSTER; + const short local_sg = sgitg - cluster*NSG_CLUSTER; + const int nb = args.n/NB; + const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0; + + device const half4 *ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const half4 *) + (weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half)); + } + + float sumf_mv[NR0] = { 0.f }; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = local_sg*NF + ix; + for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) { + float4 yl4[NF4]; + FOR_UNROLL (short i = 0; i < NF4; ++i) { + yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale; + } + FOR_UNROLL (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]), yl4[i]); + } + sumf_mv[row] += sumq; + } + } + + threadgroup float *cluster_shmem[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + cluster_shmem[row] = mv_shmem + + ((uint)cluster*NR0 + row)*NW; + if (local_sg == 0) { + cluster_shmem[row][tiisg] = 0.0f; + } + sumf_mv[row] = simd_sum(sumf_mv[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + cluster_shmem[row][local_sg] = sumf_mv[row]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device volatile float *mixes_f32 = + (device volatile float *)dst; + if (local_sg == 0) { + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const float tot = simd_sum(cluster_shmem[row][tiisg]); + if (tiisg == 0 && r0 + row < args.out_dim) { + mixes_f32[r0 + row] = tot; + } + } + } + + // The first producer group owns mix[0:4]. After materializing and + // reloading those values, fold the established 1024-thread HC collapse + // and RMS reduction over this group's 512 physical threads as two + // independent virtual slices. This retains the original 32-partial tree. + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + const uint tid = (uint)sgitg * (uint)NW + (uint)tiisg; + threadgroup float *pre_shmem = norm_shmem + 32u + 4u*NW; + threadgroup float *sum_shmem = pre_shmem + 4; + + if (tgpig.x == 0) { + device float *out = (device float *)split; + if (tid == 0) { + const float4 pre_z = + *((device volatile const float4 *)mixes_f32) * hc_scale[0] + + *((device const float4 *)hc_base); + const float4 pre = + 1.0f / (1.0f + exp(-pre_z)) + split_args.eps; + *((device float4 *)out) = pre; + pre_shmem[0] = pre.x; + pre_shmem[1] = pre.y; + pre_shmem[2] = pre.z; + pre_shmem[3] = pre.w; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint n4_collapse = uint(split_args.n_embd) >> 2; + const uint i0 = tid; + const uint i1 = tid + 512u; + device const float4 *x0 = (device const float4 *)( + x + 0 * split_args.nb_x1); + device const float4 *x1 = (device const float4 *)( + x + 1 * split_args.nb_x1); + device const float4 *x2 = (device const float4 *)( + x + 2 * split_args.nb_x1); + device const float4 *x3 = (device const float4 *)( + x + 3 * split_args.nb_x1); + + float4 v0 = 0.0f; + v0 += x0[i0] * pre_shmem[0]; + v0 += x1[i0] * pre_shmem[1]; + v0 += x2[i0] * pre_shmem[2]; + v0 += x3[i0] * pre_shmem[3]; + float sum0 = simd_sum(dot(v0, v0)); + + float4 v1 = 0.0f; + if (i1 < n4_collapse) { + v1 += x0[i1] * pre_shmem[0]; + v1 += x1[i1] * pre_shmem[1]; + v1 += x2[i1] * pre_shmem[2]; + v1 += x3[i1] * pre_shmem[3]; + } + float sum1 = simd_sum(dot(v1, v1)); + if (tiisg == 0) { + sum_shmem[sgitg] = sum0; + sum_shmem[sgitg + 16] = sum1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float sumf = sum_shmem[tiisg]; + sumf = simd_sum(sumf); + const float norm_arg = + sumf / float(split_args.n_embd) + split_args.norm_eps; + const float norm_scale = rsqrt(norm_arg); + device float4 *dst4 = (device float4 *)collapse_dst; + device const float4 *w4 = (device const float4 *)norm_weight; + device float4 *norm4 = (device float4 *)norm_dst; + dst4[i0] = v0; + norm4[i0] = (v0 * norm_scale) * w4[i0]; + if (i1 < n4_collapse) { + dst4[i1] = v1; + norm4[i1] = (v1 * norm_scale) * w4[i1]; + } + } else if (tgpig.x == 1 && tid == 0) { + device float *out = (device float *)split; + const float4 post_z = + *((device volatile const float4 *)(mixes_f32 + 4)) * hc_scale[1] + + *((device const float4 *)(hc_base + 4)); + *((device float4 *)(out + 4)) = 2.0f / (1.0f + exp(-post_z)); + } + + // Groups 2..5 own exactly the comb range consumed by the + // continuation. Their four-way completion overlaps TG0's independent + // pre-collapse/RMS epilogue. Every writer crosses the uniform publish + // fence; only lane zero then participates in the completion protocol. + atomic_thread_fence(mem_flags::mem_device, + memory_order_seq_cst, + thread_scope_device); + if (tgpig.x < 2 || tid != 0) { + return; + } + + const uint old = atomic_fetch_add_explicit( + completion, 1u, memory_order_relaxed); + if (old + 1u != 4u) { + return; + } + atomic_thread_fence(mem_flags::mem_device, + memory_order_seq_cst, + thread_scope_device); + ds4_hc_comb_weights4_exact( + split_args, mixes_f32, hc_scale, hc_base, + (device float *)split); + atomic_thread_fence(mem_flags::mem_device, + memory_order_seq_cst, + thread_scope_device); + atomic_store_explicit(completion, 0u, memory_order_relaxed); +} diff --git a/metal/dsv4_kv.metal b/metal/dsv4_kv.metal index 4f44a43..dd602ce 100644 --- a/metal/dsv4_kv.metal +++ b/metal/dsv4_kv.metal @@ -319,6 +319,228 @@ kernel void kernel_dsv4_compressor_pack_ratio4( } } +// Decode already holds the complete ratio-4 recurrent window in state layout: +// eight rows of two head_dim planes. Pack the previous plane from rows 0..3 +// and the current plane from rows 4..7 directly into the transposed [head_dim, +// 8] layout consumed by the exact GGML softmax/multiply/sum sequence. KV and +// score move together; no arithmetic or reduction order changes. +kernel void kernel_dsv4_compressor_pack_ratio4_decode_ggml( + constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args, + device const uint * state_kv, + device const uint * state_score, + device uint * packed_kv, + device uint * packed_score, + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + if (row >= 8u || args.head_dim == 0u || args.n_threads == 0u) { + return; + } + + const uint64_t state_row_stride = 2ull * args.head_dim; + const uint64_t src_plane = row >= 4u ? args.head_dim : 0u; + for (uint col = tid; col < args.head_dim; col += args.n_threads) { + const uint64_t src = (uint64_t)row * state_row_stride + + src_plane + col; + const uint64_t dst = (uint64_t)col * 8u + row; + packed_kv[dst] = state_kv[src]; + packed_score[dst] = state_score[src]; + } +} + +// Exact decode specialization for the first two operations in GGML's +// softmax -> multiply -> sum_rows compressor reduction. The normalized +// softmax values are deliberately materialized in device memory and reloaded +// after a device barrier before the in-place product, preserving the dispatch +// boundary's float store/load semantics. The final sum remains the standalone +// eight-thread sum_rows kernel: changing a 32-thread group into the original +// eight-thread reduction inside this kernel would make its threadgroup +// barriers non-uniform or alter simd_sum's active-lane topology. +kernel void kernel_dsv4_compressor_exact_softmax_product_ratio4( + constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args, + device const float * packed_kv, + device const float * packed_score, + device float * softmax, + device float * product, + threadgroup float * softmax_scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + if (row >= args.head_dim || args.n_comp != 1u || + args.n_threads != 32u) { + return; + } + + device const float4 * score4 = + (device const float4 *)(packed_score + (uint64_t)row * 8u); + device float4 * softmax4 = + (device float4 *)(softmax + (uint64_t)row * 8u); + const float scale = (float)args.replay; + const float zero = (float)(args.n_comp - 1u); + + // Match kernel_soft_max_f32_4(width=8, nth=32) literally. Only lanes zero + // and one own float4s, while all 32 lanes participate in both reductions. + float4 lmax4 = -INFINITY; + for (int i00 = (int)tid; i00 < 2; i00 += 32) { + lmax4 = fmax(lmax4, score4[i00] * scale + (float4)zero); + } + + const float lmax = + MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3])); + const float max_val = simd_max(lmax); + + float4 lsum4 = 0.0f; + for (int i00 = (int)tid; i00 < 2; i00 += 32) { + const float4 exp_score4 = + exp((score4[i00] * scale + (float4)zero) - max_val); + lsum4 += exp_score4; + softmax4[i00] = exp_score4; + } + + const float lsum = + lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3]; + threadgroup_barrier(mem_flags::mem_none); + const float sum = simd_sum(lsum); + const float inv_sum = 1.0f / sum; + + for (int i00 = (int)tid; i00 < 2; i00 += 32) { + softmax4[i00] *= inv_sum; + } + + // Force the same normalized-softmax device store/reload boundary that the + // separate multiply dispatch observes. + threadgroup_barrier(mem_flags::mem_device); + device volatile const float * reloaded_softmax = + (device volatile const float *)(softmax + (uint64_t)row * 8u); + device const float * kv_row = packed_kv + (uint64_t)row * 8u; + device float * product_row = product + (uint64_t)row * 8u; + + // Match kernel_bin_fuse_f32_f32_f32(width=8, nth=4): four lanes each + // process their low element followed by the element four positions later. + if (tid < 4u) { + for (uint i0 = tid; i0 < 8u; i0 += 4u) { + float value = kv_row[i0]; + value *= reloaded_softmax[i0]; + product_row[i0] = value; + } + } + + // All 32 lanes reach the final device barrier. The following standalone + // sum_rows dispatch performs the required global reload and exact TG8 + // two-stage simd_sum topology. + threadgroup_barrier(mem_flags::mem_device); + (void)softmax_scratch; +} + +// Exact one-dispatch ratio-4 decode pool. This specializes the three-dispatch +// pack -> exact softmax/product -> sum_rows chain above without changing any +// floating-point operation or reduction topology. The normalized softmax and +// product are still materialized and volatile-reloaded through device memory. +// The two simd_sum calls in the final reduction execute under an eight-lane +// active mask, exactly matching kernel_sum_rows_f32_f32's original TG8. +kernel void kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml( + constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args, + device const float * state_kv, + device const float * state_score, + device float * softmax, + device float * product, + device float * dst, + threadgroup float * sum_scratch [[threadgroup(0)]], + uint col [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + if (col >= args.head_dim || args.n_comp != 1u || + args.n_threads != 32u) { + return; + } + + const uint64_t state_row_stride = 2ull * args.head_dim; + const float scale = (float)args.replay; + const float zero = (float)(args.n_comp - 1u); + + // Match the packed float4 ownership: lane 0 owns rows 0..3 and lane 1 + // rows 4..7. The gather itself is an integer-addressed bit-preserving load. + float4 score_values = -INFINITY; + if (tid < 2u) { + const uint row0 = 4u * tid; + for (uint j = 0u; j < 4u; ++j) { + const uint row = row0 + j; + const uint64_t src = (uint64_t)row * state_row_stride + + (row >= 4u ? args.head_dim : 0u) + col; + score_values[j] = state_score[src]; + } + } + + const uint64_t scratch_base = (uint64_t)col * 8u; + device float4 * softmax4 = + (device float4 *)(softmax + scratch_base); + + // Verbatim kernel_soft_max_f32_4(width=8, nth=32) arithmetic. + float4 lmax4 = -INFINITY; + for (int i00 = (int)tid; i00 < 2; i00 += 32) { + lmax4 = fmax(lmax4, score_values * scale + (float4)zero); + } + const float lmax = + MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3])); + const float max_val = simd_max(lmax); + + float4 lsum4 = 0.0f; + for (int i00 = (int)tid; i00 < 2; i00 += 32) { + const float4 exp_score4 = + exp((score_values * scale + (float4)zero) - max_val); + lsum4 += exp_score4; + softmax4[i00] = exp_score4; + } + const float lsum = + lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3]; + threadgroup_barrier(mem_flags::mem_none); + const float sum = simd_sum(lsum); + const float inv_sum = 1.0f / sum; + for (int i00 = (int)tid; i00 < 2; i00 += 32) { + softmax4[i00] *= inv_sum; + } + + threadgroup_barrier(mem_flags::mem_device); + device volatile const float * reloaded_softmax = + (device volatile const float *)(softmax + scratch_base); + device float * product_row = product + scratch_base; + + // Verbatim width=8, TG4 multiply ownership: low element, then +4. + if (tid < 4u) { + for (uint i0 = tid; i0 < 8u; i0 += 4u) { + const uint64_t src = (uint64_t)i0 * state_row_stride + + (i0 >= 4u ? args.head_dim : 0u) + col; + float value = state_kv[src]; + value *= reloaded_softmax[i0]; + product_row[i0] = value; + } + } + + // Preserve the product dispatch's device store/reload boundary. + threadgroup_barrier(mem_flags::mem_device); + device volatile const float * reloaded_product = + (device volatile const float *)product_row; + + // Reproduce kernel_sum_rows_f32_f32(width=8, TG8) literally. MSL defines + // simdgroup collectives over active lanes, so the branch recreates the + // original eight-lane partial SIMD group inside this 32-thread group. + sum_scratch[tid] = 0.0f; + float row_sum = 0.0f; + if (tid < 8u) { + row_sum += reloaded_product[tid]; + row_sum = simd_sum(row_sum); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0u) { + sum_scratch[0] = row_sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 8u) { + row_sum = sum_scratch[tid]; + row_sum = simd_sum(row_sum); + if (tid == 0u) { + dst[col] = row_sum; + } + } +} + // Ratio-4 compression keeps two 4-row halves of recurrent state. After an // emitted compressed row, the second half becomes the next window's previous // half. The old encoder expressed this as four generic copies; this DS4-specific diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index b24167a..f24ca11 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -70,7 +70,7 @@ struct ds4_metal_args_dsv4_indexed_attention { uint32_t window; uint32_t ratio; uint32_t comp_kv_f16; - uint32_t pad0; + uint32_t n_splits; uint64_t q_token_stride; uint64_t q_head_stride; uint64_t raw_row_stride; @@ -4915,6 +4915,270 @@ kernel void kernel_dsv4_router_finalize_weights_one_simd( } } +// M3 decode specialization that materializes the probability +// transform in device memory before running the exact SIMD selection and +// weight normalization above. The volatile reload after the device barrier +// pins the same float store/load boundary as the standalone transform dispatch. +kernel void kernel_dsv4_router_transform_finalize_weights_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *logits, + device float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + device float *weights, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + if (tid < 64) { + device const float4 *s = (device const float4 *)logits; + device float4 *d = (device float4 *)probs; + const float4 x = s[tid]; + const float4 sp = select(log(1.0f + exp(x)), x, x > 20.0f); + d[tid] = sqrt(sp); + } + threadgroup_barrier(mem_flags::mem_device); + device volatile const float *reloaded_probs = + (device volatile const float *)probs; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = reloaded_probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } + threadgroup_barrier(mem_flags::mem_device); + + threadgroup volatile float *norm_scratch = + (threadgroup volatile float *)scratch; + if (tid == 0) { + device const int32_t *s = selected; + norm_scratch[0] = 0.0f; + for (uint i = 0; i < 6; i++) { + norm_scratch[0] = + norm_scratch[0] + reloaded_probs[s[i]]; + } + norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f); + norm_scratch[1] = 1.5f / norm_scratch[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 6) { + device const int32_t *s = selected; + weights[tid] = reloaded_probs[s[tid]] * norm_scratch[1]; + } +} + +kernel void kernel_dsv4_router_project_select_fused( + constant ds4_metal_args_mul_mv & args, + constant ds4_metal_args_dsv4_router_select_one & select_args, + device const char * src0_router, + device const char * src1, + device float * logits, + device float * probs, + device const float * bias, + device int32_t * selected, + device float * weights, + device atomic_uint * completion, + threadgroup char * shmem_raw [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + uint3 tpitg [[thread_position_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 uint tid = tpitg.x; + const int nb = args.ne00/NB; + const int r0 = tgpig.x*NR0; + device const float4 *y4 = (device const float4 *)src1; + device const half4 *ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const half4 *) + (src0_router + (uint64_t)(r0 + row)*args.nb01); + } + 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 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]), yl4[i]); + } + sumf[row] += sumq; + } + yb4 += NSG*NF*NW/4; + } + helper_mv_reduce_and_write(logits, sumf, r0, args.ne01, + tiisg, sgitg, shmem_raw); + + threadgroup float *scratch = (threadgroup float *)shmem_raw; + 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 == 127u ? 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); + + if (tid < 64) { + device volatile const float4 *s = + (device volatile const float4 *)logits; + device float4 *d = (device float4 *)probs; + const float4 xv = s[tid]; + const float4 sp = select(log(1.0f + exp(xv)), xv, xv > 20.0f); + d[tid] = sqrt(sp); + } + threadgroup_barrier(mem_flags::mem_device); + device volatile const float *reloaded_probs = + (device volatile const float *)probs; + + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = reloaded_probs[tid]; + float score = select_args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + if (tid < 6) selected[tid] = idx; + threadgroup_barrier(mem_flags::mem_device); + threadgroup volatile float *norm_scratch = + (threadgroup volatile float *)scratch; + if (tid == 0) { + norm_scratch[0] = 0.0f; + for (uint i = 0; i < 6; ++i) { + norm_scratch[0] = norm_scratch[0] + reloaded_probs[selected[i]]; + } + norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f); + norm_scratch[1] = 1.5f / norm_scratch[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 6) { + weights[tid] = reloaded_probs[selected[tid]] * norm_scratch[1]; + } + 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); + } +} + + // Fills the dense compressed-attention mask with -inf. The selected top-k rows // are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch. kernel void kernel_dsv4_topk_mask( @@ -5319,6 +5583,117 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8( dst4[lane + 96] = o3 * inv_s; } +// Each simdgroup owns two heads and updates both from one staged K/V row. +// This doubles row reuse without increasing the 256-thread workgroup. +kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual( + constant ds4_metal_args_dsv4_indexed_attention &args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const uint token = tgpig.x; + const uint head0 = tgpig.y*16u + (uint)sg; + const uint head1 = head0 + 8u; + if (token >= args.n_tokens || head0 >= args.n_head) return; + + device const float4 *qa = (device const float4 *)(q + + (uint64_t)token*args.q_token_stride + + (uint64_t)head0*args.q_head_stride); + half4 qa0 = (half4)qa[lane + 0]; + half4 qa1 = (half4)qa[lane + 32]; + half4 qa2 = (half4)qa[lane + 64]; + half4 qa3 = (half4)qa[lane + 96]; + half4 qb0 = half4(0.0h), qb1 = half4(0.0h); + half4 qb2 = half4(0.0h), qb3 = half4(0.0h); + if (head1 < args.n_head) { + device const float4 *qb = (device const float4 *)(q + + (uint64_t)token*args.q_token_stride + + (uint64_t)head1*args.q_head_stride); + qb0 = (half4)qb[lane + 0]; + qb1 = (half4)qb[lane + 32]; + qb2 = (half4)qb[lane + 64]; + qb3 = (half4)qb[lane + 96]; + } + + float Ma = -FLT_MAX/2.0f, Sa = 0.0f; + float Mb = -FLT_MAX/2.0f, Sb = 0.0f; + float4 ao0 = 0.0f, ao1 = 0.0f, ao2 = 0.0f, ao3 = 0.0f; + float4 bo0 = 0.0f, bo1 = 0.0f, bo2 = 0.0f, bo3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + const uint first = max(first_raw_pos, window_first); + const uint last = min(qpos, raw_last_pos); + if (first <= last) { + for (uint pos = first; pos <= last; pos++) { + const uint logical = pos - first_raw_pos; + const uint row = (args.raw_start + logical)%args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row*args.raw_row_stride); + if (tid < 128) kv_shared[tid] = (half4)src[tid]; + threadgroup_barrier(mem_flags::mem_threadgroup); + dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3, + args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3); + if (head1 < args.n_head) { + dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3, + args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + const uint visible = min((qpos + 1u)/args.ratio, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token*args.topk_token_stride); + for (uint i = 0; i < args.top_k; i++) { + const int32_t idx = row_topk[i]; + if (idx < 0) continue; + if ((uint)idx >= visible) break; + if (tid < 128) { + kv_shared[tid] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, (uint)idx, tid, args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3, + args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3); + if (head1 < args.n_head) { + dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3, + args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head0], + Ma, Sa, ao0, ao1, ao2, ao3); + const float ia = Sa == 0.0f ? 0.0f : 1.0f/Sa; + device float4 *da = (device float4 *)(dst + + (uint64_t)token*args.dst_token_stride + + (uint64_t)head0*args.dst_head_stride); + da[lane + 0] = ao0*ia; da[lane + 32] = ao1*ia; + da[lane + 64] = ao2*ia; da[lane + 96] = ao3*ia; + if (head1 < args.n_head) { + dsv4_attend_sink(((device const float *)sinks)[head1], + Mb, Sb, bo0, bo1, bo2, bo3); + const float ib = Sb == 0.0f ? 0.0f : 1.0f/Sb; + device float4 *db = (device float4 *)(dst + + (uint64_t)token*args.dst_token_stride + + (uint64_t)head1*args.dst_head_stride); + db[lane + 0] = bo0*ib; db[lane + 32] = bo1*ib; + db[lane + 64] = bo2*ib; db[lane + 96] = bo3*ib; + } +} + // Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8. // Generation attends one token at a time, so the ratio-4 indexed path spends a // visible amount of time repeatedly staging the same K/V row for the eight @@ -5450,6 +5825,190 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb16( dst4[lane + 96] = o3 * inv_s; } +// Long-context decode specialization of the indexed mixed-attention path. +// +// The ordinary heads8 kernel reuses each K/V row across eight heads, but only +// launches one threadgroup per head group. Long-context decode therefore has +// too little parallel work while each group scans its raw and selected rows. +// This kernel retains the same eight-head reuse while splitting that row +// sequence across args.n_splits workgroups. A second kernel merges the online +// softmax partials and applies the attention sink. +kernel void kernel_dsv4_indexed_mixed_attention_heads8_split( + constant ds4_metal_args_dsv4_indexed_attention & args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device char *tmp, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint rows_per_block = 16u; + constexpr uint vecs_per_row = 128u; + + const uint token = tgpig.x; + const uint head = tgpig.y * 8u + (uint)sg; + const uint split = tgpig.z; + const uint n_splits = args.n_splits; + if (token >= args.n_tokens || head >= args.n_head || + n_splits < 2u || n_splits > 31u || split >= n_splits) { + return; + } + + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + const half4 q0 = (half4)q4[lane + 0]; + const half4 q1 = (half4)q4[lane + 32]; + const half4 q2 = (half4)q4[lane + 64]; + const half4 q3 = (half4)q4[lane + 96]; + + float M = -FLT_MAX/2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + const uint raw_first = max(first_raw_pos, window_first); + const uint raw_last = min(qpos, raw_last_pos); + const uint raw_count = raw_first <= raw_last ? + raw_last - raw_first + 1u : 0u; + const uint total_rows = raw_count + args.top_k; + const uint rows_per_split = + (total_rows + n_splits - 1u) / n_splits; + const uint split_first = min(split * rows_per_split, total_rows); + const uint split_last = min(split_first + rows_per_split, total_rows); + const uint visible = min((qpos + 1u) / args.ratio, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token * args.topk_token_stride); + + for (uint seq0 = split_first; seq0 < split_last; + seq0 += rows_per_block) { + const uint n_rows = min(rows_per_block, split_last - seq0); + for (uint off = (uint)tid; + off < n_rows * vecs_per_row; + off += 256u) { + const uint r = off / vecs_per_row; + const uint c = off - r * vecs_per_row; + const uint seq = seq0 + r; + half4 value = half4(0.0h); + if (seq < raw_count) { + const uint pos = raw_first + seq; + const uint logical = pos - first_raw_pos; + const uint row = (args.raw_start + logical) % args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row * args.raw_row_stride); + value = (half4)src[c]; + } else { + const int32_t idx = row_topk[seq - raw_count]; + if (idx >= 0 && (uint)idx < visible) { + value = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + (uint)idx, + c, + args.comp_kv_f16 != 0u); + } + } + kv_shared[off] = value; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + const uint seq = seq0 + r; + bool valid = true; + if (seq >= raw_count) { + const int32_t idx = row_topk[seq - raw_count]; + valid = idx >= 0 && (uint)idx < visible; + } + if (valid) { + dsv4_attend_shared_h4_row_at(kv_shared, + r, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const uint64_t n_rows = (uint64_t)args.n_tokens * args.n_head; + const uint64_t row = (uint64_t)token * args.n_head + head; + device float4 *partials = (device float4 *)tmp; + partials[(row * vecs_per_row + lane + 0u) * n_splits + split] = o0; + partials[(row * vecs_per_row + lane + 32u) * n_splits + split] = o1; + partials[(row * vecs_per_row + lane + 64u) * n_splits + split] = o2; + partials[(row * vecs_per_row + lane + 96u) * n_splits + split] = o3; + + if (lane == 0u) { + device float *stats = (device float *)(partials + + n_rows * vecs_per_row * n_splits); + const uint64_t stat = (row * n_splits + split) * 2u; + stats[stat + 0u] = S; + stats[stat + 1u] = M; + } +} + +kernel void kernel_dsv4_indexed_mixed_attention_heads8_split_reduce( + constant ds4_metal_args_dsv4_indexed_attention & args, + device const char *tmp, + device const char *sinks, + device char *dst, + uint tgpig [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint vecs_per_row = 128u; + const uint n_splits = args.n_splits; + const uint64_t n_rows = (uint64_t)args.n_tokens * args.n_head; + const uint64_t row = tgpig; + if (row >= n_rows || n_splits < 2u || n_splits > 31u) { + return; + } + + device const float4 *partials = (device const float4 *)tmp; + device const float *stats = (device const float *)(partials + + n_rows * vecs_per_row * n_splits); + float part_sum = 0.0f; + float part_max = -FLT_MAX/2.0f; + if ((uint)lane < n_splits) { + const uint64_t stat = (row * n_splits + (uint)lane) * 2u; + part_sum = stats[stat + 0u]; + part_max = stats[stat + 1u]; + } else if ((uint)lane == n_splits) { + const uint head = (uint)(row % args.n_head); + part_sum = 1.0f; + part_max = ((device const float *)sinks)[head]; + } + + const float global_max = simd_max(part_max); + const float part_scale = part_sum > 0.0f ? + exp(part_max - global_max) : 0.0f; + const float total_sum = simd_sum(part_sum * part_scale); + const float inv_sum = total_sum > 0.0f ? 1.0f / total_sum : 0.0f; + + device float4 *out = (device float4 *)dst + row * vecs_per_row; + for (uint i = (uint)sg; i < vecs_per_row; i += 4u) { + float4 value = float4(0.0f); + if ((uint)lane < n_splits) { + value = partials[(row * vecs_per_row + i) * n_splits + + (uint)lane] * part_scale; + } + value = simd_sum(value); + if (lane == 0u) { + out[i] = value * inv_sum; + } + } +} + static inline float dsv4_indexer_dot128_shared_q( float4 c0, float4 c1, diff --git a/metal/dsv4_rope.metal b/metal/dsv4_rope.metal index 94a83c7..279365c 100644 --- a/metal/dsv4_rope.metal +++ b/metal/dsv4_rope.metal @@ -41,6 +41,23 @@ struct ds4_metal_args_dsv4_rope_affine_pair { float beta_slow; }; +struct ds4_metal_args_dsv4_head_norm_rope { + int32_t n_head; + int32_t head_dim; + int32_t head_dim4; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t pos0; + int32_t inverse; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + static float rope_yarn_ramp(const float low, const float high, const int i0) { const float y = (i0 / 2 - low) / max(0.001f, high - low); return 1.0f - min(1.0f, max(0.0f, y)); @@ -327,36 +344,109 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4( *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; } -// DS4 positions are always affine within one RoPE dispatch. This variant -// reconstructs the same wrapped int32 position in-kernel, avoiding the host -// position array and its buffer binding while preserving the pair lane mapping -// and all floating-point operations of the specialization above. -kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine( - constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]], - device const char * src0 [[buffer(1)]], - device char * dst [[buffer(4)]], - uint tid [[thread_index_in_threadgroup]], - ushort3 ntg [[threads_per_threadgroup]], - uint3 tgpig [[threadgroup_position_in_grid]]) { - const int i1 = tgpig[0]; - const int i2 = tgpig[1]; +// Fuses the per-head RMSNorm and partial Q RoPE while retaining the standalone +// norm reduction tree and the mode-0 RoPE lane mapping. +kernel void kernel_dsv4_head_rms_norm_rope_tail_f32( + constant ds4_metal_args_dsv4_head_norm_rope & args, + device char * xraw, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort3 tpitg [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + if (sgitg == 0) { + shmem_f32[tiisg] = 0.0f; + } + + const uint head = tgpig.x; + const uint tok = tgpig.y; + device float4 * x4 = (device float4 *)xraw + + ((uint64_t)tok * (uint64_t)args.n_head + head) * + (uint64_t)args.head_dim4; + + float sumf = 0.0f; + for (int i00 = tpitg.x; i00 < args.head_dim4; i00 += ntg.x) { + sumf += dot(x4[i00], x4[i00]); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + shmem_f32[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = simd_sum(shmem_f32[tiisg]); + const float scale = 1.0f / sqrt(sumf / args.head_dim + args.eps); const int n_nope = args.head_dim - args.n_dims; if (n_nope < 0) { return; } + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, + args.beta_fast, args.beta_slow, corr_dims); + const float theta_base = (float)(args.pos0 + (int)tok); + const float inv_ndims = -1.0f / args.n_dims; + device float * xs = (device float *)x4; + + for (int i0 = tpitg.x; i0 < args.head_dim; i0 += ntg.x) { + if (i0 < n_nope) { + xs[i0] = xs[i0] * scale; + continue; + } + const int r = i0 - n_nope; + if ((r & 1) != 0) { + continue; + } +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = + theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = + theta_base * pow(args.freq_base, inv_ndims * r); +#endif + float cos_theta; + float sin_theta; + rope_yarn(theta, args.freq_scale, corr_dims, r, + args.ext_factor, args.attn_factor, + &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const float x0 = xs[i0] * scale; + const float x1 = xs[i0 + 1] * scale; + xs[i0] = x0 * cos_theta - x1 * sin_theta; + xs[i0 + 1] = x0 * sin_theta + x1 * cos_theta; + } +} + +// DS4 positions are always affine within one RoPE dispatch. This variant +// reconstructs the same wrapped int32 position in-kernel, avoiding the host +// position array and its buffer binding while preserving the pair lane mapping +// and all floating-point operations of the specialization above. + +/* Shared, deliberately noinline so that every caller gets bit-identical + * trigonometric codegen. The header note about tiny trig codegen changes + * flipping sampled tokens is exactly why this body must be compiled once and + * shared rather than inlined separately into each kernel. */ +static __attribute__((noinline)) void ds4_rope_tail_pair_affine_row( + constant ds4_metal_args_dsv4_rope_affine_pair & args, + device const char * src_base, + device char * dst_base, + int n_nope, + uint raw_pos, + uint tid, + uint nthreads) { float corr_dims[2]; rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); - const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step; - const float theta_base = (float)as_type(raw_pos); + const float theta_base = (float)as_type(raw_pos); const float inv_ndims = -1.f/args.n_dims; - device const char * src_base = - src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; - device char * dst_base = - dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; - for (int r = tid; r < args.n_dims; r += ntg.x) { + for (int r = tid; r < args.n_dims; r += nthreads) { if ((r & 1) != 0) { continue; } @@ -381,5 +471,413 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine( *((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta; *((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta; + }} + +kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine( + constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]], + device const char * src0 [[buffer(1)]], + device char * dst [[buffer(4)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int n_nope = args.head_dim - args.n_dims; + if (n_nope < 0) { + return; + } + const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step; + device const char * src_base = + src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; + device char * dst_base = + dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; + ds4_rope_tail_pair_affine_row(args, src_base, dst_base, n_nope, raw_pos, tid, ntg.x); + +} + +// Decode-only fusion of the KV RoPE tail with the FP8/raw finalizer. Both were +// already single 64-thread threadgroups on the same row, back to back, so the +// pair cost two dispatches (~12.4 us) to touch 2 KB. The RoPE body below is a +// verbatim copy of kernel_dsv4_rope_tail_f32_inplace_pair_affine specialised to +// the decode grid (one head, one token, so i1 = i2 = 0) and the finalizer body +// is a verbatim copy of kernel_dsv4_kv_fp8_store_f32. The barrier between them +// is required because RoPE writes element pairs across lanes while the raw copy +// reads them per lane. Arithmetic, order and rounding are unchanged; the header +// warning above about trigonometric codegen still applies, so this kernel is +// gated and verified against full-vocabulary logits before promotion. +kernel void kernel_dsv4_kv_rope_fp8_store_f32( + constant ds4_metal_args_dsv4_kv_fp8_store & args, + constant ds4_metal_args_dsv4_rope_affine_pair & rope, + device float * kv, + device float * raw_cache, + threadgroup float * scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]]) { + { + const int rope_n_nope = rope.head_dim - rope.n_dims; + if (rope_n_nope < 0) { + return; + } + ds4_rope_tail_pair_affine_row(rope, + (device const char *)kv, + (device char *)kv, + rope_n_nope, + rope.pos0, + tid, + 64u); + } + /* The RoPE helper writes device-memory pairs that different lanes read + * below. A threadgroup-only fence does not make those cross-lane device + * writes visible. */ + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + { + + const int head_dim = args.head_dim; + const int n_rot = args.n_rot; + const int n_nope = head_dim - n_rot; + if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) { + return; + } + + device float * raw = raw_cache + (int64_t)args.raw_row * head_dim; + + for (int off = 0; off < n_nope; off += 64) { + float v = 0.0f; + if (off + (int)tid < n_nope) { + v = kv[off + tid]; + scratch[tid] = abs(v); + } else { + scratch[tid] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tid < stride) { + scratch[tid] = max(scratch[tid], scratch[tid + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float amax = max(scratch[0], 1.0e-4f); + const float fp8_scale = exp2(ceil(log2(amax / 448.0f))); + if (off + (int)tid < n_nope) { + const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale; + kv[off + tid] = q; + // Diagnostic only: skip the FP16 round-trip that normally matches the + // half-typed FlashAttention KV buffer's precision. With this enabled the + // indexer will see higher-precision raw values than FlashAttention does, + // which is informative but not a production-ready setting. +#ifdef DS4_METAL_KV_RAW_F32 + raw[off + tid] = q; +#else + raw[off + tid] = (float)((half)q); +#endif + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (int i = n_nope + tid; i < head_dim; i += 64) { +#ifdef DS4_METAL_KV_RAW_F32 + raw[i] = kv[i]; +#else + raw[i] = (float)((half)kv[i]); +#endif + } } } + +/* Decode-only sibling of kernel_flash_attn_ext_vec_reduce that also applies the + * inverse RoPE tail to the row it just produced, removing a whole dispatch per + * layer. Each threadgroup owns one head's entire 512-float row, so the RoPE is + * an intra-threadgroup dependency: reduce, barrier, rotate. Both halves call the + * same shared noinline helpers the standalone kernels use, so the arithmetic and + * its codegen are identical to running the two dispatches back to back. */ +kernel void kernel_flash_attn_ext_vec_reduce_rope( + constant ds4_metal_args_flash_attn_ext_vec_reduce & args, + device const char * htmp, + device char * dst, + constant ds4_metal_args_dsv4_rope_affine_pair & rope, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg, + (short)FC_flash_attn_ext_vec_reduce_NWG, + (short)FC_flash_attn_ext_vec_reduce_DV); + + threadgroup_barrier(mem_flags::mem_device); + + const int n_nope = rope.head_dim - rope.n_dims; + if (n_nope < 0) { + return; + } + device char * row = dst + (uint64_t)tgpig * rope.row_bytes; + ds4_rope_tail_pair_affine_row(rope, + (device const char *)row, + row, + n_nope, + rope.pos0, + tiitg, + (uint)(32 * FC_flash_attn_ext_vec_reduce_NWG)); +} + +struct ds4_metal_args_dsv4_comp_finalize { + ds4_metal_args_dsv4_rope_affine_pair rope; + float rms_eps; + uint32_t pad0; +}; + +/* Decode-only emit-path fusion. Every ratio-th token, each layer finalizes + * one freshly pooled compressor row per compressor: RMS norm, RoPE tail, and + * then the FP8 round-trip + F16 commit copy (attention, 512 floats) or the + * Hadamard+FP4 QAT (indexer, 128 floats). Those were seven single-row + * dispatches; this kernel is one dispatch with two threadgroups. + * + * Each phase reproduces its standalone kernel bit-exactly: + * - norm: kernel_rms_norm_mul_f32_4's tree (float4 lanes, simd_sum, zero- + * padded 32-slot cross-simdgroup reduce); 512 uses 128 virtual threads on + * simdgroups 0-3, 128 uses 32 virtual threads on simdgroup 0. + * - rope: ds4_rope_tail_pair_affine_row verbatim (lanes 0-63, nthreads=64). + * - fp8: kernel_dsv4_fp8_kv_quantize_f32's 64-lane shmem max tree and + * round-trip, src==dst so the verbatim tail copy is a no-op and dropped. + * - commit: per-element f32->f16 conversion (value-wise exact). + * - qat: kernel_dsv4_indexer_hadamard_fp4_f32's butterfly and per-32 amax + * tree on lanes 0-127. + * Threads outside a phase's virtual width still execute every barrier, so + * threadgroup barriers stay uniform across the 256-thread threadgroup. */ +kernel void kernel_dsv4_comp_row_finalize_f32( + constant ds4_metal_args_dsv4_comp_finalize & args [[buffer(0)]], + device float * attn_row [[buffer(1)]], + device const float * attn_norm_w [[buffer(2)]], + device char * attn_cache [[buffer(3)]], + device float * index_row [[buffer(4)]], + device const float * index_norm_w [[buffer(5)]], + device float * attn_state_kv [[buffer(6)]], + device float * attn_state_score [[buffer(7)]], + device float * index_state_kv [[buffer(8)]], + device float * index_state_score [[buffer(9)]], + threadgroup float * shmem [[threadgroup(0)]], + uint tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constant ds4_metal_args_dsv4_rope_affine_pair & rope_args = args.rope; + + if (tgpig == 0) { + /* -------- attention compressor row (512 floats) -------- */ + { + device float4 * y4 = (device float4 *)attn_row; + device const float4 * x4 = (device const float4 *)attn_row; + device const float4 * w4 = (device const float4 *)attn_norm_w; + if (sgitg == 0) { + shmem[tiisg] = 0.0f; + } + float sumf = 0.0f; + if (tiitg < 128) { + sumf = dot(x4[tiitg], x4[tiitg]); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiitg < 128 && tiisg == 0) { + shmem[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total = 0.0f; + if (tiitg < 128) { + total = simd_sum(shmem[tiisg]); + } + const float mean = total / 512.0f; + const float scale = 1.0f/sqrt(mean + args.rms_eps); + if (tiitg < 128) { + y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg]; + } + } + threadgroup_barrier(mem_flags::mem_device); + ds4_rope_tail_pair_affine_row(rope_args, + (device const char *)attn_row, + (device char *)attn_row, + 512 - rope_args.n_dims, + rope_args.pos0, + tiitg, + 64u); + threadgroup_barrier(mem_flags::mem_device); + for (int off = 0; off < 512 - rope_args.n_dims; off += 64) { + float v = 0.0f; + if (tiitg < 64) { + v = attn_row[off + tiitg]; + shmem[tiitg] = abs(v); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tiitg < stride) { + shmem[tiitg] = max(shmem[tiitg], shmem[tiitg + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float amax = max(shmem[0], 1.0e-4f); + const float scale = exp2(ceil(log2(amax / 448.0f))); + if (tiitg < 64) { + const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale; + attn_row[off + tiitg] = q; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + threadgroup_barrier(mem_flags::mem_device); + if (tiitg < 128) { + device const float4 * x4 = (device const float4 *)attn_row; + device half4 * o4 = (device half4 *)attn_cache; + const float4 v = x4[tiitg]; + o4[tiitg] = half4(v); + } + return; + } + + if (tgpig >= 2u) { + /* Ratio-4 state shifts for both compressors (elementwise row move, + * so the flat gid mapping is bit-exact): 4*1024 attention elements + * then 4*256 indexer elements. */ + const uint gid = (tgpig - 2u) * 256u + tiitg; + const uint n0 = 4u * 1024u; + if (gid < n0) { + attn_state_kv[gid] = attn_state_kv[n0 + gid]; + attn_state_score[gid] = attn_state_score[n0 + gid]; + return; + } + const uint gid1 = gid - n0; + const uint n1 = 4u * 256u; + if (gid1 >= n1) return; + index_state_kv[gid1] = index_state_kv[n1 + gid1]; + index_state_score[gid1] = index_state_score[n1 + gid1]; + return; + } + + /* -------- indexer compressor row (128 floats) -------- */ + { + device float4 * y4 = (device float4 *)index_row; + device const float4 * x4 = (device const float4 *)index_row; + device const float4 * w4 = (device const float4 *)index_norm_w; + if (sgitg == 0) { + shmem[tiisg] = 0.0f; + } + float sumf = 0.0f; + if (tiitg < 32) { + sumf = dot(x4[tiitg], x4[tiitg]); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + shmem[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total = 0.0f; + if (tiitg < 32) { + total = simd_sum(shmem[tiisg]); + } + const float mean = total / 128.0f; + const float scale = 1.0f/sqrt(mean + args.rms_eps); + if (tiitg < 32) { + y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg]; + } + } + threadgroup_barrier(mem_flags::mem_device); + ds4_rope_tail_pair_affine_row(rope_args, + (device const char *)index_row, + (device char *)index_row, + 128 - rope_args.n_dims, + rope_args.pos0, + tiitg, + 64u); + threadgroup_barrier(mem_flags::mem_device); + { + threadgroup float *vals = shmem; + threadgroup float *absbuf = shmem + 128; + if (tiitg < 128) { + vals[tiitg] = index_row[tiitg]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = 1u; stride < 128u; stride <<= 1u) { + if (tiitg < 128 && (tiitg & stride) == 0u) { + const uint base = (tiitg & ~(2u * stride - 1u)) + (tiitg & (stride - 1u)); + const float a = vals[base]; + const float b = vals[base + stride]; + vals[base] = a + b; + vals[base + stride] = a - b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float v = 0.0f; + if (tiitg < 128) { + v = vals[tiitg] * 0.08838834764831845f; + absbuf[tiitg] = abs(v); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint block = tiitg >> 5u; + const uint lane = tiitg & 31u; + const uint block_base = block * 32u; + for (uint stride = 16u; stride > 0u; stride >>= 1u) { + if (tiitg < 128 && lane < stride) { + absbuf[block_base + lane] = max(absbuf[block_base + lane], + absbuf[block_base + lane + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tiitg < 128) { + const float amax = max(absbuf[block_base], 7.052966104933725e-38f); + const float scale = exp2(ceil(log2(amax / 6.0f))); + index_row[tiitg] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale; + } + } +} + +// Host-visible packed FlashAttention + exact inverse-RoPE decode kernel. +kernel void kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512( + constant ds4_metal_args_flash_attn_ext_vec & args [[buffer(0)]], + device const char * q [[buffer(1)]], + device const char * k [[buffer(2)]], + device const char * v [[buffer(3)]], + device const char * mask [[buffer(4)]], + device const char * sinks [[buffer(5)]], + device const char * pad [[buffer(6)]], + device char * dst [[buffer(7)]], + constant ds4_metal_args_dsv4_rope_affine_pair & rope + [[buffer(8)]], + threadgroup char * shmem [[threadgroup(0)]], + uint head [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + /* Uniform specialization guard; host applies the same eligibility gate. */ + if (!FC_flash_attn_ext_vec_has_mask || + !FC_flash_attn_ext_vec_has_sinks || + FC_flash_attn_ext_vec_has_bias || + FC_flash_attn_ext_vec_has_scap || + FC_flash_attn_ext_vec_nsg != 1 || + FC_flash_attn_ext_vec_nwg != 32 || + FC_flash_attn_ext_vec_ns10 != 512 || + FC_flash_attn_ext_vec_ns20 != 512 || + args.ne01 != 1 || args.ne02 != 64 || args.ne03 != 1 || + args.ne_12_2 != 1 || args.ne_12_3 != 1 || + args.ne31 != 1 || args.ne32 != 1 || args.ne33 != 1 || + args.ne11 <= 0 || args.ne11 > 1024 || head >= (uint)args.ne02 || + args.nb02 != 2048 || args.nb11 != 1024 || args.nb21 != 1024 || + rope.head_dim != 512 || rope.n_dims != 64 || + rope.row_bytes != 2048 || rope.inverse == 0) { + return; + } + + ds4_flash_attn_vec_packed8_reduce_f16_512( + args, q, k, v, mask, sinks, pad, dst, shmem, + head, tiisg, sgitg); + + /* Same producer/consumer boundary as the current reduce+RoPE kernel. */ + threadgroup_barrier(mem_flags::mem_device); + + const int n_nope = rope.head_dim - rope.n_dims; + device char * row = dst + (uint64_t)head * rope.row_bytes; + ds4_rope_tail_pair_affine_row(rope, + (device const char *)row, + row, + n_nope, + rope.pos0, + tiitg, + 32u * 32u); +} diff --git a/metal/flash_attn.metal b/metal/flash_attn.metal index 900d539..e3232a7 100644 --- a/metal/flash_attn.metal +++ b/metal/flash_attn.metal @@ -1398,24 +1398,26 @@ constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_A // Reduces split-K decode FlashAttention partials. It combines each workgroup's // output vector and softmax (sum,max) pair into the final attention result. -kernel void kernel_flash_attn_ext_vec_reduce( +/* Shared and deliberately noinline so the split-K reduction is compiled once and + * every caller gets identical codegen. The RoPE-fused sibling in dsv4_rope.metal + * calls this same body, which is what keeps the fusion bit-exact. */ +static __attribute__((noinline)) void ds4_flash_attn_vec_reduce_row( constant ds4_metal_args_flash_attn_ext_vec_reduce & args, device const char * htmp, device char * dst, - uint tgpig[[threadgroup_position_in_grid]], - ushort tiisg[[thread_index_in_simdgroup]], - ushort sgitg[[simdgroup_index_in_threadgroup]]) { -#define NWG (FC_flash_attn_ext_vec_reduce_NWG) -#define DV (FC_flash_attn_ext_vec_reduce_DV) - + uint tgpig, + ushort tiisg, + ushort sgitg, + short NWG_, + short DV_) { const uint64_t rid = tgpig; const short iwg = tiisg; - device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV*NWG; + device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV_*NWG_; - float S = ss[rid*(2*NWG) + 2*iwg + 0]; - float M = ss[rid*(2*NWG) + 2*iwg + 1]; + float S = ss[rid*(2*NWG_) + 2*iwg + 0]; + float M = ss[rid*(2*NWG_) + 2*iwg + 1]; const float m = simd_max(M); const float ms = exp(M - m); @@ -1423,19 +1425,268 @@ kernel void kernel_flash_attn_ext_vec_reduce( S = simd_sum(S*ms); S = S == 0.0f ? 0.0f : 1.0f/S; - const short DV4 = DV/4; + const short DV4 = DV_/4; - device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG; + device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG_; device float4 * dst4 = (device float4 *) dst + rid*DV4; - for (short i = sgitg; i < DV4; i += NWG) { - const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms); + for (short i = sgitg; i < DV4; i += NWG_) { + const float4 v = simd_sum(htmp4[i*NWG_ + iwg]*ms); if (iwg == 0) { dst4[i] = v*S; } } - -#undef NWG -#undef DV +} + +kernel void kernel_flash_attn_ext_vec_reduce( + constant ds4_metal_args_flash_attn_ext_vec_reduce & args, + device const char * htmp, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg, + (short)FC_flash_attn_ext_vec_reduce_NWG, + (short)FC_flash_attn_ext_vec_reduce_DV); +} + +// M5 decode specialization: time-slice all 32 split-K workgroups through eight +// physical simdgroups, then reduce through the same 32-lane topology without a +// device partial buffer. The host gate fixes the exact F16 512-wide geometry. +static inline void ds4_flash_attn_vec_packed8_reduce_f16_512( + constant ds4_metal_args_flash_attn_ext_vec & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * sinks, + device const char * pad, + device char * dst, + threadgroup char * shmem, + uint head, + ushort tiisg, + ushort sgitg) { + constexpr short NW = 32; + constexpr short C = 32; + constexpr short NSG = 8; + constexpr short NWG = 32; + constexpr short DK4 = 128; + constexpr short DV4 = 128; + constexpr short SH = 128; + + /* 24,448 dynamic bytes: shared Q, eight score/mask banks, all 32 + * split-local weights and stats, sink scales, and a padded 32x33 F32 + * float4 partial plane. */ + threadgroup half4 *q_shared = (threadgroup half4 *)shmem; + threadgroup half *score_banks = + (threadgroup half *)(q_shared + DK4); + threadgroup volatile float *weights = + (threadgroup volatile float *)(score_banks + NSG * SH); + threadgroup volatile float *stats = weights + NWG * C; + threadgroup volatile float *sink_scale = stats + 2 * NWG; + threadgroup volatile float4 *partial_plane = + (threadgroup volatile float4 *)(sink_scale + NWG); + + const short lane = (short)tiisg; + threadgroup half *bank = score_banks + (short)sgitg * SH; + threadgroup float *ss = (threadgroup float *)bank; + threadgroup half *sm = bank + 2 * C; + + device const float4 *q4 = + (device const float4 *)(q + (uint64_t)head * args.nb02); + if (sgitg == 0) { + for (short i = lane; i < DK4; i += NW) { + q_shared[i] = (half4)q4[i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + /* Eight physical simdgroups time-slice the exact 32 legacy split-K + * workgroups. The official <=1024-key gate gives each virtual split at + * most one 32-row block, so its value partial can be formed later from + * these materialized weights without changing online-softmax order. */ + for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) { + float S = 0.0f; + float M = -FLT_MAX / 2; + float out_scale = 1.0f; + const int ic_original = (int)iwg * C; + + weights[(uint)iwg * C + (uint)lane] = 0.0f; + ss[lane] = 0.0f; + sm[lane] = (half)0.0h; + simdgroup_barrier(mem_flags::mem_threadgroup); + + if (ic_original < args.ne11) { + device const char *k_block = k; + device const char *v_block = v; + device const half *pm = (device const half *)mask; + int ic = ic_original; + + if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) { + k_block = pad; + const uint64_t k_pad_bytes = + args.nb11 * (uint64_t)C * + (uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3; + const uint64_t v_pad_bytes = + args.nb21 * (uint64_t)C * + (uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3; + if (FC_flash_attn_ext_vec_shared_kvpad) { + v_block = k_block; + pm = (device const half *)(k_block + + k_pad_bytes + v_pad_bytes); + } else { + v_block = k_block + k_pad_bytes; + pm = (device const half *)(v_block + v_pad_bytes); + } + ic = 0; + } + + sm[lane] = pm[ic + lane]; + if (simd_max(sm[lane]) > -MAXHALF) { + device const half4 *pk4 = + (device const half4 *)(k_block + + (uint64_t)ic * args.nb11); + threadgroup const half4 *pq4 = q_shared; + pk4 += lane; + pq4 += lane; + + float lane_mqk = 0.0f; + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + float mqk = 0.0f; + FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) { + mqk += dot((float4)pk4[cc * DK4 + ii * NW], + (float4)pq4[ii * NW]); + } + mqk = simd_sum(mqk); + if (lane == cc) { + lane_mqk = mqk; + } + } + + ss[lane] = fma(lane_mqk, args.scale, + (float)sm[lane]); + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float old_m = M; + const float score = ss[lane]; + M = simd_max(max(M, score)); + const float ms = exp(old_m - M); + const float vs = exp(score - M); + S = S * ms + simd_sum(vs); + ss[lane] = vs; + simdgroup_barrier(mem_flags::mem_threadgroup); + + weights[(uint)iwg * C + (uint)lane] = ss[lane]; + } + + if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) { + const float old_m = M; + const float sink = lane == 0 + ? ((device const float *)sinks)[head] + : -FLT_MAX / 2; + M = simd_max(max(M, sink)); + const float ms = exp(old_m - M); + const float vs = exp(sink - M); + S = S * ms + simd_sum(vs); + out_scale = ms; + } + } else if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) { + const float old_m = M; + const float sink = lane == 0 + ? ((device const float *)sinks)[head] + : -FLT_MAX / 2; + M = simd_max(max(M, sink)); + const float ms = exp(old_m - M); + const float vs = exp(sink - M); + S = S * ms + simd_sum(vs); + out_scale = ms; + } + + if (lane == 0) { + stats[2 * (uint)iwg + 0] = S; + stats[2 * (uint)iwg + 1] = M; + sink_scale[(uint)iwg] = out_scale; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + /* Recreate the legacy 32-lane reducer exactly: lane is the virtual split + * index, including neutral idle splits in their original tree positions. */ + const short split = lane; + float reduce_S = stats[2 * (uint)split + 0]; + float reduce_M = stats[2 * (uint)split + 1]; + const float reduce_max = simd_max(reduce_M); + const float reduce_ms = exp(reduce_M - reduce_max); + reduce_S = simd_sum(reduce_S * reduce_ms); + const float reduce_inv = + reduce_S == 0.0f ? 0.0f : 1.0f / reduce_S; + + device float4 *dst4 = + (device float4 *)(dst + + (uint64_t)head * 512u * sizeof(float)); + + /* Form one 32-float4 output quadrant at a time. During production each + * physical simdgroup time-slices four virtual splits while SIMD lanes are + * contiguous output columns, exactly matching the legacy V loads and + * cc-major accumulation. A padded 33-column plane avoids a 32-way TG-bank + * conflict when the reducer transposes lanes back to virtual splits. */ + for (short quadrant = 0; quadrant < 4; ++quadrant) { + for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) { + float4 lo = float4(0.0f); + const int ic_original = (int)iwg * C; + if (ic_original < args.ne11) { + device const char *v_block = v; + int ic = ic_original; + if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) { + device const char *k_block = pad; + const uint64_t k_pad_bytes = + args.nb11 * (uint64_t)C * + (uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3; + if (FC_flash_attn_ext_vec_shared_kvpad) { + v_block = k_block; + } else { + v_block = k_block + k_pad_bytes; + } + ic = 0; + } + + device const half4 *pv4 = + (device const half4 *)(v_block + + (uint64_t)ic * args.nb21); + threadgroup volatile float *split_weights = + weights + (uint)iwg * C; + const short oc = quadrant * NW + lane; + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + lo += float4(pv4[cc * DV4 + oc]) * + float4(split_weights[cc]); + } + + float4 acc = float4(0.0f); + acc += lo; + if (iwg == 0) { + acc *= sink_scale[0]; + } + lo = acc; + } + partial_plane[(uint)iwg * 33u + (uint)lane] = lo; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + /* lane is now the legacy split index and each physical simdgroup + * reduces four output columns through the identical simd_sum tree. */ + for (short out_lane = (short)sgitg; out_lane < NW; out_lane += NSG) { + const float4 materialized = + (float4)partial_plane[(uint)lane * 33u + (uint)out_lane]; + const float4 reduced = simd_sum(materialized * reduce_ms); + if (lane == 0) { + dst4[quadrant * NW + out_lane] = reduced * reduce_inv; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } } diff --git a/metal/moe.metal b/metal/moe.metal index c1a19e8..7128d47 100644 --- a/metal/moe.metal +++ b/metal/moe.metal @@ -3,6 +3,7 @@ #ifndef QK_K #define QK_K 256 #endif +#define QK_MXFP4 32 #define N_R0_Q2_K 4 #define N_R0_GLM_Q2_PAIR2_K 1 #define N_R0_Q4_K 2 @@ -13,6 +14,280 @@ #define N_R0_Q5_K 4 #define N_R0_Q6_K 2 #define N_R0_IQ2_XXS 4 +#define N_R0_MXFP4 2 + +static constant float ds4_metal_mxfp4_values[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, +}; + +static inline float ds4_metal_e8m0_to_f32(uchar e) { + const uint bits = e == 0 ? 0x00400000u : (uint)e << 23; + return as_type(bits); +} + +// BEGIN GENERATED MXFP4 HALF LUT +// Generated by metal/generate_mxfp4_half_lut.py; do not edit by hand. +// Index = (E8M0 exponent byte << 4) | E2M1 code. +static constant ushort ds4_metal_mxfp4_half_lut[256][16] = { + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0001u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8001u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0001u, 0x0001u, 0x0002u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8001u, 0x8001u, 0x8002u }, + { 0x0000u, 0x0000u, 0x0000u, 0x0001u, 0x0001u, 0x0002u, 0x0002u, 0x0003u, 0x8000u, 0x8000u, 0x8000u, 0x8001u, 0x8001u, 0x8002u, 0x8002u, 0x8003u }, + { 0x0000u, 0x0000u, 0x0001u, 0x0002u, 0x0002u, 0x0003u, 0x0004u, 0x0006u, 0x8000u, 0x8000u, 0x8001u, 0x8002u, 0x8002u, 0x8003u, 0x8004u, 0x8006u }, + { 0x0000u, 0x0001u, 0x0002u, 0x0003u, 0x0004u, 0x0006u, 0x0008u, 0x000cu, 0x8000u, 0x8001u, 0x8002u, 0x8003u, 0x8004u, 0x8006u, 0x8008u, 0x800cu }, + { 0x0000u, 0x0002u, 0x0004u, 0x0006u, 0x0008u, 0x000cu, 0x0010u, 0x0018u, 0x8000u, 0x8002u, 0x8004u, 0x8006u, 0x8008u, 0x800cu, 0x8010u, 0x8018u }, + { 0x0000u, 0x0004u, 0x0008u, 0x000cu, 0x0010u, 0x0018u, 0x0020u, 0x0030u, 0x8000u, 0x8004u, 0x8008u, 0x800cu, 0x8010u, 0x8018u, 0x8020u, 0x8030u }, + { 0x0000u, 0x0008u, 0x0010u, 0x0018u, 0x0020u, 0x0030u, 0x0040u, 0x0060u, 0x8000u, 0x8008u, 0x8010u, 0x8018u, 0x8020u, 0x8030u, 0x8040u, 0x8060u }, + { 0x0000u, 0x0010u, 0x0020u, 0x0030u, 0x0040u, 0x0060u, 0x0080u, 0x00c0u, 0x8000u, 0x8010u, 0x8020u, 0x8030u, 0x8040u, 0x8060u, 0x8080u, 0x80c0u }, + { 0x0000u, 0x0020u, 0x0040u, 0x0060u, 0x0080u, 0x00c0u, 0x0100u, 0x0180u, 0x8000u, 0x8020u, 0x8040u, 0x8060u, 0x8080u, 0x80c0u, 0x8100u, 0x8180u }, + { 0x0000u, 0x0040u, 0x0080u, 0x00c0u, 0x0100u, 0x0180u, 0x0200u, 0x0300u, 0x8000u, 0x8040u, 0x8080u, 0x80c0u, 0x8100u, 0x8180u, 0x8200u, 0x8300u }, + { 0x0000u, 0x0080u, 0x0100u, 0x0180u, 0x0200u, 0x0300u, 0x0400u, 0x0600u, 0x8000u, 0x8080u, 0x8100u, 0x8180u, 0x8200u, 0x8300u, 0x8400u, 0x8600u }, + { 0x0000u, 0x0100u, 0x0200u, 0x0300u, 0x0400u, 0x0600u, 0x0800u, 0x0a00u, 0x8000u, 0x8100u, 0x8200u, 0x8300u, 0x8400u, 0x8600u, 0x8800u, 0x8a00u }, + { 0x0000u, 0x0200u, 0x0400u, 0x0600u, 0x0800u, 0x0a00u, 0x0c00u, 0x0e00u, 0x8000u, 0x8200u, 0x8400u, 0x8600u, 0x8800u, 0x8a00u, 0x8c00u, 0x8e00u }, + { 0x0000u, 0x0400u, 0x0800u, 0x0a00u, 0x0c00u, 0x0e00u, 0x1000u, 0x1200u, 0x8000u, 0x8400u, 0x8800u, 0x8a00u, 0x8c00u, 0x8e00u, 0x9000u, 0x9200u }, + { 0x0000u, 0x0800u, 0x0c00u, 0x0e00u, 0x1000u, 0x1200u, 0x1400u, 0x1600u, 0x8000u, 0x8800u, 0x8c00u, 0x8e00u, 0x9000u, 0x9200u, 0x9400u, 0x9600u }, + { 0x0000u, 0x0c00u, 0x1000u, 0x1200u, 0x1400u, 0x1600u, 0x1800u, 0x1a00u, 0x8000u, 0x8c00u, 0x9000u, 0x9200u, 0x9400u, 0x9600u, 0x9800u, 0x9a00u }, + { 0x0000u, 0x1000u, 0x1400u, 0x1600u, 0x1800u, 0x1a00u, 0x1c00u, 0x1e00u, 0x8000u, 0x9000u, 0x9400u, 0x9600u, 0x9800u, 0x9a00u, 0x9c00u, 0x9e00u }, + { 0x0000u, 0x1400u, 0x1800u, 0x1a00u, 0x1c00u, 0x1e00u, 0x2000u, 0x2200u, 0x8000u, 0x9400u, 0x9800u, 0x9a00u, 0x9c00u, 0x9e00u, 0xa000u, 0xa200u }, + { 0x0000u, 0x1800u, 0x1c00u, 0x1e00u, 0x2000u, 0x2200u, 0x2400u, 0x2600u, 0x8000u, 0x9800u, 0x9c00u, 0x9e00u, 0xa000u, 0xa200u, 0xa400u, 0xa600u }, + { 0x0000u, 0x1c00u, 0x2000u, 0x2200u, 0x2400u, 0x2600u, 0x2800u, 0x2a00u, 0x8000u, 0x9c00u, 0xa000u, 0xa200u, 0xa400u, 0xa600u, 0xa800u, 0xaa00u }, + { 0x0000u, 0x2000u, 0x2400u, 0x2600u, 0x2800u, 0x2a00u, 0x2c00u, 0x2e00u, 0x8000u, 0xa000u, 0xa400u, 0xa600u, 0xa800u, 0xaa00u, 0xac00u, 0xae00u }, + { 0x0000u, 0x2400u, 0x2800u, 0x2a00u, 0x2c00u, 0x2e00u, 0x3000u, 0x3200u, 0x8000u, 0xa400u, 0xa800u, 0xaa00u, 0xac00u, 0xae00u, 0xb000u, 0xb200u }, + { 0x0000u, 0x2800u, 0x2c00u, 0x2e00u, 0x3000u, 0x3200u, 0x3400u, 0x3600u, 0x8000u, 0xa800u, 0xac00u, 0xae00u, 0xb000u, 0xb200u, 0xb400u, 0xb600u }, + { 0x0000u, 0x2c00u, 0x3000u, 0x3200u, 0x3400u, 0x3600u, 0x3800u, 0x3a00u, 0x8000u, 0xac00u, 0xb000u, 0xb200u, 0xb400u, 0xb600u, 0xb800u, 0xba00u }, + { 0x0000u, 0x3000u, 0x3400u, 0x3600u, 0x3800u, 0x3a00u, 0x3c00u, 0x3e00u, 0x8000u, 0xb000u, 0xb400u, 0xb600u, 0xb800u, 0xba00u, 0xbc00u, 0xbe00u }, + { 0x0000u, 0x3400u, 0x3800u, 0x3a00u, 0x3c00u, 0x3e00u, 0x4000u, 0x4200u, 0x8000u, 0xb400u, 0xb800u, 0xba00u, 0xbc00u, 0xbe00u, 0xc000u, 0xc200u }, + { 0x0000u, 0x3800u, 0x3c00u, 0x3e00u, 0x4000u, 0x4200u, 0x4400u, 0x4600u, 0x8000u, 0xb800u, 0xbc00u, 0xbe00u, 0xc000u, 0xc200u, 0xc400u, 0xc600u }, + { 0x0000u, 0x3c00u, 0x4000u, 0x4200u, 0x4400u, 0x4600u, 0x4800u, 0x4a00u, 0x8000u, 0xbc00u, 0xc000u, 0xc200u, 0xc400u, 0xc600u, 0xc800u, 0xca00u }, + { 0x0000u, 0x4000u, 0x4400u, 0x4600u, 0x4800u, 0x4a00u, 0x4c00u, 0x4e00u, 0x8000u, 0xc000u, 0xc400u, 0xc600u, 0xc800u, 0xca00u, 0xcc00u, 0xce00u }, + { 0x0000u, 0x4400u, 0x4800u, 0x4a00u, 0x4c00u, 0x4e00u, 0x5000u, 0x5200u, 0x8000u, 0xc400u, 0xc800u, 0xca00u, 0xcc00u, 0xce00u, 0xd000u, 0xd200u }, + { 0x0000u, 0x4800u, 0x4c00u, 0x4e00u, 0x5000u, 0x5200u, 0x5400u, 0x5600u, 0x8000u, 0xc800u, 0xcc00u, 0xce00u, 0xd000u, 0xd200u, 0xd400u, 0xd600u }, + { 0x0000u, 0x4c00u, 0x5000u, 0x5200u, 0x5400u, 0x5600u, 0x5800u, 0x5a00u, 0x8000u, 0xcc00u, 0xd000u, 0xd200u, 0xd400u, 0xd600u, 0xd800u, 0xda00u }, + { 0x0000u, 0x5000u, 0x5400u, 0x5600u, 0x5800u, 0x5a00u, 0x5c00u, 0x5e00u, 0x8000u, 0xd000u, 0xd400u, 0xd600u, 0xd800u, 0xda00u, 0xdc00u, 0xde00u }, + { 0x0000u, 0x5400u, 0x5800u, 0x5a00u, 0x5c00u, 0x5e00u, 0x6000u, 0x6200u, 0x8000u, 0xd400u, 0xd800u, 0xda00u, 0xdc00u, 0xde00u, 0xe000u, 0xe200u }, + { 0x0000u, 0x5800u, 0x5c00u, 0x5e00u, 0x6000u, 0x6200u, 0x6400u, 0x6600u, 0x8000u, 0xd800u, 0xdc00u, 0xde00u, 0xe000u, 0xe200u, 0xe400u, 0xe600u }, + { 0x0000u, 0x5c00u, 0x6000u, 0x6200u, 0x6400u, 0x6600u, 0x6800u, 0x6a00u, 0x8000u, 0xdc00u, 0xe000u, 0xe200u, 0xe400u, 0xe600u, 0xe800u, 0xea00u }, + { 0x0000u, 0x6000u, 0x6400u, 0x6600u, 0x6800u, 0x6a00u, 0x6c00u, 0x6e00u, 0x8000u, 0xe000u, 0xe400u, 0xe600u, 0xe800u, 0xea00u, 0xec00u, 0xee00u }, + { 0x0000u, 0x6400u, 0x6800u, 0x6a00u, 0x6c00u, 0x6e00u, 0x7000u, 0x7200u, 0x8000u, 0xe400u, 0xe800u, 0xea00u, 0xec00u, 0xee00u, 0xf000u, 0xf200u }, + { 0x0000u, 0x6800u, 0x6c00u, 0x6e00u, 0x7000u, 0x7200u, 0x7400u, 0x7600u, 0x8000u, 0xe800u, 0xec00u, 0xee00u, 0xf000u, 0xf200u, 0xf400u, 0xf600u }, + { 0x0000u, 0x6c00u, 0x7000u, 0x7200u, 0x7400u, 0x7600u, 0x7800u, 0x7a00u, 0x8000u, 0xec00u, 0xf000u, 0xf200u, 0xf400u, 0xf600u, 0xf800u, 0xfa00u }, + { 0x0000u, 0x7000u, 0x7400u, 0x7600u, 0x7800u, 0x7a00u, 0x7c00u, 0x7c00u, 0x8000u, 0xf000u, 0xf400u, 0xf600u, 0xf800u, 0xfa00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7400u, 0x7800u, 0x7a00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xf400u, 0xf800u, 0xfa00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7800u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xf800u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, + { 0x7e00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7e00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u }, +}; +// END GENERATED MXFP4 HALF LUT static constant uchar ds4_metal_kmask_iq2xs[8] = { 1, 2, 4, 8, 16, 32, 64, 128 @@ -134,6 +409,11 @@ struct block_iq2_xxs { ushort qs[QK_K/8]; }; +struct block_mxfp4 { + uchar e; + uchar qs[QK_MXFP4/2]; +}; + struct ds4_metal_glm_routed_moe_args { uint32_t in_dim; uint32_t mid_dim; @@ -1061,186 +1341,6 @@ kernel void kernel_glm_q4_K_addr_pair_swiglu_f32_masked( tgpig, slot, token, selected_off, 0, tiisg, sgitg); } -kernel void kernel_glm_q4_K_slots6_pair_swiglu_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *gate0, - device const char *gate1, - device const char *gate2, - device const char *gate3, - device const char *gate4, - device const char *gate5, - device const char *up0, - device const char *up1, - device const char *up2, - device const char *up3, - device const char *up4, - device const char *up5, - device const float *x, - device const float *weights, - device float *mid, - threadgroup float *scratch [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const uint slot = tgpig.y; - const uint token = tgpig.z; - if (slot >= args.n_expert_used || token >= args.n_tokens) return; - - device const char *gate_cur = gate0; - device const char *up_cur = up0; - switch (slot) { - case 1: gate_cur = gate1; up_cur = up1; break; - case 2: gate_cur = gate2; up_cur = up2; break; - case 3: gate_cur = gate3; up_cur = up3; break; - case 4: gate_cur = gate4; up_cur = up4; break; - case 5: gate_cur = gate5; up_cur = up5; break; - default: break; - } - - const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; - glm_q4_K_pair_swiglu_simd_f32_impl( - args, gate_cur, up_cur, x, weights, mid, scratch, - tgpig, slot, token, selected_off, 0, tiisg, sgitg); -} - -kernel void kernel_glm_q4_K_slots8_pair_swiglu_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *gate0, - device const char *gate1, - device const char *gate2, - device const char *gate3, - device const char *gate4, - device const char *gate5, - device const char *gate6, - device const char *gate7, - device const char *up0, - device const char *up1, - device const char *up2, - device const char *up3, - device const char *up4, - device const char *up5, - device const char *up6, - device const char *up7, - device const float *x, - device const float *weights, - device float *mid, - threadgroup float *scratch [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const uint slot = tgpig.y; - const uint token = tgpig.z; - if (slot >= args.n_expert_used || token >= args.n_tokens) return; - - device const char *gate_cur = gate0; - device const char *up_cur = up0; - switch (slot) { - case 1: gate_cur = gate1; up_cur = up1; break; - case 2: gate_cur = gate2; up_cur = up2; break; - case 3: gate_cur = gate3; up_cur = up3; break; - case 4: gate_cur = gate4; up_cur = up4; break; - case 5: gate_cur = gate5; up_cur = up5; break; - case 6: gate_cur = gate6; up_cur = up6; break; - case 7: gate_cur = gate7; up_cur = up7; break; - default: break; - } - - const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; - glm_q4_K_pair_swiglu_simd_f32_impl( - args, gate_cur, up_cur, x, weights, mid, scratch, - tgpig, slot, token, selected_off, 0, tiisg, sgitg); -} - -kernel void kernel_glm_q4_K_slots6_pair_swiglu4_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *gate0, - device const char *gate1, - device const char *gate2, - device const char *gate3, - device const char *gate4, - device const char *gate5, - device const char *up0, - device const char *up1, - device const char *up2, - device const char *up3, - device const char *up4, - device const char *up5, - device const float *x, - device const float *weights, - device float *mid, - threadgroup float *scratch [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const uint slot = tgpig.y; - const uint token = tgpig.z; - if (slot >= args.n_expert_used || token >= args.n_tokens) return; - - device const char *gate_cur = gate0; - device const char *up_cur = up0; - switch (slot) { - case 1: gate_cur = gate1; up_cur = up1; break; - case 2: gate_cur = gate2; up_cur = up2; break; - case 3: gate_cur = gate3; up_cur = up3; break; - case 4: gate_cur = gate4; up_cur = up4; break; - case 5: gate_cur = gate5; up_cur = up5; break; - default: break; - } - - const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; - glm_q4_K_pair_swiglu_simd_f32_impl( - args, gate_cur, up_cur, x, weights, mid, scratch, - tgpig, slot, token, selected_off, 0, tiisg, sgitg); -} - -kernel void kernel_glm_q4_K_slots8_pair_swiglu4_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *gate0, - device const char *gate1, - device const char *gate2, - device const char *gate3, - device const char *gate4, - device const char *gate5, - device const char *gate6, - device const char *gate7, - device const char *up0, - device const char *up1, - device const char *up2, - device const char *up3, - device const char *up4, - device const char *up5, - device const char *up6, - device const char *up7, - device const float *x, - device const float *weights, - device float *mid, - threadgroup float *scratch [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const uint slot = tgpig.y; - const uint token = tgpig.z; - if (slot >= args.n_expert_used || token >= args.n_tokens) return; - - device const char *gate_cur = gate0; - device const char *up_cur = up0; - switch (slot) { - case 1: gate_cur = gate1; up_cur = up1; break; - case 2: gate_cur = gate2; up_cur = up2; break; - case 3: gate_cur = gate3; up_cur = up3; break; - case 4: gate_cur = gate4; up_cur = up4; break; - case 5: gate_cur = gate5; up_cur = up5; break; - case 6: gate_cur = gate6; up_cur = up6; break; - case 7: gate_cur = gate7; up_cur = up7; break; - default: break; - } - - const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; - glm_q4_K_pair_swiglu_simd_f32_impl( - args, gate_cur, up_cur, x, weights, mid, scratch, - tgpig, slot, token, selected_off, 0, tiisg, sgitg); -} - kernel void kernel_glm_q4_K_pair_swiglu4_f32( constant ds4_metal_glm_routed_moe_args &args, device const char *gate, @@ -1521,96 +1621,6 @@ kernel void kernel_glm_q5_K_pair_swiglu_f32( tgpig, slot, token, selected_off, expert, tiisg, sgitg); } -kernel void kernel_glm_q5_K_slots6_pair_swiglu_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *gate0, - device const char *gate1, - device const char *gate2, - device const char *gate3, - device const char *gate4, - device const char *gate5, - device const char *up0, - device const char *up1, - device const char *up2, - device const char *up3, - device const char *up4, - device const char *up5, - device const float *x, - device const float *weights, - device float *mid, - threadgroup float *scratch [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const uint slot = tgpig.y; - const uint token = tgpig.z; - if (slot >= args.n_expert_used || token >= args.n_tokens) return; - - device const char *gate_cur = gate0; - device const char *up_cur = up0; - switch (slot) { - case 1: gate_cur = gate1; up_cur = up1; break; - case 2: gate_cur = gate2; up_cur = up2; break; - case 3: gate_cur = gate3; up_cur = up3; break; - case 4: gate_cur = gate4; up_cur = up4; break; - case 5: gate_cur = gate5; up_cur = up5; break; - default: break; - } - - const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; - glm_q5_K_pair_swiglu_f32_impl( - args, gate_cur, up_cur, x, weights, mid, scratch, - tgpig, slot, token, selected_off, 0, tiisg, sgitg); -} - -kernel void kernel_glm_q5_K_slots8_pair_swiglu_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *gate0, - device const char *gate1, - device const char *gate2, - device const char *gate3, - device const char *gate4, - device const char *gate5, - device const char *gate6, - device const char *gate7, - device const char *up0, - device const char *up1, - device const char *up2, - device const char *up3, - device const char *up4, - device const char *up5, - device const char *up6, - device const char *up7, - device const float *x, - device const float *weights, - device float *mid, - threadgroup float *scratch [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const uint slot = tgpig.y; - const uint token = tgpig.z; - if (slot >= args.n_expert_used || token >= args.n_tokens) return; - - device const char *gate_cur = gate0; - device const char *up_cur = up0; - switch (slot) { - case 1: gate_cur = gate1; up_cur = up1; break; - case 2: gate_cur = gate2; up_cur = up2; break; - case 3: gate_cur = gate3; up_cur = up3; break; - case 4: gate_cur = gate4; up_cur = up4; break; - case 5: gate_cur = gate5; up_cur = up5; break; - case 6: gate_cur = gate6; up_cur = up6; break; - case 7: gate_cur = gate7; up_cur = up7; break; - default: break; - } - - const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; - glm_q5_K_pair_swiglu_f32_impl( - args, gate_cur, up_cur, x, weights, mid, scratch, - tgpig, slot, token, selected_off, 0, tiisg, sgitg); -} - kernel void kernel_glm_q5_K_pair_swiglu_mapped_f32( constant ds4_metal_glm_routed_moe_args &args, device const char *gate, @@ -2350,468 +2360,6 @@ kernel void kernel_glm_q6_K_down_f32( } } -kernel void kernel_glm_q5_K_slots6_down_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *down0, - device const char *down1, - device const char *down2, - device const char *down3, - device const char *down4, - device const char *down5, - device const float *mid, - device float *out, - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const short NSG = 2; - constexpr uint16_t kmask1 = 0x3f3f; - constexpr uint16_t kmask2 = 0x0f0f; - constexpr uint16_t kmask3 = 0xc0c0; - - const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_K; - const uint token = tgpig.y; - if (row0 >= args.out_dim || token >= args.n_tokens) return; - - const short ix = tiisg / 8; - const short it = tiisg % 8; - const short iq = it / 4; - const short ir = it % 4; - const uint bit0 = 2u * (uint)iq; - const uint bit1 = bit0 + 1u; - const uint bit2 = bit0 + 4u; - const uint bit3 = bit0 + 5u; - const int nb = args.mid_dim / QK_K; - - float sumf[N_R0_Q5_K] = {0.f}; - uint16_t sc16[4]; - thread const uint8_t *sc8 = (thread const uint8_t *)sc16; - const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; - for (uint slot = 0; slot < args.n_expert_used; slot++) { - device const char *down_cur = down0; - switch (slot) { - case 1: down_cur = down1; break; - case 2: down_cur = down2; break; - case 3: down_cur = down3; break; - case 4: down_cur = down4; break; - case 5: down_cur = down5; break; - default: break; - } - - device const block_q5_K *x = - (device const block_q5_K *)(down_cur + - (uint64_t)row0 * args.down_row_bytes); - device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; - device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; - - for (int ib = ix; ib < nb; ib += 4) { - float yl[16]; - float yh[16]; - float4 sumy = {0.f, 0.f, 0.f, 0.f}; - - for (short i = 0; i < 8; ++i) { - yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; - yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; - yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; - yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; - } - - device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; - device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; - device const uint16_t *qh = (device const uint16_t *)x[ib].qh + 4 * ir; - device const half *dh = &x[ib].d; - - for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { - sc16[0] = sc[0] & kmask1; - sc16[1] = sc[2] & kmask1; - sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); - sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); - - device const uint16_t *q2 = q1 + 32; - float4 acc = {0.f, 0.f, 0.f, 0.f}; - - FOR_UNROLL (short i = 0; i < 4; ++i) { - const uint ql1 = (uint)q1[i]; - const uint ql2 = (uint)q2[i]; - const uint hb = (uint)qh[i]; - acc[0] += yl[2 * i + 0] * - (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + - yl[2 * i + 1] * - (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); - acc[1] += yl[2 * i + 8] * - (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + - yl[2 * i + 9] * - (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); - acc[2] += yh[2 * i + 0] * - (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + - yh[2 * i + 1] * - (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); - acc[3] += yh[2 * i + 8] * - (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + - yh[2 * i + 9] * - (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); - } - - sumf[row] += dh[0] * (acc[0] * sc8[0] + acc[1] * sc8[1] + - acc[2] * sc8[4] + acc[3] * sc8[5]) - - dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + - sumy[2] * sc8[6] + sumy[3] * sc8[7]); - - q1 += args.down_row_bytes / 2; - qh += args.down_row_bytes / 2; - sc += args.down_row_bytes / 2; - dh += args.down_row_bytes / 2; - } - - y4 += 4 * QK_K; - } - } - - for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { - const float sum_all = simd_sum(sumf[row]); - if (tiisg == 0u) { - out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; - } - } -} - -kernel void kernel_glm_q5_K_slots8_down_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *down0, - device const char *down1, - device const char *down2, - device const char *down3, - device const char *down4, - device const char *down5, - device const char *down6, - device const char *down7, - device const float *mid, - device float *out, - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const short NSG = 2; - constexpr uint16_t kmask1 = 0x3f3f; - constexpr uint16_t kmask2 = 0x0f0f; - constexpr uint16_t kmask3 = 0xc0c0; - - const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_K; - const uint token = tgpig.y; - if (row0 >= args.out_dim || token >= args.n_tokens) return; - - const short ix = tiisg / 8; - const short it = tiisg % 8; - const short iq = it / 4; - const short ir = it % 4; - const uint bit0 = 2u * (uint)iq; - const uint bit1 = bit0 + 1u; - const uint bit2 = bit0 + 4u; - const uint bit3 = bit0 + 5u; - const int nb = args.mid_dim / QK_K; - - float sumf[N_R0_Q5_K] = {0.f}; - uint16_t sc16[4]; - thread const uint8_t *sc8 = (thread const uint8_t *)sc16; - const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; - for (uint slot = 0; slot < args.n_expert_used; slot++) { - device const char *down_cur = down0; - switch (slot) { - case 1: down_cur = down1; break; - case 2: down_cur = down2; break; - case 3: down_cur = down3; break; - case 4: down_cur = down4; break; - case 5: down_cur = down5; break; - case 6: down_cur = down6; break; - case 7: down_cur = down7; break; - default: break; - } - - device const block_q5_K *x = - (device const block_q5_K *)(down_cur + - (uint64_t)row0 * args.down_row_bytes); - device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; - device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; - - for (int ib = ix; ib < nb; ib += 4) { - float yl[16]; - float yh[16]; - float4 sumy = {0.f, 0.f, 0.f, 0.f}; - - for (short i = 0; i < 8; ++i) { - yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; - yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; - yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; - yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; - } - - device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; - device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; - device const uint16_t *qh = (device const uint16_t *)x[ib].qh + 4 * ir; - device const half *dh = &x[ib].d; - - for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { - sc16[0] = sc[0] & kmask1; - sc16[1] = sc[2] & kmask1; - sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); - sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); - - device const uint16_t *q2 = q1 + 32; - float4 acc = {0.f, 0.f, 0.f, 0.f}; - - FOR_UNROLL (short i = 0; i < 4; ++i) { - const uint ql1 = (uint)q1[i]; - const uint ql2 = (uint)q2[i]; - const uint hb = (uint)qh[i]; - acc[0] += yl[2 * i + 0] * - (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + - yl[2 * i + 1] * - (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); - acc[1] += yl[2 * i + 8] * - (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + - yl[2 * i + 9] * - (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); - acc[2] += yh[2 * i + 0] * - (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + - yh[2 * i + 1] * - (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); - acc[3] += yh[2 * i + 8] * - (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + - yh[2 * i + 9] * - (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); - } - - sumf[row] += dh[0] * (acc[0] * sc8[0] + acc[1] * sc8[1] + - acc[2] * sc8[4] + acc[3] * sc8[5]) - - dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + - sumy[2] * sc8[6] + sumy[3] * sc8[7]); - - q1 += args.down_row_bytes / 2; - qh += args.down_row_bytes / 2; - sc += args.down_row_bytes / 2; - dh += args.down_row_bytes / 2; - } - - y4 += 4 * QK_K; - } - } - - for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { - const float sum_all = simd_sum(sumf[row]); - if (tiisg == 0u) { - out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; - } - } -} - -kernel void kernel_glm_q6_K_slots6_down_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *down0, - device const char *down1, - device const char *down2, - device const char *down3, - device const char *down4, - device const char *down5, - device const float *mid, - device float *out, - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const short NSG = 2; - constexpr uint kmask1 = 0x03u; - constexpr uint kmask2 = 0x0Cu; - constexpr uint kmask3 = 0x30u; - constexpr uint kmask4 = 0xC0u; - - const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q6_K; - const uint token = tgpig.y; - if (row0 >= args.out_dim || token >= args.n_tokens) return; - - const int nb = args.mid_dim / QK_K; - float sumf[N_R0_Q6_K] = {0.f}; - float yl[16]; - const short tid = tiisg / 2; - const short ix = tiisg % 2; - const short ip = tid / 8; - const short il = tid % 8; - const short l0 = 4 * il; - const short is = 8 * ip + l0 / 16; - const short y_offset = 128 * ip + l0; - const short q_offset_l = 64 * ip + l0; - const short q_offset_h = 32 * ip + l0; - - const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; - for (uint slot = 0; slot < args.n_expert_used; slot++) { - device const char *down_cur = down0; - switch (slot) { - case 1: down_cur = down1; break; - case 2: down_cur = down2; break; - case 3: down_cur = down3; break; - case 4: down_cur = down4; break; - case 5: down_cur = down5; break; - default: break; - } - - device const block_q6_K *x = - (device const block_q6_K *)(down_cur + - (uint64_t)row0 * args.down_row_bytes); - device const float *yy = mid + mid_base + (uint64_t)slot * args.mid_dim; - - for (int ib = ix; ib < nb; ib += 2) { - device const uchar *q1 = x[ib].ql + q_offset_l; - device const uchar *q2 = q1 + 32; - device const uchar *qh = x[ib].qh + q_offset_h; - device const char *sc = x[ib].scales + is; - device const half *dh = &x[ib].d; - device const float *y = yy + (uint64_t)ib * QK_K + y_offset; - - for (short l = 0; l < 4; ++l) { - yl[4 * l + 0] = y[l + 0]; - yl[4 * l + 1] = y[l + 32]; - yl[4 * l + 2] = y[l + 64]; - yl[4 * l + 3] = y[l + 96]; - } - - for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { - float4 sums = {0.f, 0.f, 0.f, 0.f}; - - FOR_UNROLL (short l = 0; l < 4; ++l) { - const uint h = (uint)qh[l]; - sums[0] += yl[4 * l + 0] * - (float)((int)((q1[l] & 0x0Fu) | ((h & kmask1) << 4u)) - 32); - sums[1] += yl[4 * l + 1] * - (float)((int)((q2[l] & 0x0Fu) | ((h & kmask2) << 2u)) - 32); - sums[2] += yl[4 * l + 2] * - (float)((int)((q1[l] >> 4u) | (h & kmask3)) - 32); - sums[3] += yl[4 * l + 3] * - (float)((int)((q2[l] >> 4u) | ((h & kmask4) >> 2u)) - 32); - } - - sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + - sums[2] * sc[4] + sums[3] * sc[6]); - - q1 += args.down_row_bytes; - q2 += args.down_row_bytes; - qh += args.down_row_bytes; - sc += args.down_row_bytes; - dh += args.down_row_bytes / 2; - } - } - } - - for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { - const float sum_all = simd_sum(sumf[row]); - if (tiisg == 0u) { - out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; - } - } -} - -kernel void kernel_glm_q6_K_slots8_down_f32( - constant ds4_metal_glm_routed_moe_args &args, - device const char *down0, - device const char *down1, - device const char *down2, - device const char *down3, - device const char *down4, - device const char *down5, - device const char *down6, - device const char *down7, - device const float *mid, - device float *out, - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - const short NSG = 2; - constexpr uint kmask1 = 0x03u; - constexpr uint kmask2 = 0x0Cu; - constexpr uint kmask3 = 0x30u; - constexpr uint kmask4 = 0xC0u; - - const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q6_K; - const uint token = tgpig.y; - if (row0 >= args.out_dim || token >= args.n_tokens) return; - - const int nb = args.mid_dim / QK_K; - float sumf[N_R0_Q6_K] = {0.f}; - float yl[16]; - const short tid = tiisg / 2; - const short ix = tiisg % 2; - const short ip = tid / 8; - const short il = tid % 8; - const short l0 = 4 * il; - const short is = 8 * ip + l0 / 16; - const short y_offset = 128 * ip + l0; - const short q_offset_l = 64 * ip + l0; - const short q_offset_h = 32 * ip + l0; - - const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; - for (uint slot = 0; slot < args.n_expert_used; slot++) { - device const char *down_cur = down0; - switch (slot) { - case 1: down_cur = down1; break; - case 2: down_cur = down2; break; - case 3: down_cur = down3; break; - case 4: down_cur = down4; break; - case 5: down_cur = down5; break; - case 6: down_cur = down6; break; - case 7: down_cur = down7; break; - default: break; - } - - device const block_q6_K *x = - (device const block_q6_K *)(down_cur + - (uint64_t)row0 * args.down_row_bytes); - device const float *yy = mid + mid_base + (uint64_t)slot * args.mid_dim; - - for (int ib = ix; ib < nb; ib += 2) { - device const uchar *q1 = x[ib].ql + q_offset_l; - device const uchar *q2 = q1 + 32; - device const uchar *qh = x[ib].qh + q_offset_h; - device const char *sc = x[ib].scales + is; - device const half *dh = &x[ib].d; - device const float *y = yy + (uint64_t)ib * QK_K + y_offset; - - for (short l = 0; l < 4; ++l) { - yl[4 * l + 0] = y[l + 0]; - yl[4 * l + 1] = y[l + 32]; - yl[4 * l + 2] = y[l + 64]; - yl[4 * l + 3] = y[l + 96]; - } - - for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { - float4 sums = {0.f, 0.f, 0.f, 0.f}; - - FOR_UNROLL (short l = 0; l < 4; ++l) { - const uint h = (uint)qh[l]; - sums[0] += yl[4 * l + 0] * - (float)((int)((q1[l] & 0x0Fu) | ((h & kmask1) << 4u)) - 32); - sums[1] += yl[4 * l + 1] * - (float)((int)((q2[l] & 0x0Fu) | ((h & kmask2) << 2u)) - 32); - sums[2] += yl[4 * l + 2] * - (float)((int)((q1[l] >> 4u) | (h & kmask3)) - 32); - sums[3] += yl[4 * l + 3] * - (float)((int)((q2[l] >> 4u) | ((h & kmask4) >> 2u)) - 32); - } - - sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + - sums[2] * sc[4] + sums[3] * sc[6]); - - q1 += args.down_row_bytes; - q2 += args.down_row_bytes; - qh += args.down_row_bytes; - sc += args.down_row_bytes; - dh += args.down_row_bytes / 2; - } - } - } - - for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { - const float sum_all = simd_sum(sumf[row]); - if (tiisg == 0u) { - out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; - } - } -} - template void dequantize_q4_K(device const block_q4_K *xb, short il, thread type4x4 ®) { device const uchar *q = xb->qs; @@ -2833,6 +2381,111 @@ void dequantize_q4_K(device const block_q4_K *xb, short il, thread type4x4 ®) } } +template +void dequantize_mxfp4(device const block_mxfp4 *xb, short il, thread type4x4 ®) { + const float d = ds4_metal_e8m0_to_f32(xb->e); + const uint shift = il == 0 ? 0u : 4u; + FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) { + const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu; + reg[i/4][i%4] = d * ds4_metal_mxfp4_values[q]; + } +} + +// The resident MXFP4 prefill down projection stages its weight tile as half. +// All finite E8M0 x E2M1 products can therefore be rounded once in this exact +// binary16 table. E8M0 0xff falls back to the established arithmetic because +// infinity-times-zero produces a GPU-canonicalized NaN rather than a portable +// table payload. +static __attribute__((noinline)) half ds4_metal_mxfp4_half_fallback_value( + uchar e, + uint q) { + half result; + result = ds4_metal_e8m0_to_f32(e) * ds4_metal_mxfp4_values[q]; + return result; +} + +static inline half ds4_metal_mxfp4_half_lut_value(uchar e, uint q) { + if (e == 0xffu) { + return ds4_metal_mxfp4_half_fallback_value(e, q); + } + return as_type(ds4_metal_mxfp4_half_lut[(uint)e][q]); +} + +void dequantize_mxfp4_half_lut( + device const block_mxfp4 *xb, + short il, + thread half4x4 ®) { + const uchar e = xb->e; + const uint shift = il == 0 ? 0u : 4u; + if (e == 0xffu) { + FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) { + const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu; + reg[i/4][i%4] = ds4_metal_mxfp4_half_fallback_value(e, q); + } + return; + } + FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) { + const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu; + reg[i/4][i%4] = as_type( + ds4_metal_mxfp4_half_lut[(uint)e][q]); + } +} + +// Exact half-domain scaling for the resident MXFP4 prefill pair tile. E8M0 is +// a pure power of two and every E2M1 magnitude is exact in binary16, so for +// E8M0 bytes in [103, 142] the half product is the once-rounded value of the +// float product: half(d) is exact there and half(d) * half(v) rounds exactly +// once, exactly like half(float(d) * float(v)). Outside that band (subnormal +// or overflowing halves, where a second rounding could appear) the established +// float arithmetic is kept. The 32-byte table replaces both the 64-byte float +// table and the per-element f32->f16 narrowing. +static constant half ds4_metal_mxfp4_half_values[16] = { + 0.0h, 0.5h, 1.0h, 1.5h, 2.0h, 3.0h, 4.0h, 6.0h, + -0.0h, -0.5h, -1.0h, -1.5h, -2.0h, -3.0h, -4.0h, -6.0h, +}; + +void dequantize_mxfp4_half_scale( + device const block_mxfp4 *xb, + short il, + thread half4x4 ®) { + const uchar e = xb->e; + const uint shift = il == 0 ? 0u : 4u; + if (e < 103u || e > 142u) { + const float d = ds4_metal_e8m0_to_f32(e); + FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) { + const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu; + reg[i/4][i%4] = d * ds4_metal_mxfp4_values[q]; + } + return; + } + const half dh = (half)ds4_metal_e8m0_to_f32(e); + FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) { + const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu; + reg[i/4][i%4] = dh * ds4_metal_mxfp4_half_values[q]; + } +} + +// Test-only raw-bit oracle for the half-scale path: covers all 4096 +// exponent/code pairs including both sides of the band boundary. +// Test-only raw-bit oracle. Keeping both conversions in one GPU invocation +// covers all 4096 exponent/code pairs, including the 0xff fallback row, +// without depending on host floating-point conversion behavior. +kernel void kernel_test_mxfp4_down_half_lut( + device ushort *legacy [[buffer(0)]], + device ushort *lut [[buffer(1)]], + uint tid [[thread_position_in_grid]]) { + if (tid >= 4096u) { + return; + } + const uchar e = (uchar)(tid >> 4u); + const uint q = tid & 0x0fu; + half legacy_value; + legacy_value = ds4_metal_e8m0_to_f32(e) * + ds4_metal_mxfp4_values[q]; + legacy[tid] = as_type(legacy_value); + lut[tid] = as_type(ds4_metal_mxfp4_half_lut_value(e, q)); +} + template void dequantize_q5_K(device const block_q5_K *xb, short il, thread type4x4 ®) { const short group = il / 2; @@ -3208,6 +2861,72 @@ void kernel_mul_mv_q4_K_f32_impl( (void)shmem; } +template +void kernel_mul_mv_mxfp4_f32_impl( + args_t 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; + const int nb = args.ne00 / QK_MXFP4; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint i12 = tgpig.z % args.ne12; + const uint i13 = tgpig.z / args.ne12; + const uint64_t offset0 = (uint64_t)first_row * args.nb01 + + (uint64_t)(i12 / args.r2) * args.nb02 + (uint64_t)(i13 / args.r3) * args.nb03; + const uint64_t offset1 = (uint64_t)tgpig.y * args.nb11 + + (uint64_t)i12 * args.nb12 + (uint64_t)i13 * args.nb13; + + device const block_mxfp4 *x = (device const block_mxfp4 *)(src0 + offset0); + device const float *y = (device const float *)(src1 + offset1); + const int row_blocks = (int)(args.nb01 / sizeof(block_mxfp4)); + const short ix = tiisg / 2; + const short it = tiisg & 1; + + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float sumf[nr0] = {0.f}; + device const float *yb = y + ix * QK_MXFP4 + it * 8; + for (int ib = ix; ib < nb; ib += 16) { + device const float4 *y4 = (device const float4 *)yb; + const float4 yl0 = y4[0]; + const float4 yl1 = y4[4]; + const float4 yl2 = y4[1]; + const float4 yl3 = y4[5]; + + FOR_UNROLL (short row = 0; row < nr0; row++) { + device const block_mxfp4 &xb = x[row * row_blocks + ib]; + device const uchar *q = xb.qs + 8 * it; + float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15], + lut[q[2] & 15], lut[q[3] & 15]); + acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4], + lut[q[2] >> 4], lut[q[3] >> 4]); + acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15], + lut[q[6] & 15], lut[q[7] & 15]); + acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4], + lut[q[6] >> 4], lut[q[7] >> 4]); + sumf[row] += ds4_metal_e8m0_to_f32(xb.e) * + ((acc.x + acc.y) + (acc.z + acc.w)); + } + yb += 16 * QK_MXFP4; + } + + device float *dst_f32 = (device float *)dst + + (uint64_t)tgpig.z * args.ne0 * args.ne1 + (uint64_t)tgpig.y * args.ne0; + FOR_UNROLL (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + const float sum = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum; + } + } +} + template void kernel_mul_mv_q8_K_f32_impl( args_t args, @@ -3606,6 +3325,7 @@ template [[host_name("kernel_mul_mv_id_q2_K_f32")]] kernel kernel_mul_mv_id_q template [[host_name("kernel_mul_mv_id_q4_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q8_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_mxfp4_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; // Plain dense Q4_K matvec on the classic impl. The mul_mv_ext family used // by the generic dense path tops out around 220 GB/s on M5 for the GLM @@ -3938,29 +3658,184 @@ kernel void kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32( const float c = act.clamp_value; const float route_weight = route_w[0]; - for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) { - const float sum_gate = simd_sum(sumg[row]); - const float sum_up = simd_sum(sumu[row]); - if (tiisg == 0) { - const uint out_row = first_row + row; - const float gate = sum_gate * 0.25f; - const float up = sum_up * 0.25f; - float g = gate; - float u = up; - if (c > 1.0e-6f) { - g = min(g, c); - u = clamp(u, -c, c); - } - dst_gate_f32[out_row] = gate; - dst_up_f32[out_row] = up; - const float silu = g / (1.0f + exp(-g)); - dst_mid_f32[out_row] = silu * u * route_weight; + float4 reduced_gate; + float4 reduced_up; + for (int row = 0; row < N_R0_IQ2_XXS; ++row) { + reduced_gate[row] = simd_sum(sumg[row]); + reduced_up[row] = simd_sum(sumu[row]); + } + if (tiisg < N_R0_IQ2_XXS && first_row + tiisg < args.ne0) { + const uint out_row = first_row + tiisg; + const float gate = reduced_gate[tiisg] * 0.25f; + const float up = reduced_up[tiisg] * 0.25f; + float g = gate; + float u = up; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); } + dst_gate_f32[out_row] = gate; + dst_up_f32[out_row] = up; + const float silu = g / (1.0f + exp(-g)); + dst_mid_f32[out_row] = silu * u * route_weight; } (void)tiitg; } +kernel void kernel_mul_mv_id_iq2_xxs_pair_swiglu_pack2_overlap_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short NSG = 4; + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *) (ids + iid1 * args.nbi1))[idx]; + if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) return; + const int i02b = i02 - args.tp_expert_base; + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS; + const int nb32 = nb * (QK_K / 32); + + device const block_iq2_xxs *xg = + (device const block_iq2_xxs *)(src0_gate + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01); + device const block_iq2_xxs *xu = + (device const block_iq2_xxs *)(src0_up + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01); + device const float *y = + (device const float *)(src1 + i11 * args.nb11 + i12 * args.nb12); + + float yl[32]; + float sumg[N_R0_IQ2_XXS] = {0.f}; + float sumu[N_R0_IQ2_XXS] = {0.f}; + + const int ix = tiisg; + device const float *y4 = y + 32 * ix; + bool yl_preloaded = sgitg >= 2; + if (yl_preloaded) { + for (short i = 0; i < 32; ++i) { + yl[i] = y4[i]; + } + } + + threadgroup uint64_t *svalues = (threadgroup uint64_t *)(shmem); + threadgroup uint8_t *ssigns = (threadgroup uint8_t *)(svalues + 256); + if (sgitg < 2) { + int nval = 4; + int pos = (32 * sgitg + tiisg) * nval; + for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i]; + nval = 2; + pos = (32 * sgitg + tiisg) * nval; + for (int i = 0; i < nval; ++i) ssigns[pos + i] = ds4_metal_ksigns_iq2xs[pos + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int ib32 = ix; ib32 < nb32; ib32 += 32) { + if (!yl_preloaded) { + for (short i = 0; i < 32; ++i) { + yl[i] = y4[i]; + } + } + yl_preloaded = false; + + const int ibl = ib32 / (QK_K / 32); + const int ib = ib32 % (QK_K / 32); + + device const block_iq2_xxs *xgr = xg + ibl; + device const block_iq2_xxs *xur = xu + ibl; + device const uint16_t *qg = xgr->qs + 4 * ib; + device const uint16_t *qu = xur->qs + 4 * ib; + device const half *dhg = &xgr->d; + device const half *dhu = &xur->d; + + for (short row = 0; row < N_R0_IQ2_XXS; row++) { + device const uint8_t *aux8g = (device const uint8_t *)qg; + device const uint8_t *aux8u = (device const uint8_t *)qu; + const uint32_t aux32g = qg[2] | (qg[3] << 16); + const uint32_t aux32u = qu[2] | (qu[3] << 16); + const float dg = (float)dhg[0] * (0.5f + (aux32g >> 28)); + const float du = (float)dhu[0] * (0.5f + (aux32u >> 28)); + + float sg = 0; + float su = 0; + for (short l = 0; l < 4; ++l) { + const threadgroup uint8_t *gridg = (const threadgroup uint8_t *)(svalues + aux8g[l]); + const threadgroup uint8_t *gridu = (const threadgroup uint8_t *)(svalues + aux8u[l]); + const uint8_t signg = ssigns[(aux32g >> 7 * l) & 127]; + const uint8_t signu = ssigns[(aux32u >> 7 * l) & 127]; + for (short j = 0; j < 8; ++j) { + const float v = yl[8 * l + j]; + sg += v * gridg[j] * (signg & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + su += v * gridu[j] * (signu & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + } + } + sumg[row] += dg * sg; + sumu[row] += du * su; + + dhg += args.nb01 / 2; + dhu += args.nb01 / 2; + qg += args.nb01 / 2; + qu += args.nb01 / 2; + } + + y4 += 32 * 32; + } + + device float *dst_gate_f32 = + (device float *)dst_gate + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0; + device float *dst_up_f32 = + (device float *)dst_up + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *dst_mid_f32 = + (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = + (device const float *)(weights + pair_row * act.weight_stride); + + const float c = act.clamp_value; + const float route_weight = route_w[0]; + float4 reduced_gate; + float4 reduced_up; + for (int row = 0; row < N_R0_IQ2_XXS; ++row) { + reduced_gate[row] = simd_sum(sumg[row]); + reduced_up[row] = simd_sum(sumu[row]); + } + if (tiisg < N_R0_IQ2_XXS && first_row + tiisg < args.ne0) { + const uint out_row = first_row + tiisg; + const float gate = reduced_gate[tiisg] * 0.25f; + const float up = reduced_up[tiisg] * 0.25f; + float g = gate; + float u = up; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + dst_gate_f32[out_row] = gate; + dst_up_f32[out_row] = up; + const float silu = g / (1.0f + exp(-g)); + dst_mid_f32[out_row] = silu * u * route_weight; + } + + (void)tiitg; +} + + kernel void kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32( constant ds4_metal_args_mul_mv_id & args, constant ds4_metal_dsv4_moe_swiglu_weight_args & act, @@ -4579,6 +4454,376 @@ kernel void kernel_mul_mv_id_q4_K_pair_swiglu_f32( (void)tiitg; } +template +void kernel_mul_mv_mxfp4_pair_swiglu_impl( + args_t args, + constant ds4_metal_dsv4_moe_swiglu_weight_args &act, + device const char *src0_gate, + device const char *src0_up, + device const char *src1, + device char *dst_gate, + device char *dst_up, + device char *dst_mid, + float route_weight, + threadgroup char *shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_MXFP4; + const int nb = args.ne00 / QK_MXFP4; + const int row_blocks = (int)(args.nb01 / sizeof(block_mxfp4)); + const short ix = tiisg / 2; + const short it = tiisg & 1; + + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + device const block_mxfp4 *xg = + (device const block_mxfp4 *)(src0_gate + (uint64_t)first_row * args.nb01); + device const block_mxfp4 *xu = + (device const block_mxfp4 *)(src0_up + (uint64_t)first_row * args.nb01); + device const float *yb = (device const float *)src1 + ix * QK_MXFP4 + it * 8; + float sumg[N_R0_MXFP4] = {0.f}; + float sumu[N_R0_MXFP4] = {0.f}; + + for (int ib = ix; ib < nb; ib += 16) { + device const float4 *y4 = (device const float4 *)yb; + const float4 yl0 = y4[0]; + const float4 yl1 = y4[4]; + const float4 yl2 = y4[1]; + const float4 yl3 = y4[5]; + + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + device const block_mxfp4 &bg = xg[row * row_blocks + ib]; + device const block_mxfp4 &bu = xu[row * row_blocks + ib]; + device const uchar *qg = bg.qs + 8 * it; + device const uchar *qu = bu.qs + 8 * it; + + float4 ag = yl0 * float4(lut[qg[0] & 15], lut[qg[1] & 15], + lut[qg[2] & 15], lut[qg[3] & 15]); + ag += yl1 * float4(lut[qg[0] >> 4], lut[qg[1] >> 4], + lut[qg[2] >> 4], lut[qg[3] >> 4]); + ag += yl2 * float4(lut[qg[4] & 15], lut[qg[5] & 15], + lut[qg[6] & 15], lut[qg[7] & 15]); + ag += yl3 * float4(lut[qg[4] >> 4], lut[qg[5] >> 4], + lut[qg[6] >> 4], lut[qg[7] >> 4]); + + float4 au = yl0 * float4(lut[qu[0] & 15], lut[qu[1] & 15], + lut[qu[2] & 15], lut[qu[3] & 15]); + au += yl1 * float4(lut[qu[0] >> 4], lut[qu[1] >> 4], + lut[qu[2] >> 4], lut[qu[3] >> 4]); + au += yl2 * float4(lut[qu[4] & 15], lut[qu[5] & 15], + lut[qu[6] & 15], lut[qu[7] & 15]); + au += yl3 * float4(lut[qu[4] >> 4], lut[qu[5] >> 4], + lut[qu[6] >> 4], lut[qu[7] >> 4]); + + sumg[row] += ds4_metal_e8m0_to_f32(bg.e) * + ((ag.x + ag.y) + (ag.z + ag.w)); + sumu[row] += ds4_metal_e8m0_to_f32(bu.e) * + ((au.x + au.y) + (au.z + au.w)); + } + yb += 16 * QK_MXFP4; + } + + device float *gate_f32 = (device float *)dst_gate; + device float *up_f32 = (device float *)dst_up; + device float *mid_f32 = (device float *)dst_mid; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + if (first_row + row < args.ne0) { + const float gate = simd_sum(sumg[row]); + const float up = simd_sum(sumu[row]); + if (tiisg == 0) { + const uint out_row = first_row + row; + float g = gate; + float u = up; + if (act.clamp_value > 1.0e-6f) { + g = min(g, act.clamp_value); + u = clamp(u, -act.clamp_value, act.clamp_value); + } + gate_f32[out_row] = gate; + up_f32[out_row] = up; + mid_f32[out_row] = (g / (1.0f + exp(-g))) * u * route_weight; + } + } + } +} + +kernel void kernel_mul_mv_id_mxfp4_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id &args, + constant ds4_metal_dsv4_moe_swiglu_weight_args &act, + device const char *src0_gate, + device const char *src0_up, + device const char *src1, + device char *dst_gate, + device char *dst_up, + device char *dst_mid, + device const char *ids, + device const char *weights, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + const int32_t expert = ((device const int32_t *)(ids + (uint64_t)iid1 * args.nbi1))[idx]; + if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) return; + + const uint64_t pair_row = (uint64_t)iid1 * args.nei0 + (uint64_t)idx; + device const float *route = + (device const float *)(weights + pair_row * act.weight_stride); + device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float); + device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float); + device char *mid_cur = dst_mid + pair_row * act.mid_row_stride; + device const char *x_cur = src1 + (uint64_t)(idx % args.ne11) * args.nb11 + + (uint64_t)iid1 * args.nb12; + device const char *gate_expert = src0_gate + + (int64_t)(expert - args.tp_expert_base) * args.nb02; + device const char *up_expert = src0_up + + (int64_t)(expert - args.tp_expert_base) * args.nb02; + tgpig.z = 0; + kernel_mul_mv_mxfp4_pair_swiglu_impl(args, act, gate_expert, up_expert, + x_cur, gate_cur, up_cur, mid_cur, + route[0], shmem, tgpig, tiisg, sgitg); + (void)tiitg; +} + +kernel void kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_f32( + constant ds4_metal_args_mul_mv_id &args, + constant ds4_metal_dsv4_moe_swiglu_weight_args &act, + device const char *src0_gate, + device const char *src0_up, + device const char *src1, + device char *dst_gate, + device char *dst_up, + device char *dst_mid, + device const char *ids, + device const char *weights, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const int idx = (int)tgpig.z; + const int32_t expert = ((device const int32_t *)ids)[idx]; + const uint64_t pair_row = (uint64_t)idx; + device const float *route = + (device const float *)(weights + pair_row * act.weight_stride); + device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float); + device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float); + device char *mid_cur = dst_mid + pair_row * act.mid_row_stride; + device const char *gate_expert = + src0_gate + (int64_t)expert * args.nb02; + device const char *up_expert = + src0_up + (int64_t)expert * args.nb02; + tgpig.z = 0; + kernel_mul_mv_mxfp4_pair_swiglu_impl(args, act, gate_expert, up_expert, + src1, gate_cur, up_cur, mid_cur, + route[0], shmem, tgpig, tiisg, sgitg); + (void)tiitg; +} + +/* Exact-shape sibling of the fixed-route decode pair-SwiGLU kernel. The host + * selects it only after proving expert_in_dim == 4096 and gate/up row bytes == + * 2176, so the per-lane K walk has a compile-time trip count (128 blocks / 16 + * lane pairs = 8 steps) and the row stride is a literal. That lets the + * compiler unroll the walk and keep several independent expert loads in + * flight; the per-lane block order, the per-block accumulate order and the + * simd_sum tree are byte-identical to kernel_mul_mv_mxfp4_pair_swiglu_impl. */ +#define DS4_MXFP4_PAIR_STATIC_NB 128 +#define DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS 128 + +template +void kernel_mul_mv_mxfp4_pair_swiglu_static_impl( + args_t args, + constant ds4_metal_dsv4_moe_swiglu_weight_args &act, + device const char *src0_gate, + device const char *src0_up, + device const char *src1, + device char *dst_gate, + device char *dst_up, + device char *dst_mid, + float route_weight, + threadgroup char *shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_MXFP4; + const short ix = tiisg / 2; + const short it = tiisg & 1; + + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + device const block_mxfp4 *xg = + (device const block_mxfp4 *)(src0_gate + (uint64_t)first_row * args.nb01); + device const block_mxfp4 *xu = + (device const block_mxfp4 *)(src0_up + (uint64_t)first_row * args.nb01); + device const float *yb = (device const float *)src1 + ix * QK_MXFP4 + it * 8; + float sumg[N_R0_MXFP4] = {0.f}; + float sumu[N_R0_MXFP4] = {0.f}; + + for (int ib = ix; ib < DS4_MXFP4_PAIR_STATIC_NB; ib += 16) { + device const float4 *y4 = (device const float4 *)yb; + const float4 yl0 = y4[0]; + const float4 yl1 = y4[4]; + const float4 yl2 = y4[1]; + const float4 yl3 = y4[5]; + + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + device const block_mxfp4 &bg = + xg[row * DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS + ib]; + device const block_mxfp4 &bu = + xu[row * DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS + ib]; + device const uchar *qg = bg.qs + 8 * it; + device const uchar *qu = bu.qs + 8 * it; + + float4 ag = yl0 * float4(lut[qg[0] & 15], lut[qg[1] & 15], + lut[qg[2] & 15], lut[qg[3] & 15]); + ag += yl1 * float4(lut[qg[0] >> 4], lut[qg[1] >> 4], + lut[qg[2] >> 4], lut[qg[3] >> 4]); + ag += yl2 * float4(lut[qg[4] & 15], lut[qg[5] & 15], + lut[qg[6] & 15], lut[qg[7] & 15]); + ag += yl3 * float4(lut[qg[4] >> 4], lut[qg[5] >> 4], + lut[qg[6] >> 4], lut[qg[7] >> 4]); + + float4 au = yl0 * float4(lut[qu[0] & 15], lut[qu[1] & 15], + lut[qu[2] & 15], lut[qu[3] & 15]); + au += yl1 * float4(lut[qu[0] >> 4], lut[qu[1] >> 4], + lut[qu[2] >> 4], lut[qu[3] >> 4]); + au += yl2 * float4(lut[qu[4] & 15], lut[qu[5] & 15], + lut[qu[6] & 15], lut[qu[7] & 15]); + au += yl3 * float4(lut[qu[4] >> 4], lut[qu[5] >> 4], + lut[qu[6] >> 4], lut[qu[7] >> 4]); + + sumg[row] += ds4_metal_e8m0_to_f32(bg.e) * + ((ag.x + ag.y) + (ag.z + ag.w)); + sumu[row] += ds4_metal_e8m0_to_f32(bu.e) * + ((au.x + au.y) + (au.z + au.w)); + } + yb += 16 * QK_MXFP4; + } + + device float *gate_f32 = (device float *)dst_gate; + device float *up_f32 = (device float *)dst_up; + device float *mid_f32 = (device float *)dst_mid; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + if (first_row + row < args.ne0) { + const float gate = simd_sum(sumg[row]); + const float up = simd_sum(sumu[row]); + if (tiisg == 0) { + const uint out_row = first_row + row; + float g = gate; + float u = up; + if (act.clamp_value > 1.0e-6f) { + g = min(g, act.clamp_value); + u = clamp(u, -act.clamp_value, act.clamp_value); + } + gate_f32[out_row] = gate; + up_f32[out_row] = up; + mid_f32[out_row] = (g / (1.0f + exp(-g))) * u * route_weight; + } + } + } +} + +kernel void kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_f32( + constant ds4_metal_args_mul_mv_id &args, + constant ds4_metal_dsv4_moe_swiglu_weight_args &act, + device const char *src0_gate, + device const char *src0_up, + device const char *src1, + device char *dst_gate, + device char *dst_up, + device char *dst_mid, + device const char *ids, + device const char *weights, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const int idx = (int)tgpig.z; + const int32_t expert = ((device const int32_t *)ids)[idx]; + const uint64_t pair_row = (uint64_t)idx; + device const float *route = + (device const float *)(weights + pair_row * act.weight_stride); + device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float); + device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float); + device char *mid_cur = dst_mid + pair_row * act.mid_row_stride; + device const char *gate_expert = + src0_gate + (int64_t)expert * args.nb02; + device const char *up_expert = + src0_up + (int64_t)expert * args.nb02; + tgpig.z = 0; + /* Defensive: the host proves this shape before selecting the pipeline, so + * the dynamic sibling here is unreachable in production. Keeping it makes + * a mis-selected pipeline produce identical results instead of garbage. */ + if (args.ne00 == DS4_MXFP4_PAIR_STATIC_NB * QK_MXFP4 && + args.nb01 == (uint64_t)DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS * + sizeof(block_mxfp4)) { + kernel_mul_mv_mxfp4_pair_swiglu_static_impl( + args, act, gate_expert, up_expert, src1, gate_cur, up_cur, mid_cur, + route[0], shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_mxfp4_pair_swiglu_impl( + args, act, gate_expert, up_expert, src1, gate_cur, up_cur, mid_cur, + route[0], shmem, tgpig, tiisg, sgitg); + } + (void)tiitg; +} + +kernel void kernel_mul_mv_slots6_mxfp4_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id &args, + constant ds4_metal_dsv4_moe_swiglu_weight_args &act, + device const char *gate0, device const char *gate1, + device const char *gate2, device const char *gate3, + device const char *gate4, device const char *gate5, + device const char *up0, device const char *up1, + device const char *up2, device const char *up3, + device const char *up4, device const char *up5, + device const char *src1, + device char *dst_gate, + device char *dst_up, + device char *dst_mid, + device const char *weights, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + device const char *gate_expert = gate0; + device const char *up_expert = up0; + switch (idx) { + case 1: gate_expert = gate1; up_expert = up1; break; + case 2: gate_expert = gate2; up_expert = up2; break; + case 3: gate_expert = gate3; up_expert = up3; break; + case 4: gate_expert = gate4; up_expert = up4; break; + case 5: gate_expert = gate5; up_expert = up5; break; + default: break; + } + + const uint64_t pair_row = (uint64_t)iid1 * args.nei0 + (uint64_t)idx; + device const float *route = + (device const float *)(weights + pair_row * act.weight_stride); + device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float); + device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float); + device char *mid_cur = dst_mid + pair_row * act.mid_row_stride; + device const char *x_cur = src1 + (uint64_t)(idx % args.ne11) * args.nb11 + + (uint64_t)iid1 * args.nb12; + tgpig.z = 0; + kernel_mul_mv_mxfp4_pair_swiglu_impl(args, act, gate_expert, up_expert, + x_cur, gate_cur, up_cur, mid_cur, + route[0], shmem, tgpig, tiisg, sgitg); + (void)tiitg; +} + kernel void kernel_mul_mv_table_q4_K_pair_swiglu_f32( constant ds4_metal_args_mul_mv_id & args, constant ds4_metal_dsv4_moe_swiglu_weight_args & act, @@ -5986,6 +6231,392 @@ kernel void kernel_mul_mv_addr_q2_K_sum6_masked_f32( (void)tgpig; } +static inline float2 ds4_mxfp4_accumulate_rows( + device const char *src0, + uint64_t row_bytes, + device const float *y, + uint32_t n_cols, + uint32_t first_row, + uint32_t n_rows, + threadgroup const float *lut, + ushort tiisg) { + const int nb = (int)(n_cols / QK_MXFP4); + const int row_blocks = (int)(row_bytes / sizeof(block_mxfp4)); + const short ix = tiisg / 2; + const short it = tiisg & 1; + device const block_mxfp4 *x = + (device const block_mxfp4 *)(src0 + (uint64_t)first_row * row_bytes); + device const float *yb = y + ix * QK_MXFP4 + it * 8; + float2 sums = 0.0f; + + for (int ib = ix; ib < nb; ib += 16) { + device const float4 *y4 = (device const float4 *)yb; + const float4 yl0 = y4[0]; + const float4 yl1 = y4[4]; + const float4 yl2 = y4[1]; + const float4 yl3 = y4[5]; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + if (first_row + row < n_rows) { + device const block_mxfp4 &b = x[row * row_blocks + ib]; + device const uchar *q = b.qs + 8 * it; + float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15], + lut[q[2] & 15], lut[q[3] & 15]); + acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4], + lut[q[2] >> 4], lut[q[3] >> 4]); + acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15], + lut[q[6] & 15], lut[q[7] & 15]); + acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4], + lut[q[6] >> 4], lut[q[7] >> 4]); + sums[row] += ds4_metal_e8m0_to_f32(b.e) * + ((acc.x + acc.y) + (acc.z + acc.w)); + } + } + yb += 16 * QK_MXFP4; + } + return sums; +} + +// Exact-shape sibling for the resident Flash decode down projection. Its +// NR2/nsg1 grid covers all 4096 rows without a tail, so retain the established +// per-lane K walk and arithmetic while omitting only the redundant row guard. +static inline float2 ds4_mxfp4_accumulate_full_rows( + device const char *src0, + uint64_t row_bytes, + device const float *y, + uint32_t n_cols, + uint32_t first_row, + threadgroup const float *lut, + ushort tiisg) { + const int nb = (int)(n_cols / QK_MXFP4); + const int row_blocks = (int)(row_bytes / sizeof(block_mxfp4)); + const short ix = tiisg / 2; + const short it = tiisg & 1; + device const block_mxfp4 *x = + (device const block_mxfp4 *)(src0 + (uint64_t)first_row * row_bytes); + device const float *yb = y + ix * QK_MXFP4 + it * 8; + float2 sums = 0.0f; + + for (int ib = ix; ib < nb; ib += 16) { + device const float4 *y4 = (device const float4 *)yb; + const float4 yl0 = y4[0]; + const float4 yl1 = y4[4]; + const float4 yl2 = y4[1]; + const float4 yl3 = y4[5]; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + device const block_mxfp4 &b = x[row * row_blocks + ib]; + device const uchar *q = b.qs + 8 * it; + float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15], + lut[q[2] & 15], lut[q[3] & 15]); + acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4], + lut[q[2] >> 4], lut[q[3] >> 4]); + acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15], + lut[q[6] & 15], lut[q[7] & 15]); + acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4], + lut[q[6] >> 4], lut[q[7] >> 4]); + sums[row] += ds4_metal_e8m0_to_f32(b.e) * + ((acc.x + acc.y) + (acc.z + acc.w)); + } + yb += 16 * QK_MXFP4; + } + return sums; +} + +kernel void kernel_mul_mv_id_mxfp4_sum6_f32( + constant ds4_metal_args_mul_mv_id &args, + device const char *src0s, + device const char *src1, + device char *dst, + device const char *ids, + device const char *add_in, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4); + const uint32_t token = tgpig.y; + device const int32_t *token_ids = + (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float2 sumf = 0.0f; + for (int slot = 0; slot < args.nei0; slot++) { + const int32_t expert = token_ids[slot]; + if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue; + device const char *expert_base = src0s + + (int64_t)(expert - args.tp_expert_base) * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_rows(expert_base, args.nb01, y, + args.ne00, first_row, args.ne0, + lut, tiisg); + } + + device float *out = (device float *)(dst + (uint64_t)token * args.nb1); + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + if (first_row + row < (uint32_t)args.ne0) { + const float value = simd_sum(sumf[row]); + if (tiisg == 0) { + out[first_row + row] = value + + (args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f); + } + } + } + (void)tiitg; +} + +kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_f32( + constant ds4_metal_args_mul_mv_id &args, + device const char *src0s, + device const char *src1, + device char *dst, + device const char *ids, + device const char *add_in, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4); + device const int32_t *token_ids = (device const int32_t *)ids; + device const char *token_src1 = src1; + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float2 sumf = 0.0f; + for (int slot = 0; slot < args.nei0; slot++) { + const int32_t expert = token_ids[slot]; + device const char *expert_base = + src0s + (int64_t)expert * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_rows(expert_base, args.nb01, y, + args.ne00, first_row, args.ne0, + lut, tiisg); + } + + device float *out = (device float *)dst; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + if (first_row + row < (uint32_t)args.ne0) { + const float value = simd_sum(sumf[row]); + if (tiisg == 0) { + out[first_row + row] = value + + (args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f); + } + } + } + (void)tiitg; +} + +kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_f32( + constant ds4_metal_args_mul_mv_id &args, + device const char *src0s, + device const char *src1, + device char *dst, + device const char *ids, + device const char *add_in, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4); + device const int32_t *token_ids = (device const int32_t *)ids; + device const char *token_src1 = src1; + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float2 sumf = 0.0f; + for (int slot = 0; slot < args.nei0; slot++) { + const int32_t expert = token_ids[slot]; + device const char *expert_base = + src0s + (int64_t)expert * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_full_rows( + expert_base, args.nb01, y, args.ne00, first_row, lut, tiisg); + } + + device float *out = (device float *)dst; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + const float value = simd_sum(sumf[row]); + if (tiisg == 0) { + out[first_row + row] = value + + (args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f); + } + } + (void)tiitg; +} + +/* Exact-shape sibling of the fixed-route full-rows decode down projection. + * The host proves ne00 == 2048, row bytes == 1088, nei0 == 6 and ne0 == 4096 + * before selecting it, so the K walk (64 blocks / 16 lane pairs = 4 steps), + * the row stride and the routed-slot count are all literals here. The lane + * mapping, the block visit order, the per-block accumulate order and the + * simd_sum tree are byte-identical to ds4_mxfp4_accumulate_full_rows. */ +#define DS4_MXFP4_DOWN_STATIC_NB 64 +#define DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS 64 +#define DS4_MXFP4_DOWN_STATIC_SLOTS 6 + +static inline float2 ds4_mxfp4_accumulate_full_rows_static( + device const char *src0, + device const float *y, + uint32_t first_row, + threadgroup const float *lut, + ushort tiisg) { + const short ix = tiisg / 2; + const short it = tiisg & 1; + device const block_mxfp4 *x = + (device const block_mxfp4 *)(src0 + + (uint64_t)first_row * (DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS * + sizeof(block_mxfp4))); + device const float *yb = y + ix * QK_MXFP4 + it * 8; + float2 sums = 0.0f; + + for (int ib = ix; ib < DS4_MXFP4_DOWN_STATIC_NB; ib += 16) { + device const float4 *y4 = (device const float4 *)yb; + const float4 yl0 = y4[0]; + const float4 yl1 = y4[4]; + const float4 yl2 = y4[1]; + const float4 yl3 = y4[5]; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + device const block_mxfp4 &b = + x[row * DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS + ib]; + device const uchar *q = b.qs + 8 * it; + float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15], + lut[q[2] & 15], lut[q[3] & 15]); + acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4], + lut[q[2] >> 4], lut[q[3] >> 4]); + acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15], + lut[q[6] & 15], lut[q[7] & 15]); + acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4], + lut[q[6] >> 4], lut[q[7] >> 4]); + sums[row] += ds4_metal_e8m0_to_f32(b.e) * + ((acc.x + acc.y) + (acc.z + acc.w)); + } + yb += 16 * QK_MXFP4; + } + return sums; +} + +kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32( + constant ds4_metal_args_mul_mv_id &args, + device const char *src0s, + device const char *src1, + device char *dst, + device const char *ids, + device const char *add_in, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4); + device const int32_t *token_ids = (device const int32_t *)ids; + device const char *token_src1 = src1; + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + /* Defensive: unreachable in production because the host proves the shape + * before selecting this pipeline; identical results either way. */ + const bool static_shape = + args.ne00 == DS4_MXFP4_DOWN_STATIC_NB * QK_MXFP4 && + args.nb01 == (uint64_t)DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS * + sizeof(block_mxfp4) && + args.nei0 == DS4_MXFP4_DOWN_STATIC_SLOTS; + + float2 sumf = 0.0f; + if (static_shape) { + for (short slot = 0; slot < DS4_MXFP4_DOWN_STATIC_SLOTS; slot++) { + const int32_t expert = token_ids[slot]; + device const char *expert_base = + src0s + (int64_t)expert * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_full_rows_static( + expert_base, y, first_row, lut, tiisg); + } + } else { + for (int slot = 0; slot < args.nei0; slot++) { + const int32_t expert = token_ids[slot]; + device const char *expert_base = + src0s + (int64_t)expert * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_full_rows( + expert_base, args.nb01, y, args.ne00, first_row, lut, tiisg); + } + } + + device float *out = (device float *)dst; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + const float value = simd_sum(sumf[row]); + if (tiisg == 0) { + out[first_row + row] = value + + (args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f); + } + } + (void)tiitg; +} + +kernel void kernel_mul_mv_slots6_mxfp4_sum6_f32( + constant ds4_metal_args_mul_mv_id &args, + device const char *src00, device const char *src01, + device const char *src02, device const char *src03, + device const char *src04, device const char *src05, + device const char *src1, + device char *dst, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4); + const uint32_t token = tgpig.y; + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float2 sumf = 0.0f; + for (int slot = 0; slot < 6; slot++) { + device const char *expert_base = src00; + switch (slot) { + case 1: expert_base = src01; break; + case 2: expert_base = src02; break; + case 3: expert_base = src03; break; + case 4: expert_base = src04; break; + case 5: expert_base = src05; break; + default: break; + } + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_rows(expert_base, args.nb01, y, + args.ne00, first_row, args.ne0, + lut, tiisg); + } + + device float *out = (device float *)(dst + (uint64_t)token * args.nb1); + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + if (first_row + row < (uint32_t)args.ne0) { + const float value = simd_sum(sumf[row]); + if (tiisg == 0) out[first_row + row] = value; + } + } + (void)tiitg; +} + kernel void kernel_mul_mv_id_q4_K_sum6_f32( constant ds4_metal_args_mul_mv_id & args, device const char * src0s, @@ -6986,6 +7617,7 @@ kernel void kernel_mul_mm_id_map0( device const char * src2, device char * htpe, device char * hids, + device char * work, threadgroup char * shmem [[threadgroup(0)]], ushort tpitg[[thread_position_in_threadgroup]], ushort ntg[[threads_per_threadgroup]]) { @@ -7032,6 +7664,29 @@ kernel void kernel_mul_mm_id_map0( device uint32_t * tpe_u32 = (device uint32_t *) (htpe); tpe_u32[ide] = n_all; + + // Reuse the route-id staging memory after the map is complete to build a + // compact list of non-empty 32-row matmul tiles. The old dispatch covered + // every possible token tile for every expert, even though most experts + // receive only a small fraction of the prompt rows. + threadgroup uint16_t * tile_counts = (threadgroup uint16_t *) shmem; + const uint16_t n_tiles = (uint16_t)((n_all + 31u) / 32u); + tile_counts[ide] = n_tiles; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint32_t tile_base = 0; + for (ushort i = 0; i < ide; i++) { + tile_base += tile_counts[i]; + } + + device uint32_t * work_count = (device uint32_t *) work; + device uint2 * work_items = (device uint2 *)(work + 8); + for (uint32_t tile = 0; tile < n_tiles; tile++) { + work_items[tile_base + tile] = uint2((uint32_t)ide, tile * 32u); + } + if (ide + 1u == ntg) { + work_count[0] = tile_base + n_tiles; + } } typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t; @@ -7048,10 +7703,97 @@ template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_ template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>; template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>; +// Token-centric map builder for the common six-route MXFP4 prefill shape. +// Each selected route reserves one expert-local row instead of making every +// expert scan every token. Expert-local order is intentionally unspecified: +// hids carries the original token/slot id, so every downstream row still +// writes its distinct fixed destination. The work trailer exactly matches the +// established direct-dispatch ABI used by the mapped pair and down kernels. +template +kernel void kernel_mul_mm_id_map_scatter_work( + constant ds4_metal_args_mul_mm_id_map0 & args, + device const char * src2, + device char * htpe, + device char * hids, + device char * work, + threadgroup char * shmem [[threadgroup(0)]], + ushort tpitg[[thread_position_in_threadgroup]], + ushort ntg[[threads_per_threadgroup]]) { + threadgroup atomic_uint * counts = + (threadgroup atomic_uint *) shmem; + threadgroup uint16_t * tile_counts = + (threadgroup uint16_t *)(shmem + + (uint32_t)args.ne02*sizeof(uint32_t)); + device uint32_t * tpe_u32 = (device uint32_t *) htpe; + device int32_t * ids_i32 = (device int32_t *) hids; + + if (tpitg < args.ne02) { + atomic_store_explicit(counts + tpitg, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int i21 = tpitg; i21 < args.ne21; i21 += ntg) { + device const int32_t * src2_i32 = + (device const int32_t *)(src2 + i21*args.nb21); + + #pragma unroll(ne20) + for (short i20 = 0; i20 < ne20; i20++) { + const int32_t expert = src2_i32[i20]; + if ((uint32_t)expert >= (uint32_t)args.ne02) { + continue; + } + + const uint32_t row = atomic_fetch_add_explicit( + counts + expert, 1u, memory_order_relaxed); + // Production top-k selections are unique. Keep malformed or + // synthetic duplicates from exceeding the fixed expert slice. + if (row < (uint32_t)args.ne21) { + ids_i32[(uint32_t)expert*args.ne21 + row] = + i21*ne20 + i20; + } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const short ide = tpitg; + const uint32_t n_all = min( + atomic_load_explicit(counts + ide, memory_order_relaxed), + (uint32_t)args.ne21); + tpe_u32[ide] = n_all; + + const uint16_t n_tiles = (uint16_t)((n_all + 31u) / 32u); + tile_counts[ide] = n_tiles; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint32_t tile_base = 0; + for (ushort i = 0; i < ide; i++) { + tile_base += tile_counts[i]; + } + + device uint32_t * work_count = (device uint32_t *) work; + device uint2 * work_items = (device uint2 *)(work + 8); + for (uint32_t tile = 0; tile < n_tiles; tile++) { + work_items[tile_base + tile] = + uint2((uint32_t)ide, tile * 32u); + } + if (ide + 1u == ntg) { + // Direct consumers launch the padded work capacity and cull against + // this exact count. Zero work must overwrite any stale prior value. + work_count[0] = tile_base + n_tiles; + } +} + +typedef decltype(kernel_mul_mm_id_map_scatter_work<6>) + kernel_mul_mm_id_map_scatter_work_t; +template [[host_name("kernel_mul_mm_id_map_scatter_work_ne20_6")]] +kernel kernel_mul_mm_id_map_scatter_work_t + kernel_mul_mm_id_map_scatter_work<6>; + // Batched routed-expert matmul. It reads the expert-major map produced above, // loads selected expert weights, and writes results back to token-major slots // so the DS4 FFN can apply SwiGLU, weighting, and the down projection. -template +template kernel void kernel_mul_mm_id( constant ds4_metal_args_mul_mm_id & args, device const char * src0, @@ -7059,6 +7801,7 @@ kernel void kernel_mul_mm_id( device const char * htpe, device const char * hids, device char * dst, + device const char * work, threadgroup char * shmem [[threadgroup(0)]], uint3 tgpig[[threadgroup_position_in_grid]], ushort tiitg[[thread_index_in_threadgroup]], @@ -7075,9 +7818,16 @@ kernel void kernel_mul_mm_id( threadgroup S0 * sa = (threadgroup S0 *)(shmem); threadgroup S1 * sb = (threadgroup S1 *)(shmem + SA_BYTES); - const int im = tgpig.z; + device const uint32_t * work_count = (device const uint32_t *) work; + const uint32_t work_index = tgpig.x; + if (work_index >= work_count[0]) { + return; + } + device const uint2 * work_items = (device const uint2 *)(work + 8); + const uint2 item = work_items[work_index]; + const int im = (int)item.x; const int r0 = tgpig.y*NR0; - const int r1 = tgpig.x*NR1; + const int r1 = (int)item.y; device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); device const int32_t * ids_i32 = (device const int32_t *) (hids); @@ -7090,6 +7840,11 @@ kernel void kernel_mul_mm_id( const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + // SIMDgroups 0/1 own routed rows 0..15 and SIMDgroups 2/3 own rows + // 16..31. Keep all threads in staging and at every threadgroup barrier, + // but let the second row-half skip MMA on short final expert tiles. + const bool mma_active = + !CULL_TAIL_SIMDGROUPS || 16*(short)(sgitg/2) < nr1; if (!ds4_tp_owns_expert(im, args.ne02, args.tp_rank, args.tp_world)) { /* Unowned expert under the TP split: zero this tile's output rows so @@ -7212,27 +7967,29 @@ kernel void kernel_mul_mm_id( threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2)); threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2)); - FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { - simdgroup_barrier(mem_flags::mem_none); + if (mma_active) { + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); - FOR_UNROLL (short i = 0; i < 4; i++) { - simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++){ + simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); + } + + lsma += 8*64; + lsmb += 4*64; } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 2; i++) { - simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 8; i++){ - simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); - } - - lsma += 8*64; - lsmb += 4*64; } } @@ -7240,8 +7997,10 @@ kernel void kernel_mul_mm_id( threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; - for (short i = 0; i < 8; i++) { - simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + if (mma_active) { + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -7483,7 +8242,7 @@ kernel void kernel_mul_mm_id_addr( // each output keeps the exact MMA accumulation order of the separate GEMMs, // and the epilogue matches kernel_dsv4_moe_swiglu_weight_f16, so the fused // result is bit-identical to the unfused path. -template +template kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( constant ds4_metal_args_mul_mm_id & args, constant ds4_metal_dsv4_moe_swiglu_weight_args & act, @@ -7494,6 +8253,7 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( device const char * hids, device char * dst_mid, device const char * weights, + device const char * work, threadgroup char * shmem [[threadgroup(0)]], uint3 tgpig[[threadgroup_position_in_grid]], ushort tiitg[[thread_index_in_threadgroup]], @@ -7509,9 +8269,16 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( constexpr int NL0 = NK/16; constexpr int NL1 = NK/8; - const int im = tgpig.z; + device const uint32_t * work_count = (device const uint32_t *) work; + const uint32_t work_index = tgpig.x; + if (work_index >= work_count[0]) { + return; + } + device const uint2 * work_items = (device const uint2 *)(work + 8); + const uint2 item = work_items[work_index]; + const int im = (int)item.x; const int r0 = tgpig.y*NR0; - const int r1 = tgpig.x*NR1; + const int r1 = (int)item.y; device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); device const int32_t * ids_i32 = (device const int32_t *) (hids); @@ -7524,6 +8291,11 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + // SIMDgroups 0/1 own routed rows 0..15 and SIMDgroups 2/3 own rows + // 16..31. Keep every thread in staging and at every threadgroup barrier, + // but let the second row-half skip MMA on short final expert tiles. + const bool mma_active = + !CULL_TAIL_SIMDGROUPS || 16*(short)(sgitg/2) < nr1; const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; @@ -7600,27 +8372,29 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( threadgroup const half * lsma_up = (sa_up + 4*64*(sgitg%2)); threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); - FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { - simdgroup_barrier(mem_flags::mem_none); + if (mma_active) { + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); - FOR_UNROLL (short i = 0; i < 4; i++) { - simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false); - simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false); + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false); + simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false); + } + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]); + simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]); + } + + lsma_gate += 8*64; + lsma_up += 8*64; + lsmb += 4*64; } - FOR_UNROLL (short i = 0; i < 2; i++) { - simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 8; i++) { - simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]); - simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]); - } - - lsma_gate += 8*64; - lsma_up += 8*64; - lsmb += 4*64; } il = (il + 2 < nl) ? il + 2 : il % 2; @@ -7638,9 +8412,11 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( threadgroup float * temp_up_str = temp_up + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; - for (short i = 0; i < 8; i++) { - simdgroup_store(mc_gate[i], temp_gate_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); - simdgroup_store(mc_up[i], temp_up_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + if (mma_active) { + for (short i = 0; i < 8; i++) { + simdgroup_store(mc_gate[i], temp_gate_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + simdgroup_store(mc_up[i], temp_up_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -7674,18 +8450,251 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( } } +// MXFP4 resident-prefill specialization with a compact 32x32 +// output/routed-row tile. Two SIMDgroups share the same four 8-column A +// blocks and each owns one 16-row half of B. The per-output dequantization, +// MMA K sequence, and SwiGLU epilogue are identical to the 64x32 kernel +// above. On a short final work item, SIMDgroup 1 remains in staging and at +// every threadgroup barrier but skips its MMA/store when rows 16..31 are +// empty. +template +kernel void kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl( + constant ds4_metal_args_mul_mm_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device const char * htpe, + device const char * hids, + device char * dst_mid, + device const char * weights, + device const char * work, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + threadgroup half *sa_gate = (threadgroup half *)(shmem); + threadgroup half *sa_up = (threadgroup half *)(shmem + 2048); + threadgroup half *sb = (threadgroup half *)(shmem + 4096); + + constexpr int NR0 = 32; + constexpr int NR1 = 32; + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + device const uint32_t * work_count = (device const uint32_t *) work; + const uint32_t work_index = tgpig.x; + if (work_index >= work_count[0]) { + return; + } + device const uint2 * work_items = (device const uint2 *)(work + 8); + const uint2 item = work_items[work_index]; + const int im = (int)item.x; + const int r0 = tgpig.y*NR0; + const int r1 = (int)item.y; + + device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); + device const int32_t * ids_i32 = (device const int32_t *) (hids); + + const int32_t neh1 = tpe_u32[im]; + if (r1 >= neh1) { + return; + } + + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + const bool mma_active = 16*(short)sgitg < nr1; + + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; + const short row_b0 = (short)tiitg/NL1; + const short row_b1 = row_b0 + 16; + const short lr1_b0 = row_b0 < nr1 ? row_b0 : nr1 - 1; + const short lr1_b1 = row_b1 < nr1 ? row_b1 : nr1 - 1; + + const short il0 = (tiitg % NL0); + short il = il0; + + const int id_b0 = ids_i32[im*args.ne21 + r1 + lr1_b0]; + const int id_b1 = ids_i32[im*args.ne21 + r1 + lr1_b1]; + + const short i11_b0 = (id_b0 % args.ne20) % args.ne11; + const short i12_b0 = (id_b0 / args.ne20); + const short i11_b1 = (id_b1 % args.ne20) % args.ne11; + const short i12_b1 = (id_b1 / args.ne20); + const short i13 = 0; + + const uint64_t offset0 = im*args.nb02 + i13*args.nb03; + const short offset1 = il0/nl; + + device const block_q * xg = + (device const block_q *)(src0_gate + args.nb01*(r0 + lr0) + offset0) + offset1; + device const block_q * xu = + (device const block_q *)(src0_up + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const float * y_b0 = (device const float *)(src1 + + args.nb13*i13 + + args.nb12*i12_b0 + + args.nb11*i11_b0 + + args.nb10*iy); + device const float * y_b1 = (device const float *)(src1 + + args.nb13*i13 + + args.nb12*i12_b1 + + args.nb11*i11_b1 + + args.nb10*iy); + + simdgroup_half8x8 ma_g[4]; + simdgroup_half8x8 ma_u[4]; + simdgroup_half8x8 mb[2]; + + simdgroup_float8x8 mc_gate[8]; + simdgroup_float8x8 mc_up[8]; + + for (short i = 0; i < 8; i++) { + mc_gate[i] = make_filled_simdgroup_matrix(0.f); + mc_up[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + half4x4 temp_gate; + dequantize_func(xg, il, temp_gate); + half4x4 temp_up; + dequantize_func(xu, il, temp_up); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Sixty-four threads stage all 32 B rows by loading one row in each + // 16-row half. The shared-memory block order matches the 64x32 path. + { + const short sx_b = (tiitg%NL1); + const short sy_b0 = row_b0/8; + const short sy_b1 = row_b1/8; + const short ly_b = row_b0%8; + const short ib_b0 = 4*sx_b + sy_b0; + const short ib_b1 = 4*sx_b + sy_b1; + *(threadgroup half2x4 *)(sb + 64*ib_b0 + 8*ly_b) = + (half2x4)(*((device float2x4 *) y_b0)); + *(threadgroup half2x4 *)(sb + 64*ib_b1 + 8*ly_b) = + (half2x4)(*((device float2x4 *) y_b1)); + } + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + const short ib = 4*sx + sy; + *(sa_gate + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4]; + *(sa_up + 64*ib + 8*ly + lx) = temp_up[i/4][i%4]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const half * lsma_gate = sa_gate; + threadgroup const half * lsma_up = sa_up; + threadgroup const half * lsmb = sb + 2*64*sgitg; + + if (mma_active) { + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false); + simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false); + } + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]); + simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]); + } + + lsma_gate += 4*64; + lsma_up += 4*64; + lsmb += 4*64; + } + } + + il = (il + 2 < nl) ? il + 2 : il % 2; + xg = (il < 2) ? xg + (2 + nl - 1)/nl : xg; + xu = (il < 2) ? xu + (2 + nl - 1)/nl : xu; + y_b0 += NK; + y_b1 += NK; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_gate = (threadgroup float *) shmem; + threadgroup float * temp_up = temp_gate + NR0*NR1; + threadgroup float * temp_gate_str = temp_gate + 16*sgitg*NR0; + threadgroup float * temp_up_str = temp_up + 16*sgitg*NR0; + + if (mma_active) { + for (short i = 0; i < 8; i++) { + simdgroup_store(mc_gate[i], temp_gate_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + simdgroup_store(mc_up[i], temp_up_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float c = act.clamp_value; + for (short j = sgitg; j < nr1; j += 2) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + + device half *D = (device half *)(dst_mid + + ((uint64_t)idt*args.ne1 + (uint64_t)ide)*act.mid_row_stride) + r0; + device const float *w = (device const float *)(weights + (uint64_t)idj*act.weight_stride); + const float route_weight = w[0]; + + threadgroup float *Cg = temp_gate + j*NR0; + threadgroup float *Cu = temp_up + j*NR0; + + int i = tiisg; + for (; i < nr0; i += 32) { + float g = Cg[i]; + float u = Cu[i]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + D[i] = (half)(silu * u * route_weight); + } + } +} + typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_iq2; typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_q4; -typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_q5; +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_mxfp4; +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_mxfp4_tail_cull; +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl) mul_mm_id_pair_swiglu_f16_mxfp4_compact_tail; // Host-visible fused routed pair matmuls for the DS4 expert quant formats. template [[host_name("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_iq2 kernel_mul_mm_id_pair_swiglu_f16_impl; template [[host_name("kernel_mul_mm_id_q4_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q4 kernel_mul_mm_id_pair_swiglu_f16_impl; -template [[host_name("kernel_mul_mm_id_q5_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q5 kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4 kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_half_scale")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4 kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull_half_scale")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_tail_cull kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull_half_scale")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_compact_tail kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl; +template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_tail_cull kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_compact_tail kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs; -typedef decltype(kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_ff32; +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_tail_cull; +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4>) mul_mm_id_mxfp4_f16_rhs_half_lut; +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_half_lut_tail_cull; typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_addr; typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_addr_f16_rhs; @@ -7696,34 +8705,41 @@ template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id ker template [[host_name("kernel_mul_mm_id_q5_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_q6_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_q8_0_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q5_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q6_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, half, half4x4, half, half2x4>; -template [[host_name("kernel_mul_mm_id_q4_K_ff32")]] kernel mul_mm_id_ff32 kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>; -template [[host_name("kernel_mul_mm_id_q5_K_ff32")]] kernel mul_mm_id_ff32 kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>; -template [[host_name("kernel_mul_mm_id_q6_K_ff32")]] kernel mul_mm_id_ff32 kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_mxfp4_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_mxfp4_f16_tail_cull")]] kernel mul_mm_id_mxfp4_f16_rhs_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4, true>; +template [[host_name("kernel_mul_mm_id_mxfp4_f16_half_lut")]] kernel mul_mm_id_mxfp4_f16_rhs_half_lut kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_mxfp4_f16_half_lut_tail_cull")]] kernel mul_mm_id_mxfp4_f16_rhs_half_lut_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4, true>; template [[host_name("kernel_mul_mm_id_addr_q2_K_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_addr_q4_K_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_addr_mxfp4_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_addr_q2_K_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_addr_q4_K_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_addr_mxfp4_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4>; #ifdef DS4_METAL_HAS_TENSOR -// Attention-output low-rank projection retained for Metal4 prefill. It uses -// the same direct-RHS idea as dense matmul: dequantize the Q8_0 low projection -// weights to a half tile, then let TensorOps read the dense head activations -// directly. Only the 64-token direct-RHS instantiation is exported because the -// staged-RHS and 32-token variants were benchmark-only experiments. +// Attention-output low-rank projection retained for Metal4 prefill. It uses +// the same direct-RHS idea as dense matmul: dequantize the Q8_0 or Q4_K low +// projection weights to a half tile, then let TensorOps read the dense head +// activations directly. Only the 64-token direct-RHS instantiations are +// exported because the staged-RHS and 32-token variants were benchmark-only +// experiments. // // Full tiles (the host dispatch guarantee for aligned batches) skip all bounds // work. The weight tile is double-buffered: the next k-step's dequant // overlaps the current cooperative matmul, so the k-loop needs one // threadgroup barrier per step instead of two. -template -kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( +template +kernel void kernel_attn_out_low_mpp_direct_rhs( constant ds4_metal_args_mul_mm_id & args, device const char * srcA, device const char * srcB, @@ -7778,24 +8794,24 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( const short k_base = k_chunk*16; if (full_tile) { - device const block_q8_0 *row_ptr = - (device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02); + device const block_q *row_ptr = + (device const block_q *)(srcA + args.nb01*(r0 + row) + group*args.nb02); half4x4 temp_a; - dequantize_q8_0_pairs(row_ptr + k_pos/32, (k_pos/16)%2, temp_a); + dequantize_func(row_ptr + k_pos/(16*qnl), (k_pos/16)%qnl, temp_a); threadgroup half4 *dst4 = (threadgroup half4 *)(buf + row*NK + k_base); dst4[0] = temp_a[0]; dst4[1] = temp_a[1]; dst4[2] = temp_a[2]; dst4[3] = temp_a[3]; } else if (r0 + row < M) { - const int block_idx = k_pos/32; - const short il = (k_pos/16)%2; - device const block_q8_0 *row_ptr = - (device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02); + const int block_idx = k_pos/(16*qnl); + const short il = (k_pos/16)%qnl; + device const block_q *row_ptr = + (device const block_q *)(srcA + args.nb01*(r0 + row) + group*args.nb02); half4x4 temp_a; - dequantize_q8_0_pairs(row_ptr + block_idx, il, temp_a); + dequantize_func(row_ptr + block_idx, il, temp_a); FOR_UNROLL (short i = 0; i < 16; i++) { buf[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (half)0; } @@ -7836,9 +8852,228 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( } } -typedef decltype(kernel_attn_out_low_q8_0_mpp_direct_rhs<64>) attn_out_low_q8_0_mpp_direct_rhs_n64_t; +// Routed-expert grouped matmul on the Metal4 TensorOps/MPP pipeline. The +// barrier after mm.run prevents the next K iteration from replacing staged +// tiles while the cooperative matmul still reads them. +template +kernel void kernel_mul_mm_id_mpp( + constant ds4_metal_args_mul_mm_id & args, + device const char * src0, + device const char * src1, + device const char * htpe, + device const char * hids, + device char * dst, + device const char * work, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + threadgroup S0 * sa = (threadgroup S0 *)(shmem); + threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096); + threadgroup float *sc = (threadgroup float *)shmem; -template [[host_name("kernel_attn_out_low_q8_0_mpp_direct_rhs_n64")]] kernel attn_out_low_q8_0_mpp_direct_rhs_n64_t kernel_attn_out_low_q8_0_mpp_direct_rhs<64>; + constexpr int NR0 = 64; + constexpr int NR1 = 32; + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + device const uint32_t *work_count = (device const uint32_t *)work; + const uint32_t work_index = tgpig.x; + if (work_index >= work_count[0]) { + return; + } + device const uint2 *work_items = (device const uint2 *)(work + 8); + const uint2 item = work_items[work_index]; + const int im = (int)item.x; + const int r0 = tgpig.y*NR0; + const int r1 = (int)item.y; + + device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); + device const int32_t * ids_i32 = (device const int32_t *) (hids); + + const int32_t neh1 = tpe_u32[im]; + + if (r1 >= neh1) { + return; + } + + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + + if (!ds4_tp_owns_expert(im, args.ne02, args.tp_rank, args.tp_world)) { + for (short j = sgitg; j < nr1; j += 4) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + device float *D = (device float *)dst + r0 + ide*args.ne0 + + idt*args.ne1*args.ne0; + for (int i = tiisg; i < nr0; i += 32) D[i] = 0.0f; + } + return; + } + + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; + + const short il0 = (tiitg % NL0); + short il = il0; + + const int id = ids_i32[im*args.ne21 + r1 + lr1]; + + const short i11 = (id % args.ne20) % args.ne11; + const short i12 = (id / args.ne20); + const short i13 = 0; + + const uint64_t offset0 = + (uint64_t)(im - args.tp_expert_base)*args.nb02 + i13*args.nb03; + const short offset1 = il0/nl; + + device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const T1 * y = (device const T1 *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*i11 + + args.nb10*iy); + + auto tA = tensor(sa, dextents(NK, NR0)); + auto tB = tensor(sb, dextents(NR1, NK)); + + matmul2d< + matmul2d_descriptor(NR1, NR0, NK, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups<4>> mm; + + auto cT = mm.template get_destination_cooperative_tensor(); + + #pragma unroll + for (uint16_t i = 0; i < cT.get_capacity(); ++i) { + if (cT.is_valid_element(i)) { + cT[i] = 0.0f; + } + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + if (is_same::value && FC_mul_mm_bc_inp) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + const short lx = i%8; + const short ly = (tiitg/NL0)%8; + + *(sa + NK*(8*sy + ly) + 8*sx + lx) = + loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0; + } + } else { + S0_4x4 temp_a; + dequantize_func(x, il, temp_a); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + const short lx = i%8; + const short ly = (tiitg/NL0)%8; + + *(sa + NK*(8*sy + ly) + 8*sx + lx) = temp_a[i/4][i%4]; + } + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + const short lx = i; + const short ly = (tiitg/NL1)%8; + + *(sb + NK*(8*sy + ly) + 8*sx + lx) = + loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + const short ly = (tiitg/NL1)%8; + + *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = + (S1_2x4)(*((device T1_2x4 *) y)); + } + + il = (il + 2 < nl) ? il + 2 : il % 2; + x = (il < 2) ? x + (2 + nl - 1)/nl : x; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + auto sA = tA.slice(0, 0); + auto sB = tB.slice(0, 0); + mm.run(sB, sA, cT); + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + auto tC = tensor(sc, dextents(NR0, NR1)); + cT.store(tC); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short j = tiitg/32; j < nr1; j += 4) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + + device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = (threadgroup float *) shmem + j*NR0; + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = tiisg; + for (; i < nr0/4; i += 32) { + *(D4 + i) = *(C4 + i); + } + + i = (4*(nr0/4)) + tiisg; + for (; i < nr0; i += 32) { + *(D + i) = *(C + i); + } + } +} + + + +typedef decltype(kernel_mul_mm_id_mpp) mul_mm_id_mpp_t; +typedef decltype(kernel_mul_mm_id_mpp) mul_mm_id_mpp_f16_rhs_t; + +template [[host_name("kernel_mul_mm_id_iq2_xxs_f32_mpp")]] kernel mul_mm_id_mpp_t kernel_mul_mm_id_mpp; +template [[host_name("kernel_mul_mm_id_q2_K_f16_mpp")]] kernel mul_mm_id_mpp_f16_rhs_t kernel_mul_mm_id_mpp; +template [[host_name("kernel_mul_mm_id_iq2_xxs_f16_mpp")]] kernel mul_mm_id_mpp_f16_rhs_t kernel_mul_mm_id_mpp; + +typedef decltype(kernel_attn_out_low_mpp_direct_rhs< + block_q8_0, 2, dequantize_q8_0_pairs, 64>) + attn_out_low_q8_0_mpp_direct_rhs_n64_t; +typedef decltype(kernel_attn_out_low_mpp_direct_rhs< + block_q4_K, QK_NL, dequantize_q4_K, 64>) + attn_out_low_q4_K_mpp_direct_rhs_n64_t; + +template [[host_name("kernel_attn_out_low_q8_0_mpp_direct_rhs_n64")]] +kernel attn_out_low_q8_0_mpp_direct_rhs_n64_t +kernel_attn_out_low_mpp_direct_rhs< + block_q8_0, 2, dequantize_q8_0_pairs, 64>; +template [[host_name("kernel_attn_out_low_q4_K_mpp_direct_rhs_n64")]] +kernel attn_out_low_q4_K_mpp_direct_rhs_n64_t +kernel_attn_out_low_mpp_direct_rhs< + block_q4_K, QK_NL, dequantize_q4_K, 64>; #endif diff --git a/metal/norm.metal b/metal/norm.metal index 4bc72f2..25a8805 100644 --- a/metal/norm.metal +++ b/metal/norm.metal @@ -241,3 +241,145 @@ kernel void kernel_dsv4_qkv_rms_norm_f32_4( y[i] = (x[i] * scale) * w[i]; } } + +// Decode-only triple fusion: the q/kv RMS norm, the KV RoPE tail, and the +// FP8/raw finalizer were three back-to-back dispatches on the same rows. +// The q threadgroup is byte-identical to kernel_dsv4_qkv_rms_norm_f32_4. +// The kv threadgroup continues with the shared affine-row RoPE helper (lane +// mapping preserved: r == lane on the first 64 lanes) and a verbatim copy of +// kernel_dsv4_kv_fp8_store_f32 with its work predicated to the first 64 +// lanes (barriers stay uniform across the whole threadgroup). Arithmetic, +// order and rounding are unchanged; gated and verified against +// full-vocabulary logits before promotion. +kernel void kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32( + constant ds4_metal_args_qkv_rms_norm & args, + constant ds4_metal_args_dsv4_rope_affine_pair & rope, + constant ds4_metal_args_dsv4_kv_fp8_store & store, + device const float4 * q_src, + device const float4 * q_weight, + device float4 * q_dst, + device const float4 * kv_src, + device const float4 * kv_weight, + device float4 * kv_dst, + device float * raw_cache, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + if (sgitg == 0) { + shmem_f32[tiisg] = 0.0f; + } + + const uint row = tgpig.x; + const bool kv_task = tgpig.y != 0; + const int n = kv_task ? args.kv_n : args.q_n; + const int n4 = kv_task ? args.kv_n4 : args.q_n4; + const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4); + + device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4; + device const float4 * w = kv_task ? kv_weight : q_weight; + device float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4; + + float sumf = 0.0f; + for (int i = tpitg.x; i < n4; i += ntg.x) { + const float4 v = x[i]; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + shmem_f32[sgitg] = sumf; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = shmem_f32[tiisg]; + sumf = simd_sum(sumf); + +#ifdef DS4_METAL_NORM_RSQRT_DISABLE + const float scale = 1.0f / sqrt(sumf / float(n) + args.eps); +#else + const float scale = rsqrt(sumf / float(n) + args.eps); +#endif + + for (int i = tpitg.x; i < n4; i += ntg.x) { + y[i] = (x[i] * scale) * w[i]; + } + + if (!kv_task) { + return; + } + + // KV RoPE tail in place, then the FP8/raw finalizer (verbatim bodies). + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + + device char *kv_row = (device char *)(kv_dst + row * row_stride4); + const int rope_n_nope = rope.head_dim - rope.n_dims; + if (rope_n_nope < 0) { + return; + } + ds4_rope_tail_pair_affine_row(rope, + (device const char *)kv_row, + kv_row, + rope_n_nope, + rope.pos0, + tpitg.x, + ntg.x); + + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + + const int head_dim = store.head_dim; + const int n_rot = store.n_rot; + const int n_nope = head_dim - n_rot; + if (head_dim <= 0 || n_rot < 0 || n_nope < 0) { + return; + } + const uint tid = tpitg.x; + + device float *kv = (device float *)kv_row; + device float *raw = raw_cache + (int64_t)store.raw_row * head_dim; + threadgroup float *scratch = shmem_f32 + 32; + + for (int off = 0; off < n_nope; off += 64) { + float v = 0.0f; + if (tid < 64u && off + (int)tid < n_nope) { + v = kv[off + tid]; + scratch[tid] = abs(v); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tid < stride) { + scratch[tid] = max(scratch[tid], scratch[tid + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float amax = max(scratch[0], 1.0e-4f); + const float fp8_scale = exp2(ceil(log2(amax / 448.0f))); + if (tid < 64u && off + (int)tid < n_nope) { + const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale; + kv[off + tid] = q; +#ifdef DS4_METAL_KV_RAW_F32 + raw[off + tid] = q; +#else + raw[off + tid] = (float)((half)q); +#endif + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid < 64u) { + for (int i = n_nope + tid; i < head_dim; i += 64) { +#ifdef DS4_METAL_KV_RAW_F32 + raw[i] = kv[i]; +#else + raw[i] = (float)((half)kv[i]); +#endif + } + } +} diff --git a/native/metal/README.md b/native/metal/README.md index 0f2a5f5..9e777cc 100644 --- a/native/metal/README.md +++ b/native/metal/README.md @@ -1,7 +1,7 @@ # Vendored DS4 Metal boundary -These files are a one-time snapshot of the DS4 Metal boundary from commit -`efdadd41e20134af4f3381e1ed90e96fe4faef6f`: +These files track the DS4 Metal boundary audited at commit +`8db89fe083ae4d17c9a2428ccd29803d3ae8f577` (2026-08-28): - `ds4_metal.m` - `ds4.h` diff --git a/native/metal/ds4.h b/native/metal/ds4.h index 42bc116..3c9e066 100644 --- a/native/metal/ds4.h +++ b/native/metal/ds4.h @@ -150,6 +150,7 @@ typedef struct { bool glm_mtp_timing; bool dspark; bool dspark_strict; + bool dspark_exact_sampling; bool dspark_confidence_threshold_set; bool cuda_tensor_parallel; bool ssd_streaming; @@ -158,6 +159,8 @@ typedef struct { bool inspect_only; /* Multi-GPU placement uses this to price per-layer KV storage. */ int placement_ctx_hint; + /* Number of independently allocated session graphs/caches to reserve. */ + int placement_session_count_hint; /* Server batch mode serializes execution and can share prefill scratch. */ bool share_session_prefill_workspace; bool first_token_test; @@ -371,6 +374,30 @@ int ds4_test_sample_logits(const float *logits, uint32_t n_vocab, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, float *prob_scratch); +int ds4_test_sampling_probabilities(const float *logits, uint32_t n_vocab, + float temperature, int top_k, + float top_p, float min_p, float *probs); +int ds4_test_speculative_sample(const float *target_logits, + const float *draft_logits, + uint32_t n_vocab, + float temperature, + int top_k, + float top_p, + float min_p, + uint64_t *rng, + float *target_probs, + float *draft_probs); +int ds4_test_speculative_delta_sample(const float *target_logits, + uint32_t n_vocab, + int draft_token, + float temperature, + int top_k, + float top_p, + float min_p, + uint64_t *rng, + float *target_probs); +int ds4_test_argmax_excluding_logits(const float *logits, uint32_t n_vocab, + int excluded_id); uint64_t ds4_test_mixed_native_count(void); #endif int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k); @@ -403,6 +430,15 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, int max_tokens, int eos_token, int *accepted, int accepted_cap, char *err, size_t errlen); +/* Evaluate one already-sampled target token and speculatively extend it. + * Positive-temperature DSpark normally commits greedily verified draft + * tokens; dspark_exact_sampling selects exact stochastic p/q acceptance. */ +int ds4_session_eval_speculative(ds4_session *s, int first_token, + int max_tokens, int eos_token, + float temperature, int top_k, + float top_p, float min_p, uint64_t *rng, + int *accepted, int accepted_cap, + char *err, size_t errlen); /* TP worker side of a mirrored speculative-verify block: run its half of the * batch verify for KV side effects, then obey the leader's commit frame * (keep, or roll back and replay). Only called from ds4_tp_worker_run. */ diff --git a/native/metal/ds4_gpu.h b/native/metal/ds4_gpu.h index 6111016..2fe34b3 100644 --- a/native/metal/ds4_gpu.h +++ b/native/metal/ds4_gpu.h @@ -78,6 +78,24 @@ int ds4_gpu_begin_commands(void); int ds4_gpu_flush_encoder(void); int ds4_gpu_flush_commands(void); int ds4_gpu_commands_active(void); +#ifdef __APPLE__ +int ds4_gpu_parallel_ffn_finish(void); +void ds4_gpu_parallel_ffn_abort(void); +int ds4_gpu_parallel_ffn_start( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *shared_out, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t model_dim, + uint32_t shared_dim, + const ds4_gpu_tensor *x, + float clamp); +#endif int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value); int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label); int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label); @@ -95,6 +113,10 @@ int ds4_gpu_synchronize(void); int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size); int ds4_gpu_set_model_fd(int fd); int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map); +int ds4_gpu_build_derived_artifacts(const void *model_map, uint64_t model_size, + const char *model_path); +int ds4_gpu_model_range_replaced(const void *model_map, uint64_t offset, + uint64_t bytes); int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes); int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes); int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label); @@ -141,7 +163,30 @@ void ds4_gpu_set_glm_model(bool enabled); void ds4_gpu_set_ssd_streaming(bool enabled); void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled); #ifdef __APPLE__ +int ds4_gpu_device_is_pre_m5_apple_silicon(void); +int ds4_gpu_device_is_m5_apple_silicon(void); +int ds4_gpu_set_decode_pipeline_fast_lookup(int enabled); +/* Strict test oracle for the fixed decode mul_mv pipeline lookup cache. */ +int ds4_gpu_test_decode_pipeline_fast_lookup(void); +/* Strict test oracle for the extended decode mul_mv_ext (nsg + nxpsg) cache. */ +int ds4_gpu_test_decode_pipeline_fast_lookup_ext(void); +/* Strict test oracle for the generated resident-prefill MXFP4 half LUT. */ +int ds4_gpu_test_mxfp4_down_half_lut(uint16_t *legacy_bits, + uint16_t *lut_bits); +enum { + DS4_GPU_TEST_MXFP4_PAIR_TAIL_CULL = 1u << 0, + DS4_GPU_TEST_MXFP4_PAIR_COMPACT_TILE = 1u << 1, + DS4_GPU_TEST_MXFP4_MAP_SCATTER = 1u << 2, + DS4_GPU_TEST_MXFP4_DOWN_TAIL_CULL = 1u << 3, + DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT = 1u << 4, + DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4 = 1u << 5, + DS4_GPU_TEST_HC_RMS_SCALE_PROJ = 1u << 6, +}; +void ds4_gpu_test_set_flags(uint32_t flags); void ds4_gpu_release_zero_prefix_prefill_mask_cache(void); +#else +static inline int ds4_gpu_device_is_pre_m5_apple_silicon(void) { return 0; } +static inline int ds4_gpu_device_is_m5_apple_silicon(void) { return 0; } #endif void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts); void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes); @@ -216,6 +261,7 @@ int ds4_gpu_stream_expert_cache_seed_from_layer_selected( uint32_t n_tokens, uint32_t n_seed_tokens, uint32_t n_selected); +int ds4_gpu_stream_expert_cache_finish_pending_batch(void); int ds4_gpu_stream_expert_cache_release_layer_cache(void); #endif int ds4_gpu_stream_expert_cache_seed_experts( @@ -223,6 +269,14 @@ int ds4_gpu_stream_expert_cache_seed_experts( const int32_t *expert_ids, const uint32_t *expert_priorities, uint32_t n_experts); +#ifdef __APPLE__ +/* Seed from mapped weights with blits appended to the active command buffer. */ +int ds4_gpu_stream_expert_cache_seed_experts_gpu_copy( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts); +#endif void ds4_gpu_print_memory_report(const char *label); /* Tensor-parallel per-layer gates (Metal only). The encoder calls @@ -659,6 +713,36 @@ int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( uint64_t out_dim, const ds4_gpu_tensor *x, float clamp); + +int ds4_gpu_router_shared_gate_up_q8_0_tensor( + ds4_gpu_tensor *router_logits, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t router_weight_offset, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t router_out_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp, + bool router_only); +#ifdef __APPLE__ +int ds4_gpu_router_project_select_fused_tensor( + ds4_gpu_tensor *router_logits, + ds4_gpu_tensor *probs, + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + const void *model_map, + uint64_t model_size, + uint64_t router_weight_offset, + uint64_t bias_offset, + bool has_bias, + const ds4_gpu_tensor *x); +#endif int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( ds4_gpu_tensor *mid, const void *model_map, @@ -736,6 +820,20 @@ int ds4_gpu_matmul_f16_tensor( const ds4_gpu_tensor *x, uint64_t n_tok); +/* CUDA batch path: fold an input RMS normalization into the FP16 activation + * conversion used by the following projection. Returns 0 without touching + * out when the optimized path is unavailable. */ +int ds4_gpu_matmul_f16_rms_fold_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float norm_eps); + /* Exact multi-row form of the DeepSeek 4096x256 F16 router projection. */ int ds4_gpu_matmul_f16_router_rows_exact_tensor( ds4_gpu_tensor *out, @@ -777,6 +875,95 @@ int ds4_gpu_matmul_f16_pair_compressor_store_tensor( uint32_t ratio, uint32_t pos); +int ds4_gpu_matmul_f16_quad_compressor_store_tensor( + ds4_gpu_tensor *out0_kv, + ds4_gpu_tensor *out0_score, + ds4_gpu_tensor *out1_kv, + ds4_gpu_tensor *out1_score, + ds4_gpu_tensor *state0_kv, + ds4_gpu_tensor *state0_score, + ds4_gpu_tensor *state1_kv, + ds4_gpu_tensor *state1_score, + const void *model_map, + uint64_t model_size, + uint64_t weight0_kv_offset, + uint64_t weight0_score_offset, + uint64_t weight1_kv_offset, + uint64_t weight1_score_offset, + uint64_t ape0_offset, + uint32_t ape0_type, + uint64_t ape1_offset, + uint32_t ape1_type, + uint64_t in_dim, + uint32_t width0, + uint32_t width1, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos); + +/* Decode-only M5 fusion: emit-path compressor row finalize (norm + rope + + * fp8/commit + indexer qat) in one dispatch. Bit-exact vs the separate + * dispatches. Returns 1 when fused, 0 to fall back. */ +int ds4_gpu_dsv4_comp_row_finalize_tensor( + ds4_gpu_tensor *attn_stage, + ds4_gpu_tensor *attn_cache, + uint32_t attn_comp_row, + uint64_t attn_norm_offset, + ds4_gpu_tensor *index_cache, + uint32_t index_comp_row, + uint64_t index_norm_offset, + ds4_gpu_tensor *attn_state_kv, + ds4_gpu_tensor *attn_state_score, + ds4_gpu_tensor *index_state_kv, + ds4_gpu_tensor *index_state_score, + const void *model_map, + uint64_t model_size, + uint32_t pos, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps); + +/* Decode-only M5 fusion: q_a/kv Q8 pair projection + F16 quad compressor + * projection/store in one dispatch. Bit-exact vs the separate dispatches. + * Returns 1 when fused, 0 to fall back, -1 on error. */ +int ds4_gpu_qkv_pair_quad_compressor_store_tensor( + ds4_gpu_tensor *qr, + ds4_gpu_tensor *kv_raw, + ds4_gpu_tensor *out0_kv, + ds4_gpu_tensor *out0_score, + ds4_gpu_tensor *out1_kv, + ds4_gpu_tensor *out1_score, + ds4_gpu_tensor *state0_kv, + ds4_gpu_tensor *state0_score, + ds4_gpu_tensor *state1_kv, + ds4_gpu_tensor *state1_score, + const void *model_map, + uint64_t model_size, + uint64_t q_a_offset, + uint64_t kv_offset, + uint64_t weight0_kv_offset, + uint64_t weight0_score_offset, + uint64_t weight1_kv_offset, + uint64_t weight1_score_offset, + uint64_t ape0_offset, + uint32_t ape0_type, + uint64_t ape1_offset, + uint32_t ape1_type, + uint32_t in_dim, + uint32_t q_rank, + uint32_t kv_dim, + uint32_t width0, + uint32_t width1, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos); + int ds4_gpu_matmul_f32_tensor( ds4_gpu_tensor *out, const void *model_map, @@ -857,6 +1044,31 @@ int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( uint32_t rows, float eps); +int ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_out, + const ds4_gpu_tensor *kv, + uint64_t kv_weight_offset, + uint32_t kv_n, + ds4_gpu_tensor *raw_cache, + uint64_t raw_cap, + uint32_t raw_row, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps); + int ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor( ds4_gpu_tensor *q_out, const ds4_gpu_tensor *q, @@ -942,6 +1154,8 @@ int ds4_gpu_dsv4_indexer_qat_tensor( uint32_t n_rows, uint32_t head_dim); + + int ds4_gpu_rope_tail_tensor( ds4_gpu_tensor *x, uint32_t n_tok, @@ -1609,7 +1823,9 @@ int ds4_gpu_compressor_update_tensor( float beta_fast, float beta_slow, float rms_eps, - bool state_already_stored); + bool state_already_stored, + bool decode_one_token, + bool defer_finalize); int ds4_gpu_compressor_store_batch_tensor( const ds4_gpu_tensor *kv, @@ -2409,18 +2625,6 @@ int ds4_gpu_hc_weighted_sum_tensor( uint32_t n_embd, uint32_t n_hc); -int ds4_gpu_hc_weighted_sum_norm_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *norm_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *weights, - const void *model_map, - uint64_t model_size, - uint64_t norm_weight_offset, - uint32_t n_embd, - uint32_t n_hc, - float norm_eps); - int ds4_gpu_hc_weighted_sum_split_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, @@ -2461,6 +2665,17 @@ int ds4_gpu_hc_split_weighted_sum_norm_tensor( float eps, float norm_eps); +int ds4_gpu_hc_rms_norm_mix_f16_available(void); +int ds4_gpu_hc_rms_norm_mix_f16_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + uint32_t out_dim, + float eps); + /* Batched HC RMSNorm followed by its narrow F16 mixer projection. On the * tuned Metal path, scale_scratch stores one float per row instead of the * full normalized HC tensor; other shapes retain the established fallback. */ @@ -2476,6 +2691,29 @@ int ds4_gpu_hc_rms_scale_project_f16_tensor( uint32_t n_rows, float eps); +#ifdef __APPLE__ +int ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( + ds4_gpu_tensor *mix, + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t mix_weight_offset, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n, + uint32_t mix_dim, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float hc_eps, + float norm_eps); + +#endif int ds4_gpu_output_hc_weights_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *pre, @@ -2612,6 +2850,34 @@ int ds4_gpu_matmul_q8_0_hc_expand_tensor( uint32_t n_embd, uint32_t n_hc); +/* Decode-island CUDA graph capture (CUDA backend; Metal/ROCm/CPU stub it + * out and stay eager). Design ported from the Entrpi/ds4 batched-serving + * fork's per-layer decode graph capture. The key identifies a captured + * island: layer, island index, and the activation buffers whose addresses + * the captured kernels bake in. ds4_cuda.cu mirrors this struct + * byte-for-byte (it does not include this header); keep both in sync. */ +typedef struct ds4_decode_graph_key { + uint32_t il; + uint32_t island; /* 0: layer top to pre-rope; 1: attn-out to layer end */ + uint32_t variant; + uint32_t _pad; + void *cur_hc; + void *after_attn_hc; + void *after_ffn_hc; + void *attn_norm; +} ds4_decode_graph_key; + +int ds4_gpu_decode_graphs_supported(void); +/* 1: replayed (island already executed; skip encoding it) + * 0: capturing (encode the island, then call _end) + * -1: run eagerly */ +int ds4_gpu_decode_graph_begin(const ds4_decode_graph_key *key); +/* 0: capture committed and launched; -1: capture failed (entry retired; + * the caller must re-encode the island eagerly -- no work was executed). */ +int ds4_gpu_decode_graph_end(const ds4_decode_graph_key *key); +void ds4_gpu_decode_graph_abort(const ds4_decode_graph_key *key); +void ds4_gpu_decode_graphs_invalidate(void); + #ifdef __cplusplus } #endif diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index bcb4480..4f94437 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -42,6 +42,7 @@ enum { DS4_METAL_TENSOR_Q6_K = 14, DS4_METAL_TENSOR_Q8_K = 15, DS4_METAL_TENSOR_IQ2_XXS = 16, + DS4_METAL_TENSOR_MXFP4 = 39, }; @class DS4MetalQ4ExpertTable; @@ -51,7 +52,9 @@ static id g_queue; static id g_library; static id g_batch_cb; static id g_batch_enc; +static BOOL g_batch_encoder_concurrent; static BOOL g_batch_has_work; +static void ds4_gpu_parallel_ffn_reset_state(BOOL close_encoder); static NSMutableArray> *g_pending_cbs; static id g_selected_readback_event; static uint64_t g_selected_readback_event_value; @@ -85,12 +88,15 @@ static id g_rms_norm_plain_pipeline; static id g_add_rms_norm_pipeline; static id g_rms_norm_scale_pipeline; static id g_dsv4_qkv_rms_norm_pipeline; +static id g_dsv4_head_rms_norm_rope_tail_pipeline; +static bool g_use_dsv4_head_rms_norm_rope_tail_pipeline; static id g_hc_split_sinkhorn_pipeline; static id g_hc_split_weighted_sum_pipeline; static id g_hc_split_weighted_sum_norm_pipeline; +static id g_dsv4_hc_producer_pre_norm_pipeline; static id g_hc_weighted_sum_pipeline; -static id g_hc_weighted_sum_norm_pipeline; static id g_output_hc_weights4_pipeline; +static uint32_t g_test_flags; static id g_hc_expand_pipeline; static id g_unary_sigmoid_pipeline; static id g_unary_silu_pipeline; @@ -105,6 +111,7 @@ static id g_bin_div_row_pipeline; static id g_moe_mul_mv_id_iq2_xxs_pipeline; static id g_moe_mul_mv_id_iq2_xxs_pair_pipeline; static id g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline; static id g_moe_mul_mv_id_q2_k_pipeline; static id g_moe_mul_mv_id_q2_k_sum6_pipeline; static id g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; @@ -124,6 +131,20 @@ static id g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipel static id g_moe_mul_mv_slots6_q2_k_sum6_pipeline; static id g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; static id g_moe_mul_mv_slots6_q4_k_sum6_pipeline; +static id g_moe_mul_mv_id_mxfp4_pipeline; +static id g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline; +static id g_moe_mul_mv_id_mxfp4_sum6_pipeline; +static id g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple; +static id g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple; +static id g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1; +static id g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline; +static id g_moe_mul_mv_slots6_mxfp4_sum6_pipeline; static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline; static id g_moe_mul_mv_addr_iq2_xxs_pipeline; static id g_moe_mul_mv_addr_q2_k_sum6_pipeline; @@ -145,8 +166,12 @@ static id g_rope_tail_inplace_pair_affine_pipeline; static id g_dsv4_fp8_kv_quantize_pipeline; static id g_dsv4_indexer_qat_pipeline; static id g_dsv4_kv_fp8_store_pipeline; +static id g_dsv4_kv_rope_fp8_store_pipeline; static id g_dsv4_ratio4_shift_pipeline; static id g_dsv4_compressor_pack_ratio4_pipeline; +static id g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline; +static id g_dsv4_compressor_exact_softmax_product_pipeline; +static id g_dsv4_compressor_exact_pool_ratio4_pipeline; static id g_dsv4_softmax_pool_ratio4_direct_pipeline; static id g_dsv4_softmax_pool_pipeline; static id g_soft_max_f32_pipeline; @@ -162,10 +187,18 @@ static id g_dsv4_compressor_store_one_pipeline; static id g_dsv4_sort_i32_rows_asc_pipeline; static id g_dsv4_indexed_attention_heads8_pipeline; static id g_dsv4_indexed_attention_heads8_rb16_pipeline; +static id g_dsv4_indexed_attention_heads16_dual_pipeline; +static id g_dsv4_indexed_attention_heads8_split_pipeline; +static id g_dsv4_indexed_attention_heads8_split_reduce_pipeline; static id g_dsv4_softplus_sqrt_pipeline; static id g_dsv4_router_finalize_one_pipeline; static id g_dsv4_router_finalize_one_simd_pipeline; static id g_dsv4_router_finalize_weights_one_simd_pipeline; +static id g_dsv4_router_transform_finalize_weights_one_simd_pipeline; +static NSCache> *g_dsv4_completion_cache; +static __weak id g_dsv4_hc_producer_last_mix_buffer; +static NSUInteger g_dsv4_hc_producer_last_mix_offset; +static id g_dsv4_hc_producer_last_completion; static id g_dsv4_router_weights_one_pipeline; static id g_glm_router_select_one_pipeline; static id g_glm_kv_lora_rms_norm_pipeline; @@ -229,6 +262,39 @@ static id g_glm_q6_k_down_f32_pipeline; static id g_dsv4_router_weights_batch_pipeline; static id g_dsv4_hc_expand4_pipeline; static NSMutableDictionary> *g_pipeline_cache; + +enum { + DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS = 64, + DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES = 96, +}; +/* + * Sentinel nxpsg stored by the single-constant ds4_gpu_get_mul_mv_pipeline + * path, which has no second SIMD-group constant. The extended + * ds4_gpu_get_mul_mv_ext_pipeline path derives nxpsg as a positive divisor of + * the 32-row output tile, so INT16_MIN can never collide with a real value. + */ +#define DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE INT16_MIN +_Static_assert( + (DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS & + (DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS - 1u)) == 0, + "decode pipeline fast cache slot count must be a power of two"); + +typedef struct { + id __strong pipeline; + uint64_t hash; + int16_t nsg; + int16_t nxpsg; + uint16_t name_len; + bool used; + char name[DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES]; +} ds4_gpu_decode_pipeline_fast_cache_entry; + +static ds4_gpu_decode_pipeline_fast_cache_entry + g_decode_pipeline_fast_cache[DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS]; +/* Command encoding and the existing NSMutableDictionary pipeline cache are + * backend-serialized; this hot-path mirror intentionally shares that model. */ +static bool g_decode_pipeline_fast_lookup_active; +static uint32_t g_decode_pipeline_fast_cache_entries; static NSMutableDictionary> *g_model_buffer_cache; static NSMutableDictionary *g_q4_expert_table_cache; static NSMutableDictionary *g_q4_expert_layer_residency_cache; @@ -641,6 +707,14 @@ static uint32_t g_stream_expert_cache_slab_count; static uint32_t g_stream_expert_cache_slab_total_slots; static uint32_t g_stream_expert_cache_free_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES]; static uint32_t g_stream_expert_cache_free_slot_count; +static uint8_t g_stream_expert_cache_slab_slot_locked[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES]; +static uint32_t g_stream_expert_cache_mlock_budget_cap; +static uint8_t g_stream_expert_cache_mlock_relief_applied; +static uint64_t g_stream_expert_cache_mlock_bytes; +static uint64_t g_stream_expert_cache_mlock_fail_bytes; +static uint64_t g_stream_expert_cache_mlock_failures; +static double g_stream_expert_cache_mlock_ms; +static int g_stream_expert_cache_mlock_warned; static uint64_t g_stream_expert_cache_slab_slot_bytes; static uint64_t g_stream_expert_cache_cb_seq; static uint64_t g_stream_expert_cache_done_seq; @@ -908,7 +982,11 @@ static id ds4_gpu_command_buffer(int *owned) { static id ds4_gpu_compute_encoder(id cb) { if (g_batch_cb && cb == g_batch_cb) { g_batch_has_work = YES; - if (!g_batch_enc) g_batch_enc = [cb computeCommandEncoder]; + if (!g_batch_enc) { + g_batch_enc = g_batch_encoder_concurrent + ? [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent] + : [cb computeCommandEncoder]; + } return g_batch_enc; } return [cb computeCommandEncoder]; @@ -929,6 +1007,17 @@ static void ds4_gpu_close_batch_encoder(void) { static double g_gpu_busy_accum; static uint64_t g_gpu_busy_cbs; +/* A failed command buffer can leave a cross-threadgroup arrival counter at an + * arbitrary partial value. Drop cached ownership instead of CPU-resetting + * buffers that another in-flight command buffer might still reference; bound + * resources remain retained by their command buffers/transient list. */ +static void ds4_gpu_invalidate_completion_counters(void) { + [g_dsv4_completion_cache removeAllObjects]; + g_dsv4_hc_producer_last_mix_buffer = nil; + g_dsv4_hc_producer_last_mix_offset = 0; + g_dsv4_hc_producer_last_completion = nil; +} + static int ds4_gpu_wait_command_buffer(id cb, const char *label) { [cb waitUntilCompleted]; if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) { @@ -943,6 +1032,7 @@ static int ds4_gpu_wait_command_buffer(id cb, const char *labe if (cb.status == MTLCommandBufferStatusError) { fprintf(stderr, "ds4: Metal %s failed: %s\n", label, [[cb.error localizedDescription] UTF8String]); + ds4_gpu_invalidate_completion_counters(); return 0; } return 1; @@ -1937,12 +2027,11 @@ static id ds4_gpu_new_transient_buffer(NSUInteger bytes, const char * } static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) { - if (getenv("DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL || + if (getenv("DS4_METAL_DISABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL || getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) { return 0; } - return ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL; + return ds4_gpu_device_name_contains("M3"); } static ds4_gpu_zero_prefix_prefill_mask_cache_entry * @@ -2302,6 +2391,33 @@ static int ds4_gpu_device_name_contains(const char *needle) { return g_metal_device_name[0] != '\0' && strstr(g_metal_device_name, needle) != NULL; } +int ds4_gpu_device_is_pre_m5_apple_silicon(void) { + return strncmp(g_metal_device_name, "Apple M", 7) == 0 && + g_metal_device_name[7] >= '1' && + g_metal_device_name[7] <= '4' && + (g_metal_device_name[8] == '\0' || + g_metal_device_name[8] == ' '); +} + +int ds4_gpu_device_is_m5_apple_silicon(void) { + return strncmp(g_metal_device_name, "Apple M5", 8) == 0 && + (g_metal_device_name[8] == '\0' || + g_metal_device_name[8] == ' '); +} + +static bool ds4_gpu_ported_m5_decode_feature_enabled( + const char *pre_m5_disable_env, + const char *m5_disable_env) { + if (m5_disable_env && getenv(m5_disable_env) != NULL) return false; + const bool pre_m5 = ds4_gpu_device_is_pre_m5_apple_silicon(); + if (pre_m5 && + (getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") != NULL || + (pre_m5_disable_env && getenv(pre_m5_disable_env) != NULL))) { + return false; + } + return pre_m5 || ds4_gpu_device_is_m5_apple_silicon(); +} + static int ds4_gpu_compile_tensor_probe(void) { #if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 if (!g_device) return 0; @@ -2508,12 +2624,161 @@ static const char *ds4_gpu_mul_mm_id_map0_name(uint32_t ne20) { } } +static void ds4_gpu_decode_pipeline_fast_cache_reset(void) { + for (uint32_t i = 0; i < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS; i++) { + ds4_gpu_decode_pipeline_fast_cache_entry *entry = + &g_decode_pipeline_fast_cache[i]; + /* ARC must see every strong release; do not memset this table. */ + entry->pipeline = nil; + entry->hash = 0; + entry->nsg = 0; + entry->name_len = 0; + entry->used = false; + entry->name[0] = '\0'; + } + g_decode_pipeline_fast_lookup_active = false; + g_decode_pipeline_fast_cache_entries = 0; +} + +static bool ds4_gpu_decode_pipeline_fast_key( + const char *function_name, + int16_t nsg, + int16_t nxpsg, + uint16_t *name_len_out, + uint64_t *hash_out) { + if (!function_name || !name_len_out || !hash_out) return false; + + uint64_t hash = UINT64_C(14695981039346656037); + size_t name_len = 0; + while (function_name[name_len] != '\0') { + if (name_len + 1u >= DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES) { + return false; + } + hash ^= (uint8_t)function_name[name_len++]; + hash *= UINT64_C(1099511628211); + } + const uint16_t nsg_bits = (uint16_t)nsg; + hash ^= (uint8_t)nsg_bits; + hash *= UINT64_C(1099511628211); + hash ^= (uint8_t)(nsg_bits >> 8u); + hash *= UINT64_C(1099511628211); + const uint16_t nxpsg_bits = (uint16_t)nxpsg; + hash ^= (uint8_t)nxpsg_bits; + hash *= UINT64_C(1099511628211); + hash ^= (uint8_t)(nxpsg_bits >> 8u); + hash *= UINT64_C(1099511628211); + + *name_len_out = (uint16_t)name_len; + *hash_out = hash; + return true; +} + +static id ds4_gpu_decode_pipeline_fast_cache_lookup( + const char *function_name, + int16_t nsg, + int16_t nxpsg, + uint16_t name_len, + uint64_t hash) { + const uint32_t mask = DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS - 1u; + uint32_t slot = (uint32_t)hash & mask; + for (uint32_t probe = 0; + probe < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS; + probe++, slot = (slot + 1u) & mask) { + ds4_gpu_decode_pipeline_fast_cache_entry *entry = + &g_decode_pipeline_fast_cache[slot]; + if (!entry->used) break; + if (entry->hash == hash && + entry->nsg == nsg && + entry->nxpsg == nxpsg && + entry->name_len == name_len && + memcmp(entry->name, function_name, name_len) == 0) { + return entry->pipeline; + } + } + return nil; +} + +static void ds4_gpu_decode_pipeline_fast_cache_insert( + const char *function_name, + int16_t nsg, + int16_t nxpsg, + uint16_t name_len, + uint64_t hash, + id pipeline) { + if (!pipeline || + g_decode_pipeline_fast_cache_entries >= + DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS) { + return; + } + + const uint32_t mask = DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS - 1u; + uint32_t slot = (uint32_t)hash & mask; + for (uint32_t probe = 0; + probe < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS; + probe++, slot = (slot + 1u) & mask) { + ds4_gpu_decode_pipeline_fast_cache_entry *entry = + &g_decode_pipeline_fast_cache[slot]; + if (entry->used) { + if (entry->hash == hash && + entry->nsg == nsg && + entry->nxpsg == nxpsg && + entry->name_len == name_len && + memcmp(entry->name, function_name, name_len) == 0) { + entry->pipeline = pipeline; + return; + } + continue; + } + memcpy(entry->name, function_name, name_len); + entry->name[name_len] = '\0'; + entry->hash = hash; + entry->nsg = nsg; + entry->nxpsg = nxpsg; + entry->name_len = name_len; + entry->pipeline = pipeline; + entry->used = true; + g_decode_pipeline_fast_cache_entries++; + return; + } +} + +int ds4_gpu_set_decode_pipeline_fast_lookup(int enabled) { + const int previous = g_decode_pipeline_fast_lookup_active ? 1 : 0; + g_decode_pipeline_fast_lookup_active = enabled != 0; + return previous; +} + static id ds4_gpu_get_mul_mv_pipeline( const char *function_name, int16_t nsg) { + uint16_t fast_name_len = 0; + uint64_t fast_hash = 0; + const bool fast_key_valid = + g_decode_pipeline_fast_lookup_active && + ds4_gpu_decode_pipeline_fast_key( + function_name, nsg, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + &fast_name_len, &fast_hash); + if (fast_key_valid) { + id fast_cached = + ds4_gpu_decode_pipeline_fast_cache_lookup( + function_name, nsg, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + fast_name_len, fast_hash); + if (fast_cached) return fast_cached; + } + NSString *key = [NSString stringWithFormat:@"%s_nsg=%d", function_name, (int)nsg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + if (fast_key_valid) { + ds4_gpu_decode_pipeline_fast_cache_insert( + function_name, nsg, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + fast_name_len, fast_hash, cached); + } + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; @@ -2538,17 +2803,256 @@ static id ds4_gpu_get_mul_mv_pipeline( } [g_pipeline_cache setObject:pipeline forKey:key]; + if (fast_key_valid) { + ds4_gpu_decode_pipeline_fast_cache_insert( + function_name, nsg, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + fast_name_len, fast_hash, pipeline); + } return pipeline; } +/* The ordinary mul-mv cache key covers function name and nsg only. Keep the + * descriptor-hinted PSO separate so a cache hit cannot erase this compiler + * contract or substitute it for the fallback pipeline. */ +static id ds4_gpu_new_mul_mv_tg_multiple_pipeline( + const char *function_name, + int16_t nsg) { + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s tg-multiple function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + MTLComputePipelineDescriptor *descriptor = [[MTLComputePipelineDescriptor alloc] init]; + descriptor.label = name; + descriptor.computeFunction = fn; + descriptor.threadGroupSizeIsMultipleOfThreadExecutionWidth = YES; + + error = nil; + id pipeline = + [g_device newComputePipelineStateWithDescriptor:descriptor + options:MTLPipelineOptionNone + reflection:nil + error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s tg-multiple pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + return pipeline; +} + +int ds4_gpu_test_decode_pipeline_fast_lookup(void) { + if (!g_initialized || !g_pipeline_cache || !g_library || !g_device) { + return 0; + } + + static const char base_name[] = "kernel_mul_mv_id_mxfp4_f32"; + char copied_name[sizeof(base_name)]; + memcpy(copied_name, base_name, sizeof(base_name)); + + ds4_gpu_decode_pipeline_fast_cache_reset(); + id baseline = + ds4_gpu_get_mul_mv_pipeline(base_name, 2); + if (!baseline || + g_decode_pipeline_fast_cache_entries != 0u) { + ds4_gpu_decode_pipeline_fast_cache_reset(); + return 0; + } + + (void)ds4_gpu_set_decode_pipeline_fast_lookup(1); + id populated = + ds4_gpu_get_mul_mv_pipeline(base_name, 2); + id hit = + ds4_gpu_get_mul_mv_pipeline(copied_name, 2); + uint16_t copied_name_len = 0; + uint64_t copied_hash = 0; + const bool copied_key_valid = + ds4_gpu_decode_pipeline_fast_key( + copied_name, 2, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + &copied_name_len, &copied_hash); + id direct_hit = + copied_key_valid ? + ds4_gpu_decode_pipeline_fast_cache_lookup( + copied_name, 2, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + copied_name_len, copied_hash) : nil; + id distinct_name = + ds4_gpu_get_mul_mv_pipeline( + "kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 2); + id distinct_nsg = + ds4_gpu_get_mul_mv_pipeline(base_name, 4); + const bool populated_and_hit = + populated == baseline && + hit == baseline && + direct_hit == baseline && + distinct_name != nil && + distinct_nsg != nil && + g_decode_pipeline_fast_cache_entries == 3u; + + (void)ds4_gpu_set_decode_pipeline_fast_lookup(0); + const uint32_t entries_before_disabled_lookup = + g_decode_pipeline_fast_cache_entries; + id disabled = + ds4_gpu_get_mul_mv_pipeline(base_name, 2); + const bool inactive_unchanged = + disabled == baseline && + g_decode_pipeline_fast_cache_entries == entries_before_disabled_lookup; + + ds4_gpu_decode_pipeline_fast_cache_reset(); + (void)ds4_gpu_set_decode_pipeline_fast_lookup(1); + bool filled = true; + for (uint32_t i = 0; + i < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS && filled; + i++) { + char synthetic_name[DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES]; + const int written = snprintf( + synthetic_name, sizeof(synthetic_name), + "ds4_test_decode_pipeline_fast_%u", i); + uint16_t synthetic_name_len = 0; + uint64_t synthetic_hash = 0; + filled = written > 0 && + (size_t)written < sizeof(synthetic_name) && + ds4_gpu_decode_pipeline_fast_key( + synthetic_name, 2, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + &synthetic_name_len, &synthetic_hash); + if (filled) { + ds4_gpu_decode_pipeline_fast_cache_insert( + synthetic_name, 2, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + synthetic_name_len, + synthetic_hash, baseline); + } + } + id full_fallback = + ds4_gpu_get_mul_mv_pipeline(base_name, 2); + const bool full_table_unchanged = + filled && full_fallback == baseline && + g_decode_pipeline_fast_cache_entries == + DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS; + + ds4_gpu_decode_pipeline_fast_cache_reset(); + return populated_and_hit && inactive_unchanged && full_table_unchanged; +} + +/* + * Exercises the second SIMD-group constant key shape used by + * ds4_gpu_get_mul_mv_ext_pipeline (name, nsg, nxpsg). Confirms the extended key + * populates and hits, distinguishes a real nxpsg from the single-constant + * sentinel, and coexists with the single-constant mv entries in the shared + * 64-slot table. Uses baseline as a stand-in pipeline object so the guard does + * not depend on compiling a specific extended-kernel function. + */ +int ds4_gpu_test_decode_pipeline_fast_lookup_ext(void) { + if (!g_initialized || !g_pipeline_cache || !g_library || !g_device) { + return 0; + } + + static const char base_name[] = "kernel_mul_mv_id_mxfp4_f32"; + char copied_name[sizeof(base_name)]; + memcpy(copied_name, base_name, sizeof(base_name)); + + id baseline = + ds4_gpu_get_mul_mv_pipeline(base_name, 2); + if (!baseline) return 0; + + ds4_gpu_decode_pipeline_fast_cache_reset(); + (void)ds4_gpu_set_decode_pipeline_fast_lookup(1); + + const int16_t ext_nsg = 2; + const int16_t ext_nxpsg = 4; + uint16_t ext_name_len = 0; + uint64_t ext_hash = 0; + const bool ext_key_valid = + ds4_gpu_decode_pipeline_fast_key( + copied_name, ext_nsg, ext_nxpsg, &ext_name_len, &ext_hash); + id ext_populated = nil; + if (ext_key_valid) { + ds4_gpu_decode_pipeline_fast_cache_insert( + copied_name, ext_nsg, ext_nxpsg, ext_name_len, ext_hash, baseline); + ext_populated = ds4_gpu_decode_pipeline_fast_cache_lookup( + copied_name, ext_nsg, ext_nxpsg, ext_name_len, ext_hash); + } + + /* + * The single-constant sentinel must not match an extended key even when the + * name and nsg agree, proving the two key shapes never alias. + */ + uint16_t none_name_len = 0; + uint64_t none_hash = 0; + const bool none_key_valid = + ds4_gpu_decode_pipeline_fast_key( + copied_name, ext_nsg, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + &none_name_len, &none_hash); + id none_hit = + none_key_valid ? + ds4_gpu_decode_pipeline_fast_cache_lookup( + copied_name, ext_nsg, + DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE, + none_name_len, none_hash) : nil; + + const int16_t alt_nxpsg = 8; + uint16_t alt_name_len = 0; + uint64_t alt_hash = 0; + const bool alt_key_valid = + ds4_gpu_decode_pipeline_fast_key( + copied_name, ext_nsg, alt_nxpsg, &alt_name_len, &alt_hash); + id alt_hit = + alt_key_valid ? + ds4_gpu_decode_pipeline_fast_cache_lookup( + copied_name, ext_nsg, alt_nxpsg, alt_name_len, alt_hash) : nil; + + const bool ext_populated_and_hit = + ext_key_valid && + ext_populated == baseline && + g_decode_pipeline_fast_cache_entries == 1u && + none_hit == nil && + alt_hit == nil; + + (void)ds4_gpu_set_decode_pipeline_fast_lookup(0); + ds4_gpu_decode_pipeline_fast_cache_reset(); + return ext_populated_and_hit; +} + static id ds4_gpu_get_mul_mv_ext_pipeline( const char *function_name, int16_t nsg, int16_t nxpsg) { + uint16_t fast_name_len = 0; + uint64_t fast_hash = 0; + const bool fast_key_valid = + g_decode_pipeline_fast_lookup_active && + ds4_gpu_decode_pipeline_fast_key( + function_name, nsg, nxpsg, &fast_name_len, &fast_hash); + if (fast_key_valid) { + id fast_cached = + ds4_gpu_decode_pipeline_fast_cache_lookup( + function_name, nsg, nxpsg, fast_name_len, fast_hash); + if (fast_cached) return fast_cached; + } + NSString *key = [NSString stringWithFormat:@"%s_nsg=%d_nxpsg=%d", function_name, (int)nsg, (int)nxpsg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + if (fast_key_valid) { + ds4_gpu_decode_pipeline_fast_cache_insert( + function_name, nsg, nxpsg, fast_name_len, fast_hash, cached); + } + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; @@ -2574,16 +3078,42 @@ static id ds4_gpu_get_mul_mv_ext_pipeline( } [g_pipeline_cache setObject:pipeline forKey:key]; + if (fast_key_valid) { + ds4_gpu_decode_pipeline_fast_cache_insert( + function_name, nsg, nxpsg, fast_name_len, fast_hash, pipeline); + } return pipeline; } static id ds4_gpu_get_flash_attn_pad_pipeline( bool has_mask, int32_t ncpsg) { + /* + * Decode calls this once per layer with identical arguments, so memoize + * the last hit and skip the NSString key + dictionary lookup on the hot + * path. The generic cache below remains the fallback for new variants. + * The rollback switch restores the dictionary path for same-binary A/B. + */ + static struct { + bool m; + int32_t nc; + id pipeline; + } memo; + const bool memo_disabled = + getenv("DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PAD_BLK_MEMO") != NULL; + if (!memo_disabled && memo.pipeline && memo.m == has_mask && memo.nc == ncpsg) { + return memo.pipeline; + } + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_pad_mask=%d_ncpsg=%d", has_mask ? 1 : 0, (int)ncpsg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + if (!memo_disabled) { + memo = (typeof(memo)){ has_mask, ncpsg, cached }; + } + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:100]; @@ -2608,16 +3138,41 @@ static id ds4_gpu_get_flash_attn_pad_pipeline( } [g_pipeline_cache setObject:pipeline forKey:key]; + if (!memo_disabled) { + memo = (typeof(memo)){ has_mask, ncpsg, pipeline }; + } return pipeline; } static id ds4_gpu_get_flash_attn_blk_pipeline( int32_t nqptg, int32_t ncpsg) { + /* + * Decode calls this once per layer with identical arguments, so memoize + * the last hit and skip the NSString key + dictionary lookup on the hot + * path. The generic cache below remains the fallback for new variants. + * The rollback switch restores the dictionary path for same-binary A/B. + */ + static struct { + int32_t nq; + int32_t nc; + id pipeline; + } memo; + const bool memo_disabled = + getenv("DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PAD_BLK_MEMO") != NULL; + if (!memo_disabled && memo.pipeline && memo.nq == nqptg && memo.nc == ncpsg) { + return memo.pipeline; + } + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_blk_nqptg=%d_ncpsg=%d", (int)nqptg, (int)ncpsg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + if (!memo_disabled) { + memo = (typeof(memo)){ nqptg, ncpsg, cached }; + } + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&nqptg type:MTLDataTypeInt atIndex:224]; @@ -2642,6 +3197,9 @@ static id ds4_gpu_get_flash_attn_blk_pipeline( } [g_pipeline_cache setObject:pipeline forKey:key]; + if (!memo_disabled) { + memo = (typeof(memo)){ nqptg, ncpsg, pipeline }; + } return pipeline; } @@ -2656,6 +3214,29 @@ static id ds4_gpu_get_flash_attn_pipeline( int32_t ns10, int32_t ns20, int32_t nsg) { + /* + * Prefill and batched decode call this once per layer with identical + * arguments, so memoize the last hit and skip the NSString key + dictionary + * lookup on the hot path. The generic cache below remains the fallback for + * new variants. The rollback switch restores the dictionary path for + * same-binary A/B. + */ + static struct { + const char *fn; + bool m, s, b, c, k, bc; + int32_t n10, n20, sg; + id pipeline; + } memo; + const bool memo_disabled = + getenv("DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_BATCHED_MEMO") != NULL; + if (!memo_disabled && memo.pipeline && memo.fn != NULL && + strcmp(memo.fn, function_name) == 0 && + memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias && + memo.c == has_scap && memo.k == has_kvpad && memo.bc == bc_mask && + memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg) { + return memo.pipeline; + } + NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d", function_name, has_mask ? 1 : 0, @@ -2668,7 +3249,14 @@ static id ds4_gpu_get_flash_attn_pipeline( (int)ns20, (int)nsg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + if (!memo_disabled) { + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, bc_mask, ns10, ns20, + nsg, cached }; + } + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300]; @@ -2701,6 +3289,11 @@ static id ds4_gpu_get_flash_attn_pipeline( } [g_pipeline_cache setObject:pipeline forKey:key]; + if (!memo_disabled) { + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, bc_mask, ns10, ns20, + nsg, pipeline }; + } return pipeline; } @@ -2792,6 +3385,56 @@ static id ds4_gpu_get_flash_attn_vec_pipeline( return pipeline; } +/* Pipeline for the RoPE-fused decode reduce. Same function constants as the + * plain reduce so the split-K geometry is identical. */ +static id ds4_gpu_get_flash_attn_reduce_rope_pipeline( + int32_t dv, + int32_t nwg) { + static int32_t memo_dv, memo_nwg; + static id memo_pipeline; + if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) { + return memo_pipeline; + } + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_rope_dv=%d_nwg=%d", + (int)dv, (int)nwg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) { + memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached; + return cached; + } + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500]; + [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501]; + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce_rope" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce_rope not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce_rope pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + [g_pipeline_cache setObject:pipeline forKey:key]; + memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline; + return pipeline; +} + +int ds4_gpu_decode_attn_rope_fuse_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_rope_tail_inplace_pair_affine_pipeline == nil) return 0; + if (getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") != NULL) return 0; + if (getenv("DS4_METAL_DISABLE_AFFINE_ROPE_PAIR") != NULL) return 0; + if (!ds4_gpu_device_name_contains("M3") && !ds4_gpu_device_name_contains("M5")) return 0; + return 1; +} + static id ds4_gpu_get_flash_attn_reduce_pipeline( int32_t dv, int32_t nwg) { @@ -3391,6 +4034,15 @@ void ds4_gpu_print_memory_report(const char *label) { (unsigned long long)g_stream_expert_cache_buffer_allocs, (unsigned long long)g_stream_expert_cache_buffer_reuses); } + if (g_stream_expert_cache_mlock_bytes != 0 || + g_stream_expert_cache_mlock_failures != 0) { + fprintf(stderr, + "ds4: streaming expert buffer mlock locked=%.2f GiB failed=%.2f GiB failures=%llu time=%.3f ms\n", + ds4_gpu_gib(g_stream_expert_cache_mlock_bytes), + ds4_gpu_gib(g_stream_expert_cache_mlock_fail_bytes), + (unsigned long long)g_stream_expert_cache_mlock_failures, + g_stream_expert_cache_mlock_ms); + } if (ds4_gpu_stream_expert_timing_summary_enabled()) { const ds4_gpu_stream_expert_timing_snapshot total = ds4_gpu_stream_expert_timing_current(); @@ -4333,15 +4985,6 @@ static ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) { const uint64_t default_nsg = ds4_gpu_tp_world_is_two() ? 2u : 4u; const int16_t nsg = (int16_t)ds4_gpu_env_u64("DS4_METAL_Q8_MV_NSG", default_nsg, 1u, 8u); - const uint64_t rows = ds4_gpu_env_u64("DS4_METAL_Q8_MV_ROWS", 2u, 2u, 4u); - if (rows >= 4u) { - return (ds4_gpu_mv_dispatch) { - .function_name = "kernel_mul_mv_q8_0_f32_r4", - .nsg = nsg, - .nr0 = 4, - .smem = 32u * 4u * sizeof(float), - }; - } return (ds4_gpu_mv_dispatch) { .function_name = "kernel_mul_mv_q8_0_f32", .nsg = nsg, @@ -4593,21 +5236,6 @@ typedef struct { uint64_t nb1; } ds4_gpu_hc_weighted_sum_args; -typedef struct { - int64_t n_embd; - int64_t n_hc; - int64_t n_tokens; - uint64_t nb_x0; - uint64_t nb_x1; - uint64_t nb_x2; - uint64_t nb_w0; - uint64_t nb_w1; - uint64_t nb0; - uint64_t nb1; - uint64_t nb_norm1; - float norm_eps; -} ds4_gpu_hc_weighted_sum_norm_args; - typedef struct { float post_scale; float eps; @@ -4771,18 +5399,7 @@ static int ds4_gpu_encode_attn_out_low_q8_direct( NSUInteger nsg, bool rows_per_group_is_nr0); -static int ds4_gpu_encode_attn_out_low_q8_mpp( - id cb, - id pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off); - -static int ds4_gpu_encode_attn_out_low_q8_mpp( +static int ds4_gpu_encode_attn_out_low_mpp( id cb, id pipeline, const ds4_gpu_mul_mm_id_args *mm_args, @@ -4999,6 +5616,52 @@ typedef struct { float beta_slow; } ds4_gpu_rope_affine_pair_args; +/* Matches ds4_metal_args_dsv4_comp_finalize in dsv4_rope.metal. */ +typedef struct { + ds4_gpu_rope_affine_pair_args rope; + float rms_eps; + uint32_t pad0; +} ds4_gpu_comp_finalize_args; + +/* Set by ds4.c immediately before the decode attention call when the inverse + * RoPE tail is deferred into the reduce kernel. Cleared by the encoder. */ +static ds4_gpu_rope_affine_pair_args g_decode_attn_rope_args; +static int g_decode_attn_rope_fuse; +/* Set only by the encoder that actually applied the deferred rotation, so ds4.c + * can fall back to the standalone RoPE on any attention path that does not + * consume it (for example ratio-0 layers that take the raw-heads encoder). */ +static int g_decode_attn_rope_fuse_used; + +void ds4_gpu_set_decode_attn_rope_fuse( + uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, + bool inverse, float freq_base, float freq_scale, float ext_factor, + float attn_factor, float beta_fast, float beta_slow) { + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + g_decode_attn_rope_args = (ds4_gpu_rope_affine_pair_args) { + .row_bytes = row_bytes, + .token_bytes = row_bytes, + .head_dim = (int32_t)head_dim, + .n_dims = (int32_t)n_rot, + .n_ctx_orig = (int32_t)n_ctx_orig, + .inverse = inverse ? 1 : 0, + .pos0 = pos0, + .pos_step = 1, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + g_decode_attn_rope_fuse = 1; + g_decode_attn_rope_fuse_used = 0; +} + +int ds4_gpu_decode_attn_rope_fuse_used(void) { + return g_decode_attn_rope_fuse_used; +} + + _Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64, "Metal affine RoPE argument ABI changed"); @@ -5057,43 +5720,30 @@ static int ds4_gpu_encode_rope_tail_inplace( const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u; const bool lane_compatible = tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u; - const bool force_affine_position = - getenv("DS4_METAL_ENABLE_AFFINE_ROPE_PAIR") != NULL; const bool use_inplace_pair = g_rope_tail_inplace_pair_pipeline != nil && + getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") == NULL && args->mode == 0 && !args->src2 && lane_compatible && (ds4_gpu_device_name_contains("M3") || - (ds4_gpu_device_name_contains("M5") && n_tok == 1u) || - getenv("DS4_METAL_ENABLE_INPLACE_ROPE_PAIR") != NULL || - force_affine_position) && - getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL; + (ds4_gpu_device_name_contains("M5") && n_tok == 1u)); const bool use_shared_coeff = - use_inplace_pair && !force_affine_position && + use_inplace_pair && g_rope_tail_inplace_pair_shared4_pipeline != nil && + getenv("DS4_METAL_DISABLE_SHARED_ROPE_COEFF") == NULL && /* The 256-thread grouped schedule helps long prefill, but reduces the * per-head parallelism that short batches and decode rely on. */ - tail_threads == 64u && n_head >= 4u && n_tok >= 32u && - getenv("DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF") == NULL; - if (force_affine_position && - g_rope_tail_inplace_pair_affine_pipeline == nil && - getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL && - getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL) { - fprintf(stderr, - "ds4: forced affine-position RoPE pipeline is unavailable\n"); - return 0; - } + tail_threads == 64u && n_head >= 4u && n_tok >= 32u; /* Keep long prefill on the proven shared4 kernel. Reconstructing affine * positions inside its coefficient cohort perturbs YaRN fast-math codegen; * the compact affine specialization is exact for the decode pair schedule. */ const bool use_affine_position = use_inplace_pair && !use_shared_coeff && g_rope_tail_inplace_pair_affine_pipeline != nil && - (n_tok == 1u || force_affine_position) && + getenv("DS4_METAL_DISABLE_AFFINE_ROPE_PAIR") == NULL && + n_tok == 1u && (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || - force_affine_position) && - getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL; + ds4_gpu_device_name_contains("M5")); int32_t pos_stack[256]; int32_t *pos = NULL; @@ -5653,7 +6303,7 @@ typedef struct { uint32_t window; uint32_t ratio; uint32_t comp_kv_f16; - uint32_t pad0; + uint32_t n_splits; uint64_t q_token_stride; uint64_t q_head_stride; uint64_t raw_row_stride; @@ -5718,6 +6368,7 @@ int ds4_gpu_init(void) { if (g_initialized) return 1; @autoreleasepool { + ds4_gpu_decode_pipeline_fast_cache_reset(); g_device = MTLCreateSystemDefaultDevice(); if (!g_device) { fprintf(stderr, "ds4: Metal device not available\n"); @@ -5739,14 +6390,18 @@ int ds4_gpu_init(void) { g_q4_expert_table_cache = [NSMutableDictionary dictionary]; g_q4_expert_layer_residency_cache = [NSMutableDictionary dictionary]; g_pipeline_cache = [NSMutableDictionary dictionary]; + g_dsv4_completion_cache = [NSCache new]; + g_dsv4_completion_cache.countLimit = 256u; g_transient_buffers = [NSMutableArray array]; g_pending_cbs = [NSMutableArray array]; if (!g_model_buffer_cache || !g_q4_expert_table_cache || !g_q4_expert_layer_residency_cache || - !g_pipeline_cache || !g_transient_buffers || !g_pending_cbs) { + !g_pipeline_cache || !g_dsv4_completion_cache || + !g_transient_buffers || !g_pending_cbs) { fprintf(stderr, "ds4: Metal bookkeeping allocation failed\n"); g_pending_cbs = nil; g_transient_buffers = nil; + g_dsv4_completion_cache = nil; g_pipeline_cache = nil; g_q4_expert_layer_residency_cache = nil; g_q4_expert_table_cache = nil; @@ -6155,6 +6810,23 @@ int ds4_gpu_init(void) { return 0; } + fn = [library newFunctionWithName:@"kernel_dsv4_kv_rope_fp8_store_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_kv_rope_fp8_store_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + error = nil; + g_dsv4_kv_rope_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_kv_rope_fp8_store_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_kv_rope_fp8_store_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + fn = [library newFunctionWithName:@"kernel_dsv4_ratio4_shift_f32"]; if (!fn) { fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 function not found\n"); @@ -6464,6 +7136,29 @@ int ds4_gpu_init(void) { return 0; } + fn = [library newFunctionWithName:@"kernel_dsv4_head_rms_norm_rope_tail_f32"]; + if (!fn) { + fprintf(stderr, + "ds4: Metal kernel_dsv4_head_rms_norm_rope_tail_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + error = nil; + g_dsv4_head_rms_norm_rope_tail_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_head_rms_norm_rope_tail_pipeline) { + fprintf(stderr, + "ds4: Metal kernel_dsv4_head_rms_norm_rope_tail_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_use_dsv4_head_rms_norm_rope_tail_pipeline = + ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_HEAD_RMS_ROPE_PIPELINE_STATIC") == NULL; + MTLFunctionConstantValues *moe_mv_id_constants = [[MTLFunctionConstantValues alloc] init]; int16_t moe_mv_id_nsg = 2; [moe_mv_id_constants setConstantValue:&moe_mv_id_nsg type:MTLDataTypeShort atIndex:600]; @@ -6528,6 +7223,21 @@ int ds4_gpu_init(void) { return 0; } + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_pack2_overlap_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (fn) { + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + } + if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline) { + fprintf(stderr, + "ds4: optional Metal IQ2 pair pack2 pipeline unavailable: %s\n", + error ? [[error localizedDescription] UTF8String] : + "function not found"); + } + error = nil; fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_f32" constantValues:moe_mv_id_constants @@ -7028,6 +7738,71 @@ int ds4_gpu_init(void) { return 0; } + g_moe_mul_mv_id_mxfp4_pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_f32", 2); + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 2); + g_moe_mul_mv_id_mxfp4_sum6_pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_sum6_f32", 2); + /* Single-simdgroup variants of the MXFP4 decode MoE kernels. The + * per-row per-lane arithmetic is identical; only the row-to-simdgroup + * mapping changes, so outputs remain bit-exact. Smaller threadgroups + * measurably improve GPU scheduling granularity on these + * latency-bound small-grid matvecs. Pipelines are built + * unconditionally; the dispatch-time selection gate reads the + * environment so A/B harnesses can toggle it after engine creation. */ + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 1); + g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_sum6_f32", 1); + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 1); + g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_sum6_f32", 1); + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_f32", 1); + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_sum6_fixed_route_f32", 1); + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_f32", 1); + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_f32", 1); + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32", 1); + if (!g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple || + !g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple || + !g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1) { + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_slots6_mxfp4_pair_swiglu_f32", 2); + g_moe_mul_mv_slots6_mxfp4_sum6_pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_slots6_mxfp4_sum6_f32", 2); + if (!g_moe_mul_mv_id_mxfp4_pipeline || + !g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline || + !g_moe_mul_mv_id_mxfp4_sum6_pipeline || + !g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline || + !g_moe_mul_mv_slots6_mxfp4_sum6_pipeline) { + g_queue = nil; + g_device = nil; + return 0; + } + error = nil; fn = [library newFunctionWithName:@"kernel_q4_gather_slots6"]; if (!fn) { @@ -7414,21 +8189,6 @@ int ds4_gpu_init(void) { return 0; } - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum_norm4"]; - if (fn) { - g_hc_weighted_sum_norm_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_weighted_sum_norm_pipeline) { - fprintf(stderr, - "ds4: optional Metal output HC sum/RMSNorm pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, - "ds4: optional Metal output HC sum/RMSNorm kernel unavailable\n"); - } - error = nil; fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"]; if (fn) { @@ -7662,6 +8422,15 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one"); g_dsv4_compressor_pack_ratio4_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4"); + g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline = + ds4_gpu_get_pipeline( + "kernel_dsv4_compressor_pack_ratio4_decode_ggml"); + g_dsv4_compressor_exact_softmax_product_pipeline = + ds4_gpu_get_pipeline( + "kernel_dsv4_compressor_exact_softmax_product_ratio4"); + g_dsv4_compressor_exact_pool_ratio4_pipeline = + ds4_gpu_get_pipeline( + "kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml"); g_dsv4_softmax_pool_ratio4_direct_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct"); g_rms_norm_scale_pipeline = @@ -7672,6 +8441,12 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8"); g_dsv4_indexed_attention_heads8_rb16_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16"); + g_dsv4_indexed_attention_heads16_dual_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual"); + g_dsv4_indexed_attention_heads8_split_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split"); + g_dsv4_indexed_attention_heads8_split_reduce_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split_reduce"); g_dsv4_softplus_sqrt_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4"); g_dsv4_router_finalize_one_pipeline = @@ -7680,6 +8455,9 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd"); g_dsv4_router_finalize_weights_one_simd_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd"); + g_dsv4_router_transform_finalize_weights_one_simd_pipeline = + ds4_gpu_get_pipeline( + "kernel_dsv4_router_transform_finalize_weights_one_simd"); g_dsv4_router_weights_one_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one"); g_glm_router_select_one_pipeline = @@ -7809,6 +8587,9 @@ int ds4_gpu_init(void) { !g_dsv4_sort_i32_rows_asc_pipeline || !g_dsv4_indexed_attention_heads8_pipeline || !g_dsv4_indexed_attention_heads8_rb16_pipeline || + !g_dsv4_indexed_attention_heads16_dual_pipeline || + !g_dsv4_indexed_attention_heads8_split_pipeline || + !g_dsv4_indexed_attention_heads8_split_reduce_pipeline || !g_dsv4_softplus_sqrt_pipeline || !g_dsv4_router_finalize_one_pipeline || !g_dsv4_router_weights_one_pipeline || @@ -7883,6 +8664,48 @@ int ds4_gpu_init(void) { return 1; } +int ds4_gpu_test_mxfp4_down_half_lut(uint16_t *legacy_bits, + uint16_t *lut_bits) { + if (!legacy_bits || !lut_bits) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + + @autoreleasepool { + const NSUInteger count = 4096u; + const NSUInteger bytes = count * sizeof(uint16_t); + id pipeline = + ds4_gpu_get_pipeline("kernel_test_mxfp4_down_half_lut"); + id legacy = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + id lut = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + id cb = ds4_gpu_new_command_buffer(); + if (!pipeline || !legacy || !lut || !cb) { + fprintf(stderr, "ds4: Metal MXFP4 half-LUT test setup failed\n"); + return 0; + } + + id enc = [cb computeCommandEncoder]; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBuffer:legacy offset:0 atIndex:0]; + [enc setBuffer:lut offset:0 atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake((count + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)]; + [enc endEncoding]; + [cb commit]; + if (!ds4_gpu_wait_command_buffer(cb, "MXFP4 half-LUT raw-bit test")) { + return 0; + } + memcpy(legacy_bits, [legacy contents], bytes); + memcpy(lut_bits, [lut contents], bytes); + } + return 1; +} + +void ds4_gpu_test_set_flags(uint32_t flags) { + g_test_flags = flags; +} + ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { if (!g_initialized && !ds4_gpu_init()) return NULL; if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL; @@ -8162,7 +8985,15 @@ int ds4_gpu_pack_slot_rows_f32_tensor( int ds4_gpu_begin_commands(void) { if (!g_initialized && !ds4_gpu_init()) return 0; + /* A failed concurrent FFN must never affect the next command batch. */ + ds4_gpu_parallel_ffn_reset_state(YES); if (g_batch_cb) return 0; + /* Refresh once per command batch so same-engine A/B runs can toggle the + * static PSO without paying for environment lookups in every layer. */ + g_use_dsv4_head_rms_norm_rope_tail_pipeline = + !g_ssd_streaming_mode && + ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_HEAD_RMS_ROPE_PIPELINE_STATIC") == NULL; g_batch_cb = ds4_gpu_new_command_buffer(); g_batch_has_work = NO; if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); @@ -8171,6 +9002,7 @@ int ds4_gpu_begin_commands(void) { int ds4_gpu_flush_encoder(void) { if (!g_initialized && !ds4_gpu_init()) return 0; + ds4_gpu_parallel_ffn_reset_state(YES); if (!g_batch_cb) return 0; ds4_gpu_close_batch_encoder(); return 1; @@ -8178,6 +9010,7 @@ int ds4_gpu_flush_encoder(void) { int ds4_gpu_flush_commands(void) { if (!g_initialized && !ds4_gpu_init()) return 0; + ds4_gpu_parallel_ffn_reset_state(YES); if (!g_batch_cb) return 0; ds4_gpu_close_batch_encoder(); @@ -8203,6 +9036,311 @@ int ds4_gpu_commands_active(void) { return g_batch_cb != nil; } +/* Exact M5 full-FFN overlap inside one concurrent compute encoder. Shared + * gate/up and routed IQ2 pair-SwiGLU launch together; explicit level barriers + * precede the routed Q2 and shared Q8 down consumers. */ +static id g_parallel_q8_pipeline; +static id g_parallel_q8_weight; +static id g_parallel_q8_x; +static id g_parallel_q8_out; +static NSUInteger g_parallel_q8_weight_offset; +static NSUInteger g_parallel_q8_x_offset; +static NSUInteger g_parallel_q8_out_offset; +static ds4_gpu_q8_0_matvec_args g_parallel_q8_args; +static id g_parallel_gate_up_pipeline; +static id g_parallel_gate_weight; +static id g_parallel_up_weight; +static id g_parallel_gate_x; +static id g_parallel_gate_out; +static id g_parallel_up_out; +static id g_parallel_mid_out; +static NSUInteger g_parallel_gate_weight_offset; +static NSUInteger g_parallel_up_weight_offset; +static NSUInteger g_parallel_gate_x_offset; +static NSUInteger g_parallel_gate_out_offset; +static NSUInteger g_parallel_up_out_offset; +static NSUInteger g_parallel_mid_out_offset; +static ds4_gpu_q8_0_matvec_args g_parallel_gate_up_args; +static float g_parallel_gate_up_clamp; +static NSUInteger g_parallel_gate_up_nsg; +static NSUInteger g_parallel_gate_up_nr0; +static NSUInteger g_parallel_gate_up_smem; +static int g_parallel_ffn_mode; /* 2: gate/up + down */ +static int g_parallel_ffn_stage; +static BOOL g_parallel_q8_pending; +static BOOL g_parallel_q8_encoded; + +/* Reset is deliberately idempotent. Closing the concurrent encoder preserves + * work already encoded, while clearing every admission/reference field keeps + * a failed FFN path from turning later ordinary dispatches concurrent. */ +static void ds4_gpu_parallel_ffn_reset_state(BOOL close_encoder) { + if (close_encoder && + (g_batch_encoder_concurrent || g_parallel_q8_pending || + g_parallel_q8_encoded || g_parallel_ffn_mode != 0 || + g_parallel_ffn_stage != 0)) { + ds4_gpu_close_batch_encoder(); + } + + g_batch_encoder_concurrent = NO; + g_parallel_q8_pending = NO; + g_parallel_q8_encoded = NO; + g_parallel_ffn_mode = 0; + g_parallel_ffn_stage = 0; + + g_parallel_q8_pipeline = nil; + g_parallel_q8_weight = nil; + g_parallel_q8_x = nil; + g_parallel_q8_out = nil; + g_parallel_q8_weight_offset = 0; + g_parallel_q8_x_offset = 0; + g_parallel_q8_out_offset = 0; + g_parallel_q8_args = (ds4_gpu_q8_0_matvec_args){0}; + + g_parallel_gate_up_pipeline = nil; + g_parallel_gate_weight = nil; + g_parallel_up_weight = nil; + g_parallel_gate_x = nil; + g_parallel_gate_out = nil; + g_parallel_up_out = nil; + g_parallel_mid_out = nil; + g_parallel_gate_weight_offset = 0; + g_parallel_up_weight_offset = 0; + g_parallel_gate_x_offset = 0; + g_parallel_gate_out_offset = 0; + g_parallel_up_out_offset = 0; + g_parallel_mid_out_offset = 0; + g_parallel_gate_up_args = (ds4_gpu_q8_0_matvec_args){0}; + g_parallel_gate_up_clamp = 0.0f; + g_parallel_gate_up_nsg = 0; + g_parallel_gate_up_nr0 = 0; + g_parallel_gate_up_smem = 0; +} + +void ds4_gpu_parallel_ffn_abort(void) { + ds4_gpu_parallel_ffn_reset_state(YES); +} + +/* Clang cleanup makes every early return from the large generic routed-MoE + * wrapper abort an armed concurrent FFN without touching its many + * established fallback branches. */ +static void ds4_gpu_parallel_ffn_scope_cleanup(BOOL *armed) { + if (armed && *armed && g_parallel_q8_pending) { + ds4_gpu_parallel_ffn_abort(); + } +} + +int ds4_gpu_parallel_ffn_start( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *shared_out, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t model_dim, + uint32_t shared_dim, + const ds4_gpu_tensor *x, + float clamp) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb || g_parallel_q8_pending || g_batch_encoder_concurrent || + !gate || !up || !mid || !shared_out || !x || !model_map || + model_dim == 0 || shared_dim == 0 || + (model_dim & 31u) != 0 || (shared_dim & 31u) != 0 || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(shared_out); + if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < (uint64_t)model_dim * sizeof(float) || + ds4_gpu_tensor_bytes(gate) < (uint64_t)shared_dim * sizeof(float) || + ds4_gpu_tensor_bytes(up) < (uint64_t)shared_dim * sizeof(float) || + ds4_gpu_tensor_bytes(mid) < (uint64_t)shared_dim * sizeof(float) || + ds4_gpu_tensor_bytes(shared_out) < (uint64_t)model_dim * sizeof(float)) { + return 0; + } + + const uint64_t gate_row_bytes = ((uint64_t)model_dim / 32u) * 34u; + const uint64_t gate_weight_bytes = (uint64_t)shared_dim * gate_row_bytes; + const uint64_t down_row_bytes = ((uint64_t)shared_dim / 32u) * 34u; + const uint64_t down_weight_bytes = (uint64_t)model_dim * down_row_bytes; + if (gate_offset > model_size || gate_weight_bytes > model_size - gate_offset || + up_offset > model_size || gate_weight_bytes > model_size - up_offset || + down_offset > model_size || down_weight_bytes > model_size - down_offset) { + return 0; + } + + uint64_t gate_inner = 0, up_inner = 0, down_inner = 0; + id gate_wbuf = ds4_gpu_wrap_model_range( + model_map, model_size, gate_offset, gate_weight_bytes, &gate_inner); + id up_wbuf = ds4_gpu_wrap_model_range( + model_map, model_size, up_offset, gate_weight_bytes, &up_inner); + id down_wbuf = ds4_gpu_wrap_model_range( + model_map, model_size, down_offset, down_weight_bytes, &down_inner); + if (!gate_wbuf || !up_wbuf || !down_wbuf) return 0; + + ds4_gpu_mv_dispatch gate_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + const char *gate_fn = "kernel_dsv4_shared_gate_up_swiglu_q8_0"; + id gate_pipeline = + ds4_gpu_get_mul_mv_pipeline(gate_fn, gate_dispatch.nsg); + id down_pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32", 4); + if (!gate_pipeline || !down_pipeline || + down_pipeline.maxTotalThreadsPerThreadgroup < 128u) { + return 0; + } + + ds4_gpu_close_batch_encoder(); + g_batch_encoder_concurrent = YES; + + g_parallel_gate_up_pipeline = gate_pipeline; + g_parallel_gate_weight = gate_wbuf; + g_parallel_up_weight = up_wbuf; + g_parallel_gate_x = xbuf; + g_parallel_gate_out = gatebuf; + g_parallel_up_out = upbuf; + g_parallel_mid_out = midbuf; + g_parallel_gate_weight_offset = (NSUInteger)gate_inner; + g_parallel_up_weight_offset = (NSUInteger)up_inner; + g_parallel_gate_x_offset = ds4_gpu_tensor_offset(x); + g_parallel_gate_out_offset = ds4_gpu_tensor_offset(gate); + g_parallel_up_out_offset = ds4_gpu_tensor_offset(up); + g_parallel_mid_out_offset = ds4_gpu_tensor_offset(mid); + g_parallel_gate_up_args = + ds4_gpu_make_q8_0_mv_args(model_dim, shared_dim); + g_parallel_gate_up_args.nr0 = gate_dispatch.nr0; + g_parallel_gate_up_clamp = clamp; + g_parallel_gate_up_nsg = gate_dispatch.nsg; + g_parallel_gate_up_nr0 = gate_dispatch.nr0; + g_parallel_gate_up_smem = gate_dispatch.smem; + + g_parallel_q8_pipeline = down_pipeline; + g_parallel_q8_weight = down_wbuf; + g_parallel_q8_x = midbuf; + g_parallel_q8_out = outbuf; + g_parallel_q8_weight_offset = (NSUInteger)down_inner; + g_parallel_q8_x_offset = ds4_gpu_tensor_offset(mid); + g_parallel_q8_out_offset = ds4_gpu_tensor_offset(shared_out); + g_parallel_q8_args = + ds4_gpu_make_q8_0_mv_args(shared_dim, model_dim); + g_parallel_q8_args.nr0 = 2; + + g_parallel_ffn_mode = 2; + g_parallel_ffn_stage = 0; + g_parallel_q8_pending = YES; + g_parallel_q8_encoded = NO; + return 1; +} + +static void ds4_gpu_encode_parallel_q8_down( + id enc) { + [enc setComputePipelineState:g_parallel_q8_pipeline]; + [enc setBytes:&g_parallel_q8_args + length:sizeof(g_parallel_q8_args) + atIndex:0]; + [enc setBuffer:g_parallel_q8_weight + offset:g_parallel_q8_weight_offset + atIndex:1]; + [enc setBuffer:g_parallel_q8_x offset:g_parallel_q8_x_offset atIndex:2]; + [enc setBuffer:g_parallel_q8_out offset:g_parallel_q8_out_offset atIndex:3]; + [enc setThreadgroupMemoryLength:32u * 2u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)g_parallel_q8_args.ne0 + 1u) / 2u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; +} + +static int ds4_gpu_parallel_q8_matvec_encode_pending( + id cb, + id routed_mid) { + if (!g_parallel_q8_pending || g_parallel_ffn_mode != 2 || + g_parallel_ffn_stage != 0 || !g_batch_encoder_concurrent || + !g_batch_cb || cb != g_batch_cb) { + return 0; + } + id enc = ds4_gpu_compute_encoder(cb); + if (!enc || enc.dispatchType != MTLDispatchTypeConcurrent) return 0; + + [enc setComputePipelineState:g_parallel_gate_up_pipeline]; + [enc setBytes:&g_parallel_gate_up_args + length:sizeof(g_parallel_gate_up_args) + atIndex:0]; + [enc setBuffer:g_parallel_gate_weight + offset:g_parallel_gate_weight_offset + atIndex:1]; + [enc setBuffer:g_parallel_up_weight + offset:g_parallel_up_weight_offset + atIndex:2]; + [enc setBuffer:g_parallel_gate_x + offset:g_parallel_gate_x_offset + atIndex:3]; + [enc setBuffer:g_parallel_gate_out + offset:g_parallel_gate_out_offset + atIndex:4]; + [enc setBuffer:g_parallel_up_out + offset:g_parallel_up_out_offset + atIndex:5]; + [enc setBuffer:g_parallel_mid_out + offset:g_parallel_mid_out_offset + atIndex:6]; + [enc setBytes:&g_parallel_gate_up_clamp + length:sizeof(g_parallel_gate_up_clamp) + atIndex:7]; + [enc setThreadgroupMemoryLength:2u * g_parallel_gate_up_smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)g_parallel_gate_up_args.ne0 + + g_parallel_gate_up_nr0 - 1u) / g_parallel_gate_up_nr0, + 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, g_parallel_gate_up_nsg, 1)]; + g_parallel_ffn_stage = 1; + + /* Metal defines either barrier form as an execution barrier for every + * earlier dispatch in this concurrent encoder. The resource list narrows + * visibility, not completion. This is therefore the explicit level break: + * routed pair-SwiGLU and shared gate/up must finish before + * the routed Q2 sum6 consumer can start. */ + id barrier_resources[2] = { + routed_mid, + g_parallel_mid_out, + }; + [enc memoryBarrierWithResources:barrier_resources count:2u]; + return 1; +} + +static int ds4_gpu_parallel_ffn_encode_second_stage( + id cb) { + if (!g_parallel_q8_pending || g_parallel_ffn_mode != 2 || + g_parallel_ffn_stage != 1 || !g_batch_encoder_concurrent || + !g_batch_cb || cb != g_batch_cb) { + return 0; + } + id enc = ds4_gpu_compute_encoder(cb); + if (!enc || enc.dispatchType != MTLDispatchTypeConcurrent) return 0; + ds4_gpu_encode_parallel_q8_down(enc); + /* No final memory barrier is required: finish() ends this concurrent + * encoder before the serial HC encoder is created, and a Metal command + * buffer is an ordered list of command encoders. */ + g_parallel_ffn_stage = 2; + g_parallel_q8_encoded = YES; + return 1; +} + +int ds4_gpu_parallel_ffn_finish(void) { + const int completed = + g_parallel_q8_pending && g_parallel_q8_encoded && + g_parallel_ffn_stage == 2 && g_batch_encoder_concurrent; + /* Reset even on an incomplete join. This makes the error path just as + * safe and idempotent as an explicit abort. */ + ds4_gpu_parallel_ffn_reset_state(YES); + return completed; +} + + static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label) { const char *what = label ? label : "streaming expert cache in-flight"; if (g_batch_cb && ds4_gpu_flush_commands() == 0) return 0; @@ -8217,6 +9355,7 @@ int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { if (!event_value) return 0; *event_value = 0; if (!g_initialized && !ds4_gpu_init()) return 0; + ds4_gpu_parallel_ffn_reset_state(YES); if (!g_batch_cb) return 0; if (@available(macOS 12.0, *)) { @@ -8242,6 +9381,7 @@ int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) { if (!g_initialized && !ds4_gpu_init()) return 0; + ds4_gpu_parallel_ffn_reset_state(YES); if (!g_batch_cb || event_value == 0) return 0; if (@available(macOS 12.0, *)) { @@ -8882,6 +10022,7 @@ int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label static int ds4_gpu_signal_batch_and_wait_event(const char *label) { if (!g_initialized && !ds4_gpu_init()) return 0; + ds4_gpu_parallel_ffn_reset_state(YES); if (!g_batch_cb) return 0; if (@available(macOS 12.0, *)) { @@ -8943,7 +10084,11 @@ static int ds4_gpu_signal_batch_and_wait_event(const char *label) { } int ds4_gpu_end_commands(void) { - if (!g_batch_cb) return 0; + if (!g_batch_cb) { + ds4_gpu_parallel_ffn_reset_state(YES); + return 0; + } + ds4_gpu_parallel_ffn_reset_state(YES); ds4_gpu_close_batch_encoder(); id cb = g_batch_cb; g_batch_cb = nil; @@ -9000,6 +10145,7 @@ static int ds4_gpu_flash_attn_stage_profile_boundary( int ds4_gpu_synchronize(void) { if (!g_initialized && !ds4_gpu_init()) return 0; if (g_batch_cb) return ds4_gpu_end_commands(); + ds4_gpu_parallel_ffn_reset_state(YES); if ([g_pending_cbs count] != 0) { int ok = ds4_gpu_wait_pending_command_buffers("synchronize"); [g_transient_buffers removeAllObjects]; @@ -9016,6 +10162,8 @@ void ds4_gpu_cleanup(void) { if (!g_initialized) return; @autoreleasepool { + ds4_gpu_decode_pipeline_fast_cache_reset(); + ds4_gpu_parallel_ffn_reset_state(YES); if (g_batch_cb) { ds4_gpu_close_batch_encoder(); [g_batch_cb commit]; @@ -9086,16 +10234,19 @@ void ds4_gpu_cleanup(void) { g_add_rms_norm_pipeline = nil; g_rms_norm_scale_pipeline = nil; g_dsv4_qkv_rms_norm_pipeline = nil; + g_dsv4_head_rms_norm_rope_tail_pipeline = nil; + g_use_dsv4_head_rms_norm_rope_tail_pipeline = false; g_hc_split_sinkhorn_pipeline = nil; g_hc_split_weighted_sum_pipeline = nil; g_hc_split_weighted_sum_norm_pipeline = nil; + g_dsv4_hc_producer_pre_norm_pipeline = nil; g_hc_weighted_sum_pipeline = nil; - g_hc_weighted_sum_norm_pipeline = nil; g_output_hc_weights4_pipeline = nil; g_hc_expand_pipeline = nil; g_moe_mul_mv_id_iq2_xxs_pipeline = nil; g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil; g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline = nil; g_moe_mul_mv_id_q2_k_pipeline = nil; g_moe_mul_mv_id_q2_k_sum6_pipeline = nil; g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = nil; @@ -9115,6 +10266,20 @@ void ds4_gpu_cleanup(void) { g_moe_mul_mv_slots6_q2_k_sum6_pipeline = nil; g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = nil; g_moe_mul_mv_slots6_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_id_mxfp4_pipeline = nil; + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline = nil; + g_moe_mul_mv_id_mxfp4_sum6_pipeline = nil; + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple = nil; + g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple = nil; + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 = nil; + g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline = nil; + g_moe_mul_mv_slots6_mxfp4_sum6_pipeline = nil; g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil; g_moe_mul_mv_addr_iq2_xxs_pipeline = nil; g_moe_mul_mv_addr_q2_k_sum6_pipeline = nil; @@ -9136,8 +10301,12 @@ void ds4_gpu_cleanup(void) { g_dsv4_fp8_kv_quantize_pipeline = nil; g_dsv4_indexer_qat_pipeline = nil; g_dsv4_kv_fp8_store_pipeline = nil; + g_dsv4_kv_rope_fp8_store_pipeline = nil; g_dsv4_ratio4_shift_pipeline = nil; g_dsv4_compressor_pack_ratio4_pipeline = nil; + g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline = nil; + g_dsv4_compressor_exact_softmax_product_pipeline = nil; + g_dsv4_compressor_exact_pool_ratio4_pipeline = nil; g_dsv4_softmax_pool_ratio4_direct_pipeline = nil; g_dsv4_softmax_pool_pipeline = nil; g_soft_max_f32_pipeline = nil; @@ -9153,10 +10322,19 @@ void ds4_gpu_cleanup(void) { g_dsv4_sort_i32_rows_asc_pipeline = nil; g_dsv4_indexed_attention_heads8_pipeline = nil; g_dsv4_indexed_attention_heads8_rb16_pipeline = nil; + g_dsv4_indexed_attention_heads16_dual_pipeline = nil; + g_dsv4_indexed_attention_heads8_split_pipeline = nil; + g_dsv4_indexed_attention_heads8_split_reduce_pipeline = nil; g_dsv4_softplus_sqrt_pipeline = nil; g_dsv4_router_finalize_one_pipeline = nil; g_dsv4_router_finalize_one_simd_pipeline = nil; g_dsv4_router_finalize_weights_one_simd_pipeline = nil; + g_dsv4_router_transform_finalize_weights_one_simd_pipeline = nil; + g_dsv4_hc_producer_last_completion = nil; + g_dsv4_hc_producer_last_mix_buffer = nil; + g_dsv4_hc_producer_last_mix_offset = 0; + [g_dsv4_completion_cache removeAllObjects]; + g_dsv4_completion_cache = nil; g_dsv4_router_weights_one_pipeline = nil; g_glm_router_select_one_pipeline = nil; g_glm_kv_lora_rms_norm_pipeline = nil; @@ -10417,70 +11595,6 @@ static id ds4_gpu_wrap_model_exact_range_owned( DS4_GPU_EXACT_VIEW_OWNED); } -static id ds4_gpu_wrap_q8_decode_model_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t n_tokens, - uint64_t *inner_offset) { - const uint64_t exact_decode_max_mib = - ds4_gpu_env_u64("DS4_METAL_Q8_DECODE_EXACT_VIEW_MAX_MIB", - 1024u, - 1u, - 4096u); - const uint64_t exact_decode_max_bytes = - exact_decode_max_mib * 1024ull * 1024ull; - const bool exact_decode_weight_view = - n_tokens == 1u && - len <= exact_decode_max_bytes && - getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; - return exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - offset, - len, - inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - offset, - len, - inner_offset); -} - -static id ds4_gpu_wrap_f32_decode_model_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t n_tokens, - uint64_t *inner_offset) { - const uint64_t exact_decode_max_mib = - ds4_gpu_env_u64("DS4_METAL_F32_DECODE_EXACT_VIEW_MAX_MIB", - 64u, - 1u, - 4096u); - const uint64_t exact_decode_max_bytes = - exact_decode_max_mib * 1024ull * 1024ull; - const bool exact_decode_weight_view = - n_tokens == 1u && - len <= exact_decode_max_bytes && - getenv("DS4_METAL_ENABLE_F32_DECODE_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_F32_DECODE_EXACT_VIEWS") == NULL; - return exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - offset, - len, - inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - offset, - len, - inner_offset); -} - uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { @@ -10536,7 +11650,13 @@ static uint32_t ds4_gpu_stream_expert_cache_requested_budget(void) { } static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void) { - return ds4_gpu_stream_expert_cache_requested_budget(); + uint32_t budget = ds4_gpu_stream_expert_cache_requested_budget(); + if (budget != 0 && + g_stream_expert_cache_mlock_budget_cap != 0 && + budget > g_stream_expert_cache_mlock_budget_cap) { + budget = g_stream_expert_cache_mlock_budget_cap; + } + return budget; } static uint32_t ds4_gpu_stream_expert_cache_effective_cap( @@ -11172,6 +12292,104 @@ static int ds4_gpu_stream_expert_pread_tasks( return ok; } +static void ds4_gpu_stream_expert_cache_warn_mlock_failure( + uint64_t failed_len, + int err) { + if (g_stream_expert_cache_mlock_warned) return; + g_stream_expert_cache_mlock_warned = 1; + + const uint64_t gib = 1024ull * 1024ull * 1024ull; + const uint32_t budget = ds4_gpu_stream_expert_cache_requested_budget(); + uint64_t requested = 0; + if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) { + requested = + budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ? + UINT64_MAX : + (uint64_t)budget * g_stream_expert_cache_expert_bytes; + } + + uint64_t suggested_gib = g_stream_expert_cache_mlock_bytes / gib; + if (suggested_gib > 1) suggested_gib--; + + fprintf(stderr, + "ds4: warning: streaming expert cache could not mlock all buffers\n"); + if (requested != 0) { + fprintf(stderr, + "ds4: requested cache: %u experts / %.2f GiB\n", + budget, + ds4_gpu_gib(requested)); + } else { + fprintf(stderr, + "ds4: requested cache: %u experts\n", + budget); + } + fprintf(stderr, + "ds4: locked so far: %.2f GiB\n", + ds4_gpu_gib(g_stream_expert_cache_mlock_bytes)); + fprintf(stderr, + "ds4: failed buffer: %.2f MiB (%s)\n", + ds4_gpu_mib(failed_len), + err != 0 ? strerror(err) : "mlock unavailable"); + fprintf(stderr, + "ds4: macOS may page unlocked expert buffers, causing poor or unstable speed\n"); + if (g_stream_expert_cache_mlock_budget_cap != 0 && + g_stream_expert_cache_expert_bytes != 0) { + const uint64_t capped_bytes = + g_stream_expert_cache_mlock_budget_cap > + UINT64_MAX / g_stream_expert_cache_expert_bytes ? + UINT64_MAX : + (uint64_t)g_stream_expert_cache_mlock_budget_cap * + g_stream_expert_cache_expert_bytes; + fprintf(stderr, + "ds4: using locked cache cap: %u experts / %.2f GiB\n", + g_stream_expert_cache_mlock_budget_cap, + ds4_gpu_gib(capped_bytes)); + } + if (suggested_gib != 0) { + fprintf(stderr, + "ds4: try: --ssd-streaming-cache-experts %" PRIu64 "GB\n", + suggested_gib); + } else { + fprintf(stderr, + "ds4: try a smaller --ssd-streaming-cache-experts NGB budget\n"); + } +} + +static uint32_t ds4_gpu_stream_expert_slab_locked_slot_count(void) { + uint32_t count = 0; + uint32_t total = g_stream_expert_cache_slab_total_slots; + if (total > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + total = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; + } + for (uint32_t slot = 0; slot < total; slot++) { + if (g_stream_expert_cache_slab_slot_locked[slot]) count++; + } + return count; +} + +static void ds4_gpu_stream_expert_cache_cap_budget_to_locked(void) { + uint32_t cap = g_stream_expert_cache_entry_count; + const uint32_t locked_slots = ds4_gpu_stream_expert_slab_locked_slot_count(); + if (locked_slots != 0 && locked_slots < cap) cap = locked_slots; + const uint64_t gib = 1024ull * 1024ull * 1024ull; + uint64_t safe_gib = g_stream_expert_cache_mlock_bytes / gib; + if (safe_gib > 1) safe_gib--; + if (safe_gib != 0 && g_stream_expert_cache_expert_bytes != 0) { + uint64_t safe_bytes = + safe_gib > UINT64_MAX / gib ? UINT64_MAX : safe_gib * gib; + uint64_t safe_cap64 = safe_bytes / g_stream_expert_cache_expert_bytes; + if (safe_cap64 > UINT32_MAX) safe_cap64 = UINT32_MAX; + if (safe_cap64 != 0 && safe_cap64 < cap) { + cap = (uint32_t)safe_cap64; + } + } + if (cap == 0) return; + if (g_stream_expert_cache_mlock_budget_cap == 0 || + cap < g_stream_expert_cache_mlock_budget_cap) { + g_stream_expert_cache_mlock_budget_cap = cap; + } +} + static id ds4_gpu_stream_expert_alloc_buffer( uint64_t len, NSString *label) { @@ -11191,9 +12409,54 @@ static id ds4_gpu_stream_expert_alloc_buffer( } buffer.label = label; g_stream_expert_cache_buffer_allocs++; + if (g_ssd_streaming_mode) { + void *ptr = [buffer contents]; + const NSUInteger n = [buffer length]; + const double t0 = ds4_gpu_now_ms(); + if (ptr && n != 0 && mlock(ptr, (size_t)n) == 0) { + const double dt = ds4_gpu_now_ms() - t0; + g_stream_expert_cache_mlock_ms += dt; + if (g_stream_expert_cache_mlock_bytes > UINT64_MAX - (uint64_t)n) { + g_stream_expert_cache_mlock_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_mlock_bytes += (uint64_t)n; + } + } else { + const double dt = ds4_gpu_now_ms() - t0; + g_stream_expert_cache_mlock_ms += dt; + g_stream_expert_cache_mlock_failures++; + if (g_stream_expert_cache_mlock_fail_bytes > UINT64_MAX - (uint64_t)n) { + g_stream_expert_cache_mlock_fail_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_mlock_fail_bytes += (uint64_t)n; + } + const int err = ptr && n != 0 ? errno : 0; + ds4_gpu_stream_expert_cache_cap_budget_to_locked(); + ds4_gpu_stream_expert_cache_warn_mlock_failure((uint64_t)n, err); + if (getenv("DS4_METAL_STREAMING_EXPERT_BUFFER_MLOCK_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert buffer mlock failed len=%.2f MiB: %s\n", + ds4_gpu_mib((uint64_t)n), + err != 0 ? strerror(err) : "mlock unavailable"); + } + return nil; + } + } return buffer; } +static void ds4_gpu_stream_expert_unlock_explicit_buffer(id buffer) { + if (!g_ssd_streaming_mode || !buffer) return; + void *ptr = [buffer contents]; + const NSUInteger n = [buffer length]; + if (!ptr || n == 0 || munlock(ptr, (size_t)n) != 0) return; + if (g_stream_expert_cache_mlock_bytes >= (uint64_t)n) { + g_stream_expert_cache_mlock_bytes -= (uint64_t)n; + } else { + g_stream_expert_cache_mlock_bytes = 0; + } +} + static int ds4_gpu_stream_expert_combined_buffer_enabled(void) { return g_ssd_streaming_mode && getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_COMBINED_BUFFER") == NULL; @@ -11296,6 +12559,87 @@ static void ds4_gpu_stream_expert_slab_push_free_slot(uint32_t slot) { slot; } +static int ds4_gpu_stream_expert_slab_lock_slot(uint32_t slot) { + if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES || + g_stream_expert_cache_slab_slot_locked[slot]) { + return 1; + } + uint32_t slab = UINT32_MAX; + uint64_t base = 0; + if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) || + slab >= g_stream_expert_cache_slab_count || + !g_stream_expert_cache_slabs[slab] || + g_stream_expert_cache_slab_slot_bytes == 0 || + base > (uint64_t)NSUIntegerMax) { + return 0; + } + + void *contents = [g_stream_expert_cache_slabs[slab] contents]; + if (!contents) return 0; + const double t0 = ds4_gpu_now_ms(); + void *ptr = (uint8_t *)contents + (NSUInteger)base; + const size_t n = (size_t)g_stream_expert_cache_slab_slot_bytes; + if (mlock(ptr, n) == 0) { + const double dt = ds4_gpu_now_ms() - t0; + g_stream_expert_cache_mlock_ms += dt; + if (g_stream_expert_cache_mlock_bytes > + UINT64_MAX - g_stream_expert_cache_slab_slot_bytes) { + g_stream_expert_cache_mlock_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_mlock_bytes += + g_stream_expert_cache_slab_slot_bytes; + } + g_stream_expert_cache_slab_slot_locked[slot] = 1; + return 1; + } + + const double dt = ds4_gpu_now_ms() - t0; + g_stream_expert_cache_mlock_ms += dt; + g_stream_expert_cache_mlock_failures++; + if (g_stream_expert_cache_mlock_fail_bytes > + UINT64_MAX - g_stream_expert_cache_slab_slot_bytes) { + g_stream_expert_cache_mlock_fail_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_mlock_fail_bytes += + g_stream_expert_cache_slab_slot_bytes; + } + ds4_gpu_stream_expert_cache_cap_budget_to_locked(); + ds4_gpu_stream_expert_cache_warn_mlock_failure( + g_stream_expert_cache_slab_slot_bytes, + errno); + return 0; +} + +static int ds4_gpu_stream_expert_slab_unlock_slot(uint32_t slot) { + if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES || + !g_stream_expert_cache_slab_slot_locked[slot]) { + return 1; + } + uint32_t slab = UINT32_MAX; + uint64_t base = 0; + if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) || + slab >= g_stream_expert_cache_slab_count || + !g_stream_expert_cache_slabs[slab] || + g_stream_expert_cache_slab_slot_bytes == 0 || + base > (uint64_t)NSUIntegerMax) { + return 0; + } + + void *contents = [g_stream_expert_cache_slabs[slab] contents]; + if (!contents) return 0; + void *ptr = (uint8_t *)contents + (NSUInteger)base; + const size_t n = (size_t)g_stream_expert_cache_slab_slot_bytes; + if (munlock(ptr, n) != 0) return 0; + + g_stream_expert_cache_slab_slot_locked[slot] = 0; + if (g_stream_expert_cache_mlock_bytes >= g_stream_expert_cache_slab_slot_bytes) { + g_stream_expert_cache_mlock_bytes -= g_stream_expert_cache_slab_slot_bytes; + } else { + g_stream_expert_cache_mlock_bytes = 0; + } + return 1; +} + static int ds4_gpu_stream_expert_slab_slot_buffers( uint32_t slot, uint64_t gate_expert_bytes, @@ -11320,6 +12664,7 @@ static int ds4_gpu_stream_expert_slab_slot_buffers( return 0; } id b = g_stream_expert_cache_slabs[slab]; + if (!ds4_gpu_stream_expert_slab_lock_slot(slot)) return 0; *gate_buf = b; *up_buf = b; *down_buf = b; @@ -12382,6 +13727,15 @@ static void ds4_gpu_stream_expert_cache_clear_entry_internal( reuse->down_inner = e->down_inner; } else if (e->slab_backed && recycle_slab_slot) { ds4_gpu_stream_expert_slab_push_free_slot(e->slab_slot); + } else if (!e->slab_backed) { + ds4_gpu_stream_expert_unlock_explicit_buffer(e->gate_buffer); + if (e->up_buffer != e->gate_buffer) { + ds4_gpu_stream_expert_unlock_explicit_buffer(e->up_buffer); + } + if (e->down_buffer != e->gate_buffer && + e->down_buffer != e->up_buffer) { + ds4_gpu_stream_expert_unlock_explicit_buffer(e->down_buffer); + } } e->gate_buffer = nil; e->up_buffer = nil; @@ -12470,6 +13824,9 @@ static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) { g_stream_expert_cache_slab_total_slots = 0; g_stream_expert_cache_free_slot_count = 0; g_stream_expert_cache_slab_slot_bytes = 0; + memset(g_stream_expert_cache_slab_slot_locked, + 0, + sizeof(g_stream_expert_cache_slab_slot_locked)); if (reset_stats) { g_stream_expert_cache_hits = 0; g_stream_expert_cache_misses = 0; @@ -12480,6 +13837,13 @@ static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) { g_stream_expert_cache_willneed_advise_bytes = 0; g_stream_expert_cache_pread_bytes = 0; g_stream_expert_cache_pread_ms = 0.0; + g_stream_expert_cache_mlock_bytes = 0; + g_stream_expert_cache_mlock_fail_bytes = 0; + g_stream_expert_cache_mlock_failures = 0; + g_stream_expert_cache_mlock_ms = 0.0; + g_stream_expert_cache_mlock_warned = 0; + g_stream_expert_cache_mlock_budget_cap = 0; + g_stream_expert_cache_mlock_relief_applied = 0; g_stream_expert_cache_buffer_allocs = 0; g_stream_expert_cache_buffer_reuses = 0; g_stream_expert_cache_decode_tokens = 0; @@ -12915,6 +14279,140 @@ static int ds4_gpu_stream_expert_batch_reuse_enabled( return slot_bytes <= 16ull * 1024ull * 1024ull; } +static uint32_t ds4_gpu_stream_expert_cache_release_mlock_margin( + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect) { + if (g_stream_expert_cache_mlock_relief_applied || + g_stream_expert_cache_slab_slot_bytes == 0) { + return 0; + } + + const uint32_t locked_before = + ds4_gpu_stream_expert_slab_locked_slot_count(); + if (locked_before == 0) return 0; + + const uint32_t target_release = (locked_before + 9u) / 10u; + uint32_t released = 0; + while (released < target_release) { + uint32_t victim_layer = UINT32_MAX; + uint32_t victim_expert = UINT32_MAX; + uint32_t victim_slot = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid || + !e->slab_backed || + e->slab_slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES || + !g_stream_expert_cache_slab_slot_locked[e->slab_slot] || + ds4_gpu_stream_expert_cache_entry_inflight(e) || + ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim_layer = layer; + victim_expert = expert; + victim_slot = e->slab_slot; + } + } + } + + if (victim_layer == UINT32_MAX || + victim_expert == UINT32_MAX || + victim_slot == UINT32_MAX) { + break; + } + + ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer, + victim_expert, + 1, + 0, + NULL); + if (g_stream_expert_cache[victim_layer][victim_expert].valid) { + break; + } + if (ds4_gpu_stream_expert_slab_unlock_slot(victim_slot)) { + released++; + } + } + + if (released == 0) return 0; + g_stream_expert_cache_mlock_relief_applied = 1; + + uint32_t cap = g_stream_expert_cache_entry_count; + const uint32_t locked_after = + ds4_gpu_stream_expert_slab_locked_slot_count(); + if (locked_after != 0 && locked_after < cap) cap = locked_after; + if (cap != 0 && + (g_stream_expert_cache_mlock_budget_cap == 0 || + cap < g_stream_expert_cache_mlock_budget_cap)) { + g_stream_expert_cache_mlock_budget_cap = cap; + } + + const uint64_t released_bytes = + (uint64_t)released * g_stream_expert_cache_slab_slot_bytes; + fprintf(stderr, + "ds4: released locked cache margin: %u experts / %.2f GiB; " + "runtime cache cap now %u experts\n", + released, + ds4_gpu_gib(released_bytes), + g_stream_expert_cache_mlock_budget_cap); + return released; +} + +static int ds4_gpu_stream_expert_cache_take_capped_reusable( + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + if (g_stream_expert_cache_mlock_budget_cap == 0) return 0; + + ds4_gpu_stream_expert_cache_release_mlock_margin(protect_layer, + protect_ids, + n_protect); + ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 }; + if (!ds4_gpu_stream_expert_cache_take_reusable(1, + protect_layer, + protect_ids, + n_protect, + gate_expert_bytes, + down_expert_bytes, + &reuse)) { + return 0; + } + *gate_buf = reuse.gate_buffer; + *up_buf = reuse.up_buffer; + *down_buf = reuse.down_buffer; + *gate_inner = reuse.gate_inner; + *up_inner = reuse.up_inner; + *down_inner = reuse.down_inner; + return *gate_buf && *up_buf && *down_buf; +} + static int ds4_gpu_stream_expert_cache_prepare_load_buffers( uint32_t layer, uint32_t expert, @@ -13006,6 +14504,13 @@ static int ds4_gpu_stream_expert_cache_prepare_load_buffers( (uint64_t)NSUIntegerMax) { return 0; } + if (g_stream_expert_cache_mlock_budget_cap != 0) { + return ds4_gpu_stream_expert_cache_take_capped_reusable( + protect_layer, protect_ids, n_protect, + gate_expert_bytes, down_expert_bytes, + gate_buf, up_buf, down_buf, + gate_inner, up_inner, down_inner); + } if (ds4_gpu_stream_expert_alloc_slab_slot(gate_expert_bytes, down_expert_bytes, gate_buf, @@ -13016,13 +14521,26 @@ static int ds4_gpu_stream_expert_cache_prepare_load_buffers( down_inner)) { return 1; } + if (g_stream_expert_cache_mlock_budget_cap != 0) { + return ds4_gpu_stream_expert_cache_take_capped_reusable( + protect_layer, protect_ids, n_protect, + gate_expert_bytes, down_expert_bytes, + gate_buf, up_buf, down_buf, + gate_inner, up_inner, down_inner); + } const uint64_t up_off = gate_expert_bytes; const uint64_t down_off = gate_expert_bytes * 2ull; const uint64_t combined_bytes = down_off + down_expert_bytes; id combined = ds4_gpu_stream_expert_alloc_buffer(combined_bytes, @"ds4_stream_expert_combined"); - if (!combined) return 0; + if (!combined) { + return ds4_gpu_stream_expert_cache_take_capped_reusable( + protect_layer, protect_ids, n_protect, + gate_expert_bytes, down_expert_bytes, + gate_buf, up_buf, down_buf, + gate_inner, up_inner, down_inner); + } *gate_buf = combined; *up_buf = combined; *down_buf = combined; @@ -13032,13 +14550,29 @@ static int ds4_gpu_stream_expert_cache_prepare_load_buffers( return 1; } - *gate_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, - @"ds4_stream_expert_gate"); - *up_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, - @"ds4_stream_expert_up"); - *down_buf = ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes, - @"ds4_stream_expert_down"); - return *gate_buf && *up_buf && *down_buf; + id explicit_gate = + ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, + @"ds4_stream_expert_gate"); + id explicit_up = + ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, + @"ds4_stream_expert_up"); + id explicit_down = + ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes, + @"ds4_stream_expert_down"); + if (!explicit_gate || !explicit_up || !explicit_down) { + ds4_gpu_stream_expert_unlock_explicit_buffer(explicit_gate); + ds4_gpu_stream_expert_unlock_explicit_buffer(explicit_up); + ds4_gpu_stream_expert_unlock_explicit_buffer(explicit_down); + return ds4_gpu_stream_expert_cache_take_capped_reusable( + protect_layer, protect_ids, n_protect, + gate_expert_bytes, down_expert_bytes, + gate_buf, up_buf, down_buf, + gate_inner, up_inner, down_inner); + } + *gate_buf = explicit_gate; + *up_buf = explicit_up; + *down_buf = explicit_down; + return 1; } static void ds4_gpu_stream_expert_cache_prune_global( @@ -13879,7 +15413,7 @@ int ds4_gpu_stream_expert_cache_begin_selected_load( return 1; } -static int ds4_gpu_stream_expert_cache_load_selected_missing( +static int ds4_gpu_stream_expert_cache_load_selected_missing_with_source( const void *model_map, uint64_t model_size, uint32_t layer, @@ -13892,6 +15426,7 @@ static int ds4_gpu_stream_expert_cache_load_selected_missing( uint64_t gate_expert_bytes, uint64_t down_expert_bytes, uint32_t missing_mask, + int gpu_copy_source, ds4_gpu_stream_expert_cache_entry **entries) { if (!g_ssd_streaming_mode || !model_map || @@ -14048,12 +15583,14 @@ static int ds4_gpu_stream_expert_cache_load_selected_missing( const int force_reuse = cache_budget != 0 && reserved_entries >= cache_budget; - ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot], - down_expert_bytes); + if (!gpu_copy_source) { + ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot], + down_expert_bytes); + } if (load_i < batch_reuse_count && batch_reuse[load_i].gate_buffer && @@ -14104,7 +15641,17 @@ static int ds4_gpu_stream_expert_cache_load_selected_missing( uint8_t *down_dst = (uint8_t *)[down_bufs[load_i] contents] + down_inners[load_i]; if (!gate_dst || !up_dst || !down_dst) return 0; - + if (gpu_copy_source && + (gate_abs_offsets[slot] > model_size || + gate_expert_bytes > model_size - gate_abs_offsets[slot] || + up_abs_offsets[slot] > model_size || + gate_expert_bytes > model_size - up_abs_offsets[slot] || + down_abs_offsets[slot] > model_size || + down_expert_bytes > model_size - down_abs_offsets[slot])) { + fprintf(stderr, + "ds4: Metal streaming mapped expert source is outside the model mapping\n"); + return 0; + } const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { .offset = gate_abs_offsets[slot], @@ -14136,27 +15683,113 @@ static int ds4_gpu_stream_expert_cache_load_selected_missing( uint64_t read_bytes = 0; double read_ms = 0.0; - const int ok = ds4_gpu_stream_expert_pread_tasks(tasks, - n_tasks, - &read_bytes, - &read_ms); + int ok = 1; + if (gpu_copy_source) { + if (!g_batch_cb || + gate_expert_bytes > (uint64_t)NSUIntegerMax || + down_expert_bytes > (uint64_t)NSUIntegerMax) { + return 0; + } + __strong id gate_srcs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id up_srcs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id down_srcs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger gate_src_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger up_src_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger down_src_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t slot = load_slots[load_i]; + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + gate_srcs[load_i] = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_abs_offsets[slot], + gate_expert_bytes, + &gate_inner); + up_srcs[load_i] = ds4_gpu_wrap_model_range(model_map, + model_size, + up_abs_offsets[slot], + gate_expert_bytes, + &up_inner); + down_srcs[load_i] = ds4_gpu_wrap_model_range(model_map, + model_size, + down_abs_offsets[slot], + down_expert_bytes, + &down_inner); + if (!gate_srcs[load_i] || !up_srcs[load_i] || !down_srcs[load_i] || + gate_inner > (uint64_t)NSUIntegerMax || + up_inner > (uint64_t)NSUIntegerMax || + down_inner > (uint64_t)NSUIntegerMax) { + return 0; + } + gate_src_inners[load_i] = (NSUInteger)gate_inner; + up_src_inners[load_i] = (NSUInteger)up_inner; + down_src_inners[load_i] = (NSUInteger)down_inner; + } + + ds4_gpu_close_batch_encoder(); + id blit = [g_batch_cb blitCommandEncoder]; + if (!blit) return 0; + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + [blit copyFromBuffer:gate_srcs[load_i] + sourceOffset:gate_src_inners[load_i] + toBuffer:gate_bufs[load_i] + destinationOffset:gate_inners[load_i] + size:(NSUInteger)gate_expert_bytes]; + [blit copyFromBuffer:up_srcs[load_i] + sourceOffset:up_src_inners[load_i] + toBuffer:up_bufs[load_i] + destinationOffset:up_inners[load_i] + size:(NSUInteger)gate_expert_bytes]; + [blit copyFromBuffer:down_srcs[load_i] + sourceOffset:down_src_inners[load_i] + toBuffer:down_bufs[load_i] + destinationOffset:down_inners[load_i] + size:(NSUInteger)down_expert_bytes]; + const uint64_t bytes = + gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2u ? + UINT64_MAX : gate_expert_bytes * 2u + down_expert_bytes; + read_bytes = read_bytes > UINT64_MAX - bytes ? + UINT64_MAX : read_bytes + bytes; + } + [blit endEncoding]; + g_batch_has_work = YES; + } else { + ok = ds4_gpu_stream_expert_pread_tasks(tasks, + n_tasks, + &read_bytes, + &read_ms); + } if (!ok) return 0; if (load_timing) { load_t0 = ds4_gpu_now_ms(); } - ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + if (!gpu_copy_source) { + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + } - for (uint32_t load_i = 0; load_i < n_loads; load_i++) { - [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; - [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; - [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; + if (!gpu_copy_source) { + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; + } } if (load_timing) { const double now_ms = ds4_gpu_now_ms(); load_modify_ms = now_ms - load_t0; load_t0 = now_ms; } - if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + if (gpu_copy_source && + getenv("DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert GPU cache copy layer=%u " + "experts=%u bytes=%.2f GiB\n", + layer, + n_loads, + ds4_gpu_gib(read_bytes)); + } else if (!gpu_copy_source && + getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { fprintf(stderr, "ds4: Metal streaming expert parallel pread layer=%u experts=%u tensors=%u " "threads=%u bytes=%.2f GiB wall=%.3f ms\n", @@ -14201,7 +15834,7 @@ static int ds4_gpu_stream_expert_cache_load_selected_missing( for (uint32_t i = 0; i < n_selected; i++) { if ((missing_mask & (1u << i)) != 0 && !entries[i]) return 0; } - if (load_timing) { + if (load_timing && !gpu_copy_source) { load_install_ms = ds4_gpu_now_ms() - load_t0; ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, read_ms, @@ -14211,6 +15844,37 @@ static int ds4_gpu_stream_expert_cache_load_selected_missing( return 1; } +static int ds4_gpu_stream_expert_cache_load_selected_missing( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_total_expert, + uint32_t n_selected, + const uint64_t *gate_abs_offsets, + const uint64_t *up_abs_offsets, + const uint64_t *down_abs_offsets, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + uint32_t missing_mask, + ds4_gpu_stream_expert_cache_entry **entries) { + return ds4_gpu_stream_expert_cache_load_selected_missing_with_source( + model_map, + model_size, + layer, + selected_ids, + n_total_expert, + n_selected, + gate_abs_offsets, + up_abs_offsets, + down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + missing_mask, + 0, + entries); +} + static void ds4_gpu_glm_stream_selected_prefetch_set( const ds4_gpu_stream_expert_table *table, const int32_t *selected_ids, @@ -14905,11 +16569,12 @@ int ds4_gpu_stream_expert_cache_seed_selected( return 1; } -int ds4_gpu_stream_expert_cache_seed_experts( +static int ds4_gpu_stream_expert_cache_seed_experts_impl( const ds4_gpu_stream_expert_table *table, const int32_t *expert_ids, const uint32_t *expert_priorities, - uint32_t n_experts) { + uint32_t n_experts, + int gpu_copy_source) { if (!g_ssd_streaming_mode) return 1; if (!table) return 0; const void *model_map = table->model_map; @@ -15004,19 +16669,21 @@ int ds4_gpu_stream_expert_cache_seed_experts( } if (missing_mask != 0 && - !ds4_gpu_stream_expert_cache_load_selected_missing(model_map, - model_size, - layer, - selected_ids, - n_total_expert, - batch, - gate_abs_offsets, - up_abs_offsets, - down_abs_offsets, - gate_expert_bytes, - down_expert_bytes, - missing_mask, - entries)) { + !ds4_gpu_stream_expert_cache_load_selected_missing_with_source( + model_map, + model_size, + layer, + selected_ids, + n_total_expert, + batch, + gate_abs_offsets, + up_abs_offsets, + down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + missing_mask, + gpu_copy_source, + entries)) { return 0; } if (expert_priorities) { @@ -15042,6 +16709,29 @@ int ds4_gpu_stream_expert_cache_seed_experts( return 1; } +int ds4_gpu_stream_expert_cache_seed_experts( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts) { + return ds4_gpu_stream_expert_cache_seed_experts_impl(table, + expert_ids, + expert_priorities, + n_experts, + 0); +} + +int ds4_gpu_stream_expert_cache_seed_experts_gpu_copy( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts) { + return ds4_gpu_stream_expert_cache_seed_experts_impl(table, + expert_ids, + expert_priorities, + n_experts, + 1); +} static uint32_t ds4_gpu_q4_expert_table_group_size(uint32_t n_total_expert) { const char *env = getenv("DS4_METAL_Q4_EXPERT_TABLE_GROUP_SIZE"); if (!env || !env[0]) return 1; @@ -16217,8 +17907,7 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok, - bool prefer_decode_mpp, - bool force_model_view) { + bool prefer_decode_mpp) { if (!g_initialized && !ds4_gpu_init()) return 0; if ((in_dim & 31u) != 0 || in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { @@ -16246,18 +17935,11 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( } uint64_t inner_offset = 0; - id wbuf = force_model_view ? - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset) : - ds4_gpu_wrap_q8_decode_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - n_tok, - &inner_offset); + id wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); if (!wbuf) { return 0; } @@ -16268,8 +17950,7 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( if (n_tok == 1) { if (ds4_gpu_mpp_available() && - (prefer_decode_mpp || getenv("DS4_METAL_Q8_DECODE_MPP") != NULL) && - getenv("DS4_METAL_DISABLE_Q8_DECODE_MPP") == NULL && + prefer_decode_mpp && (in_dim % 64u) == 0) { const char *nax_fn = "kernel_mul_mm_q8_0_f32_nax_direct_rhs"; id mpp_pipeline = @@ -16302,21 +17983,6 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); if (out_dim > 65536u) mv_dispatch.nsg = 8; - const bool force_output_nr4 = - getenv("DS4_METAL_ENABLE_OUTPUT_Q8_NR4") != NULL; - const bool output_shape = - in_dim == 4096u && out_dim == 129280u; - const bool use_output_nr4 = - !g_quality_mode && (out_dim & 3u) == 0u && - (force_output_nr4 || - (output_shape && ds4_gpu_device_name_contains("M3"))) && - getenv("DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4") == NULL; - if (use_output_nr4) { - mv_dispatch.function_name = - "kernel_mul_mv_q8_0_f32_nr4"; - mv_dispatch.nr0 = 4; - mv_dispatch.smem = 32u * 4u * sizeof(float); - } mv_args.nr0 = mv_dispatch.nr0; id pipeline = ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); @@ -16383,15 +18049,28 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( * staging RHS into threadgroup memory and was the direct replacement for * the slower generic MPP prototype. */ + /* + * An unaligned token count used to send the entire projection through + * the generic kernel. Run its aligned prefix through TensorOps and + * leave only the final partial tile to the boundary-safe kernel. + * Tiny prompts do not amortize the second dispatch, while --quality + * deliberately retains the single-kernel arithmetic schedule. + */ + const bool split_nax_prefix = + !g_quality_mode && n_tok >= 192u && (n_tok % 32u) != 0u; + const uint64_t nax_rows = + (n_tok % 32u) == 0u ? n_tok : + (split_nax_prefix ? n_tok - (n_tok % 32u) : 0u); + uint64_t generic_row0 = 0u; + uint64_t generic_rows = n_tok; if (ds4_gpu_mpp_available() && - n_tok >= 32u && + nax_rows >= 32u && (in_dim % 64u) == 0 && - (out_dim % 64u) == 0 && - (n_tok % 32u) == 0) { + (out_dim % 64u) == 0) { uint64_t nax_tile_n = 32u; - if ((n_tok % 128u) == 0) { + if ((nax_rows % 128u) == 0) { nax_tile_n = 128u; - } else if ((n_tok % 64u) == 0) { + } else if ((nax_rows % 64u) == 0) { nax_tile_n = 64u; } const char *nax_fn = nax_tile_n == 128u @@ -16402,7 +18081,8 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( id pipeline = ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); if (pipeline) { - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + ds4_gpu_mul_mm_args args = + ds4_gpu_make_mm_args(in_dim, out_dim, nax_rows, row_bytes); id enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:pipeline]; @@ -16411,36 +18091,48 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(nax_rows / nax_tile_n), (NSUInteger)out_dim / 64u, 1) threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) { - return 0; + if (nax_rows == n_tok) { + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) { + return 0; + } + return 1; } - return 1; + generic_row0 = nax_rows; + generic_rows = n_tok - nax_rows; } - ds4_gpu_warn_mpp_fallback(); + if (!pipeline) ds4_gpu_warn_mpp_fallback(); } const bool bc_inp = (in_dim % 32u) != 0; - const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + const bool bc_out = + (out_dim % 64u) != 0 || (generic_rows % 32u) != 0; id pipeline = ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_q8_0_f32", bc_inp, bc_out); if (!pipeline) return 0; - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + ds4_gpu_mul_mm_args args = + ds4_gpu_make_mm_args(in_dim, out_dim, generic_rows, row_bytes); id enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:pipeline]; [enc setBytes:&args length:sizeof(args) atIndex:0]; [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:xbuf + offset:ds4_gpu_tensor_offset(x) + + (NSUInteger)(generic_row0 * in_dim * sizeof(float)) + atIndex:2]; + [enc setBuffer:outbuf + offset:ds4_gpu_tensor_offset(out) + + (NSUInteger)(generic_row0 * out_dim * sizeof(float)) + atIndex:3]; [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)generic_rows + 31u) / 32u, ((NSUInteger)out_dim + 63u) / 64u, 1) threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; @@ -16503,7 +18195,7 @@ int ds4_gpu_matmul_q8_0_tensor( const double profile_t0 = profile_prefill ? ds4_gpu_now_ms() : 0.0; int ok = ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, weight_offset, in_dim, out_dim, - x, n_tok, false, false); + x, n_tok, false); if (profile_prefill) { if (split_batch_for_profile && ds4_gpu_end_commands() == 0) { ok = 0; @@ -16562,8 +18254,8 @@ int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( } uint64_t inner_offset = 0; - id wbuf = ds4_gpu_wrap_q8_decode_model_range( - model_map, model_size, weight_offset, weight_bytes, 1u, + id wbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_offset, weight_bytes, &inner_offset); if (!wbuf) return 0; @@ -16617,7 +18309,7 @@ int ds4_gpu_matmul_q8_0_decode_mpp_tensor( uint64_t n_tok) { return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, weight_offset, in_dim, out_dim, - x, n_tok, true, false); + x, n_tok, true); } int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor( @@ -16631,7 +18323,7 @@ int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor( uint64_t n_tok) { return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, weight_offset, in_dim, out_dim, - x, n_tok, true, true); + x, n_tok, true); } static const char *ds4_gpu_q4_mv_ext_name(uint32_t weight_type, int16_t r1ptg) { @@ -16692,8 +18384,7 @@ static int ds4_gpu_matmul_quant_impl_tensor( uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok, - bool prefer_decode_mpp, - bool force_model_view) { + bool prefer_decode_mpp) { if (weight_type == DS4_METAL_TENSOR_Q8_0) { return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, @@ -16703,8 +18394,7 @@ static int ds4_gpu_matmul_quant_impl_tensor( out_dim, x, n_tok, - prefer_decode_mpp, - force_model_view); + prefer_decode_mpp); } if (!g_initialized && !ds4_gpu_init()) return 0; if (!out || !x || !model_map || @@ -16741,17 +18431,11 @@ static int ds4_gpu_matmul_quant_impl_tensor( } uint64_t inner_offset = 0; - id wbuf = force_model_view ? - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); + id wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); if (!wbuf) return 0; int owned = 0; @@ -16926,7 +18610,6 @@ int ds4_gpu_matmul_quant_tensor( out_dim, x, n_tok, - false, false); } @@ -16940,11 +18623,7 @@ int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor( uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { - /* GLM decode dense matvecs: the classic Q8_0 kernels stream 570-613 - * GB/s on the GLM shapes while the MPP/nax matrix path measures ~150 - * GB/s at n_tok=1 (TP decode 8.75 -> 16.64 t/s on the IQ2+Q8 gguf). - * MPP stays available as an opt-in via DS4_METAL_Q8_DECODE_MPP. */ - const bool prefer_mpp = getenv("DS4_METAL_Q8_DECODE_MPP") != NULL; + /* Classic Q8_0 matvec is substantially faster than MPP at n_tok=1. */ return ds4_gpu_matmul_quant_impl_tensor(out, model_map, model_size, @@ -16954,8 +18633,7 @@ int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor( out_dim, x, n_tok, - prefer_mpp, - true); + false); } int ds4_gpu_matmul_quant_rows_scalar_tensor( @@ -17205,8 +18883,7 @@ static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl( uint64_t out_dim, const ds4_gpu_tensor *x, float clamp, - int store_gate_up, - bool force_model_view) { + int store_gate_up) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!mid || !x || !model_map || (store_gate_up && (!gate || !up)) || @@ -17243,46 +18920,26 @@ static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl( return 0; } - const bool exact_decode_views = - !force_model_view && - getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; uint64_t gate_inner = 0; uint64_t up_inner = 0; - id gate_wbuf = exact_decode_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - gate_offset, - weight_bytes, - &gate_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset, - weight_bytes, - &gate_inner); - id up_wbuf = exact_decode_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - up_offset, - weight_bytes, - &up_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset, - weight_bytes, - &up_inner); + 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(in_dim, out_dim); ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); args.nr0 = mv_dispatch.nr0; const char *fn_name = store_gate_up ? - (mv_dispatch.nr0 >= 4 ? - "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : - "kernel_dsv4_shared_gate_up_swiglu_q8_0") : - (mv_dispatch.nr0 >= 4 ? - "kernel_dsv4_shared_mid_swiglu_q8_0_r4" : - "kernel_dsv4_shared_mid_swiglu_q8_0"); + "kernel_dsv4_shared_gate_up_swiglu_q8_0" : + "kernel_dsv4_shared_mid_swiglu_q8_0"; id pipeline = ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); if (!pipeline) return 0; @@ -17325,6 +18982,217 @@ static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl( return 1; } +/* Decode-only fusion of the router logits matvec with the shared-expert + * gate/up SwiGLU: one dispatch instead of two on the same normalized FFN + * input. Bit-exact by construction (see metal/dense.metal). Returns 1 on + * success, 0 when the shape is unsupported (caller falls back), -1 on + * error. */ +int ds4_gpu_router_shared_gate_up_q8_0_tensor( + ds4_gpu_tensor *router_logits, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t router_weight_offset, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t router_out_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp, + bool router_only) { + if (!g_initialized && !ds4_gpu_init()) return -1; + if (!router_logits || !gate || !up || !mid || !model_map || !x || + (in_dim & 31u) != 0 || in_dim > UINT32_MAX || + router_out_dim > UINT32_MAX || out_dim > UINT32_MAX || + (router_out_dim & 1u) != 0 || (out_dim & 3u) != 0 || + !isfinite(clamp) || clamp < 0.0f) { + return -1; + } + if (in_dim != 4096u) return 0; + + @autoreleasepool { + const uint64_t router_row_bytes = in_dim * sizeof(uint16_t); + const uint64_t router_weight_bytes = router_row_bytes * router_out_dim; + const uint64_t q8_row_bytes = (in_dim / 32u) * 34u; + const uint64_t q8_weight_bytes = q8_row_bytes * out_dim; + if (router_weight_offset > model_size || + router_weight_bytes > model_size - router_weight_offset || + gate_offset > model_size || q8_weight_bytes > model_size - gate_offset || + up_offset > model_size || q8_weight_bytes > model_size - up_offset) { + return -1; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id rbuf = ds4_gpu_tensor_buffer(router_logits); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + if (!xbuf || !rbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(router_logits) < router_out_dim * sizeof(float) || + ds4_gpu_tensor_bytes(gate) < out_dim * sizeof(float) || + ds4_gpu_tensor_bytes(up) < out_dim * sizeof(float) || + ds4_gpu_tensor_bytes(mid) < out_dim * sizeof(float)) { + return -1; + } + + uint64_t router_inner = 0, gate_inner = 0, up_inner = 0; + id router_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, router_weight_offset, router_weight_bytes, &router_inner); + id gate_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, q8_weight_bytes, &gate_inner); + id up_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, q8_weight_bytes, &up_inner); + if (!router_wbuf || !gate_wbuf || !up_wbuf) return -1; + + ds4_gpu_f16_matvec_args rargs = ds4_gpu_make_f16_mv_args(in_dim, router_out_dim); + rargs.nr0 = 2; + ds4_gpu_q8_0_matvec_args sargs = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + sargs.nr0 = 2; + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_shared_gate_up_q8_0"); + if (!pipeline) return -1; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return -1; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&rargs length:sizeof(rargs) atIndex:0]; + [enc setBytes:&sargs length:sizeof(sargs) atIndex:1]; + [enc setBuffer:router_wbuf offset:(NSUInteger)router_inner atIndex:2]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:3]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:4]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:5]; + [enc setBuffer:rbuf offset:ds4_gpu_tensor_offset(router_logits) atIndex:6]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:7]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:8]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:9]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:10]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0]; + const NSUInteger router_tgs = ((NSUInteger)router_out_dim + 1u) / 2u; + const NSUInteger shared_tgs = ((NSUInteger)out_dim / 2u + 1u) / 2u; + [enc dispatchThreadgroups:MTLSizeMake(router_tgs + + (router_only ? 0u : shared_tgs), + 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "router + shared gate/up fused")) { + return -1; + } + } + + return 1; +} + +int ds4_gpu_router_project_select_fused_tensor( + ds4_gpu_tensor *router_logits, + ds4_gpu_tensor *probs, + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + const void *model_map, + uint64_t model_size, + uint64_t router_weight_offset, + uint64_t bias_offset, + bool has_bias, + const ds4_gpu_tensor *x) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!router_logits || !probs || !selected || !weights || !model_map || !x) { + return 0; + } + @autoreleasepool { + const uint32_t in_dim = 4096u; + const uint32_t n_expert = 256u; + const uint32_t n_used = 6u; + const uint64_t row_bytes = (uint64_t)in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * n_expert; + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (router_weight_offset > model_size || + weight_bytes > model_size - router_weight_offset || + (has_bias && (bias_offset > model_size || + bias_bytes > model_size - bias_offset))) { + return 0; + } + id xbuf = ds4_gpu_tensor_buffer(x); + id logitsbuf = ds4_gpu_tensor_buffer(router_logits); + id probsbuf = ds4_gpu_tensor_buffer(probs); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + if (!xbuf || !logitsbuf || !probsbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(router_logits) < bias_bytes || + ds4_gpu_tensor_bytes(probs) < bias_bytes || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_used * sizeof(int32_t) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_used * sizeof(float)) { + return 0; + } + uint64_t router_inner = 0; + uint64_t bias_inner = 0; + id routerbuf = ds4_gpu_wrap_model_range( + model_map, model_size, router_weight_offset, + weight_bytes, &router_inner); + id biasbuf = has_bias ? ds4_gpu_wrap_model_range( + model_map, model_size, bias_offset, bias_bytes, &bias_inner) : nil; + if (!routerbuf || (has_bias && !biasbuf)) return 0; + id pipeline = ds4_gpu_get_pipeline( + "kernel_dsv4_router_project_select_fused"); + if (!pipeline || pipeline.maxTotalThreadsPerThreadgroup < 256u) return 0; + if (!g_dsv4_completion_cache) return 0; + NSString *completion_key = [NSString stringWithFormat:@"router:%p:%llu", + (void *)logitsbuf, + (unsigned long long)ds4_gpu_tensor_offset(router_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]; + ds4_gpu_f16_matvec_args rargs = + ds4_gpu_make_f16_mv_args(in_dim, n_expert); + rargs.nr0 = 2; + ds4_gpu_dsv4_router_select_one_args select_args = { + .has_bias = has_bias ? 1u : 0u, + .hash_mode = 0u, + .use_token_buffer = 0u, + .token = 0u, + .hash_rows = 0u, + }; + const float zero_f32 = 0.0f; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&rargs length:sizeof(rargs) atIndex:0]; + [enc setBytes:&select_args length:sizeof(select_args) atIndex:1]; + [enc setBuffer:routerbuf offset:(NSUInteger)router_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(router_logits) atIndex:4]; + [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; + if (has_bias) { + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:6]; + } else { + [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:6]; + } + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:7]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:8]; + [enc setBuffer:completion offset:0 atIndex:9]; + [enc setThreadgroupMemoryLength:1024u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(128, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer( + cb, owned, "router project/select fused")) return -1; + } + return 1; +} + int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( ds4_gpu_tensor *gate, ds4_gpu_tensor *up, @@ -17348,8 +19216,7 @@ int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( out_dim, x, clamp, - 1, - false); + 1); } int ds4_gpu_shared_mid_swiglu_q8_0_tensor( @@ -17373,8 +19240,7 @@ int ds4_gpu_shared_mid_swiglu_q8_0_tensor( out_dim, x, clamp, - 0, - true); + 0); } int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( @@ -17400,8 +19266,7 @@ int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( out_dim, x, clamp, - 1, - true); + 1); } int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( @@ -17605,9 +19470,7 @@ int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( args.nb12 = n_tok * x_row_bytes; args.nb13 = args.nb12; args.ne1 = (int32_t)n_tok; - const char *fn_name = mv_dispatch.nr0 >= 4 ? - "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : - "kernel_dsv4_shared_gate_up_swiglu_q8_0"; + const char *fn_name = "kernel_dsv4_shared_gate_up_swiglu_q8_0"; id pipeline = ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); if (!pipeline) return 0; @@ -17677,12 +19540,11 @@ int ds4_gpu_matmul_f16_tensor( uint64_t inner_offset = 0; id wbuf = - ds4_gpu_wrap_f32_decode_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - n_tok, - &inner_offset); + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); if (!wbuf) return 0; int owned = 0; @@ -17927,12 +19789,9 @@ int ds4_gpu_matmul_f16_pair_compressor_store_tensor( uint32_t ratio, uint32_t pos) { if (!g_initialized && !ds4_gpu_init()) return -1; - const bool force = - getenv("DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE") != NULL; if ((g_quality_mode || (!ds4_gpu_device_name_contains("M3") && - !ds4_gpu_device_name_contains("M5") && !force)) || - getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE") != NULL || + !ds4_gpu_device_name_contains("M5"))) || getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL || getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) { return 0; @@ -18049,6 +19908,464 @@ int ds4_gpu_matmul_f16_pair_compressor_store_tensor( return 1; } +/* Quad variant of the paired compressor projection: the attention compressor + * and indexer compressor pairs share the input activation and F16 matvec + * shape, so one dispatch covers all four matrices. Bit-exact by + * construction (see metal/dense.metal). Returns 1 when the fused dispatch + * ran, 0 when the caller should use the separate paths, -1 on error. */ +int ds4_gpu_matmul_f16_quad_compressor_store_tensor( + ds4_gpu_tensor *out0_kv, + ds4_gpu_tensor *out0_score, + ds4_gpu_tensor *out1_kv, + ds4_gpu_tensor *out1_score, + ds4_gpu_tensor *state0_kv, + ds4_gpu_tensor *state0_score, + ds4_gpu_tensor *state1_kv, + ds4_gpu_tensor *state1_score, + const void *model_map, + uint64_t model_size, + uint64_t weight0_kv_offset, + uint64_t weight0_score_offset, + uint64_t weight1_kv_offset, + uint64_t weight1_score_offset, + uint64_t ape0_offset, + uint32_t ape0_type, + uint64_t ape1_offset, + uint32_t ape1_type, + uint64_t in_dim, + uint32_t width0, + uint32_t width1, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos) { + if (!g_initialized && !ds4_gpu_init()) return -1; + if (!out0_kv || !out0_score || !out1_kv || !out1_score || + !state0_kv || !state0_score || !state1_kv || !state1_score || + !model_map || !x || + (ape0_type != 0u && ape0_type != 1u) || + (ape1_type != 0u && ape1_type != 1u)) { + return -1; + } + if (in_dim != 4096u || + (width0 != 512u && width0 != 1024u) || + width1 != 256u || ratio != 4u) { + return 0; + } + + @autoreleasepool { + const uint32_t state_rows = 2u * ratio; + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight0_bytes = row_bytes * width0; + const uint64_t weight1_bytes = row_bytes * width1; + const uint64_t ape0_elem = ape0_type == 1u ? sizeof(uint16_t) : sizeof(float); + const uint64_t ape1_elem = ape1_type == 1u ? sizeof(uint16_t) : sizeof(float); + const uint64_t ape0_bytes = (uint64_t)ratio * width0 * ape0_elem; + const uint64_t ape1_bytes = (uint64_t)ratio * width1 * ape1_elem; + if (weight0_kv_offset > model_size || weight0_bytes > model_size - weight0_kv_offset || + weight0_score_offset > model_size || weight0_bytes > model_size - weight0_score_offset || + weight1_kv_offset > model_size || weight1_bytes > model_size - weight1_kv_offset || + weight1_score_offset > model_size || weight1_bytes > model_size - weight1_score_offset || + ape0_offset > model_size || ape0_bytes > model_size - ape0_offset || + ape1_offset > model_size || ape1_bytes > model_size - ape1_offset) { + return -1; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id out0kvbuf = ds4_gpu_tensor_buffer(out0_kv); + id out0scbuf = ds4_gpu_tensor_buffer(out0_score); + id out1kvbuf = ds4_gpu_tensor_buffer(out1_kv); + id out1scbuf = ds4_gpu_tensor_buffer(out1_score); + id state0kvbuf = ds4_gpu_tensor_buffer(state0_kv); + id state0scbuf = ds4_gpu_tensor_buffer(state0_score); + id state1kvbuf = ds4_gpu_tensor_buffer(state1_kv); + id state1scbuf = ds4_gpu_tensor_buffer(state1_score); + if (!xbuf || !out0kvbuf || !out0scbuf || !out1kvbuf || !out1scbuf || + !state0kvbuf || !state0scbuf || !state1kvbuf || !state1scbuf || + ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out0_kv) < (uint64_t)width0 * sizeof(float) || + ds4_gpu_tensor_bytes(out0_score) < (uint64_t)width0 * sizeof(float) || + ds4_gpu_tensor_bytes(out1_kv) < (uint64_t)width1 * sizeof(float) || + ds4_gpu_tensor_bytes(out1_score) < (uint64_t)width1 * sizeof(float) || + ds4_gpu_tensor_bytes(state0_kv) < (uint64_t)state_rows * width0 * sizeof(float) || + ds4_gpu_tensor_bytes(state0_score) < (uint64_t)state_rows * width0 * sizeof(float) || + ds4_gpu_tensor_bytes(state1_kv) < (uint64_t)state_rows * width1 * sizeof(float) || + ds4_gpu_tensor_bytes(state1_score) < (uint64_t)state_rows * width1 * sizeof(float)) { + return -1; + } + + uint64_t w0kv_inner = 0, w0sc_inner = 0, w1kv_inner = 0, w1sc_inner = 0; + uint64_t ape0_inner = 0, ape1_inner = 0; + id w0kvbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_kv_offset, weight0_bytes, &w0kv_inner); + id w0scbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_score_offset, weight0_bytes, &w0sc_inner); + id w1kvbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight1_kv_offset, weight1_bytes, &w1kv_inner); + id w1scbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight1_score_offset, weight1_bytes, &w1sc_inner); + id ape0buf = ds4_gpu_wrap_model_range(model_map, model_size, ape0_offset, ape0_bytes, &ape0_inner); + id ape1buf = ds4_gpu_wrap_model_range(model_map, model_size, ape1_offset, ape1_bytes, &ape1_inner); + if (!w0kvbuf || !w0scbuf || !w1kvbuf || !w1scbuf || !ape0buf || !ape1buf) return -1; + + ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, width0); + mv_args.nr0 = 2; + ds4_gpu_dsv4_compressor_store_one_args store0_args = { + .width = width0, + .ratio = ratio, + .pos = pos, + .ape_type = ape0_type, + }; + ds4_gpu_dsv4_compressor_store_one_args store1_args = { + .width = width1, + .ratio = ratio, + .pos = pos, + .ape_type = ape1_type, + }; + id pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_mul_mv_f16_f32_quad_compressor_store_4", 8); + if (!pipeline) return -1; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return -1; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&store0_args length:sizeof(store0_args) atIndex:1]; + [enc setBytes:&store1_args length:sizeof(store1_args) atIndex:2]; + [enc setBuffer:w0kvbuf offset:(NSUInteger)w0kv_inner atIndex:3]; + [enc setBuffer:w0scbuf offset:(NSUInteger)w0sc_inner atIndex:4]; + [enc setBuffer:w1kvbuf offset:(NSUInteger)w1kv_inner atIndex:5]; + [enc setBuffer:w1scbuf offset:(NSUInteger)w1sc_inner atIndex:6]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:7]; + [enc setBuffer:out0kvbuf offset:ds4_gpu_tensor_offset(out0_kv) atIndex:8]; + [enc setBuffer:out0scbuf offset:ds4_gpu_tensor_offset(out0_score) atIndex:9]; + [enc setBuffer:out1kvbuf offset:(NSUInteger)ds4_gpu_tensor_offset(out1_kv) atIndex:10]; + [enc setBuffer:out1scbuf offset:(NSUInteger)ds4_gpu_tensor_offset(out1_score) atIndex:11]; + [enc setBuffer:ape0buf offset:(NSUInteger)ape0_inner atIndex:12]; + [enc setBuffer:ape1buf offset:(NSUInteger)ape1_inner atIndex:13]; + [enc setBuffer:state0kvbuf offset:ds4_gpu_tensor_offset(state0_kv) atIndex:14]; + [enc setBuffer:state0scbuf offset:ds4_gpu_tensor_offset(state0_score) atIndex:15]; + [enc setBuffer:state1kvbuf offset:ds4_gpu_tensor_offset(state1_kv) atIndex:16]; + [enc setBuffer:state1scbuf offset:ds4_gpu_tensor_offset(state1_score) atIndex:17]; + [enc setThreadgroupMemoryLength:32u * 2u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)width0 + (NSUInteger)width1 + 1u) / 2u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "F16 quad matvec compressor state store")) { + return -1; + } + } + + return 1; +} + +/* Decode-only emit-path fusion: finalize the freshly pooled attention and + * indexer compressor rows (norm + rope + fp8/commit + qat) in one dispatch + * instead of seven. Bit-exact vs the separate dispatches; see the kernel + * comment. Returns 1 when the fused dispatch ran, 0 to fall back. */ +int ds4_gpu_dsv4_comp_row_finalize_tensor( + ds4_gpu_tensor *attn_stage, + ds4_gpu_tensor *attn_cache, + uint32_t attn_comp_row, + uint64_t attn_norm_offset, + ds4_gpu_tensor *index_cache, + uint32_t index_comp_row, + uint64_t index_norm_offset, + ds4_gpu_tensor *attn_state_kv, + ds4_gpu_tensor *attn_state_score, + ds4_gpu_tensor *index_state_kv, + ds4_gpu_tensor *index_state_score, + const void *model_map, + uint64_t model_size, + uint32_t pos, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!attn_stage || !attn_cache || !index_cache || !model_map) return 0; + + @autoreleasepool { + const uint64_t attn_norm_bytes = 512u * sizeof(float); + const uint64_t index_norm_bytes = 128u * sizeof(float); + if (attn_norm_offset > model_size || + attn_norm_bytes > model_size - attn_norm_offset || + index_norm_offset > model_size || + index_norm_bytes > model_size - index_norm_offset) { + return 0; + } + if (!attn_state_kv || !attn_state_score || !index_state_kv || + !index_state_score) return 0; + if (ds4_gpu_tensor_bytes(attn_state_kv) < 8u * 1024u * sizeof(float) || + ds4_gpu_tensor_bytes(attn_state_score) < 8u * 1024u * sizeof(float) || + ds4_gpu_tensor_bytes(index_state_kv) < 8u * 256u * sizeof(float) || + ds4_gpu_tensor_bytes(index_state_score) < 8u * 256u * sizeof(float)) { + return 0; + } + if (ds4_gpu_tensor_bytes(attn_stage) < 512u * sizeof(float) || + ds4_gpu_tensor_bytes(index_cache) < + ((uint64_t)index_comp_row + 1u) * 128u * sizeof(float) || + ds4_gpu_tensor_bytes(attn_cache) < + ((uint64_t)attn_comp_row + 1u) * 512u * sizeof(uint16_t)) { + return 0; + } + + uint64_t attn_norm_inner = 0, index_norm_inner = 0; + id anbuf = ds4_gpu_wrap_model_range( + model_map, model_size, attn_norm_offset, attn_norm_bytes, &attn_norm_inner); + id inbuf = ds4_gpu_wrap_model_range( + model_map, model_size, index_norm_offset, index_norm_bytes, &index_norm_inner); + if (!anbuf || !inbuf) return 0; + + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_comp_row_finalize_f32"); + if (!pipeline) return 0; + + ds4_gpu_comp_finalize_args fin_args = { + .rope = { + .row_bytes = 0, + .token_bytes = 0, + .head_dim = 0, + .n_dims = (int32_t)n_rot, + .n_ctx_orig = (int32_t)n_ctx_orig, + .inverse = 0, + .pos0 = pos, + .pos_step = 0, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }, + .rms_eps = rms_eps, + .pad0 = 0, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&fin_args length:sizeof(fin_args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(attn_stage) + offset:ds4_gpu_tensor_offset(attn_stage) atIndex:1]; + [enc setBuffer:anbuf offset:(NSUInteger)attn_norm_inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(attn_cache) + offset:ds4_gpu_tensor_offset(attn_cache) + + (NSUInteger)attn_comp_row * 512u * sizeof(uint16_t) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(index_cache) + offset:ds4_gpu_tensor_offset(index_cache) + + (NSUInteger)index_comp_row * 128u * sizeof(float) atIndex:4]; + [enc setBuffer:inbuf offset:(NSUInteger)index_norm_inner atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(attn_state_kv) + offset:ds4_gpu_tensor_offset(attn_state_kv) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(attn_state_score) + offset:ds4_gpu_tensor_offset(attn_state_score) atIndex:7]; + [enc setBuffer:ds4_gpu_tensor_buffer(index_state_kv) + offset:ds4_gpu_tensor_offset(index_state_kv) atIndex:8]; + [enc setBuffer:ds4_gpu_tensor_buffer(index_state_score) + offset:ds4_gpu_tensor_offset(index_state_score) atIndex:9]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(2u + (4u * 1024u + 4u * 256u + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor row finalize")) { + return 0; + } + } + + return 1; +} + +/* Decode-only fusion: q_a/kv Q8 pair projection + F16 quad compressor + * projection/store in one dispatch (both read the same normalized input). + * Bit-exact vs the separate dispatches; see the kernel comment. Returns + * 1 when the fused dispatch ran, 0 when the caller must use the separate + * paths, -1 on error. */ +int ds4_gpu_qkv_pair_quad_compressor_store_tensor( + ds4_gpu_tensor *qr, + ds4_gpu_tensor *kv_raw, + ds4_gpu_tensor *out0_kv, + ds4_gpu_tensor *out0_score, + ds4_gpu_tensor *out1_kv, + ds4_gpu_tensor *out1_score, + ds4_gpu_tensor *state0_kv, + ds4_gpu_tensor *state0_score, + ds4_gpu_tensor *state1_kv, + ds4_gpu_tensor *state1_score, + const void *model_map, + uint64_t model_size, + uint64_t q_a_offset, + uint64_t kv_offset, + uint64_t weight0_kv_offset, + uint64_t weight0_score_offset, + uint64_t weight1_kv_offset, + uint64_t weight1_score_offset, + uint64_t ape0_offset, + uint32_t ape0_type, + uint64_t ape1_offset, + uint32_t ape1_type, + uint32_t in_dim, + uint32_t q_rank, + uint32_t kv_dim, + uint32_t width0, + uint32_t width1, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!qr || !kv_raw || !out0_kv || !out0_score || !out1_kv || !out1_score || + !state0_kv || !state0_score || !state1_kv || !state1_score || + !model_map || !x || ratio == 0u || (in_dim & 31u) != 0 || + (width0 & 1u) != 0 || (width1 & 1u) != 0) { + return 0; + } + + @autoreleasepool { + const uint64_t q8_row_bytes = (in_dim / 32u) * 34u; + const uint64_t f16_row_bytes = (uint64_t)in_dim * sizeof(uint16_t); + const uint64_t q_a_bytes = (uint64_t)q_rank * q8_row_bytes; + const uint64_t kv_bytes = (uint64_t)kv_dim * q8_row_bytes; + const uint64_t weight0_bytes = (uint64_t)width0 * f16_row_bytes; + const uint64_t weight1_bytes = (uint64_t)width1 * f16_row_bytes; + const uint64_t state_rows = ratio == 4u ? 2u * ratio : ratio; + const uint64_t ape0_bytes = (uint64_t)width0 * ratio * + (ape0_type == 1u ? 2u : 4u); + const uint64_t ape1_bytes = (uint64_t)width1 * ratio * + (ape1_type == 1u ? 2u : 4u); + if (q_a_offset > model_size || q_a_bytes > model_size - q_a_offset || + kv_offset > model_size || kv_bytes > model_size - kv_offset || + weight0_kv_offset > model_size || + weight0_bytes > model_size - weight0_kv_offset || + weight0_score_offset > model_size || + weight0_bytes > model_size - weight0_score_offset || + weight1_kv_offset > model_size || + weight1_bytes > model_size - weight1_kv_offset || + weight1_score_offset > model_size || + weight1_bytes > model_size - weight1_score_offset || + ape0_offset > model_size || ape0_bytes > model_size - ape0_offset || + ape1_offset > model_size || ape1_bytes > model_size - ape1_offset) { + return -1; + } + if (ds4_gpu_tensor_bytes(qr) < (uint64_t)q_rank * sizeof(float) || + ds4_gpu_tensor_bytes(kv_raw) < (uint64_t)kv_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out0_kv) < (uint64_t)width0 * sizeof(float) || + ds4_gpu_tensor_bytes(out0_score) < (uint64_t)width0 * sizeof(float) || + ds4_gpu_tensor_bytes(out1_kv) < (uint64_t)width1 * sizeof(float) || + ds4_gpu_tensor_bytes(out1_score) < (uint64_t)width1 * sizeof(float) || + ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(state0_kv) < state_rows * width0 * sizeof(float) || + ds4_gpu_tensor_bytes(state0_score) < state_rows * width0 * sizeof(float) || + ds4_gpu_tensor_bytes(state1_kv) < state_rows * width1 * sizeof(float) || + ds4_gpu_tensor_bytes(state1_score) < state_rows * width1 * sizeof(float)) { + return -1; + } + + uint64_t q_a_inner = 0, kv_inner = 0; + uint64_t w0kv_inner = 0, w0sc_inner = 0, w1kv_inner = 0, w1sc_inner = 0; + uint64_t ape0_inner = 0, ape1_inner = 0; + id qw0buf = ds4_gpu_wrap_model_range(model_map, model_size, q_a_offset, q_a_bytes, &q_a_inner); + id qw1buf = ds4_gpu_wrap_model_range(model_map, model_size, kv_offset, kv_bytes, &kv_inner); + id w0kvbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_kv_offset, weight0_bytes, &w0kv_inner); + id w0scbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_score_offset, weight0_bytes, &w0sc_inner); + id w1kvbuf = width1 != 0u + ? ds4_gpu_wrap_model_range(model_map, model_size, weight1_kv_offset, + weight1_bytes, &w1kv_inner) + : w0kvbuf; + id w1scbuf = width1 != 0u + ? ds4_gpu_wrap_model_range(model_map, model_size, weight1_score_offset, + weight1_bytes, &w1sc_inner) + : w0scbuf; + id ape0buf = ds4_gpu_wrap_model_range(model_map, model_size, ape0_offset, ape0_bytes, &ape0_inner); + id ape1buf = width1 != 0u + ? ds4_gpu_wrap_model_range(model_map, model_size, ape1_offset, + ape1_bytes, &ape1_inner) + : ape0buf; + if (width1 == 0u) { + w1kv_inner = w0kv_inner; + w1sc_inner = w0sc_inner; + ape1_inner = ape0_inner; + } + if (!qw0buf || !qw1buf || !w0kvbuf || !w0scbuf || + !w1kvbuf || !w1scbuf || !ape0buf || !ape1buf) return -1; + + ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, q_rank); + ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, kv_dim); + args0.nr0 = 2; + args1.nr0 = 2; + ds4_gpu_f16_matvec_args cargs = ds4_gpu_make_f16_mv_args(in_dim, width0); + cargs.nr0 = 2; + ds4_gpu_dsv4_compressor_store_one_args store0_args = { + .width = width0, + .ratio = ratio, + .pos = pos, + .ape_type = ape0_type, + }; + ds4_gpu_dsv4_compressor_store_one_args store1_args = { + .width = width1, + .ratio = ratio, + .pos = pos, + .ape_type = ape1_type, + }; + const uint32_t max_out = q_rank > kv_dim ? q_rank : kv_dim; + const uint32_t pair_vtgs = (max_out + 1u) / 2u; + id pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_dsv4_qkv_pair_quad_compressor_store_q8_0", 8); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args0 length:sizeof(args0) atIndex:0]; + [enc setBytes:&args1 length:sizeof(args1) atIndex:1]; + [enc setBytes:&cargs length:sizeof(cargs) atIndex:2]; + [enc setBytes:&store0_args length:sizeof(store0_args) atIndex:3]; + [enc setBytes:&store1_args length:sizeof(store1_args) atIndex:4]; + [enc setBytes:&pair_vtgs length:sizeof(pair_vtgs) atIndex:5]; + [enc setBuffer:qw0buf offset:(NSUInteger)q_a_inner atIndex:6]; + [enc setBuffer:qw1buf offset:(NSUInteger)kv_inner atIndex:7]; + [enc setBuffer:w0kvbuf offset:(NSUInteger)w0kv_inner atIndex:8]; + [enc setBuffer:w0scbuf offset:(NSUInteger)w0sc_inner atIndex:9]; + [enc setBuffer:w1kvbuf offset:(NSUInteger)w1kv_inner atIndex:10]; + [enc setBuffer:w1scbuf offset:(NSUInteger)w1sc_inner atIndex:11]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:12]; + [enc setBuffer:ds4_gpu_tensor_buffer(qr) offset:ds4_gpu_tensor_offset(qr) atIndex:13]; + [enc setBuffer:ds4_gpu_tensor_buffer(kv_raw) offset:ds4_gpu_tensor_offset(kv_raw) atIndex:14]; + [enc setBuffer:ds4_gpu_tensor_buffer(out0_kv) offset:ds4_gpu_tensor_offset(out0_kv) atIndex:15]; + [enc setBuffer:ds4_gpu_tensor_buffer(out0_score) offset:ds4_gpu_tensor_offset(out0_score) atIndex:16]; + [enc setBuffer:ds4_gpu_tensor_buffer(out1_kv) offset:ds4_gpu_tensor_offset(out1_kv) atIndex:17]; + [enc setBuffer:ds4_gpu_tensor_buffer(out1_score) offset:ds4_gpu_tensor_offset(out1_score) atIndex:18]; + [enc setBuffer:ape0buf offset:(NSUInteger)ape0_inner atIndex:19]; + [enc setBuffer:ape1buf offset:(NSUInteger)ape1_inner atIndex:20]; + [enc setBuffer:ds4_gpu_tensor_buffer(state0_kv) offset:ds4_gpu_tensor_offset(state0_kv) atIndex:21]; + [enc setBuffer:ds4_gpu_tensor_buffer(state0_score) offset:ds4_gpu_tensor_offset(state0_score) atIndex:22]; + [enc setBuffer:ds4_gpu_tensor_buffer(state1_kv) offset:ds4_gpu_tensor_offset(state1_kv) atIndex:23]; + [enc setBuffer:ds4_gpu_tensor_buffer(state1_score) offset:ds4_gpu_tensor_offset(state1_score) atIndex:24]; + [enc setThreadgroupMemoryLength:2u * 2u * 2u * 32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + (NSUInteger)(pair_vtgs + 1u) / 2u + + ((NSUInteger)width0 + (NSUInteger)width1 + 1u) / 2u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "qkv pair + quad compressor store")) { + return -1; + } + } + + return 1; +} + int ds4_gpu_matmul_f32_tensor( ds4_gpu_tensor *out, const void *model_map, @@ -18086,12 +20403,11 @@ int ds4_gpu_matmul_f32_tensor( uint64_t inner_offset = 0; id wbuf = - ds4_gpu_wrap_f32_decode_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - n_tok, - &inner_offset); + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); if (!wbuf) return 0; int owned = 0; @@ -18326,14 +20642,13 @@ static int ds4_gpu_hc_rms_scale_project_mode( (in_dim == 16384u || in_dim == 28672u) && out_dim == 24u; if (!hard_shape) return 0; - const bool force = - getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL; - if (getenv("DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ") != NULL || - (!ds4_gpu_device_name_contains("M3") && !force)) { + if ((!ds4_gpu_device_name_contains("M3") && + (g_test_flags & DS4_GPU_TEST_HC_RMS_SCALE_PROJ) == 0u) || + getenv("DS4_METAL_DISABLE_HC_RMS_SCALE_PROJ") != NULL) { return 0; } if (g_rms_norm_scale_pipeline == nil) { - return force ? -1 : 0; + return 0; } return 1; } @@ -18405,9 +20720,6 @@ int ds4_gpu_hc_rms_scale_project_f16_tensor( ds4_gpu_get_mul_mm_pipeline( "kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out); if (!mm_pipeline) { - if (getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL) { - return 0; - } return ds4_gpu_rms_norm_plain_rows_tensor( scale_scratch, x, in_dim, n_rows, eps) != 0 && ds4_gpu_matmul_f16_tensor( @@ -18510,23 +20822,12 @@ int ds4_gpu_rms_norm_weight_rows_tensor( return 0; } - const bool exact_decode_weight_view = - rows == 1u && - row_bytes <= (1ull << 20) && - getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; uint64_t inner_offset = 0; - id wbuf = exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset); + id wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset); if (!wbuf) return 0; ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); @@ -18584,22 +20885,12 @@ int ds4_gpu_add_rms_norm_weight_tensor( return 0; } - const bool exact_decode_weight_view = - row_bytes <= (1ull << 20) && - getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; uint64_t inner_offset = 0; - id wbuf = exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset); + id wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset); if (!wbuf) return 0; ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, 1, eps); @@ -18712,6 +21003,140 @@ int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( return 1; } +/* Decode-only triple fusion: q/kv RMS norm + KV RoPE tail + FP8/raw store. + * Same arithmetic, order and rounding as the three separate dispatches (see + * metal/norm.metal); gated and verified against full-vocabulary logits. + * Returns 1 on success, 0 on unsupported shape (caller falls back). */ +int ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_out, + const ds4_gpu_tensor *kv, + uint64_t kv_weight_offset, + uint32_t kv_n, + ds4_gpu_tensor *raw_cache, + uint64_t raw_cap, + uint32_t raw_row, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!q_out || !q || !kv_out || !kv || !raw_cache || + q_n == 0 || kv_n == 0 || (q_n & 3u) != 0 || (kv_n & 3u) != 0 || + n_rot > kv_n || (n_rot & 1u) != 0) { + return 0; + } + + @autoreleasepool { + id qbuf = ds4_gpu_tensor_buffer(q); + id qoutbuf = ds4_gpu_tensor_buffer(q_out); + id kvbuf = ds4_gpu_tensor_buffer(kv); + id kvoutbuf = ds4_gpu_tensor_buffer(kv_out); + id rawbuf = ds4_gpu_tensor_buffer(raw_cache); + + const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); + const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float); + const uint64_t raw_bytes = raw_cap * kv_row_bytes; + if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf || !rawbuf || + ds4_gpu_tensor_bytes(q) < q_row_bytes || + ds4_gpu_tensor_bytes(q_out) < q_row_bytes || + ds4_gpu_tensor_bytes(kv) < kv_row_bytes || + ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes || + ds4_gpu_tensor_bytes(raw_cache) < raw_bytes || + (uint64_t)raw_row >= raw_cap) { + fprintf(stderr, "ds4: Metal fused q/kv norm RoPE store received undersized buffers\n"); + return 0; + } + if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) { + fprintf(stderr, "ds4: Metal fused q/kv norm RoPE store weight range is outside the mapped model\n"); + return 0; + } + + uint64_t q_inner_offset = 0; + uint64_t kv_inner_offset = 0; + id q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + q_weight_offset, q_row_bytes, + &q_inner_offset); + if (!q_wbuf) return 0; + id kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + kv_weight_offset, kv_row_bytes, + &kv_inner_offset); + if (!kv_wbuf) return 0; + + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32"); + if (!pipeline) return 0; + + ds4_gpu_qkv_rms_norm_args args = { + .q_n = (int32_t)q_n, + .q_n4 = (int32_t)(q_n / 4u), + .kv_n = (int32_t)kv_n, + .kv_n4 = (int32_t)(kv_n / 4u), + .q_row_stride = q_row_bytes, + .kv_row_stride = kv_row_bytes, + .eps = eps, + }; + ds4_gpu_rope_affine_pair_args rope = { + .row_bytes = kv_row_bytes, + .token_bytes = kv_row_bytes, + .head_dim = (int32_t)kv_n, + .n_dims = (int32_t)n_rot, + .n_ctx_orig = (int32_t)n_ctx_orig, + .inverse = 0, + .pos0 = pos0, + .pos_step = 1, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + ds4_gpu_dsv4_kv_fp8_store_args store = { + .head_dim = (int32_t)kv_n, + .n_rot = (int32_t)n_rot, + .raw_row = (int32_t)raw_row, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBytes:&rope length:sizeof(rope) atIndex:1]; + [enc setBytes:&store length:sizeof(store) atIndex:2]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:3]; + [enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:4]; + [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:5]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:6]; + [enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:7]; + [enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:8]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:9]; + [enc setThreadgroupMemoryLength:(32u + 64u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 2, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv norm RoPE store")) return 0; + } + + return 1; +} + int ds4_gpu_head_rms_norm_tensor( ds4_gpu_tensor *x, uint32_t n_tok, @@ -18824,11 +21249,85 @@ int ds4_gpu_head_rms_norm_rope_tail_tensor( float beta_fast, float beta_slow, float eps) { - (void)x; (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; - (void)pos0; (void)n_ctx_orig; (void)inverse; (void)freq_base; - (void)freq_scale; (void)ext_factor; (void)attn_factor; - (void)beta_fast; (void)beta_slow; (void)eps; - return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0u || n_head == 0u || head_dim == 0u || + (head_dim & 3u) != 0u || n_rot > head_dim || + pos0 > (uint32_t)INT32_MAX - n_tok) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = + (uint64_t)n_tok * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, + "ds4: Metal fused head norm/RoPE received undersized activation buffer\n"); + return 0; + } + + id pipeline = + g_use_dsv4_head_rms_norm_rope_tail_pipeline + ? g_dsv4_head_rms_norm_rope_tail_pipeline + : ds4_gpu_get_pipeline( + "kernel_dsv4_head_rms_norm_rope_tail_f32"); + if (!pipeline) return 0; + + struct { + int32_t n_head; + int32_t head_dim; + int32_t head_dim4; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t pos0; + int32_t inverse; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + } args = { + .n_head = (int32_t)n_head, + .head_dim = (int32_t)head_dim, + .head_dim4 = (int32_t)(head_dim / 4u), + .n_dims = (int32_t)n_rot, + .n_ctx_orig = (int32_t)n_ctx_orig, + .pos0 = (int32_t)pos0, + .inverse = inverse ? 1 : 0, + .eps = eps, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake( + ds4_gpu_rms_norm_pipeline_threads(head_dim, pipeline), + 1, + 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "fused head norm/RoPE")) { + return 0; + } + } + + return 1; } int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( @@ -18956,6 +21455,7 @@ int ds4_gpu_dsv4_indexer_qat_tensor( return 1; } + static void ds4_gpu_set_rows_thread_shape( uint32_t width, NSUInteger *nth_out, @@ -19160,6 +21660,100 @@ int ds4_gpu_store_raw_kv_tensor( /* Release decode fused KV finalizer. Reference paths are selected by the C * graph driver; this Objective-C entry point always means "use the fused * Metal kernel." */ +/* The fused KV RoPE/FP8 kernel replicates the affine decode RoPE specialisation, + * so it is only valid where that specialisation is the path in use. */ +int ds4_gpu_kv_rope_fp8_fuse_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_dsv4_kv_rope_fp8_store_pipeline == nil) return 0; + if (g_rope_tail_inplace_pair_affine_pipeline == nil) return 0; + if (getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") != NULL) return 0; + if (getenv("DS4_METAL_DISABLE_AFFINE_ROPE_PAIR") != NULL) return 0; + if (!ds4_gpu_device_name_contains("M3") && !ds4_gpu_device_name_contains("M5")) return 0; + return 1; +} + +/* Decode-only fusion: one dispatch does the KV RoPE tail and the FP8/raw + * finalizer that previously cost two. Same arithmetic, same order; gated and + * verified against full-vocabulary logits. */ +int ds4_gpu_kv_rope_fp8_store_raw_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || raw_cap > INT32_MAX || + g_dsv4_kv_rope_fp8_store_pipeline == nil) { + return 0; + } + + @autoreleasepool { + id kvbuf = ds4_gpu_tensor_buffer(kv); + id rawbuf = ds4_gpu_tensor_buffer(raw_cache); + const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (!kvbuf || !rawbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal fused KV RoPE/FP8 store received undersized buffers\n"); + return 0; + } + + ds4_gpu_dsv4_kv_fp8_store_args args = { + .head_dim = (int32_t)head_dim, + .n_rot = (int32_t)n_rot, + .raw_row = (int32_t)row, + }; + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + ds4_gpu_rope_affine_pair_args rope = { + .row_bytes = row_bytes, + .token_bytes = row_bytes, + .head_dim = (int32_t)head_dim, + .n_dims = (int32_t)n_rot, + .n_ctx_orig = (int32_t)n_ctx_orig, + .inverse = inverse ? 1 : 0, + .pos0 = pos0, + .pos_step = 1, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_kv_rope_fp8_store_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBytes:&rope length:sizeof(rope) atIndex:1]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:2]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:3]; + [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "KV RoPE/FP8 store fused")) return 0; + } + + return 1; +} + int ds4_gpu_kv_fp8_store_raw_tensor( ds4_gpu_tensor *kv, ds4_gpu_tensor *raw_cache, @@ -19293,11 +21887,9 @@ static int ds4_gpu_encode_compressor_score_with_ape( } const uint32_t total_elems = (uint32_t)total_elems64; - const bool force_fused = - getenv("DS4_METAL_ENABLE_COMPRESSOR_APE_ADD") != NULL; const bool use_fused = - (ds4_gpu_device_name_contains("M3") || force_fused) && - getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD") == NULL; + ds4_gpu_device_name_contains("M3") && + getenv("DS4_METAL_DISABLE_COMPRESSOR_APE_ADD") == NULL; if (use_fused) { id pipeline = ds4_gpu_get_pipeline( ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16" @@ -19325,7 +21917,6 @@ static int ds4_gpu_encode_compressor_score_with_ape( ds4_gpu_end_compute_encoder(cb, enc); return 1; } - if (force_fused) return 0; } const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); @@ -19770,6 +22361,140 @@ static int ds4_gpu_encode_softmax_f32_contiguous( return 1; } +static int ds4_gpu_ensure_compressor_pool_ggml_scratch( + uint32_t n_rows, + uint32_t head_dim) { + const uint64_t pack_elems = (uint64_t)n_rows * head_dim; + if (n_rows == 0u || head_dim == 0u || + pack_elems > NSUIntegerMax / sizeof(float)) { + return 0; + } + const NSUInteger pack_bytes = + (NSUInteger)pack_elems * sizeof(float); + return + ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer, + &g_compressor_pool_product_bytes, + pack_bytes, + "ds4_compressor_pool_product") && + ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer, + &g_compressor_pool_score_cont_bytes, + pack_bytes, + "ds4_compressor_pool_score_cont") && + ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer, + &g_compressor_pool_softmax_bytes, + pack_bytes, + "ds4_compressor_pool_softmax"); +} + +static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml_reduce( + id cb, + ds4_gpu_tensor *out, + uint32_t n_rows, + uint32_t head_dim, + bool decode_ratio4_one_comp) { + id outbuf = ds4_gpu_tensor_buffer(out); + if (!cb || !outbuf || n_rows == 0u || head_dim == 0u || + ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float) || + !ds4_gpu_ensure_compressor_pool_ggml_scratch(n_rows, head_dim)) { + return 0; + } + + const bool use_pre_m5_exact_reduction_fusion_default = + !g_ssd_streaming_mode && + (ds4_gpu_device_is_pre_m5_apple_silicon() || + ds4_gpu_device_is_m5_apple_silicon()) && + getenv("DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_EXACT_REDUCTION_FUSION") == NULL; + const bool request_exact_reduction_fusion = + decode_ratio4_one_comp && + n_rows == 8u && + (head_dim == 128u || head_dim == 512u) && + use_pre_m5_exact_reduction_fusion_default; + id exact_reduction_fusion_pipeline = nil; + if (request_exact_reduction_fusion) { + exact_reduction_fusion_pipeline = ds4_gpu_hot_pipeline( + g_dsv4_compressor_exact_softmax_product_pipeline, + "kernel_dsv4_compressor_exact_softmax_product_ratio4"); + } + const bool use_exact_reduction_fusion = + exact_reduction_fusion_pipeline != nil && + exact_reduction_fusion_pipeline.threadExecutionWidth == 32u && + exact_reduction_fusion_pipeline.maxTotalThreadsPerThreadgroup >= 32u; + if (use_exact_reduction_fusion) { + if (getenv( + "DS4_METAL_TEST_POISON_COMPRESSOR_EXACT_REDUCTION_SCRATCH") != NULL) { + const uint32_t poison_bits = 0x7fc01234u; + float poison; + memcpy(&poison, &poison_bits, sizeof(poison)); + if (!ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_softmax_buffer, + 0, + n_rows, + head_dim, + poison)) { + return 0; + } + } + + ds4_gpu_dsv4_compressor_pack_ratio4_args args = { + .head_dim = head_dim, + .n_comp = 1u, + .replay = 1u, + .n_threads = 32u, + }; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:exact_reduction_fusion_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:1]; + [enc setBuffer:g_compressor_pool_score_cont_buffer offset:0 atIndex:2]; + [enc setBuffer:g_compressor_pool_softmax_buffer offset:0 atIndex:3]; + [enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(head_dim, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_encode_sum_rows_f32(cb, + g_compressor_pool_product_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out), + n_rows, + head_dim); + } + + /* + * Preserve the exact graph reduction sequence and its contiguous + * [head_dim, n_rows] scratch layout. Only the preceding copy/transpose + * may be fused by callers. + */ + ds4_gpu_bin_args mul_args = + ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1); + return + ds4_gpu_encode_softmax_f32_contiguous(cb, + g_compressor_pool_score_cont_buffer, + 0, + g_compressor_pool_softmax_buffer, + 0, + n_rows, + head_dim, + 1) && + ds4_gpu_encode_bin_f32_rows(cb, + g_mul_pipeline, + &mul_args, + g_compressor_pool_product_buffer, + 0, + g_compressor_pool_softmax_buffer, + 0, + g_compressor_pool_product_buffer, + 0) && + ds4_gpu_encode_sum_rows_f32(cb, + g_compressor_pool_product_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out), + n_rows, + head_dim); +} + static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml( id cb, ds4_gpu_tensor *out, @@ -19791,39 +22516,13 @@ static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml( return 0; } - const NSUInteger pack_bytes = (NSUInteger)n_rows * head_dim * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer, - &g_compressor_pool_product_bytes, - pack_bytes, - "ds4_compressor_pool_product") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer, - &g_compressor_pool_score_cont_bytes, - pack_bytes, - "ds4_compressor_pool_score_cont") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer, - &g_compressor_pool_softmax_bytes, - pack_bytes, - "ds4_compressor_pool_softmax")) { + if (!ds4_gpu_ensure_compressor_pool_ggml_scratch(n_rows, head_dim)) { return 0; } const uint64_t cont_row_stride = (uint64_t)n_rows * sizeof(float); const uint64_t cont_plane_stride = (uint64_t)head_dim * cont_row_stride; - /* - * Keep the n_comp == 1 compressor path as the unfused graph sequence: - * - * score = soft_max(contiguous(score)) - * pooled = sum_rows(contiguous(kv) * score) - * - * The fused DS4 pool kernel is mathematically equivalent, but it reduces in - * a different order. That is enough to create ~1e-6 compressor differences - * and later FP8/routing flips, so this path intentionally keeps the same - * operation boundary and memory layout as the graph. - */ - ds4_gpu_bin_args mul_args = - ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1); - return ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, kvbuf, @@ -19851,30 +22550,11 @@ static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml( score_nb2, cont_row_stride, cont_plane_stride) && - ds4_gpu_encode_softmax_f32_contiguous(cb, - g_compressor_pool_score_cont_buffer, - 0, - g_compressor_pool_softmax_buffer, - 0, - n_rows, - head_dim, - 1) && - ds4_gpu_encode_bin_f32_rows(cb, - g_mul_pipeline, - &mul_args, - g_compressor_pool_product_buffer, - 0, - g_compressor_pool_softmax_buffer, - 0, - g_compressor_pool_product_buffer, - 0) && - ds4_gpu_encode_sum_rows_f32(cb, - g_compressor_pool_product_buffer, - 0, - outbuf, - ds4_gpu_tensor_offset(out), - n_rows, - head_dim); + ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml_reduce(cb, + out, + n_rows, + head_dim, + false); } static int ds4_gpu_encode_dsv4_softmax_pool( @@ -20014,15 +22694,28 @@ static int ds4_gpu_encode_concat_f32_dim1( } static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) { - const bool force = - getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL; const bool default_shape = head_dim == 128u || head_dim == 512u; - if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION") != NULL || - (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + if (getenv("DS4_METAL_DISABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL || + !(ds4_gpu_device_name_contains("M3") && default_shape)) { return 0; } if (g_dsv4_compressor_pack_ratio4_pipeline == nil) { - return force ? -1 : 0; + return 0; + } + return 1; +} + +static int ds4_gpu_compressor_ratio4_decode_pack_mode(uint32_t head_dim) { + const bool default_shape = head_dim == 128u || head_dim == 512u; + const bool default_pre_m5 = + (ds4_gpu_device_is_pre_m5_apple_silicon() || + ds4_gpu_device_is_m5_apple_silicon()) && default_shape; + if (getenv("DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_RATIO4_DECODE_PACK_FUSION") != NULL || + !default_pre_m5) { + return 0; + } + if (g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline == nil) { + return 0; } return 1; } @@ -20047,15 +22740,13 @@ static int ds4_gpu_compressor_ratio4_direct_pool_mode( // One compressed row intentionally uses the legacy GGML reduction graph. if (n_comp <= 1u) return 0; - const bool force = - getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL; const bool default_shape = head_dim == 128u || head_dim == 512u; - if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL || - (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + if (getenv("DS4_METAL_DISABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL || + !(ds4_gpu_device_name_contains("M3") && default_shape)) { return 0; } if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) { - return force ? -1 : 0; + return 0; } return 1; } @@ -20170,13 +22861,122 @@ static int ds4_gpu_encode_compressor_pack_ratio4( return 1; } +static int ds4_gpu_encode_compressor_ratio4_decode_pack_ggml( + id cb, + ds4_gpu_tensor *out, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim) { + if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0u || + !ds4_gpu_ensure_compressor_pool_ggml_scratch(8u, head_dim)) { + return 0; + } + + const bool exact_pool_enabled = + ds4_gpu_ported_m5_decode_feature_enabled( + "DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_EXACT_POOL_RATIO4", + "DS4_METAL_DISABLE_M5_COMPRESSOR_EXACT_POOL_RATIO4"); + if ((head_dim == 128u || head_dim == 512u) && exact_pool_enabled) { + id exact_pool_pipeline = ds4_gpu_hot_pipeline( + g_dsv4_compressor_exact_pool_ratio4_pipeline, + "kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml"); + id outbuf = ds4_gpu_tensor_buffer(out); + const NSUInteger out_offset = ds4_gpu_tensor_offset(out); + const uint64_t state_bytes = 16ull * head_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)head_dim * sizeof(float); + const bool overlap = + ds4_gpu_buffer_ranges_overlap(outbuf, out_offset, out_bytes, + statekvbuf, state_kv_offset, + state_bytes) || + ds4_gpu_buffer_ranges_overlap(outbuf, out_offset, out_bytes, + statescbuf, state_score_offset, + state_bytes); + if (exact_pool_pipeline && outbuf && !overlap && + exact_pool_pipeline.threadExecutionWidth == 32u && + exact_pool_pipeline.maxTotalThreadsPerThreadgroup >= 32u) { + if (getenv( + "DS4_METAL_TEST_POISON_COMPRESSOR_EXACT_REDUCTION_SCRATCH") != NULL) { + const uint32_t poison_bits = 0x7fc01234u; + float poison; + memcpy(&poison, &poison_bits, sizeof(poison)); + if (!ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_softmax_buffer, + 0, + 8u, + head_dim, + poison)) { + return 0; + } + } + ds4_gpu_dsv4_compressor_pack_ratio4_args args = { + .head_dim = head_dim, + .n_comp = 1u, + .replay = 1u, + .n_threads = 32u, + }; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:exact_pool_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:1]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:2]; + [enc setBuffer:g_compressor_pool_softmax_buffer offset:0 atIndex:3]; + [enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:4]; + [enc setBuffer:outbuf offset:out_offset atIndex:5]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(head_dim, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + } + + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline, + "kernel_dsv4_compressor_pack_ratio4_decode_ggml"); + if (!pipeline) return 0; + + NSUInteger nth = head_dim; + if (nth > 256u) nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0u) return 0; + + ds4_gpu_dsv4_compressor_pack_ratio4_args args = { + .head_dim = head_dim, + .n_comp = 1u, + .replay = 1u, + .n_threads = (uint32_t)nth, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:1]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:2]; + [enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:3]; + [enc setBuffer:g_compressor_pool_score_cont_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(8u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml_reduce(cb, + out, + 8u, + head_dim, + true); +} + static int ds4_gpu_encode_compressor_pool( id cb, ds4_gpu_tensor *out, const ds4_gpu_tensor *state_kv, const ds4_gpu_tensor *state_score, uint32_t head_dim, - uint32_t ratio) { + uint32_t ratio, + bool decode_one_token) { id statekvbuf = ds4_gpu_tensor_buffer(state_kv); id statescbuf = ds4_gpu_tensor_buffer(state_score); if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0 || ratio == 0) return 0; @@ -20209,6 +23009,21 @@ static int ds4_gpu_encode_compressor_pool( 1); } + const int decode_pack_mode = decode_one_token + ? ds4_gpu_compressor_ratio4_decode_pack_mode(head_dim) + : 0; + if (decode_pack_mode < 0) return 0; + if (decode_pack_mode > 0) { + return ds4_gpu_encode_compressor_ratio4_decode_pack_ggml( + cb, + out, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim); + } + const NSUInteger packed_bytes = (NSUInteger)8u * head_dim * sizeof(float); if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, &g_compressor_pool_kv_bytes, @@ -21207,7 +24022,9 @@ int ds4_gpu_compressor_update_tensor( float beta_fast, float beta_slow, float rms_eps, - bool state_already_stored) { + bool state_already_stored, + bool decode_one_token, + bool defer_finalize) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || !model_map || head_dim == 0 || ratio == 0 || @@ -21295,9 +24112,14 @@ int ds4_gpu_compressor_update_tensor( state_kv, state_score, head_dim, - ratio); + ratio, + decode_one_token); if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor DS4 softmax pool"); - if (ok) { + /* When the caller runs the fused finalize dispatch afterwards, the + * pooled row is consumed un-normalized there; skip the standalone + * norm+rope so the work happens exactly once. The state shift still + * runs: it only moves the rolling state rows. */ + if (ok && !defer_finalize) { ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, comp_row_view, model_map, @@ -21307,7 +24129,7 @@ int ds4_gpu_compressor_update_tensor( 1, rms_eps) != 0; } - if (ok) { + if (ok && !defer_finalize) { const uint32_t comp_pos = pos + 1u - ratio; ok = ds4_gpu_rope_tail_tensor(comp_row_view, 1, @@ -21324,7 +24146,7 @@ int ds4_gpu_compressor_update_tensor( beta_fast, beta_slow) != 0; } - if (ok && ratio == 4u) { + if (ok && ratio == 4u && !defer_finalize) { cb = ds4_gpu_command_buffer(&owned); ok = cb && ds4_gpu_encode_compressor_shift_ratio4(cb, @@ -21532,15 +24354,15 @@ int ds4_gpu_attention_output_q8_batch_tensor( "kernel_attn_out_low_q8_0_mpp_direct_rhs_n64"; id mm_pipeline = ds4_gpu_get_mul_mm_id_pipeline(attn_out_pipeline_name, false); - ok = ds4_gpu_encode_attn_out_low_q8_mpp(cb, - mm_pipeline, - &mm_args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low)) != 0; + ok = ds4_gpu_encode_attn_out_low_mpp(cb, + mm_pipeline, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low)) != 0; if (!ok) { ds4_gpu_warn_mpp_fallback(); if (ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { @@ -21771,23 +24593,57 @@ int ds4_gpu_attention_output_q4_K_batch_tensor( (void)group_tmp; (void)low_tmp; + const uint64_t padded_n_tokens_u64 = + ((uint64_t)n_tokens + DS4_METAL_ATTN_OUT_MPP_TILE_N - 1u) / + DS4_METAL_ATTN_OUT_MPP_TILE_N * DS4_METAL_ATTN_OUT_MPP_TILE_N; + const uint32_t padded_n_tokens = + padded_n_tokens_u64 <= UINT32_MAX ? (uint32_t)padded_n_tokens_u64 : n_tokens; + /* + * TensorOps requires complete 64-row tiles. For a substantial final + * prefill chunk, compute harmless extra rows when the caller's + * workspace already covers them; the following projection consumes + * only the real rows. Small tails stay on the legacy kernel because + * padding overhead outweighs the faster arithmetic there. + */ + const bool use_mpp_padding = + n_tokens >= 256u && + padded_n_tokens > n_tokens && + ds4_gpu_tensor_bytes(heads) >= + (uint64_t)padded_n_tokens * n_groups * group_dim * sizeof(float) && + ds4_gpu_tensor_bytes(low) >= + (uint64_t)padded_n_tokens * low_dim * sizeof(float); + bool use_mpp_low = + ((n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 || + use_mpp_padding) && + ds4_gpu_use_mpp_attn_out_low_matmul(); + id mpp_low_pipeline = nil; + if (use_mpp_low) { + mpp_low_pipeline = ds4_gpu_get_mul_mm_id_pipeline( + "kernel_attn_out_low_q4_K_mpp_direct_rhs_n64", false); + if (!mpp_low_pipeline) { + ds4_gpu_warn_mpp_fallback(); + use_mpp_low = false; + } + } const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); id group_ids_buffer = nil; - if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { - group_ids_buffer = - ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids"); - } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, - &g_attn_out_group_ids_bytes, - ids_bytes, - "ds4_attention_output_group_ids")) { - group_ids_buffer = g_attn_out_group_ids_buffer; - } - if (!group_ids_buffer) return 0; + if (!use_mpp_low) { + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids"); + } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + group_ids_buffer = g_attn_out_group_ids_buffer; + } + if (!group_ids_buffer) return 0; - int32_t *ids = (int32_t *)[group_ids_buffer contents]; - for (uint32_t t = 0; t < n_tokens; t++) { - for (uint32_t group = 0; group < n_groups; group++) { - ids[(uint64_t)t * n_groups + group] = (int32_t)group; + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } } } @@ -21807,12 +24663,6 @@ int ds4_gpu_attention_output_q4_K_batch_tensor( if (!cb || owned) ok = false; if (ok) { - ds4_gpu_mul_mm_id_map_args map_args = - ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, - n_groups, - n_groups, - n_groups, - n_tokens); ds4_gpu_mul_mm_id_args mm_args = ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, (uint32_t)rank, @@ -21821,24 +24671,45 @@ int ds4_gpu_attention_output_q4_K_batch_tensor( (uint64_t)rank * row_a_bytes, n_groups, n_groups, - n_tokens); - id map_pipeline = - ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); - id mm_pipeline = - ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); - ok = ds4_gpu_encode_mul_mm_id(cb, - map_pipeline, - mm_pipeline, - &map_args, - &mm_args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - group_ids_buffer, - 0) != 0; + use_mpp_low && use_mpp_padding + ? padded_n_tokens + : n_tokens); + if (use_mpp_low) { + ok = ds4_gpu_encode_attn_out_low_mpp( + cb, + mpp_low_pipeline, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low)) != 0; + } else { + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + mm_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } } if (ok) { @@ -22885,10 +25756,6 @@ static int ds4_gpu_encode_flash_kv_stage_f16( return 0; } - const bool force = - getenv("DS4_METAL_ENABLE_GATHERED_KV_STAGE") != NULL; - const bool disabled = - getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE") != NULL; const bool require = getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL; const bool supported_shape = @@ -22906,17 +25773,18 @@ static int ds4_gpu_encode_flash_kv_stage_f16( const bool valid_grid = total_vecs64 != 0 && total_vecs64 <= UINT32_MAX; const bool eligible = - supported_shape && valid_grid && !g_quality_mode && !disabled && + supported_shape && valid_grid && !g_quality_mode && + getenv("DS4_METAL_DISABLE_GATHERED_KV_STAGE") == NULL && g_flash_kv_stage_f16_pipeline != nil && (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || force); + ds4_gpu_device_name_contains("M5")); const bool component_disabled = eligible && (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 || ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0); const bool use_fusion = eligible && !component_disabled; const bool use_pad_fusion = use_fusion && fuse_pad && mask != nil && pad != nil && - getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION") == NULL; + getenv("DS4_METAL_DISABLE_GATHERED_KV_PAD_FUSION") == NULL; if (require && supported_shape && !use_fusion) { fprintf(stderr, "ds4: required Metal gathered KV staging kernel was not selected\n"); @@ -23125,8 +25993,21 @@ static int ds4_gpu_encode_flash_attention_raw_heads( const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + nrows * (2u * (NSUInteger)nwg) * sizeof(float); - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + const bool use_persistent_zero_mask = + ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_RAW_ZERO_ATTN_MASK") == NULL && + !g_quality_mode && + !g_ssd_streaming_mode && + g_tp_split_world == 1; + id mask_buffer = nil; + if (use_persistent_zero_mask && + ds4_gpu_ensure_zero_attention_mask(mask_bytes)) { + mask_buffer = g_flash_attn_zero_mask_buffer; + } else { + mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_buffer) memset([mask_buffer contents], 0, mask_bytes); + } if (!mask_buffer || !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, &g_flash_attn_kv_bytes, @@ -23142,8 +26023,6 @@ static int ds4_gpu_encode_flash_attention_raw_heads( "ds4_flash_attn_tmp")) { return 0; } - memset([mask_buffer contents], 0, mask_bytes); - id pad_pipeline = nil; if ((n_raw % ncpsg) != 0) { pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); @@ -24654,11 +27533,27 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + const bool packed_shape = + ds4_gpu_ported_m5_decode_feature_enabled( + "DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PACKED32_REDUCE", + NULL) && + !g_quality_mode && use_mask == 0u && comp_kv_f16 != 0u && n_comp != 0u && + n_head == 64u && head_dim == 512u && nsg == 1u && nwg == 32u && + n_keys <= 1024u && g_decode_attn_rope_fuse != 0 && + g_decode_attn_rope_args.head_dim == 512 && + g_decode_attn_rope_args.n_dims == 64 && + g_decode_attn_rope_args.row_bytes == 2048 && + g_decode_attn_rope_args.inverse != 0; + const bool m5_persistent_zero_mask = + ds4_gpu_device_is_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_M5_PERSISTENT_ZERO_ATTN_MASK") == NULL && + (!packed_shape || + getenv("DS4_METAL_DISABLE_M5_PACKED_ZERO_MASK") == NULL); const bool use_persistent_zero_mask = use_mask == 0u && (ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK") != NULL) && - getenv("DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK") == NULL; + m5_persistent_zero_mask) && + getenv("DS4_METAL_DISABLE_PERSISTENT_ZERO_ATTN_MASK") == NULL; if (!(use_persistent_zero_mask ? ds4_gpu_ensure_zero_attention_mask(mask_bytes) @@ -24687,20 +27582,57 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( const bool has_kvpad = (n_keys % ncpsg) != 0; const bool use_shared_kvpad = has_kvpad && - (ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_SHARED_KV_PAD") != NULL) && - getenv("DS4_METAL_DISABLE_M3_SHARED_KV_PAD") == NULL; - id vec_pipeline = - ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", - true, true, false, false, has_kvpad, - use_shared_kvpad, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg, - (int32_t)nwg); - id reduce_pipeline = - ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); - if (!vec_pipeline || !reduce_pipeline) return 0; + ds4_gpu_device_name_contains("M3") && + getenv("DS4_METAL_DISABLE_SHARED_KV_PAD") == NULL; + const NSUInteger packed_threads = 8u * 32u; + const NSUInteger packed_shared_bytes = + 512u * sizeof(uint16_t) + + 8u * 128u * sizeof(uint16_t) + + (32u * 32u + 2u * 32u + 32u) * sizeof(float) + + 32u * 33u * 4u * sizeof(float); + const bool packed_requested = + packed_shape && + getenv("DS4_METAL_DISABLE_M5_FLASH_ATTN_PACKED32_REDUCE") == NULL; + + id packed_pipeline = nil; + if (packed_requested) { + packed_pipeline = ds4_gpu_get_flash_attn_vec_pipeline( + "kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512", + true, true, false, false, has_kvpad, use_shared_kvpad, + (int32_t)head_dim, (int32_t)head_dim, 1, 32); + } + const NSUInteger max_tgmem = [g_device maxThreadgroupMemoryLength]; + const bool use_packed = + packed_pipeline != nil && + packed_pipeline.threadExecutionWidth == 32u && + packed_pipeline.maxTotalThreadsPerThreadgroup >= packed_threads && + (max_tgmem == 0u || max_tgmem >= packed_shared_bytes); + if (packed_requested && + getenv("DS4_METAL_TRACE_M5_FLASH_ATTN_PACKED32_REDUCE") != NULL) { + fprintf(stderr, + "ds4: packed FA use=%d max_threads=%lu tew=%lu tgmem=%lu need=%lu\n", + use_packed ? 1 : 0, + (unsigned long)(packed_pipeline ? + packed_pipeline.maxTotalThreadsPerThreadgroup : 0u), + (unsigned long)(packed_pipeline ? + packed_pipeline.threadExecutionWidth : 0u), + (unsigned long)max_tgmem, + (unsigned long)packed_shared_bytes); + } + + id vec_pipeline = nil; + id reduce_pipeline = nil; + if (!use_packed) { + vec_pipeline = ds4_gpu_get_flash_attn_vec_pipeline( + "kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, has_kvpad, use_shared_kvpad, + (int32_t)head_dim, (int32_t)head_dim, + (int32_t)nsg, (int32_t)nwg); + reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, + (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + } if (!use_persistent_zero_mask && !ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) { @@ -24818,6 +27750,29 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + if (use_packed) { + id packed_enc = ds4_gpu_compute_encoder(cb); + [packed_enc setComputePipelineState:packed_pipeline]; + [packed_enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [packed_enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [packed_enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [packed_enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [packed_enc setBuffer:flash_mask_buffer offset:0 atIndex:4]; + [packed_enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [packed_enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [packed_enc setBuffer:headsbuf + offset:ds4_gpu_tensor_offset(heads) atIndex:7]; + [packed_enc setBytes:&g_decode_attn_rope_args + length:sizeof(g_decode_attn_rope_args) atIndex:8]; + [packed_enc setThreadgroupMemoryLength:packed_shared_bytes atIndex:0]; + [packed_enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(packed_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, packed_enc); + g_decode_attn_rope_fuse = 0; + g_decode_attn_rope_fuse_used = 1; + return 1; + } + id enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:vec_pipeline]; [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; @@ -24836,11 +27791,26 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( ds4_gpu_flash_attn_reduce_args reduce_args = { .nrows = (int32_t)nrows, }; + /* When ds4.c defers the inverse RoPE tail, the reduce threadgroup that owns + * a head's whole row rotates it in place, removing a dispatch per layer. */ + const int fuse_rope = g_decode_attn_rope_fuse; + g_decode_attn_rope_fuse = 0; + if (fuse_rope) g_decode_attn_rope_fuse_used = 1; + id reduce_pso = reduce_pipeline; + if (fuse_rope) { + reduce_pso = ds4_gpu_get_flash_attn_reduce_rope_pipeline((int32_t)head_dim, + (int32_t)nwg); + if (!reduce_pso) return 0; + } enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:reduce_pipeline]; + [enc setComputePipelineState:reduce_pso]; [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + if (fuse_rope) { + [enc setBytes:&g_decode_attn_rope_args + length:sizeof(g_decode_attn_rope_args) atIndex:3]; + } [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); @@ -25672,13 +28642,35 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, "kernel_dsv4_sort_i32_rows_asc"); const bool decode_one_token = n_tokens == 1u; + const bool prefill_dual_heads = + !decode_one_token && !g_quality_mode && ds4_gpu_mpp_available() && + n_head == 64u && + top_k == 512u && window == 128u && head_dim == 512u; + const uint32_t decode_splits = + decode_one_token && !g_quality_mode ? 12u : 1u; + const bool split_decode = decode_splits > 1u; id attn_pipeline = + split_decode ? + ds4_gpu_hot_pipeline( + g_dsv4_indexed_attention_heads8_split_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8_split") : decode_one_token ? ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline, "kernel_dsv4_indexed_mixed_attention_heads8_rb16") : + prefill_dual_heads ? + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads16_dual") : ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline, "kernel_dsv4_indexed_mixed_attention_heads8"); - if (!sort_pipeline || !attn_pipeline) return 0; + id split_reduce_pipeline = split_decode ? + ds4_gpu_hot_pipeline( + g_dsv4_indexed_attention_heads8_split_reduce_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8_split_reduce") : + nil; + if (!sort_pipeline || !attn_pipeline || + (split_decode && !split_reduce_pipeline)) { + return 0; + } if ((NSUInteger)top_k > sort_pipeline.maxTotalThreadsPerThreadgroup) { fprintf(stderr, "ds4: Metal indexed attention top-k exceeds sort threadgroup limit\n"); return 0; @@ -25719,7 +28711,7 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( .window = window, .ratio = ratio, .comp_kv_f16 = comp_kv_f16 ? 1u : 0u, - .pad0 = 0, + .n_splits = decode_splits, .q_token_stride = (uint64_t)n_head * row_bytes, .q_head_stride = row_bytes, .raw_row_stride = row_bytes, @@ -25747,23 +28739,72 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( ds4_gpu_end_compute_encoder(cb, enc); } - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; - [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3]; - [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer - offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0 - atIndex:4]; - [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6]; - [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) * - 128u * 4u * sizeof(uint16_t) - atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, ((NSUInteger)n_head + 7u) / 8u, 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); + if (split_decode) { + const uint64_t nrows = (uint64_t)n_tokens * n_head; + const uint64_t partial_bytes = + nrows * head_dim * decode_splits * sizeof(float); + const uint64_t stats_bytes = + nrows * decode_splits * 2u * sizeof(float); + if (partial_bytes > NSUIntegerMax - stats_bytes || + !ds4_gpu_ensure_scratch_buffer( + &g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + (NSUInteger)(partial_bytes + stats_bytes), + "ds4_dsv4_indexed_attention_split_tmp")) { + return 0; + } + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3]; + [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer + offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0 + atIndex:4]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:5]; + [enc setThreadgroupMemoryLength:16u * 128u * sizeof(uint16_t) * 4u + atIndex:0]; + [enc dispatchThreadgroups: + MTLSizeMake((NSUInteger)n_tokens, + ((NSUInteger)n_head + 7u) / 8u, + decode_splits) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:split_reduce_pipeline]; + [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:2]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } else { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3]; + [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer + offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0 + atIndex:4]; + [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6]; + [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) * + 128u * 4u * sizeof(uint16_t) + atIndex:0]; + [enc dispatchThreadgroups: + MTLSizeMake((NSUInteger)n_tokens, + ((NSUInteger)n_head + (prefill_dual_heads ? 15u : 7u)) / + (prefill_dual_heads ? 16u : 8u), + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } if (!ds4_gpu_finish_command_buffer(cb, owned, "graph indexed mixed attention heads")) return 0; } @@ -26497,6 +29538,7 @@ static uint32_t ds4_gpu_routed_mv_nr0(uint32_t type) { case DS4_METAL_TENSOR_Q8_0: return 2; case DS4_METAL_TENSOR_Q8_K: return 2; case DS4_METAL_TENSOR_Q4_K: return 2; + case DS4_METAL_TENSOR_MXFP4: return 2; case DS4_METAL_TENSOR_Q2_K: case DS4_METAL_TENSOR_IQ2_XXS: return 4; default: return 0; @@ -26510,6 +29552,7 @@ static const char *ds4_gpu_metal_tensor_type_name(uint32_t type) { case DS4_METAL_TENSOR_Q4_K: return "q4_k"; case DS4_METAL_TENSOR_Q5_K: return "q5_k"; case DS4_METAL_TENSOR_Q6_K: return "q6_k"; + case DS4_METAL_TENSOR_MXFP4: return "mxfp4"; default: return "unknown"; } } @@ -26595,6 +29638,9 @@ static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) { if (type == DS4_METAL_TENSOR_IQ2_XXS) { return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t); } + if (type == DS4_METAL_TENSOR_MXFP4) { + return 32u * sizeof(float); + } return 0; } @@ -26615,10 +29661,18 @@ static id ds4_gpu_routed_mv_pipeline(uint32_t type) { case DS4_METAL_TENSOR_IQ2_XXS: return g_moe_mul_mv_id_iq2_xxs_pipeline; case DS4_METAL_TENSOR_Q2_K: return g_moe_mul_mv_id_q2_k_pipeline; case DS4_METAL_TENSOR_Q4_K: return g_moe_mul_mv_id_q4_k_pipeline; + case DS4_METAL_TENSOR_MXFP4: return g_moe_mul_mv_id_mxfp4_pipeline; default: return nil; } } +/* TensorOps routed-MoE prefill uses bits 0/1/2 for gate/up/down. Unsupported + * tensor types keep their established kernels. --quality and the global + * Metal4 comparison switch retain the reference path. */ +static int ds4_gpu_routed_mm_mpp_mask(void) { + return ds4_gpu_mpp_available() ? 7 : 0; +} + static id ds4_gpu_routed_mm_pipeline(uint32_t type) { switch (type) { case DS4_METAL_TENSOR_Q8_0: @@ -26635,6 +29689,8 @@ static id ds4_gpu_routed_mm_pipeline(uint32_t type) { return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f32", false); case DS4_METAL_TENSOR_Q6_K: return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f32", false); + case DS4_METAL_TENSOR_MXFP4: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_mxfp4_f32", false); default: return nil; } @@ -26646,6 +29702,8 @@ static id ds4_gpu_routed_mm_addr_pipeline(uint32_t type return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f32", false); case DS4_METAL_TENSOR_Q4_K: return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f32", false); + case DS4_METAL_TENSOR_MXFP4: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_mxfp4_f32", false); default: return nil; } @@ -26667,6 +29725,8 @@ static id ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t t return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f16", false); case DS4_METAL_TENSOR_Q6_K: return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f16", false); + case DS4_METAL_TENSOR_MXFP4: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_mxfp4_f16", false); default: return nil; } @@ -26678,6 +29738,8 @@ static id ds4_gpu_routed_mm_addr_f16_rhs_pipeline(uint3 return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f16", false); case DS4_METAL_TENSOR_Q4_K: return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f16", false); + case DS4_METAL_TENSOR_MXFP4: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_mxfp4_f16", false); default: return nil; } @@ -28064,11 +31126,25 @@ static int ds4_gpu_encode_mul_mm_id_map( } const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); - const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + const NSUInteger hids_bytes = + (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); if (tpe_bytes > NSUIntegerMax - hids_bytes) return 0; + const NSUInteger work_offset = (tpe_bytes + hids_bytes + 7u) & ~7u; + const uint64_t pair_rows = + (uint64_t)(uint32_t)mm_args->ne20 * (uint32_t)mm_args->ne21; + const uint64_t work_cap = + (pair_rows + 31u * (uint32_t)mm_args->ne02 + 31u) / 32u; + const NSUInteger work_item_bytes = 2u * sizeof(uint32_t); + if (work_cap > (NSUIntegerMax - 8u) / work_item_bytes || + work_offset > NSUIntegerMax - 8u - + (NSUInteger)work_cap * work_item_bytes) { + return 0; + } + const NSUInteger total_bytes = + work_offset + 8u + (NSUInteger)work_cap * work_item_bytes; if (!ds4_gpu_ensure_scratch_buffer(&g_moe_id_map_buffer, &g_moe_id_map_bytes, - tpe_bytes + hids_bytes, + total_bytes, "ds4_moe_id_map")) { return 0; } @@ -28079,6 +31155,7 @@ static int ds4_gpu_encode_mul_mm_id_map( [enc setBuffer:ids offset:ids_off atIndex:1]; [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2]; [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:3]; + [enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:4]; [enc setThreadgroupMemoryLength:(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne20 * sizeof(uint16_t) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake((NSUInteger)mm_args->ne02, 1, 1)]; @@ -28108,15 +31185,27 @@ static int ds4_gpu_encode_mul_mm_id_mapped_tile( * The removed TensorOps variant was not semantically stable on evals, so keep * this encoder tied to the tested simdgroup kernel shape. */ - const NSUInteger tile_n = 32u; const bool use_resource_hints = getenv("DS4_METAL_MOE_MM_ID_USE_RESOURCES") != NULL && getenv("DS4_METAL_DISABLE_MOE_MM_ID_USE_RESOURCES") == NULL; const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); - if (tpe_bytes > NSUIntegerMax - hids_bytes || - g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + if (tpe_bytes > NSUIntegerMax - hids_bytes) { + return 0; + } + const NSUInteger work_offset = (tpe_bytes + hids_bytes + 7u) & ~7u; + const uint64_t pair_rows = + (uint64_t)(uint32_t)mm_args->ne20 * (uint32_t)mm_args->ne21; + const uint64_t work_cap = + (pair_rows + 31u * (uint32_t)mm_args->ne02 + 31u) / 32u; + const NSUInteger work_item_bytes = 2u * sizeof(uint32_t); + if (work_cap > NSUIntegerMax || + work_offset > NSUIntegerMax - 8u || + (NSUInteger)work_cap > + (NSUIntegerMax - work_offset - 8u) / work_item_bytes || + g_moe_id_map_bytes < + work_offset + 8u + (NSUInteger)work_cap * work_item_bytes) { return 0; } @@ -28128,13 +31217,14 @@ static int ds4_gpu_encode_mul_mm_id_mapped_tile( [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; [enc setBuffer:dst offset:dst_off atIndex:5]; + [enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:6]; if (use_resource_hints) { [enc useResource:src0 usage:MTLResourceUsageRead]; } [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)work_cap, ((NSUInteger)mm_args->ne0 + 63u) / 64u, - (NSUInteger)mm_args->ne02) + 1) threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); return 1; @@ -28203,6 +31293,7 @@ static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16( id cb, id pipeline, + bool compact_tile, const ds4_gpu_mul_mm_id_args *mm_args, const ds4_gpu_dsv4_moe_swiglu_weight_args *act_args, id gate_src0, @@ -28225,8 +31316,21 @@ static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16( const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); - if (tpe_bytes > NSUIntegerMax - hids_bytes || - g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + if (tpe_bytes > NSUIntegerMax - hids_bytes) { + return 0; + } + const NSUInteger work_offset = (tpe_bytes + hids_bytes + 7u) & ~7u; + const uint64_t pair_rows = + (uint64_t)(uint32_t)mm_args->ne20 * (uint32_t)mm_args->ne21; + const uint64_t work_cap = + (pair_rows + 31u * (uint32_t)mm_args->ne02 + 31u) / 32u; + const NSUInteger work_item_bytes = 2u * sizeof(uint32_t); + if (work_cap > NSUIntegerMax || + work_offset > NSUIntegerMax - 8u || + (NSUInteger)work_cap > + (NSUIntegerMax - work_offset - 8u) / work_item_bytes || + g_moe_id_map_bytes < + work_offset + 8u + (NSUInteger)work_cap * work_item_bytes) { return 0; } @@ -28241,11 +31345,13 @@ static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16( [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:6]; [enc setBuffer:mid offset:mid_off atIndex:7]; [enc setBuffer:weights offset:weights_off atIndex:8]; - [enc setThreadgroupMemoryLength:16384u atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + 31u) / 32u, - ((NSUInteger)mm_args->ne0 + 63u) / 64u, - (NSUInteger)mm_args->ne02) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + [enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:9]; + const NSUInteger tile_m = compact_tile ? 32u : 64u; + [enc setThreadgroupMemoryLength:compact_tile ? 8192u : 16384u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)work_cap, + ((NSUInteger)mm_args->ne0 + tile_m - 1u) / tile_m, + 1) + threadsPerThreadgroup:MTLSizeMake(compact_tile ? 64u : 128u, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); return 1; } @@ -28272,7 +31378,7 @@ static int ds4_gpu_encode_mul_mm_id_mapped( 8192u); } -static int ds4_gpu_encode_attn_out_low_q8_mpp( +static int ds4_gpu_encode_attn_out_low_mpp( id cb, id pipeline, const ds4_gpu_mul_mm_id_args *mm_args, @@ -28759,38 +31865,56 @@ static int ds4_gpu_encode_router_select( if (flash_router_fast_path && !g_quality_mode && n_tokens == 1 && getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) { - const bool force_simd_weights_fusion = - getenv("DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION") != NULL; + const bool pre_m5_device = + ds4_gpu_device_is_pre_m5_apple_silicon(); const bool use_simd_finalize = !hash_mode && g_dsv4_router_finalize_one_simd_pipeline != nil && g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u && g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && - (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || - getenv("DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE") != NULL || - force_simd_weights_fusion) && - getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE") == NULL; + (pre_m5_device || + ds4_gpu_device_name_contains("M5")) && + (!pre_m5_device || + getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_SIMD_FINALIZE") == NULL); const bool use_simd_weights_fusion = use_simd_finalize && g_dsv4_router_finalize_weights_one_simd_pipeline != nil && g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u && g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && - (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || - force_simd_weights_fusion) && - getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION") == NULL; + (pre_m5_device || + ds4_gpu_device_name_contains("M5")) && + (!pre_m5_device || + getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_SIMD_WEIGHTS_FUSION") == NULL); + const bool use_pre_m5_transform_finalize_fusion_default = + !g_ssd_streaming_mode && + (pre_m5_device || ds4_gpu_device_name_contains("M5")) && + getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_TRANSFORM_FINALIZE_FUSION") == NULL; + const bool use_transform_finalize_fusion = + use_simd_weights_fusion && + !hash_mode && + use_pre_m5_transform_finalize_fusion_default && + (pre_m5_device || + ds4_gpu_device_name_contains("M5")) && + g_dsv4_router_transform_finalize_weights_one_simd_pipeline != nil && + g_dsv4_router_transform_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_transform_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u; id softplus_sqrt_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, - "kernel_dsv4_softplus_sqrt_f32_4"); + use_transform_finalize_fusion + ? nil + : ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, + "kernel_dsv4_softplus_sqrt_f32_4"); id router_finalize_pipeline = ds4_gpu_hot_pipeline( - use_simd_weights_fusion + use_transform_finalize_fusion + ? g_dsv4_router_transform_finalize_weights_one_simd_pipeline + : use_simd_weights_fusion ? g_dsv4_router_finalize_weights_one_simd_pipeline : use_simd_finalize ? g_dsv4_router_finalize_one_simd_pipeline : g_dsv4_router_finalize_one_pipeline, - use_simd_weights_fusion + use_transform_finalize_fusion + ? "kernel_dsv4_router_transform_finalize_weights_one_simd" + : use_simd_weights_fusion ? "kernel_dsv4_router_finalize_weights_one_simd" : use_simd_finalize ? "kernel_dsv4_router_finalize_one_simd" @@ -28799,21 +31923,24 @@ static int ds4_gpu_encode_router_select( ? nil : ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, "kernel_dsv4_router_weights_one"); - if (!softplus_sqrt_pipeline || !router_finalize_pipeline || + if ((!use_transform_finalize_fusion && !softplus_sqrt_pipeline) || + !router_finalize_pipeline || (!use_simd_weights_fusion && !router_weights_pipeline)) return 0; - ok = ds4_gpu_encode_unary_f32_rows(cb, - softplus_sqrt_pipeline, - logitsbuf, - logits_off, - probsbuf, - probs_off, - n_expert, - 1, - 1, - 0.0f, - 0.0f); - if (!ok) return 0; + if (!use_transform_finalize_fusion) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + softplus_sqrt_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + 1, + 1, + 0.0f, + 0.0f); + if (!ok) return 0; + } const bool use_token_buffer = single_token == NULL; ds4_gpu_dsv4_router_select_one_args args = { @@ -28835,25 +31962,43 @@ static int ds4_gpu_encode_router_select( id enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:router_finalize_pipeline]; [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:probsbuf offset:probs_off atIndex:1]; - if (has_bias) { - [enc setBuffer:biasbuf offset:bias_off atIndex:2]; - } else { - [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2]; - } - if (hash_mode) { - [enc setBuffer:hashbuf offset:hash_off atIndex:3]; - } else { - [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3]; - } - if (use_token_buffer) { - [enc setBuffer:tokensbuf offset:tokens_off atIndex:4]; - } else { + if (use_transform_finalize_fusion) { + [enc setBuffer:logitsbuf offset:logits_off atIndex:1]; + [enc setBuffer:probsbuf offset:probs_off atIndex:2]; + if (has_bias) { + [enc setBuffer:biasbuf offset:bias_off atIndex:3]; + } else { + [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:3]; + } [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4]; - } - [enc setBuffer:selectedbuf offset:selected_off atIndex:5]; - if (use_simd_weights_fusion) { - [enc setBuffer:weightsbuf offset:weights_off atIndex:6]; + if (use_token_buffer) { + [enc setBuffer:tokensbuf offset:tokens_off atIndex:5]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:5]; + } + [enc setBuffer:selectedbuf offset:selected_off atIndex:6]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:7]; + } else { + [enc setBuffer:probsbuf offset:probs_off atIndex:1]; + if (has_bias) { + [enc setBuffer:biasbuf offset:bias_off atIndex:2]; + } else { + [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2]; + } + if (hash_mode) { + [enc setBuffer:hashbuf offset:hash_off atIndex:3]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3]; + } + if (use_token_buffer) { + [enc setBuffer:tokensbuf offset:tokens_off atIndex:4]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4]; + } + [enc setBuffer:selectedbuf offset:selected_off atIndex:5]; + if (use_simd_weights_fusion) { + [enc setBuffer:weightsbuf offset:weights_off atIndex:6]; + } } const NSUInteger router_finalize_scratch_bytes = use_simd_finalize ? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t)) @@ -28970,9 +32115,8 @@ static int ds4_gpu_encode_router_select( const bool use_batch_weights_fusion = flash_router_fast_path && !g_quality_mode && n_tokens > 1u && g_dsv4_router_weights_batch_pipeline != nil && - (ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION") != NULL) && - getenv("DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION") == NULL && + ds4_gpu_device_name_contains("M3") && + getenv("DS4_METAL_DISABLE_ROUTER_WEIGHTS_BATCH_FUSION") == NULL && getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL; if (use_batch_weights_fusion) { const float scale = expert_weight_scale; @@ -34626,6 +37770,16 @@ int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n return 1; } +/* Single-simdgroup MXFP4 decode MoE policy: the nsg=1 pair_swiglu/sum6 + * variants are the pre-M5 Apple-Silicon default for single-token decode + * (bit-exact; only the row-to-simdgroup mapping changes). Read on the + * dispatch path so A/B harnesses can toggle it after engine creation. */ +static bool ds4_gpu_mxfp4_moe_decode_nsg1_enabled(uint32_t n_tokens) { + if (n_tokens != 1) return false; + return ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_NSG1") == NULL; +} + int ds4_gpu_routed_moe_one_tensor( ds4_gpu_tensor *out, ds4_gpu_tensor *gate, @@ -34655,6 +37809,9 @@ int ds4_gpu_routed_moe_one_tensor( const ds4_gpu_tensor *add_in, uint32_t layer_index, bool force_resident) { + BOOL parallel_ffn_scope + __attribute__((cleanup(ds4_gpu_parallel_ffn_scope_cleanup))) = + g_parallel_q8_pending; if (!g_initialized && !ds4_gpu_init()) return 0; /* TP sharding: only the owned contiguous expert range is mapped, * so bind from the owned base, validate only its bytes, and tell the @@ -34815,14 +37972,105 @@ int ds4_gpu_routed_moe_one_tensor( const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); + int pair_swiglu_nsg = 2; + int down_sum6_nsg = 2; int ok = 1; const bool write_clamped_moe = getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL; + /* The selected nsg=1 encoders dispatch exactly 32 threads. */ + const bool use_mxfp4_moe_decode_tg_multiple = + ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_TG_MULTIPLE") == NULL && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + n_tokens == 1 && + n_expert == 6 && + g_tp_split_world == 1 && + add_in == NULL && + !g_quality_mode && + !write_clamped_moe && + (force_resident || !g_ssd_streaming_mode); + const bool use_mxfp4_moe_decode_fixed_route_pair = + use_mxfp4_moe_decode_tg_multiple && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_FIXED_ROUTE_PAIR") == NULL && + expert_in_dim == 4096 && + expert_mid_dim == 2048 && + n_total_expert == 256 && + gate_row_bytes == 2176 && + gate_args.nei0 == 6 && + gate_args.nei1 == 1 && + gate_args.ne11 == 1 && + gate_args.tp_world == 1 && + gate_args.tp_expert_base == 0 && + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 != nil; + const bool use_mxfp4_moe_decode_fixed_route_sum6 = + use_mxfp4_moe_decode_fixed_route_pair && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_FIXED_ROUTE_SUM6") == NULL && + out_dim == 4096 && + down_row_bytes == 1088 && + down_expert_bytes == 4456448 && + down_args.nei0 == 6 && + down_args.nei1 == 1 && + down_args.ne11 == 6 && + down_args.tp_world == 1 && + down_args.tp_expert_base == 0 && + !down_args.tp_addend && + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 != nil; + const bool use_mxfp4_moe_decode_sum6_full_rows = + use_mxfp4_moe_decode_fixed_route_sum6 && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_SUM6_FULL_ROWS") == NULL && + down_args.ne00 == 2048 && + down_args.ne01 == 4096 && + down_args.ne0 == 4096 && + down_args.nr0 == 2 && + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 != nil; + /* Compile-time trip counts for the one-token fixed-route decode MoE + * kernels. Every shape the specialized kernels assume is proven above + * (gate ne00 4096 / row bytes 2176; down ne00 2048 / row bytes 1088 / + * nei0 6), so the specialization only removes runtime loop bounds; the + * per-lane K walk and accumulate order stay byte-identical. */ + const bool use_mxfp4_moe_decode_static_trip_pair = + use_mxfp4_moe_decode_fixed_route_pair && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_STATIC_TRIP") == NULL && + gate_args.ne00 == 4096 && + gate_args.nb01 == 2176 && + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 != nil; + const bool use_mxfp4_moe_decode_static_trip_down = + use_mxfp4_moe_decode_sum6_full_rows && + use_mxfp4_moe_decode_static_trip_pair && + down_args.nb01 == 1088 && + down_args.nei0 == 6 && + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 != nil; id pair_swiglu_pipeline = nil; if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; + } else if (gate_type == DS4_METAL_TENSOR_MXFP4) { + pair_swiglu_pipeline = g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline; + if (ds4_gpu_mxfp4_moe_decode_nsg1_enabled(n_tokens)) { + pair_swiglu_pipeline = + use_mxfp4_moe_decode_static_trip_pair ? + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 : + use_mxfp4_moe_decode_fixed_route_pair ? + g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 : + (use_mxfp4_moe_decode_tg_multiple ? + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple : + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1); + pair_swiglu_nsg = 1; + } + } + const bool use_iq2_pair_pack2 = + gate_type == DS4_METAL_TENSOR_IQ2_XXS && n_tokens == 1 && + n_expert == 6 && n_total_expert == 256 && + expert_in_dim == 4096 && expert_mid_dim == 2048 && + getenv("DS4_METAL_DISABLE_M5_IQ2_PAIR_PACK2") == NULL && + ds4_gpu_device_is_m5_apple_silicon(); + if (use_iq2_pair_pack2 && + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline) { + pair_swiglu_pipeline = + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline; + pair_swiglu_nsg = 4; } const bool fuse_pair_swiglu = !g_quality_mode && @@ -34834,6 +38082,21 @@ int ds4_gpu_routed_moe_one_tensor( down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; } else if (down_type == DS4_METAL_TENSOR_Q4_K) { down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_MXFP4) { + down_sum6_pipeline = g_moe_mul_mv_id_mxfp4_sum6_pipeline; + if (ds4_gpu_mxfp4_moe_decode_nsg1_enabled(n_tokens)) { + down_sum6_pipeline = + use_mxfp4_moe_decode_static_trip_down ? + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 : + use_mxfp4_moe_decode_sum6_full_rows ? + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 : + (use_mxfp4_moe_decode_fixed_route_sum6 ? + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 : + (use_mxfp4_moe_decode_tg_multiple ? + g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple : + g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1)); + down_sum6_nsg = 1; + } } else if (down_type == DS4_METAL_TENSOR_IQ2_XXS && g_tp_split_world == 2) { /* IQ2 down-sum exists for the GLM TP resident split only; the @@ -34846,6 +38109,33 @@ int ds4_gpu_routed_moe_one_tensor( (n_expert == 6 || (n_expert == 8 && g_tp_split_world == 2)) && n_tokens == 1 && down_sum6_pipeline != nil; + + if (g_parallel_q8_pending) { + /* A concurrent encoder invalidates every implicit dependency in + * this generic function. Admit only the one path whose dependency + * graph is explicitly bracketed below: one fused resident IQ2 + * pair-SwiGLU producer and one direct resident Q2 top-6 consumer. */ + const bool narrow_parallel_route = + g_parallel_ffn_mode == 2 && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_Q2_K && + n_tokens == 1 && n_expert == 6 && n_total_expert == 256 && + expert_in_dim == 4096 && expert_mid_dim == 2048 && + out_dim == 4096 && gate_row_bytes == 1056 && + gate_expert_bytes == 2162688 && down_row_bytes == 672 && + down_expert_bytes == 2752512 && + first_expert == 0 && n_bind_expert == n_total_expert && + g_tp_split_rank == 0 && g_tp_split_world == 1 && + add_in == NULL && (force_resident || !g_ssd_streaming_mode) && + !write_clamped_moe && fuse_pair_swiglu && direct_down_sum; + if (!narrow_parallel_route) { + fprintf(stderr, + "ds4: concurrent FFN requires resident fused " + "IQ2 pair-SwiGLU + direct Q2 top-6 sum\n"); + return 0; + } + } + /* The expert-ownership split lives only in the fused id pair+sum6 * kernels; every other routed variant would silently compute full * sums on both ranks and double the combine. Fail fast instead. */ @@ -35110,6 +38400,19 @@ int ds4_gpu_routed_moe_one_tensor( g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline != nil && g_moe_mul_mv_slots6_q2_k_sum6_pipeline != nil && getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL; + const bool use_mxfp4_selected_slots = + !force_resident && + g_ssd_streaming_mode && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + n_expert == 6 && + n_tokens == 1 && + n_total_expert >= 128 && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_mxfp4_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_MXFP4_SELECTED_EXPERT_VIEWS") == NULL; const bool use_iq2_stream_addr_table = !force_resident && g_ssd_streaming_mode && @@ -35122,13 +38425,16 @@ int ds4_gpu_routed_moe_one_tensor( g_moe_mul_mv_addr_iq2_xxs_pipeline != nil && getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL; const bool use_selected_slots = - use_q4_selected_slots || use_iq2_selected_slots || use_iq2_stream_addr_table; + use_q4_selected_slots || use_iq2_selected_slots || + use_mxfp4_selected_slots || use_iq2_stream_addr_table; id slots_pair_swiglu_pipeline = use_iq2_selected_slots ? g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline : - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; + (use_mxfp4_selected_slots ? g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline : + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline); id slots_sum6_pipeline = use_iq2_selected_slots ? g_moe_mul_mv_slots6_q2_k_sum6_pipeline : - g_moe_mul_mv_slots6_q4_k_sum6_pipeline; + (use_mxfp4_selected_slots ? g_moe_mul_mv_slots6_mxfp4_sum6_pipeline : + g_moe_mul_mv_slots6_q4_k_sum6_pipeline); const char *selected_profile_env = getenv("DS4_METAL_SELECTED_PROFILE"); if (!selected_profile_env) { selected_profile_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE"); @@ -35465,7 +38771,8 @@ int ds4_gpu_routed_moe_one_tensor( g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; use_stream_expert_cache = !use_iq2_full_expert_addr_table && - (use_iq2_selected_slots || use_iq2_stream_addr_table || use_q4_selected_slots) && + (use_iq2_selected_slots || use_iq2_stream_addr_table || + use_q4_selected_slots || use_mxfp4_selected_slots) && stream_expert_cache_size_known && ds4_gpu_stream_expert_cache_effective_cap(layer_index, n_total_expert, @@ -35923,7 +39230,8 @@ int ds4_gpu_routed_moe_one_tensor( g_stream_expert_cache_evictions - selected_cache_evictions0; const char *selected_path = use_iq2_stream_addr_table ? "iq2/iq2" : - (use_iq2_selected_slots ? "iq2/q2" : "q4/q4"); + (use_iq2_selected_slots ? "iq2/q2" : + (use_mxfp4_selected_slots ? "mxfp4/mxfp4" : "q4/q4")); const char *selected_view_mode = use_stream_expert_split_deferred ? "stream-split" : use_stream_expert_masked_addr_table ? "stream-addr-mask" : @@ -36001,6 +39309,7 @@ int ds4_gpu_routed_moe_one_tensor( use_stream_expert_masked_addr_table ? "iq2_stream_addr_mask_pair_swiglu" : use_stream_expert_addr_table ? "iq2_stream_addr_pair_swiglu" : use_iq2_selected_slots ? "iq2_slots6_pair_swiglu" : + use_mxfp4_selected_slots ? "mxfp4_slots6_pair_swiglu" : use_q4_selected_slots ? "q4_slots6_pair_swiglu" : (fuse_pair_swiglu ? "pair_swiglu" : ((!g_quality_mode && @@ -36689,7 +39998,7 @@ int ds4_gpu_routed_moe_one_tensor( weightsbuf, ds4_gpu_tensor_offset(weights), gate_smem, - 2, + pair_swiglu_nsg, false); } else if (!g_quality_mode && gate_type == DS4_METAL_TENSOR_IQ2_XXS && @@ -36780,6 +40089,9 @@ int ds4_gpu_routed_moe_one_tensor( false); } DS4_METAL_PROFILE_MOE_ONE_STAGE("activation_weight"); + if (ok && g_parallel_q8_pending) { + ok = ds4_gpu_parallel_q8_matvec_encode_pending(cb, midbuf) != 0; + } id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : @@ -36999,7 +40311,7 @@ int ds4_gpu_routed_moe_one_tensor( add_in ? ds4_gpu_tensor_buffer(add_in) : nil, add_in ? ds4_gpu_tensor_offset(add_in) : 0, down_smem, - 2); + down_sum6_nsg); } else if (ok) { ok = ds4_gpu_encode_mul_mv_id(cb, down_mv_pipeline, @@ -37028,6 +40340,9 @@ int ds4_gpu_routed_moe_one_tensor( n_tokens); } DS4_METAL_PROFILE_MOE_ONE_STAGE("sum"); + if (ok && g_parallel_ffn_mode == 2 && g_parallel_ffn_stage == 1) { + ok = ds4_gpu_parallel_ffn_encode_second_stage(cb) != 0; + } if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34395); return 0; } if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34397); return 0; } @@ -37040,6 +40355,7 @@ int ds4_gpu_routed_moe_one_tensor( #undef DS4_METAL_PROFILE_MOE_ONE_STAGE } + parallel_ffn_scope = NO; return 1; } @@ -37187,7 +40503,21 @@ int ds4_gpu_routed_moe_batch_tensor( can_single_token_q4_expert_address_table || can_single_token_q4_expert_table || can_single_token_q4_selected_slots); - if (use_single_token_q4_one_tensor) { + const bool use_single_token_mxfp4_one_tensor = + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + n_tokens == 1 && + n_expert == 6 && + n_total_expert >= 128 && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline != nil && + g_moe_mul_mv_id_mxfp4_sum6_pipeline != nil && + (!g_ssd_streaming_mode || + (g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_mxfp4_sum6_pipeline != nil)); + if (use_single_token_q4_one_tensor || use_single_token_mxfp4_one_tensor) { if (mid_is_f16) *mid_is_f16 = false; return ds4_gpu_routed_moe_one_tensor(out, gate, @@ -37367,12 +40697,18 @@ int ds4_gpu_routed_moe_batch_tensor( !use_q4_batch_expert_table && !use_mm_id && ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || - (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline)); + (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline) || + (gate_type == DS4_METAL_TENSOR_MXFP4 && + g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline && + getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL)); id tiny_pair_swiglu_pipeline = nil; if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; + } else if (gate_type == DS4_METAL_TENSOR_MXFP4) { + tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline; } const bool use_tiny_pair_swiglu = use_tiny_pair_mv && @@ -37394,24 +40730,114 @@ int ds4_gpu_routed_moe_batch_tensor( !use_q4_batch_expert_table && !use_iq2_batch_selected_addr; /* - * Fused gate+up grouped matmul with the SwiGLU epilogue. The IQ2 - * variant stays opt-in behind its env flag; the Q4_K variant is the - * default path — same MMA accumulation order and epilogue math as the - * separate GEMMs + swiglu pass, so the mid tensor is bit-identical. + * Fused gate+up grouped matmul with the SwiGLU epilogue. Both IQ2 and + * Q4_K use the compact expert work list, the same MMA accumulation + * order, and the same epilogue math as separate GEMMs + SwiGLU, so the + * mid tensor is bit-identical. */ const bool use_mm_id_pair_swiglu = use_mm_id && + !(gate_type == DS4_METAL_TENSOR_IQ2_XXS && + (ds4_gpu_routed_mm_mpp_mask() & 3) == 3) && g_tp_split_world != 2 && /* pair-swiglu mm kernel lacks expert ownership */ request_mid_f16 && n_expert == 6 && ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && - down_type == DS4_METAL_TENSOR_Q2_K && - getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL) || + down_type == DS4_METAL_TENSOR_Q2_K) || (gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K)) && + down_type == DS4_METAL_TENSOR_Q4_K) || + (gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4)) && getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL && getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL; + /* + * The MXFP4 32x32 specialization uses two SIMDgroups and 8 KiB of + * threadgroup memory, and exactly culls SIMDgroup 1 on at-most-16-row + * final work items. Balanced full-model A/B at both 2K and 8K makes + * it the resident pre-M5 default for large prefill. + */ + const bool use_pre_m5_mxfp4_mm_id_pair_swiglu_compact_tile_default = + ds4_gpu_device_is_pre_m5_apple_silicon() && + !g_ssd_streaming_mode && + n_tokens >= 2048u && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_PAIR_SWIGLU_COMPACT_TILE") == NULL; + const bool use_mxfp4_mm_id_pair_swiglu_compact_tile = + use_mm_id_pair_swiglu && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + g_tp_split_world == 1 && + (use_pre_m5_mxfp4_mm_id_pair_swiglu_compact_tile_default || + (g_test_flags & DS4_GPU_TEST_MXFP4_PAIR_COMPACT_TILE) != 0u); + /* + * Token-centric map construction for the compact resident pre-M5 MXFP4 + * path. It emits the established work_count/uint2 trailer consumed + * by the existing padded direct launches and changes no arithmetic. + */ + const bool use_pre_m5_mxfp4_mm_id_map_scatter_default = + ds4_gpu_device_is_pre_m5_apple_silicon() && + n_tokens >= 2048u && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_MAP_SCATTER") == NULL; + const bool use_mxfp4_mm_id_map_scatter = + use_mxfp4_mm_id_pair_swiglu_compact_tile && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + n_expert == 6 && + !g_ssd_streaming_mode && + g_tp_split_world == 1 && + (use_pre_m5_mxfp4_mm_id_map_scatter_default || + (g_test_flags & DS4_GPU_TEST_MXFP4_MAP_SCATTER) != 0u); + /* + * On a final occupied 32-row work tile with at most 16 rows, the + * second SIMDgroup pair has no valid MMA output. The specialized + * kernel still keeps every thread in staging, barriers, and the + * scalar SwiGLU epilogue. Keep the automatic path to the resident pre-M5 + * shape and prefixes covered by the full-model A/B gate. + */ + const bool use_pre_m5_mxfp4_mm_id_pair_tail_simdgroup_cull_default = + ds4_gpu_device_is_pre_m5_apple_silicon() && + !g_ssd_streaming_mode && + n_tokens >= 2048u && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_PAIR_TAIL_SIMDGROUP_CULL") == NULL; + const bool use_mxfp4_mm_id_pair_tail_simdgroup_cull = + use_mm_id_pair_swiglu && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + g_tp_split_world == 1 && + (use_pre_m5_mxfp4_mm_id_pair_tail_simdgroup_cull_default || + (g_test_flags & DS4_GPU_TEST_MXFP4_PAIR_TAIL_CULL) != 0u); + const bool use_pre_m5_mxfp4_mm_id_down_tail_simdgroup_cull_default = + ds4_gpu_device_is_pre_m5_apple_silicon() && + !g_ssd_streaming_mode && + n_tokens >= 2048u && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_TAIL_SIMDGROUP_CULL") == NULL; + const bool use_mxfp4_mm_id_down_tail_simdgroup_cull = + use_mm_id && + request_mid_f16 && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + g_tp_split_world == 1 && + (use_pre_m5_mxfp4_mm_id_down_tail_simdgroup_cull_default || + (g_test_flags & DS4_GPU_TEST_MXFP4_DOWN_TAIL_CULL) != 0u); + /* + * Exact half-result dequantization table for the resident MXFP4 down + * projection. Balanced full-model A/B at 2K and 8K makes it the + * pre-M5 Apple-Silicon default for large prefill. + */ + const bool use_pre_m5_mxfp4_mm_id_down_half_lut_default = + ds4_gpu_device_is_pre_m5_apple_silicon() && + n_tokens >= 2048u && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_HALF_LUT") == NULL; + const bool use_mxfp4_mm_id_down_half_lut = + use_mm_id && + request_mid_f16 && + gate_type == DS4_METAL_TENSOR_MXFP4 && + down_type == DS4_METAL_TENSOR_MXFP4 && + n_expert == 6 && + !g_ssd_streaming_mode && + g_tp_split_world == 1 && + (use_pre_m5_mxfp4_mm_id_down_half_lut_default || + (g_test_flags & DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT) != 0u); if (use_mm_id) { gate_map_args = ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); @@ -37431,17 +40857,73 @@ int ds4_gpu_routed_moe_batch_tensor( down_mm_args.tp_world = g_tp_split_world; down_mm_args.tp_expert_base = tp_expert_base_host; - map_pipeline = ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + map_pipeline = ds4_gpu_get_pipeline( + use_mxfp4_mm_id_map_scatter ? + "kernel_mul_mm_id_map_scatter_work_ne20_6" : + ds4_gpu_mul_mm_id_map0_name(n_expert)); gate_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); up_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); - down_mm_pipeline = request_mid_f16 ? - ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) : - ds4_gpu_routed_mm_pipeline(down_type); + down_mm_pipeline = use_mxfp4_mm_id_down_half_lut ? + ds4_gpu_get_mul_mm_id_pipeline( + use_mxfp4_mm_id_down_tail_simdgroup_cull ? + "kernel_mul_mm_id_mxfp4_f16_half_lut_tail_cull" : + "kernel_mul_mm_id_mxfp4_f16_half_lut", + false) : + use_mxfp4_mm_id_down_tail_simdgroup_cull ? + ds4_gpu_get_mul_mm_id_pipeline( + "kernel_mul_mm_id_mxfp4_f16_tail_cull", false) : + request_mid_f16 ? + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) : + ds4_gpu_routed_mm_pipeline(down_type); + const int mpp_mask = ds4_gpu_routed_mm_mpp_mask(); + if (mpp_mask && gate_type == DS4_METAL_TENSOR_IQ2_XXS) { + id mpp = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32_mpp", false); + if (mpp) { + if (mpp_mask & 1) gate_mm_pipeline = mpp; + if (mpp_mask & 2) up_mm_pipeline = mpp; + } + } + if ((mpp_mask & 4) && request_mid_f16 && + (down_type == DS4_METAL_TENSOR_Q2_K || down_type == DS4_METAL_TENSOR_IQ2_XXS)) { + id mpp = ds4_gpu_get_mul_mm_id_pipeline( + down_type == DS4_METAL_TENSOR_Q2_K ? + "kernel_mul_mm_id_q2_K_f16_mpp" : + "kernel_mul_mm_id_iq2_xxs_f16_mpp", false); + if (mpp) down_mm_pipeline = mpp; + } if (use_mm_id_pair_swiglu) { + /* Exact half-domain block scaling for the resident MXFP4 pair + * tile: E8M0 is a power of two and every E2M1 magnitude is + * exact in binary16, so inside the guarded exponent band the + * half product is the identical once-rounded value while the + * per-element f32->f16 narrowing and the 64-byte float table + * disappear. Outside the band the kernel keeps the established + * float arithmetic. */ + const bool use_mxfp4_mm_id_pair_half_scale = + gate_type == DS4_METAL_TENSOR_MXFP4 && + ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MM_ID_PAIR_HALF_SCALE") == NULL && + !g_quality_mode && + !g_ssd_streaming_mode && + g_tp_split_world == 1; pair_swiglu_mm_pipeline = - ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mm_id_q4_K_pair_swiglu_f16" : - "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); + ds4_gpu_get_pipeline( + gate_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mm_id_q4_K_pair_swiglu_f16" : + gate_type == DS4_METAL_TENSOR_MXFP4 ? + (use_mxfp4_mm_id_pair_swiglu_compact_tile ? + (use_mxfp4_mm_id_pair_half_scale ? + "kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull_half_scale" : + "kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull") : + use_mxfp4_mm_id_pair_tail_simdgroup_cull ? + (use_mxfp4_mm_id_pair_half_scale ? + "kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull_half_scale" : + "kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull") : + (use_mxfp4_mm_id_pair_half_scale ? + "kernel_mul_mm_id_mxfp4_pair_swiglu_f16_half_scale" : + "kernel_mul_mm_id_mxfp4_pair_swiglu_f16")) : + "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); } if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline || (use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) { @@ -37691,6 +41173,8 @@ int ds4_gpu_routed_moe_batch_tensor( down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; } else if (down_type == DS4_METAL_TENSOR_Q4_K) { down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_MXFP4) { + down_sum6_pipeline = g_moe_mul_mv_id_mxfp4_sum6_pipeline; } const bool direct_down_sum = !g_quality_mode && @@ -37804,6 +41288,7 @@ int ds4_gpu_routed_moe_batch_tensor( }; ok = ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(cb, pair_swiglu_mm_pipeline, + use_mxfp4_mm_id_pair_swiglu_compact_tile, &gate_mm_args, &act_args, gate_buf, @@ -38365,123 +41850,6 @@ int ds4_gpu_hc_weighted_sum_tensor( "HC weighted sum"); } -int ds4_gpu_hc_weighted_sum_norm_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *norm_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *weights, - const void *model_map, - uint64_t model_size, - uint64_t norm_weight_offset, - uint32_t n_embd, - uint32_t n_hc, - float norm_eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !norm_out || !residual_hc || !weights || !model_map) return 0; - - const bool force = - getenv("DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; - const bool disabled = - getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") != NULL; - const bool require = - getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; - const bool supported_shape = - n_hc == 4u && (n_embd == 4096u || n_embd == 7168u); - const bool auto_shape = - n_embd == 4096u && ds4_gpu_device_name_contains("M3"); - const bool use_fusion = - supported_shape && !g_quality_mode && !disabled && - g_hc_weighted_sum_norm_pipeline != nil && - (auto_shape || force); - if (require && supported_shape && !use_fusion) { - fprintf(stderr, - "ds4: required Metal output HC sum/RMSNorm fusion was not selected\n"); - return 0; - } - if (!use_fusion) return 0; - - @autoreleasepool { - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t residual_bytes = (uint64_t)n_hc * out_row_bytes; - const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float); - id xbuf = ds4_gpu_tensor_buffer(residual_hc); - id wbuf = ds4_gpu_tensor_buffer(weights); - id outbuf = ds4_gpu_tensor_buffer(out); - id normbuf = ds4_gpu_tensor_buffer(norm_out); - if (!xbuf || !wbuf || !outbuf || !normbuf || - ds4_gpu_tensor_bytes(residual_hc) < residual_bytes || - ds4_gpu_tensor_bytes(weights) < weight_bytes || - ds4_gpu_tensor_bytes(out) != out_row_bytes || - ds4_gpu_tensor_bytes(norm_out) < out_row_bytes) { - fprintf(stderr, - "ds4: Metal output HC sum/RMSNorm fusion received invalid activation buffers\n"); - return 0; - } - if (norm_weight_offset > model_size || - out_row_bytes > model_size - norm_weight_offset) { - fprintf(stderr, - "ds4: Metal output HC sum/RMSNorm weight range is outside the mapped model\n"); - return 0; - } - - uint64_t norm_inner = 0; - id normwbuf = ds4_gpu_wrap_model_range( - model_map, model_size, norm_weight_offset, - out_row_bytes, &norm_inner); - if (!normwbuf) return 0; - - ds4_gpu_hc_weighted_sum_norm_args args = { - .n_embd = (int64_t)n_embd, - .n_hc = (int64_t)n_hc, - .n_tokens = 1, - .nb_x0 = sizeof(float), - .nb_x1 = out_row_bytes, - .nb_x2 = residual_bytes, - .nb_w0 = sizeof(float), - .nb_w1 = weight_bytes, - .nb0 = sizeof(float), - .nb1 = out_row_bytes, - .nb_norm1 = out_row_bytes, - .norm_eps = norm_eps, - }; - const NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); - const NSUInteger shared_bytes = - ((NSUInteger)n_embd + 32u) * sizeof(float); - if (nth > g_hc_weighted_sum_norm_pipeline.maxTotalThreadsPerThreadgroup || - shared_bytes > [g_device maxThreadgroupMemoryLength]) { - if (require) { - fprintf(stderr, - "ds4: required Metal output HC sum/RMSNorm fusion exceeds device limits\n"); - } - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_hc_weighted_sum_norm_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; - [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:4]; - [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer( - cb, owned, "output HC sum/RMSNorm fused")) { - return 0; - } - } - - return 1; -} - int ds4_gpu_hc_weighted_sum_split_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, @@ -38762,6 +42130,274 @@ int ds4_gpu_hc_split_weighted_sum_norm_tensor( return 1; } +typedef struct { + int32_t n; + int32_t out_dim; + float eps; +} ds4_gpu_hc_norm_mix_args; + +int ds4_gpu_hc_rms_norm_mix_f16_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + static int cached = -1; + if (cached < 0) { + cached = ds4_gpu_get_pipeline("kernel_dsv4_hc_rms_norm_mix_f16") != nil; + } + return cached; +} + +/* Fused decode HC-pre: unweighted RMSNorm over the flattened HC row plus the + * F16 HC-mix matvec in one dispatch. The kernel reproduces the standalone + * kernels' reduction trees bit-exactly (see metal/dsv4_hc.metal), so this is + * purely a dispatch-count optimization. */ +int ds4_gpu_hc_rms_norm_mix_f16_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + 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; + + @autoreleasepool { + const uint64_t row_bytes = (uint64_t)n * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal fused HC norm/mix weight range is outside the mapped model\n"); + return 0; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < (uint64_t)n * sizeof(float) || + ds4_gpu_tensor_bytes(out) < (uint64_t)out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal fused HC norm/mix received undersized activation buffers\n"); + return 0; + } + + uint64_t weight_inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, + &weight_inner); + if (!wbuf) return 0; + + const bool use_cluster2 = + ds4_gpu_device_is_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_M5_HC_NORM_MIX_CLUSTER2") == NULL; + id pipeline = ds4_gpu_get_pipeline( + use_cluster2 ? "kernel_dsv4_hc_rms_norm_mix_f16_cluster2" : + "kernel_dsv4_hc_rms_norm_mix_f16"); + if (!pipeline || + (use_cluster2 && pipeline.maxTotalThreadsPerThreadgroup < 512u)) { + return 0; + } + const NSUInteger shared_floats = + use_cluster2 ? 32u + 32u * 2u * 2u : 32u + 32u * 2u; + + ds4_gpu_hc_norm_mix_args args = { + .n = (int32_t)n, + .out_dim = (int32_t)out_dim, + .eps = eps, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)weight_inner atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:shared_floats * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups: + MTLSizeMake(use_cluster2 ? (out_dim + 3u) / 4u : + (out_dim + 1u) / 2u, + 1, 1) + threadsPerThreadgroup: + MTLSizeMake(32, use_cluster2 ? 16 : 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "fused HC norm/mix")) return 0; + } + + return 1; +} + + +int ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( + ds4_gpu_tensor *mix, + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t mix_weight_offset, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n, + uint32_t mix_dim, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float hc_eps, + float norm_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!mix || !out || !norm_out || !split || !residual_hc || !model_map || + n != 16384u || mix_dim != 24u || n_embd != 4096u || n_hc != 4u) { + return 0; + } + + @autoreleasepool { + const uint64_t mix_weight_bytes = + (uint64_t)n * (uint64_t)mix_dim * sizeof(uint16_t); + const uint64_t mix_bytes = (uint64_t)mix_dim * sizeof(float); + const uint64_t residual_bytes = (uint64_t)n * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t scale_bytes = 3u * sizeof(float); + if (mix_weight_offset > model_size || + mix_weight_bytes > model_size - mix_weight_offset || + scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset || + norm_weight_offset > model_size || + out_bytes > model_size - norm_weight_offset) { + fprintf(stderr, + "ds4: Metal compound HC producer parameter range is outside the mapped model\n"); + return 0; + } + + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id mixbuf = ds4_gpu_tensor_buffer(mix); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + if (!xbuf || !mixbuf || !splitbuf || !outbuf || !normbuf || + ds4_gpu_tensor_bytes(residual_hc) < residual_bytes || + ds4_gpu_tensor_bytes(mix) < mix_bytes || + ds4_gpu_tensor_bytes(split) < mix_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(norm_out) < out_bytes) { + fprintf(stderr, + "ds4: Metal compound HC producer received undersized activation buffers\n"); + return 0; + } + + uint64_t mix_weight_inner = 0; + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + uint64_t norm_inner = 0; + id mix_weight = ds4_gpu_wrap_model_range( + model_map, model_size, mix_weight_offset, mix_weight_bytes, + &mix_weight_inner); + id scalebuf = ds4_gpu_wrap_model_range( + model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range( + model_map, model_size, base_offset, mix_bytes, &base_inner); + id norm_weight = ds4_gpu_wrap_model_range( + model_map, model_size, norm_weight_offset, out_bytes, &norm_inner); + if (!mix_weight || !scalebuf || !basebuf || !norm_weight) return 0; + + if (!g_dsv4_hc_producer_pre_norm_pipeline) { + g_dsv4_hc_producer_pre_norm_pipeline = ds4_gpu_get_pipeline( + "kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm"); + } + id producer = + g_dsv4_hc_producer_pre_norm_pipeline; + if (!producer || producer.maxTotalThreadsPerThreadgroup < 512u) { + return 0; + } + if (!g_dsv4_completion_cache) return 0; + const NSUInteger mix_offset = ds4_gpu_tensor_offset(mix); + id completion = nil; + if (g_dsv4_hc_producer_last_mix_buffer == mixbuf && + g_dsv4_hc_producer_last_mix_offset == mix_offset) { + completion = g_dsv4_hc_producer_last_completion; + } else { + NSString *completion_key = [NSString stringWithFormat:@"%p:%llu", + (void *)mixbuf, (unsigned long long)mix_offset]; + 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_dsv4_hc_producer_last_mix_buffer = mixbuf; + g_dsv4_hc_producer_last_mix_offset = mix_offset; + g_dsv4_hc_producer_last_completion = completion; + } + [g_transient_buffers addObject:completion]; + + ds4_gpu_hc_norm_mix_args mix_args = { + .n = (int32_t)n, + .out_dim = (int32_t)mix_dim, + .eps = eps, + }; + ds4_gpu_hc_split_weighted_sum_norm_args split_args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = 1, + .mix_hc = (int64_t)mix_dim, + .nb_mix1 = mix_bytes, + .nb_split1 = mix_bytes, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = residual_bytes, + .nb0 = sizeof(float), + .nb1 = out_bytes, + .nb_norm1 = out_bytes, + .eps = hc_eps, + .norm_eps = norm_eps, + }; + + const NSUInteger shared_floats = 32u + 4u * 32u + 4u + 32u; + const NSUInteger shared_bytes = shared_floats * sizeof(float); + const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength]; + if (max_shared != 0 && shared_bytes > max_shared) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + + [enc setComputePipelineState:producer]; + [enc setBytes:&mix_args length:sizeof(mix_args) atIndex:0]; + [enc setBytes:&split_args length:sizeof(split_args) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:mix_weight offset:(NSUInteger)mix_weight_inner atIndex:3]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:4]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:5]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:6]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:7]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:8]; + [enc setBuffer:norm_weight offset:(NSUInteger)norm_inner atIndex:9]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:10]; + [enc setBuffer:completion offset:0 atIndex:11]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(6, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 16, 1)]; + + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "compound HC producer/pre-norm")) { + return -1; + } + } + return 1; +} + int ds4_gpu_output_hc_weights_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *pre, @@ -38814,17 +42450,14 @@ int ds4_gpu_output_hc_weights_tensor( &base_inner); if (!scalebuf || !basebuf) return 0; - const bool force_weights4 = - getenv("DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4") != NULL; - const bool disable_weights4 = - getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4") != NULL; const bool require_weights4 = getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL; const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u; const bool use_weights4 = - weights4_shape && !g_quality_mode && !disable_weights4 && + weights4_shape && !g_quality_mode && g_output_hc_weights4_pipeline != nil && - (ds4_gpu_device_name_contains("M3") || force_weights4) && + (ds4_gpu_device_name_contains("M3") || + (g_test_flags & DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4) != 0u) && g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u; if (require_weights4 && weights4_shape && !use_weights4) { fprintf(stderr, @@ -39606,9 +43239,20 @@ int ds4_gpu_matmul_q8_0_hc_expand_tensor( .has_add = 0, }; + const NSUInteger split_offset = ds4_gpu_tensor_offset(split); + const bool vec_hc = + (split_offset & 15u) == 0u && + getenv("DS4_METAL_DISABLE_M5_Q8_HC_VEC") == NULL && + ds4_gpu_device_is_m5_apple_silicon(); id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_q8_hc_expand4_q8_0", - mv_dispatch.nsg); + ds4_gpu_get_mul_mv_pipeline( + vec_hc ? "kernel_dsv4_q8_hc_expand4_q8_0_vec_hc" : + "kernel_dsv4_q8_hc_expand4_q8_0", + mv_dispatch.nsg); + if (!pipeline && vec_hc) { + pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_dsv4_q8_hc_expand4_q8_0", mv_dispatch.nsg); + } if (!pipeline) return 0; int owned = 0; @@ -39623,8 +43267,8 @@ int ds4_gpu_matmul_q8_0_hc_expand_tensor( [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:4]; [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:5]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:6]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7]; + [enc setBuffer:splitbuf offset:split_offset + (NSUInteger)n_hc * sizeof(float) atIndex:6]; + [enc setBuffer:splitbuf offset:split_offset + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7]; [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:8]; [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / diff --git a/src/app.rs b/src/app.rs index d441025..e0f4b0c 100644 --- a/src/app.rs +++ b/src/app.rs @@ -398,6 +398,7 @@ pub(crate) enum Message { PreferenceGlmMtpTimingChanged(bool), PreferenceDsparkConfidenceChanged(String), PreferenceDsparkStrictChanged(bool), + PreferenceDsparkExactSamplingChanged(bool), PreferenceSsdChanged(bool), PreferenceSsdColdChanged(bool), PreferenceSsdCacheChanged(String), @@ -2719,7 +2720,7 @@ mod tests { Some(Message::FocusPrevious) )); assert!(ModelChoice::DeepSeekV4Flash.supports_dspark()); - assert!(!ModelChoice::DeepSeekV4Flash0731.supports_dspark()); + assert!(ModelChoice::DeepSeekV4Flash0731.supports_dspark()); assert!(!ModelChoice::DeepSeekV4Pro.supports_dspark()); assert!(!ModelChoice::Glm52.supports_dspark()); } diff --git a/src/app/preferences.rs b/src/app/preferences.rs index 45074ab..5e39e35 100644 --- a/src/app/preferences.rs +++ b/src/app/preferences.rs @@ -45,6 +45,7 @@ pub(super) struct PreferenceDraft { pub(super) glm_mtp_timing: bool, pub(super) dspark_confidence_threshold: String, pub(super) dspark_strict: bool, + pub(super) dspark_exact_sampling: bool, pub(super) ssd_streaming: bool, pub(super) ssd_streaming_cold: bool, pub(super) ssd_cache: String, @@ -111,6 +112,7 @@ impl PreferenceDraft { glm_mtp_timing: speculative.glm_mtp_timing, dspark_confidence_threshold: optional_string(speculative.dspark_confidence_threshold), dspark_strict: speculative.dspark_strict, + dspark_exact_sampling: speculative.dspark_exact_sampling, ssd_streaming: runtime.ssd.enabled, ssd_streaming_cold: runtime.ssd.cold, ssd_cache: optional_string(runtime.ssd.cache), @@ -226,6 +228,7 @@ impl PreferenceDraft { &self.dspark_confidence_threshold, )?, dspark_strict: self.dspark_strict, + dspark_exact_sampling: self.dspark_exact_sampling, }) } @@ -267,6 +270,7 @@ impl PreferenceDraft { self.glm_mtp_timing = speculative.glm_mtp_timing; self.dspark_confidence_threshold = optional_string(speculative.dspark_confidence_threshold); self.dspark_strict = speculative.dspark_strict; + self.dspark_exact_sampling = speculative.dspark_exact_sampling; self.ssd_streaming = ssd.enabled; self.ssd_streaming_cold = ssd.cold; self.ssd_cache = optional_string(ssd.cache); @@ -693,12 +697,16 @@ impl App { self.preference_error = None; } Message::PreferenceLegacyMtpChanged(enabled) => { - self.preference_draft.legacy_mtp_enabled = - self.preference_draft.acceleration_model.supports_dspark() && enabled; + self.preference_draft.legacy_mtp_enabled = self + .preference_draft + .acceleration_model + .supports_legacy_mtp() + && enabled; if self.preference_draft.legacy_mtp_enabled { self.preference_draft.dspark_enabled = false; self.preference_draft.dspark_confidence_threshold.clear(); self.preference_draft.dspark_strict = false; + self.preference_draft.dspark_exact_sampling = false; } self.preference_error = None; } @@ -708,6 +716,7 @@ impl App { if !self.preference_draft.dspark_enabled { self.preference_draft.dspark_confidence_threshold.clear(); self.preference_draft.dspark_strict = false; + self.preference_draft.dspark_exact_sampling = false; } else { self.preference_draft.legacy_mtp_enabled = false; } @@ -903,6 +912,15 @@ impl App { } self.preference_error = None; } + Message::PreferenceDsparkExactSamplingChanged(value) => { + self.preference_draft.dspark_exact_sampling = + self.preference_draft.acceleration_model.supports_dspark() && value; + if self.preference_draft.dspark_exact_sampling { + self.preference_draft.dspark_enabled = true; + self.preference_draft.legacy_mtp_enabled = false; + } + self.preference_error = None; + } Message::PreferenceSsdChanged(value) => { self.preference_draft.ssd_streaming = value; self.preference_error = None; @@ -1011,14 +1029,14 @@ mod tests { assert_eq!(draft.context_tokens, "32768"); draft.context_tokens = "456".into(); draft - .select_generation(ModelChoice::DeepSeekV4Flash, ReasoningMode::High) + .select_generation(ModelChoice::DeepSeekV4Flash0731, ReasoningMode::High) .unwrap(); assert_eq!(draft.context_tokens, "123"); draft.select_acceleration(ModelChoice::Glm52).unwrap(); assert!(!draft.ssd_streaming); draft - .select_acceleration(ModelChoice::DeepSeekV4Flash) + .select_acceleration(ModelChoice::DeepSeekV4Flash0731) .unwrap(); assert!(draft.ssd_streaming); } diff --git a/src/app/view/preferences.rs b/src/app/view/preferences.rs index efd45f9..4d19551 100644 --- a/src/app/view/preferences.rs +++ b/src/app/view/preferences.rs @@ -6,7 +6,7 @@ impl App { let legacy_mtp_toggle: Option Message> = self .preference_draft .acceleration_model - .supports_dspark() + .supports_legacy_mtp() .then_some(Message::PreferenceLegacyMtpChanged); let legacy_mtp = hint( toggle(self.preference_draft.legacy_mtp_enabled) @@ -36,6 +36,11 @@ impl App { .acceleration_model .supports_dspark() .then_some(Message::PreferenceDsparkStrictChanged); + let dspark_exact_toggle: Option Message> = self + .preference_draft + .acceleration_model + .supports_dspark() + .then_some(Message::PreferenceDsparkExactSamplingChanged); let effective = self .preference_draft .effective_for( @@ -74,7 +79,7 @@ impl App { text_input("Automatic", &self.preference_draft.directional_steering_ffn); let mut steering_attn = text_input("0", &self.preference_draft.directional_steering_attn); let mut dspark_confidence = text_input( - "0.9 (DS4 default)", + "0.6 (DS4 default)", &self.preference_draft.dspark_confidence_threshold, ); if self.preference_draft.model != ModelChoice::Glm52 { @@ -509,7 +514,7 @@ impl App { dspark, preference_input_row( "DSpark confidence threshold", - "How sure the draft model must be, from 0 to 1, before its token is handed to the verifier. Lower forwards more guesses for more speed and more rejected work; blank uses DS4's 0.9.", + "How sure the draft model must be, from 0 to 1, before its token is handed to the verifier. Lower forwards more guesses for more speed and more rejected work; blank uses DS4's 0.6, or 0.8 for exact sampling.", dspark_confidence, ), hint( @@ -518,6 +523,12 @@ impl App { .on_toggle_maybe(dspark_strict_toggle), "Lets the draft model only propose, never decide: every token is sampled by the full model. Gives up some of the speedup in exchange for output identical to non-speculative decoding.", ), + hint( + toggle(self.preference_draft.dspark_exact_sampling) + .label("Use exact DSpark sampling") + .on_toggle_maybe(dspark_exact_toggle), + "For non-zero temperatures, applies DS4's exact acceptance and corrected rejection sampling. Off uses the faster opportunistic mode: sample a boundary token, then accept DSpark tokens only while they match the target's greedy path.", + ), text(if self.preference_draft.acceleration_model.supports_dspark() { "Legacy MTP and DSpark use separate managed support artifacts; entering a DSpark threshold or enabling strict mode selects DSpark." } else if self.preference_draft.acceleration_model == ModelChoice::Glm52 { @@ -532,7 +543,7 @@ impl App { |engine| { let settings = engine.speculative; format!( - "Engine: MTP draft {} • margin {} • legacy MTP {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}", + "Engine: MTP draft {} • margin {} • legacy MTP {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {} • exact sampling {}", settings.mtp_draft_tokens, settings.mtp_margin, if self.preference_draft.legacy_mtp_enabled { "on" } else { "off" }, @@ -542,6 +553,7 @@ impl App { settings.dspark_confidence_threshold, if settings.dspark_confidence_threshold_set { " explicit" } else { " default" }, if settings.dspark_strict { "on" } else { "off" }, + if settings.dspark_exact_sampling { "on" } else { "off" }, ) }, )) diff --git a/src/engine.rs b/src/engine.rs index 1a0a69c..b4a2e8d 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -8,7 +8,7 @@ mod validation; #[cfg(target_os = "macos")] use crate::metrics::{KvLookup, Metrics, SsdStats}; -use crate::model::ModelChoice; +use crate::model::{ModelChoice, validate_engine_artifacts}; #[cfg(target_os = "macos")] use crate::settings::TurnSettings; use crate::settings::{EngineSettings, ReasoningMode}; @@ -232,6 +232,12 @@ pub(crate) struct ModelSummary { impl Model { #[allow(dead_code)] pub(crate) fn open(settings: &EngineSettings) -> Result { + validate_engine_artifacts( + settings.model, + settings.artifacts.mtp.is_some() && !settings.speculative.dspark, + settings.speculative.dspark, + &settings.artifacts, + )?; let mut model = Self::open_main(&settings.artifacts.model, settings.model)?; if settings.execution.warm_weights { model.main.warm()?; @@ -1225,8 +1231,17 @@ impl Generator { cancelled, )? } else { - self.executor.eval(token)?; - vec![token] + self.executor.eval_speculative_sampled( + token, + generation_limit - generated_tokens, + settings.reasoning_mode, + settings.temperature, + settings.top_p, + settings.min_p, + settings.top_k, + &mut rng, + cancelled, + )? }; self.publish_execution_stats(); for token in cycle { @@ -1545,12 +1560,30 @@ fn sample( top_k: i32, rng: &mut Rng, ) -> i32 { + let probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k); + sample_probabilities(&probabilities, rng, None) +} + +#[cfg(any(target_os = "macos", test))] +fn sampling_probabilities( + logits: &[f32], + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, +) -> Vec<(usize, f32)> { + let greedy = || { + vec![( + logits + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .map_or(0, |(index, _)| index), + 1.0, + )] + }; if temperature <= 0.0 { - return logits - .iter() - .enumerate() - .max_by(|a, b| a.1.total_cmp(b.1)) - .map_or(0, |(index, _)| index as i32); + return greedy(); } let maximum = logits .iter() @@ -1558,7 +1591,7 @@ fn sample( .filter(|value| value.is_finite()) .fold(f32::NEG_INFINITY, f32::max); if !maximum.is_finite() { - return 0; + return greedy(); } let top_p = if top_p <= 0.0 || top_p > 1.0 { 1.0 @@ -1571,45 +1604,96 @@ fn sample( .enumerate() .filter(|(_, logit)| logit.is_finite()) .map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp())) - .filter(|(_, probability)| *probability >= min_p) .collect(); if probabilities.is_empty() { - return logits - .iter() - .enumerate() - .max_by(|a, b| a.1.total_cmp(b.1)) - .map_or(0, |(index, _)| index as i32); + return greedy(); } - if top_p < 1.0 || top_k > 0 { + if top_p < 1.0 || top_k > 0 || min_p > 0.0 { probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0))); if top_k > 0 { - probabilities.truncate(probabilities.len().min(top_k as usize)); + probabilities.truncate(probabilities.len().min((top_k as usize).min(1024))); } } - if top_p < 1.0 { - let total: f32 = probabilities - .iter() - .map(|(_, probability)| probability) - .sum(); - let mut kept = 0.0; - let count = probabilities - .iter() - .position(|(_, probability)| { - kept += *probability; - kept / total >= top_p - }) - .map_or(probabilities.len(), |index| index + 1); - probabilities.truncate(count); + let total: f32 = probabilities + .iter() + .map(|(_, probability)| probability) + .sum(); + let mut kept = 0.0; + let mut count = 0; + for (_, probability) in &probabilities { + if count > 0 && *probability < min_p { + break; + } + kept += *probability; + count += 1; + if kept / total >= top_p { + break; + } } - let kept_total: f32 = probabilities.iter().map(|(_, p)| p).sum(); - let mut choice = rng.unit() * kept_total; - for (token, probability) in &probabilities { + probabilities.truncate(count); + if probabilities.is_empty() || !kept.is_finite() || kept <= 0.0 { + return greedy(); + } + for (_, probability) in &mut probabilities { + *probability /= kept; + } + if top_p >= 1.0 && top_k <= 0 { + probabilities.sort_unstable_by_key(|(token, _)| *token); + } + probabilities +} + +#[cfg(any(target_os = "macos", test))] +fn sample_probabilities( + probabilities: &[(usize, f32)], + rng: &mut Rng, + excluded: Option, +) -> i32 { + let total: f32 = probabilities + .iter() + .filter(|(token, _)| Some(*token) != excluded) + .map(|(_, probability)| probability) + .sum(); + let mut choice = rng.unit() * total; + for (token, probability) in probabilities { + if Some(*token) == excluded { + continue; + } choice -= probability; if choice <= 0.0 { return *token as i32; } } - probabilities.last().map_or(0, |(token, _)| *token as i32) + probabilities + .iter() + .rev() + .find(|(token, _)| Some(*token) != excluded) + .map_or(0, |(token, _)| *token as i32) +} + +#[cfg(any(target_os = "macos", test))] +fn exact_delta_sample( + logits: &[f32], + draft: i32, + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, + rng: &mut Rng, +) -> (i32, bool) { + let mut probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k); + let draft_probability = probabilities + .iter() + .find(|(token, _)| *token == draft as usize) + .map_or(0.0, |(_, probability)| *probability); + if rng.unit() <= draft_probability { + return (draft, true); + } + probabilities.sort_unstable_by_key(|(token, _)| *token); + ( + sample_probabilities(&probabilities, rng, Some(draft as usize)), + false, + ) } #[cfg(any(target_os = "macos", test))] @@ -1665,6 +1749,37 @@ mod sampling_tests { assert_eq!(chunks, ["hello "]); } + #[test] + fn exact_delta_sampling_accepts_or_corrects_the_draft() { + let mut accept_rng = Rng::new(2); + let (accepted, was_draft) = + exact_delta_sample(&[10.0, 0.0], 0, 1.0, 1.0, 0.0, 0, &mut accept_rng); + assert_eq!((accepted, was_draft), (0, true)); + + let mut reject_rng = Rng::new(1); + let (replacement, was_draft) = + exact_delta_sample(&[0.0, 10.0], 0, 1.0, 1.0, 0.0, 0, &mut reject_rng); + assert_eq!((replacement, was_draft), (1, false)); + } + + #[test] + fn sampling_probabilities_match_ds4_filter_order() { + let probabilities = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 0.8, 0.2, 0); + assert_eq!(probabilities.len(), 2); + assert_eq!(probabilities[0].0, 1); + assert_eq!(probabilities[1].0, 2); + assert!((probabilities.iter().map(|(_, value)| value).sum::() - 1.0).abs() < 1e-6); + + let min_p_only = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 1.0, 0.2, 0); + assert_eq!( + min_p_only + .iter() + .map(|(token, _)| *token) + .collect::>(), + [1, 2] + ); + } + #[test] fn split_utf8_token_bytes_are_joined_before_decoding() { let mut generated = ChatTurn { diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 511bd1b..6f258e8 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -10,7 +10,7 @@ use profile::ExpertProfile; use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q4_K, Q8_0, Tensor as GgufTensor}; use super::validation::{DsparkConfig, SupportKind, dspark_config}; -use super::{Model, ModelFamily}; +use super::{Model, ModelFamily, Rng, exact_delta_sample}; use crate::model::ModelChoice; use crate::settings::{ EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode, @@ -407,7 +407,13 @@ impl Dspark { capture_mask: 0, cache_start: 0, cache_len: 0, - confidence_threshold: settings.dspark_confidence_threshold, + confidence_threshold: if settings.dspark_exact_sampling + && !settings.dspark_confidence_threshold_set + { + settings.dspark_confidence_threshold.max(0.8) + } else { + settings.dspark_confidence_threshold + }, strict: settings.dspark_strict || quality, drafted: 0, accepted: 0, @@ -2600,6 +2606,48 @@ struct SpecFrontier { dspark_cache_len: u32, } +struct SpecPrefixFrontier { + layers: Vec, +} + +struct BatchVerification { + tops: Vec, + logits: Vec>, + prefixes: Vec, +} + +fn capture_compression_frontier( + state: &CompressionState, + bytes: u64, + purpose: &str, +) -> Result { + let state_kv = Buffer::bytes(bytes)?; + let state_score = Buffer::bytes(bytes)?; + state_kv.copy_from(0, &state.state_kv, 0, bytes, purpose)?; + state_score.copy_from(0, &state.state_score, 0, bytes, purpose)?; + Ok(CompressionFrontier { + state_kv, + state_score, + bytes, + rows: state.rows, + }) +} + +fn restore_compression_frontier( + state: &mut CompressionState, + saved: &CompressionFrontier, + purpose: &str, +) -> Result<(), String> { + state + .state_kv + .copy_from(0, &saved.state_kv, 0, saved.bytes, purpose)?; + state + .state_score + .copy_from(0, &saved.state_score, 0, saved.bytes, purpose)?; + state.rows = saved.rows; + Ok(()) +} + impl LayerState { fn allocate(model: &Model, index: u32, context: u32, raw_cap: u32) -> Result { let shape = model.shape; @@ -2671,7 +2719,7 @@ impl Session { // SAFETY: declaration order is required because Rust drops fields in order. // `session` must release every Buffer before `_context` calls ds4_gpu_cleanup(), // and `_context` must drop before `model` unmaps memory wrapped without copying -// by native/metal/ds4_metal.m:10329. This intentionally differs from +// by `ds4_gpu_cleanup` in native/metal/ds4_metal.m. This intentionally differs from // DS4's `ds4.c` consumes this exact field order; do not reorder it. #[derive(Clone, Copy, Default)] pub(super) struct ExecutionStats { @@ -2934,31 +2982,6 @@ impl DeepSeekExecutor { } fn snapshot_spec_frontier(&self) -> Result { - fn snapshot(state: &CompressionState, bytes: u64) -> Result { - let state_kv = Buffer::bytes(bytes)?; - let state_score = Buffer::bytes(bytes)?; - state_kv.copy_from( - 0, - &state.state_kv, - 0, - bytes, - "saving speculative compressor KV state", - )?; - state_score.copy_from( - 0, - &state.state_score, - 0, - bytes, - "saving speculative compressor score state", - )?; - Ok(CompressionFrontier { - state_kv, - state_score, - bytes, - rows: state.rows, - }) - } - let shape = self.model.shape; let commands = Commands::begin()?; let layers = self @@ -2971,9 +2994,10 @@ impl DeepSeekExecutor { .as_ref() .map(|state| { let coefficient = if state.ratio == 4 { 2 } else { 1 }; - snapshot( + capture_compression_frontier( state, coefficient * coefficient * state.ratio as u64 * shape.head_dim * 4, + "saving speculative compressor state", ) }) .transpose()?; @@ -2981,7 +3005,11 @@ impl DeepSeekExecutor { .indexer .as_ref() .map(|state| { - snapshot(state, 4 * state.ratio as u64 * shape.indexer_head_dim * 4) + capture_compression_frontier( + state, + 4 * state.ratio as u64 * shape.indexer_head_dim * 4, + "saving speculative indexer state", + ) }) .transpose()?; Ok::<_, String>(LayerFrontier { @@ -3020,40 +3048,26 @@ impl DeepSeekExecutor { } fn restore_spec_frontier(&mut self, frontier: &SpecFrontier) -> Result<(), String> { - fn restore( - state: &mut CompressionState, - saved: &CompressionFrontier, - ) -> Result<(), String> { - state.state_kv.copy_from( - 0, - &saved.state_kv, - 0, - saved.bytes, - "restoring speculative compressor KV state", - )?; - state.state_score.copy_from( - 0, - &saved.state_score, - 0, - saved.bytes, - "restoring speculative compressor score state", - )?; - state.rows = saved.rows; - Ok(()) - } - if frontier.layers.len() != self.session.layers.len() { return Err("speculative frontier layer count changed".into()); } let commands = Commands::begin()?; for (layer, saved) in self.session.layers.iter_mut().zip(&frontier.layers) { match (&mut layer.compression, &saved.compression) { - (Some(state), Some(saved)) => restore(state, saved)?, + (Some(state), Some(saved)) => restore_compression_frontier( + state, + saved, + "restoring speculative compressor state", + )?, (None, None) => {} _ => return Err("speculative compressor layout changed".into()), } match (&mut layer.indexer, &saved.indexer) { - (Some(state), Some(saved)) => restore(state, saved)?, + (Some(state), Some(saved)) => restore_compression_frontier( + state, + saved, + "restoring speculative indexer state", + )?, (None, None) => {} _ => return Err("speculative indexer layout changed".into()), } @@ -3078,6 +3092,63 @@ impl DeepSeekExecutor { Ok(()) } + fn commit_spec_prefix( + &mut self, + baseline: &SpecFrontier, + prefix: &SpecPrefixFrontier, + proposals: &[i32], + logits: &[f32], + ) -> Result<(), String> { + let count = + u32::try_from(proposals.len()).map_err(|_| "speculative prefix is too large")?; + if count == 0 || prefix.layers.len() != self.session.layers.len() { + return Err("invalid speculative prefix frontier".into()); + } + let commands = Commands::begin()?; + for (layer, saved) in self.session.layers.iter_mut().zip(&prefix.layers) { + match (&mut layer.compression, &saved.compression) { + (Some(state), Some(saved)) => restore_compression_frontier( + state, + saved, + "committing speculative compressor prefix", + )?, + (None, None) => {} + _ => return Err("speculative compressor prefix layout changed".into()), + } + match (&mut layer.indexer, &saved.indexer) { + (Some(state), Some(saved)) => restore_compression_frontier( + state, + saved, + "committing speculative indexer prefix", + )?, + (None, None) => {} + _ => return Err("speculative indexer prefix layout changed".into()), + } + } + if let Some(dspark) = &mut self.dspark { + let row = u64::from(count - 1); + for slot in 0..dspark.config.target_layers.len() as u64 { + dspark.target_hidden.copy_from( + slot * self.model.shape.embd * 4, + &dspark.target_hidden_batch, + (slot * u64::from(self.session.prefill_cap) + row) * self.model.shape.embd * 4, + self.model.shape.embd * 4, + "committing speculative DSpark target prefix", + )?; + } + dspark.capture_mask = (1_u32 << dspark.config.target_layers.len()) - 1; + dspark.cache_start = baseline.dspark_cache_start; + dspark.cache_len = baseline.dspark_cache_len; + dspark.commit_proposed_prefix(count, self.session.raw_cap); + } + commands.finish()?; + self.session.position = baseline.position + count; + self.tokens.truncate(baseline.token_len); + self.tokens.extend_from_slice(proposals); + self.logits.clone_from_slice(logits); + Ok(()) + } + fn verify_target_suffix( &mut self, proposals: &[i32], @@ -3115,10 +3186,10 @@ impl DeepSeekExecutor { } let frontier = self.snapshot_spec_frontier()?; - let row_tops = match self.eval_batch_tops(proposals) { - Ok(tops) => { + let verification = match self.eval_batch_tops(proposals) { + Ok(verification) => { self.verifier_passes += 1; - tops + verification } Err(error) => { self.restore_spec_frontier(&frontier)?; @@ -3129,7 +3200,7 @@ impl DeepSeekExecutor { } }; let mut commit = 1_usize; - while commit < proposals.len() && row_tops[commit - 1] == proposals[commit] { + while commit < proposals.len() && verification.tops[commit - 1] == proposals[commit] { commit += 1; } if commit == proposals.len() { @@ -3139,6 +3210,17 @@ impl DeepSeekExecutor { return Ok(proposals.to_vec()); } + if let (Some(prefix), Some(logits)) = ( + verification.prefixes.get(commit - 1), + verification.logits.get(commit - 1), + ) { + self.commit_spec_prefix(&frontier, prefix, &proposals[..commit], logits)?; + self.verifier_ns = self + .verifier_ns + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); + return Ok(proposals[..commit].to_vec()); + } + self.restore_spec_frontier(&frontier)?; if let Some(dspark) = &mut self.dspark { dspark.commit_proposed_prefix(1, self.session.raw_cap); @@ -3157,6 +3239,133 @@ impl DeepSeekExecutor { Ok(proposals[..commit].to_vec()) } + #[allow(clippy::too_many_arguments)] + fn verify_target_suffix_stochastic( + &mut self, + proposals: &[i32], + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, + rng: &mut Rng, + cancelled: &std::sync::atomic::AtomicBool, + ) -> Result<(Vec, usize), String> { + if proposals.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) { + return Ok((Vec::new(), 0)); + } + let started = Instant::now(); + if self.quality + || proposals.len() == 1 + || self + .ssd + .as_ref() + .is_some_and(|ssd| u64::from(ssd.cache_experts) < self.model.shape.experts) + { + let mut emitted = Vec::new(); + let mut accepted = 0; + for &proposal in proposals { + let (token, was_draft) = exact_delta_sample( + &self.logits, + proposal, + temperature, + top_p, + min_p, + top_k, + rng, + ); + self.eval_target(token)?; + self.verifier_passes += 1; + emitted.push(token); + if !was_draft { + break; + } + accepted += 1; + if cancelled.load(std::sync::atomic::Ordering::Relaxed) { + break; + } + } + self.verifier_ns = self + .verifier_ns + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); + return Ok((emitted, accepted)); + } + let (first, accepted_first) = exact_delta_sample( + &self.logits, + proposals[0], + temperature, + top_p, + min_p, + top_k, + rng, + ); + if !accepted_first { + if let Some(dspark) = &mut self.dspark { + dspark.commit_proposed_prefix(1, self.session.raw_cap); + } + self.eval_target(first)?; + self.verifier_passes += 1; + self.verifier_ns = self + .verifier_ns + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); + return Ok((vec![first], 0)); + } + + let frontier = self.snapshot_spec_frontier()?; + let verification = match self.eval_batch_tops(proposals) { + Ok(verification) => { + self.verifier_passes += 1; + verification + } + Err(error) => { + self.restore_spec_frontier(&frontier)?; + return Err(error); + } + }; + let mut accepted = 1; + let mut replacement = None; + for (index, &proposal) in proposals.iter().enumerate().skip(1) { + let (token, was_draft) = exact_delta_sample( + &verification.logits[index - 1], + proposal, + temperature, + top_p, + min_p, + top_k, + rng, + ); + if !was_draft { + replacement = Some(token); + break; + } + accepted += 1; + } + if replacement.is_none() { + self.verifier_ns = self + .verifier_ns + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); + return Ok((proposals.to_vec(), accepted)); + } + + let prefix = verification + .prefixes + .get(accepted - 1) + .ok_or("missing stochastic verifier prefix")?; + let logits = verification + .logits + .get(accepted - 1) + .ok_or("missing stochastic verifier logits")?; + self.commit_spec_prefix(&frontier, prefix, &proposals[..accepted], logits)?; + let replacement = replacement.expect("replacement disappeared"); + self.eval_target(replacement)?; + self.verifier_passes += 1; + let mut emitted = proposals[..accepted].to_vec(); + emitted.push(replacement); + self.verifier_ns = self + .verifier_ns + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); + Ok((emitted, accepted)) + } + pub(super) fn eval_speculative_greedy( &mut self, first_token: i32, @@ -3269,6 +3478,81 @@ impl DeepSeekExecutor { Ok(accepted) } + #[allow(clippy::too_many_arguments)] + fn eval_speculative_sampled( + &mut self, + first_token: i32, + max_tokens: u32, + reasoning: ReasoningMode, + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, + rng: &mut Rng, + cancelled: &std::sync::atomic::AtomicBool, + ) -> Result, String> { + if self.dspark.is_none() { + self.eval_target(first_token)?; + return Ok(vec![first_token]); + } + if !self.speculative.dspark_exact_sampling { + return self.eval_speculative_greedy(first_token, max_tokens, reasoning, cancelled); + } + self.speculative_cycles += 1; + self.eval_target(first_token)?; + let mut emitted = vec![first_token]; + if self.dspark.as_ref().is_some_and(|dspark| dspark.strict) + || max_tokens <= 1 + || cancelled.load(std::sync::atomic::Ordering::Relaxed) + { + return Ok(emitted); + } + if self.ssd.is_some() { + install_speculative_model_maps(&self.model, "DSpark support mapping")?; + } + let mut dspark = self.dspark.take().expect("DSpark disappeared"); + let proposals = dspark.propose( + &self.model, + &self.weights, + first_token, + self.session.position.saturating_sub(1), + self.session.raw_cap, + ); + self.dspark = Some(dspark); + let mut proposals = proposals?; + proposals.truncate( + max_tokens + .saturating_sub(1) + .min(self.session.context.saturating_sub(self.session.position)) + as usize, + ); + if let Some(stop) = proposals + .iter() + .position(|token| self.model.is_stop_token_for_reasoning(*token, reasoning)) + { + proposals.truncate(stop + 1); + } + if proposals.len() < 2 { + self.dspark + .as_mut() + .expect("DSpark disappeared") + .commit_proposed_prefix(1, self.session.raw_cap); + return Ok(emitted); + } + let (verified, accepted) = self.verify_target_suffix_stochastic( + &proposals, + temperature, + top_p, + min_p, + top_k, + rng, + cancelled, + )?; + emitted.extend_from_slice(&verified); + self.dspark.as_mut().expect("DSpark disappeared").accepted += accepted as u64; + Ok(emitted) + } + fn legacy_mtp_draft(&mut self, token: i32, target_hc: bool) -> Result<(i32, f32), String> { let support = self .model @@ -3511,11 +3795,15 @@ impl DeepSeekExecutor { self.eval_batch_inner(tokens, false).map(|_| ()) } - fn eval_batch_tops(&mut self, tokens: &[i32]) -> Result, String> { + fn eval_batch_tops(&mut self, tokens: &[i32]) -> Result { self.eval_batch_inner(tokens, true) } - fn eval_batch_inner(&mut self, tokens: &[i32], collect_tops: bool) -> Result, String> { + fn eval_batch_inner( + &mut self, + tokens: &[i32], + collect_tops: bool, + ) -> Result { let rows = u32::try_from(tokens.len()).map_err(|_| "prefill batch is too large")?; if rows == 0 || rows > self.session.prefill_cap { return Err("prefill batch exceeds the configured prefill workspace".into()); @@ -3535,17 +3823,16 @@ impl DeepSeekExecutor { let size = self.model.main.len(); let shape = self.model.shape; let pos = self.session.position; - let batch_selected_addr = self.ssd.is_some() - && self.weights.layers.first().is_some_and(|layer| unsafe { - ds4_gpu_stream_prefill_batch_selected_addr_enabled( - rows, - shape.experts as u32, - shape.experts_used as u32, - layer.expert_gate.kind, - layer.expert_down.kind, - ) != 0 - }); - + let mut prefixes = (0..if collect_tops { rows } else { 0 }) + .map(|_| SpecPrefixFrontier { + layers: (0..shape.layers) + .map(|_| LayerFrontier { + compression: None, + indexer: None, + }) + .collect(), + }) + .collect::>(); if self.ssd.is_some() { install_deepseek_model_spans( &self.model, @@ -3580,6 +3867,16 @@ impl DeepSeekExecutor { .enumerate() { let started = Instant::now(); + let layer_selected_addr = self.ssd.is_some() + && unsafe { + ds4_gpu_stream_prefill_batch_selected_addr_enabled( + rows, + shape.experts as u32, + shape.experts_used as u32, + weights.expert_gate.kind, + weights.expert_down.kind, + ) != 0 + }; if let Some(ssd) = &self.ssd { install_deepseek_model_spans( &self.model, @@ -3587,7 +3884,7 @@ impl DeepSeekExecutor { &self.model, weights, index as u32, - batch_selected_addr, + layer_selected_addr, ssd.per_expert_bytes, )?, "DeepSeek prefill layer mapping", @@ -3606,6 +3903,7 @@ impl DeepSeekExecutor { rows, self.session.raw_cap, self.steering.as_ref(), + collect_tops.then_some(prefixes.as_mut_slice()), )?; if let Some(profile) = &mut self.profile { profile.record( @@ -3659,6 +3957,7 @@ impl DeepSeekExecutor { let output_rows = if collect_tops { rows } else { 1 }; let first_output = rows - output_rows; let mut tops = Vec::with_capacity(output_rows as usize); + let mut output_logits = Vec::with_capacity(output_rows as usize); for row in first_output..rows { let commands = Commands::begin()?; self.session.scratch.current_hc.copy_from( @@ -3672,13 +3971,20 @@ impl DeepSeekExecutor { commands.finish()?; self.session.scratch.logits.read_f32(&mut self.logits)?; tops.push(argmax(&self.logits)); + if collect_tops { + output_logits.push(self.logits.clone()); + } } self.session.position += rows; self.tokens.extend_from_slice(tokens); if let Some(profile) = &self.profile { profile.write()?; } - Ok(tops) + Ok(BatchVerification { + tops, + logits: output_logits, + prefixes, + }) } pub(super) fn logits(&self) -> &[f32] { @@ -4003,6 +4309,7 @@ impl Executor { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, crate::settings::EngineSsdSettings { enabled: false, @@ -4061,6 +4368,38 @@ impl Executor { } } + #[allow(clippy::too_many_arguments)] + pub(super) fn eval_speculative_sampled( + &mut self, + token: i32, + max_tokens: u32, + reasoning: ReasoningMode, + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, + rng: &mut Rng, + cancelled: &std::sync::atomic::AtomicBool, + ) -> Result, String> { + match self { + Self::DeepSeek(executor) => executor.eval_speculative_sampled( + token, + max_tokens, + reasoning, + temperature, + top_p, + min_p, + top_k, + rng, + cancelled, + ), + Self::Glm(executor) => { + executor.eval(token)?; + Ok(vec![token]) + } + } + } + pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { match self { Self::DeepSeek(executor) => executor.eval(token), @@ -4272,10 +4611,13 @@ fn compress_attention_batch( freq_scale: f32, ext: f32, attn_factor: f32, + mut prefixes: Option<&mut [SpecPrefixFrontier]>, + layer: usize, ) -> Result { let ratio = state.ratio; let chunk = rows / ratio; - if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) { + if prefixes.is_none() && (pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio))) + { let before = if pos == 0 { 0 } else { state.rows }; let target = s .compressed_stage @@ -4439,6 +4781,15 @@ fn compress_attention_batch( )?; state.rows += 1; } + if let Some(prefixes) = prefixes.as_deref_mut() { + let coefficient = if ratio == 4 { 2 } else { 1 }; + prefixes[row as usize].layers[layer].compression = + Some(capture_compression_frontier( + state, + coefficient * coefficient * u64::from(ratio) * shape.head_dim * 4, + "capturing speculative compressor prefix", + )?); + } } } Ok(state.rows) @@ -4459,9 +4810,12 @@ fn compress_index_batch( freq_scale: f32, ext: f32, attn_factor: f32, + mut prefixes: Option<&mut [SpecPrefixFrontier]>, + layer: usize, ) -> Result<(), String> { let ratio = state.ratio; - if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) { + if prefixes.is_none() && (pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio))) + { let before = if pos == 0 { 0 } else { state.rows }; let chunk = rows / ratio; let target = state.cache.view( @@ -4613,6 +4967,13 @@ fn compress_index_batch( )?; state.rows += 1; } + if let Some(prefixes) = prefixes.as_deref_mut() { + prefixes[row as usize].layers[layer].indexer = Some(capture_compression_frontier( + state, + 4 * u64::from(ratio) * shape.indexer_head_dim * 4, + "capturing speculative indexer prefix", + )?); + } } } Ok(()) @@ -4631,6 +4992,7 @@ fn encode_batch_layer( rows: u32, raw_cap: u32, steering: Option<&Steering>, + mut prefixes: Option<&mut [SpecPrefixFrontier]>, ) -> Result<(), String> { let hc_dim = shape.hc * shape.embd; let mix_hc = 2 * shape.hc + shape.hc * shape.hc; @@ -4897,6 +5259,8 @@ fn encode_batch_layer( freq_scale, ext, attn_factor, + prefixes.as_deref_mut(), + layer as usize, )?; } @@ -4943,6 +5307,8 @@ fn encode_batch_layer( freq_scale, ext, attn_factor, + prefixes, + layer as usize, )?; matmul_rows( &s.indexer_q, @@ -6558,48 +6924,32 @@ fn encode_output( }, "output HC weights", )?; - let fused_sum_norm = unsafe { - ds4_gpu_hc_weighted_sum_norm_tensor( - s.output_embedding.raw(), - s.output_norm.raw(), - s.current_hc.raw(), - s.output_weights.raw(), - map, - size, - w.output_norm.offset, - shape.embd as u32, - shape.hc as u32, - shape.rms_epsilon, - ) - } != 0; - if !fused_sum_norm { - call( - unsafe { - ds4_gpu_hc_weighted_sum_tensor( - s.output_embedding.raw(), - s.current_hc.raw(), - s.output_weights.raw(), - shape.embd as u32, - shape.hc as u32, - ) - }, - "output HC collapse", - )?; - call( - unsafe { - ds4_gpu_rms_norm_weight_tensor( - s.output_norm.raw(), - s.output_embedding.raw(), - map, - size, - w.output_norm.offset, - shape.embd as u32, - shape.rms_epsilon, - ) - }, - "output norm", - )?; - } + call( + unsafe { + ds4_gpu_hc_weighted_sum_tensor( + s.output_embedding.raw(), + s.current_hc.raw(), + s.output_weights.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "output HC collapse", + )?; + call( + unsafe { + ds4_gpu_rms_norm_weight_tensor( + s.output_norm.raw(), + s.output_embedding.raw(), + map, + size, + w.output_norm.offset, + shape.embd as u32, + shape.rms_epsilon, + ) + }, + "output norm", + )?; q8( &s.logits, w.output, @@ -7186,6 +7536,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: false, @@ -7323,6 +7674,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: false, @@ -7414,6 +7766,94 @@ mod tests { ); } + #[test] + #[ignore = "requires the installed 0731 target and checkpoint-specific DSpark GGUF fixtures"] + fn flash_0731_runs_exact_sampled_dspark() { + use super::{DeepSeekExecutor, argmax, configure_sources}; + use crate::engine::gguf::Gguf; + use crate::engine::validation::validate_support; + use crate::engine::{Model, Rng}; + use crate::model::{ModelChoice, validate_engine_artifacts}; + use crate::settings::{ + EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode, + }; + use std::sync::atomic::AtomicBool; + + configure_sources().unwrap(); + let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true); + validate_engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &artifacts) + .unwrap(); + let mut model = + Model::open_main(&artifacts.model, ModelChoice::DeepSeekV4Flash0731).unwrap(); + let support = Gguf::open(artifacts.mtp.as_ref().unwrap()).unwrap(); + model.support_kind = Some(validate_support(&support, &model.shape).unwrap()); + model.support = Some(support); + let prompt = model.render_conversation( + "", + &[crate::engine::ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "hi".into(), + }], + ReasoningMode::Direct, + ); + let mut executor = DeepSeekExecutor::open( + model, + 64, + false, + 64, + 100, + EngineSpeculativeSettings { + mtp_draft_tokens: 1, + mtp_margin: 3.0, + glm_mtp: false, + glm_mtp_timing: false, + dspark: true, + dspark_confidence_threshold: 0.6, + dspark_confidence_threshold_set: false, + dspark_strict: false, + dspark_exact_sampling: true, + }, + EngineSsdSettings { + enabled: false, + cold: false, + cache_experts: 0, + cache_bytes: 0, + full_layers: 0, + full_layers_set: false, + preload_experts: 0, + }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, + ) + .unwrap(); + executor.prefill(&prompt, |_| true).unwrap(); + let first = argmax(executor.logits()); + let cycle = executor + .eval_speculative_sampled( + first, + 4, + ReasoningMode::Direct, + 0.8, + 0.95, + 0.0, + 0, + &mut Rng::new(7), + &AtomicBool::new(false), + ) + .unwrap(); + assert!(!cycle.is_empty()); + assert!(executor.logits().iter().all(|logit| logit.is_finite())); + assert!(executor.session.position >= prompt.len() as u32 + cycle.len() as u32); + } + #[test] #[ignore = "requires the installed Flash, legacy MTP, and DSpark GGUF fixtures"] fn ssd_streaming_supports_legacy_mtp_and_dspark() { @@ -7468,6 +7908,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: true, @@ -7544,6 +7985,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }; let mut executor = DeepSeekExecutor::open( model, @@ -7627,6 +8069,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: true, @@ -7691,6 +8134,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: false, @@ -7775,6 +8219,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: false, @@ -7850,6 +8295,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, EngineSsdSettings { enabled: true, diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs index 7886c0a..a5f83b2 100644 --- a/src/engine/metal/glm.rs +++ b/src/engine/metal/glm.rs @@ -401,6 +401,7 @@ impl GlmExecutor { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, None, ) @@ -3412,6 +3413,7 @@ mod tests { dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, dspark_strict: false, + dspark_exact_sampling: false, }, None, ) diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 1828eab..e0c8459 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -1283,18 +1283,6 @@ unsafe extern "C" { hc: u32, eps: f32, ) -> i32; - pub(super) fn ds4_gpu_hc_weighted_sum_norm_tensor( - out: *mut GpuTensor, - norm: *mut GpuTensor, - residual: *const GpuTensor, - weights: *const GpuTensor, - map: *const c_void, - size: u64, - norm_weight: u64, - embd: u32, - hc: u32, - eps: f32, - ) -> i32; pub(super) fn ds4_gpu_hc_weighted_sum_tensor( out: *mut GpuTensor, residual: *const GpuTensor, @@ -1443,9 +1431,9 @@ impl Buffer { } pub(super) fn view(&self, offset: u64, bytes: u64) -> Result { - // SAFETY: native/metal/ds4_metal.m:7916-7940 bounds-checks the view, - // ARC-retains base_obj.buffer at :7926, and marks the view non-owning at - // :7929, so it may outlive and be freed independently of this wrapper. + // SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks + // the view, ARC-retains base_obj.buffer, and marks the view non-owning, so + // it may outlive and be freed independently of this wrapper. // Recheck those guarantees whenever the vendored Metal file is re-synced. NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) }) .map(Self) diff --git a/src/engine/validation.rs b/src/engine/validation.rs index 1ba830c..0a339ef 100644 --- a/src/engine/validation.rs +++ b/src/engine/validation.rs @@ -7,7 +7,14 @@ pub(crate) fn validate_model_artifact( ) -> Result<(), String> { if support { let model = Gguf::open(path)?; - validate_support(&model, &FLASH).map(|_| ()) + let shape = match expected { + ModelChoice::DeepSeekV4Flash => FLASH, + ModelChoice::DeepSeekV4Flash0731 => FLASH_0731, + ModelChoice::DeepSeekV4Pro | ModelChoice::Glm52 => { + return Err(format!("{expected} does not use an external support GGUF")); + } + }; + validate_support(&model, &shape).map(|_| ()) } else { let model = Model::open_main(path, expected)?; let summary = model.summary(); @@ -759,7 +766,10 @@ fn validate_glm_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> { } pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String> { - if shape.model != ModelChoice::DeepSeekV4Flash { + if !matches!( + shape.model, + ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Flash0731 + ) { return Err("DSpark support is available only for DeepSeek V4 Flash".into()); } let DsparkConfig { diff --git a/src/model.rs b/src/model.rs index 046260d..d0118cd 100644 --- a/src/model.rs +++ b/src/model.rs @@ -10,22 +10,22 @@ use std::fs; use std::path::{Path, PathBuf}; pub(crate) const MODEL_CHOICES: [ModelChoice; 4] = [ - ModelChoice::DeepSeekV4Flash, ModelChoice::DeepSeekV4Flash0731, + ModelChoice::DeepSeekV4Flash, ModelChoice::DeepSeekV4Pro, ModelChoice::Glm52, ]; -pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 6] = [ +pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [ + ManagedArtifactId::DeepSeekV4Flash0731, + ManagedArtifactId::DeepSeekV4Flash0731Dspark, ManagedArtifactId::DeepSeekV4Flash, ManagedArtifactId::DeepSeekV4FlashMtp, ManagedArtifactId::DeepSeekV4FlashDspark, - ManagedArtifactId::DeepSeekV4Flash0731, ManagedArtifactId::DeepSeekV4Pro, ManagedArtifactId::Glm52, ]; const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf"; -const DEEPSEEK_FLASH_0731_REPOSITORY: &str = "Rednalreden/DeepSeek-V4-Flash-0731-dwarfstar-q2-gguf"; const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf"; const FLASH: Artifact = Artifact { @@ -54,18 +54,26 @@ const FLASH_MTP: Artifact = Artifact { }; const FLASH_0731: Artifact = Artifact { label: "DeepSeek V4 Flash 0731 model", - file_name: "DeepSeek-V4-Flash-0731-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-imatrix.gguf", - repository: DEEPSEEK_FLASH_0731_REPOSITORY, - size: 86_720_111_520, - sha256: "0b39f9c337d6b49c77db2190556b8563abf3c5fbb98be3b58cf8d3a1db191e5f", + file_name: "DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf", + repository: DEEPSEEK_REPOSITORY, + size: 86_720_111_488, + sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0", support: Some(false), }; +const FLASH_0731_DSPARK: Artifact = Artifact { + label: "DeepSeek V4 Flash 0731 DSpark support", + file_name: "DeepSeek-V4-Flash-DSpark-support-0731.gguf", + repository: DEEPSEEK_REPOSITORY, + size: 5_989_114_272, + sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360", + support: Some(true), +}; const PRO: Artifact = Artifact { - label: "DeepSeek V4 Pro model", - file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf", + label: "DeepSeek V4 Pro 0813 model", + file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf", repository: DEEPSEEK_REPOSITORY, size: 464_627_334_560, - sha256: "a0314d9c0e16122cd60071079124a2d17185d317c55a8f95ecb3ed3506278a96", + sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193", support: Some(false), }; const GLM: Artifact = Artifact { @@ -79,9 +87,9 @@ const GLM: Artifact = Artifact { #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)] pub(crate) enum ModelChoice { - #[default] #[serde(rename = "deepseek-v4-flash")] DeepSeekV4Flash, + #[default] #[serde(rename = "deepseek-v4-flash-0731")] DeepSeekV4Flash0731, #[serde(rename = "deepseek-v4-pro")] @@ -104,10 +112,14 @@ impl ModelChoice { MODEL_CHOICES.into_iter().find(|model| model.id() == id) } - pub(crate) fn supports_dspark(self) -> bool { + pub(crate) fn supports_legacy_mtp(self) -> bool { self == Self::DeepSeekV4Flash } + pub(crate) fn supports_dspark(self) -> bool { + matches!(self, Self::DeepSeekV4Flash | Self::DeepSeekV4Flash0731) + } + fn main_artifact(self) -> &'static Artifact { match self { Self::DeepSeekV4Flash => &FLASH, @@ -117,6 +129,14 @@ impl ModelChoice { } } + fn dspark_artifact(self) -> Option<&'static Artifact> { + match self { + Self::DeepSeekV4Flash => Some(&FLASH_DSPARK), + Self::DeepSeekV4Flash0731 => Some(&FLASH_0731_DSPARK), + Self::DeepSeekV4Pro | Self::Glm52 => None, + } + } + #[cfg(test)] fn artifacts( self, @@ -125,8 +145,8 @@ impl ModelChoice { ) -> impl Iterator { [ Some(self.main_artifact()), - (self.supports_dspark() && legacy_mtp_enabled).then_some(&FLASH_MTP), - (self.supports_dspark() && dspark_enabled).then_some(&FLASH_DSPARK), + (self.supports_legacy_mtp() && legacy_mtp_enabled).then_some(&FLASH_MTP), + dspark_enabled.then(|| self.dspark_artifact()).flatten(), ] .into_iter() .flatten() @@ -147,22 +167,61 @@ pub(crate) fn engine_artifacts( ) -> EngineArtifacts { EngineArtifacts { model: model.main_artifact().path(model, models_path), - mtp: if model.supports_dspark() && legacy_mtp_enabled { + mtp: if model.supports_legacy_mtp() && legacy_mtp_enabled { Some(FLASH_MTP.path(model, models_path)) - } else if model.supports_dspark() && dspark_enabled { - Some(FLASH_DSPARK.path(model, models_path)) + } else if dspark_enabled { + model + .dspark_artifact() + .map(|artifact| artifact.path(model, models_path)) } else { None }, } } +pub(crate) fn validate_engine_artifacts( + model: ModelChoice, + legacy_mtp_enabled: bool, + dspark_enabled: bool, + artifacts: &EngineArtifacts, +) -> Result<(), String> { + if legacy_mtp_enabled && dspark_enabled { + return Err("Legacy MTP and DSpark cannot be enabled together".into()); + } + if legacy_mtp_enabled && !model.supports_legacy_mtp() { + return Err(format!("Legacy MTP is not compatible with {model}")); + } + if dspark_enabled && !model.supports_dspark() { + return Err(format!("DSpark is not compatible with {model}")); + } + model + .main_artifact() + .validate_installed_path(&artifacts.model)?; + let expected_support = if legacy_mtp_enabled { + model.supports_legacy_mtp().then_some(&FLASH_MTP) + } else if dspark_enabled { + model.dspark_artifact() + } else { + None + }; + match (expected_support, artifacts.mtp.as_deref()) { + (Some(expected), Some(path)) => expected.validate_installed_path(path), + (None, None) => Ok(()), + (Some(_), None) => Err(format!("{model} is missing its required support GGUF")), + (None, Some(path)) => Err(format!( + "{} is not compatible with the selected {model} checkpoint", + path.display() + )), + } +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum ManagedArtifactId { DeepSeekV4Flash, DeepSeekV4FlashMtp, DeepSeekV4FlashDspark, DeepSeekV4Flash0731, + DeepSeekV4Flash0731Dspark, DeepSeekV4Pro, Glm52, } @@ -173,7 +232,9 @@ impl ManagedArtifactId { Self::DeepSeekV4Flash | Self::DeepSeekV4FlashMtp | Self::DeepSeekV4FlashDspark => { ModelChoice::DeepSeekV4Flash } - Self::DeepSeekV4Flash0731 => ModelChoice::DeepSeekV4Flash0731, + Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => { + ModelChoice::DeepSeekV4Flash0731 + } Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro, Self::Glm52 => ModelChoice::Glm52, } @@ -185,6 +246,7 @@ impl ManagedArtifactId { Self::DeepSeekV4FlashMtp => &FLASH_MTP, Self::DeepSeekV4FlashDspark => &FLASH_DSPARK, Self::DeepSeekV4Flash0731 => &FLASH_0731, + Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK, Self::DeepSeekV4Pro => &PRO, Self::Glm52 => &GLM, } @@ -278,9 +340,9 @@ pub(crate) enum DownloadOutcome { impl fmt::Display for ModelChoice { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter.write_str(match self { - Self::DeepSeekV4Flash => "DeepSeek V4 Flash", + Self::DeepSeekV4Flash => "DeepSeek V4 Flash (deprecated preview)", Self::DeepSeekV4Flash0731 => "DeepSeek V4 Flash 0731", - Self::DeepSeekV4Pro => "DeepSeek V4 Pro", + Self::DeepSeekV4Pro => "DeepSeek V4 Pro 0813", Self::Glm52 => "GLM 5.2", }) } @@ -296,6 +358,36 @@ struct Artifact { } impl Artifact { + fn validate_installed_path(&self, path: &Path) -> Result<(), String> { + if path.file_name().and_then(|name| name.to_str()) != Some(self.file_name) { + return Err(format!( + "{} is not the expected {} artifact", + path.display(), + self.label + )); + } + let size = path + .metadata() + .map_err(|error| format!("Could not inspect {}: {error}", path.display()))? + .len(); + if size != self.size { + return Err(format!( + "{} has {size} bytes, expected {}", + path.display(), + self.size + )); + } + let verification = fs::read_to_string(path.with_extension("gguf.sha256")) + .map_err(|_| format!("{} has not passed checksum verification", path.display()))?; + if verification.trim() != self.sha256 { + return Err(format!( + "{} has the wrong checkpoint identity", + path.display() + )); + } + Ok(()) + } + fn path(&self, model: ModelChoice, models_path: &Path) -> PathBuf { models_path.join(model.id()).join(self.file_name) } diff --git a/src/model/transfer.rs b/src/model/transfer.rs index 23655dc..5753815 100644 --- a/src/model/transfer.rs +++ b/src/model/transfer.rs @@ -337,7 +337,7 @@ mod tests { assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448); assert_eq!( ModelChoice::DeepSeekV4Flash0731.main_artifact().size, - 86_720_111_520 + 86_720_111_488 ); assert_eq!( ModelChoice::DeepSeekV4Flash.artifacts(true, true).count(), @@ -348,7 +348,7 @@ mod tests { ModelChoice::DeepSeekV4Flash0731 .artifacts(true, true) .count(), - 1 + 2 ); let id = SystemTime::now() @@ -365,8 +365,14 @@ mod tests { engine.mtp.as_deref().and_then(Path::file_name), Some(std::ffi::OsStr::new(FLASH_DSPARK.file_name)) ); + let flash_0731 = + engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &models_path); + assert_eq!( + flash_0731.mtp.as_deref().and_then(Path::file_name), + Some(std::ffi::OsStr::new(FLASH_0731_DSPARK.file_name)) + ); assert!( - engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, true, &models_path) + engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, false, &models_path) .mtp .is_none() ); @@ -389,6 +395,36 @@ mod tests { .unwrap(), empty.sha256 ); + let installed = empty.path(ModelChoice::DeepSeekV4Flash, &models_path); + assert!(empty.validate_installed_path(&installed).is_ok()); + let wrong_name = installed.with_file_name("wrong-checkpoint.gguf"); + fs::write(&wrong_name, []).unwrap(); + fs::write(wrong_name.with_extension("gguf.sha256"), empty.sha256).unwrap(); + assert!(empty.validate_installed_path(&wrong_name).is_err()); + assert!( + validate_engine_artifacts( + ModelChoice::DeepSeekV4Flash0731, + true, + false, + &EngineArtifacts { + model: installed.clone(), + mtp: None, + }, + ) + .is_err() + ); + assert!( + validate_engine_artifacts( + ModelChoice::DeepSeekV4Pro, + false, + true, + &EngineArtifacts { + model: installed, + mtp: None, + }, + ) + .is_err() + ); fs::remove_dir_all(models_path).unwrap(); } diff --git a/src/settings.rs b/src/settings.rs index 9c0993c..2e4c910 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -30,6 +30,7 @@ pub(crate) struct SpeculativePreferences { pub(crate) dspark_enabled: bool, pub(crate) dspark_confidence_threshold: Option, pub(crate) dspark_strict: bool, + pub(crate) dspark_exact_sampling: bool, } impl Default for SpeculativePreferences { @@ -43,6 +44,7 @@ impl Default for SpeculativePreferences { dspark_enabled: false, dspark_confidence_threshold: None, dspark_strict: false, + dspark_exact_sampling: false, } } } @@ -62,13 +64,15 @@ impl SpeculativePreferences { if self.dspark_enabled && !model.supports_dspark() { return Err("DSpark is not available for the selected model.".into()); } - if self.legacy_mtp_enabled && !model.supports_dspark() { + if self.legacy_mtp_enabled && !model.supports_legacy_mtp() { return Err("Legacy MTP is not available for the selected model.".into()); } if self.legacy_mtp_enabled && self.dspark_enabled { return Err("Legacy MTP and DSpark use different support artifacts.".into()); } - if (self.dspark_confidence_threshold.is_some() || self.dspark_strict) + if (self.dspark_confidence_threshold.is_some() + || self.dspark_strict + || self.dspark_exact_sampling) && !self.dspark_enabled { return Err("DSpark tuning requires DSpark to be enabled.".into()); @@ -86,9 +90,10 @@ impl SpeculativePreferences { glm_mtp: self.glm_mtp, glm_mtp_timing: self.glm_mtp_timing, dspark: self.dspark_enabled, - dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.9), + dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.6), dspark_confidence_threshold_set: self.dspark_confidence_threshold.is_some(), dspark_strict: self.dspark_strict, + dspark_exact_sampling: self.dspark_exact_sampling, } } } @@ -103,6 +108,7 @@ pub(crate) struct EngineSpeculativeSettings { pub(crate) dspark_confidence_threshold: f32, pub(crate) dspark_confidence_threshold_set: bool, pub(crate) dspark_strict: bool, + pub(crate) dspark_exact_sampling: bool, } /// An expert count, or a whole GiB budget. Written as `4` or `64GB`, the same @@ -785,7 +791,7 @@ mod tests { let defaults = SpeculativePreferences::default(); let engine = defaults.engine_settings(); assert_eq!((engine.mtp_draft_tokens, engine.mtp_margin), (1, 3.0)); - assert_eq!(engine.dspark_confidence_threshold, 0.9); + assert_eq!(engine.dspark_confidence_threshold, 0.6); assert!(!engine.dspark_confidence_threshold_set); let tuned = SpeculativePreferences { @@ -793,6 +799,7 @@ mod tests { dspark_enabled: true, dspark_confidence_threshold: Some(0.7), dspark_strict: true, + dspark_exact_sampling: true, ..defaults }; assert!(tuned.validate(ModelChoice::DeepSeekV4Flash).is_ok()); @@ -812,6 +819,7 @@ mod tests { ..SpeculativePreferences::default() }; assert!(legacy.validate(ModelChoice::DeepSeekV4Flash).is_ok()); + assert!(legacy.validate(ModelChoice::DeepSeekV4Flash0731).is_err()); assert!(legacy.validate(ModelChoice::DeepSeekV4Pro).is_err()); assert!( SpeculativePreferences { @@ -821,6 +829,14 @@ mod tests { .validate(ModelChoice::DeepSeekV4Flash) .is_err() ); + assert!( + SpeculativePreferences { + dspark_exact_sampling: true, + ..SpeculativePreferences::default() + } + .validate(ModelChoice::DeepSeekV4Flash0731) + .is_err() + ); } #[test]