mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 14:21:06 -10:00
Metal: Q5_0 down experts, decode (moe.metal, moeu) and prompt (moe_mm v1/v2/v3), host tiles, honest gu_qk/d_qk; parity tests take Q5_0
This commit is contained in:
Niko1221 committed
1 parent
2b8179f3e6
commit
cbef1a9ff6
7 files changed
+46
-12
No files matched your search
@@ -338,6 +338,7 @@ GU_SET(T_IQ3_S)
|
||||
GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer
|
||||
GU_SET(T_Q4_K) GU_SET(T_Q5_K) GU_SET(T_Q6_K) GU_SET(T_Q8_0) // the UD-Q4/5/6_K_XL packs
|
||||
DN_SET(T_Q5_1) DN_SET(T_Q8_0)
|
||||
DN_SET(T_Q5_0) // the UD-IQ2_XS / Q2_K_XL / Q3_K_XL / IQ4_XS packs' down experts (generic Dec: d * (code - 16))
|
||||
DN_SET(T_Q2_0)
|
||||
DN_SET(T_Q2_0R)
|
||||
DN_SET(T_IQ4_NL)
|
||||
|
||||
@@ -257,7 +257,7 @@ kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* p
|
||||
#define MM_GU_SET(F) MMGU(F, 32, 64, 64, 4) MMGU(F, 64, 64, 64, 4) MMGU(F, 16, 64, 64, 2) MMGU(F, 32, 32, 64, 2)
|
||||
#define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2)
|
||||
MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S)
|
||||
MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL)
|
||||
MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) MM_DN_SET(T_Q5_0)
|
||||
|
||||
// ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory ---------------------
|
||||
// mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads.
|
||||
@@ -388,7 +388,7 @@ kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS
|
||||
MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 64, 64, 32, 4)
|
||||
// K-quants: the staged K chunk must hold whole units as DW places them (128 wide): TK 128
|
||||
#define MM_GU2_KQ(F) MMGU2(F, 16, 32, 128, 2) MMGU2(F, 16, 16, 128, 2) MMGU2(F, 32, 32, 128, 2) MMGU2(F, 16, 32, 128, 4)
|
||||
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_Q8_0)
|
||||
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)
|
||||
// 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)
|
||||
@@ -550,6 +550,21 @@ template <> struct W4<T_Q5_1> { // 32 a block: d, m, qh (32 bits), qs[16] (low
|
||||
return half4(w);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_Q5_0> { // 32 a block: d, qh (32 bits), qs[16] (low nibbles 0..15, high 16..31); w = (code - 16) d
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
|
||||
device const uchar* b = row + (k >> 5) * 22;
|
||||
const float d = rdh(b);
|
||||
const uint hb = rd32(b + 2);
|
||||
const int j = k & 31, o = j & 15, hi = j >> 4;
|
||||
float4 w;
|
||||
for (int c = 0; c < 4; ++c) {
|
||||
const uint qv = b[6 + o + c];
|
||||
const int x = int(hi ? (qv >> 4) : (qv & 15)) | int(((hb >> (j + c)) & 1) << 4);
|
||||
w[c] = float(x - 16) * d;
|
||||
}
|
||||
return half4(w);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_Q4_K> { // 256 a block: d, dmin, scales[12], qs[128] (64-chunks: low nibbles then high)
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
|
||||
device const uchar* b = row + (k >> 8) * 144;
|
||||
@@ -826,6 +841,24 @@ template <> struct W8<T_Q5_1> {
|
||||
b = half4(wb);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_Q5_0> { // 22-byte blocks (2-aligned): the code bytes as ushorts; one d and one qh a pair
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 5) * 22;
|
||||
const float d = rdh(bl);
|
||||
const uint hb = rd32(bl + 2);
|
||||
const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones
|
||||
device const ushort* p = (device const ushort*)(bl + 6 + j);
|
||||
const uint v = uint(p[0]) | (uint(p[1]) << 16);
|
||||
float4 wa, wb;
|
||||
for (int c = 0; c < 4; ++c) {
|
||||
const uint qv = (v >> (8 * c)) & 0xFFu;
|
||||
wa[c] = float((int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) - 16) * d;
|
||||
wb[c] = float((int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) - 16) * d;
|
||||
}
|
||||
a = half4(wa);
|
||||
b = half4(wb);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_Q6_K> {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 8) * 210;
|
||||
@@ -992,7 +1025,7 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS
|
||||
#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_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_Q8_0)
|
||||
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)
|
||||
|
||||
// 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
|
||||
|
||||
@@ -151,5 +151,5 @@ kernel void moe_down_u(device const uchar* W [[buffer(0)]], device const float*
|
||||
// the tiles the Q2_0 and IQ3_S packs use (window.cpp / token_chain.cpp pick 32,1,4 for gate/up; down 8,1,4 for Q2_0 and 4,2,4 otherwise)
|
||||
#define GUU_ALL(U) GUU(T_Q2_0, 32, 1, 4, U) GUU(T_IQ3_XXS, 32, 1, 4, U) GUU(T_IQ2_S, 32, 1, 4, U) GUU(T_IQ3_S, 32, 1, 4, U) \
|
||||
GUU(T_IQ4_XS, 32, 1, 4, U)
|
||||
#define DNU_ALL(U) DNU(T_Q2_0, 8, 1, 4, U) DNU(T_Q2_0R, 8, 1, 4, U) DNU(T_IQ4_NL, 4, 2, 4, U) DNU(T_Q8_0, 4, 2, 4, U) DNU(T_Q5_1, 4, 2, 4, U)
|
||||
#define DNU_ALL(U) DNU(T_Q2_0, 8, 1, 4, U) DNU(T_Q2_0R, 8, 1, 4, U) DNU(T_IQ4_NL, 4, 2, 4, U) DNU(T_Q8_0, 4, 2, 4, U) DNU(T_Q5_1, 4, 2, 4, U) DNU(T_Q5_0, 4, 2, 4, U)
|
||||
GUU_ALL(2) DNU_ALL(2)
|
||||
@@ -24,7 +24,7 @@ int gu_qk(int t) {
|
||||
case 16: case 17: case 18: case 21: case 22: case 23: case 29: case 12: case 13: case 14: return 256;
|
||||
case 42: return 64;
|
||||
case 144: case 145: return 128; // TQ2_T / TQK6 gate/up (the CUDA engine's STRATA_GU_FMTS)
|
||||
case 6: case 2: case 3: case 8: return 32;
|
||||
case 8: return 32; // Q8_0 (the UD-Q5/Q6_K_XL packs); Q5_0 / Q4_0 / Q4_1 gate/up have no Metal kernel (moe.metal GU_SET)
|
||||
default: return 0;
|
||||
}
|
||||
}
|
||||
@@ -33,7 +33,7 @@ int d_qk(int t) {
|
||||
case 23: return 256;
|
||||
case 42: return 64;
|
||||
case 144: case 146: return 128; // TQ2_T / TQK7 down
|
||||
case 20: case 7: case 6: case 2: case 3: case 8: return 32;
|
||||
case 20: case 7: case 6: case 8: return 32; // IQ4_NL, Q5_1, Q5_0 (moe.metal DN_SET since 0.2.11 UD sets), Q8_0; Q4_0 / Q4_1 down have no Metal kernel
|
||||
default: return 0;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1194,7 +1194,7 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
|
||||
// 24-36 GB/s, mac-kern 513fc615)
|
||||
const bool m4 = !m5_layout();
|
||||
const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4,
|
||||
dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2;
|
||||
dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1 || td.type == ref::Q5_0)) ? 1 : 2;
|
||||
const bool m16 = f16_enc_ && !stream_ && tg_.type == ref::Q2_0 && dq2 && !had_gu(l) && !had_dn(l);
|
||||
a.gx_gu = NFF / (gs * (32 / gl) * gr);
|
||||
a.gx_dn = N / (ds_ * (32 / dl) * dr);
|
||||
@@ -2470,7 +2470,7 @@ void Window::mtp_encode(std::vector<mtl::Dispatch>& ds, const mtl::Buffer& srcR,
|
||||
// 24-36 GB/s, mac-kern 513fc615)
|
||||
const bool m4 = !m5_layout();
|
||||
const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4,
|
||||
dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2;
|
||||
dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1 || td.type == ref::Q5_0)) ? 1 : 2;
|
||||
a.gx_gu = NFF / (gs * (32 / gl) * gr);
|
||||
a.gx_dn = N / (ds_ * (32 / dl) * dr);
|
||||
mtl::Dispatch r;
|
||||
@@ -2670,7 +2670,7 @@ std::string gu2_tile(int t) {
|
||||
if (t == ref::Q8_0 || t == ref::IQ4_XS) return "_16_64_32_2";
|
||||
return "_16_64_16_2";
|
||||
}
|
||||
std::string dn2_tile(int t) { return t == ref::IQ4_NL || t == ref::Q5_1 || t == ref::Q8_0 ? "_16_64_32_2" : "_16_64_16_2"; }
|
||||
std::string dn2_tile(int t) { return t == ref::IQ4_NL || t == ref::Q5_1 || t == ref::Q5_0 || t == ref::Q8_0 ? "_16_64_32_2" : "_16_64_16_2"; }
|
||||
} // namespace
|
||||
|
||||
struct Window::PromptGemm {
|
||||
|
||||
@@ -50,7 +50,7 @@ struct MoeArgs {
|
||||
};
|
||||
|
||||
int type_of(const std::string& s) {
|
||||
for (int t : {ref::TQ2_T, ref::TQK6, ref::TQK7, ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1})
|
||||
for (int t : {ref::TQ2_T, ref::TQK6, ref::TQK7, ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1, ref::Q5_0})
|
||||
if (s == ref::type_name(t)) return t;
|
||||
std::fprintf(stderr, "unknown type %s\n", s.c_str());
|
||||
std::exit(2);
|
||||
|
||||
@@ -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::TQ2_T, ref::TQK6, ref::TQK7, ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1})
|
||||
for (int t : {ref::TQ2_T, ref::TQK6, ref::TQK7, ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1, ref::Q5_0})
|
||||
if (s == ref::type_name(t)) return t;
|
||||
std::exit(2);
|
||||
}
|
||||
@@ -90,7 +90,7 @@ int main(int argc, char** argv) {
|
||||
mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"),
|
||||
pclear = cx.pipeline(lib, "mm_clear"), pcomb = cx.pipeline(lib, "mm_combine");
|
||||
if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; }
|
||||
if (dn == ref::IQ4_NL && DTK < 32) DTK = 32;
|
||||
if ((dn == ref::IQ4_NL || dn == ref::Q5_0) && DTK < 32) DTK = 32; // 32-value units (the v2 / v3 down tiles of IQ4_NL / Q5_0 need TK >= 32)
|
||||
const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG);
|
||||
const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG);
|
||||
mtl::Constants probe;
|
||||
|
||||
Reference in new issue
Block a user