Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 18:35:04 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

metal: IQ4_XS routed-expert gate/up for decode (moe_gu_23_*) and prompt (mm_gu3 W4<IQ4_XS>, mm_gu2 TK >= 32 tiles) - the IQ3_S pack's one IQ4_XS layer; parity on the M5 and the M4

Browse Source
This commit is contained in:
Niko1221 committed 2026-10-12 19:41:13 +02:00
1 parent 5565a347de
commit b6480437d2
5 files changed
+26 -4
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

metal/kernels/moe.metal
+1
View File
@@ -302,6 +302,7 @@ GU_SET(T_IQ2_XXS)
GU_SET(T_IQ2_XS)
GU_SET(T_IQ2_S)
GU_SET(T_IQ3_S)
GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer
DN_SET(T_Q2_0)
DN_SET(T_Q2_0R)
DN_SET(T_IQ4_NL)
metal/kernels/moe_mm.metal
+20 -1
View File
@@ -365,6 +365,9 @@ kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS
#define MMDN2(F, M, N, K, S) MMDN2_(F, M, N, K, S)
#define MM_GU2_SET(F) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 32, 32, 32, 2) MMGU2(F, 32, 64, 16, 2) MMGU2(F, 64, 32, 16, 2) MMGU2(F, 16, 32, 32, 2) MMGU2(F, 16, 64, 16, 2) MMGU2(F, 16, 128, 32, 4) MMGU2(F, 32, 64, 16, 4) MMGU2(F, 16, 32, 32, 1) MMGU2(F, 16, 64, 32, 4) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 64, 4) MMGU2(F, 64, 64, 64, 4) MMGU2(F, 16, 32, 64, 2) MMGU2(F, 32, 64, 32, 4) MMGU2(F, 64, 64, 32, 4) MMGU2(F, 16, 64, 32, 2)
#define MM_DN2_SET(F) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 32, 32, 32, 2) MMDN2(F, 32, 64, 16, 2) MMDN2(F, 64, 32, 16, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 16, 64, 16, 2) MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 16, 4) MMDN2(F, 16, 32, 32, 1) MMDN2(F, 16, 64, 32, 4) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 32, 64, 64, 4) MMDN2(F, 64, 64, 64, 4) MMDN2(F, 16, 32, 64, 2) MMDN2(F, 32, 64, 32, 4) MMDN2(F, 64, 64, 32, 4) MMDN2(F, 16, 64, 32, 2)
// 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)
MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_XS) MM_GU2_SET(T_IQ2_S) MM_GU2_SET(T_IQ3_S)
MM_DN2_SET(T_Q2_0)
MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4)
@@ -457,6 +460,22 @@ template <> struct W4<T_IQ3_S> {
return w4_signed(tab[i], hh ? s >> 4 : s, sc);
}
};
template <> struct W4<T_IQ4_XS> { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 136;
const int j = k & 255, ib = j >> 5, o = j & 15, hi = (j >> 4) & 1;
const uint sh = *(device const ushort*)(b + 2);
const int ls = int((b[4 + (ib >> 1)] >> (4 * (ib & 1))) & 15) | int(((sh >> (2 * ib)) & 3) << 4);
const float s = rdh(b) * float(ls - 32);
device const uchar* qs = b + 8 + 16 * ib;
float4 q;
for (int c = 0; c < 4; ++c) {
const half2 w = as_type<half2>(tab[qs[o + c]]);
q[c] = float(hi ? w.y : w.x);
}
return half4(q * s);
}
};
template <> struct W4<T_IQ4_NL> { // 32 a block: d, 16 bytes; low nibbles are elements 0..15, high 16..31
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 5) * 18;
@@ -682,5 +701,5 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS
MMGU3(F, 16, 32, 16, 2) MMGU3(F, 32, 32, 32, 2) MMGU3(F, 16, 64, 32, 1) MMGU3(F, 16, 64, 32, 4) MMGU3(F, 16, 32, 16, 4) MMGU3(F, 16, 32, 32, 5) MMGU3(F, 16, 32, 16, 5) MMGU3(F, 16, 16, 32, 4) MMGU3(F, 16, 16, 32, 8) MMGU3(F, 16, 16, 16, 8)
#define MM3_DN(F) MMDN3(F, 16, 64, 32, 2) MMDN3(F, 16, 32, 32, 4) MMDN3(F, 32, 64, 32, 2) MMDN3(F, 16, 64, 64, 2) \
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_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_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)
metal/tests/expert_parity.cpp
+1 -1
View File
@@ -36,7 +36,7 @@ struct MoeArgs {
};
int type_of(const std::string& s) {
for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL})
for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS})
if (s == ref::type_name(t)) return t;
std::fprintf(stderr, "unknown type %s\n", s.c_str());
std::exit(2);
metal/tests/moe_mm_parity.cpp
+1 -1
View File
@@ -23,7 +23,7 @@ struct MmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; };
struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; };
double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; }
int type_of(const std::string& s) {
for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL})
for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS})
if (s == ref::type_name(t)) return t;
std::exit(2);
}
metal/tests/prompt_chain.cpp
+3 -1
View File
@@ -49,6 +49,7 @@ struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32
struct ArgArgs { uint32_t n_vocab, parts; };
double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; }
bool fmt_unit32(int t) { return t == ref::IQ4_XS || t == ref::IQ4_NL; } // mm_gu2 needs a K chunk >= 32
int unit_of(int t) { return t == ref::Q2_0 || t == ref::IQ3_S || t == ref::BF16 ? 16 : 32; }
} // namespace
@@ -433,7 +434,8 @@ int main(int argc, char** argv) {
push(g, "experts");
mtl::Dispatch gu;
gu.p = gu_v3 ? P(lib, "mm_gu3_" + std::to_string(eg.type) + "_" + gu3_for(eg.type), mtl::Constants().b(7, k4))
: P(lib, "mm_gu2_" + std::to_string(eg.type) + "_16_64_16_2", mtl::Constants().i(5, 0));
: P(lib, "mm_gu2_" + std::to_string(eg.type) + (fmt_unit32(eg.type) ? "_16_64_32_2" : "_16_64_16_2"),
mtl::Constants().i(5, 0));
gu.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
gu.tg[0] = gu_v3 ? 32 * gu3[3] : 64;
gu.indirect = &meta;