Metal: instantiate the K-quant / IQ down kernels only for a model that has them

Measured on the M5 Max (cold start, same parity binaries, kernel sources of
main against the branch): the extra instantiations cost moe_mm.metal's library
about +0.5 s on every start (5.7 -> 6.2 s) and moe.metal's about +0.65 s on a
cold shader cache.  A Flash-Next set should not pay for types it never uses.

The new instantiations (DN_SET IQ4_XS / IQ3_S, DNK_SET, MM_DN2_KQ and the other
down tiles, MM3_DN, MM6_DN1) sit under #ifdef STRATA_KQ_DOWN.  window.cpp
kq_down_prelude defines it when any layer's ffn_down_exps is Q3_K, Q4_K, Q5_K,
Q6_K, IQ3_S or IQ4_XS, at the three compile sites of moe.metal / moe_mm.metal
(decode window, prompt GEMM setup, prompt chain).  The parity tests do the same
from the down type.  For today's sets the kernel text is what it was.

metal/kernels/moe.metal

@@ -446,10 +446,12 @@ DN_SET(T_Q2_0R)
 DN_SET(T_IQ4_NL)
 // UD files at expert widths divisible by 256: IQ4_XS / IQ3_S downs read raw h like the gate/up's raw-kind formats (the generic
 // X / Dec); the K-quants take moe_down_k (XD above).  Same host names, moe_down_<type>_<L>_<R>_<S>.
+#ifdef STRATA_KQ_DOWN   // the host defines it when a layer's down is one of these types (window.cpp kq_down_prelude): today's sets compile exactly what they did
 DN_SET(T_IQ4_XS) DN_SET(T_IQ3_S)
 #define DNK(F, L, R, S) template [[host_name("moe_down_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_down_k<F, L, R, S>) moe_down_k<F, L, R, S>;
 #define DNK_SET(F) DNK(F, 8, 1, 4) DNK(F, 8, 2, 4) DNK(F, 8, 4, 4) DNK(F, 8, 1, 8) DNK(F, 8, 2, 8) DNK(F, 4, 1, 4) DNK(F, 4, 2, 4)                    DNK(F, 4, 4, 4) DNK(F, 16, 1, 4) DNK(F, 16, 2, 4) DNK(F, 32, 1, 4) DNK(F, 2, 2, 4)
 DNK_SET(T_Q4_K) DNK_SET(T_Q5_K) DNK_SET(T_Q6_K) DNK_SET(T_Q3_K)
+#endif   // STRATA_KQ_DOWN
 // Gyro (agentionai's trellis experts): TQK6 / TQ2_T gate/up, TQK7 / TQ2_T down (the CUDA engine's STRATA_GU_FMTS / D_FMTS)
 GU_SET(T_TQK6) GU_SET(T_TQ2T)
 DN_SET(T_TQK7) DN_SET(T_TQ2T)

metal/kernels/moe_mm.metal

@@ -391,11 +391,13 @@ kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS
 MM_GU2_KQ(T_Q4_K) MM_GU2_KQ(T_Q5_K) MM_GU2_KQ(T_Q6_K) MM_GU2_U32(T_Q8_0) MM_DN2_U32(T_Q5_1) MM_DN2_U32(T_Q5_0) MM_DN2_U32(T_Q8_0)
 // K-quant / IQ4_XS / IQ3_S down experts (a UD file's downs at expert widths divisible by 256: Q4-KQ): the staged K chunk
 // follows the same rule as the gate/up's (Q4_K / Q5_K / Q6_K 128, Q3_K and IQ4_XS 32, IQ3_S 16)
+#ifdef STRATA_KQ_DOWN   // (see moe.metal)
 #define MM_DN2_KQ(F) MMDN2(F, 16, 16, 128, 2) MMDN2(F, 16, 32, 128, 2) MMDN2(F, 32, 32, 128, 2) MMDN2(F, 16, 32, 128, 4)
 MM_DN2_KQ(T_Q4_K) MM_DN2_KQ(T_Q5_K) MM_DN2_KQ(T_Q6_K)
 MMDN2(T_IQ4_XS, 16, 64, 32, 2) MMDN2(T_IQ4_XS, 16, 64, 64, 2) MMDN2(T_IQ4_XS, 16, 32, 32, 2)
 MMDN2(T_Q3_K, 16, 64, 32, 2) MMDN2(T_Q3_K, 16, 32, 32, 2)
 MMDN2(T_IQ3_S, 16, 64, 16, 2) MMDN2(T_IQ3_S, 16, 64, 32, 2) MMDN2(T_IQ3_S, 16, 32, 32, 2)
+#endif
 // IQ4_XS (UNIT 32): the K chunk must be >= 32
 MMGU2(T_IQ4_XS, 16, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 64, 64, 2) MMGU2(T_IQ4_XS, 32, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 32, 32, 2)
 MMGU2(T_IQ4_XS, 64, 64, 32, 4)
@@ -1050,7 +1052,9 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS
     MMDN3(F, 16, 32, 16, 2) MMDN3(F, 32, 32, 32, 2) MMDN3(F, 16, 64, 32, 1) MMDN3(F, 16, 64, 32, 4)     MMDN3(F, 16, 64, 16, 2) MMDN3(F, 16, 32, 16, 4) MMDN3(F, 16, 64, 32, 5) MMDN3(F, 16, 64, 16, 4)
 MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S)
 MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q5_0) MM3_DN(T_Q8_0)
-MM3_DN(T_Q4_K) MM3_DN(T_Q5_K) MM3_DN(T_Q6_K) MM3_DN(T_Q3_K) MM3_DN(T_IQ4_XS) MM3_DN(T_IQ3_S)   // Q4-KQ: K-quant / IQ downs of UD files at expert widths divisible by 256
+#ifdef STRATA_KQ_DOWN
+MM3_DN(T_Q4_K) MM3_DN(T_Q5_K) MM3_DN(T_Q6_K) MM3_DN(T_Q3_K) MM3_DN(T_IQ4_XS) MM3_DN(T_IQ3_S)
+#endif
 
 // mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token
 // tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times
@@ -1427,4 +1431,6 @@ MM6_DN(T_Q5_1) MM6_DN(T_Q8_0)
 #define MM6_DN1(F) MMDN6(F, 4, 32, 32, 4)
 MM6_GU1(T_Q6_K) MM6_GU1(T_Q8_0) MM6_GU1(T_IQ3_XXS) MM6_GU1(T_IQ2_S) MM6_GU1(T_IQ4_XS) MM6_GU1(T_IQ3_S) MM6_GU1(T_IQ2_XS) MM6_GU1(T_IQ2_XXS)
 MM6_DN1(T_Q5_0) MM6_DN1(T_IQ4_NL) MM6_DN1(T_Q2_0)
-MM6_DN1(T_Q4_K) MM6_DN1(T_Q5_K) MM6_DN1(T_Q6_K) MM6_DN1(T_Q3_K) MM6_DN1(T_IQ4_XS) MM6_DN1(T_IQ3_S)   // Q4-KQ
+#ifdef STRATA_KQ_DOWN
+MM6_DN1(T_Q4_K) MM6_DN1(T_Q5_K) MM6_DN1(T_Q6_K) MM6_DN1(T_Q3_K) MM6_DN1(T_IQ4_XS) MM6_DN1(T_IQ3_S)
+#endif

metal/src/engine/prompt_chain.inc

@@ -52,7 +52,7 @@ bool Window::pc_setup(int C) {
     const double t0 = now_ms();
     if (b.C == 0) {
         b.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal",
-                              "gemv.metal", "prompt.metal"});
+                              "gemv.metal", "prompt.metal"}, kq_down_prelude(f_));
         b.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
         b.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
         b.k4 = coop_k4();

metal/src/engine/window.cpp

@@ -97,6 +97,20 @@ namespace {
 
 int g_slot_cap = 0;   // --batch N: the window keeps N + 1 sequences' running states side by side (set_batch_slots)
 
+// Q4-KQ: the kernel sources define STRATA_KQ_DOWN to instantiate the K-quant / IQ4_XS / IQ3_S down experts (moe.metal, moe_mm.metal).
+// The host asks for them only when a layer's down tensor has one of those types, so a Flash-Next set (Q5_1 / Q5_0 / Q8_0 / IQ4_NL / Q2_0
+// downs) compiles the same kernel text as before and pays nothing for them.
+std::string kq_down_prelude(const gguf::File& f) {
+    for (int l = 0; l < 128; ++l) {
+        const gguf::Tensor* t = f.find("blk." + std::to_string(l) + ".ffn_down_exps.weight");
+        if (!t) continue;
+        if (t->type == ref::Q3_K || t->type == ref::Q4_K || t->type == ref::Q5_K || t->type == ref::Q6_K || t->type == ref::IQ3_S ||
+            t->type == ref::IQ4_XS)
+            return "#define STRATA_KQ_DOWN 1\n";
+    }
+    return "";
+}
+
 constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8, NVOCAB_MAX = 248320;
 constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4,
               P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8;
@@ -704,7 +718,7 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta
     n_vocab_ = (int64_t) T("output.weight").shape[1];
     if (const gguf::Tensor* e0 = f_.find("blk.0.ffn_gate_exps.weight")) gyro_dense_ = e0->type >= 144;
     lib_ = cx_->compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", "gr2.metal",
-                         "glue.metal"}, metal_f16() ? "#define STRATA_METAL_F16 1\n" : "");
+                         "glue.metal"}, (metal_f16() ? std::string("#define STRATA_METAL_F16 1\n") : std::string()) + kq_down_prelude(f_));
     lgdn_ = cx_->compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n");
     lqsa_ = cx_->compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n");
     lsmp_ = cx_->compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n");
@@ -2864,7 +2878,7 @@ bool Window::prompt_gemm_init(std::string& err) {
     p.MT = C;
     p.QS = std::min(C, 256);
     p.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal",
-                          "prompt.metal"});
+                          "prompt.metal"}, kq_down_prelude(f_));
     p.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
     p.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
     p.k4 = ::coop_k4(*cx_);

metal/tests/expert_parity.cpp

@@ -282,7 +282,10 @@ int main(int argc, char** argv) {
         std::fprintf(stderr, "EXPERT_F16: Q2_0 gate/up and down only, without --dn-prep and Hadamard\n");
         return 2;
     }
-    const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}, g_f16 ? "#define STRATA_METAL_F16 1\n" : "");
+    // the host defines STRATA_KQ_DOWN only for a model whose downs are K-quant / IQ4_XS / IQ3_S (window.cpp kq_down_prelude)
+    const bool kq_dn = dn == ref::Q3_K || dn == ref::Q4_K || dn == ref::Q5_K || dn == ref::Q6_K || dn == ref::IQ3_S || dn == ref::IQ4_XS;
+    const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"},
+                               std::string(g_f16 ? "#define STRATA_METAL_F16 1\n" : "") + (kq_dn ? "#define STRATA_KQ_DOWN 1\n" : ""));
     const Layout L = layout(gu, dn);
     // Gyro: the pack's hadamard.txt (the signs of the widths 2560 and 640) -> the engine's CPU spec (the reference)
     mtl::Buffer hsx, hsh;

metal/tests/moe_mm_parity.cpp

@@ -58,7 +58,9 @@ int main(int argc, char** argv) {
     if (v4) TM = 16 * MS;   // the plan's tile height; mm_gu4's name carries MS instead
     if (v6) TM = 16;   // ragged tiles: the padding is 16 rows, the group size is MS
     mtl::Context cx;
-    const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"});
+    // the host defines STRATA_KQ_DOWN only for a model whose downs are K-quant / IQ4_XS / IQ3_S (window.cpp kq_down_prelude)
+    const bool kq_dn = dn == ref::Q3_K || dn == ref::Q4_K || dn == ref::Q5_K || dn == ref::Q6_K || dn == ref::IQ3_S || dn == ref::IQ4_XS;
+    const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}, kq_dn ? "#define STRATA_KQ_DOWN 1\n" : "");
     const int NE = bench ? 512 : 64, LAYERS = bench ? 4 : 1;
     const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF);
     const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row;