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: QSA prompt attention correct on the M4 (K = 32 runs, threadgroup-staged store, fresh per-chunk accumulator)

Browse Source 
- qsa_attn_mm: the P V matmul in K = 32 cell runs (K = 64 runs computed garbage on the M4), the output through
  threadgroup memory (a direct matmul2d store into 258-float rows lost the threadgroup on the M4); also faster on the
  M5 (8.1 -> 7.4 ms at pos 0)
- qsa_flash: the same K = 32 runs (was a hang on the M4) and a fresh per-chunk accumulator folded into the running
  one (on the M4 a run does not accumulate into elements rescaled in place); now about equal to attn_mm + merge on
  the M5 (9.4 vs 9.5 ms at pos 0, 20.1 vs 20.3 at 30000) and the prompt_chain default there (PC_QSA=flash|mm|v1);
  the M4 defaults to v1 attention (faster there) with the GEMM score
- qsa_attn_mm bisection probe (function constant 8, QSA_STOP), qsa_prompt_parity QSA_DEBUG / pipeline report;
  coop_layout prints the destination layout too
This commit is contained in:
Niko1221 committed 2026-10-12 16:57:55 +02:00
1 parent 6f6c555b17
commit dbc2e2fd8f
4 files changed
+87 -23
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

metal/kernels/mm_bench.metal
+7
View File
@@ -61,6 +61,13 @@ kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_inde
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>();
auto dd = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(l)>,
__tensor_ops_detail::__remove_addrspace_t<decltype(r)>, float>();
for (uint16_t i = 0; i < dd.get_capacity() && i < 128; ++i) {
auto ix = dd.get_multidimensional_index(i);
device int* o = out + ((2 * 32 + lane) * 128 + i) * 4;
o[0] = ix[0]; o[1] = ix[1]; o[2] = dd.get_capacity();
}
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;
metal/kernels/qsa_prompt.metal
+50 -19
View File
@@ -16,6 +16,11 @@ constant bool COOP_K4 [[function_constant(7)]];
constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true;
#endif
// a bisection probe (function constant 8): stop at stage n and write m = n (1: after the ids, 2: after S, 3: after
// the softmax); 0 = normal
constant int QSA_STOP [[function_constant(8)]];
constant int qsa_stop = is_function_constant_defined(QSA_STOP) ? QSA_STOP : 0;
// 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)]],
@@ -42,6 +47,7 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char*
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);
if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; }
// ---- S = Q K^T
{
constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
@@ -81,6 +87,7 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char*
acc.store(tS);
}
simdgroup_barrier(mem_flags::mem_threadgroup);
if (qsa_stop == 2) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 2.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; }
// ---- softmax pieces per head (lane = head for the reductions; 12 heads)
float mh = -INFINITY, lh = 0.0f;
if (lane < HPG) {
@@ -94,26 +101,42 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char*
for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f);
}
simdgroup_barrier(mem_flags::mem_threadgroup);
if (qsa_stop == 3) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 3.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; }
// ---- acc = P V, 64 dims at a time
{
constexpr auto desc = matmul2d_descriptor(16, 64, QCH, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
// K = 32 cells a run (two runs a 64-cell chunk): K = 64 runs compute garbage on the M4
constexpr auto desc = matmul2d_descriptor(16, 64, 32, 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>();
auto tPh = tP.template slice<32, 16>(0, 0);
auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tPh)>, __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;
for (int c0 = 0; c0 < QCH; c0 += 32) {
#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);
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);
}
tPh = tP.template slice<32, 16>(c0, 0);
op.run(tPh, rv, acc);
}
op.run(tP, rv, acc);
auto mO = tO.slice(d0, 0);
acc.store(mO);
if (qsa_stop == 4) { // the bisection: the old direct store into the 258-float rows
auto mO = tO.slice(d0, 0);
acc.store(mO);
continue;
}
// through threadgroup memory (s is free now): a matmul2d store into rows of QD + 2 floats loses the
// whole threadgroup on the M4
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tT(s, dextents<int32_t, 2>(64, 16));
simdgroup_barrier(mem_flags::mem_threadgroup);
acc.store(tT);
simdgroup_barrier(mem_flags::mem_threadgroup);
for (int it = lane; it < HPG * 64; it += 32) pp[(it >> 6) * (QD + 2) + d0 + (it & 63)] = s[it];
}
}
if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; }
@@ -198,17 +221,20 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq
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);
constexpr auto dO = matmul2d_descriptor(16, 64, 32, 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)>,
auto tPh = tP.template slice<32, 16>(0, 0);
auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tPh)>,
__tensor_ops_detail::__remove_addrspace_t<decltype(rv)>, float>();
auto tot = acc; // the running output; acc takes one chunk (fresh each chunk: on the M4 a run does not accumulate
// into elements rescaled in place)
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
if (acc.is_valid_element(i)) acc[i] = 0.0f;
if (acc.is_valid_element(i)) tot[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);
@@ -282,14 +308,19 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq
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]];
if (acc.is_valid_element(i)) acc[i] = 0.0f;
for (int c0 = 0; c0 < QCH; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4)
#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);
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(vsc[(uint(id) * QKV + g) * 4 + sg])) : half(0.0f);
}
tPh = tP.template slice<32, 16>(c0, 0);
opO.run(tPh, rv, acc);
}
opO.run(tP, rv, acc);
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
@@ -299,7 +330,7 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq
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))));
xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate))));
}
}
}
metal/tests/prompt_chain.cpp
+15 -3
View File
@@ -126,8 +126,9 @@ int main(int argc, char** argv) {
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");
// flash (one pass on the matrix units, no partials) | mm (split-K on the matrix units + merge) | v1 (qsa_attn,
// fp32 + merge). The M5: flash (= mm + merge or better); the M4: v1 attention (faster there) with the GEMM score.
const std::string qsa_mode = qsa_env ? qsa_env : (k4 ? "flash" : "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");
@@ -353,7 +354,7 @@ 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")
if (!std::getenv("PC_QSA_SCORE")) // the GEMM score unless PC_QSA_SCORE (= the scalar qsa_score)
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");
@@ -366,6 +367,17 @@ int main(int argc, char** argv) {
// 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 == "flash") {
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_flash", 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(qfull), mtl::Arg::B(xqh), mtl::Arg::V(a)},
2, T, 128, "qsa attn");
gemm(xqh, Wout, bo, T, N, 6144, "qsa proj");
++qi;
goto moe;
}
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");
metal/tests/qsa_prompt_parity.cpp
+15 -1
View File
@@ -38,9 +38,11 @@ int main(int argc, char** argv) {
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))),
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)).i(8, std::getenv("QSA_STOP") ? std::atoi(std::getenv("QSA_STOP")) : 0)),
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);
std::printf("pipelines: qsa_attn_mm max threads %d, tg mem %d; qsa_flash max threads %d, tg mem %d\n", pmm.max_threads,
pmm.static_tg_mem, pf.max_threads, pf.static_tg_mem);
int fails = 0;
for (uint32_t pos0 : {0u, 30000u}) {
// ids: within the budget the identity; past it 2051 distinct ascending cells < n_kv
@@ -81,6 +83,18 @@ int main(int argc, char** argv) {
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);
if (std::getenv("QSA_DEBUG")) {
// how much of the partials the matrix-unit kernel wrote, and token 5's first chunk (m, l, acc[0..3])
size_t written = 0, nonfinite = 0;
const float* pp = part.as<float>();
for (size_t i = 0; i < (size_t)T * NCH * QH * (QD + 2); ++i) {
if (!std::isnan(pp[i])) ++written;
if (!std::isfinite(pp[i]) && !std::isnan(pp[i])) ++nonfinite;
}
const float* q = pp + ((size_t)5 * NCH * QH) * (QD + 2);
std::printf("debug pos0 %u: %zu of %zu partial floats written (%zu inf); t5 h0: m %g l %g acc %g %g %g %g\n", pos0,
written, (size_t)T * NCH * QH * (QD + 2), nonfinite, q[QD], q[QD + 1], q[0], q[1], q[2], q[3]);
}
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];