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<uint32_t>()[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<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 / (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);
+ 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<uint32_t>()[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);