[Engine] support Qwen3.8-Flash-Next Coder (256 experts) on Apple Silicon Metal

.gitignore

@@ -29,6 +29,7 @@ td.log
 /.venv/
 /build/
 /third_party/llama.cpp
+/third_party/metal-cpp/
 /models/
 /packs/
 /mtp/

metal/kernels/dense_mm.metal

@@ -3,8 +3,10 @@
 //   dq_f16<F>  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(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
 #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
 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(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
 template <int TM, int TN, int SG>
 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<F>) dq_f16<F>;
 #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(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
 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
+

metal/kernels/mm_bench.metal

@@ -2,6 +2,7 @@
 //   mpp_mm<T, TM, TN, SG>   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<T>                simdgroup_matrix 8x8, a 32 x 32 tile per simdgroup (4 x 4 accumulators), 4 simdgroups
+#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
 #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
 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 <typename T>
@@ -53,6 +57,7 @@ template [[host_name("sg_mm_half")]] kernel decltype(sg_mm<half>) sg_mm<half>;
 template [[host_name("sg_mm_float")]] kernel decltype(sg_mm<float>) sg_mm<float>;
 
 
+#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
 // 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 <int TM, int TN, int TK>
@@ -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

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);

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

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);