From 84e19bd6910c58e88a37f29799e3062fe9e013ef Mon Sep 17 00:00:00 2001 From: Aaron Z Date: Wed, 21 Oct 2026 16:03:12 -0700 Subject: [PATCH 1/2] [Engine] support Qwen3.8-Flash-Next Coder (256 experts) on Apple Silicon Metal Support Qwen3.8-Flash-Next Coder variants with 256 experts (including IQ1_M) on Apple Silicon Mac (Metal backend). - metal/kernels/router.metal, moe.metal, prompt.metal: add 256-expert routing and MoE dispatch paths while preserving 512-expert paths for standard models. - metal/src/engine/window.cpp, stream.cpp, stream.h, prompt_chain.inc: dynamically read expert count from model metadata instead of assuming static 512. - metal/src/engine/window.cpp, kernels/dense_mm.metal: guard MetalPerformancePrimitives and add supports_metal4() check to fall back cleanly to decode window prefill on macOS 15 / M1-M4 hardware. - setup.py, metal/setup_mac.py, strata_control/mac_fit.py: add support for family 'coder' and model 'IQ1_M' on macOS. --- .gitignore | 1 + metal/kernels/dense_mm.metal | 7 ++ metal/kernels/mm_bench.metal | 6 ++ metal/kernels/moe.metal | 114 ++++++++++++++++++++---------- metal/kernels/prompt.metal | 46 ++++++++++++ metal/kernels/router.metal | 114 ++++++++++++++++++++---------- metal/setup_mac.py | 34 ++++++--- metal/src/engine/prompt_chain.inc | 18 ++--- metal/src/engine/stream.cpp | 75 +++++++++++--------- metal/src/engine/stream.h | 2 + metal/src/engine/window.cpp | 69 +++++++++--------- metal/tests/coop_probe.h | 1 + metal/tools/stream_experts.py | 22 +++--- setup.py | 15 ++-- strata_control/mac_fit.py | 9 ++- 15 files changed, 364 insertions(+), 169 deletions(-) diff --git a/.gitignore b/.gitignore index 65457fdc042..6b1da6b047e 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ td.log /.venv/ /build/ /third_party/llama.cpp +/third_party/metal-cpp/ /models/ /packs/ /mtp/ diff --git a/metal/kernels/dense_mm.metal b/metal/kernels/dense_mm.metal index 45e29acec88..059c8c19904 100644 --- a/metal/kernels/dense_mm.metal +++ b/metal/kernels/dense_mm.metal @@ -3,8 +3,10 @@ // dq_f16 a quantized matrix -> half rows (ref_quant's values rounded to fp16), once per chunk // gemm_f16 Y[t][n] = Xh[t] . Wh[n] (fp32) by MetalPerformancePrimitives matmul2d, TM x TN tiles, K dynamic; // rows of Y past n_out / t past n_tok are not written (the tensors carry the extents) +#if __has_include() #include using namespace mpp::tensor_ops; +#endif kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], uint gid [[thread_position_in_grid]]) { @@ -45,6 +47,7 @@ kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer struct GemmArgs { int M, N, K; int pad; }; // Y is M rows of N floats +#if __has_include() template kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device float* Y [[buffer(2)]], constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { @@ -58,12 +61,16 @@ kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM); op.run(mA, mB, mY); } +#endif #define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; #define DQ_(F) DQ(F) DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) +#if __has_include() template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; +#endif + diff --git a/metal/kernels/mm_bench.metal b/metal/kernels/mm_bench.metal index c04659fc5ee..b26e6a52abf 100644 --- a/metal/kernels/mm_bench.metal +++ b/metal/kernels/mm_bench.metal @@ -2,6 +2,7 @@ // mpp_mm Metal 4 MetalPerformancePrimitives matmul2d (the M5's neural accelerators), a TM x TN // tile of C per threadgroup, K dynamic, A and B straight from device memory // sg_mm simdgroup_matrix 8x8, a 32 x 32 tile per simdgroup (4 x 4 accumulators), 4 simdgroups +#if __has_include() #include using namespace mpp::tensor_ops; @@ -29,6 +30,9 @@ MPP(half, 64, 16, 2) MPP(half, 64, 16, 4) MPP(half, 128, 16, 4) MPP(half, 128, 3 MPP(half, 32, 32, 2) MPP(half, 16, 32, 1) MPP(half, 16, 16, 1) MPP(half, 128, 48, 4) MPP(half, 64, 48, 4) MPP(float, 64, 32, 4) MPP(float, 64, 64, 4) MPP(float, 128, 64, 4) MPP(bfloat, 64, 64, 4) MPP(bfloat, 128, 64, 4) +#else +struct MmArgs { int M, N, K; }; +#endif // simdgroup_matrix: threadgroup tile 64 x 64 (4 simdgroups x 32 x 32), K in steps of 8, operands from device memory template @@ -53,6 +57,7 @@ template [[host_name("sg_mm_half")]] kernel decltype(sg_mm) sg_mm; template [[host_name("sg_mm_float")]] kernel decltype(sg_mm) sg_mm; +#if __has_include() // the element layout of single-simdgroup matmul2d cooperative tensors (right input, left input, destination) for a // TM x TN x TK descriptor: out[((which * 32 + lane) * 128 + i) * 4 + (0, 1)] = the two indices, [2] = capacity template @@ -83,3 +88,4 @@ template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 3 template [[host_name("coop_layout_32_32_16")]] kernel decltype(coop_layout<32, 32, 16>) coop_layout<32, 32, 16>; template [[host_name("coop_layout_16_64_32")]] kernel decltype(coop_layout<16, 64, 32>) coop_layout<16, 64, 32>; template [[host_name("coop_layout_16_32_32")]] kernel decltype(coop_layout<16, 32, 32>) coop_layout<16, 32, 32>; +#endif diff --git a/metal/kernels/moe.metal b/metal/kernels/moe.metal index e3f0b1d68ff..738a99c2810 100644 --- a/metal/kernels/moe.metal +++ b/metal/kernels/moe.metal @@ -119,45 +119,87 @@ kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { threadgroup int s_grp[32], s_ent[32]; if (sg < a.n_tok) { - device const float* l = logits + sg * 512; - float v[16]; - for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; - float mx = -INFINITY; - for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); - for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); - float sum = 0.0f; - for (int i = 0; i < 16; ++i) { - v[i] = exp(v[i] - mx); - sum += v[i]; - } - for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); - const float rcp = 1.0f / sum; - for (int i = 0; i < 16; ++i) { - v[i] *= rcp; - if (isnan(v[i])) v[i] = -FLT_MAX; - } - float selected = 0.0f, ssum = 0.0f; - for (int rank = 0; rank < 10; ++rank) { - float best = v[0]; - int ex = lane; - for (int i = 1; i < 16; ++i) - if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } - for (ushort m = 16; m > 0; m >>= 1) { - const float ob = simd_shuffle_xor(best, m); - const int oe = simd_shuffle_xor(ex, m); - if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + if (a.n_expert == 256) { + device const float* l = logits + sg * 256; + float v[8]; + for (int i = 0; i < 8; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 8; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 8; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 8; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 8; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } else { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; } - if ((ex & 31) == lane) { - v[ex / 32] = -INFINITY; - ids[sg * 10 + rank] = ex; - ssum += best; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; } - if (rank == lane) selected = best; + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; } - for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); - ssum = max(ssum, 6.103515625e-5f); - const float inv = 1.0f / ssum; - if (lane < 10) w[sg * 10 + lane] = selected * inv; } threadgroup_barrier(mem_flags::mem_device); plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); diff --git a/metal/kernels/prompt.metal b/metal/kernels/prompt.metal index 64833b0a711..0feeb50615e 100644 --- a/metal/kernels/prompt.metal +++ b/metal/kernels/prompt.metal @@ -55,6 +55,52 @@ kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); } +kernel void route_topk_256(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const uint t = tg.x * 8 + sg; + if (t >= n_tok) return; + device const float* l = logits + ulong(t) * 256; + float v[8]; + for (int i = 0; i < 8; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 8; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 8; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 8; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 8; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[t * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); +} + struct EpiArgs { uint n_tok; uint n_embd; // 2560 diff --git a/metal/kernels/router.metal b/metal/kernels/router.metal index e5c75274779..3674f54f18b 100644 --- a/metal/kernels/router.metal +++ b/metal/kernels/router.metal @@ -9,45 +9,87 @@ static inline void route_body(device const float* logits, device int* ids, devic constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, threadgroup int* s_ent) { if (sg < a.n_tok) { - device const float* l = logits + sg * 512; - float v[16]; - for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; - float mx = -INFINITY; - for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); - for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); - float sum = 0.0f; - for (int i = 0; i < 16; ++i) { - v[i] = exp(v[i] - mx); - sum += v[i]; - } - for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); - const float rcp = 1.0f / sum; - for (int i = 0; i < 16; ++i) { - v[i] *= rcp; - if (isnan(v[i])) v[i] = -FLT_MAX; - } - float selected = 0.0f, ssum = 0.0f; - for (int rank = 0; rank < 10; ++rank) { - float best = v[0]; - int ex = lane; - for (int i = 1; i < 16; ++i) - if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } - for (ushort m = 16; m > 0; m >>= 1) { - const float ob = simd_shuffle_xor(best, m); - const int oe = simd_shuffle_xor(ex, m); - if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + if (a.n_expert == 256) { + device const float* l = logits + sg * 256; + float v[8]; + for (int i = 0; i < 8; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 8; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 8; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 8; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 8; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } else { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; } - if ((ex & 31) == lane) { - v[ex / 32] = -INFINITY; - ids[sg * 10 + rank] = ex; - ssum += best; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; } - if (rank == lane) selected = best; + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; } - for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); - ssum = max(ssum, 6.103515625e-5f); - const float inv = 1.0f / ssum; - if (lane < 10) w[sg * 10 + lane] = selected * inv; } threadgroup_barrier(mem_flags::mem_device); plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); diff --git a/metal/setup_mac.py b/metal/setup_mac.py index 774ac3b29df..42dd2077562 100644 --- a/metal/setup_mac.py +++ b/metal/setup_mac.py @@ -66,11 +66,13 @@ def model_shard_names(model: str) -> list[str]: """The model's files in order (shard 1 first).""" n = MAC_MODELS[model].get("shards", 2) - if MAC_MODELS[model].get("family") == "gyro": # agentionai's one file (the n-gram table is inside it) + fam_name = MAC_MODELS[model].get("family", "qwen") + if fam_name == "gyro": # agentionai's one file (the n-gram table is inside it) return [S.model_file(S.FAMILIES["gyro"], model, 1)] - if MAC_MODELS[model].get("family") == "unsloth": + if fam_name == "unsloth": return [f"Qwen3.8-Flash-Next-{model}-{i:05d}-of-{n:05d}.gguf" for i in range(1, n + 1)] - return [S.model_file(S.FAMILIES["qwen"], model, i) for i in range(1, n + 1)] + return [S.model_file(S.FAMILIES[fam_name], model, i) for i in range(1, n + 1)] + def shard_sha(name: str): @@ -326,8 +328,8 @@ def start(cfg_path: Path, port: int | None, open_browser: bool = True) -> int: # ------------------------------------------------------------------------------------------------ main def main(argv=None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) - ap.add_argument("--family", choices=["qwen", "unsloth", "gyro"], default=None, - help="qwen = Qwen3.8-Flash-Next (GSQ-RCO), unsloth = Unsloth's UD-*_K_XL, gyro = agentionai's Gyro-S / Gyro-M " + ap.add_argument("--family", choices=["qwen", "unsloth", "gyro", "coder"], default=None, + help="qwen = Qwen3.8-Flash-Next (GSQ-RCO), coder = Qwen3.8-Flash-Next Coder, unsloth = Unsloth's UD-*_K_XL, gyro = agentionai's Gyro-S / Gyro-M " "(follows --model)") ap.add_argument("--model", choices=list(MAC_MODELS)) ap.add_argument("--context", type=int) @@ -392,7 +394,12 @@ def main(argv=None) -> int: if a.model: model = a.model else: - names = [m for m in MAC_MODELS if not MAC_MODELS[m].get("missing")] # --model offers the others + target_family = a.family or "qwen" + names = [m for m in MAC_MODELS if not MAC_MODELS[m].get("missing") and MAC_MODELS[m].get("family", "qwen") == target_family] + if not names: + names = [m for m in MAC_MODELS if not MAC_MODELS[m].get("missing")] + if rec not in names: + rec = names[0] for i, m in enumerate(names, 1): S.say(f" {i}) {m:8} {MAC_ABOUT[m]} (maps {MAC_MODELS[m]['shard1_gb']:.0f} GB" + ("" if MAC_MODELS[m]["tested"] else ", not measured on a Mac") + ")" @@ -453,10 +460,17 @@ def main(argv=None) -> int: models_dir.mkdir(parents=True, exist_ok=True) files = [models_dir / n for n in model_shard_names(model)] for s in files: - if s.exists() and S.done(s): - S.ok(f"{s.name} already downloaded") - continue - S.download(fam["hf"].format(q=model) + s.name, s) + if s.exists() and not S.done(s) and S.whole_shard(s): + try: + S.mark(s, "whole (checked against its own tensor directory)") + except OSError as e: + S.warn(f"{s.name} is whole but its finish mark cannot be written ({e})") + if not a.gguf_dir: + for s in files: + if s.exists() and S.done(s): + S.ok(f"{s.name} already downloaded") + continue + S.download(fam["hf"].format(q=model) + s.name, s) S.check_shards(files) for s in files: # the Unsloth files: pinned sizes and SHA-256 if shard_sha(s.name): diff --git a/metal/src/engine/prompt_chain.inc b/metal/src/engine/prompt_chain.inc index f9b24133c79..db1ee40d97a 100644 --- a/metal/src/engine/prompt_chain.inc +++ b/metal/src/engine/prompt_chain.inc @@ -25,6 +25,7 @@ int pc_unit_of(int t) { bool Window::coop_k4() { if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0; + if (!cx_->supports_metal4()) return false; const int lib = cx_->compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n"); mtl::Pipeline p = cx_->pipeline(lib, "coop_layout_16_64_32"); mtl::Buffer out = cx_->alloc(4 * 3 * 32 * 128 * 4); @@ -75,16 +76,16 @@ bool Window::pc_setup(int C) { b.qcur = F32(MT * 24 * 256); b.qcurh = F16(MT * 24 * 256); b.qidx = F32(MT * 4 * 128); b.scores = F32(MT * MB); b.qids = cx_->alloc(4 * MT * PC_QB); b.part = F32(MT * PC_QCHUNKS * 24 * 258); b.XQ = F32(4 * MT * 6144 + MT * 384); b.xqh = F16(MT * 6144); - b.logits = F32(MT * NEXP); b.ids = cx_->alloc(4 * MT * TOPK); b.wts = F32(MT * TOPK); - b.bcnt = cx_->alloc(4 * (size_t) nblk_max * NEXP); b.boff = cx_->alloc(4 * (size_t) nblk_max * NEXP); - b.eoff = cx_->alloc(4 * NEXP); b.tiles = cx_->alloc(8 * (size_t) (b.maxrows / PC_TM + 1)); b.meta = cx_->alloc(64); + b.logits = F32(MT * n_expert_); b.ids = cx_->alloc(4 * MT * TOPK); b.wts = F32(MT * TOPK); + b.bcnt = cx_->alloc(4 * (size_t) nblk_max * n_expert_); b.boff = cx_->alloc(4 * (size_t) nblk_max * n_expert_); + b.eoff = cx_->alloc(4 * n_expert_); b.tiles = cx_->alloc(8 * (size_t) (b.maxrows / PC_TM + 1)); b.meta = cx_->alloc(64); b.stok = cx_->alloc(4 * (size_t) b.maxrows); b.sw = F32(b.maxrows); b.pos = cx_->alloc(4 * MT * TOPK); b.XS = F16((size_t) b.maxrows * N); b.HS = F16((size_t) b.maxrows * NFF); b.YE = F32((size_t) b.maxrows * N); b.yr = F32(MT * N); b.gsh = F32(MT * NFF); b.ush = F32(MT * NFF); b.hh = F16(MT * NFF); b.ys = F32(MT * N); b.sgv = F32(MT); b.WD = F16((size_t) (LR + HC) * HC * N); b.WU = F16((size_t) HC * N * LR); b.Wqkv = F16((size_t) 10240 * N); b.Wz = F16((size_t) 6144 * N); b.Wal = F16((size_t) 48 * N); b.Wbe = F16((size_t) 48 * N); b.Wout = F16((size_t) N * 6144); b.Wq = F16((size_t) 12288 * N); b.Wk = F16((size_t) 512 * N); b.Wv = F16((size_t) 512 * N); b.Wiq = F16((size_t) 512 * N); - b.Wik = F16((size_t) 128 * N); b.Wr = F16((size_t) NEXP * N); b.Wsg = F16((size_t) N); b.Wshg = F16((size_t) NFF * N); + b.Wik = F16((size_t) 128 * N); b.Wr = F16((size_t) n_expert_ * N); b.Wsg = F16((size_t) N); b.Wshg = F16((size_t) NFF * N); b.Wshu = F16((size_t) NFF * N); b.Wshd = F16((size_t) N * NFF); // the PLE block's rows and projections b.pemb = F32(MT * N); b.pembh = F16(MT * N); b.pkey = F32(MT * 10240); b.pval = F32(MT * N); b.keyn = F32(MT * 10240); @@ -274,16 +275,16 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& dq(T(bl + "ffn_gate_shexp.weight"), b.Wshg); dq(T(bl + "ffn_up_shexp.weight"), b.Wshu); dq(T(bl + "ffn_down_shexp.weight"), b.Wshd); - gemm(b.mixh, b.Wr, b.logits, nT, NEXP, N); + gemm(b.mixh, b.Wr, b.logits, nT, n_expert_, N); const uint32_t nt = (uint32_t) nT; - D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256); + D(P_(b.lib, n_expert_ == 256 ? "route_topk_256" : "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256); const int nent = nT * TOPK, nblk = (nent + 1023) / 1024; - PcMmPlanArgs pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) PC_TM, (uint32_t) nblk}; + PcMmPlanArgs pa{(uint32_t) nent, TOPK, (uint32_t) n_expert_, (uint32_t) PC_TM, (uint32_t) nblk}; const uint32_t nclr = (uint32_t) b.maxrows; D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256); D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta), - mtl::Arg::V(pa)}, 1, 1, NEXP); + mtl::Arg::V(pa)}, 1, 1, n_expert_); D(P_(b.lib, "mm_place"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.stok), mtl::Arg::B(b.sw), mtl::Arg::B(b.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); PcMmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), (uint32_t) ref::row_bytes(ed.type, NFF), @@ -403,6 +404,7 @@ bool Window::prompt(const int64_t* tokens, int64_t n, int64_t pos0, const char* e = std::getenv("STRATA_METAL_PROMPT_CHUNK"); return std::max(64, std::min(8192, e ? std::atoi(e) : 2048)); }(); + if (!cx_ || !cx_->supports_metal4()) { err = "metal prompt: prompt kernels require Metal 4"; return false; } if (!pc_setup((int) std::min(C, std::max(n, 64)))) { err = "metal prompt: setup failed"; return false; } const double t0 = now_ms(); pb_.gpu_ms = pb_.host_ms = pb_.mtp_ms = 0; diff --git a/metal/src/engine/stream.cpp b/metal/src/engine/stream.cpp index 4d90b985334..7525ec56a73 100644 --- a/metal/src/engine/stream.cpp +++ b/metal/src/engine/stream.cpp @@ -21,7 +21,7 @@ namespace strata::metal_engine { namespace { -constexpr int NL = 48, NE = 512; +constexpr int NL = 48; constexpr uint64_t PAGE = 16384; double now_ms() { return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); } int env_int(const char* n, int d) { @@ -62,6 +62,17 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { if (fd_ < 0) { err = "metal stream: cannot open " + f.path; return false; } fcntl(fd_, F_NOCACHE, 1); fcntl(fd_, F_RDAHEAD, 0); + const gguf::Tensor* r0 = f.find("blk.0.ffn_gate_inp.weight"); + if (r0 && r0->shape.size() >= 2 && (r0->shape[1] == 256 || r0->shape[1] == 512)) { + ne_ = (int) r0->shape[1]; + } else { + auto it = f.meta_num.find("qwen4exp.expert_count"); + if (it != f.meta_num.end() && ((int) it->second == 256 || (int) it->second == 512)) { + ne_ = (int) it->second; + } else { + ne_ = 512; + } + } lay_.resize(NL); for (int l = 0; l < NL; ++l) { const std::string b = "blk." + std::to_string(l) + "."; @@ -69,7 +80,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight"); const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight"); if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; } - lay_[(size_t) l] = {g->offset, g->bytes / NE, u->offset, u->bytes / NE, d->offset, d->bytes / NE, -1}; + lay_[(size_t) l] = {g->offset, g->bytes / (uint64_t) ne_, u->offset, u->bytes / (uint64_t) ne_, d->offset, d->bytes / (uint64_t) ne_, -1}; } // the expert-contiguous copy beside the GGUF, when it is this GGUF's (sizes in its header) and its first and // last experts match the GGUF's bytes @@ -88,7 +99,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { uint32_t ver = 0, nl = 0, ne = 0; uint64_t src = 0; std::memcpy(&ver, h + 4, 4); std::memcpy(&nl, h + 8, 4); std::memcpy(&ne, h + 12, 4); std::memcpy(&src, h + 16, 8); - ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) NE && src == (uint64_t) st.st_size; + ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) ne_ && src == (uint64_t) st.st_size; for (int l = 0; ok && l < NL; ++l) { const Lay& x = lay_[(size_t) l]; ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds; @@ -97,7 +108,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { if (ok) { // spot check: layer 0's first and layer 47's last expert, every slice for (const int l : {0, NL - 1}) { const Lay& x = lay_[(size_t) l]; - const int e = l == 0 ? 0 : NE - 1; + const int e = l == 0 ? 0 : ne_ - 1; std::vector a(x.gs + x.us + x.ds), b(a.size()); ok = ok && ::pread(cfd, a.data(), a.size(), (off_t) (off[(size_t) l] + (uint64_t) e * a.size())) == (ssize_t) a.size() && ::pread(fd_, b.data(), x.gs, (off_t) (x.g + (uint64_t) e * x.gs)) == (ssize_t) x.gs && @@ -259,7 +270,7 @@ int ExpertStream::victim(int ci) { if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_) { Slot& x = slots_[(size_t) s]; if (x.layer >= 0) { - const int key = x.layer * NE + x.e; + const int key = x.layer * ne_ + x.e; where_[(size_t) key] = -1; smap_.as()[key] = 0xFFFFFFFFu; } @@ -321,7 +332,7 @@ void ExpertStream::io_main() { std::lock_guard g(mu_); Slot& x = slots_[(size_t) j.slot]; if (--x.left == 0) { - smap_.as()[j.layer * NE + j.e] = code_[(size_t) j.slot]; + smap_.as()[j.layer * ne_ + j.e] = code_[(size_t) j.slot]; x.state = 0; --outstanding_; ++reads_; @@ -341,7 +352,7 @@ void ExpertStream::submit(int s, int l, int e, bool urgent) { const int parts = cfd_ >= 0 ? split_ : 3 * split_; x.left = parts; ++outstanding_; - where_[(size_t) (l * NE + e)] = s; + where_[(size_t) (l * ne_ + e)] = s; for (int p = 0; p < parts; ++p) (urgent ? urgent_q_ : pre_q_).push_back({s, l, e, p}); } @@ -349,7 +360,7 @@ void ExpertStream::submit(int s, int l, int e, bool urgent) { void ExpertStream::cancel_slot(int s) { Slot& x = slots_[(size_t) s]; for (auto it = pre_q_.begin(); it != pre_q_.end();) it = it->slot == s ? pre_q_.erase(it) : it + 1; - where_[(size_t) (x.layer * NE + x.e)] = -1; + where_[(size_t) (x.layer * ne_ + x.e)] = -1; x.layer = x.e = -1; x.state = 0; x.pre = false; @@ -375,17 +386,17 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri for (auto& k : cls_) k.n = 0; // a second try (after a shrink) lays the arena out again // the ranking: the profile's pairs, then every pair it misses, layer-interleaved std::vector> rank; - std::vector seen((size_t) NL * NE, 0); + std::vector seen((size_t) NL * ne_, 0); if (FILE* f = std::fopen(profile.c_str(), "rb")) { char magic[4]; uint32_t h[5]; if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 && - h[1] == (uint32_t) NL && h[2] == (uint32_t) NE) { + h[1] == (uint32_t) NL && h[2] == (uint32_t) ne_) { for (uint32_t i = 0; i < h[4]; ++i) { uint16_t p[2]; if (std::fread(p, 2, 2, f) != 2) break; - if (p[0] < NL && p[1] < NE && !seen[(size_t) p[0] * NE + p[1]]) { - seen[(size_t) p[0] * NE + p[1]] = 1; + if (p[0] < NL && p[1] < ne_ && !seen[(size_t) p[0] * ne_ + p[1]]) { + seen[(size_t) p[0] * ne_ + p[1]] = 1; rank.push_back({p[0], p[1]}); } } @@ -395,9 +406,9 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\n", profile.c_str()); } const size_t from_profile = rank.size(); - for (int e = 0; e < NE; ++e) + for (int e = 0; e < ne_; ++e) for (int l = 0; l < NL; ++l) - if (!seen[(size_t) l * NE + e]) rank.push_back({l, e}); + if (!seen[(size_t) l * ne_ + e]) rank.push_back({l, e}); // the pool of each size class: at least min_slots (a prompt chunk's layer needs up to 512 at once), then the // top-ranked pairs that fit in the budget uint64_t used = 0; @@ -465,11 +476,11 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri code_[(size_t) s] = ((uint32_t) at[(size_t) s].first << 28) | (uint32_t) at[(size_t) s].second; host_[(size_t) s] = chunks_[(size_t) at[(size_t) s].first].as() + at[(size_t) s].second * PAGE; } - smap_ = cx.alloc((size_t) NL * NE * 4); - std::memset(smap_.ptr, 0xFF, (size_t) NL * NE * 4); + smap_ = cx.alloc((size_t) NL * ne_ * 4); + std::memset(smap_.ptr, 0xFF, (size_t) NL * ne_ * 4); slots_ = std::vector((size_t) base); retired_.assign((size_t) base, 0); - where_.assign((size_t) NL * NE, -1); + where_.assign((size_t) NL * ne_, -1); for (auto& k : cls_) { for (int s = k.base; s < k.base + k.n; ++s) { slots_[(size_t) s].prev = s == k.base ? -1 : s - 1; @@ -528,7 +539,7 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri } } // the LRU order = the rank order (best at the front) - for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]); + for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * ne_ + fill[i].second)]); } if (stopped) std::fprintf(stderr, "strata metal stream: the arena's fill stopped after %.2f GB: macOS's free memory reached what it " @@ -577,10 +588,10 @@ void ExpertStream::ahead(int ln, int T, const float* pred, int K) { std::pair pw[8 * 32]; int np = 0; for (int t = 0; t < T; ++t) { - const float* lg = pred + (size_t) t * NE; + const float* lg = pred + (size_t) t * ne_; int top[32]; int nt = 0; - for (int e = 0; e < NE; ++e) { + for (int e = 0; e < ne_; ++e) { if (nt < K) { top[nt++] = e; continue; } int m = 0; for (int j = 1; j < K; ++j) if (lg[top[j]] < lg[top[m]]) m = j; @@ -598,7 +609,7 @@ void ExpertStream::ahead(int ln, int T, const float* pred, int K) { { std::lock_guard g(mu_); for (int j = 0; j < np; ++j) { - const int key = ln * NE + pw[j].second; + const int key = ln * ne_ + pw[j].second; const int s0 = where_[(size_t) key]; if (s0 >= 0) { lru_front(s0); continue; } const int s = victim(lay_[(size_t) ln].cls); @@ -619,18 +630,18 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl const bool prompt = T > 8; // a prompt chunk (the GEMM path): nearly every expert of the layer is used if (l == 0) { if (prompt) ++p_chunks_; else ++windows_; } ++stamp_; // the slots this layer uses are pinned until the next handshake - bool used[NE] = {}; - int want[NE]; + bool used[512] = {}; + int want[512]; int nw = 0; for (int i = 0; i < T * topk; ++i) { const int e = ids[i]; - if (e < 0 || e >= NE || used[e]) continue; + if (e < 0 || e >= ne_ || used[e]) continue; used[e] = true; want[nw++] = e; } if (fake_) { for (int j = 0; j < nw; ++j) - if (where_[(size_t) (l * NE + want[j])] < 0) smap_.as()[l * NE + want[j]] = code_[0]; + if (where_[(size_t) (l * ne_ + want[j])] < 0) smap_.as()[l * ne_ + want[j]] = code_[0]; host_ms_ += now_ms() - t0; return; } @@ -647,9 +658,9 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl for (int s : drop) cancel_slot(s); } for (int j = 0; j < nw; ++j) // pass 1: pin what is resident (or on its way), so no miss below evicts it - if (where_[(size_t) (l * NE + want[j])] >= 0) slots_[(size_t) where_[(size_t) (l * NE + want[j])]].pin = stamp_; + if (where_[(size_t) (l * ne_ + want[j])] >= 0) slots_[(size_t) where_[(size_t) (l * ne_ + want[j])]].pin = stamp_; for (int j = 0; j < nw; ++j) { - const int e = want[j], key = l * NE + e; + const int e = want[j], key = l * ne_ + e; uint64_t& lk = prompt ? p_lookups_ : lookups_; ++lk; int s = where_[(size_t) key]; @@ -694,8 +705,8 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl bool any = false; std::lock_guard g(mu_); const int ln = l + 1; - for (int e = 0; e < NE; ++e) { - if (where_[(size_t) (ln * NE + e)] >= 0) continue; + for (int e = 0; e < ne_; ++e) { + if (where_[(size_t) (ln * ne_ + e)] >= 0) continue; const int s = victim(lay_[(size_t) ln].cls); if (s < 0) break; submit(s, ln, e, false); @@ -728,12 +739,12 @@ uint8_t* ExpertStream::region(uint64_t bytes) { const uint64_t sz = cls_[(size_t) cls_of(s)].pages * PAGE; if (h + sz > lo && h < end) { region_slots_.push_back(s); first = std::min(first, (uint8_t*) h); } } - { // a class left with fewer slots than a prompt chunk's layer can use (NE) would stop with "no free slot": a small + { // a class left with fewer slots than a prompt chunk's layer can use (ne_) would stop with "no free slot": a small // arena (a Mac with little free memory at start) gets the prompt's buffers as their own allocation instead std::vector stolen(cls_.size(), 0); for (int s : region_slots_) ++stolen[(size_t) cls_of(s)]; bool tight = false; - for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < NE); + for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < ne_); if (tight) return standalone_region(bytes); } region_ = lo; @@ -774,7 +785,7 @@ void ExpertStream::hold_region(bool hold) { for (int s : region_slots_) { Slot& x = slots_[(size_t) s]; if (x.layer >= 0) { - const int key = x.layer * NE + x.e; + const int key = x.layer * ne_ + x.e; where_[(size_t) key] = -1; smap_.as()[key] = 0xFFFFFFFFu; } diff --git a/metal/src/engine/stream.h b/metal/src/engine/stream.h index 1a72af5f362..afde20989a8 100644 --- a/metal/src/engine/stream.h +++ b/metal/src/engine/stream.h @@ -57,6 +57,7 @@ class ExpertStream { // the arena's buffers (at most 4, each within the device's maxBufferLength): chunk(i) for i >= count is chunk 0 const mtl::Buffer& chunk(int i) const { return chunks_[(size_t) i < chunks_.size() ? (size_t) i : 0]; } const mtl::Buffer& smap() const { return smap_; } + int n_expert() const { return ne_; } uint64_t up_in_slot(int l) const { return lay_[(size_t) l].gs; } uint64_t down_in_slot(int l) const { return lay_[(size_t) l].gs + lay_[(size_t) l].us; } // The host's turn at layer l: ids = the window's T x topk routed ids; pred = T x 512 router logits of layer l + 1 @@ -118,6 +119,7 @@ class ExpertStream { bool stop_ = false; int outstanding_ = 0; int nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2; + int ne_ = 512; bool prefetch_ = true, fake_ = false, pre_mru_ = false; // stats uint64_t stamp_ = 1; // the handshake count: slots pinned with the current one are in use on the GPU diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp index add35c0e9f3..0ee8b1ee138 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -127,6 +127,7 @@ class Window { int64_t vocab() const { return n_vocab_; } bool ready() const { return cx_ != nullptr; } bool streaming() const { return stream_; } + bool supports_metal4() const { return cx_ != nullptr && cx_->supports_metal4(); } // --serve's conversation checkpoints: the window's running state (the GDN recurrences and conv histories, the // PLE history, each QSA layer's indexer tail and tail-block key); the KV and pooled keys are positional bool checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err); @@ -273,8 +274,9 @@ class Window { // ---- SSD streaming (stream.h): the dense weights in W_, the routed experts in the arena std::unique_ptr st_; bool stream_ = false, st_pred_ = false; + int n_expert_ = 512; std::vector* cuts_ = nullptr; - mtl::Buffer predL_, predL2_; // 48 x MAXT x 512: layer l + 1's (l + 2's) router logits from layer l's FFN input + mtl::Buffer predL_, predL2_; // 48 x MAXT x n_expert_: layer l + 1's (l + 2's) router logits from layer l's FFN input int seg_T_ = 0; bool stream_arena(std::string& err); double run_ds(Cached& c); @@ -302,8 +304,8 @@ class Window { int win_n_[MAXT + 1] = {}; double win_ms_[MAXT + 1] = {}, win_gpu_[MAXT + 1] = {}; void at_cut(int s) { - st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as(), st_pred_ && !cut_ids_ ? predL_.as() + (size_t) s * MAXT * 512 : nullptr, - predL2_.ptr && !cut_ids_ ? predL2_.as() + (size_t) s * MAXT * 512 : nullptr); + st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as(), st_pred_ && !cut_ids_ ? predL_.as() + (size_t) s * MAXT * n_expert_ : nullptr, + predL2_.ptr && !cut_ids_ ? predL2_.as() + (size_t) s * MAXT * n_expert_ : nullptr); } private: std::vector ple_hist_; // NG_HIST x hc_dim normalized rows, row-fastest @@ -437,10 +439,11 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta ss_ = &ss; max_t_ = std::min(max_t, MAXT); max_cells_ = max_cells; - if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) { + if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || (g.n_expert != 512 && g.n_expert != 256) || g.n_layers != 48) { err = "metal window: the geometry differs from Qwen3.8-Flash-Next's"; return false; } + n_expert_ = (int) g.n_expert; eps_ = 1e-6f; const double t0 = now_ms(); if (const long rp = repack_pid_running(shard1); rp > 0 && !(std::getenv("STRATA_METAL_IGNORE_REPACK") && std::string(std::getenv("STRATA_METAL_IGNORE_REPACK")) == "1")) { @@ -501,7 +504,7 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta auto A = [&](size_t bytes) { return cx_->alloc(bytes); }; R_ = A(4 * MAXT * HC * N); bo_ = A(4 * MAXT * N); inj_[0] = A(4 * MAXT * HC); inj_[1] = A(4 * MAXT * HC); xn_ = A(4 * MAXT * HC * N); lo_ = A(4 * MAXT * LR); qkv_ = A(4 * MAXT * 10240); z_ = A(4 * MAXT * 6144); - al_ = A(4 * MAXT * 48); be_ = A(4 * MAXT * 48); yg_ = A(4 * MAXT * 6144); logits_ = A(4 * MAXT * 512); + al_ = A(4 * MAXT * 48); be_ = A(4 * MAXT * 48); yg_ = A(4 * MAXT * 6144); logits_ = A(4 * MAXT * n_expert_); ids_ = A(4 * MAXT * TOPK); wts_ = A(4 * MAXT * TOPK); plan_ = A(4 * P_INTS); h_ = A(4 * MAXT * TOPK * NFF); ye_ = A(4 * MAXT * TOPK * N); gsh_ = A(4 * MAXT * NFF); ush_ = A(4 * MAXT * NFF); ysh_ = A(4 * MAXT * N); sg_ = A(4 * MAXT * 4); qfull_ = A(4 * MAXT * 12288); kcur_ = A(4 * MAXT * 512); vcur_ = A(4 * MAXT * 512); @@ -624,8 +627,8 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta if (!st_->load_dense(f_, *cx_, skip, W_, err)) return false; const char* e = std::getenv("STRATA_METAL_PREFETCH"); st_pred_ = !(e && std::atoi(e) <= 0); - if (st_pred_) predL_ = A(4 * (size_t) 48 * MAXT * 512); - if (st_pred_ && st_->ahead2() > 0) predL2_ = A(4 * (size_t) 48 * MAXT * 512); + if (st_pred_) predL_ = A(4 * (size_t) 48 * MAXT * n_expert_); + if (st_pred_ && st_->ahead2() > 0) predL2_ = A(4 * (size_t) 48 * MAXT * n_expert_); } ple_on_ = ss.ple.ready() && f_.find("blk.1.ple_key.weight") != nullptr; if (ple_on_ && T("blk.1.ple_conv1d.weight").type == ref::F32) { @@ -882,17 +885,17 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 const gguf::Tensor& trt = T(b + "ffn_gate_inp.weight"); const bool rf = routef_ && (trt.type == ref::BF16 || trt.type == ref::F32); if (!rf) { - gemv(ds, trt, XF_, logits_, 512, ntok); + gemv(ds, trt, XF_, logits_, n_expert_, ntok); nb_ = true; } if (stream_ && st_pred_ && l + 1 < 48) { // the next layer's routing guessed from this layer's FFN input nb_ = !rf; // beside the router GEMV; with router_route it is the first reader of XF_ here: a barrier - gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, 512, ntok, - (size_t) l * MAXT * 512 * 4); + gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, n_expert_, ntok, + (size_t) l * MAXT * n_expert_ * 4); if (predL2_.ptr && l + 2 < 48) { nb_ = true; - gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, 512, ntok, - (size_t) l * MAXT * 512 * 4); + gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, n_expert_, ntok, + (size_t) l * MAXT * n_expert_ * 4); } nb_ = true; } @@ -922,7 +925,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 a.n_ff = NFF; a.n_tok = ntok; a.topk = TOPK; - a.n_expert = NEXP; + a.n_expert = n_expert_; const bool dq2 = td.type == ref::Q2_0; // the trellis gate/up: 16 lanes a row, 2 rows a lane (expert_parity --bench --sweep, M4, TQK6, 10 experts: 166.6 us a layer // against 174.5 for 32,2,4 and 196.8 for 32,1,4; also the fastest at 2 and 3 tokens) @@ -939,17 +942,17 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 if (rf) { // the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs r.p = P(lib_, "router_route_" + std::to_string(trt.type), mtl::Constants().i(0, ntok), "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok)); - GemvArgs ga{(uint32_t) N, 512, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512}; + GemvArgs ga{(uint32_t) N, (uint32_t) n_expert_, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, (uint32_t) n_expert_}; r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4), mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; - r.grid[0] = 512 / 16; - r.tg[0] = 512; + r.grid[0] = n_expert_ / 16; + r.tg[0] = n_expert_; nb_ = true; } else { r.p = P(lib_, "moe_route", {}, "moe_route"); r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; - r.tg[0] = 512; + r.tg[0] = n_expert_; } push(ds, r); { @@ -981,7 +984,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; if (stream_) for (int c = -1; c < 3; ++c) - g.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); + g.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * n_expert_ * 4) : mtl::Arg::B(st_->chunk(c + 1))); g.tg[0] = 32 * gs; g.indirect = &plan_; g.ind_off = P_IGU * 4; @@ -995,7 +998,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 dd.args = {mtl::Arg::B(EW, ebase), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; if (stream_) for (int c = -1; c < 3; ++c) - dd.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); + dd.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * n_expert_ * 4) : mtl::Arg::B(st_->chunk(c + 1))); dd.tg[0] = 32 * ds_; dd.indirect = &plan_; dd.ind_off = P_IDN * 4; @@ -1319,7 +1322,7 @@ bool Window::stream_arena(std::string& err) { "resident of 128)" : ""); } const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM"); - const bool pg_on = !(epg && std::string(epg) == "0"); + const bool pg_on = !(epg && std::string(epg) == "0") && cx_->supports_metal4(); bool minimal = false; // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps for (int tries = 0;; ++tries) { const double planned = budget; @@ -2146,6 +2149,7 @@ bool Window::prompt_gemm_on(int64_t n) const { bool Window::prompt_gemm_init(std::string& err) { if (pg_) return true; + if (!cx_ || !cx_->supports_metal4()) { err = "metal: prompt gemm requires Metal 4"; return false; } if (stream_ && !st_->ready()) return stream_arena(err); // it makes the prompt buffers after the arena const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK"); // Streaming: every chunk makes one pass over the layer's experts that are not in the arena, so a longer chunk @@ -2172,7 +2176,7 @@ bool Window::prompt_gemm_init(std::string& err) { const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA"); // v1 | flash (default: flash with the M5's layout) p.flash = eq ? std::string(eq) == "flash" : p.k4; const size_t MT = (size_t) C; - p.maxrows = C * TOPK + NEXP * p.TM; + p.maxrows = C * TOPK + n_expert_ * p.TM; p.nblk_max = (C * TOPK + 1023) / 1024; const int MB = (int) max_cells_ / 4 + 2; // the buffers: measured first, then placed - on a streaming Mac inside the expert arena (a region of its slots @@ -2202,9 +2206,9 @@ bool Window::prompt_gemm_init(std::string& err) { p.qids = AB(4 * MT * 2051); p.part = p.flash ? F32(16) : F32((size_t) p.QS * 33 * 24 * 258); p.XQ = F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384); p.xqh = F16(MT * 6144); - p.logits = F32(MT * NEXP); p.ids = AB(4 * MT * TOPK); p.wts = F32(MT * TOPK); - p.bcnt = AB(4 * (size_t) p.nblk_max * NEXP); p.boff = AB(4 * (size_t) p.nblk_max * NEXP); - p.eoff = AB(4 * NEXP); p.tiles = AB(8 * (mr / p.TM + 1)); p.meta = AB(64); + p.logits = F32(MT * n_expert_); p.ids = AB(4 * MT * TOPK); p.wts = F32(MT * TOPK); + p.bcnt = AB(4 * (size_t) p.nblk_max * n_expert_); p.boff = AB(4 * (size_t) p.nblk_max * n_expert_); + p.eoff = AB(4 * n_expert_); p.tiles = AB(8 * (mr / p.TM + 1)); p.meta = AB(64); p.stok = AB(4 * mr); p.sw = F32(mr); p.pos = AB(4 * MT * TOPK); p.XS = F16(mr * N); p.HS = F16(mr * NFF); p.YE = F32(mr * N); p.yr = F32(MT * N); p.gsh = F32(MT * NFF); p.ush = F32(MT * NFF); p.hh = F16(MT * NFF); p.ys = F32(MT * N); p.sgv = F32(MT); @@ -2212,7 +2216,7 @@ bool Window::prompt_gemm_init(std::string& err) { p.WD = F16((size_t) (LR + HC) * HC * N); p.WU = F16((size_t) HC * N * LR); p.Wqkv = F16((size_t) 10240 * N); p.Wz = F16((size_t) 6144 * N); p.Wal = F16((size_t) 48 * N); p.Wbe = F16((size_t) 48 * N); p.Wout = F16((size_t) N * 6144); p.Wq = F16((size_t) 12288 * N); p.Wk = F16((size_t) 512 * N); p.Wv = F16((size_t) 512 * N); p.Wiq = F16((size_t) 512 * N); - p.Wik = F16((size_t) 128 * N); p.Wr = F16((size_t) NEXP * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); + p.Wik = F16((size_t) 128 * N); p.Wr = F16((size_t) n_expert_ * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); p.Wshu = F16((size_t) NFF * N); p.Wshd = F16((size_t) N * NFF); if (ple_on_) { p.emb = F32(MT * N); p.embh = F16(MT * N); p.pkey = F32(MT * 10240); p.pval = F32(MT * N); p.keyn = F32(MT * 10240); @@ -2427,16 +2431,16 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri dq(tg_, p.Wshg); dq(tu_, p.Wshu); dq(td_, p.Wshd); - gemm(p.mixh, p.Wr, p.logits, T, NEXP, N); + gemm(p.mixh, p.Wr, p.logits, T, n_expert_, N); const uint32_t nt = (uint32_t) T; - D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256); + D(PP(p.lib, n_expert_ == 256 ? "route_topk_256" : "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256); const int nent = T * TOPK, nblk = (nent + 1023) / 1024; - PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk}; + PgPlan pa{(uint32_t) nent, TOPK, (uint32_t) n_expert_, (uint32_t) p.TM, (uint32_t) nblk}; const uint32_t nclr = (uint32_t) p.maxrows; D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256); D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles), - mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, NEXP); + mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, n_expert_); D(PP(p.lib, "mm_place"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.stok), mtl::Arg::B(p.sw), mtl::Arg::B(p.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); PgMm ma{eg.offset, TN(b + "ffn_up_exps.weight").offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), @@ -2460,7 +2464,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri const mtl::Buffer& EW = stream_ ? st_->chunk(0) : W_; auto arena_args = [&](mtl::Dispatch& d) { if (!stream_) return; - d.args.push_back(mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4)); + d.args.push_back(mtl::Arg::B(st_->smap(), (size_t) l * n_expert_ * 4)); for (int c = 1; c < 4; ++c) d.args.push_back(mtl::Arg::B(st_->chunk(c))); }; const std::string ak = stream_ ? "#arena" : ""; @@ -2691,8 +2695,9 @@ bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, if (g_window == nullptr) { err = "metal prompt: the decode window is not up yet"; return false; } static const bool windows = std::getenv("STRATA_METAL_PROMPT") && std::string(std::getenv("STRATA_METAL_PROMPT")) == "window"; static const int64_t min_rows = std::getenv("STRATA_METAL_PROMPT_MIN") ? std::atoll(std::getenv("STRATA_METAL_PROMPT_MIN")) : 32; - if (!windows && !g_window->streaming() && n >= min_rows) return g_window->prompt(tokens, n, pos0, on_chunk, err); - if (!windows && g_window->streaming()) { + const bool m4 = g_window->supports_metal4(); + if (!windows && m4 && !g_window->streaming() && n >= min_rows) return g_window->prompt(tokens, n, pos0, on_chunk, err); + if (!windows && m4 && g_window->streaming()) { const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM"); const char* em = std::getenv("STRATA_METAL_PROMPT_GEMM_MIN"); const int64_t mn = em ? std::atoll(em) : 64; diff --git a/metal/tests/coop_probe.h b/metal/tests/coop_probe.h index dc781ff4661..7ca2334dfd0 100644 --- a/metal/tests/coop_probe.h +++ b/metal/tests/coop_probe.h @@ -9,6 +9,7 @@ inline bool coop_k4(mtl::Context& cx) { if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0; + if (!cx.supports_metal4()) return false; const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n"); for (const char* tile : {"coop_layout_16_64_32", "coop_layout_16_32_32"}) { mtl::Pipeline p = cx.pipeline(lib, tile); diff --git a/metal/tools/stream_experts.py b/metal/tools/stream_experts.py index 4eedafbe616..4a1b4df58ae 100644 --- a/metal/tools/stream_experts.py +++ b/metal/tools/stream_experts.py @@ -32,23 +32,25 @@ def plan(gguf: Path): first data offset.""" rd = R.Reader(gguf) tab = {t["name"]: t for t in rd.tensors} + g0 = tab.get("blk.0.ffn_gate_exps.weight") + ne = g0["dims"][-1] if (g0 and len(g0["dims"]) > 1) else NE lay = [] for l in range(NL): g, u, d = (tab[f"blk.{l}.ffn_{k}_exps.weight"] for k in ("gate", "up", "down")) - lay.append((g, u, d, g["bytes"] // NE, u["bytes"] // NE, d["bytes"] // NE)) + lay.append((g, u, d, g["bytes"] // ne, u["bytes"] // ne, d["bytes"] // ne)) head_len = 64 + NL * 32 data0 = (head_len + ALIGN - 1) // ALIGN * ALIGN offs, at = [], data0 for g, u, d, gs, us, ds in lay: offs.append(at) - at += NE * (gs + us + ds) + at += ne * (gs + us + ds) at = (at + ALIGN - 1) // ALIGN * ALIGN - return rd, lay, offs, at, data0 + return rd, lay, offs, at, data0, ne def size_bytes(gguf: Path) -> int: """The size the file for this GGUF will have (setup checks the free disk with it).""" - rd, _, _, at, _ = plan(gguf) + rd, _, _, at, _, _ = plan(gguf) rd.f.close() return at @@ -60,8 +62,8 @@ def main() -> int: ap.add_argument("--verify", action="store_true") a = ap.parse_args() out = a.out or a.gguf.with_name(a.gguf.name + ".experts") - rd, lay, offs, at, data0 = plan(a.gguf) - head = struct.pack("<4sIIIQ", b"STRX", 1, NL, NE, os.path.getsize(a.gguf)).ljust(64, b"\0") + rd, lay, offs, at, data0, ne = plan(a.gguf) + head = struct.pack("<4sIIIQ", b"STRX", 1, NL, ne, os.path.getsize(a.gguf)).ljust(64, b"\0") for (g, u, d, gs, us, ds), o in zip(lay, offs): head += struct.pack(" int: src.seek(u["abs"]); U = src.read(u["bytes"]) src.seek(d["abs"]); D = src.read(d["bytes"]) f.seek(o) - f.write(b"".join(G[e * gs:(e + 1) * gs] + U[e * us:(e + 1) * us] + D[e * ds:(e + 1) * ds] for e in range(NE))) + f.write(b"".join(G[e * gs:(e + 1) * gs] + U[e * us:(e + 1) * us] + D[e * ds:(e + 1) * ds] for e in range(ne))) print(f"\r layer {l + 1} / {NL}", end="", flush=True) f.truncate(at) print() @@ -85,14 +87,14 @@ def main() -> int: src.seek(t["abs"]) h1.update(src.read(t["bytes"])) f.seek(o) - blob = f.read(NE * (gs + us + ds)) + blob = f.read(ne * (gs + us + ds)) step = gs + us + ds for part, size, base in ((0, gs, 0), (1, us, gs), (2, ds, gs + us)): - h2.update(b"".join(blob[e * step + base:e * step + base + size] for e in range(NE))) + h2.update(b"".join(blob[e * step + base:e * step + base + size] for e in range(ne))) if h1.digest() != h2.digest(): tmp.unlink() sys.exit(f"verify FAILED at layer {l} (output removed)") - print(f"verify: {NL} layers x {NE} experts identical to the GGUF") + print(f"verify: {NL} layers x {ne} experts identical to the GGUF") os.replace(tmp, out) print(f"{out}: {at / 2**30:.2f} GiB") return 0 diff --git a/setup.py b/setup.py index 0c4b1b927ab..da2608d7415 100644 --- a/setup.py +++ b/setup.py @@ -8642,12 +8642,19 @@ def mac_control_setup(a) -> int: fail(problems[0]) step(2, "your choices") - app_sizes = [m for m, d in MF.MAC_MODELS.items() if d.get("family", "qwen") in ("qwen", "gyro")] # Gyro: one file, no repack - if a.family not in (None, "qwen", "gyro"): + is_coder = a.family == "coder" or a.model == "IQ1_M" + target_family = "coder" if is_coder else ("gyro" if a.family == "gyro" or (a.model and MF.MAC_MODELS.get(a.model, {}).get("family") == "gyro") else None) + if target_family: + app_sizes = [m for m, d in MF.MAC_MODELS.items() if d.get("family", "qwen") == target_family] + else: + app_sizes = [m for m, d in MF.MAC_MODELS.items() if d.get("family", "qwen") in ("qwen", "gyro", "coder")] + if a.family not in (None, "qwen", "gyro", "coder"): fail(f"the app on a Mac runs the Qwen3.8-Flash-Next sizes ({', '.join(app_sizes)}); {a.family} is not one of them", - "choose Qwen3.8-Flash-Next") + "choose Qwen3.8-Flash-Next, coder or gyro") rec, notes = MF.recommend_model(ram) - if rec not in app_sizes: + if target_family == "coder": + rec = "IQ1_M" + elif rec not in app_sizes: rec = "Q2_0" model = a.model or rec if model not in app_sizes: diff --git a/strata_control/mac_fit.py b/strata_control/mac_fit.py index a273c741692..90f0356d4d7 100644 --- a/strata_control/mac_fit.py +++ b/strata_control/mac_fit.py @@ -20,6 +20,8 @@ "Q2_0": {"shard1_gb": 37.6, "tested": True}, "IQ3_XXS": {"shard1_gb": 47.0, "tested": True}, "IQ3_S": {"shard1_gb": 54.8, "tested": True}, + # ISTA-DASLab's Coder: half the experts (256), shard 1 is the 48-layer model, shard 2 is the n-gram table + "IQ1_M": {"shard1_gb": 29.6, "tested": True, "family": "coder"}, # Unsloth's ~4-6-bit files (family "unsloth", 4-6 shards cutting through layers): setup repacks the experts and # dense tensors into one GGUF for the GPU (metal/tools/metal_repack.py); shard1_gb is that file, the n-gram # table stays in its shard on the SSD @@ -61,6 +63,7 @@ CONTEXTS = [8192, 32768, 65536, 131072] MAC_ABOUT = {"Q2_0": "2-bit, the fastest", "IQ3_XXS": "3-bit i-quant, better quality", "IQ3_S": "3.5-bit i-quant, the best quality (matches the full model), the slowest", + "IQ1_M": "Coder ~3.5-bit (half experts, code/agentic focus)", "UD-Q4_K_XL": "Unsloth ~4.5-bit (111 GB download)", "UD-Q5_K_XL": "Unsloth ~5.5-bit (158 GB download)", "UD-Q6_K_XL": "Unsloth ~6.5-bit (169 GB download)", "Gyro-S": "agentionai's trellis experts, the smallest (58 GB download, the n-gram table is inside it)", @@ -188,11 +191,15 @@ def stream_mtp(stream: bool, ram_gb: float) -> bool: def engine_args(pack: Path, shards: list, rt: Path, ctx: int, stream: bool = False, ram_gb: float = 0.0, - profile: Path | str = "data/expert-profile.bin") -> list[str]: + profile: Path | str | None = None) -> list[str]: """The Metal defaults: --spec 2 with the MTP draft layer (measured best on the M5 Max, docs/MAC.md) and the Q8_0 hyper-connection read (the engine's default on Metal; STRATA_HC_Q8=0 in the config's env gives the exact BF16 read). Streaming from the SSD: one token per step, no draft layer (measured faster on the 32 GB M4); a streaming Mac with STREAM_MTP_FROM_GB or more keeps the draft layer and uses --spec STREAM_SPEC.""" + if profile is None or str(profile).endswith("expert-profile.bin"): + is_coder = any("coder" in str(s).lower() for s in shards) or (pack and "coder" in str(pack).lower()) + pname = "expert-profile-coder.bin" if is_coder else "expert-profile.bin" + profile = (profile.parent / pname) if isinstance(profile, Path) else pname mtp = stream_mtp(stream, ram_gb) args = ["--pack", str(pack), "--native", str(shards[0]), "--ple-gguf", str(shards[1])] if shards[0].name.endswith("-metal.gguf"): # the repacked file: it is its own dense GGUF From f171715ea757fc8c64a49be5d1a7753197175416 Mon Sep 17 00:00:00 2001 From: Aaron Z Date: Sun, 25 Oct 2026 14:33:46 -0700 Subject: [PATCH 2/2] [Engine] Metal: fix fused router tg size for 256 experts, wire streaming expert profile and log prompt path --- metal/src/engine/window.cpp | 26 +++++++++++++++++++++----- src/program/generate.cpp | 6 +++++- strata_control/mac_fit.py | 17 ++++++++++++----- 3 files changed, 38 insertions(+), 11 deletions(-) diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp index 0ee8b1ee138..28b8ba6dbd1 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -947,7 +947,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; r.grid[0] = n_expert_ / 16; - r.tg[0] = n_expert_; + r.tg[0] = 512; // 16 simdgroups to compute 16 rows per threadgroup, even when n_expert_ == 256 nb_ = true; } else { r.p = P(lib_, "moe_route", {}, "moe_route"); @@ -1230,6 +1230,8 @@ double Window::run_ds(Cached& c) { namespace { Window* g_stream_window = nullptr; +std::string g_shard1; +std::string g_expert_profile; // What the process holds now, all of it (task_info TASK_VM_INFO): phys_footprint is its anonymous + wired + compressed memory // (not the clean file pages it only maps), resident_size is the RSS (those pages too). Bytes; 0 when the call fails. struct ProcMem { @@ -1305,7 +1307,11 @@ bool Window::stream_arena(std::string& err) { } std::string prof; if (const char* e = std::getenv("STRATA_METAL_EXPERT_PROFILE")) prof = e; - else prof = strata_metal::resource(nullptr, "data", STRATA_METAL_DATA_DEFAULT) + "/expert-profile.bin"; + else if (!g_expert_profile.empty()) prof = g_expert_profile; + else { + const std::string name = n_expert_ == 256 ? "/expert-profile-coder.bin" : "/expert-profile.bin"; + prof = strata_metal::resource(nullptr, "data", STRATA_METAL_DATA_DEFAULT) + name; + } std::fprintf(stderr, "strata metal stream: %.2f GB allocated before the arena (the process holds %.2f GB, %.2f GB resident); " "RAM %.1f GB (%.1f GB available, %.1f GB kept for macOS), GPU working set %.1f GB\n", used / gb, pm0.footprint / gb, pm0.rss / gb, phys / gb, avail / gb, keep / gb, cx_->recommended_working_set() / gb); @@ -2583,12 +2589,10 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri } // ---------------------------------------------------------------------------------------------- the hook's entry points -namespace { -std::string g_shard1; -} void window_register(void* w); bool mtp_load_pending(Window* w, std::string& err); void set_model_shard(const char* shard1) { g_shard1 = shard1 ? shard1 : ""; } +void set_expert_profile(const char* prof) { g_expert_profile = prof ? prof : ""; } // One decode window per process: the prompt path (Prefill::init) may bring it up before the Verifier does. namespace { @@ -2696,6 +2700,18 @@ bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, static const bool windows = std::getenv("STRATA_METAL_PROMPT") && std::string(std::getenv("STRATA_METAL_PROMPT")) == "window"; static const int64_t min_rows = std::getenv("STRATA_METAL_PROMPT_MIN") ? std::atoll(std::getenv("STRATA_METAL_PROMPT_MIN")) : 32; const bool m4 = g_window->supports_metal4(); + static bool s_logged_path = false; + if (!s_logged_path) { + s_logged_path = true; + if (!windows && m4 && !g_window->streaming()) { + std::fprintf(stderr, "strata metal prompt: Metal 4 resident prefill path\n"); + } else if (!windows && m4 && g_window->streaming()) { + std::fprintf(stderr, "strata metal prompt: Metal 4 streaming prefill path\n"); + } else { + std::fprintf(stderr, "strata metal prompt: universal decode window prefill path%s\n", + !m4 ? " (no Metal 4 on this OS/device)" : (windows ? " (STRATA_METAL_PROMPT=window)" : "")); + } + } if (!windows && m4 && !g_window->streaming() && n >= min_rows) return g_window->prompt(tokens, n, pos0, on_chunk, err); if (!windows && m4 && g_window->streaming()) { const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM"); diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 956d2c56254..3ee316b4203 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -129,7 +129,10 @@ #include #if defined(STRATA_GPU_METAL) -namespace strata::metal_engine { void set_model_shard(const char* shard1); } // metal/src/engine/window.cpp +namespace strata::metal_engine { +void set_model_shard(const char* shard1); // metal/src/engine/window.cpp +void set_expert_profile(const char* prof); +} #endif #include @@ -3971,6 +3974,7 @@ int main(int argc, char** argv) { const char* isa_floor = strata::kernels::cpu::isa_floor_build(); #if defined(STRATA_GPU_METAL) // Apple Silicon: every expert is on the GPU; ggml-cpu (NEON) is the only CPU expert code strata::metal_engine::set_model_shard(o.native_preset.c_str()); // the decode window reads the model GGUF itself + strata::metal_engine::set_expert_profile(o.expert_profile.c_str()); o.mmap_experts = true; // ... experts included: no resident copy of them (the GGUF shards in place) if (false) { #else diff --git a/strata_control/mac_fit.py b/strata_control/mac_fit.py index 90f0356d4d7..6c7045584ce 100644 --- a/strata_control/mac_fit.py +++ b/strata_control/mac_fit.py @@ -191,15 +191,22 @@ def stream_mtp(stream: bool, ram_gb: float) -> bool: def engine_args(pack: Path, shards: list, rt: Path, ctx: int, stream: bool = False, ram_gb: float = 0.0, - profile: Path | str | None = None) -> list[str]: + profile: Path | str = "data/expert-profile.bin") -> list[str]: """The Metal defaults: --spec 2 with the MTP draft layer (measured best on the M5 Max, docs/MAC.md) and the Q8_0 hyper-connection read (the engine's default on Metal; STRATA_HC_Q8=0 in the config's env gives the exact BF16 read). Streaming from the SSD: one token per step, no draft layer (measured faster on the 32 GB M4); a streaming Mac with STREAM_MTP_FROM_GB or more keeps the draft layer and uses --spec STREAM_SPEC.""" - if profile is None or str(profile).endswith("expert-profile.bin"): - is_coder = any("coder" in str(s).lower() for s in shards) or (pack and "coder" in str(pack).lower()) - pname = "expert-profile-coder.bin" if is_coder else "expert-profile.bin" - profile = (profile.parent / pname) if isinstance(profile, Path) else pname + if profile is None: + profile = "data/expert-profile.bin" + is_coder = any("coder" in str(s).lower() for s in shards) or (pack and "coder" in str(pack).lower()) + if is_coder and str(profile).endswith("expert-profile.bin"): + pname = "expert-profile-coder.bin" + if isinstance(profile, Path): + profile = profile.parent / pname + elif "/" in str(profile): + profile = str(Path(profile).parent / pname) + else: + profile = pname mtp = stream_mtp(stream, ram_gb) args = ["--pack", str(pack), "--native", str(shards[0]), "--ple-gguf", str(shards[1])] if shards[0].name.endswith("-metal.gguf"): # the repacked file: it is its own dense GGUF