diff --git a/CHANGELOG.md b/CHANGELOG.md index 9206562..a5aebd1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,19 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - `skinny_matmul`: x @ w.T at token widths 1..16 against small-N large-K nn.Linear-layout weights, 4-8x faster than the stock GEMM at widths 2..16 (router gates, indexer projections at speculative verify widths). +- `hc_front_reduce` / `hc_front_expand_reduce` / `hc_sinkhorn_collapse` / + `hc_expand`: fused deepseek4 hyper-connection glue for single-token + decode; replaces ~176 python kernel launches per step with 4 native ops. +- `get_cb_caps` / `set_cb_caps`: runtime read/write of MLX's command + buffer split caps, so a server can run coarse buffers during decode and + fine buffers during deep prefill. + +### Changed +- iq2_xxs / iq2_xs / iq2_s / iq3_s MoE gather decode is 9-12% faster per + call (hoisted block scale, byte-indexed grids); the ext mat-vec at + verify widths 2..8 gains 7-10% on the same codecs. +- Score-mixed MoE down gather gains a slot-parallel kernel at decode + scale (bit-identical; KQ_MOE_SP forces either form). ## [0.3.9] diff --git a/CMakeLists.txt b/CMakeLists.txt index 6c82681..40ec913 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -127,6 +127,8 @@ target_sources( ${CMAKE_CURRENT_LIST_DIR}/src/kquant_dsa_qat.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_moe_glu.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_norm_fused.cpp + ${CMAKE_CURRENT_LIST_DIR}/src/kquant_hc_glue.cpp + ${CMAKE_CURRENT_LIST_DIR}/src/kquant_cb_caps.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_route_shed.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_skinny_mv.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_gather.cpp @@ -209,6 +211,7 @@ if(MLX_BUILD_METAL) ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu.metal ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu_kq.metal ${CMAKE_CURRENT_LIST_DIR}/metal/kq_norm_fused.metal + ${CMAKE_CURRENT_LIST_DIR}/metal/kq_hc_glue.metal ${CMAKE_CURRENT_LIST_DIR}/metal/kq_route_shed.metal ${CMAKE_CURRENT_LIST_DIR}/metal/kq_skinny_mv.metal INCLUDE_DIRS diff --git a/bindings.cpp b/bindings.cpp index 87d6d27..5acfadb 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -1151,6 +1152,140 @@ NB_MODULE(_ext, m) { array: same shape and dtype as a. )"); + m.def( + "hc_front_reduce", + &mlx_kquant::hc_front_reduce, + "x"_a, + "fn"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + Hyper-connection front reduction for the fused M=1 decode route: + the 24 mix dots of x against fn plus the row sum of squares + (deferred rms factor). hc_mult 4 only. + + Args: + x (array): [..., 4, D] streams, float16/bfloat16, D % 8 == 0, + D <= 8192. + fn (array): [24, 4 * D] float32 mix matrix. + + Returns: + tuple: (mixes_raw f32 [..., 24], sumsq f32 [..., 1]). + )"); + + m.def( + "hc_front_expand_reduce", + &mlx_kquant::hc_front_expand_reduce, + "x_sub"_a, + "resid"_a, + "post"_a, + "comb"_a, + "fn"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + The previous cycle's hc_expand fused ahead of the front reduction: + one dispatch expands (x_sub, resid, post, comb) to h, writes it, + and reduces the mix dots and sum of squares of h. + + Args: + x_sub (array): [..., D] sublayer output. + resid (array): [..., 4, D] residual streams. + post (array): [..., 4] float32. + comb (array): [..., 4, 4] float32. + fn (array): [24, 4 * D] float32 mix matrix. + + Returns: + tuple: (h [..., 4, D], mixes_raw f32 [..., 24], + sumsq f32 [..., 1]); h is bit-identical to the unfused expand. + )"); + + m.def( + "hc_sinkhorn_collapse", + &mlx_kquant::hc_sinkhorn_collapse, + "x"_a, + "mixes_raw"_a, + "sumsq"_a, + "scale"_a, + "base"_a, + "w"_a, + "iters"_a, + "hc_eps"_a, + "norm_eps"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + Sinkhorn mix normalization plus stream collapse with the sublayer + RMSNorm folded into the single output rounding. The deferred front + rms factor enters through sumsq and multiplies the three scales. + + Args: + x (array): [..., 4, D] streams, float16/bfloat16. + mixes_raw (array): [..., 24] float32 from the front reduction. + sumsq (array): [..., 1] float32 row sum of squares. + scale (array): [3] float32 pre/post/comb scales. + base (array): [24] float32 mix biases. + w (array): [D] sublayer norm weight, same dtype as x. + iters (int): sinkhorn iterations. + hc_eps (float): sinkhorn epsilon. + norm_eps (float): rms_norm epsilon. + + Returns: + tuple: (collapsed [..., D], post f32 [..., 4], + comb f32 [..., 4, 4]). + )"); + + m.def( + "hc_expand", + &mlx_kquant::hc_expand, + "x"_a, + "resid"_a, + "post"_a, + "comb"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + Expand the sublayer output back to four streams: + out[i] = post[i] * x + sum_j comb[j][i] * resid[j]. + + Args: + x (array): [..., D] sublayer output. + resid (array): [..., 4, D] residual streams. + post (array): [..., 4] float32. + comb (array): [..., 4, 4] float32. + + Returns: + array: [..., 4, D], dtype of resid. + )"); + + m.def( + "get_cb_caps", + &mlx_kquant::get_cb_caps, + R"( + Read MLX's live command-buffer split caps. + + Returns: + tuple: (max_ops_per_buffer, max_mb_per_buffer). + )"); + + m.def( + "set_cb_caps", + &mlx_kquant::set_cb_caps, + "max_ops"_a, + "max_mb"_a, + R"( + Set MLX's command-buffer split caps at runtime. The env knobs + latch at device init; decode wants coarse buffers, deep prefill + fine ones, so servers flip these per phase. + + Args: + max_ops (int): ops per command buffer, in [1, 2^30]. + max_mb (int): MB per command buffer, in [1, 2^30]. + + Returns: + tuple: the previous (max_ops, max_mb). + )"); + m.def( "skinny_matmul", &mlx_kquant::skinny_matmul, diff --git a/metal/kq_hc_glue.metal b/metal/kq_hc_glue.metal new file mode 100644 index 0000000..e7c9d14 --- /dev/null +++ b/metal/kq_hc_glue.metal @@ -0,0 +1,14 @@ +// clang-format off +// Hyper-connection M=1 glue kernel instantiations; see kq_hc_glue.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/kq_hc_glue.h" + +instantiate_kernel("kq_hc_front_reduce_bfloat16_t", kq_hc_front_reduce, bfloat16_t) +instantiate_kernel("kq_hc_front_reduce_float16_t", kq_hc_front_reduce, float16_t) +instantiate_kernel("kq_hc_front_expand_reduce_bfloat16_t", kq_hc_front_expand_reduce, bfloat16_t) +instantiate_kernel("kq_hc_front_expand_reduce_float16_t", kq_hc_front_expand_reduce, float16_t) +instantiate_kernel("kq_hc_sinkhorn_collapse_bfloat16_t", kq_hc_sinkhorn_collapse, bfloat16_t) +instantiate_kernel("kq_hc_sinkhorn_collapse_float16_t", kq_hc_sinkhorn_collapse, float16_t) +instantiate_kernel("kq_hc_expand_bfloat16_t", kq_hc_expand, bfloat16_t) +instantiate_kernel("kq_hc_expand_float16_t", kq_hc_expand, float16_t) + // clang-format on diff --git a/metal/kq_moe_glu_kq.metal b/metal/kq_moe_glu_kq.metal index 89a6372..28a4150 100644 --- a/metal/kq_moe_glu_kq.metal +++ b/metal/kq_moe_glu_kq.metal @@ -129,10 +129,18 @@ instantiate_kq_moe_glu_kq_fine(q8_0, float16_t) "kq_" #codec "_gather_qmv_mix_ns" sfx "_float16_t", \ kq_ext_gather_qmv_mix_ns, float16_t, traits, nx) +// The slot-parallel variant is NX = 8 only (wide K-lanes measured +// flat-to-negative; sp multiplies threads without shortening K-chains). #define instantiate_kq_ext_mix_ns(codec, traits) \ instantiate_kq_ext_mix_ns_nx(codec, traits, 8, "") \ instantiate_kq_ext_mix_ns_nx(codec, traits, 16, "_nx16") \ - instantiate_kq_ext_mix_ns_nx(codec, traits, 32, "_nx32") + instantiate_kq_ext_mix_ns_nx(codec, traits, 32, "_nx32") \ + instantiate_kernel( \ + "kq_" #codec "_gather_qmv_mix_ns_sp_bfloat16_t", \ + kq_ext_gather_qmv_mix_ns_sp, bfloat16_t, traits, 8) \ + instantiate_kernel( \ + "kq_" #codec "_gather_qmv_mix_ns_sp_float16_t", \ + kq_ext_gather_qmv_mix_ns_sp, float16_t, traits, 8) // Biased experts (gpt-oss): per-(expert, out_dim) f32 biases fused into the // GLU epilogue / qmv store. Only the clamped-SwiGLU epilogue is emitted -- diff --git a/metal/mlx/backend/metal/kernels/kq_hc_glue.h b/metal/mlx/backend/metal/kernels/kq_hc_glue.h new file mode 100644 index 0000000..c3ececf --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_hc_glue.h @@ -0,0 +1,390 @@ +// Fused hyper-connection glue kernels for the deepseek4 M=1 decode route. +// Four streams (hc_mult 4) baked in; D is the per-stream hidden size. +// Numerics mirror the certified gmlx JIT kernels exactly: f32 accumulate, +// fast::exp sinkhorn, round-before-use in the fused expand, single +// rounding at each T-dtype write. +// +// kq_hc_front_reduce: mixes_raw[m] = dot(x, fn[m]), plus the row +// sum of squares (deferred rms factor). One +// threadgroup per (row, m) with m == MIX for +// the sumsq lane. +// kq_hc_front_expand_reduce: the previous cycle's expand recomputed +// ahead of the same reduction; the sumsq +// threadgroup also writes the expanded h. +// kq_hc_sinkhorn_collapse: sinkhorn mix normalization plus collapse to +// one stream with the sublayer RMSNorm folded +// into the output. One threadgroup per row. +// kq_hc_expand: pre/comb expand of the sublayer output back +// to four streams. Two threadgroups per row. + +#define KQ_HC 4 +#define KQ_HC_MIX ((2 + KQ_HC) * KQ_HC) +#define KQ_HC_MAX_CHUNKS 8 + +template +[[kernel]] void kq_hc_front_reduce( + const device T* x [[buffer(0)]], + const device float* fn [[buffer(1)]], + device float* mixes_raw [[buffer(2)]], + device float* sumsq [[buffer(3)]], + const constant int& D [[buffer(4)]], + uint tg [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + const uint lane = tid % 32; + const uint sg = tid / 32; + const int KTOT = KQ_HC * D; + + const uint row = tg / (KQ_HC_MIX + 1); + const uint m = tg % (KQ_HC_MIX + 1); + + const device T* xr = x + (int64_t)row * KTOT; + using T4 = vec; + const device T4* x4 = (const device T4*)xr; + + float acc = 0.0f; + if (m < (uint)KQ_HC_MIX) { + const device float4* f4 = (const device float4*)(fn + (int64_t)m * KTOT); + for (uint k = tid; k < (uint)(KTOT / 4); k += 256) { + float4 xv = float4(x4[k]); + float4 fv = f4[k]; + acc = fma(xv.x, fv.x, acc); + acc = fma(xv.y, fv.y, acc); + acc = fma(xv.z, fv.z, acc); + acc = fma(xv.w, fv.w, acc); + } + } else { + for (uint k = tid; k < (uint)(KTOT / 4); k += 256) { + float4 xv = float4(x4[k]); + acc = fma(xv.x, xv.x, acc); + acc = fma(xv.y, xv.y, acc); + acc = fma(xv.z, xv.z, acc); + acc = fma(xv.w, xv.w, acc); + } + } + + threadgroup float partial[8]; + acc = simd_sum(acc); + if (lane == 0) { + partial[sg] = acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float v = (lane < 8) ? partial[lane] : 0.0f; + v = simd_sum(v); + if (lane == 0) { + if (m < (uint)KQ_HC_MIX) { + mixes_raw[row * KQ_HC_MIX + m] = v; + } else { + sumsq[row] = v; + } + } + } +} + +template +[[kernel]] void kq_hc_front_expand_reduce( + const device T* x_sub [[buffer(0)]], + const device T* resid [[buffer(1)]], + const device float* post [[buffer(2)]], + const device float* comb [[buffer(3)]], + const device float* fn [[buffer(4)]], + device T* h_out [[buffer(5)]], + device float* mixes_raw [[buffer(6)]], + device float* sumsq [[buffer(7)]], + const constant int& D [[buffer(8)]], + uint tg [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + const uint lane = tid % 32; + const uint sg = tid / 32; + const int KTOT = KQ_HC * D; + const uint D4q = (uint)D / 4; + + const uint row = tg / (KQ_HC_MIX + 1); + const uint m = tg % (KQ_HC_MIX + 1); + + const device T* xs = x_sub + (int64_t)row * D; + const device T* rr = resid + (int64_t)row * KTOT; + device T* hout = h_out + (int64_t)row * KTOT; + + const uint pb = row * 4, cb = row * 16; + + using T4 = vec; + const device T4* xs4 = (const device T4*)xs; + const device T4* rr4 = (const device T4*)rr; + device T4* h4 = (device T4*)hout; + const device float4* f4 = (const device float4*)(fn + (int64_t)m * KTOT); + + float acc = 0.0f; + for (uint i = 0; i < (uint)KQ_HC; ++i) { + const float pi = post[pb + i]; + const float c0 = comb[cb + 0 * 4 + i]; + const float c1 = comb[cb + 1 * 4 + i]; + const float c2 = comb[cb + 2 * 4 + i]; + const float c3 = comb[cb + 3 * 4 + i]; + for (uint d4 = tid; d4 < D4q; d4 += 256) { + uint k = i * D4q + d4; + float4 xv = float4(xs4[d4]); + float4 r0 = float4(rr4[0 * D4q + d4]); + float4 r1 = float4(rr4[1 * D4q + d4]); + float4 r2 = float4(rr4[2 * D4q + d4]); + float4 r3 = float4(rr4[3 * D4q + d4]); + float4 e = + fma(float4(pi), + xv, + fma(float4(c0), + r0, + fma(float4(c1), r1, fma(float4(c2), r2, float4(c3) * r3)))); + T4 hv = T4(e); + float4 hf = float4(hv); + if (m < (uint)KQ_HC_MIX) { + float4 fv = f4[k]; + acc = fma(hf.x, fv.x, acc); + acc = fma(hf.y, fv.y, acc); + acc = fma(hf.z, fv.z, acc); + acc = fma(hf.w, fv.w, acc); + } else { + h4[k] = hv; + acc = fma(hf.x, hf.x, acc); + acc = fma(hf.y, hf.y, acc); + acc = fma(hf.z, hf.z, acc); + acc = fma(hf.w, hf.w, acc); + } + } + } + + threadgroup float partial[8]; + acc = simd_sum(acc); + if (lane == 0) { + partial[sg] = acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float v = (lane < 8) ? partial[lane] : 0.0f; + v = simd_sum(v); + if (lane == 0) { + if (m < (uint)KQ_HC_MIX) { + mixes_raw[row * KQ_HC_MIX + m] = v; + } else { + sumsq[row] = v; + } + } + } +} + +template +[[kernel]] void kq_hc_sinkhorn_collapse( + const device T* x [[buffer(0)]], + const device float* mixes_raw [[buffer(1)]], + const device float* sumsq [[buffer(2)]], + const device float* scale [[buffer(3)]], + const device float* base [[buffer(4)]], + const device T* w [[buffer(5)]], + device T* collapsed [[buffer(6)]], + device float* post [[buffer(7)]], + device float* comb [[buffer(8)]], + const constant int& D [[buffer(9)]], + const constant int& iters [[buffer(10)]], + const constant float& hc_eps [[buffer(11)]], + const constant float& norm_eps [[buffer(12)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + const uint lane = tid % 32; + const uint sg = tid / 32; + const int BASE_OFF = 2 * KQ_HC; + const float EPS = hc_eps; + const float NEPS = norm_eps; + + const device float* mix = mixes_raw + row * KQ_HC_MIX; + device float* post_out = post + row * KQ_HC; + device float* comb_out = comb + row * KQ_HC * KQ_HC; + + const float factor = metal::rsqrt(sumsq[row] / (float)(KQ_HC * D) + NEPS); + + threadgroup float pre_shared[KQ_HC]; + threadgroup float ssq_shared[8]; + threadgroup float inv_shared[1]; + + if (sg == 0) { + const float pre_scale = scale[0] * factor; + const float post_scale = scale[1] * factor; + const float comb_scale = scale[2] * factor; + + const float active = (lane < (uint)KQ_HC) ? 1.0f : 0.0f; + const uint llane = metal::min(lane, (uint)(KQ_HC - 1)); + + float pre_z = mix[llane] * pre_scale + base[llane]; + float post_z = mix[KQ_HC + llane] * post_scale + base[KQ_HC + llane]; + float pre_v = 1.0f / (1.0f + metal::fast::exp(-pre_z)) + EPS; + float post_v = 2.0f / (1.0f + metal::fast::exp(-post_z)); + + if (lane < (uint)KQ_HC) { + pre_shared[lane] = pre_v; + post_out[lane] = post_v; + } + + float4 v = + (*(const device float4*)(mix + BASE_OFF + llane * KQ_HC) * comb_scale + + *(const device float4*)(base + BASE_OFF + llane * KQ_HC)) * + active; + + float row_max = metal::max(metal::max(v.x, v.y), metal::max(v.z, v.w)); + float4 e = metal::fast::exp(v - row_max) * active; + float4 r = e * (1.0f / (e.x + e.y + e.z + e.w + EPS)) + EPS * active; + + float4 col_inv = 1.0f / + (float4(simd_sum(r.x), simd_sum(r.y), simd_sum(r.z), simd_sum(r.w)) + + EPS); + r *= col_inv; + + for (int iter = 1; iter < iters; ++iter) { + r *= (1.0f / (r.x + r.y + r.z + r.w + EPS)) * active; + col_inv = 1.0f / + (float4(simd_sum(r.x), simd_sum(r.y), simd_sum(r.z), simd_sum(r.w)) + + EPS); + r *= col_inv; + } + + if (lane < (uint)KQ_HC) { + *(device float4*)(comb_out + lane * KQ_HC) = r; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float p0 = pre_shared[0]; + const float p1 = pre_shared[1]; + const float p2 = pre_shared[2]; + const float p3 = pre_shared[3]; + + const device T* x_row = x + (int64_t)row * (KQ_HC * D); + device T* out_row = collapsed + (int64_t)row * D; + + using T4 = vec; + const device T4* x_row0 = (const device T4*)(x_row + 0 * D); + const device T4* x_row1 = (const device T4*)(x_row + 1 * D); + const device T4* x_row2 = (const device T4*)(x_row + 2 * D); + const device T4* x_row3 = (const device T4*)(x_row + 3 * D); + device T4* out4 = (device T4*)out_row; + + const uint D4 = (uint)D / 4; + const uint chunks = (D4 + 255) / 256; + + float4 vals[KQ_HC_MAX_CHUNKS]; + float ssq = 0.0f; + for (uint c = 0; c < chunks; ++c) { + uint d4 = c * 256 + tid; + float4 result = float4(0.0f); + if (d4 < D4) { + float4 x0 = float4(x_row0[d4]); + float4 x1 = float4(x_row1[d4]); + float4 x2 = float4(x_row2[d4]); + float4 x3 = float4(x_row3[d4]); + result = + fma(float4(p0), + x0, + fma(float4(p1), x1, fma(float4(p2), x2, float4(p3) * x3))); + ssq += result.x * result.x + result.y * result.y + result.z * result.z + + result.w * result.w; + } + vals[c] = result; + } + + ssq = simd_sum(ssq); + if (lane == 0) { + ssq_shared[sg] = ssq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float v = (lane < 8) ? ssq_shared[lane] : 0.0f; + v = simd_sum(v); + if (lane == 0) { + inv_shared[0] = metal::rsqrt(v / (float)D + NEPS); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float inv = inv_shared[0]; + + for (uint c = 0; c < chunks; ++c) { + uint d4 = c * 256 + tid; + if (d4 < D4) { + uint d = d4 * 4; + float4 wv = float4( + (float)w[d], (float)w[d + 1], (float)w[d + 2], (float)w[d + 3]); + out4[d4] = T4(vals[c] * inv * wv); + } + } +} + +template +[[kernel]] void kq_hc_expand( + const device T* x [[buffer(0)]], + const device T* resid [[buffer(1)]], + const device float* post [[buffer(2)]], + const device float* comb [[buffer(3)]], + device T* out [[buffer(4)]], + const constant int& D [[buffer(5)]], + uint tg [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + const uint NTG = 2; + const uint row = tg / NTG; + const uint sub = tg % NTG; + + const device T* xr = x + (int64_t)row * D; + const device T* rr = resid + (int64_t)row * 4 * D; + device T* orow = out + (int64_t)row * 4 * D; + + const uint pb = row * 4, cb = row * 16; + float p0 = post[pb + 0], p1 = post[pb + 1]; + float p2 = post[pb + 2], p3 = post[pb + 3]; + // comb is [j][i]; expand applies comb^T: sum_j comb[j][i] * res[j] + float c00 = comb[cb + 0], c01 = comb[cb + 1]; + float c02 = comb[cb + 2], c03 = comb[cb + 3]; + float c10 = comb[cb + 4], c11 = comb[cb + 5]; + float c12 = comb[cb + 6], c13 = comb[cb + 7]; + float c20 = comb[cb + 8], c21 = comb[cb + 9]; + float c22 = comb[cb + 10], c23 = comb[cb + 11]; + float c30 = comb[cb + 12], c31 = comb[cb + 13]; + float c32 = comb[cb + 14], c33 = comb[cb + 15]; + + const uint SPAN = (uint)D / NTG; + const uint d0 = sub * SPAN; + using T4 = vec; + const device T4* x4 = (const device T4*)(xr + d0); + const device T4* r04 = (const device T4*)(rr + 0 * D + d0); + const device T4* r14 = (const device T4*)(rr + 1 * D + d0); + const device T4* r24 = (const device T4*)(rr + 2 * D + d0); + const device T4* r34 = (const device T4*)(rr + 3 * D + d0); + device T4* o04 = (device T4*)(orow + 0 * D + d0); + device T4* o14 = (device T4*)(orow + 1 * D + d0); + device T4* o24 = (device T4*)(orow + 2 * D + d0); + device T4* o34 = (device T4*)(orow + 3 * D + d0); + for (uint k = tid; k < SPAN / 4; k += 256) { + float4 xv = float4(x4[k]); + float4 r0 = float4(r04[k]), r1 = float4(r14[k]); + float4 r2 = float4(r24[k]), r3 = float4(r34[k]); + o04[k] = T4( + fma(float4(p0), + xv, + fma(float4(c00), + r0, + fma(float4(c10), r1, fma(float4(c20), r2, float4(c30) * r3))))); + o14[k] = T4( + fma(float4(p1), + xv, + fma(float4(c01), + r0, + fma(float4(c11), r1, fma(float4(c21), r2, float4(c31) * r3))))); + o24[k] = T4( + fma(float4(p2), + xv, + fma(float4(c02), + r0, + fma(float4(c12), r1, fma(float4(c22), r2, float4(c32) * r3))))); + o34[k] = T4( + fma(float4(p3), + xv, + fma(float4(c03), + r0, + fma(float4(c13), r1, fma(float4(c23), r2, float4(c33) * r3))))); + } +} diff --git a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h index a70306e..feed7a8 100644 --- a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h +++ b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h @@ -730,7 +730,9 @@ template // --------------------------------------------------------------------------- // Codec-matrix kernels: one generic implementation per family, templated on // the Ext codec traits from kq_quantized*.h (superblock, block_bytes, -// deq_chunk16(block, il, reg) -> 16 weights in natural order). Thread +// deq_chunk16(block, il, reg) -> 16 weights in natural order). Row loops +// use the deq_chunk16s form and fold the returned scale once per chunk +// dot. Thread // mapping follows kq_mv_ext_impl, templated on the K-lane width NX: the 32 // simdgroup lanes split into NX K-lanes x (32 / NX) output rows (each thread // owns one row); the K-reduction is a log2(NX)-step simd_shuffle_down within @@ -773,12 +775,14 @@ METAL_FUNC float kq_ext_row_partial( const device uint8_t* block = w_row + (int64_t)(ich / chpb) * Codec::block_bytes; float4x4 lw; - KqTgLuts::deq_chunk16(block, short(ich % chpb), lw, luts); + float sc; + KqTgLuts::deq_chunk16s(block, short(ich % chpb), lw, luts, sc); const device T* xp = x + ich * 16; - acc += dot(lw[0], float4(*(const device vec*)(xp + 0))) + - dot(lw[1], float4(*(const device vec*)(xp + 4))) + - dot(lw[2], float4(*(const device vec*)(xp + 8))) + - dot(lw[3], float4(*(const device vec*)(xp + 12))); + acc += sc * + (dot(lw[0], float4(*(const device vec*)(xp + 0))) + + dot(lw[1], float4(*(const device vec*)(xp + 4))) + + dot(lw[2], float4(*(const device vec*)(xp + 8))) + + dot(lw[3], float4(*(const device vec*)(xp + 12)))); } return acc; } @@ -808,10 +812,13 @@ METAL_FUNC float2 kq_ext_glu_row_partial( const float4 a2 = float4(*(const device vec*)(xp + 8)); const float4 a3 = float4(*(const device vec*)(xp + 12)); float4x4 lw; - KqTgLuts::deq_chunk16(g_row + boff, cch, lw, luts); - acc.x += dot(lw[0], a0) + dot(lw[1], a1) + dot(lw[2], a2) + dot(lw[3], a3); - KqTgLuts::deq_chunk16(u_row + boff, cch, lw, luts); - acc.y += dot(lw[0], a0) + dot(lw[1], a1) + dot(lw[2], a2) + dot(lw[3], a3); + float sc; + KqTgLuts::deq_chunk16s(g_row + boff, cch, lw, luts, sc); + acc.x += sc * + (dot(lw[0], a0) + dot(lw[1], a1) + dot(lw[2], a2) + dot(lw[3], a3)); + KqTgLuts::deq_chunk16s(u_row + boff, cch, lw, luts, sc); + acc.y += sc * + (dot(lw[0], a0) + dot(lw[1], a1) + dot(lw[2], a2) + dot(lw[3], a3)); } return acc; } @@ -1076,6 +1083,72 @@ template } } +// Slot-parallel mix_ns: same math as kq_ext_gather_qmv_mix_ns with the S +// slot dots spread across S simdgroup pairs instead of a per-thread loop. +// The loop kernel launches N / 8 threadgroups at decode (T = 1); the +// per-thread slot loop leaves the device underfilled and the solo op runs +// at ~2/3 of its cross-call-overlapped bandwidth -- chained probe calls +// recover the gap, the real serialized decode graph does not. Widening +// K-lanes (NX = 16/32) shortens per-thread chains and measured +// flat-to-negative; this mapping keeps the chunk chains at NX = 8 length +// and multiplies resident threads by S. Each simdgroup owns one +// (slot, row-half); raw lane partials stage through threadgroup memory and +// the slot-0 simdgroup pair replays the loop kernel's serial score-FMA +// chain and 3-step lane reduce, so outputs are bit-identical to it. +// Dispatch: group (32, 2 * S, 1), grid (N / (2 * RPS), 1, T); host gates +// S <= KQ_MOE_SP_MAX_S. +#define KQ_MOE_SP_MAX_S 16 + +template +[[kernel]] void kq_ext_gather_qmv_mix_ns_sp( + const device uint8_t* w [[buffer(0)]], + const device T* h [[buffer(1)]], + const device uint32_t* indices [[buffer(2)]], + const device float* scores [[buffer(3)]], + device T* out [[buffer(4)]], + const constant int& K [[buffer(5)]], + const constant int& N [[buffer(6)]], + const constant int& S [[buffer(7)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tptg [[threads_per_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int RPS = 32 / NX; + const short tx = short(simd_lid % NX); + const short ty = short(simd_lid / NX); + const int slot = int(simd_gid >> 1); + const short lrow = short((simd_gid & 1) * RPS) + ty; + const int out_row = tid.x * (2 * RPS) + lrow; + + threadgroup uint4 kq_luts_v[(KqTgLuts::bytes + 15) / 16 + 1]; + threadgroup uint8_t* kq_luts = + reinterpret_cast(kq_luts_v); + if (KqTgLuts::bytes > 0) { + KqTgLuts::stage( + kq_luts, ushort(simd_gid * 32 + simd_lid), ushort(tptg.x * tptg.y)); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + threadgroup float parts[2 * RPS][NX][KQ_MOE_SP_MAX_S]; + + const int expert = int(indices[tid.z * S + slot]); + const device T* xs = h + ((int64_t)tid.z * S + slot) * K; + parts[lrow][tx][slot] = kq_ext_row_partial( + w, xs, (int64_t)expert * N + out_row, K, tx, kq_luts); + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (slot == 0) { + float result = 0.0f; + for (int s = 0; s < S; s++) { + result += scores[tid.z * S + s] * parts[lrow][tx][s]; + } + result = kq_ext_reduce(result); + if (tx == 0) { + out[(int64_t)tid.z * N + out_row] = static_cast(result); + } + } +} + template [[kernel]] void kq_ext_gather_qmv_mix( const device uint8_t* w [[buffer(0)]], diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index 36f9841..0d0449d 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -564,6 +564,21 @@ METAL_FUNC void kq_adjust_matrix_offsets( y += tid.z * output_stride; } +// Scaled chunk dequant for the ext mat-vec loops, reading constant-space +// tables. Same contract as KqTgLuts::deq_chunk16s without staging. Codecs +// with no hoistable scale pass scale = 1 and unchanged weights. +template +struct KqExtDeq { + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + thread float& scale) { + Codec::deq_chunk16(block, il, reg); + scale = 1.0f; + } +}; + // --------------------------------------------------------------------------- // Flat-with-M verify mat-vec (port of ggml-metal kernel_mul_mv_ext_q4x4), // codec-agnostic. The per-row qmv puts M on grid_dims.x, so each of the M rows @@ -627,7 +642,8 @@ METAL_FUNC void kq_mv_ext_impl( const device uint8_t* block = w_row + static_cast(ib) * Codec::block_bytes; float4x4 lx; - Codec::deq_chunk16(block, cch, lx); + float sc; + KqExtDeq::deq_chunk16s(block, cch, lx, sc); #pragma unroll for (short ir1 = 0; ir1 < r1ptg; ++ir1) { const device T* yp = y_col[ir1]; @@ -635,8 +651,8 @@ METAL_FUNC void kq_mv_ext_impl( const float4 a1 = float4(*(const device vec*)(yp + 4)); const float4 a2 = float4(*(const device vec*)(yp + 8)); const float4 a3 = float4(*(const device vec*)(yp + 12)); - sumf[ir1] += - dot(lx[0], a0) + dot(lx[1], a1) + dot(lx[2], a2) + dot(lx[3], a3); + sumf[ir1] += sc * + (dot(lx[0], a0) + dot(lx[1], a1) + dot(lx[2], a2) + dot(lx[3], a3)); y_col[ir1] += nxpsg * 16; } } diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index 91781f6..383b7f9 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -191,6 +191,35 @@ struct KqIq3_xxsExt { } }; +template <> +struct KqExtDeq { + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + thread float& scale) { + const int ib32 = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const device uint8_t* qs = block + KQ_IQ3_XXS_QS_OFFSET + ib32 * 8; + const device uint8_t* gas = block + KQ_IQ3_XXS_GAS_OFFSET + ib32 * 4; + const uint32_t aux32 = uint32_t(gas[0]) | (uint32_t(gas[1]) << 8) | + (uint32_t(gas[2]) << 16) | (uint32_t(gas[3]) << 24); + scale = d * (0.5f + float(aux32 >> 28)) * 0.5f; +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; + const float4 v0 = float4(as_type(iq3xxs_grid[qs[2 * l]])); + const float4 v1 = float4(as_type(iq3xxs_grid[qs[2 * l + 1]])); + reg[2 * t] = + select(v0, -v0, bool4(signs & 1, signs & 2, signs & 4, signs & 8)); + reg[2 * t + 1] = select( + v1, -v1, bool4(signs & 16, signs & 32, signs & 64, signs & 128)); + } + } +}; + template [[kernel]] void kq_iq3_xxs_mv_ext( const device uint8_t* w, @@ -306,6 +335,39 @@ struct KqIq3_sExt { } }; +template <> +struct KqExtDeq { + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + thread float& scale) { + const int s = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const device uint8_t* qs = block + KQ_IQ3_S_QS_OFFSET + s * 8; + const device uint8_t* signs = block + KQ_IQ3_S_SIGNS_OFFSET + s * 4; + const device uint8_t* scales = block + KQ_IQ3_S_SCALES_OFFSET; + const int sc_nib = (scales[s / 2] >> (4 * (s & 1))) & 0xf; + scale = d * float(1 + 2 * sc_nib); + const int qhb = int(block[KQ_IQ3_S_QH_OFFSET + s]); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const int hi0 = (qhb << (8 - 2 * l)) & 256; + const int hi1 = (qhb << (7 - 2 * l)) & 256; + const float4 v0 = + float4(as_type(iq3s_grid[int(qs[2 * l]) | hi0])); + const float4 v1 = + float4(as_type(iq3s_grid[int(qs[2 * l + 1]) | hi1])); + const uint8_t sb = signs[l]; + reg[2 * t] = select(v0, -v0, bool4(sb & 1, sb & 2, sb & 4, sb & 8)); + reg[2 * t + 1] = + select(v1, -v1, bool4(sb & 16, sb & 32, sb & 64, sb & 128)); + } + } +}; + template [[kernel]] void kq_iq3_s_mv_ext( const device uint8_t* w, @@ -407,6 +469,38 @@ struct KqIq2_xxsExt { } }; +template <> +struct KqExtDeq { + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + thread float& scale) { + const int ib32 = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const device uint8_t* qs = block + KQ_IQ2_XXS_QS_OFFSET + ib32 * 8; + const uint32_t signbits = uint32_t(qs[4]) | (uint32_t(qs[5]) << 8) | + (uint32_t(qs[6]) << 16) | (uint32_t(qs[7]) << 24); + scale = d * (0.5f + float(signbits >> 28)) * 0.25f; + // One u64 grid load + vector uchar4 -> float4 conversions + selects + // replace the per-byte load/convert/select chain. Integer-exact: + // outputs bit-identical to the scalar form. +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const uint64_t ge = iq2xxs_grid[qs[l]]; + const uint8_t signs = ksigns_iq2xs[(signbits >> (7 * l)) & 127]; + const float4 v0 = float4(as_type(uint32_t(ge))); + const float4 v1 = float4(as_type(uint32_t(ge >> 32))); + reg[2 * t] = + select(v0, -v0, bool4(signs & 1, signs & 2, signs & 4, signs & 8)); + reg[2 * t + 1] = select( + v1, -v1, bool4(signs & 16, signs & 32, signs & 64, signs & 128)); + } + } +}; + template [[kernel]] void kq_iq2_xxs_mv_ext( const device uint8_t* w, @@ -512,6 +606,36 @@ struct KqIq2_xsExt { } }; +template <> +struct KqExtDeq { + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + thread float& scale) { + const int ib32 = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const uint8_t sc = block[KQ_IQ2_XS_SCALES_OFFSET + ib32]; + const int sc_nib = (lbase < 2) ? (sc & 0xf) : (sc >> 4); + scale = d * (0.5f + float(sc_nib)) * 0.25f; + const device uint8_t* qp = block + KQ_IQ2_XS_QS_OFFSET + ib32 * 8; +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const uint q = uint(qp[2 * l]) | (uint(qp[2 * l + 1]) << 8); + const uint64_t ge = iq2xs_grid[q & 511]; + const uint8_t signs = ksigns_iq2xs[q >> 9]; + const float4 v0 = float4(as_type(uint32_t(ge))); + const float4 v1 = float4(as_type(uint32_t(ge >> 32))); + reg[2 * t] = + select(v0, -v0, bool4(signs & 1, signs & 2, signs & 4, signs & 8)); + reg[2 * t + 1] = select( + v1, -v1, bool4(signs & 16, signs & 32, signs & 64, signs & 128)); + } + } +}; + template [[kernel]] void kq_iq2_xs_mv_ext( const device uint8_t* w, @@ -623,6 +747,37 @@ struct KqIq2_sExt { } }; +template <> +struct KqExtDeq { + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + thread float& scale) { + const int ib32 = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const device uint8_t* qs = block + KQ_IQ2_S_QS_OFFSET; + const device uint8_t* signs = block + KQ_IQ2_S_SIGNS_OFFSET; + const uint8_t sc = block[KQ_IQ2_S_SCALES_OFFSET + ib32]; + const int sc_nib = (lbase < 2) ? (sc & 0xf) : (sc >> 4); + scale = d * (0.5f + float(sc_nib)) * 0.25f; + const int qhb = int(block[KQ_IQ2_S_QH_OFFSET + ib32]); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const int qi = int(qs[ib32 * 4 + l]) | ((qhb << (8 - 2 * l)) & 0x300); + const uint64_t ge = iq2s_grid[qi]; + const uint8_t sb = signs[ib32 * 4 + l]; + const float4 v0 = float4(as_type(uint32_t(ge))); + const float4 v1 = float4(as_type(uint32_t(ge >> 32))); + reg[2 * t] = select(v0, -v0, bool4(sb & 1, sb & 2, sb & 4, sb & 8)); + reg[2 * t + 1] = + select(v1, -v1, bool4(sb & 16, sb & 32, sb & 64, sb & 128)); + } + } +}; + template [[kernel]] void kq_iq2_s_mv_ext( const device uint8_t* w, @@ -4517,6 +4672,10 @@ template // memory once per threadgroup removes that latency; kmask_iq2xs is 1 << j // and folds away. Codecs without LUTs keep bytes = 0 and the passthrough // three-argument deq. +// +// deq_chunk16s writes the weights sign-applied without the block scale +// and returns the scale for the caller to fold once per chunk dot. +// Codecs with no hoistable scale pass scale = 1 and unchanged weights. template struct KqTgLuts { MLX_MTL_CONST int bytes = 0; @@ -4528,6 +4687,15 @@ struct KqTgLuts { const threadgroup uint8_t*) { Codec::deq_chunk16(block, il, reg); } + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t*, + thread float& scale) { + Codec::deq_chunk16(block, il, reg); + scale = 1.0f; + } }; template <> @@ -4547,11 +4715,12 @@ struct KqTgLuts { d32[512 + i] = signs32[i]; } } - static METAL_FUNC void deq_chunk16( + static METAL_FUNC void deq_chunk16s( const device uint8_t* block, short il, thread float4x4& reg, - const threadgroup uint8_t* luts) { + const threadgroup uint8_t* luts, + thread float& scale) { const threadgroup uint64_t* grid = reinterpret_cast(luts); const threadgroup uint8_t* ksigns = luts + 2048; @@ -4561,15 +4730,30 @@ struct KqTgLuts { const device uint8_t* qs = block + KQ_IQ2_XXS_QS_OFFSET + ib32 * 8; const uint32_t signbits = uint32_t(qs[4]) | (uint32_t(qs[5]) << 8) | (uint32_t(qs[6]) << 16) | (uint32_t(qs[7]) << 24); - const float db = d * (0.5f + float(signbits >> 28)) * 0.25f; + scale = d * (0.5f + float(signbits >> 28)) * 0.25f; #pragma unroll - for (int i = 0; i < 16; ++i) { - const int l = lbase + i / 8; - const int j = i % 8; + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const threadgroup uint8_t* gb = + reinterpret_cast(grid + qs[l]); const uint8_t signs = ksigns[(signbits >> (7 * l)) & 127]; - const uint8_t gb = (grid[qs[l]] >> (8 * j)) & 0xff; - const float sgn = (signs & (1 << j)) ? -1.0f : 1.0f; - reg[i / 4][i % 4] = db * float(gb) * sgn; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float v = float(gb[j]); + reg[2 * t + j / 4][j % 4] = (signs & (1 << j)) ? -v : v; + } + } + } + static METAL_FUNC void deq_chunk16( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts) { + float scale; + deq_chunk16s(block, il, reg, luts, scale); +#pragma unroll + for (int i = 0; i < 4; ++i) { + reg[i] *= scale; } } }; @@ -4589,11 +4773,12 @@ struct KqTgLuts { d32[256 + i] = signs32[i]; } } - static METAL_FUNC void deq_chunk16( + static METAL_FUNC void deq_chunk16s( const device uint8_t* block, short il, thread float4x4& reg, - const threadgroup uint8_t* luts) { + const threadgroup uint8_t* luts, + thread float& scale) { const threadgroup uint32_t* grid = reinterpret_cast(luts); const threadgroup uint8_t* ksigns = luts + 1024; @@ -4604,17 +4789,206 @@ struct KqTgLuts { const device uint8_t* gas = block + KQ_IQ3_XXS_GAS_OFFSET + ib32 * 4; const uint32_t aux32 = uint32_t(gas[0]) | (uint32_t(gas[1]) << 8) | (uint32_t(gas[2]) << 16) | (uint32_t(gas[3]) << 24); - const float db = d * (0.5f + float(aux32 >> 28)) * 0.5f; + scale = d * (0.5f + float(aux32 >> 28)) * 0.5f; #pragma unroll - for (int i = 0; i < 16; ++i) { - const int l = lbase + i / 8; - const int sub = i % 8; + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; const uint8_t signs = ksigns[(aux32 >> (7 * l)) & 127]; - const int qi = (sub < 4) ? int(qs[2 * l]) : int(qs[2 * l + 1]); - const int bytej = (sub < 4) ? sub : (sub - 4); - const uint8_t gb = (grid[qi] >> (8 * bytej)) & 0xff; - const float sgn = (signs & (1 << sub)) ? -1.0f : 1.0f; - reg[i / 4][i % 4] = db * float(gb) * sgn; + const threadgroup uint8_t* g0 = + reinterpret_cast(grid + qs[2 * l]); + const threadgroup uint8_t* g1 = + reinterpret_cast(grid + qs[2 * l + 1]); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float v0 = float(g0[j]); + const float v1 = float(g1[j]); + reg[2 * t][j] = (signs & (1 << j)) ? -v0 : v0; + reg[2 * t + 1][j] = (signs & (1 << (j + 4))) ? -v1 : v1; + } + } + } + static METAL_FUNC void deq_chunk16( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts) { + float scale; + deq_chunk16s(block, il, reg, luts, scale); +#pragma unroll + for (int i = 0; i < 4; ++i) { + reg[i] *= scale; + } + } +}; + +template <> +struct KqTgLuts { + MLX_MTL_CONST int bytes = 4096 + 128; // u64 grid[512] | ksigns[128] + static METAL_FUNC void + stage(threadgroup uint8_t* dst, ushort lin, ushort n_threads) { + threadgroup uint32_t* d32 = reinterpret_cast(dst); + const constant uint32_t* grid32 = + reinterpret_cast(iq2xs_grid); + for (int i = lin; i < 1024; i += n_threads) { + d32[i] = grid32[i]; + } + const constant uint32_t* signs32 = + reinterpret_cast(ksigns_iq2xs); + for (int i = lin; i < 32; i += n_threads) { + d32[1024 + i] = signs32[i]; + } + } + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts, + thread float& scale) { + const threadgroup uint64_t* grid = + reinterpret_cast(luts); + const threadgroup uint8_t* ksigns = luts + 4096; + const int ib32 = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const uint8_t sc = block[KQ_IQ2_XS_SCALES_OFFSET + ib32]; + const int sc_nib = (lbase < 2) ? (sc & 0xf) : (sc >> 4); + scale = d * (0.5f + float(sc_nib)) * 0.25f; + const device uint8_t* qp = block + KQ_IQ2_XS_QS_OFFSET + ib32 * 8; +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const uint q = uint(qp[2 * l]) | (uint(qp[2 * l + 1]) << 8); + const threadgroup uint8_t* gb = + reinterpret_cast(grid + (q & 511)); + const uint8_t signs = ksigns[q >> 9]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float v = float(gb[j]); + reg[2 * t + j / 4][j % 4] = (signs & (1 << j)) ? -v : v; + } + } + } + static METAL_FUNC void deq_chunk16( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts) { + float scale; + deq_chunk16s(block, il, reg, luts, scale); +#pragma unroll + for (int i = 0; i < 4; ++i) { + reg[i] *= scale; + } + } +}; + +template <> +struct KqTgLuts { + MLX_MTL_CONST int bytes = 2048; // u32 grid[512], signs live in the block + static METAL_FUNC void + stage(threadgroup uint8_t* dst, ushort lin, ushort n_threads) { + threadgroup uint32_t* d32 = reinterpret_cast(dst); + for (int i = lin; i < 512; i += n_threads) { + d32[i] = iq3s_grid[i]; + } + } + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts, + thread float& scale) { + const threadgroup uint32_t* grid = + reinterpret_cast(luts); + const int s = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const device uint8_t* qs = block + KQ_IQ3_S_QS_OFFSET + s * 8; + const device uint8_t* signs = block + KQ_IQ3_S_SIGNS_OFFSET + s * 4; + const device uint8_t* scales = block + KQ_IQ3_S_SCALES_OFFSET; + const int sc_nib = (scales[s / 2] >> (4 * (s & 1))) & 0xf; + scale = d * float(1 + 2 * sc_nib); + const int qhb = int(block[KQ_IQ3_S_QH_OFFSET + s]); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const int hi0 = (qhb << (8 - 2 * l)) & 256; + const int hi1 = (qhb << (7 - 2 * l)) & 256; + const threadgroup uint8_t* g0 = + reinterpret_cast( + grid + (int(qs[2 * l]) | hi0)); + const threadgroup uint8_t* g1 = + reinterpret_cast( + grid + (int(qs[2 * l + 1]) | hi1)); + const uint8_t sb = signs[l]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float v0 = float(g0[j]); + const float v1 = float(g1[j]); + reg[2 * t][j] = (sb & (1 << j)) ? -v0 : v0; + reg[2 * t + 1][j] = (sb & (1 << (j + 4))) ? -v1 : v1; + } + } + } + static METAL_FUNC void deq_chunk16( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts) { + float scale; + deq_chunk16s(block, il, reg, luts, scale); +#pragma unroll + for (int i = 0; i < 4; ++i) { + reg[i] *= scale; + } + } +}; + +// iq2_s keeps its 8 KB grid in constant memory (staging it would cost +// occupancy) and hoists only the scale. +template <> +struct KqTgLuts { + MLX_MTL_CONST int bytes = 0; + static METAL_FUNC void stage(threadgroup uint8_t*, ushort, ushort) {} + static METAL_FUNC void deq_chunk16s( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t*, + thread float& scale) { + const int ib32 = il / 2; + const int lbase = (il & 1) * 2; + const float d = float(*(const device half*)block); + const device uint8_t* qs = block + KQ_IQ2_S_QS_OFFSET; + const device uint8_t* signs = block + KQ_IQ2_S_SIGNS_OFFSET; + const uint8_t sc = block[KQ_IQ2_S_SCALES_OFFSET + ib32]; + const int sc_nib = (lbase < 2) ? (sc & 0xf) : (sc >> 4); + scale = d * (0.5f + float(sc_nib)) * 0.25f; + const int qhb = int(block[KQ_IQ2_S_QH_OFFSET + ib32]); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int l = lbase + t; + const int qi = int(qs[ib32 * 4 + l]) | ((qhb << (8 - 2 * l)) & 0x300); + const constant uint8_t* gb = + reinterpret_cast(iq2s_grid + qi); + const uint8_t sb = signs[ib32 * 4 + l]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float v = float(gb[j]); + reg[2 * t + j / 4][j % 4] = (sb & (1 << j)) ? -v : v; + } + } + } + static METAL_FUNC void deq_chunk16( + const device uint8_t* block, + short il, + thread float4x4& reg, + const threadgroup uint8_t* luts) { + float scale; + deq_chunk16s(block, il, reg, luts, scale); +#pragma unroll + for (int i = 0; i < 4; ++i) { + reg[i] *= scale; } } }; diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index 133276d..3c811aa 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -48,6 +48,11 @@ gather_qmv_mix_bias, gather_qmv_mix_kq, gather_qmv_mix_ns_kq, + get_cb_caps, + hc_expand, + hc_front_expand_reduce, + hc_front_reduce, + hc_sinkhorn_collapse, load_gguf, metallib_dir, metallib_loads, @@ -71,6 +76,7 @@ sdpa_decode_gqa_paged, sdpa_fa_verify, sdpa_vector, + set_cb_caps, shared_event_create, shared_event_destroy, shared_event_read, @@ -110,6 +116,11 @@ "gather_qmv_kq", "gather_qmv_mix_kq", "gather_qmv_mix_ns_kq", + "get_cb_caps", + "hc_expand", + "hc_front_expand_reduce", + "hc_front_reduce", + "hc_sinkhorn_collapse", "load_gguf", "metallib_dir", "metallib_loads", @@ -130,6 +141,7 @@ "sdpa_decode_gqa_paged", "sdpa_fa_verify", "sdpa_vector", + "set_cb_caps", "skinny_matmul", "residency_commit", "residency_erase", diff --git a/src/kquant.h b/src/kquant.h index 42fea2f..f330cfa 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -5,6 +5,7 @@ #include #include +#include #include #include "mlx/ops.h" @@ -611,6 +612,51 @@ mx::array rmsnorm2_add( float eps, mx::StreamOrDevice s = {}); +// Fused deepseek4 hyper-connection glue for the single-token decode route +// (hc_mult 4 only; see kq_hc_glue.h). x is [..., 4, D] with D % 8 == 0 and +// D <= 8192; fn is float32 [24, 4 * D]. Returns {mixes_raw f32 [..., 24], +// sumsq f32 [..., 1]}. +std::vector +hc_front_reduce(mx::array x, mx::array fn, mx::StreamOrDevice s = {}); + +// The previous cycle's expand fused ahead of the same front reduction. +// x_sub is [..., D], resid [..., 4, D], post f32 [..., 4], comb f32 +// [..., 4, 4]. Returns {h [..., 4, D], mixes_raw f32 [..., 24], +// sumsq f32 [..., 1]}; h is bit-identical to hc_expand of the same carry. +std::vector hc_front_expand_reduce( + mx::array x_sub, + mx::array resid, + mx::array post, + mx::array comb, + mx::array fn, + mx::StreamOrDevice s = {}); + +// Sinkhorn mix normalization + collapse to one stream with the sublayer +// RMSNorm (weight w, eps norm_eps) folded into the single output rounding. +// The deferred front rms factor enters through sumsq. scale is f32 [3], +// base f32 [24]. Returns {collapsed [..., D], post f32 [..., 4], +// comb f32 [..., 4, 4]}. +std::vector hc_sinkhorn_collapse( + mx::array x, + mx::array mixes_raw, + mx::array sumsq, + mx::array scale, + mx::array base, + mx::array w, + int iters, + float hc_eps, + float norm_eps, + mx::StreamOrDevice s = {}); + +// Expand the sublayer output x [..., D] back over resid [..., 4, D] with +// the pre/comb coefficients. Returns [..., 4, D]. +mx::array hc_expand( + mx::array x, + mx::array resid, + mx::array post, + mx::array comb, + mx::StreamOrDevice s = {}); + // y = x @ w.T for token widths 1..16 against a small-N, large-K weight in // nn.Linear layout ([N, K]). x is [..., M, K] with 1 <= M <= 16 and // K % 4 == 0; x float16/bfloat16/float32, w matching x or float32; output @@ -1549,6 +1595,96 @@ class KQuantRMSNorm2Add : public mx::Primitive { float eps_; }; +// Fused hyper-connection glue primitives (see hc_front_reduce and friends). +// Inference-only, GPU-only. +class KQuantHcFrontReduce : public mx::Primitive { + public: + explicit KQuantHcFrontReduce(mx::Stream stream) : mx::Primitive(stream) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "KQuantHcFrontReduce"; + } + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } +}; + +class KQuantHcFrontExpandReduce : public mx::Primitive { + public: + explicit KQuantHcFrontExpandReduce(mx::Stream stream) + : mx::Primitive(stream) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "KQuantHcFrontExpandReduce"; + } + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } +}; + +class KQuantHcSinkhornCollapse : public mx::Primitive { + public: + explicit KQuantHcSinkhornCollapse( + mx::Stream stream, + int iters, + float hc_eps, + float norm_eps) + : mx::Primitive(stream), + iters_(iters), + hc_eps_(hc_eps), + norm_eps_(norm_eps) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "KQuantHcSinkhornCollapse"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + private: + int iters_; + float hc_eps_; + float norm_eps_; +}; + +class KQuantHcExpand : public mx::Primitive { + public: + explicit KQuantHcExpand(mx::Stream stream) : mx::Primitive(stream) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "KQuantHcExpand"; + } + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } +}; + // Skinny matmul y = x @ w.T at token widths 1..16 (see skinny_matmul). // Inference-only. class KQuantSkinnyMV : public mx::Primitive { @@ -1861,4 +1997,10 @@ class KQuantEventWait : public mx::Primitive { uint64_t value_; }; +// Runtime read/write of MLX's command-buffer split caps (ops, MB per +// buffer). set returns the previous pair. GPU-only; see +// kquant_cb_caps.cpp for the phase-flip rationale. +std::pair get_cb_caps(); +std::pair set_cb_caps(int max_ops, int max_mb); + } // namespace mlx_kquant diff --git a/src/kquant_cb_caps.cpp b/src/kquant_cb_caps.cpp new file mode 100644 index 0000000..a812be7 --- /dev/null +++ b/src/kquant_cb_caps.cpp @@ -0,0 +1,70 @@ +// Runtime access to MLX's command-buffer split caps (max ops and MB per +// buffer). The MLX_MAX_*_PER_BUFFER env knobs latch at Metal device init, +// but decode wants coarse buffers (submission otherwise blocks on +// in-flight drain) while deep prefill needs fine ones (a giant buffer +// holds every layer's transients live at once and can exhaust GPU +// memory), so serving flips the caps per phase. +// +// This is the only TU that includes device.h with access relaxed. The +// pinned mlx wheel ships the exact header this compiles against, so the +// member offsets match the wheel's libmlx; set_cb_caps still refuses to +// write when the current values read implausible, as a layout tripwire. +#include +#include + +#ifdef _METAL_ +#define private public +#include "mlx/backend/metal/device.h" +#undef private +#endif + +#include "mlx/device.h" + +namespace mlx_kquant { + +#ifdef _METAL_ + +namespace { + +bool plausible(int ops, int mb) { + return ops > 0 && ops <= (1 << 30) && mb > 0 && mb <= (1 << 30); +} + +} // namespace + +std::pair get_cb_caps() { + auto& d = mlx::core::metal::device(mlx::core::Device::gpu); + auto [ops, mb] = d.get_max_ops_mb_per_buffer(); + return {ops, mb}; +} + +std::pair set_cb_caps(int max_ops, int max_mb) { + if (!plausible(max_ops, max_mb)) { + throw std::invalid_argument( + "[mlx_kquant.set_cb_caps] caps must be in [1, 2^30]."); + } + auto& d = mlx::core::metal::device(mlx::core::Device::gpu); + auto [prev_ops, prev_mb] = d.get_max_ops_mb_per_buffer(); + if (!plausible(prev_ops, prev_mb)) { + throw std::runtime_error( + "[mlx_kquant.set_cb_caps] current caps read implausible; mlx " + "device layout drift, refusing to write."); + } + d.max_ops_per_buffer_ = max_ops; + d.max_mb_per_buffer_ = max_mb; + return {prev_ops, prev_mb}; +} + +#else // !_METAL_ + +std::pair get_cb_caps() { + throw std::runtime_error("[mlx_kquant.get_cb_caps] requires Metal."); +} + +std::pair set_cb_caps(int, int) { + throw std::runtime_error("[mlx_kquant.set_cb_caps] requires Metal."); +} + +#endif + +} // namespace mlx_kquant diff --git a/src/kquant_hc_glue.cpp b/src/kquant_hc_glue.cpp new file mode 100644 index 0000000..eaec8f4 --- /dev/null +++ b/src/kquant_hc_glue.cpp @@ -0,0 +1,414 @@ +// Fused deepseek4 hyper-connection glue ops for the single-token decode +// route (see kq_hc_glue.h for the kernel shapes). Four streams (hc_mult 4) +// baked; the gmlx caller gates on that. GPU only, like the dsa ops. +#include +#include + +#include "kquant.h" +#include "kquant_internal.h" // kq_type_string + +#include "mlx/ops.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" // kq_get_kernel +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +constexpr int HC = 4; +constexpr int MIX = (2 + HC) * HC; + +mx::array prep_hc_act( + const mx::array& x, + const char* op, + const char* what, + mx::Stream s) { + auto dt = x.dtype(); + if (dt != mx::float16 && dt != mx::bfloat16) { + throw std::invalid_argument( + std::string(op) + " " + what + " must be float16 or bfloat16."); + } + return mx::contiguous(x, false, s); +} + +mx::array prep_hc_f32( + const mx::array& a, + const char* op, + const char* what, + mx::Stream s) { + if (a.dtype() != mx::float32) { + throw std::invalid_argument( + std::string(op) + " " + what + " must be float32."); + } + return mx::contiguous(a, false, s); +} + +// [..., 4, D] stream tensor: validates the stream axis and the D limits the +// kernels assume, returns D. +int check_streams(const mx::array& x, const char* op, const char* what) { + if (x.ndim() < 2 || x.shape(-2) != HC) { + throw std::invalid_argument( + std::string(op) + " " + what + " must be [..., 4, D]."); + } + int D = x.shape(-1); + if (D % 8 != 0 || D > 8 * 1024) { + throw std::invalid_argument( + std::string(op) + " D must be a multiple of 8 and at most 8192."); + } + return D; +} + +void check_fn(const mx::array& fn, int D, const char* op) { + if (fn.ndim() != 2 || fn.shape(0) != MIX || fn.shape(1) != HC * D) { + throw std::invalid_argument( + std::string(op) + " fn must be [24, 4 * D] float32."); + } +} + +} // namespace + +#ifdef _METAL_ + +void KQuantHcFrontReduce::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + for (auto& out : outputs) { + out.set_data(mx::allocator::malloc(out.nbytes())); + } + const auto& x = inputs[0]; + int D = x.shape(-1); + int rows = int(x.size() / (HC * D)); + + std::string kname = "kq_hc_front_reduce_" + kq_type_string(x.dtype()); + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(x, 0); + ce.set_input_array(inputs[1], 1); + ce.set_output_array(outputs[0], 2); + ce.set_output_array(outputs[1], 3); + ce.set_bytes(D, 4); + ce.dispatch_threadgroups( + MTL::Size(rows * (MIX + 1), 1, 1), MTL::Size(256, 1, 1)); +} + +void KQuantHcFrontExpandReduce::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + for (auto& out : outputs) { + out.set_data(mx::allocator::malloc(out.nbytes())); + } + const auto& resid = inputs[1]; + int D = resid.shape(-1); + int rows = int(resid.size() / (HC * D)); + + std::string kname = + "kq_hc_front_expand_reduce_" + kq_type_string(resid.dtype()); + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(inputs[0], 0); + ce.set_input_array(resid, 1); + ce.set_input_array(inputs[2], 2); + ce.set_input_array(inputs[3], 3); + ce.set_input_array(inputs[4], 4); + ce.set_output_array(outputs[0], 5); + ce.set_output_array(outputs[1], 6); + ce.set_output_array(outputs[2], 7); + ce.set_bytes(D, 8); + ce.dispatch_threadgroups( + MTL::Size(rows * (MIX + 1), 1, 1), MTL::Size(256, 1, 1)); +} + +void KQuantHcSinkhornCollapse::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + for (auto& out : outputs) { + out.set_data(mx::allocator::malloc(out.nbytes())); + } + const auto& x = inputs[0]; + int D = x.shape(-1); + int rows = int(x.size() / (HC * D)); + + std::string kname = "kq_hc_sinkhorn_collapse_" + kq_type_string(x.dtype()); + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(x, 0); + ce.set_input_array(inputs[1], 1); + ce.set_input_array(inputs[2], 2); + ce.set_input_array(inputs[3], 3); + ce.set_input_array(inputs[4], 4); + ce.set_input_array(inputs[5], 5); + ce.set_output_array(outputs[0], 6); + ce.set_output_array(outputs[1], 7); + ce.set_output_array(outputs[2], 8); + ce.set_bytes(D, 9); + ce.set_bytes(iters_, 10); + ce.set_bytes(hc_eps_, 11); + ce.set_bytes(norm_eps_, 12); + ce.dispatch_threadgroups(MTL::Size(rows, 1, 1), MTL::Size(256, 1, 1)); +} + +void KQuantHcExpand::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + auto& out = outputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + const auto& resid = inputs[1]; + int D = resid.shape(-1); + int rows = int(resid.size() / (HC * D)); + + std::string kname = "kq_hc_expand_" + kq_type_string(resid.dtype()); + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(inputs[0], 0); + ce.set_input_array(resid, 1); + ce.set_input_array(inputs[2], 2); + ce.set_input_array(inputs[3], 3); + ce.set_output_array(out, 4); + ce.set_bytes(D, 5); + ce.dispatch_threadgroups(MTL::Size(rows * 2, 1, 1), MTL::Size(256, 1, 1)); +} + +#else // !_METAL_ + +void KQuantHcFrontReduce::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant.hc_front_reduce] requires Metal."); +} + +void KQuantHcFrontExpandReduce::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.hc_front_expand_reduce] requires Metal."); +} + +void KQuantHcSinkhornCollapse::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant.hc_sinkhorn_collapse] requires Metal."); +} + +void KQuantHcExpand::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant.hc_expand] requires Metal."); +} + +#endif + +void KQuantHcFrontReduce::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.hc_front_reduce] has no CPU implementation."); +} + +void KQuantHcFrontExpandReduce::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.hc_front_expand_reduce] has no CPU implementation."); +} + +void KQuantHcSinkhornCollapse::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.hc_sinkhorn_collapse] has no CPU implementation."); +} + +void KQuantHcExpand::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant.hc_expand] has no CPU implementation."); +} + +bool KQuantHcSinkhornCollapse::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return iters_ == o.iters_ && hc_eps_ == o.hc_eps_ && norm_eps_ == o.norm_eps_; +} + +std::vector +hc_front_reduce(mx::array x, mx::array fn, mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.hc_front_reduce]"; + int D = check_streams(x, op, "x"); + check_fn(fn, D, op); + auto x_c = prep_hc_act(x, op, "x", s); + auto fn_c = prep_hc_f32(fn, op, "fn", s); + + auto lead = x.shape(); + lead.pop_back(); + lead.pop_back(); + auto mix_shape = lead; + mix_shape.push_back(MIX); + auto ssq_shape = lead; + ssq_shape.push_back(1); + return mx::array::make_arrays( + {std::move(mix_shape), std::move(ssq_shape)}, + {mx::float32, mx::float32}, + std::make_shared(s), + {std::move(x_c), std::move(fn_c)}); +} + +std::vector hc_front_expand_reduce( + mx::array x_sub, + mx::array resid, + mx::array post, + mx::array comb, + mx::array fn, + mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.hc_front_expand_reduce]"; + int D = check_streams(resid, op, "resid"); + check_fn(fn, D, op); + if (x_sub.shape(-1) != D || x_sub.size() != resid.size() / HC) { + throw std::invalid_argument( + std::string(op) + " x_sub must be [..., D] matching resid."); + } + if (x_sub.dtype() != resid.dtype()) { + throw std::invalid_argument( + std::string(op) + " x_sub and resid dtypes must match."); + } + int64_t rows = resid.size() / (int64_t(HC) * D); + if (int64_t(post.size()) != rows * HC || + int64_t(comb.size()) != rows * HC * HC) { + throw std::invalid_argument( + std::string(op) + " post/comb must be [..., 4] / [..., 4, 4]."); + } + auto xs_c = prep_hc_act(x_sub, op, "x_sub", s); + auto r_c = prep_hc_act(resid, op, "resid", s); + auto p_c = prep_hc_f32(post, op, "post", s); + auto c_c = prep_hc_f32(comb, op, "comb", s); + auto fn_c = prep_hc_f32(fn, op, "fn", s); + + auto lead = resid.shape(); + lead.pop_back(); + lead.pop_back(); + auto mix_shape = lead; + mix_shape.push_back(MIX); + auto ssq_shape = lead; + ssq_shape.push_back(1); + return mx::array::make_arrays( + {resid.shape(), std::move(mix_shape), std::move(ssq_shape)}, + {resid.dtype(), mx::float32, mx::float32}, + std::make_shared(s), + {std::move(xs_c), + std::move(r_c), + std::move(p_c), + std::move(c_c), + std::move(fn_c)}); +} + +std::vector hc_sinkhorn_collapse( + mx::array x, + mx::array mixes_raw, + mx::array sumsq, + mx::array scale, + mx::array base, + mx::array w, + int iters, + float hc_eps, + float norm_eps, + mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.hc_sinkhorn_collapse]"; + int D = check_streams(x, op, "x"); + int64_t rows = x.size() / (int64_t(HC) * D); + if (int64_t(mixes_raw.size()) != rows * MIX || + int64_t(sumsq.size()) != rows) { + throw std::invalid_argument( + std::string(op) + " mixes_raw/sumsq must be [..., 24] / [..., 1]."); + } + if (scale.size() != 3 || int(base.size()) != MIX) { + throw std::invalid_argument( + std::string(op) + " scale must be [3] and base [24]."); + } + if (w.ndim() != 1 || w.shape(0) != D || w.dtype() != x.dtype()) { + throw std::invalid_argument( + std::string(op) + " w must be [D] in the activation dtype."); + } + if (iters < 1) { + throw std::invalid_argument(std::string(op) + " iters must be >= 1."); + } + auto x_c = prep_hc_act(x, op, "x", s); + auto m_c = prep_hc_f32(mixes_raw, op, "mixes_raw", s); + auto q_c = prep_hc_f32(sumsq, op, "sumsq", s); + auto sc_c = prep_hc_f32(scale, op, "scale", s); + auto b_c = prep_hc_f32(base, op, "base", s); + auto w_c = mx::contiguous(w, false, s); + + auto lead = x.shape(); + lead.pop_back(); + lead.pop_back(); + auto col_shape = lead; + col_shape.push_back(D); + auto post_shape = lead; + post_shape.push_back(HC); + auto comb_shape = lead; + comb_shape.push_back(HC); + comb_shape.push_back(HC); + return mx::array::make_arrays( + {std::move(col_shape), std::move(post_shape), std::move(comb_shape)}, + {x.dtype(), mx::float32, mx::float32}, + std::make_shared(s, iters, hc_eps, norm_eps), + {std::move(x_c), + std::move(m_c), + std::move(q_c), + std::move(sc_c), + std::move(b_c), + std::move(w_c)}); +} + +mx::array hc_expand( + mx::array x, + mx::array resid, + mx::array post, + mx::array comb, + mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.hc_expand]"; + int D = check_streams(resid, op, "resid"); + if (x.shape(-1) != D || x.size() != resid.size() / HC || + x.dtype() != resid.dtype()) { + throw std::invalid_argument( + std::string(op) + " x must be [..., D] matching resid."); + } + int64_t rows = resid.size() / (int64_t(HC) * D); + if (int64_t(post.size()) != rows * HC || + int64_t(comb.size()) != rows * HC * HC) { + throw std::invalid_argument( + std::string(op) + " post/comb must be [..., 4] / [..., 4, 4]."); + } + auto x_c = prep_hc_act(x, op, "x", s); + auto r_c = prep_hc_act(resid, op, "resid", s); + auto p_c = prep_hc_f32(post, op, "post", s); + auto c_c = prep_hc_f32(comb, op, "comb", s); + + return mx::array( + resid.shape(), + resid.dtype(), + std::make_shared(s), + {std::move(x_c), std::move(r_c), std::move(p_c), std::move(c_c)}); +} + +} // namespace mlx_kquant diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index 7e2ac82..e48d9c6 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -71,6 +71,28 @@ inline const char* kq_nx_suffix(int nx) { return nx == 32 ? "_nx32" : (nx == 16 ? "_nx16" : ""); } +// Slot-parallel mix_ns (_sp): the S slot dots spread across S simdgroup +// pairs, multiplying resident threads by S without shortening per-thread +// K-chains (the widening lever that measured flat-to-negative on the +// single-stream gathers). Outputs are bit-identical to the loop kernel. +// Solo dispatch -5% at the V4-Flash down shape (the loop kernel's launch +// ramp is occupancy-shy); E2E -0.1% pipelined / -0.5% naive, direction +// consistent across 8 runs x 2 loop regimes x both ABA orders. Default: on +// when the coarse grid underfills the device (decode-scale launches); +// prefill-scale grids keep the loop kernel. KQ_MOE_SP=1/0 forces on/off; +// read live once set so in-process A/Bs can flip arms. +inline bool kq_moe_sp(int64_t coarse_tgs, int S) { + if (S < 2 || S > 16) { + return false; // threadgroup is 64 * S threads; 1024 cap => S <= 16 + } + static const bool has_env = std::getenv("KQ_MOE_SP") != nullptr; + if (has_env) { + const char* e = std::getenv("KQ_MOE_SP"); + return e != nullptr && std::atoi(e) != 0; + } + return coarse_tgs < 2048; +} + // KQ_MOE_NX_LOG=1: print each fused-MoE kernel name once (dispatch audit). inline void kq_moe_log_kname(const std::string& kname) { static const bool log = std::getenv("KQ_MOE_NX_LOG") != nullptr; @@ -273,9 +295,10 @@ void KQuantMoEGLUKQ::eval_gpu( int K = x.shape(-1); const int nx = kq_moe_pick_nx((int64_t)N * R * T, K, true); - std::string kname = "kq_" + kq_gather_stem_nx(kquant_type_, K, nx) + - "_moe_glu_gather_" + (biased ? "bias_" : "") + act_ + kq_nx_suffix(nx) + - "_" + kq_type_string(x.dtype()); + const std::string stem = kq_gather_stem_nx(kquant_type_, K, nx); + std::string kname = "kq_" + stem + "_moe_glu_gather_" + + (biased ? "bias_" : "") + act_ + kq_nx_suffix(nx) + "_" + + kq_type_string(x.dtype()); kq_moe_log_kname(kname); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -474,10 +497,11 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( // mix_ns is generic for every codec (no tuned variants) -- plain names. // No fine tier: the Ext fine variants measured E2E-neutral and were - // dropped. + // dropped. Decode-scale launches route to the slot-parallel variant. const int nx = kq_moe_pick_nx((int64_t)N * T, K, false); + const bool sp = nx == 8 && kq_moe_sp((int64_t)T * (N / 8), S); std::string kname = "kq_" + kquant_type_ + "_gather_qmv_mix_ns" + - kq_nx_suffix(nx) + "_" + kq_type_string(x.dtype()); + (sp ? "_sp" : kq_nx_suffix(nx)) + "_" + kq_type_string(x.dtype()); kq_moe_log_kname(kname); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -490,8 +514,8 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( ce.set_bytes(K, 5); ce.set_bytes(N, 6); ce.set_bytes(S, 7); - MTL::Size group_dims(32, 2, 1); - MTL::Size grid_dims(N / (64 / nx), 1, T); + MTL::Size group_dims(32, sp ? 2 * S : 2, 1); + MTL::Size grid_dims(N / (sp ? 8 : (64 / nx)), 1, T); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/tests/test_cb_caps.py b/tests/test_cb_caps.py new file mode 100644 index 0000000..fc065d6 --- /dev/null +++ b/tests/test_cb_caps.py @@ -0,0 +1,36 @@ +"""get_cb_caps/set_cb_caps roundtrip against the live Metal device. + +Metal-only (the device singleton): skipped under KQUANT_FORCE_CPU. + +Usage: test_cb_caps.py +""" + +from __future__ import annotations + +import os + +import pytest + +import mlx_kquant as kq + +pytestmark = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="cb caps live on the Metal device; no CPU path.", +) + + +def test_roundtrip(): + ops, mb = kq.get_cb_caps() + assert ops > 0 and mb > 0 + prev = kq.set_cb_caps(ops + 7, mb + 13) + assert prev == (ops, mb) + assert kq.get_cb_caps() == (ops + 7, mb + 13) + assert kq.set_cb_caps(ops, mb) == (ops + 7, mb + 13) + assert kq.get_cb_caps() == (ops, mb) + + +def test_rejects_implausible(): + with pytest.raises(ValueError): + kq.set_cb_caps(0, 40) + with pytest.raises(ValueError): + kq.set_cb_caps(400, -1) diff --git a/tests/test_hc_glue.py b/tests/test_hc_glue.py new file mode 100644 index 0000000..b85639c --- /dev/null +++ b/tests/test_hc_glue.py @@ -0,0 +1,157 @@ +"""Fused hyper-connection glue ops vs float64 numpy references. + +The kernels accumulate in f32 and round once at each write, so references +are computed in f64 and compared with an f32-reduction tolerance. The +fused front_expand_reduce is specified bit-identical to hc_expand followed +by hc_front_reduce and is checked exactly against that composition. + +Metal-only kernels (eval_cpu throws): skipped under KQUANT_FORCE_CPU. + +Usage: test_hc_glue.py +""" + +from __future__ import annotations + +import os + +import mlx.core as mx +import numpy as np +import pytest + +import mlx_kquant as kq + +HC = 4 +MIX = 24 +ITERS = 20 +HC_EPS = 1e-6 +NORM_EPS = 1e-6 + +pytestmark = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="hc glue ops are Metal-only kernels; no CPU path.", +) + + +def _mk(seed, D=2048, dtype=mx.bfloat16): + mx.random.seed(seed) + x = (mx.random.normal((1, 1, HC, D)) * 0.05).astype(dtype) + fn = (mx.random.normal((MIX, HC * D)) * 0.02).astype(mx.float32) + scale = mx.array([1.1, 0.9, 1.3], dtype=mx.float32) + base = (mx.random.normal((MIX,)) * 0.1).astype(mx.float32) + w = (mx.random.normal((D,)) * 0.1 + 1.0).astype(dtype) + mx.eval(x, fn, scale, base, w) + return x, fn, scale, base, w + + +def _np64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + +def _sinkhorn_ref(mix, ssq, scale, base, D): + factor = 1.0 / np.sqrt(ssq / (HC * D) + NORM_EPS) + pre = 1.0 / (1.0 + np.exp(-(mix[:HC] * scale[0] * factor + base[:HC]))) + pre = pre + HC_EPS + post = 2.0 / ( + 1.0 + np.exp(-(mix[HC : 2 * HC] * scale[1] * factor + base[HC : 2 * HC])) + ) + v = (mix[2 * HC :] * scale[2] * factor + base[2 * HC :]).reshape(HC, HC) + e = np.exp(v - v.max(axis=1, keepdims=True)) + r = e / (e.sum(axis=1, keepdims=True) + HC_EPS) + HC_EPS + r = r / (r.sum(axis=0, keepdims=True) + HC_EPS) + for _ in range(1, ITERS): + r = r / (r.sum(axis=1, keepdims=True) + HC_EPS) + r = r / (r.sum(axis=0, keepdims=True) + HC_EPS) + return pre, post, r + + +@pytest.mark.parametrize("D", [2048, 4096]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) +def test_front_reduce_and_collapse(D, dtype): + x, fn, scale, base, w = _mk(7, D, dtype) + mr, ssq = kq.hc_front_reduce(x, fn) + col, post, comb = kq.hc_sinkhorn_collapse( + x, mr, ssq, scale, base, w, iters=ITERS, hc_eps=HC_EPS, norm_eps=NORM_EPS + ) + mx.eval(mr, ssq, col, post, comb) + + xf = _np64(x).reshape(HC * D) + fnf = _np64(fn) + mr_ref = fnf @ xf + ssq_ref = float(xf @ xf) + assert np.abs(np.array(mr).ravel() - mr_ref).max() < 1e-3 + assert abs(float(ssq.item()) - ssq_ref) / ssq_ref < 1e-5 + + pre_ref, post_ref, comb_ref = _sinkhorn_ref( + mr_ref, ssq_ref, _np64(scale), _np64(base), D + ) + assert np.abs(np.array(post).ravel() - post_ref).max() < 1e-4 + assert np.abs(np.array(comb).reshape(HC, HC) - comb_ref).max() < 5e-4 + + xs = _np64(x).reshape(HC, D) + collapsed = pre_ref @ xs + inv = 1.0 / np.sqrt((collapsed * collapsed).mean() + NORM_EPS) + col_ref = collapsed * inv * _np64(w) + tol = 2e-2 if dtype == mx.bfloat16 else 5e-3 + denom = np.abs(col_ref).max() + 1e-6 + assert np.abs(_np64(col).ravel() - col_ref).max() / denom < tol + + +@pytest.mark.parametrize("D", [2048, 4096]) +def test_expand(D): + x, fn, scale, base, w = _mk(11, D) + mr, ssq = kq.hc_front_reduce(x, fn) + col, post, comb = kq.hc_sinkhorn_collapse( + x, mr, ssq, scale, base, w, iters=ITERS, hc_eps=HC_EPS, norm_eps=NORM_EPS + ) + out = kq.hc_expand(col, x, post, comb) + mx.eval(out) + + xf = _np64(col).reshape(D) + rf = _np64(x).reshape(HC, D) + pf = np.array(post).ravel().astype(np.float64) + cf = np.array(comb).reshape(HC, HC).astype(np.float64) + ref = pf[:, None] * xf[None, :] + cf.T @ rf + denom = np.abs(ref).max() + 1e-6 + assert np.abs(_np64(out).reshape(HC, D) - ref).max() / denom < 2e-2 + + +@pytest.mark.parametrize("D", [2048, 4096]) +def test_front_expand_reduce_matches_composition(D): + # D a multiple of 1024 keeps the fused kernel's reduction order + # identical to hc_front_reduce, which the bit-exact claim needs. + x, fn, scale, base, w = _mk(13, D) + mr, ssq = kq.hc_front_reduce(x, fn) + col, post, comb = kq.hc_sinkhorn_collapse( + x, mr, ssq, scale, base, w, iters=ITERS, hc_eps=HC_EPS, norm_eps=NORM_EPS + ) + + h_ref = kq.hc_expand(col, x, post, comb) + mr_ref, ssq_ref = kq.hc_front_reduce(h_ref, fn) + h, mr2, ssq2 = kq.hc_front_expand_reduce(col, x, post, comb, fn) + mx.eval(h_ref, mr_ref, ssq_ref, h, mr2, ssq2) + + assert np.array_equal( + np.array(h.astype(mx.float32)), np.array(h_ref.astype(mx.float32)) + ) + assert np.array_equal(np.array(mr2), np.array(mr_ref)) + assert np.array_equal(np.array(ssq2), np.array(ssq_ref)) + + +def test_input_validation(): + x, fn, scale, base, w = _mk(17, 2048) + with pytest.raises(ValueError): + kq.hc_front_reduce(x[..., :2, :], fn) + with pytest.raises(ValueError): + kq.hc_front_reduce(x, fn[:, :100]) + with pytest.raises(ValueError): + kq.hc_sinkhorn_collapse( + x, + mx.zeros((1, 1, MIX)), + mx.zeros((1, 1, 1)), + scale, + base, + w.astype(mx.float32), + iters=ITERS, + hc_eps=HC_EPS, + norm_eps=NORM_EPS, + ) diff --git a/tests/test_moe_glu_variants.py b/tests/test_moe_glu_variants.py new file mode 100644 index 0000000..8f5a290 --- /dev/null +++ b/tests/test_moe_glu_variants.py @@ -0,0 +1,70 @@ +"""Slot-parallel mix_ns (_sp) A/B. + +The variant restructures parallelism only (the slot loop spreads onto +simdgroup pairs) and must be bit-identical to the loop kernel. Each arm +runs in a subprocess so the live-read env latch (KQ_MOE_SP) sees the +variable from the first dispatch. +""" + +import os +import subprocess +import sys + +import numpy as np +import pytest + +pytestmark = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="fused MoE gathers are Metal-only kernels; no CPU path.", +) + +_SNIPPET = r""" +import sys +import numpy as np +import mlx.core as mx +import mlx_kquant as kq + +codec, out_path = sys.argv[1], sys.argv[2] +rng = np.random.default_rng(11) +E, N, K, T, S = 32, 64, 512, 2, 6 +bpb = {"q2_k": 84, "q4_k": 144, "q8_0": 34, "iq2_xxs": 66}[codec] +wpb = 32 if codec == "q8_0" else 256 +nb = E * N * (K // wpb) +wire = rng.integers(0, 256, size=(nb, bpb), dtype=np.uint8) +d = rng.uniform(0.004, 0.01, nb).astype(np.float16) +off = 80 if codec == "q2_k" else 0 +wire[:, off:off + 2] = d.view(np.uint8).reshape(nb, 2) +dmin_off = {"q2_k": 82, "q4_k": 2}.get(codec) +if dmin_off is not None: + dm = rng.uniform(0.001, 0.004, nb).astype(np.float16) + wire[:, dmin_off:dmin_off + 2] = dm.view(np.uint8).reshape(nb, 2) +w = mx.array(wire.reshape(E, N, (K // wpb) * bpb)) +h = mx.array((rng.standard_normal((T, S, K)) * 0.05).astype(np.float16)) +x = mx.array((rng.standard_normal((T, K)) * 0.05).astype(np.float16)) +inds = mx.array(rng.integers(0, E, size=(T, S)).astype(np.uint32)) +sc = mx.array(rng.uniform(0.05, 0.9, size=(T, S)).astype(np.float32)) +mix = kq.gather_qmv_mix_ns_kq(h, w, codec, inds, sc) +glu = kq.moe_glu_gather_kq(x, w, w, codec, inds, act="silu") +mx.eval(mix, glu) +np.savez(out_path, mix=np.array(mix.astype(mx.float32)), + glu=np.array(glu.astype(mx.float32))) +""" + + +@pytest.mark.parametrize("codec", ["q2_k", "q4_k", "q8_0", "iq2_xxs"]) +def test_sp_bit_identical(codec, tmp_path): + outs = {} + for arm, env in ( + ("base", {"KQ_MOE_SP": "0"}), + ("variant", {"KQ_MOE_SP": "1"}), + ): + f = tmp_path / f"{arm}.npz" + subprocess.run( + [sys.executable, "-c", _SNIPPET, codec, str(f)], + check=True, + env={**os.environ, **env}, + ) + outs[arm] = np.load(f) + for key in ("mix", "glu"): + a, b = outs["base"][key], outs["variant"][key] + assert np.array_equal(a, b), f"{codec} {key} not bit-identical"