diff --git a/metal/kernels/qsa2.metal b/metal/kernels/qsa2.metal index d5b594c..30439a1 100644 --- a/metal/kernels/qsa2.metal +++ b/metal/kernels/qsa2.metal @@ -120,7 +120,15 @@ kernel void qsa_merge2(device const float* part [[buffer(0)]], device const floa } float o = 0.0f; device const float* q = pp + tid; - for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + int k = 0; + for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { + const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; + o = fma(p0, ek[k], o); + o = fma(p1, ek[k + 1], o); + o = fma(p2, ek[k + 2], o); + o = fma(p3, ek[k + 3], o); + } + for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); threadgroup_barrier(mem_flags::mem_threadgroup); o = o / L_sh; const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; diff --git a/metal/kernels/qsa_prompt.metal b/metal/kernels/qsa_prompt.metal index 444e0c4..2ee4a4b 100644 --- a/metal/kernels/qsa_prompt.metal +++ b/metal/kernels/qsa_prompt.metal @@ -82,8 +82,11 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* 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); + // the row is clamped and the select is on the value: a select with the device load in the arm + // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); + rk[i] = id >= 0 ? half(v) : half(0.0f); } } else #pragma clang loop unroll(full) @@ -136,7 +139,9 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* for (uint16_t i = 0; i < rv.get_capacity(); ++i) { const auto ix = rv.get_multidimensional_index(i); // (cell, dim) const int c = c0 + 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); + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); + rv[i] = id >= 0 ? half(v) : half(0.0f); } } else { uint vo[8]; @@ -332,8 +337,9 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq 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(ks_tg[ix[1]][k >> 6])) : half(0.0f); + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6]); + rk[i] = id >= 0 ? half(v) : half(0.0f); } } else #pragma clang loop unroll(full) @@ -398,7 +404,9 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq for (uint16_t i = 0; i < rv.get_capacity(); ++i) { const auto ix = rv.get_multidimensional_index(i); // (cell, dim) const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; - rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs_tg[c][sg])) : half(0.0f); + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs_tg[c][sg]); + rv[i] = id >= 0 ? half(v) : half(0.0f); } } else { // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row diff --git a/metal/tests/moe_mm_parity.cpp b/metal/tests/moe_mm_parity.cpp index 8e06a77..a44397c 100644 --- a/metal/tests/moe_mm_parity.cpp +++ b/metal/tests/moe_mm_parity.cpp @@ -176,7 +176,7 @@ int main(int argc, char** argv) { argc > 3 ? std::atoi(argv[3]) : 512, argc > 4 ? std::atoi(argv[4]) : N_EMBD); } - 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; + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 64, SG = 2, DTN = -1, DTK = -1, DSG = -1; bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; int MS = 1; bool gu_only = false, dn_only = false; @@ -283,17 +283,17 @@ int main(int argc, char** argv) { if (gemm) { MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; - if (v2) { + if (v2 || v3) { const uint32_t ne = N_EMBD; mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); } 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)}; + if (v2 || v3) 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 / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[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 / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d); + mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 || v3 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * ((v2 || v3) ? 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);