mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 22:59:00 -10:00
Merge f-mac-kern 6f6c555b (prompt experts / QSA on matrix-unit operands, COOP_K4)
This commit is contained in:
Niko1221 committed
commit
68387f1a08
8 files changed
+908
-14
No files matched your search
@@ -52,3 +52,26 @@ kernel void sg_mm(device const T* A [[buffer(0)]], device const T* B [[buffer(1)
|
||||
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>;
|
||||
|
||||
|
||||
// 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>
|
||||
kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_index_in_simdgroup]]) {
|
||||
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
matmul2d<desc, execution_simdgroup> op;
|
||||
auto r = op.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto l = op.template get_left_input_cooperative_tensor<half, half, float>();
|
||||
for (uint16_t i = 0; i < r.get_capacity() && i < 128; ++i) {
|
||||
auto ix = r.get_multidimensional_index(i);
|
||||
device int* o = out + ((0 * 32 + lane) * 128 + i) * 4;
|
||||
o[0] = ix[0]; o[1] = ix[1]; o[2] = r.get_capacity();
|
||||
}
|
||||
for (uint16_t i = 0; i < l.get_capacity() && i < 128; ++i) {
|
||||
auto ix = l.get_multidimensional_index(i);
|
||||
device int* o = out + ((1 * 32 + lane) * 128 + i) * 4;
|
||||
o[0] = ix[0]; o[1] = ix[1]; o[2] = l.get_capacity();
|
||||
}
|
||||
}
|
||||
template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 32, 16>) coop_layout<16, 32, 16>;
|
||||
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>;
|
||||
@@ -368,3 +368,314 @@ kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS
|
||||
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)
|
||||
|
||||
// ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope),
|
||||
// no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN.
|
||||
// The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values
|
||||
// at a time (W4<F>). Activations: the gathered fp16 rows from device memory.
|
||||
#ifndef STRATA_COOP_K4
|
||||
#define STRATA_COOP_K4
|
||||
// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's
|
||||
// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own.
|
||||
constant bool COOP_K4 [[function_constant(7)]];
|
||||
constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true;
|
||||
#endif
|
||||
template <int F> struct W4;
|
||||
// the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16
|
||||
static inline half4 w4_signed(uint g, uint s, float sc) {
|
||||
const uchar4 v = as_type<uchar4>(g);
|
||||
const float4 q = float4((s & 1) ? -float(v[0]) : float(v[0]), (s & 2) ? -float(v[1]) : float(v[1]),
|
||||
(s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3]));
|
||||
return half4(q * sc);
|
||||
}
|
||||
template <> struct W4<T_Q2_0> { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4)
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
|
||||
device const uchar* b = row + (k >> 6) * 18;
|
||||
const float d = float(*(device const half*)b);
|
||||
const uint q = b[2 + ((k & 63) >> 2)];
|
||||
return half4(float4(float(int(q & 3) - 1), float(int((q >> 2) & 3) - 1), float(int((q >> 4) & 3) - 1),
|
||||
float(int(q >> 6) - 1)) * d);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_IQ3_XXS> {
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
|
||||
device const uchar* b = row + (k >> 8) * 98;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint aux = rd32(b + 2 + 64 + 4 * ib);
|
||||
threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256);
|
||||
const uint s = ks[(aux >> (7 * l)) & 127];
|
||||
const float sc = rdh(b) * (0.5f + float(aux >> 28)) * 0.5f;
|
||||
return w4_signed(tab[b[2 + 8 * ib + 2 * l + hh]], hh ? s >> 4 : s, sc);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_IQ2_XXS> {
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
|
||||
device const uchar* b = row + (k >> 8) * 66;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4);
|
||||
threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512);
|
||||
const uint idx = (a0 >> (8 * l)) & 255, s = ks[(a1 >> (7 * l)) & 127];
|
||||
const float sc = rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f;
|
||||
return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_IQ2_XS> {
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
|
||||
device const uchar* b = row + (k >> 8) * 74;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint v = ((device const ushort*)(b + 2))[4 * ib + l];
|
||||
threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024);
|
||||
const uint idx = v & 511, s = ks[v >> 9];
|
||||
const uint scb = b[66 + ib];
|
||||
const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f;
|
||||
return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_IQ2_S> {
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
|
||||
device const uchar* b = row + (k >> 8) * 82;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint qh = b[66 + ib];
|
||||
const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300);
|
||||
const uint s = b[2 + 32 + 4 * ib + l];
|
||||
const uint scb = b[74 + ib];
|
||||
const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f;
|
||||
return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc);
|
||||
}
|
||||
};
|
||||
template <> struct W4<T_IQ3_S> {
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
|
||||
device const uchar* b = row + (k >> 8) * 110;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint qh = b[66 + ib];
|
||||
const uint i = hh ? (uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256))
|
||||
: (uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256));
|
||||
const uint s = b[74 + 4 * ib + l];
|
||||
const uint scb = b[106 + (ib >> 1)];
|
||||
const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15));
|
||||
return w4_signed(tab[i], hh ? s >> 4 : s, sc);
|
||||
}
|
||||
};
|
||||
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;
|
||||
const int j = k & 31, hi = j >> 4, o = j & 15;
|
||||
const float d = rdh(b);
|
||||
float4 q;
|
||||
for (int c = 0; c < 4; ++c) {
|
||||
const half2 w = as_type<half2>(tab[b[2 + o + c]]);
|
||||
q[c] = float(hi ? w.y : w.x);
|
||||
}
|
||||
return half4(q * d);
|
||||
}
|
||||
};
|
||||
|
||||
// W8<F>: the runs k .. k + 3 and k + 16 .. k + 19 of a row (k % 32 < 16): one 32-group, so the per-group fields (scale,
|
||||
// sign word) are read once. The generic one is two W4.
|
||||
template <int F> struct W8 {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
|
||||
a = W4<F>::run(row, k, tab);
|
||||
b = W4<F>::run(row, k + 16, tab);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_IQ3_XXS> {
|
||||
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) * 98;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint aux = rd32(bl + 2 + 64 + 4 * ib);
|
||||
threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256);
|
||||
const float sc = rdh(bl) * (0.5f + float(aux >> 28)) * 0.5f;
|
||||
const uint s0 = ks[(aux >> (7 * l)) & 127], s1 = ks[(aux >> (7 * (l + 2))) & 127];
|
||||
device const uchar* qs = bl + 2 + 8 * ib + hh;
|
||||
a = w4_signed(tab[qs[2 * l]], hh ? s0 >> 4 : s0, sc);
|
||||
b = w4_signed(tab[qs[2 * l + 4]], hh ? s1 >> 4 : s1, sc);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_IQ2_XXS> {
|
||||
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) * 66;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint a0 = rd32(bl + 2 + 8 * ib), a1 = rd32(bl + 2 + 8 * ib + 4);
|
||||
threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512);
|
||||
const float sc = rdh(bl) * (0.5f + float(a1 >> 28)) * 0.25f;
|
||||
const uint i0 = (a0 >> (8 * l)) & 255, i1 = (a0 >> (8 * (l + 2))) & 255;
|
||||
const uint s0 = ks[(a1 >> (7 * l)) & 127], s1 = ks[(a1 >> (7 * (l + 2))) & 127];
|
||||
a = w4_signed(tab[2 * i0 + hh], hh ? s0 >> 4 : s0, sc);
|
||||
b = w4_signed(tab[2 * i1 + hh], hh ? s1 >> 4 : s1, sc);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_IQ2_XS> {
|
||||
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) * 74;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
device const ushort* q = (device const ushort*)(bl + 2) + 4 * ib;
|
||||
const uint v0 = q[l], v1 = q[l + 2];
|
||||
threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024);
|
||||
const uint scb = bl[66 + ib];
|
||||
const float d = rdh(bl);
|
||||
const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f;
|
||||
const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f;
|
||||
const uint s0 = ks[v0 >> 9], s1 = ks[v1 >> 9];
|
||||
a = w4_signed(tab[2 * (v0 & 511) + hh], hh ? s0 >> 4 : s0, sc0);
|
||||
b = w4_signed(tab[2 * (v1 & 511) + hh], hh ? s1 >> 4 : s1, sc1);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_IQ2_S> {
|
||||
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) * 82;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint qh = bl[66 + ib], scb = bl[74 + ib];
|
||||
const float d = rdh(bl);
|
||||
const uint i0 = uint(bl[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300);
|
||||
const uint i1 = uint(bl[2 + 4 * ib + l + 2]) | ((qh << (8 - 2 * (l + 2))) & 0x300);
|
||||
const uint s0 = bl[2 + 32 + 4 * ib + l], s1 = bl[2 + 32 + 4 * ib + l + 2];
|
||||
const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f;
|
||||
const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f;
|
||||
a = w4_signed(tab[2 * i0 + hh], hh ? s0 >> 4 : s0, sc0);
|
||||
b = w4_signed(tab[2 * i1 + hh], hh ? s1 >> 4 : s1, sc1);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_IQ3_S> {
|
||||
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) * 110;
|
||||
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
|
||||
const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)];
|
||||
const float sc = rdh(bl) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15));
|
||||
const int l1 = l + 2;
|
||||
const uint i0 = hh ? (uint(bl[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256))
|
||||
: (uint(bl[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256));
|
||||
const uint i1 = hh ? (uint(bl[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256))
|
||||
: (uint(bl[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256));
|
||||
const uint s0 = bl[74 + 4 * ib + l], s1 = bl[74 + 4 * ib + l1];
|
||||
a = w4_signed(tab[i0], hh ? s0 >> 4 : s0, sc);
|
||||
b = w4_signed(tab[i1], hh ? s1 >> 4 : s1, sc);
|
||||
}
|
||||
};
|
||||
|
||||
template <int F, int TM, int TN, int TK, int SG>
|
||||
kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]],
|
||||
device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) {
|
||||
threadgroup float ot[SG][TM * TN];
|
||||
threadgroup uint tab[fmt<F>::TAB];
|
||||
load_tab<F>(tab, tid, SG * 32);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const int2 tl = tiles[tg.y];
|
||||
const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN;
|
||||
device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row;
|
||||
device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row;
|
||||
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
matmul2d<desc, execution_simdgroup> op;
|
||||
tensor<device half, dextents<int32_t, 2>, tensor_inline> tX((device half*)XS, dextents<int32_t, 2>(a.n_embd, 1 << 30));
|
||||
auto rg = op.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto ru = op.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto tA = tX.template slice<TK, TM>(0, row0);
|
||||
auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rg)>, float>();
|
||||
auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rg)>, float>();
|
||||
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
|
||||
if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; }
|
||||
for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) {
|
||||
const uint16_t cap = rg.get_capacity(), half_cap = cap / 2;
|
||||
if (!coop_k4) {
|
||||
for (uint16_t i = 0; i < cap; ++i) {
|
||||
const auto ix = rg.get_multidimensional_index(i);
|
||||
const int k = k0 + ix[0], n = ix[1];
|
||||
rg[i] = W4<F>::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3];
|
||||
ru[i] = W4<F>::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3];
|
||||
}
|
||||
} else
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < half_cap; i += 4) {
|
||||
const auto ix = rg.get_multidimensional_index(i);
|
||||
const auto iy = rg.get_multidimensional_index(i + half_cap);
|
||||
const int k = k0 + ix[0], n = ix[1];
|
||||
device const uchar* rgw = g0 + ulong(n) * a.gu_row;
|
||||
device const uchar* ruw = u0 + ulong(n) * a.gu_row;
|
||||
half4 wg, wu, wg2, wu2;
|
||||
if (TK == 32 && iy[0] == ix[0] + 16 && iy[1] == ix[1] && (ix[0] & 31) < 16) {
|
||||
// the lane's two runs of this row sit in one 32-group (the layout: element i + cap / 2 is k + 16)
|
||||
W8<F>::run(rgw, k, tab, wg, wg2);
|
||||
W8<F>::run(ruw, k, tab, wu, wu2);
|
||||
} else {
|
||||
const int k2 = k0 + iy[0];
|
||||
device const uchar* rgw2 = g0 + ulong(iy[1]) * a.gu_row;
|
||||
device const uchar* ruw2 = u0 + ulong(iy[1]) * a.gu_row;
|
||||
wg = W4<F>::run(rgw, k, tab); wu = W4<F>::run(ruw, k, tab);
|
||||
wg2 = W4<F>::run(rgw2, k2, tab); wu2 = W4<F>::run(ruw2, k2, tab);
|
||||
}
|
||||
rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w;
|
||||
ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w;
|
||||
const uint16_t j = i + half_cap;
|
||||
rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w;
|
||||
ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w;
|
||||
}
|
||||
tA = tX.template slice<TK, TM>(k0, row0);
|
||||
op.run(tA, rg, accg);
|
||||
op.run(tA, ru, accu);
|
||||
}
|
||||
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
|
||||
if (accg.is_valid_element(i)) {
|
||||
const float g = accg[i];
|
||||
accg[i] = g / (1.0f + exp(-g)) * accu[i];
|
||||
}
|
||||
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot[sg], dextents<int32_t, 2>(TN, TM));
|
||||
accg.store(tO);
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int it = lane; it < TM * TN; it += 32) {
|
||||
const int m = it / TN, f = it - m * TN;
|
||||
HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int F, int TM, int TN, int TK, int SG>
|
||||
kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]],
|
||||
device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort tid [[thread_index_in_threadgroup]]) {
|
||||
threadgroup uint tab[fmt<F>::TAB];
|
||||
load_tab<F>(tab, tid, SG * 32);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const int2 tl = tiles[tg.y];
|
||||
const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN;
|
||||
device const uchar* w0 = W + a.down_off + ulong(e) * a.d_stride + ulong(n0) * a.d_row;
|
||||
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
matmul2d<desc, execution_simdgroup> op;
|
||||
tensor<device half, dextents<int32_t, 2>, tensor_inline> tH((device half*)HS, dextents<int32_t, 2>(a.n_ff, 1 << 30));
|
||||
auto rw = op.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto tA = tH.template slice<TK, TM>(0, row0);
|
||||
auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rw)>, float>();
|
||||
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
|
||||
if (acc.is_valid_element(i)) acc[i] = 0.0f;
|
||||
for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) {
|
||||
if (!coop_k4) {
|
||||
for (uint16_t i = 0; i < rw.get_capacity(); ++i) {
|
||||
const auto ix = rw.get_multidimensional_index(i);
|
||||
const int k = k0 + ix[0];
|
||||
rw[i] = W4<F>::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3];
|
||||
}
|
||||
} else
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < rw.get_capacity(); i += 4) {
|
||||
const auto ix = rw.get_multidimensional_index(i);
|
||||
const half4 w = W4<F>::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab);
|
||||
rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w;
|
||||
}
|
||||
tA = tH.template slice<TK, TM>(k0, row0);
|
||||
op.run(tA, rw, acc);
|
||||
}
|
||||
tensor<device float, dextents<int32_t, 2>, tensor_inline> tY(YE, dextents<int32_t, 2>(a.n_embd, 1 << 30));
|
||||
acc.store(tY.slice(n0, row0));
|
||||
}
|
||||
|
||||
#define MMGU3_(F, M, N, K, S) template [[host_name("mm_gu3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu3<F, M, N, K, S>) mm_gu3<F, M, N, K, S>;
|
||||
#define MMDN3_(F, M, N, K, S) template [[host_name("mm_down3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down3<F, M, N, K, S>) mm_down3<F, M, N, K, S>;
|
||||
#define MMGU3(F, M, N, K, S) MMGU3_(F, M, N, K, S)
|
||||
#define MMDN3(F, M, N, K, S) MMDN3_(F, M, N, K, S)
|
||||
#define MM3_ONE(F, M, N, K, S) MMGU3(F, M, N, K, S) MMDN3(F, M, N, K, S)
|
||||
#define MM3_GU(F) MMGU3(F, 16, 64, 32, 2) MMGU3(F, 16, 32, 32, 4) MMGU3(F, 32, 64, 32, 2) MMGU3(F, 16, 64, 64, 2) \
|
||||
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_DN(T_Q2_0) MM3_DN(T_IQ4_NL)
|
||||
@@ -0,0 +1,305 @@
|
||||
// metal/kernels/qsa_prompt.metal - QSA attention for a PROMPT chunk on the matrix units (after qsa.metal: QH, QKV, QD,
|
||||
// QBUDGET, QCH, QsaArgs). The same split-K partials as qsa_attn (part [t][chunk][QH][QD + 2]: acc, then m, l), so
|
||||
// qsa_merge is unchanged; per (chunk of QCH = 64 selected cells, kv head g, token t) one simdgroup:
|
||||
// S[12 x 64] = Qh[12 x 256] . K^T / 16 single-simdgroup matmul2d, K (int8 * fp16 scale per 64) decoded into
|
||||
// the right-input cooperative tensor, 32 dims a step
|
||||
// per head: m = max, p = exp(s - m), l = sum p (fp32, through threadgroup memory), P -> fp16
|
||||
// acc[12 x 256] = P[12 x 64] . V[64 x 256] four 64-dim matmuls, V decoded into the right input
|
||||
// The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32.
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
using namespace mpp::tensor_ops;
|
||||
#ifndef STRATA_COOP_K4
|
||||
#define STRATA_COOP_K4
|
||||
// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's
|
||||
// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own.
|
||||
constant bool COOP_K4 [[function_constant(7)]];
|
||||
constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true;
|
||||
#endif
|
||||
|
||||
// grid (chunks, QKV, n_tok), 32 threads. qh: the q heads as fp16 [t][QH][QD] (qcur rounded).
|
||||
kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]],
|
||||
device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]],
|
||||
device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]],
|
||||
device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) {
|
||||
constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH;
|
||||
threadgroup float s[16 * QCH];
|
||||
threadgroup half p[16 * QCH];
|
||||
threadgroup int cid[QCH];
|
||||
threadgroup half vs[QCH][4];
|
||||
const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z);
|
||||
const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET));
|
||||
const int c0 = chunk * QCH;
|
||||
device float* pp = part + ((uint(t) * NCH + chunk) * QH + g * HPG) * (QD + 2);
|
||||
if (uint(c0) >= width) {
|
||||
if (lane < HPG) { pp[lane * (QD + 2) + QD] = -INFINITY; pp[lane * (QD + 2) + QD + 1] = 0.0f; }
|
||||
return;
|
||||
}
|
||||
const int nc = min(QCH, int(width) - c0);
|
||||
for (int c = lane; c < QCH; c += 32) {
|
||||
const int id = c < nc ? ids[t * QBUDGET + c0 + c] : -1;
|
||||
cid[c] = id;
|
||||
for (int j = 0; j < 4; ++j) vs[c][j] = id >= 0 ? vsc[(uint(id) * QKV + g) * 4 + j] : half(0.0f);
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// ---- S = Q K^T
|
||||
{
|
||||
constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
matmul2d<desc, execution_simdgroup> op;
|
||||
tensor<device half, dextents<int32_t, 2>, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD,
|
||||
dextents<int32_t, 2>(QD, HPG));
|
||||
auto rk = op.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto tA = tQ.template slice<32, 16>(0, 0);
|
||||
auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rk)>, float>();
|
||||
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
|
||||
if (acc.is_valid_element(i)) acc[i] = 0.0f;
|
||||
for (int k0 = 0; k0 < QD; k0 += 32) {
|
||||
if (!coop_k4) {
|
||||
for (uint16_t i = 0; i < rk.get_capacity(); ++i) {
|
||||
const auto ix = rk.get_multidimensional_index(i);
|
||||
const int k = k0 + ix[0], id = cid[ix[1]];
|
||||
const uint row = uint(id) * QKV + g;
|
||||
rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f);
|
||||
}
|
||||
} else
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < rk.get_capacity(); i += 4) {
|
||||
const auto ix = rk.get_multidimensional_index(i);
|
||||
const int k = k0 + ix[0], c = ix[1], id = cid[c];
|
||||
half4 w = half4(0.0h);
|
||||
if (id >= 0) {
|
||||
const uint row = uint(id) * QKV + g;
|
||||
const char4 q4 = *(device const char4*)(kq + row * QD + k);
|
||||
w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)]));
|
||||
}
|
||||
rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w;
|
||||
}
|
||||
tA = tQ.template slice<32, 16>(k0, 0);
|
||||
op.run(tA, rk, acc);
|
||||
}
|
||||
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tS(s, dextents<int32_t, 2>(QCH, 16));
|
||||
acc.store(tS);
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// ---- softmax pieces per head (lane = head for the reductions; 12 heads)
|
||||
float mh = -INFINITY, lh = 0.0f;
|
||||
if (lane < HPG) {
|
||||
for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f);
|
||||
for (int c = 0; c < QCH; ++c) {
|
||||
const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f;
|
||||
lh += e;
|
||||
p[lane * QCH + c] = half(e);
|
||||
}
|
||||
} else if (lane < 16) {
|
||||
for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f);
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// ---- acc = P V, 64 dims at a time
|
||||
{
|
||||
constexpr auto desc = matmul2d_descriptor(16, 64, QCH, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
matmul2d<desc, execution_simdgroup> op;
|
||||
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tP(p, dextents<int32_t, 2>(QCH, 16));
|
||||
tensor<device float, dextents<int32_t, 2>, tensor_inline> tO(pp, dextents<int32_t, 2>(QD + 2, HPG));
|
||||
for (int d0 = 0; d0 < QD; d0 += 64) {
|
||||
auto rv = op.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tP)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rv)>, float>();
|
||||
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
|
||||
if (acc.is_valid_element(i)) acc[i] = 0.0f;
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < rv.get_capacity(); ++i) {
|
||||
const auto ix = rv.get_multidimensional_index(i); // (cell, dim)
|
||||
const int c = ix[0], d = d0 + ix[1], id = cid[c];
|
||||
rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f);
|
||||
}
|
||||
op.run(tP, rv, acc);
|
||||
auto mO = tO.slice(d0, 0);
|
||||
acc.store(mO);
|
||||
}
|
||||
}
|
||||
if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; }
|
||||
}
|
||||
|
||||
// qcur (fp32) -> fp16 for qsa_attn_mm
|
||||
kernel void qsa_q_f16(device const float* q [[buffer(0)]], device half* qh [[buffer(1)]], constant uint& n [[buffer(2)]],
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid < n) qh[gid] = half(q[gid]);
|
||||
}
|
||||
|
||||
// qsa_score as a GEMM (fp32 operands, so the selection sees qsa_score's values to fp32 rounding): S[(t, h)][b] =
|
||||
// qidx[t][h] . pooled[b] for 16 tokens x IH heads (64 rows) by 64 blocks a threadgroup; the epilogue sums relu over
|
||||
// the heads. The tail block n_bid (the `dead` key, + 1e9 when it has cells) and blocks past it are left to the
|
||||
// scalar path: the tail is scored here by one simdgroup a token; blocks > n_bid are not written (qsa_topk reads
|
||||
// b <= n_bid). grid (ceil(max blocks / 64), ceil(T / 16)), 128 threads.
|
||||
kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]],
|
||||
device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]],
|
||||
constant QsaArgs& a [[buffer(4)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
|
||||
threadgroup float ot[64 * 64];
|
||||
const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16;
|
||||
const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score
|
||||
tensor<device float, dextents<int32_t, 2>, tensor_inline> tQ((device float*)qidx, dextents<int32_t, 2>(ID, int(a.n_tok) * IH));
|
||||
tensor<device float, dextents<int32_t, 2>, tensor_inline> tK((device float*)pooled, dextents<int32_t, 2>(ID, int(nb_all)));
|
||||
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot, dextents<int32_t, 2>(64, 64));
|
||||
constexpr auto desc = matmul2d_descriptor(64, 64, dynamic_length_v<int>, false, true, false);
|
||||
matmul2d<desc, execution_simdgroups<4>> op;
|
||||
auto mA = tQ.slice(0, t0 * IH);
|
||||
auto mB = tK.slice(0, b0);
|
||||
op.run(mA, mB, tO);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint stride = a.max_cells / 4 + 2;
|
||||
for (int idx = tid; idx < 16 * 64; idx += 128) {
|
||||
const int tt = idx >> 6, j = idx & 63, t = t0 + tt, b = b0 + j;
|
||||
if (t >= int(a.n_tok)) continue;
|
||||
const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2;
|
||||
if (uint(b) >= n_bid) continue;
|
||||
float s = 0.0f;
|
||||
for (int h = 0; h < IH; ++h) s += max(ot[(tt * IH + h) * 64 + j], 0.0f);
|
||||
scores[uint(t) * stride + b] = s;
|
||||
}
|
||||
// the tail block of each token in this threadgroup's block range (the first column threadgroup does it)
|
||||
if (tg.x == 0)
|
||||
for (int tt = sg; tt < 16; tt += 4) {
|
||||
const int t = t0 + tt;
|
||||
if (t >= int(a.n_tok)) break;
|
||||
const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2;
|
||||
float k[4];
|
||||
for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i];
|
||||
float s = 0.0f;
|
||||
for (int h = 0; h < IH; ++h) {
|
||||
float d = 0.0f;
|
||||
for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d);
|
||||
d = simd_sum(d);
|
||||
s += max(d, 0.0f);
|
||||
}
|
||||
if ((n_kv & 3) != 0) s += 1e9f;
|
||||
if (lane == 0) scores[uint(t) * stride + n_bid] = s;
|
||||
}
|
||||
}
|
||||
|
||||
// qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the
|
||||
// width selected cells in chunks of 64 with an online softmax; 4 simdgroups, simdgroup s owns dims 64 s .. 64 s + 63
|
||||
// (its slice of S's K sum and its 64 output dims). Out: xq = (acc / l) * sigmoid(gate) as fp16 (attn_output's GEMM
|
||||
// operand). grid (QKV, n_tok), 128 threads.
|
||||
kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]],
|
||||
device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]],
|
||||
device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]],
|
||||
device const float* qfull [[buffer(6)]], device half* xqh [[buffer(7)]], constant QsaArgs& a [[buffer(8)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
|
||||
constexpr int HPG = QH / QKV;
|
||||
threadgroup float sp[4][16 * QCH];
|
||||
threadgroup half p[16 * QCH];
|
||||
threadgroup int cid[QCH];
|
||||
threadgroup float vs[QCH];
|
||||
threadgroup float fac[16], lsum[16], mrun[16];
|
||||
const int g = int(tg.x), t = int(tg.y);
|
||||
const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET));
|
||||
const int nch = int((width + QCH - 1) / QCH);
|
||||
if (tid < 16) { mrun[tid] = -INFINITY; lsum[tid] = 0.0f; }
|
||||
constexpr auto dS = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
constexpr auto dO = matmul2d_descriptor(16, 64, QCH, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
|
||||
matmul2d<dS, execution_simdgroup> opS;
|
||||
matmul2d<dO, execution_simdgroup> opO;
|
||||
tensor<device half, dextents<int32_t, 2>, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD,
|
||||
dextents<int32_t, 2>(QD, HPG));
|
||||
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tP(p, dextents<int32_t, 2>(QCH, 16));
|
||||
auto rv = opO.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tP)>,
|
||||
__tensor_ops_detail::__remove_addrspace_t<decltype(rv)>, float>();
|
||||
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
|
||||
if (acc.is_valid_element(i)) acc[i] = 0.0f;
|
||||
const int d0 = sg * 64;
|
||||
for (int ch = 0; ch < nch; ++ch) {
|
||||
const int c0 = ch * QCH, nc = min(QCH, int(width) - c0);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid < QCH) {
|
||||
const int id = tid < nc ? ids[t * QBUDGET + c0 + tid] : -1;
|
||||
cid[tid] = id;
|
||||
vs[tid] = id >= 0 ? float(vsc[(uint(id) * QKV + g) * 4 + sg * 0]) : 0.0f; // (per-simdgroup scale below)
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// S partial over this simdgroup's 64 dims
|
||||
{
|
||||
auto rk = opS.template get_right_input_cooperative_tensor<half, half, float>();
|
||||
auto tA = tQ.template slice<32, 16>(d0, 0);
|
||||
auto s = opS.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>,
|
||||
__tensor_ops_detail::__remove_addrspace_t<decltype(rk)>, float>();
|
||||
for (uint16_t i = 0; i < s.get_capacity(); ++i)
|
||||
if (s.is_valid_element(i)) s[i] = 0.0f;
|
||||
for (int k0 = d0; k0 < d0 + 64; k0 += 32) {
|
||||
if (!coop_k4) {
|
||||
for (uint16_t i = 0; i < rk.get_capacity(); ++i) {
|
||||
const auto ix = rk.get_multidimensional_index(i);
|
||||
const int k = k0 + ix[0], id = cid[ix[1]];
|
||||
const uint row = uint(id) * QKV + g;
|
||||
rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f);
|
||||
}
|
||||
} else
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < rk.get_capacity(); i += 4) {
|
||||
const auto ix = rk.get_multidimensional_index(i);
|
||||
const int k = k0 + ix[0], c = ix[1], id = cid[c];
|
||||
half4 w = half4(0.0h);
|
||||
if (id >= 0) {
|
||||
const uint row = uint(id) * QKV + g;
|
||||
const char4 q4 = *(device const char4*)(kq + row * QD + k);
|
||||
w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)]));
|
||||
}
|
||||
rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w;
|
||||
}
|
||||
tA = tQ.template slice<32, 16>(k0, 0);
|
||||
opS.run(tA, rk, s);
|
||||
}
|
||||
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tS(sp[sg], dextents<int32_t, 2>(QCH, 16));
|
||||
s.store(tS);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// online softmax per head row: simdgroup 0, lane = row
|
||||
if (sg == 0 && lane < 16) {
|
||||
if (lane < HPG) {
|
||||
float m = mrun[lane];
|
||||
for (int c = 0; c < nc; ++c) {
|
||||
const int q = lane * QCH + c;
|
||||
m = max(m, (((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f));
|
||||
}
|
||||
const float f = exp(mrun[lane] - m);
|
||||
float l = lsum[lane] * f;
|
||||
for (int c = 0; c < QCH; ++c) {
|
||||
const int q = lane * QCH + c;
|
||||
const float e = c < nc ? exp(((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f - m) : 0.0f;
|
||||
l += e;
|
||||
p[q] = half(e);
|
||||
}
|
||||
mrun[lane] = m;
|
||||
lsum[lane] = l;
|
||||
fac[lane] = f;
|
||||
} else {
|
||||
for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f);
|
||||
fac[lane] = 0.0f;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims
|
||||
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
|
||||
if (acc.is_valid_element(i)) acc[i] *= fac[acc.get_multidimensional_index(i)[1]];
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < rv.get_capacity(); ++i) {
|
||||
const auto ix = rv.get_multidimensional_index(i); // (cell, dim)
|
||||
const int c = ix[0], d = d0 + ix[1], id = cid[c];
|
||||
rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vsc[(uint(id) * QKV + g) * 4 + sg])) : half(0.0f);
|
||||
}
|
||||
opO.run(tP, rv, acc);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
|
||||
if (acc.is_valid_element(i)) {
|
||||
const auto ix = acc.get_multidimensional_index(i); // (dim, row)
|
||||
const int h = ix[1], d = d0 + ix[0];
|
||||
if (h < HPG) {
|
||||
const int hh = g * HPG + h;
|
||||
const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d];
|
||||
xqh[(uint(t) * QH + hh) * QD + d] = half(acc[i] / lsum[h] * (1.0f / (1.0f + exp(-gate))));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
// metal/tests/coop_layout.cpp - which (index 0, index 1) elements of a single-simdgroup matmul2d's right input, left
|
||||
// input and destination cooperative tensors each lane holds (mm_bench.metal coop_layout).
|
||||
// coop_layout [TM_TN_TK]
|
||||
#include "mtl_rt.h"
|
||||
|
||||
#include <cstdio>
|
||||
#include <string>
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
const std::string cfg = argc > 1 ? argv[1] : "16_32_16";
|
||||
mtl::Context cx;
|
||||
const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
mtl::Pipeline p = cx.pipeline(lib, "coop_layout_" + cfg);
|
||||
mtl::Buffer out = cx.alloc(4 * 3 * 32 * 128 * 4);
|
||||
mtl::Dispatch d;
|
||||
d.p = &p;
|
||||
d.args = {mtl::Arg::B(out)};
|
||||
d.tg[0] = 32;
|
||||
cx.run({d});
|
||||
const int* o = out.as<int>();
|
||||
const char* nm[3] = {"right", "left", "dest"};
|
||||
for (int w = 0; w < 3; ++w) {
|
||||
std::printf("%s (tile %s):\n", nm[w], cfg.c_str());
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
const int* q = o + (w * 32 + l) * 128 * 4;
|
||||
std::printf(" l%-2d cap %d:", l, q[2]);
|
||||
for (int i = 0; i < q[2] && i < 128; ++i) std::printf(" (%d,%d)", q[i * 4], q[i * 4 + 1]);
|
||||
std::printf("\n");
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// metal/tests/coop_probe.h - which element layout this GPU gives matmul2d's single-simdgroup right-input cooperative
|
||||
// tensor: true = runs of 4 consecutive k at one n per lane (the M5), false = anything else (the M4: runs of 2 n at one
|
||||
// k). The kernels that fill right inputs per lane (mm_gu3, mm_down3, qsa_attn_mm, qsa_flash) take it as function
|
||||
// constant 7. COOP_K4=0/1 in the environment overrides.
|
||||
#pragma once
|
||||
#include "mtl_rt.h"
|
||||
|
||||
#include <cstdlib>
|
||||
|
||||
inline bool coop_k4(mtl::Context& cx) {
|
||||
if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0;
|
||||
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);
|
||||
mtl::Dispatch d;
|
||||
d.p = &p;
|
||||
d.args = {mtl::Arg::B(out)};
|
||||
d.tg[0] = 32;
|
||||
cx.run({d});
|
||||
// every lane: elements 4j .. 4j + 3 are (k, n), (k + 1, n), (k + 2, n), (k + 3, n)
|
||||
const int* o = out.as<int>();
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
const int* q = o + l * 128 * 4;
|
||||
const int cap = q[2];
|
||||
for (int i = 0; i + 3 < cap; i += 4)
|
||||
for (int j = 1; j < 4; ++j)
|
||||
if (q[(i + j) * 4] != q[i * 4] + j || q[(i + j) * 4 + 1] != q[i * 4 + 1]) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -4,6 +4,7 @@
|
||||
//
|
||||
// moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4] parity (64 experts)
|
||||
// moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..] per-layer time, 512 experts
|
||||
#include "coop_probe.h"
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
|
||||
@@ -31,7 +32,7 @@ float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); }
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1;
|
||||
bool bench = false, v2 = false;
|
||||
bool bench = false, v2 = false, v3 = false;
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
std::string a = argv[i];
|
||||
if (a == "--gu") gu = type_of(argv[++i]);
|
||||
@@ -41,6 +42,7 @@ int main(int argc, char** argv) {
|
||||
else if (a == "--dtile") std::sscanf(argv[++i], "%d,%d,%d", &DTN, &DTK, &DSG); // the down kernel's TN,TK,SG
|
||||
else if (a == "--bench") bench = true;
|
||||
else if (a == "--v2") v2 = true;
|
||||
else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup)
|
||||
}
|
||||
if (bench && T == 256) T = 2048;
|
||||
mtl::Context cx;
|
||||
@@ -89,8 +91,12 @@ int main(int argc, char** argv) {
|
||||
const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG);
|
||||
mtl::Constants probe;
|
||||
probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0);
|
||||
mtl::Pipeline pgu = cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe),
|
||||
pdn = cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)),
|
||||
const bool k4 = coop_k4(cx);
|
||||
if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)");
|
||||
mtl::Pipeline pgu = v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4))
|
||||
: cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe),
|
||||
pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4))
|
||||
: cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)),
|
||||
pgat = cx.pipeline(lib, "mm_gather");
|
||||
MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk};
|
||||
auto encode = [&](int layer, std::vector<mtl::Dispatch>& ds, bool plan, bool gemm) {
|
||||
@@ -113,9 +119,9 @@ int main(int argc, char** argv) {
|
||||
mtl::Dispatch g; g.p = &pgu;
|
||||
if (v2) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
|
||||
else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
|
||||
g.grid[0] = N_FF / TN; g.grid[1] = meta.as<uint32_t>()[0]; g.tg[0] = 32 * SG; ds.push_back(g);
|
||||
g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as<uint32_t>()[0]; g.tg[0] = 32 * SG; ds.push_back(g);
|
||||
mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
|
||||
d.grid[0] = N_EMBD / (v2 ? DTN : TN); d.grid[1] = meta.as<uint32_t>()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d);
|
||||
d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as<uint32_t>()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d);
|
||||
const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0};
|
||||
mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)};
|
||||
cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb);
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
// Every dense matrix is dequantized to fp16 (dq_f16) once a chunk, before its GEMM. After the last chunk: the final
|
||||
// HC read and the LM head on the LAST row only (the decode kernels: gr_norm / gr_down / gr_up, GEMV, argmax).
|
||||
// Left out: the PLE path and the MTP draft.
|
||||
#include "coop_probe.h"
|
||||
#include "gguf_lite.h"
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
@@ -96,10 +97,12 @@ int main(int argc, char** argv) {
|
||||
const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal",
|
||||
"gemv.metal", "prompt.metal"});
|
||||
const int lgdn = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
const int lqsa = cx.compile({"qsa.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
const int lqsa = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
const int lsmp = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
std::printf("kernels compiled in %.1f s\n", mtl::now_s() - tc0);
|
||||
std::map<std::string, mtl::Pipeline> pcache;
|
||||
const bool k4 = coop_k4(cx);
|
||||
std::printf("matmul2d right-input layout: %s\n", k4 ? "runs of 4 k (per-lane run decoders)" : "other (per-element decode)");
|
||||
auto P = [&](int l, const std::string& fn, const mtl::Constants& c = {}, const std::string& key = "") -> const mtl::Pipeline* {
|
||||
const std::string k = key.empty() ? fn : key;
|
||||
auto it = pcache.find(k);
|
||||
@@ -122,10 +125,30 @@ int main(int argc, char** argv) {
|
||||
ggb = F32(MT * 48 * 2), gkq = F32(MT * 16), go = F32(MT * 6144);
|
||||
const char* hc_env = std::getenv("PC_HC");
|
||||
const std::string hc_mode = hc_env ? hc_env : "fused"; // fused | plain
|
||||
const char* qsa_env = std::getenv("PC_QSA");
|
||||
// mm (matrix units; not on the M4's layout yet: qsa_attn_mm writes nothing there) | v1 (qsa_attn, fp32)
|
||||
const std::string qsa_mode = qsa_env ? qsa_env : (k4 ? "mm" : "v1");
|
||||
// the expert GEMMs: "v2" (threadgroup-staged weights) or "v3:TM,TN,TK,SG" (register operands, TN rows a simdgroup)
|
||||
const char* gu_env = std::getenv("PC_GU");
|
||||
const char* dn_env = std::getenv("PC_DN");
|
||||
// defaults (M5 Max, moe_mm_parity --v3 --bench): gate/up 16,32,32,4 (every format), down 16,64,32,2
|
||||
// without the M5's right-input layout (the M4) the per-element v3 decode is slower than v2: v2 there
|
||||
const std::string gu_mode = gu_env ? gu_env : (k4 ? "v3:auto" : "v2"), dn_mode = dn_env ? dn_env : (k4 ? "v3:16,64,32,2" : "v2");
|
||||
auto v3_tile = [](const std::string& m, int* t) { std::sscanf(m.c_str() + 3, "%d,%d,%d,%d", &t[0], &t[1], &t[2], &t[3]); };
|
||||
int gu3[4] = {16, 64, 16, 2}, dn3[4] = {16, 64, 16, 2};
|
||||
if (gu_mode.rfind("v3:", 0) == 0 && gu_mode != "v3:auto") v3_tile(gu_mode, gu3);
|
||||
else if (gu_mode == "v3:auto") { gu3[0] = 16; gu3[1] = 32; gu3[2] = 32; gu3[3] = 4; }
|
||||
auto tl3 = [](const int* t) { return std::to_string(t[0]) + "_" + std::to_string(t[1]) + "_" + std::to_string(t[2]) + "_" + std::to_string(t[3]); };
|
||||
auto gu3_for = [&](int type) -> std::string {
|
||||
(void)type; // one tile for every format since the paired-run decoders (W8)
|
||||
return tl3(gu3);
|
||||
};
|
||||
if (dn_mode.rfind("v3:", 0) == 0) v3_tile(dn_mode, dn3);
|
||||
const bool gu_v3 = gu_mode.rfind("v3:", 0) == 0, dn_v3 = dn_mode.rfind("v3:", 0) == 0;
|
||||
const char* gdn_env = std::getenv("PC_GDN");
|
||||
const std::string gdn_mode = gdn_env ? gdn_env : "scanr_4";
|
||||
mtl::Buffer qfull = F32(MT * 12288), kcur = F32(MT * 512), vcur = F32(MT * 512), iqr = F32(MT * 512), ikr = F32(MT * 128),
|
||||
qcur = F32(MT * 24 * 256), qidx = F32(MT * 4 * 128), scores = F32(MT * MB), qids = cx.alloc(4 * MT * QB),
|
||||
qcur = F32(MT * 24 * 256), qcurh = F16(MT * 24 * 256), qidx = F32(MT * 4 * 128), scores = F32(MT * MB), qids = cx.alloc(4 * MT * QB),
|
||||
part = F32(MT * QCHUNKS * 24 * 258), XQ = F32(4 * MT * 6144 + MT * 384), xqh = F16(MT * 6144);
|
||||
mtl::Buffer logits = F32(MT * NEXP), ids = cx.alloc(4 * MT * TOPK), wts = F32(MT * TOPK),
|
||||
bcnt = cx.alloc(4 * (size_t)nblk_max * NEXP), boff = cx.alloc(4 * (size_t)nblk_max * NEXP), eoff = cx.alloc(4 * NEXP),
|
||||
@@ -177,8 +200,8 @@ int main(int argc, char** argv) {
|
||||
ref::dequant(te.type, (const uint8_t*)f.map + te.offset + id * rb, E.data() + t * N, N);
|
||||
}
|
||||
}
|
||||
meta.as<uint32_t>()[2] = NFF / 64;
|
||||
meta.as<uint32_t>()[5] = N / 64;
|
||||
meta.as<uint32_t>()[2] = gu_v3 ? NFF / (gu3[1] * gu3[3]) : NFF / 64;
|
||||
meta.as<uint32_t>()[5] = dn_v3 ? N / (dn3[1] * dn3[3]) : N / 64;
|
||||
|
||||
// ---- the chunk's dispatches
|
||||
std::vector<mtl::Dispatch> ds;
|
||||
@@ -330,14 +353,26 @@ int main(int argc, char** argv) {
|
||||
mtl::Arg::B(pooled[qi]), mtl::Arg::V(a)},
|
||||
33, T, 256, "qsa prep");
|
||||
if (pos0 + T > QB) { // some token's context is past the budget: select
|
||||
if (qsa_mode == "mm")
|
||||
D(P(lqsa, "qsa_score_mm"), {mtl::Arg::B(qidx), mtl::Arg::B(pooled[qi]), mtl::Arg::B(dead[qi]),
|
||||
mtl::Arg::B(scores), mtl::Arg::V(a)},
|
||||
((pos0 + T) / 4 + 1 + 63) / 64, (T + 15) / 16, 128, "qsa score");
|
||||
else
|
||||
D(P(lqsa, "qsa_score"), {mtl::Arg::B(qidx), mtl::Arg::B(pooled[qi]), mtl::Arg::B(dead[qi]), mtl::Arg::B(scores),
|
||||
mtl::Arg::V(a)},
|
||||
((pos0 + T) / 4 + 1 + 7) / 8, T, 256, "qsa sel");
|
||||
((pos0 + T) / 4 + 1 + 7) / 8, T, 256, "qsa score");
|
||||
D(P(lqsa, "qsa_topk"), {mtl::Arg::B(scores), mtl::Arg::B(qids), mtl::Arg::V(a)}, T, 1, 1024, "qsa sel");
|
||||
} else {
|
||||
// within the budget the selection is the identity: qsa_topk only writes 0..n_kv-1
|
||||
D(P(lqsa, "qsa_topk"), {mtl::Arg::B(scores), mtl::Arg::B(qids), mtl::Arg::V(a)}, T, 1, 1024, "qsa sel");
|
||||
}
|
||||
if (qsa_mode == "mm") {
|
||||
const uint32_t nq = (uint32_t)T * 24 * 256;
|
||||
D(P(lqsa, "qsa_q_f16"), {mtl::Arg::B(qcur), mtl::Arg::B(qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256, "qsa attn");
|
||||
D(P(lqsa, "qsa_attn_mm", mtl::Constants().b(7, k4)), {mtl::Arg::B(qcurh), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]),
|
||||
mtl::Arg::B(vsc[qi]), mtl::Arg::B(qids), mtl::Arg::B(part), mtl::Arg::V(a)},
|
||||
QCHUNKS, 2, 32, "qsa attn", T);
|
||||
} else
|
||||
D(P(lqsa, "qsa_attn"), {mtl::Arg::B(qcur), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]),
|
||||
mtl::Arg::B(vsc[qi]), mtl::Arg::B(qids), mtl::Arg::B(part), mtl::Arg::V(a)},
|
||||
QCHUNKS, 2, 256, "qsa attn", T);
|
||||
@@ -385,16 +420,18 @@ int main(int argc, char** argv) {
|
||||
g.ind_off = 8 * 4;
|
||||
push(g, "experts");
|
||||
mtl::Dispatch gu;
|
||||
gu.p = P(lib, "mm_gu2_" + std::to_string(eg.type) + "_16_64_16_2", mtl::Constants().i(5, 0));
|
||||
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));
|
||||
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] = 64;
|
||||
gu.tg[0] = gu_v3 ? 32 * gu3[3] : 64;
|
||||
gu.indirect = &meta;
|
||||
gu.ind_off = 2 * 4;
|
||||
push(gu, "experts gu");
|
||||
mtl::Dispatch dn;
|
||||
dn.p = P(lib, "mm_down2_" + std::to_string(ed.type) + (ed.type == ref::IQ4_NL ? "_16_64_32_2" : "_16_64_16_2"));
|
||||
dn.p = dn_v3 ? P(lib, "mm_down3_" + std::to_string(ed.type) + "_" + tl3(dn3), mtl::Constants().b(7, k4))
|
||||
: P(lib, "mm_down2_" + std::to_string(ed.type) + (ed.type == ref::IQ4_NL ? "_16_64_32_2" : "_16_64_16_2"));
|
||||
dn.args = {mtl::Arg::B(W), mtl::Arg::B(HS), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
|
||||
dn.tg[0] = 64;
|
||||
dn.tg[0] = dn_v3 ? 32 * dn3[3] : 64;
|
||||
dn.indirect = &meta;
|
||||
dn.ind_off = 5 * 4;
|
||||
push(dn, "experts dn");
|
||||
@@ -467,6 +504,10 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
if (breakdown) {
|
||||
// each class alone on the first chunk (pos0 0) and on a chunk at pos0 = PMAX - C (QSA past the budget)
|
||||
if (PMAX > C) { // the classes of a chunk past the budget (score) too
|
||||
run_chunk(std::max(0, PMAX - C), C, false);
|
||||
all_cls.insert(cls.begin(), cls.end());
|
||||
}
|
||||
for (int p0 : {0, std::max(0, PMAX - C)}) {
|
||||
const std::string keep_only = only;
|
||||
std::map<std::string, double> by;
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
// metal/tests/qsa_prompt_parity.cpp - qsa_attn_mm (qsa_prompt.metal, the matrix units, fp16 operands) against qsa_attn
|
||||
// (fp32, itself checked by qsa_parity), both followed by qsa_merge: the attention output of T query tokens on random
|
||||
// int8 caches, within the budget (all cells, causal) and past it (2051 selected cells per token); then the speed.
|
||||
// qsa_prompt_parity [--bench]
|
||||
#include "coop_probe.h"
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
constexpr int QH = 24, QKV = 2, QD = 256, QB = 2051, NCH = 33;
|
||||
struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; };
|
||||
double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; }
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
const bool bench = argc > 1 && std::string(argv[1]) == "--bench";
|
||||
std::setvbuf(stdout, nullptr, _IOLBF, 0);
|
||||
const bool noflash = std::getenv("QSA_NOFLASH") != nullptr;
|
||||
mtl::Context cx;
|
||||
const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
std::printf("device %s\n", cx.device_name().c_str());
|
||||
const int T = bench ? 2048 : 96, MAXC = 40000;
|
||||
mtl::Buffer kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4),
|
||||
vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), qcur = cx.alloc(4 * (size_t)T * QH * QD), qh = cx.alloc(2 * (size_t)T * QH * QD),
|
||||
qfull = cx.alloc(4 * (size_t)T * QH * 2 * QD), ids = cx.alloc(4 * (size_t)T * QB),
|
||||
part = cx.alloc(4 * (size_t)T * NCH * QH * (QD + 2)), xa = cx.alloc(4 * ((size_t)4 * T * QH * QD + (size_t)T * QH * QD / 16)),
|
||||
xb = cx.alloc(4 * ((size_t)4 * T * QH * QD + (size_t)T * QH * QD / 16));
|
||||
ref::Rng r(5);
|
||||
for (size_t i = 0; i < kq.bytes; ++i) { kq.as<int8_t>()[i] = (int8_t)((int)(r.next() % 255) - 127); vq.as<int8_t>()[i] = (int8_t)((int)(r.next() % 255) - 127); }
|
||||
for (size_t i = 0; i < ksc.bytes / 2; ++i) { ksc.as<uint16_t>()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); vsc.as<uint16_t>()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); }
|
||||
for (size_t i = 0; i < (size_t)T * QH * QD; ++i) qcur.as<float>()[i] = r.sym() * 3.0f;
|
||||
for (size_t i = 0; i < (size_t)T * QH * 2 * QD; ++i) qfull.as<float>()[i] = r.sym();
|
||||
mtl::Pipeline pa = cx.pipeline(lib, "qsa_attn"), pm = cx.pipeline(lib, "qsa_merge"), pmm = cx.pipeline(lib, "qsa_attn_mm", mtl::Constants().b(7, coop_k4(cx))),
|
||||
pq = cx.pipeline(lib, "qsa_q_f16"), pf = cx.pipeline(lib, "qsa_flash", mtl::Constants().b(7, coop_k4(cx)));
|
||||
mtl::Buffer xf = cx.alloc(2 * (size_t)T * QH * QD);
|
||||
int fails = 0;
|
||||
for (uint32_t pos0 : {0u, 30000u}) {
|
||||
// ids: within the budget the identity; past it 2051 distinct ascending cells < n_kv
|
||||
for (int t = 0; t < T; ++t) {
|
||||
const uint32_t n_kv = pos0 + t + 1;
|
||||
int* o = ids.as<int>() + (size_t)t * QB;
|
||||
if (n_kv <= (uint32_t)QB) { for (uint32_t c = 0; c < n_kv; ++c) o[c] = (int)c; continue; }
|
||||
std::set<int> s;
|
||||
while ((int)s.size() < QB) s.insert((int)(r.next() % n_kv));
|
||||
int j = 0;
|
||||
for (int c : s) o[j++] = c;
|
||||
}
|
||||
QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f};
|
||||
auto D = [&](const mtl::Pipeline& p, std::vector<mtl::Arg> args, uint32_t gx, uint32_t gy, uint32_t gz, uint32_t tg) {
|
||||
mtl::Dispatch d; d.p = &p; d.args = std::move(args); d.grid[0] = gx; d.grid[1] = gy; d.grid[2] = gz; d.tg[0] = tg; return d;
|
||||
};
|
||||
const uint32_t nq = T * QH * QD;
|
||||
std::vector<mtl::Dispatch> ref_ds = {
|
||||
D(pa, {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}, NCH, QKV, T, 256),
|
||||
D(pm, {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(xa), mtl::Arg::V(a)}, QH, T, 1, 256)};
|
||||
std::vector<mtl::Dispatch> mm_ds = {
|
||||
D(pq, {mtl::Arg::B(qcur), mtl::Arg::B(qh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 1, 256),
|
||||
D(pmm, {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}, NCH, QKV, T, 32),
|
||||
D(pm, {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(xb), mtl::Arg::V(a)}, QH, T, 1, 256)};
|
||||
const mtl::Dispatch fl = D(pf, {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc),
|
||||
mtl::Arg::B(ids), mtl::Arg::B(qfull), mtl::Arg::B(xf), mtl::Arg::V(a)}, QKV, T, 1, 128);
|
||||
if (bench) {
|
||||
auto timeit = [&](const std::vector<mtl::Dispatch>& ds) {
|
||||
cx.warm(ds, 0.3);
|
||||
std::vector<double> t;
|
||||
for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds));
|
||||
return median(t);
|
||||
};
|
||||
std::printf("pos0 %5u, T %d: qsa_attn %.2f ms, qsa_attn_mm %.2f ms, merge %.2f ms; qsa_flash (no merge) %.2f ms\n", pos0,
|
||||
T, timeit({ref_ds[0]}) * 1e3, timeit({mm_ds[1]}) * 1e3, timeit({ref_ds[1]}) * 1e3, timeit({fl}) * 1e3);
|
||||
continue;
|
||||
}
|
||||
cx.run(ref_ds);
|
||||
std::fill(part.as<float>(), part.as<float>() + part.bytes / 4, std::nanf("")); // a kernel that writes nothing shows
|
||||
cx.run(mm_ds);
|
||||
double se = 0, sr = 0;
|
||||
for (size_t i = 0; i < (size_t)T * QH * QD; ++i) {
|
||||
const double e = xb.as<float>()[i] - xa.as<float>()[i];
|
||||
se += e * e;
|
||||
sr += (double)xa.as<float>()[i] * xa.as<float>()[i];
|
||||
}
|
||||
const double rel = std::sqrt(se / sr);
|
||||
const bool ok = rel < 3e-3;
|
||||
fails += !ok;
|
||||
std::printf("pos0 %5u, T %d: qsa_attn_mm vs qsa_attn (after merge) rel L2 %.2e %s\n", pos0, T, rel, ok ? "PASS" : "FAIL");
|
||||
if (noflash) continue;
|
||||
cx.run({fl});
|
||||
double sf = 0;
|
||||
for (size_t i = 0; i < (size_t)T * QH * QD; ++i) {
|
||||
const double e = ref::f16_to_f32(xf.as<uint16_t>()[i]) - xa.as<float>()[i];
|
||||
sf += e * e;
|
||||
}
|
||||
const double relf = std::sqrt(sf / sr);
|
||||
const bool okf = relf < 3e-3;
|
||||
fails += !okf;
|
||||
std::printf("pos0 %5u, T %d: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 %.2e %s\n", pos0, T, relf, okf ? "PASS" : "FAIL");
|
||||
}
|
||||
// qsa_score_mm vs qsa_score at pos0 30000: the scores of every block <= n_bid
|
||||
{
|
||||
const uint32_t pos0 = 30000, MB = MAXC / 4 + 2, nb = (pos0 + T) / 4 + 1;
|
||||
mtl::Buffer qidx = cx.alloc(4 * (size_t)T * 4 * 128), pooled = cx.alloc(4 * (size_t)MB * 128), dead = cx.alloc(4 * 128),
|
||||
sa = cx.alloc(4 * (size_t)T * MB), sb = cx.alloc(4 * (size_t)T * MB);
|
||||
for (size_t i = 0; i < (size_t)T * 4 * 128; ++i) qidx.as<float>()[i] = r.sym();
|
||||
for (size_t i = 0; i < (size_t)MB * 128; ++i) pooled.as<float>()[i] = r.sym();
|
||||
for (int i = 0; i < 128; ++i) dead.as<float>()[i] = r.sym();
|
||||
QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f};
|
||||
mtl::Pipeline ps = cx.pipeline(lib, "qsa_score"), psm = cx.pipeline(lib, "qsa_score_mm");
|
||||
mtl::Dispatch d1; d1.p = &ps; d1.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(sa), mtl::Arg::V(a)};
|
||||
d1.grid[0] = (nb + 7) / 8; d1.grid[1] = T; d1.tg[0] = 256;
|
||||
mtl::Dispatch d2; d2.p = &psm; d2.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(sb), mtl::Arg::V(a)};
|
||||
d2.grid[0] = (nb + 63) / 64; d2.grid[1] = (T + 15) / 16; d2.tg[0] = 128;
|
||||
if (bench) {
|
||||
auto timeit = [&](const mtl::Dispatch& d) {
|
||||
cx.warm({d}, 0.3);
|
||||
std::vector<double> t;
|
||||
for (int i = 0; i < 5; ++i) t.push_back(cx.run({d}));
|
||||
return median(t);
|
||||
};
|
||||
std::printf("score pos0 %u, T %d: qsa_score %.2f ms, qsa_score_mm %.2f ms\n", pos0, T, timeit(d1) * 1e3, timeit(d2) * 1e3);
|
||||
} else {
|
||||
cx.run({d1});
|
||||
cx.run({d2});
|
||||
double worst = 0;
|
||||
for (int t = 0; t < T; ++t) {
|
||||
const uint32_t nbid = (pos0 + t + 1) / 4;
|
||||
for (uint32_t b = 0; b <= nbid; ++b) {
|
||||
const double x = sa.as<float>()[(size_t)t * MB + b], y = sb.as<float>()[(size_t)t * MB + b];
|
||||
worst = std::max(worst, std::fabs(x - y) / (std::fabs(x) + 1.0)); // scores are O(10)
|
||||
}
|
||||
}
|
||||
const bool ok = worst < 1e-5;
|
||||
fails += !ok;
|
||||
std::printf("score pos0 %u, T %d: qsa_score_mm vs qsa_score worst rel %.2e %s\n", pos0, T, worst, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
}
|
||||
if (!bench) std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS");
|
||||
return fails ? 1 : 0;
|
||||
}
|
||||
Reference in new issue
Block a user