mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 18:35:04 -10:00
metal: Q8_0 HC read the default in token_chain (STRATA_HC_Q8=0 = exact BF16); prompt chain, GDN prompt scan, fused HC prompt read
- token_chain: the HC read defaults to the Q8_0 copies through the 3-dispatch read (q8_3); STRATA_HC_Q8=0 selects the exact BF16 path (--hc still overrides) - prompt.metal: route_topk, gr_epi_down/up, swiglu_h, moe_out, gr_norm_h (fp16 xn), gemm_hc_up (the HC up GEMM with the mixing epilogue; w_up rows interleaved by dq_f16's new perm field) - gdn_prompt.metal: gdn_pre (parallel conv/norms/gates) -> gdn_scan_r (state scan, one barrier a token) -> gdn_post - dense_mm.metal: DqArgs gains `perm` (4th field) - moe_mm.metal: mm_scan zero-fills the unused simdgroup slots; indirect grids for gather/gate-up/down in meta; MM_PROBE bench constant (function constant 5: pipelines of mm_gu2 need it set) - tests: prompt_chain (whole prompt in chunks, class breakdown), prompt_parity, gdn_prompt_parity; mm_bench tiles
This commit is contained in:
Niko1221 committed
1 parent
773b44b922
commit
4775d23372
12 files changed
+1431
-13
No files matched your search
@@ -15,6 +15,8 @@ struct DqArgs {
|
||||
uint n_in; // values per row
|
||||
uint n_out; // rows
|
||||
uint row_bytes;
|
||||
uint perm; // 0: row r -> r; p > 0: row r = c * (n_out / p) + d -> d * p + c (the HC up matrix with the p
|
||||
// streams of an output interleaved, for gemm_hc_up's epilogue)
|
||||
};
|
||||
|
||||
// one thread per unit: grid (units per row / 64, n_out), 64 threads
|
||||
@@ -32,7 +34,9 @@ kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer
|
||||
float w[UNIT];
|
||||
int p0, p1;
|
||||
DW<F>::run(W + ulong(r) * a.row_bytes, u, k0, w, p0, p1, tab);
|
||||
device half* o = Wh + ulong(r) * a.n_in + k0;
|
||||
const int q = a.perm ? int(a.n_out / a.perm) : 1;
|
||||
const int ro = a.perm ? (r % q) * int(a.perm) + r / q : r;
|
||||
device half* o = Wh + ulong(ro) * a.n_in + k0;
|
||||
for (int j = 0; j < H; ++j) {
|
||||
o[p0 + j] = half(w[j]);
|
||||
o[p1 + j] = half(w[H + j]);
|
||||
|
||||
@@ -0,0 +1,323 @@
|
||||
// metal/kernels/gdn_prompt.metal - a GDN layer's recurrence for a PROMPT chunk of T tokens, in three passes: the work
|
||||
// that does not touch the state runs parallel over the tokens; only the state update is a scan, and the columns of a
|
||||
// head's state are independent in it (column j needs k, q, v_j, the gate and beta), so a lane owns one column.
|
||||
//
|
||||
// gdn_pre per (token, head): conv (4 taps; history oldest first) -> SiLU; q, k L2 norm * 1/sqrt(128) (the decode
|
||||
// step's arithmetic); kq = k . q per key head; gate = exp(softplus(alpha + dt) * a), beta = sigmoid(beta);
|
||||
// the chunk's last 3 conv inputs -> hist_out
|
||||
// gdn_scan per (value head, 32 columns x CB simdgroups): the state column in registers (128 floats a lane);
|
||||
// per token a = S^T k, b = S^T q (old state), delta = (v - g a) beta, o = (g b + delta kq) / sqrt(128)
|
||||
// (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory
|
||||
// gdn_post per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand)
|
||||
// and, with `f32`, fp32
|
||||
//
|
||||
// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k.
|
||||
// The kernels are new entries (nothing of gdn.metal changes).
|
||||
|
||||
struct GdnPArgs {
|
||||
uint n_tok;
|
||||
uint h_k; // 16
|
||||
uint h_v; // 48
|
||||
uint channels; // 10240
|
||||
float eps; // 1e-6
|
||||
uint f32; // gdn_post: also write fp32 y
|
||||
};
|
||||
|
||||
static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t,
|
||||
int channels) {
|
||||
float s = 0.0f;
|
||||
for (int k = 0; k < 4; ++k) {
|
||||
const int src = t - 3 + k;
|
||||
const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)];
|
||||
s = fma(x, cw[ch * 4 + k], s);
|
||||
}
|
||||
s = s + 0.0f;
|
||||
return s / (1.0f + exp(-s));
|
||||
}
|
||||
|
||||
// grid (h_k + h_v / 2, T), 256 threads. x < h_k: q (threads 0..127) and k (128..255) of key head x; else value heads
|
||||
// 2 (x - h_k) and + 1.
|
||||
kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]],
|
||||
device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]],
|
||||
device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]],
|
||||
device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]],
|
||||
device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]],
|
||||
device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]],
|
||||
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 S = 128;
|
||||
threadgroup float part[8];
|
||||
threadgroup float prod[S];
|
||||
const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok);
|
||||
int ch;
|
||||
if (x < hk_n) ch = tid < 128 ? x * S + tid : hk_n * S + x * S + (tid - 128);
|
||||
else ch = 2 * hk_n * S + (2 * (x - hk_n) + (tid >> 7)) * S + (tid & 127);
|
||||
const float cv = gp_conv(qkv, hist, conv_w, ch, t, C);
|
||||
if (t == T - 1)
|
||||
for (int k = 0; k < 3; ++k) {
|
||||
const int src = T - 3 + k;
|
||||
hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)];
|
||||
}
|
||||
if (x < hk_n) {
|
||||
const float ssq = simd_sum(cv * cv);
|
||||
if (lane == 0) part[sg] = ssq;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const int base = tid < 128 ? 0 : 4;
|
||||
const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]);
|
||||
const float rs = rsqrt(tot / float(S) + a.eps / float(S));
|
||||
const float nv = (rs * cv) * (1.0f / sqrt(float(S)));
|
||||
if (tid < 128) {
|
||||
qn[(t * hk_n + x) * S + tid] = nv;
|
||||
prod[tid] = nv;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid >= 128) {
|
||||
kn[(t * hk_n + x) * S + tid - 128] = nv;
|
||||
const float p = simd_sum(nv * prod[tid - 128]);
|
||||
if (lane == 0) part[sg] = p; // simdgroups 4..7 (q's slots 0..3 are read already)
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid == 0) kq[t * hk_n + x] = (part[4] + part[5]) + (part[6] + part[7]);
|
||||
} else {
|
||||
const int h = 2 * (x - hk_n) + (tid >> 7);
|
||||
vc[(t * hv + h) * S + (tid & 127)] = cv;
|
||||
if ((tid & 127) == 0) {
|
||||
const float v = alpha[t * hv + h] + dt[h];
|
||||
const float e = exp(v), u = 1.0f + e;
|
||||
const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f));
|
||||
gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * hv + h])));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// grid (h_v * 4 / CB), CB * 32 threads; TB tokens of q and k staged at a time
|
||||
template <int CB, int TB>
|
||||
kernel void gdn_scan_t(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]],
|
||||
device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]],
|
||||
device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]],
|
||||
device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]],
|
||||
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 S = 128, NB = 4 / CB;
|
||||
threadgroup float4 qs[TB][32], ks[TB][32];
|
||||
const int h = int(tg.x) / NB, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k);
|
||||
const int j = ((int(tg.x) % NB) * CB + sg) * 32 + lane;
|
||||
const int T = int(a.n_tok);
|
||||
const float scale = 1.0f / sqrt(float(S));
|
||||
float st[S];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int i = 0; i < S; ++i) st[i] = state[(i * hv + h) * S + j];
|
||||
for (int t0 = 0; t0 < T; t0 += TB) {
|
||||
const int nb = min(TB, T - t0);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int idx = tid; idx < nb * 32; idx += CB * 32) {
|
||||
const int tt = idx >> 5, w = idx & 31;
|
||||
qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w];
|
||||
ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int tt = 0; tt < nb; ++tt) {
|
||||
const int t = t0 + tt;
|
||||
const float2 g_b = gb[t * hv + h];
|
||||
const float kqv = kq[t * hkn + hk];
|
||||
const float v = vc[(t * hv + h) * S + j];
|
||||
float ka = 0.0f, kb = 0.0f, qa = 0.0f, qb = 0.0f;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int w = 0; w < 32; w += 2) {
|
||||
const float4 k0 = ks[tt][w], k1 = ks[tt][w + 1], q0 = qs[tt][w], q1 = qs[tt][w + 1];
|
||||
ka = fma(st[4 * w], k0.x, ka); ka = fma(st[4 * w + 1], k0.y, ka);
|
||||
ka = fma(st[4 * w + 2], k0.z, ka); ka = fma(st[4 * w + 3], k0.w, ka);
|
||||
kb = fma(st[4 * w + 4], k1.x, kb); kb = fma(st[4 * w + 5], k1.y, kb);
|
||||
kb = fma(st[4 * w + 6], k1.z, kb); kb = fma(st[4 * w + 7], k1.w, kb);
|
||||
qa = fma(st[4 * w], q0.x, qa); qa = fma(st[4 * w + 1], q0.y, qa);
|
||||
qa = fma(st[4 * w + 2], q0.z, qa); qa = fma(st[4 * w + 3], q0.w, qa);
|
||||
qb = fma(st[4 * w + 4], q1.x, qb); qb = fma(st[4 * w + 5], q1.y, qb);
|
||||
qb = fma(st[4 * w + 6], q1.z, qb); qb = fma(st[4 * w + 7], q1.w, qb);
|
||||
}
|
||||
const float g = g_b.x, kv = ka + kb, qv = qa + qb;
|
||||
const float delta = (v - g * kv) * g_b.y;
|
||||
o[(t * hv + h) * S + j] = fma(g, qv, delta * kqv) * scale;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int w = 0; w < 32; ++w) {
|
||||
const float4 k4 = ks[tt][w];
|
||||
st[4 * w] = fma(g, st[4 * w], k4.x * delta);
|
||||
st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta);
|
||||
st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta);
|
||||
st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta);
|
||||
}
|
||||
}
|
||||
}
|
||||
#pragma clang loop unroll(full)
|
||||
for (int i = 0; i < S; ++i) state[(i * hv + h) * S + j] = st[i];
|
||||
}
|
||||
|
||||
// a probe for the bench (function constant 6): 1 = no per-token exchange (wrong results), 2 = no state update
|
||||
constant int GDN_PROBE [[function_constant(6)]];
|
||||
constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0;
|
||||
|
||||
// the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32
|
||||
// threads; lane = column (32 columns a threadgroup); per token the IB partial dots are exchanged through threadgroup
|
||||
// memory (double-buffered: one barrier a token)
|
||||
template <int IB, int TB>
|
||||
kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]],
|
||||
device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]],
|
||||
device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]],
|
||||
device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]],
|
||||
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 S = 128, R = S / IB, R4 = R / 4;
|
||||
threadgroup float4 qs[TB][32], ks[TB][32];
|
||||
threadgroup float2 red[2][IB][32];
|
||||
threadgroup float vs[TB][32];
|
||||
threadgroup float4 gk[TB]; // gate, beta, kq
|
||||
const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k);
|
||||
const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R;
|
||||
const int T = int(a.n_tok);
|
||||
const float scale = 1.0f / sqrt(float(S));
|
||||
float st[R];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int i = 0; i < R; ++i) st[i] = state[((i0 + i) * hv + h) * S + j];
|
||||
for (int t0 = 0; t0 < T; t0 += TB) {
|
||||
const int nb = min(TB, T - t0);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int idx = tid; idx < nb * 32; idx += IB * 32) {
|
||||
const int tt = idx >> 5, w = idx & 31;
|
||||
qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w];
|
||||
ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w];
|
||||
vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w];
|
||||
if (w == 0) {
|
||||
const float2 g_b = gb[(t0 + tt) * hv + h];
|
||||
gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f);
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int tt = 0; tt < nb; ++tt) {
|
||||
const int t = t0 + tt, slot = t & 1;
|
||||
float ka = 0.0f, qa = 0.0f;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int w = 0; w < R4; ++w) {
|
||||
const float4 k4 = ks[tt][i0 / 4 + w], q4 = qs[tt][i0 / 4 + w];
|
||||
ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka);
|
||||
ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka);
|
||||
qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa);
|
||||
qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa);
|
||||
}
|
||||
const float4 g_b = gk[tt];
|
||||
const float kqv = g_b.z;
|
||||
const float v = vs[tt][lane];
|
||||
float2 s = float2(ka, qa);
|
||||
if (gdn_probe != 1) {
|
||||
red[slot][sg][lane] = s;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
s = red[slot][0][lane];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int b = 1; b < IB; ++b) s += red[slot][b][lane];
|
||||
}
|
||||
const float g = g_b.x;
|
||||
const float delta = (v - g * s.x) * g_b.y;
|
||||
if (sg == 0) o[(t * hv + h) * S + j] = fma(g, s.y, delta * kqv) * scale;
|
||||
if (gdn_probe != 2)
|
||||
#pragma clang loop unroll(full)
|
||||
for (int w = 0; w < R4; ++w) {
|
||||
const float4 k4 = ks[tt][i0 / 4 + w];
|
||||
st[4 * w] = fma(g, st[4 * w], k4.x * delta);
|
||||
st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta);
|
||||
st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta);
|
||||
st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta);
|
||||
}
|
||||
}
|
||||
}
|
||||
#pragma clang loop unroll(full)
|
||||
for (int i = 0; i < R; ++i) state[((i0 + i) * hv + h) * S + j] = st[i];
|
||||
}
|
||||
|
||||
// gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check
|
||||
kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]],
|
||||
device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]],
|
||||
device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]],
|
||||
device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]],
|
||||
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 S = 128, IB = 8, TB = 16;
|
||||
threadgroup float4 qs[TB][32], ks[TB][32];
|
||||
threadgroup float2 red[2][IB][32];
|
||||
threadgroup float vs[TB][32];
|
||||
threadgroup float4 gk[TB];
|
||||
const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k);
|
||||
const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * 16, w0 = sg * 4;
|
||||
const int T = int(a.n_tok);
|
||||
const float scale = 1.0f / sqrt(float(S));
|
||||
device float* sp = state + (i0 * hv + h) * S + j;
|
||||
const int rs = hv * S;
|
||||
float4 s0 = float4(sp[0], sp[rs], sp[2 * rs], sp[3 * rs]);
|
||||
float4 s1 = float4(sp[4 * rs], sp[5 * rs], sp[6 * rs], sp[7 * rs]);
|
||||
float4 s2 = float4(sp[8 * rs], sp[9 * rs], sp[10 * rs], sp[11 * rs]);
|
||||
float4 s3 = float4(sp[12 * rs], sp[13 * rs], sp[14 * rs], sp[15 * rs]);
|
||||
for (int t0 = 0; t0 < T; t0 += TB) {
|
||||
const int nb = min(TB, T - t0);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int idx = tid; idx < nb * 32; idx += IB * 32) {
|
||||
const int tt = idx >> 5, w = idx & 31;
|
||||
qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w];
|
||||
ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w];
|
||||
vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w];
|
||||
if (w == 0) {
|
||||
const float2 g_b = gb[(t0 + tt) * hv + h];
|
||||
gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f);
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int tt = 0; tt < nb; ++tt) {
|
||||
const int t = t0 + tt, slot = t & 1;
|
||||
const float4 k0 = ks[tt][w0], k1 = ks[tt][w0 + 1], k2 = ks[tt][w0 + 2], k3 = ks[tt][w0 + 3];
|
||||
const float4 q0 = qs[tt][w0], q1 = qs[tt][w0 + 1], q2 = qs[tt][w0 + 2], q3 = qs[tt][w0 + 3];
|
||||
const float ka = (dot(s0, k0) + dot(s1, k1)) + (dot(s2, k2) + dot(s3, k3));
|
||||
const float qa = (dot(s0, q0) + dot(s1, q1)) + (dot(s2, q2) + dot(s3, q3));
|
||||
red[slot][sg][lane] = float2(ka, qa);
|
||||
const float4 g_b = gk[tt];
|
||||
const float v = vs[tt][lane];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
float2 r = red[slot][0][lane];
|
||||
for (int b = 1; b < IB; ++b) r += red[slot][b][lane];
|
||||
const float g = g_b.x;
|
||||
const float delta = (v - g * r.x) * g_b.y;
|
||||
if (sg == 0) o[(t * hv + h) * S + j] = fma(g, r.y, delta * g_b.z) * scale;
|
||||
s0 = fma(float4(g), s0, k0 * delta);
|
||||
s1 = fma(float4(g), s1, k1 * delta);
|
||||
s2 = fma(float4(g), s2, k2 * delta);
|
||||
s3 = fma(float4(g), s3, k3 * delta);
|
||||
}
|
||||
}
|
||||
sp[0] = s0.x; sp[rs] = s0.y; sp[2 * rs] = s0.z; sp[3 * rs] = s0.w;
|
||||
sp[4 * rs] = s1.x; sp[5 * rs] = s1.y; sp[6 * rs] = s1.z; sp[7 * rs] = s1.w;
|
||||
sp[8 * rs] = s2.x; sp[9 * rs] = s2.y; sp[10 * rs] = s2.z; sp[11 * rs] = s2.w;
|
||||
sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w;
|
||||
}
|
||||
|
||||
// grid (h_v, T), 128 threads
|
||||
kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]],
|
||||
device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]],
|
||||
constant GdnPArgs& a [[buffer(5)]],
|
||||
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 part[4];
|
||||
const int h = int(tg.x), t = int(tg.y), hv = int(a.h_v);
|
||||
const uint i = uint(t * hv + h) * 128 + tid;
|
||||
const float ov = o[i];
|
||||
const float ss = simd_sum(ov * ov);
|
||||
if (lane == 0) part[sg] = ss;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const float tot = (part[0] + part[1]) + (part[2] + part[3]);
|
||||
const float rs = rsqrt(tot / 128.0f + a.eps);
|
||||
const float yv = ((rs * ov) * gamma[tid]) * (1.0f / (1.0f + exp(-z[i])));
|
||||
yh[i] = half(yv);
|
||||
if (a.f32) y[i] = yv;
|
||||
}
|
||||
|
||||
template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>;
|
||||
template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>;
|
||||
template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>;
|
||||
template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 16>) gdn_scan_r<2, 16>;
|
||||
template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>;
|
||||
template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 16>) gdn_scan_r<8, 16>;
|
||||
template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>;
|
||||
@@ -25,6 +25,8 @@ kernel void mpp_mm(device T* A [[buffer(0)]], device T* B [[buffer(1)]], device
|
||||
#define MPP(T, M, N, S) template [[host_name("mpp_mm_" #T "_" #M "_" #N "_" #S)]] kernel decltype(mpp_mm<T, M, N, S>) mpp_mm<T, M, N, S>;
|
||||
MPP(half, 64, 32, 4) MPP(half, 64, 64, 4) MPP(half, 128, 64, 4) MPP(half, 32, 64, 4) MPP(half, 64, 128, 4)
|
||||
MPP(half, 16, 64, 2) MPP(half, 128, 128, 4)
|
||||
MPP(half, 64, 16, 2) MPP(half, 64, 16, 4) MPP(half, 128, 16, 4) MPP(half, 128, 32, 4) MPP(half, 64, 32, 2) MPP(half, 32, 16, 2)
|
||||
MPP(half, 32, 32, 2) MPP(half, 16, 32, 1) MPP(half, 16, 16, 1) MPP(half, 128, 48, 4) MPP(half, 64, 48, 4)
|
||||
MPP(float, 64, 32, 4) MPP(float, 64, 64, 4) MPP(float, 128, 64, 4)
|
||||
MPP(bfloat, 64, 64, 4) MPP(bfloat, 128, 64, 4)
|
||||
|
||||
@@ -49,3 +51,4 @@ 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>;
|
||||
|
||||
@@ -42,14 +42,18 @@ kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buf
|
||||
}
|
||||
|
||||
// one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e);
|
||||
// eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows,
|
||||
// meta[2..4] = an indirect grid (gx, n_tiles, 1) for the GEMMs, gx given in meta[2] by the host.
|
||||
// eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows;
|
||||
// indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the
|
||||
// down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather.
|
||||
kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]],
|
||||
device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]],
|
||||
constant MmPlanArgs& a [[buffer(5)]],
|
||||
ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]]) {
|
||||
threadgroup uint s_rows[32], s_tiles[32];
|
||||
// simdgroups past the launched ones must count 0 (512 threads fill only 16 of the 32 slots)
|
||||
if (tid < 32) { s_rows[tid] = 0; s_tiles[tid] = 0; }
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint e = tid;
|
||||
uint tot = 0;
|
||||
if (e < a.n_expert)
|
||||
@@ -75,8 +79,13 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b
|
||||
if (tid == 0) {
|
||||
meta[0] = all_tiles;
|
||||
meta[1] = all_rows;
|
||||
meta[3] = all_tiles;
|
||||
meta[3] = all_tiles; // indirect grid of the gate/up GEMM: (meta[2] given, tiles, 1)
|
||||
meta[4] = 1;
|
||||
meta[6] = all_tiles; // of the down GEMM: (meta[5] given, tiles, 1)
|
||||
meta[7] = 1;
|
||||
meta[8] = all_rows; // of the gather: (rows, 1, 1)
|
||||
meta[9] = 1;
|
||||
meta[10] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -260,6 +269,10 @@ kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt
|
||||
for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f);
|
||||
}
|
||||
|
||||
// a probe for the bench (function constant 5): 1 = the weights not decoded (a cheap fill), 2 = no matrix multiply
|
||||
constant int MM_PROBE [[function_constant(5)]];
|
||||
constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0;
|
||||
|
||||
template <int F, int TM, int TN, int TK, int SG>
|
||||
kernel void mm_gu2(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)]],
|
||||
@@ -286,12 +299,18 @@ kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[
|
||||
if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; }
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) {
|
||||
stage_w<F, TN, TK>(g0, a.gu_row, k0, wg, tab, tid, NTH);
|
||||
stage_w<F, TN, TK>(u0, a.gu_row, k0, wu, tab, tid, NTH);
|
||||
if (mm_probe == 1) {
|
||||
for (int it = tid; it < TN * TK; it += NTH) { wg[it] = half(k0 & 7); wu[it] = half(it & 7); }
|
||||
} else {
|
||||
stage_w<F, TN, TK>(g0, a.gu_row, k0, wg, tab, tid, NTH);
|
||||
stage_w<F, TN, TK>(u0, a.gu_row, k0, wu, tab, tid, NTH);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
tA = tX.template slice<TK, TM>(k0, row0);
|
||||
op.run(tA, tG, accg);
|
||||
op.run(tA, tU, accu);
|
||||
if (mm_probe != 2) {
|
||||
op.run(tA, tG, accg);
|
||||
op.run(tA, tU, accu);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
// metal/kernels/prompt.metal - the elementwise steps of a prompt chunk around the GEMMs (after quant/units.metal):
|
||||
//
|
||||
// route_topk per token: softmax over the 512 logits, top 10 (ties: the lower expert), weights renormalised
|
||||
// over the selected with a 2^-14 floor - moe_route's rule; one simdgroup a token
|
||||
// gr_epi_down after the HC down GEMM (V = xn . [w_down; w_inject], T x (hc_lr + hc)): lo_h = half(silu(v / hc)),
|
||||
// inject = v
|
||||
// gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and
|
||||
// half (the next GEMMs' operand)
|
||||
// swiglu_h h = half(silu(g) * u)
|
||||
// moe_out bo = y_routed + sigmoid(sg) * y_shared
|
||||
|
||||
kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
|
||||
device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]]) {
|
||||
const uint t = tg.x * 8 + sg;
|
||||
if (t >= n_tok) return;
|
||||
device const float* l = logits + ulong(t) * 512;
|
||||
float v[16];
|
||||
for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
|
||||
float mx = -INFINITY;
|
||||
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
|
||||
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
|
||||
float sum = 0.0f;
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
v[i] = exp(v[i] - mx);
|
||||
sum += v[i];
|
||||
}
|
||||
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
|
||||
const float rcp = 1.0f / sum;
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
v[i] *= rcp;
|
||||
if (isnan(v[i])) v[i] = -FLT_MAX;
|
||||
}
|
||||
float selected = 0.0f, ssum = 0.0f;
|
||||
for (int rank = 0; rank < 10; ++rank) {
|
||||
float best = v[0];
|
||||
int ex = lane;
|
||||
for (int i = 1; i < 16; ++i)
|
||||
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
|
||||
for (ushort m = 16; m > 0; m >>= 1) {
|
||||
const float ob = simd_shuffle_xor(best, m);
|
||||
const int oe = simd_shuffle_xor(ex, m);
|
||||
if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
|
||||
}
|
||||
if ((ex & 31) == lane) {
|
||||
v[ex / 32] = -INFINITY;
|
||||
ids[t * 10 + rank] = ex;
|
||||
ssum += best;
|
||||
}
|
||||
if (rank == lane) selected = best;
|
||||
}
|
||||
for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m);
|
||||
ssum = max(ssum, 6.103515625e-5f);
|
||||
if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum);
|
||||
}
|
||||
|
||||
struct EpiArgs {
|
||||
uint n_tok;
|
||||
uint n_embd; // 2560
|
||||
uint hc; // 4
|
||||
uint hc_lr; // 320
|
||||
};
|
||||
|
||||
// grid (T * (hc_lr + hc) / 256)
|
||||
kernel void gr_epi_down(device const float* V [[buffer(0)]], device half* lo_h [[buffer(1)]],
|
||||
device float* inject [[buffer(2)]], constant EpiArgs& a [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint nr = a.hc_lr + a.hc, t = gid / nr, r = gid % nr;
|
||||
if (t >= a.n_tok) return;
|
||||
const float v = V[gid];
|
||||
if (r < a.hc_lr) {
|
||||
const float z = v / float(a.hc);
|
||||
lo_h[t * a.hc_lr + r] = half(z / (1.0f + exp(-z)));
|
||||
} else {
|
||||
inject[t * a.hc + r - a.hc_lr] = v;
|
||||
}
|
||||
}
|
||||
|
||||
// grid (T * n_embd / 256)
|
||||
kernel void gr_epi_up(device const float* G [[buffer(0)]], device const float* xn [[buffer(1)]],
|
||||
device float* mixed [[buffer(2)]], device half* mixed_h [[buffer(3)]],
|
||||
constant EpiArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) {
|
||||
const uint n = a.n_embd, t = gid / n, d = gid % n;
|
||||
if (t >= a.n_tok) return;
|
||||
float m = 0.0f;
|
||||
for (uint c = 0; c < a.hc; ++c) {
|
||||
const float g = G[ulong(t) * a.hc * n + c * n + d];
|
||||
m = fma(xn[ulong(t) * a.hc * n + c * n + d], 1.0f / (1.0f + exp(-g)), m);
|
||||
}
|
||||
m = m / float(a.hc);
|
||||
mixed[gid] = m;
|
||||
mixed_h[gid] = half(m);
|
||||
}
|
||||
|
||||
kernel void swiglu_h(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device half* h [[buffer(2)]],
|
||||
constant uint& n [[buffer(3)]], uint gid [[thread_position_in_grid]]) {
|
||||
if (gid >= n) return;
|
||||
const float gv = g[gid];
|
||||
h[gid] = half(gv / (1.0f + exp(-gv)) * u[gid]);
|
||||
}
|
||||
|
||||
// grid (T * n_embd / 256); sg: T values
|
||||
kernel void moe_out(device const float* yr [[buffer(0)]], device const float* ys [[buffer(1)]],
|
||||
device const float* sg [[buffer(2)]], device float* bo [[buffer(3)]], constant uint2& a [[buffer(4)]],
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint t = gid / a.y;
|
||||
if (t >= a.x) return;
|
||||
bo[gid] = fma(1.0f / (1.0f + exp(-sg[t])), ys[gid], yr[gid]);
|
||||
}
|
||||
|
||||
// ---- the HC read with fewer passes (after dense_mm.metal: GemmArgs, MetalPerformancePrimitives) ---------------------
|
||||
// gr_norm_h: gr_norm with xn written as fp16 only (the GEMM operand; gemm_hc_up's epilogue reads it too).
|
||||
// one threadgroup per (stream c, token t), 256 threads; args as gr_norm's (GrArgs, gr.metal)
|
||||
kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]],
|
||||
device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]],
|
||||
device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]],
|
||||
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 part[8];
|
||||
const uint c = tg.x, t = tg.y, n = a.n_embd;
|
||||
device float* r = R + (t * a.hc + c) * n;
|
||||
const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f;
|
||||
float ss = 0.0f;
|
||||
for (uint d = tid; d < n; d += 256) {
|
||||
float v = r[d];
|
||||
if (a.apply) {
|
||||
v = fma(bo[t * n + d], g, v);
|
||||
r[d] = v;
|
||||
}
|
||||
ss = fma(v, v, ss);
|
||||
}
|
||||
ss = simd_sum(ss);
|
||||
if (lane == 0) part[sg] = ss;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
float tot = 0.0f;
|
||||
for (int i = 0; i < 8; ++i) tot += part[i];
|
||||
const float rs = rsqrt(tot / float(n) + a.eps);
|
||||
device half* o = xnh + (t * a.hc + c) * n;
|
||||
device const float* w = w_norm + c * n;
|
||||
for (uint d = tid; d < n; d += 256) o[d] = half(r[d] * w[d] * rs);
|
||||
}
|
||||
|
||||
// gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column
|
||||
// tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] *
|
||||
// sigmoid(G[t][d * hc + c]) -> fp32 and fp16. G is never written. grid (hc * n / TN, ceil(T / TM)), 32 * SG threads.
|
||||
template <int TM, int TN, int SG>
|
||||
kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]],
|
||||
device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]],
|
||||
constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]],
|
||||
ushort tid [[thread_index_in_threadgroup]]) {
|
||||
threadgroup float ot[TM * TN];
|
||||
tensor<device half, dextents<int32_t, 2>, tensor_inline> tA(A, dextents<int32_t, 2>(g.K, g.M));
|
||||
tensor<device half, dextents<int32_t, 2>, tensor_inline> tB(B, dextents<int32_t, 2>(g.K, g.N));
|
||||
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot, dextents<int32_t, 2>(TN, TM));
|
||||
constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>, false, true, false);
|
||||
matmul2d<desc, execution_simdgroups<SG>> op;
|
||||
auto mA = tA.slice(0, int(tg.y) * TM);
|
||||
auto mB = tB.slice(0, int(tg.x) * TN);
|
||||
op.run(mA, mB, tO);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc;
|
||||
const int t0 = int(tg.y) * TM, d0 = int(tg.x) * DPT;
|
||||
for (int idx = tid; idx < TM * DPT; idx += SG * 32) {
|
||||
const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m;
|
||||
if (t >= int(e.n_tok)) continue;
|
||||
float s = 0.0f;
|
||||
for (int c = 0; c < hc; ++c)
|
||||
s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s);
|
||||
s = s / float(hc);
|
||||
mixed[ulong(t) * n + d0 + dd] = s;
|
||||
mixh[ulong(t) * n + d0 + dd] = half(s);
|
||||
}
|
||||
}
|
||||
template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>;
|
||||
template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>;
|
||||
template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>;
|
||||
@@ -14,7 +14,7 @@
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
struct DqArgs { uint32_t n_in, n_out, row_bytes; };
|
||||
struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; };
|
||||
struct GemmArgs { int M, N, K, pad; };
|
||||
double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; }
|
||||
float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); }
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
// metal/tests/gdn_prompt_parity.cpp - the prompt GDN passes (gdn_prompt.metal: gdn_pre -> gdn_scan -> gdn_post)
|
||||
// against the decode step run over the same T tokens in order (gdn.metal's gdn_step, GDN_NTOK = T, itself checked
|
||||
// against the CPU reference by gdn_parity): y, the final state and the conv history. The scan reassociates
|
||||
// o = S_new^T q as g S^T q + delta (k . q), so the match is to rounding, not bitwise. Then the speed for T = 2048.
|
||||
// gdn_prompt_parity [--bench]
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
constexpr int HK = 16, HV = 48, CH = 10240, S = 128;
|
||||
struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; };
|
||||
struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; };
|
||||
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";
|
||||
mtl::Context cx;
|
||||
const int lib = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
std::printf("device %s\n", cx.device_name().c_str());
|
||||
const int TMAX = 2048;
|
||||
auto F = [&](size_t n) { return cx.alloc(4 * n); };
|
||||
mtl::Buffer qkv = F((size_t)TMAX * CH), al = F(TMAX * HV), be = F(TMAX * HV), z = F((size_t)TMAX * HV * S), cw = F(CH * 4),
|
||||
dt = F(HV), sa = F(HV), gam = F(S), st0 = F(S * HV * S), stA = F(S * HV * S), stB = F(S * HV * S),
|
||||
h0 = F(CH * 3), hA = F(CH * 3), hA2 = F(CH * 3), hB = F(CH * 3), yA = F((size_t)TMAX * HV * S),
|
||||
yB = F((size_t)TMAX * HV * S), yh = cx.alloc(2 * (size_t)TMAX * HV * S), qn = F((size_t)TMAX * HK * S),
|
||||
kn = F((size_t)TMAX * HK * S), vc = F((size_t)TMAX * HV * S), gb = F(2 * TMAX * HV), kq = F(TMAX * HK),
|
||||
o = F((size_t)TMAX * HV * S);
|
||||
ref::Rng r(21);
|
||||
for (size_t i = 0; i < (size_t)TMAX * CH; ++i) qkv.as<float>()[i] = r.sym() * 2.0f;
|
||||
for (int i = 0; i < TMAX * HV; ++i) { al.as<float>()[i] = r.sym() * 3.0f; be.as<float>()[i] = r.sym() * 2.0f; }
|
||||
for (size_t i = 0; i < (size_t)TMAX * HV * S; ++i) z.as<float>()[i] = r.sym() * 2.0f;
|
||||
for (int i = 0; i < CH * 4; ++i) cw.as<float>()[i] = r.sym() * 0.6f;
|
||||
for (int i = 0; i < HV; ++i) { dt.as<float>()[i] = r.sym(); sa.as<float>()[i] = -0.2f - r.uni() * 2.0f; }
|
||||
for (int i = 0; i < S; ++i) gam.as<float>()[i] = 0.5f + r.uni();
|
||||
for (int i = 0; i < S * HV * S; ++i) st0.as<float>()[i] = r.sym() * 0.1f;
|
||||
for (int i = 0; i < CH * 3; ++i) h0.as<float>()[i] = r.sym();
|
||||
auto prompt_ds = [&](int T, const mtl::Pipeline& pre, const mtl::Pipeline& scan, const mtl::Pipeline& post, int cb,
|
||||
std::vector<mtl::Dispatch>& ds) {
|
||||
GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u};
|
||||
mtl::Dispatch d1;
|
||||
d1.p = ⪯
|
||||
d1.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt),
|
||||
mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq),
|
||||
mtl::Arg::B(hA2), mtl::Arg::V(a)};
|
||||
d1.grid[0] = HK + HV / 2;
|
||||
d1.grid[1] = T;
|
||||
d1.tg[0] = 256;
|
||||
ds.push_back(d1);
|
||||
mtl::Dispatch d2;
|
||||
d2.p = &scan;
|
||||
d2.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq),
|
||||
mtl::Arg::B(o), mtl::Arg::V(a)};
|
||||
// cb > 0: gdn_scan_<cb> (cb column blocks of 32 a threadgroup); cb < 0: gdn_scanr_<-cb> (-cb row blocks)
|
||||
d2.grid[0] = cb > 0 ? HV * 4 / cb : HV * 4;
|
||||
d2.tg[0] = 32 * (cb > 0 ? cb : -cb);
|
||||
ds.push_back(d2);
|
||||
mtl::Dispatch d3;
|
||||
d3.p = &post;
|
||||
d3.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)};
|
||||
d3.grid[0] = HV;
|
||||
d3.grid[1] = T;
|
||||
d3.tg[0] = 128;
|
||||
ds.push_back(d3);
|
||||
};
|
||||
mtl::Pipeline pre = cx.pipeline(lib, "gdn_pre"), post = cx.pipeline(lib, "gdn_post");
|
||||
int fails = 0;
|
||||
if (!bench)
|
||||
for (int T : {1, 3, 37, 300})
|
||||
for (int cb : {1, 4, -2, -4, -8, -9}) {
|
||||
mtl::Constants pc;
|
||||
if (cb < 0 && cb != -9) pc.i(6, 0);
|
||||
mtl::Pipeline scan = cx.pipeline(lib, cb == -9 ? std::string("gdn_scan_v") : cb > 0 ? "gdn_scan_" + std::to_string(cb) : "gdn_scanr_" + std::to_string(-cb), pc);
|
||||
if (cb == -9) cb = -8;
|
||||
mtl::Pipeline step = cx.pipeline(lib, "gdn_step", mtl::Constants().i(2, T));
|
||||
std::copy(st0.as<float>(), st0.as<float>() + S * HV * S, stA.as<float>());
|
||||
std::copy(st0.as<float>(), st0.as<float>() + S * HV * S, stB.as<float>());
|
||||
std::copy(h0.as<float>(), h0.as<float>() + CH * 3, hA.as<float>());
|
||||
std::copy(h0.as<float>(), h0.as<float>() + CH * 3, hB.as<float>());
|
||||
GdnArgs da{HK, HV, CH, 1, 1e-6f, 0};
|
||||
mtl::Dispatch d;
|
||||
d.p = &step;
|
||||
d.args = {mtl::Arg::B(stA), mtl::Arg::B(hA), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(z),
|
||||
mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gam), mtl::Arg::B(yA), mtl::Arg::V(da),
|
||||
mtl::Arg::B(hA2)};
|
||||
d.grid[0] = HV;
|
||||
d.tg[0] = 512;
|
||||
cx.run({d});
|
||||
std::vector<float> hdec(hA2.as<float>(), hA2.as<float>() + CH * 3);
|
||||
std::vector<mtl::Dispatch> ds;
|
||||
prompt_ds(T, pre, scan, post, cb, ds);
|
||||
cx.run(ds);
|
||||
auto rel = [](const float* a, const float* b, size_t n) {
|
||||
double se = 0, sr = 0;
|
||||
for (size_t i = 0; i < n; ++i) { se += (double)(a[i] - b[i]) * (a[i] - b[i]); sr += (double)b[i] * b[i]; }
|
||||
return std::sqrt(se / std::max(sr, 1e-30));
|
||||
};
|
||||
const double ry = rel(yB.as<float>(), yA.as<float>(), (size_t)T * HV * S),
|
||||
rs = rel(stB.as<float>(), stA.as<float>(), S * HV * S), rh = rel(hA2.as<float>(), hdec.data(), CH * 3);
|
||||
double ryh = 0, syh = 0;
|
||||
for (size_t i = 0; i < (size_t)T * HV * S; ++i) {
|
||||
const double e = ref::f16_to_f32(yh.as<uint16_t>()[i]) - yA.as<float>()[i];
|
||||
ryh += e * e;
|
||||
syh += (double)yA.as<float>()[i] * yA.as<float>()[i];
|
||||
}
|
||||
const bool ok = ry < 1e-5 && rs < 1e-5 && rh == 0 && std::sqrt(ryh / syh) < 1e-3;
|
||||
fails += !ok;
|
||||
std::printf("T %3d, scan CB %d: y rel %.2e (fp16 %.2e), state rel %.2e, history %s %s\n", T, cb, ry,
|
||||
std::sqrt(ryh / syh), rs, rh == 0 ? "exact" : "DIFFERS", ok ? "PASS" : "FAIL");
|
||||
}
|
||||
if (fails || !bench) {
|
||||
std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS");
|
||||
return fails ? 1 : 0;
|
||||
}
|
||||
// speed, T = 2048: the decode step over the chunk vs the three passes (each scan width)
|
||||
const int T = 2048;
|
||||
mtl::Pipeline step = cx.pipeline(lib, "gdn_step", mtl::Constants().i(2, T));
|
||||
GdnArgs da{HK, HV, CH, 1, 1e-6f, 0};
|
||||
mtl::Dispatch d;
|
||||
d.p = &step;
|
||||
d.args = {mtl::Arg::B(stA), mtl::Arg::B(hA), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(z),
|
||||
mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gam), mtl::Arg::B(yA), mtl::Arg::V(da), mtl::Arg::B(hA2)};
|
||||
d.grid[0] = HV;
|
||||
d.tg[0] = 512;
|
||||
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("T %d: decode step over the chunk %.2f ms\n", T, timeit({d}) * 1e3);
|
||||
for (const char* nm : {"gdn_scan_v", "gdn_scanr_4", "gdn_scanr_8", "gdn_scan_v", "gdn_scanr_8"}) {
|
||||
const int cb = nm[9] == 'v' ? -8 : nm[8] == 'r' ? -(nm[10] - '0') : nm[9] - '0';
|
||||
mtl::Constants pc;
|
||||
if (cb < 0 && nm[9] != 'v') pc.i(6, std::getenv("GDN_PROBE") ? std::atoi(std::getenv("GDN_PROBE")) : 0);
|
||||
mtl::Pipeline scan = cx.pipeline(lib, nm, pc);
|
||||
std::vector<mtl::Dispatch> ds;
|
||||
prompt_ds(T, pre, scan, post, cb, ds);
|
||||
const double all = timeit(ds), tp = timeit({ds[0]}), ts = timeit({ds[1]}), to = timeit({ds[2]});
|
||||
std::printf("T %d, %-13s: pre %.3f + scan %.3f + post %.3f ms = %.3f ms together\n", T, nm, tp * 1e3, ts * 1e3,
|
||||
to * 1e3, all * 1e3);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -7,6 +7,7 @@
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
@@ -40,15 +41,27 @@ int main() {
|
||||
{"mpp_mm_half_16_64_2", 2, 16, 64, 2, true}, {"mpp_mm_bfloat_64_64_4", 3, 64, 64, 4, true},
|
||||
{"mpp_mm_bfloat_128_64_4", 3, 128, 64, 4, true}, {"mpp_mm_float_64_32_4", 4, 64, 32, 4, true},
|
||||
{"mpp_mm_float_64_64_4", 4, 64, 64, 4, true}, {"mpp_mm_float_128_64_4", 4, 128, 64, 4, true},
|
||||
{"sg_mm_half", 2, 64, 64, 4, false}, {"sg_mm_float", 4, 64, 64, 4, false}};
|
||||
{"sg_mm_half", 2, 64, 64, 4, false}, {"sg_mm_float", 4, 64, 64, 4, false},
|
||||
{"mpp_mm_half_64_16_2", 2, 64, 16, 2, true}, {"mpp_mm_half_64_16_4", 2, 64, 16, 4, true},
|
||||
{"mpp_mm_half_128_16_4", 2, 128, 16, 4, true}, {"mpp_mm_half_128_32_4", 2, 128, 32, 4, true},
|
||||
{"mpp_mm_half_64_32_2", 2, 64, 32, 2, true}, {"mpp_mm_half_32_16_2", 2, 32, 16, 2, true},
|
||||
{"mpp_mm_half_32_32_2", 2, 32, 32, 2, true}, {"mpp_mm_half_16_32_1", 2, 16, 32, 1, true},
|
||||
{"mpp_mm_half_16_16_1", 2, 16, 16, 1, true}, {"mpp_mm_half_128_48_4", 2, 128, 48, 4, true},
|
||||
{"mpp_mm_half_64_48_4", 2, 64, 48, 4, true}};
|
||||
struct Shape { const char* name; int M, N, K; };
|
||||
const std::vector<Shape> shapes = {{"square 4096^3", 4096, 4096, 4096},
|
||||
{"expert gu 128 tok x 1280 x 2560", 128, 1280, 2560},
|
||||
{"expert down 128 tok x 2560 x 640", 128, 2560, 640},
|
||||
{"dense qkv 2048 tok x 10240 x 2560", 2048, 4096, 2560},
|
||||
{"expert gu 16 tok x 1280 x 2560", 16, 1280, 2560}};
|
||||
{"expert gu 16 tok x 1280 x 2560", 16, 1280, 2560},
|
||||
{"expert gu^T 1280 x 48 tok x 2560", 1280, 48, 2560},
|
||||
{"expert gu^T 1280 x 16 tok x 2560", 1280, 16, 2560},
|
||||
{"many gu^T 12800 x 48 tok x 2560", 12800, 48, 2560},
|
||||
{"many gu 48 tok x 12800 x 2560", 48, 12800, 2560}};
|
||||
int filled = 0;
|
||||
const char* only = std::getenv("MMB_ONLY"); // a substring of the kernel names to run
|
||||
for (const K& k : ks) {
|
||||
if (only && k.name.find(only) == std::string::npos) continue;
|
||||
if (filled != k.elem) { fill(k.elem); filled = k.elem; }
|
||||
mtl::Pipeline p = cx.pipeline(lib, k.name);
|
||||
// correctness on a small product (M = tm*2, N = tn*2, K = 256) against the CPU
|
||||
|
||||
@@ -87,7 +87,9 @@ int main(int argc, char** argv) {
|
||||
if (dn == ref::IQ4_NL && DTK < 32) DTK = 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::Pipeline pgu = cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl),
|
||||
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)),
|
||||
pgat = cx.pipeline(lib, "mm_gather");
|
||||
MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk};
|
||||
|
||||
@@ -0,0 +1,500 @@
|
||||
// metal/tests/prompt_chain.cpp - a whole PROMPT through every layer on the REAL weights of a GGUF (mmapped, wrapped
|
||||
// no-copy), in chunks of C tokens (one command buffer a chunk): the GPU time of prompt processing with the prompt
|
||||
// kernels of metal/. A timing of the real byte and FLOP stream, not a numerical check: the tokens are random ids
|
||||
// (their real embedding rows start every hyper-connection stream), the parity programs check the kernels.
|
||||
//
|
||||
// prompt_chain <shard-1.gguf> [--prompt P,P,..] [--chunk C] [--breakdown] [--only CLASS] [--skip CLASS]
|
||||
//
|
||||
// Per layer and chunk (T tokens; activations fp32 rows, fp16 copies as the matrix units' operands):
|
||||
// HC read gr_norm_h (fp16 xn) -> GEMM [w_down; w_inject] -> gr_epi_down -> gemm_hc_up (the GEMM with the
|
||||
// mixing epilogue); PC_HC=plain: gr_norm -> act_f16 -> GEMM -> gr_epi_down -> GEMM -> gr_epi_up
|
||||
// GDN GEMMs qkv / gate / alpha / beta -> gdn_pre (parallel) -> gdn_scan_r (the state scan) -> gdn_post
|
||||
// (output norm and gate, fp16) -> GEMM out; PC_GDN=step: the decode step over the T tokens instead
|
||||
// QSA GEMMs q / k / v / indexer q, k -> qsa_prep (T tokens) -> score + top-k (past the budget) -> attn ->
|
||||
// merge -> GEMM out
|
||||
// MoE GEMM router -> route_topk -> plan (hist, scan, clear, place) -> gather -> mm_gu2 -> mm_down2 ->
|
||||
// combine; shared expert GEMMs gate, up -> swiglu_h -> GEMM down; GEMM shared gate; moe_out
|
||||
// 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 "gguf_lite.h"
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <sys/mman.h>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, NV = 248320, QB = 2051, QCHUNKS = 33;
|
||||
struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; };
|
||||
struct EpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; };
|
||||
struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; };
|
||||
struct GemmArgs { int M, N, K, pad; };
|
||||
struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; };
|
||||
struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; };
|
||||
struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; };
|
||||
struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; };
|
||||
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; };
|
||||
struct ArgArgs { uint32_t n_vocab, parts; };
|
||||
|
||||
double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; }
|
||||
int unit_of(int t) { return t == ref::Q2_0 || t == ref::IQ3_S || t == ref::BF16 ? 16 : 32; }
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
if (argc < 2) { std::fprintf(stderr, "usage: prompt_chain <shard-1.gguf> [--prompt P,..] [--chunk C] [--breakdown]\n"); return 2; }
|
||||
std::vector<int> prompts = {1024, 4096, 32768};
|
||||
int C = 2048, reps = 2;
|
||||
bool breakdown = false;
|
||||
std::string only, skip;
|
||||
for (int i = 2; i < argc; ++i) {
|
||||
std::string a = argv[i];
|
||||
if (a == "--prompt") {
|
||||
prompts.clear();
|
||||
for (char* s = argv[++i]; *s;) {
|
||||
prompts.push_back((int)std::strtol(s, &s, 10));
|
||||
if (*s == ',') ++s;
|
||||
}
|
||||
} else if (a == "--chunk") C = std::atoi(argv[++i]);
|
||||
else if (a == "--reps") reps = std::atoi(argv[++i]);
|
||||
else if (a == "--breakdown") breakdown = true;
|
||||
else if (a == "--only") only = argv[++i];
|
||||
else if (a == "--skip") skip = argv[++i];
|
||||
}
|
||||
const int PMAX = *std::max_element(prompts.begin(), prompts.end());
|
||||
gguf::File f;
|
||||
std::string err;
|
||||
if (!f.open(argv[1], err)) { std::fprintf(stderr, "%s\n", err.c_str()); return 2; }
|
||||
mtl::Context cx;
|
||||
std::printf("device %s; %s: %zu tensors, %.1f GB; chunk %d\n", cx.device_name().c_str(), argv[1], f.tensors.size(),
|
||||
f.file_bytes / 1e9, C);
|
||||
{
|
||||
const double t0 = mtl::now_s();
|
||||
madvise(f.map, f.map_bytes, MADV_WILLNEED);
|
||||
volatile uint64_t s = 0;
|
||||
for (size_t o = 0; o < f.file_bytes; o += 16384) s += ((const uint8_t*)f.map)[o];
|
||||
std::printf("file touched in %.1f s\n", mtl::now_s() - t0);
|
||||
}
|
||||
mtl::Buffer W = cx.wrap(f.map, f.map_bytes);
|
||||
auto T_ = [&](const std::string& name) -> const gguf::Tensor& {
|
||||
const gguf::Tensor* t = f.find(name);
|
||||
if (!t) { std::fprintf(stderr, "missing tensor %s\n", name.c_str()); std::exit(2); }
|
||||
return *t;
|
||||
};
|
||||
const double tc0 = mtl::now_s();
|
||||
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 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;
|
||||
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);
|
||||
if (it == pcache.end()) it = pcache.emplace(k, cx.pipeline(l, fn, c)).first;
|
||||
return &it->second;
|
||||
};
|
||||
|
||||
// ---- buffers (MT = the chunk's rows)
|
||||
const size_t MT = C;
|
||||
const int MAXC = (PMAX + 64 + 63) / 64 * 64, MB = MAXC / 4 + 2;
|
||||
const int TM = 16;
|
||||
const int maxrows = (int)MT * TOPK + NEXP * TM, nblk_max = ((int)MT * TOPK + 1023) / 1024;
|
||||
auto F32 = [&](size_t n) { return cx.alloc(4 * n); };
|
||||
auto F16 = [&](size_t n) { return cx.alloc(2 * n); };
|
||||
mtl::Buffer R = F32(MT * HC * N), bo = F32(MT * N), inj[2] = {F32(MT * HC), F32(MT * HC)}, xn = F32(MT * HC * N),
|
||||
xnh = F16(MT * HC * N), V = F32(MT * (LR + HC)), loh = F16(MT * LR), G = F32(MT * HC * N), mixed = F32(MT * N),
|
||||
mixh = F16(MT * N);
|
||||
mtl::Buffer qkv = F32(MT * 10240), z = F32(MT * 6144), al = F32(MT * 48), be = F32(MT * 48), yg = F32(MT * 6144),
|
||||
ygh = F16(MT * 6144), gqn = F32(MT * 16 * 128), gkn = F32(MT * 16 * 128), gvc = F32(MT * 6144),
|
||||
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* 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),
|
||||
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),
|
||||
tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows),
|
||||
sw = F32(maxrows), pos = cx.alloc(4 * MT * TOPK), XS = F16((size_t)maxrows * N), HS = F16((size_t)maxrows * NFF),
|
||||
YE = F32((size_t)maxrows * N), yr = F32(MT * N), gsh = F32(MT * NFF), ush = F32(MT * NFF), hh = F16(MT * NFF),
|
||||
ys = F32(MT * N), sgv = F32(MT);
|
||||
// fp16 weight scratch, one per role
|
||||
mtl::Buffer WD = F16((size_t)(LR + HC) * HC * N), WU = F16((size_t)HC * N * LR), Wqkv = F16((size_t)10240 * N),
|
||||
Wz = F16((size_t)6144 * N), Wal = F16((size_t)48 * N), Wbe = F16((size_t)48 * N), Wout = F16((size_t)N * 6144),
|
||||
Wq = F16((size_t)12288 * N), Wk = F16((size_t)512 * N), Wv = F16((size_t)512 * N), Wiq = F16((size_t)512 * N),
|
||||
Wik = F16((size_t)128 * N), Wr = F16((size_t)NEXP * N), Wsg = F16((size_t)N), Wshg = F16((size_t)NFF * N),
|
||||
Wshu = F16((size_t)NFF * N), Wshd = F16((size_t)N * NFF);
|
||||
// the last row's final read and head
|
||||
mtl::Buffer xn1 = F32(HC * N), lo1 = F32(LR), inj1 = F32(HC), X1 = F32(4 * N + N / 16), lm = F32(NV), apv = F32(120),
|
||||
api = cx.alloc(4 * 120), tokout = cx.alloc(16);
|
||||
mtl::Buffer cs = cx.alloc(8 * (size_t)(MAXC + 16) * 32);
|
||||
for (int p = 0; p < MAXC + 16; ++p)
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
const double th = (double)p * std::pow(1e7, -2.0 * i / 64.0);
|
||||
cs.as<float>()[(p * 32 + i) * 2] = (float)std::cos(th);
|
||||
cs.as<float>()[(p * 32 + i) * 2 + 1] = (float)std::sin(th);
|
||||
}
|
||||
std::vector<mtl::Buffer> gst, ghist[2], kq, vq, ksc, vsc, tail, dead, pooled;
|
||||
for (int l = 0; l < 48; ++l) {
|
||||
if (l % 4 != 3) {
|
||||
gst.push_back(F32(128 * 48 * 128));
|
||||
ghist[0].push_back(F32(10240 * 3));
|
||||
ghist[1].push_back(F32(10240 * 3));
|
||||
} else {
|
||||
kq.push_back(cx.alloc((size_t)MAXC * 2 * 256));
|
||||
vq.push_back(cx.alloc((size_t)MAXC * 2 * 256));
|
||||
ksc.push_back(cx.alloc(2 * (size_t)MAXC * 2 * 4));
|
||||
vsc.push_back(cx.alloc(2 * (size_t)MAXC * 2 * 4));
|
||||
tail.push_back(F32(4 * 128));
|
||||
dead.push_back(F32(128));
|
||||
pooled.push_back(F32((size_t)MB * 128));
|
||||
}
|
||||
}
|
||||
std::vector<int> gpar(36, 0);
|
||||
// embeddings of random token ids (the real rows)
|
||||
std::vector<float> E(MT * N);
|
||||
{
|
||||
const gguf::Tensor& te = T_("token_embd.weight");
|
||||
const size_t rb = ref::row_bytes(te.type, N);
|
||||
ref::Rng r(17);
|
||||
for (size_t t = 0; t < MT; ++t) {
|
||||
const uint64_t id = r.next() % 200000;
|
||||
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;
|
||||
|
||||
// ---- the chunk's dispatches
|
||||
std::vector<mtl::Dispatch> ds;
|
||||
std::vector<std::string> cls;
|
||||
auto push = [&](const mtl::Dispatch& d, const std::string& c) {
|
||||
if (!only.empty() && c != only) return;
|
||||
if (!skip.empty() && c == skip) return;
|
||||
ds.push_back(d);
|
||||
cls.push_back(c);
|
||||
};
|
||||
auto D = [&](const mtl::Pipeline* p, std::vector<mtl::Arg> a, uint32_t gx, uint32_t gy, uint32_t tg, const std::string& c,
|
||||
uint32_t gz = 1) {
|
||||
mtl::Dispatch d;
|
||||
d.p = p;
|
||||
d.args = std::move(a);
|
||||
d.grid[0] = gx;
|
||||
d.grid[1] = gy;
|
||||
d.grid[2] = gz;
|
||||
d.tg[0] = tg;
|
||||
push(d, c);
|
||||
};
|
||||
const char* gk_env = std::getenv("PC_GEMM");
|
||||
auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K, const std::string& c) {
|
||||
int tm = 64, tn = Nn <= 64 ? 32 : 64;
|
||||
if (gk_env && Nn > 64) std::sscanf(gk_env, "%d,%d", &tm, &tn);
|
||||
const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4";
|
||||
GemmArgs g{M, Nn, K, 0};
|
||||
D(P(lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm,
|
||||
128, c);
|
||||
};
|
||||
auto dq = [&](const gguf::Tensor& t, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) {
|
||||
const int n_in = (int)t.shape[0], n_out = t.shape.size() > 1 ? (int)t.shape[1] : 1;
|
||||
DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ref::row_bytes(t.type, n_in), perm};
|
||||
D(P(lib, "dq_f16_" + std::to_string(t.type)), {mtl::Arg::B(W, t.offset), mtl::Arg::B(out, out_off), mtl::Arg::V(a)},
|
||||
(n_in / unit_of(t.type) + 63) / 64, n_out, 64, "dq");
|
||||
};
|
||||
auto act16 = [&](const mtl::Buffer& x, const mtl::Buffer& xh, size_t n, const std::string& c) {
|
||||
const uint32_t nn = (uint32_t)n;
|
||||
D(P(lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256, c);
|
||||
};
|
||||
int cur_inj = 0;
|
||||
// the HC read of T rows: mixed (fp32) and mixh (fp16)
|
||||
auto hc_read = [&](const std::string& pre, bool apply, int T) {
|
||||
const gguf::Tensor &tn = T_(pre + "_norm.weight"), &tdn = T_(pre + "_down.weight"), &tin = T_(pre + "_inject.weight"),
|
||||
&tup = T_(pre + "_up.weight");
|
||||
dq(tdn, WD);
|
||||
dq(tin, WD, (size_t)2 * LR * HC * N);
|
||||
GrArgs a{N, HC, LR, (uint32_t)T, apply ? 1u : 0u, 1e-6f};
|
||||
EpiArgs ea{(uint32_t)T, N, HC, LR};
|
||||
if (hc_mode == "fused") {
|
||||
// fp16 xn only; the up GEMM's epilogue mixes (w_up rows interleaved by the dequantization)
|
||||
dq(tup, WU, 0, HC);
|
||||
D(P(lib, "gr_norm_h"), {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset),
|
||||
mtl::Arg::B(xnh), mtl::Arg::V(a)}, HC, T, 256, "hc");
|
||||
gemm(xnh, WD, V, T, LR + HC, HC * N, "hc gemm");
|
||||
D(P(lib, "gr_epi_down"), {mtl::Arg::B(V), mtl::Arg::B(loh), mtl::Arg::B(inj[1 - cur_inj]), mtl::Arg::V(ea)},
|
||||
(T * (LR + HC) + 255) / 256, 1, 256, "hc");
|
||||
GemmArgs g{T, HC * N, LR, 0};
|
||||
D(P(lib, "gemm_hc_up_64_64_4"), {mtl::Arg::B(loh), mtl::Arg::B(WU), mtl::Arg::B(xnh), mtl::Arg::B(mixed),
|
||||
mtl::Arg::B(mixh), mtl::Arg::V(g), mtl::Arg::V(ea)},
|
||||
HC * N / 64, (T + 63) / 64, 128, "hc gemm");
|
||||
cur_inj = 1 - cur_inj;
|
||||
return;
|
||||
}
|
||||
dq(tup, WU);
|
||||
D(P(lib, "gr_norm"), {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset),
|
||||
mtl::Arg::B(xn), mtl::Arg::V(a)}, HC, T, 256, "hc");
|
||||
act16(xn, xnh, (size_t)T * HC * N, "hc");
|
||||
gemm(xnh, WD, V, T, LR + HC, HC * N, "hc gemm");
|
||||
D(P(lib, "gr_epi_down"), {mtl::Arg::B(V), mtl::Arg::B(loh), mtl::Arg::B(inj[1 - cur_inj]), mtl::Arg::V(ea)},
|
||||
(T * (LR + HC) + 255) / 256, 1, 256, "hc");
|
||||
gemm(loh, WU, G, T, HC * N, LR, "hc gemm");
|
||||
D(P(lib, "gr_epi_up"), {mtl::Arg::B(G), mtl::Arg::B(xn), mtl::Arg::B(mixed), mtl::Arg::B(mixh), mtl::Arg::V(ea)},
|
||||
(T * N + 255) / 256, 1, 256, "hc");
|
||||
cur_inj = 1 - cur_inj;
|
||||
};
|
||||
auto build = [&](int pos0, int T, bool head) {
|
||||
ds.clear();
|
||||
cls.clear();
|
||||
int gi = 0, qi = 0;
|
||||
for (int l = 0; l < 48; ++l) {
|
||||
const std::string b = "blk." + std::to_string(l) + ".";
|
||||
hc_read(b + "hc_attn", l > 0, T);
|
||||
if (l % 4 != 3) {
|
||||
const gguf::Tensor &tq = T_(b + "attn_qkv.weight"), &tz = T_(b + "attn_gate.weight"),
|
||||
&ta = T_(b + "ssm_alpha.weight"), &tb = T_(b + "ssm_beta.weight"), &to = T_(b + "ssm_out.weight");
|
||||
dq(tq, Wqkv);
|
||||
dq(tz, Wz);
|
||||
dq(ta, Wal);
|
||||
dq(tb, Wbe);
|
||||
dq(to, Wout);
|
||||
gemm(mixh, Wqkv, qkv, T, 10240, N, "gdn proj");
|
||||
gemm(mixh, Wz, z, T, 6144, N, "gdn proj");
|
||||
gemm(mixh, Wal, al, T, 48, N, "gdn proj");
|
||||
gemm(mixh, Wbe, be, T, 48, N, "gdn proj");
|
||||
GdnArgs a{16, 48, 10240, 1, 1e-6f, 0};
|
||||
const int par = gpar[gi];
|
||||
if (gdn_mode != "step") {
|
||||
GdnPArgs ga{(uint32_t)T, 16, 48, 10240, 1e-6f, 0};
|
||||
D(P(lgdn, "gdn_pre"),
|
||||
{mtl::Arg::B(qkv), mtl::Arg::B(ghist[par][gi]), mtl::Arg::B(W, T_(b + "ssm_conv1d.weight").offset),
|
||||
mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(W, T_(b + "ssm_dt.bias").offset),
|
||||
mtl::Arg::B(W, T_(b + "ssm_a").offset), mtl::Arg::B(gqn), mtl::Arg::B(gkn), mtl::Arg::B(gvc), mtl::Arg::B(ggb),
|
||||
mtl::Arg::B(gkq), mtl::Arg::B(ghist[1 - par][gi]), mtl::Arg::V(ga)},
|
||||
16 + 24, T, 256, "gdn scan");
|
||||
const int ib = gdn_mode == "scan_v" ? 8 : std::atoi(gdn_mode.c_str() + 6);
|
||||
D(P(lgdn, "gdn_" + gdn_mode, gdn_mode == "scan_v" ? mtl::Constants() : mtl::Constants().i(6, 0)),
|
||||
{mtl::Arg::B(gst[gi]), mtl::Arg::B(gqn), mtl::Arg::B(gkn), mtl::Arg::B(gvc), mtl::Arg::B(ggb), mtl::Arg::B(gkq),
|
||||
mtl::Arg::B(go), mtl::Arg::V(ga)},
|
||||
48 * 4, 1, 32 * ib, "gdn scan");
|
||||
D(P(lgdn, "gdn_post"), {mtl::Arg::B(go), mtl::Arg::B(z), mtl::Arg::B(W, T_(b + "ssm_norm.weight").offset),
|
||||
mtl::Arg::B(ygh), mtl::Arg::B(yg), mtl::Arg::V(ga)},
|
||||
48, T, 128, "gdn scan");
|
||||
gemm(ygh, Wout, bo, T, N, 6144, "gdn proj");
|
||||
++gi;
|
||||
goto moe;
|
||||
}
|
||||
D(P(lgdn, "gdn_step", mtl::Constants().i(2, T), "gdn_step#" + std::to_string(T)),
|
||||
{mtl::Arg::B(gst[gi]), mtl::Arg::B(ghist[par][gi]), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be),
|
||||
mtl::Arg::B(z), mtl::Arg::B(W, T_(b + "ssm_conv1d.weight").offset), mtl::Arg::B(W, T_(b + "ssm_dt.bias").offset),
|
||||
mtl::Arg::B(W, T_(b + "ssm_a").offset), mtl::Arg::B(W, T_(b + "ssm_norm.weight").offset), mtl::Arg::B(yg),
|
||||
mtl::Arg::V(a), mtl::Arg::B(ghist[1 - par][gi])},
|
||||
48, 1, 512, "gdn scan");
|
||||
act16(yg, ygh, (size_t)T * 6144, "gdn proj");
|
||||
gemm(ygh, Wout, bo, T, N, 6144, "gdn proj");
|
||||
++gi;
|
||||
} else {
|
||||
const gguf::Tensor &tq = T_(b + "attn_q.weight"), &tk = T_(b + "attn_k.weight"), &tv = T_(b + "attn_v.weight"),
|
||||
&tiq = T_(b + "indexer.q_proj.weight"), &tik = T_(b + "indexer.k_proj.weight"),
|
||||
&to = T_(b + "attn_output.weight");
|
||||
dq(tq, Wq);
|
||||
dq(tk, Wk);
|
||||
dq(tv, Wv);
|
||||
dq(tiq, Wiq);
|
||||
dq(tik, Wik);
|
||||
dq(to, Wout);
|
||||
gemm(mixh, Wq, qfull, T, 12288, N, "qsa proj");
|
||||
gemm(mixh, Wk, kcur, T, 512, N, "qsa proj");
|
||||
gemm(mixh, Wv, vcur, T, 512, N, "qsa proj");
|
||||
gemm(mixh, Wiq, iqr, T, 512, N, "qsa proj");
|
||||
gemm(mixh, Wik, ikr, T, 128, N, "qsa proj");
|
||||
QsaArgs a{(uint32_t)pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f};
|
||||
D(P(lqsa, "qsa_prep"),
|
||||
{mtl::Arg::B(qfull), mtl::Arg::B(kcur), mtl::Arg::B(vcur), mtl::Arg::B(iqr), mtl::Arg::B(ikr),
|
||||
mtl::Arg::B(W, T_(b + "attn_q_norm.weight").offset), mtl::Arg::B(W, T_(b + "attn_k_norm.weight").offset),
|
||||
mtl::Arg::B(W, T_(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W, T_(b + "indexer.k_norm.weight").offset),
|
||||
mtl::Arg::B(cs), mtl::Arg::B(qcur), mtl::Arg::B(qidx), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]),
|
||||
mtl::Arg::B(ksc[qi]), mtl::Arg::B(vsc[qi]), mtl::Arg::B(tail[qi]), mtl::Arg::B(dead[qi]),
|
||||
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
|
||||
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");
|
||||
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");
|
||||
}
|
||||
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);
|
||||
D(P(lqsa, "qsa_merge"), {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(XQ), mtl::Arg::V(a)}, 24, T, 256,
|
||||
"qsa merge");
|
||||
act16(XQ, xqh, (size_t)T * 6144, "qsa proj");
|
||||
gemm(xqh, Wout, bo, T, N, 6144, "qsa proj");
|
||||
++qi;
|
||||
}
|
||||
moe:
|
||||
// ---- MoE
|
||||
hc_read(b + "hc_ffn", true, T);
|
||||
const gguf::Tensor &tr = T_(b + "ffn_gate_inp.weight"), &tsg = T_(b + "ffn_gate_inp_shexp.weight"),
|
||||
&tg_ = T_(b + "ffn_gate_shexp.weight"), &tu_ = T_(b + "ffn_up_shexp.weight"),
|
||||
&td_ = T_(b + "ffn_down_shexp.weight"), &eg = T_(b + "ffn_gate_exps.weight"),
|
||||
&eu = T_(b + "ffn_up_exps.weight"), &ed = T_(b + "ffn_down_exps.weight");
|
||||
dq(tr, Wr);
|
||||
dq(tsg, Wsg);
|
||||
dq(tg_, Wshg);
|
||||
dq(tu_, Wshu);
|
||||
dq(td_, Wshd);
|
||||
gemm(mixh, Wr, logits, T, NEXP, N, "router");
|
||||
const uint32_t nt = T;
|
||||
D(P(lib, "route_topk"), {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1,
|
||||
256, "router");
|
||||
const int nent = T * TOPK, nblk = (nent + 1023) / 1024;
|
||||
MmPlanArgs pa{(uint32_t)nent, TOPK, NEXP, (uint32_t)TM, (uint32_t)nblk};
|
||||
const uint32_t nclr = maxrows;
|
||||
D(P(lib, "mm_clear"), {mtl::Arg::B(stok), mtl::Arg::V(nclr)}, (maxrows + 255) / 256, 1, 256, "experts");
|
||||
D(P(lib, "mm_hist"), {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts");
|
||||
D(P(lib, "mm_scan"), {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta),
|
||||
mtl::Arg::V(pa)}, 1, 1, NEXP, "experts");
|
||||
D(P(lib, "mm_place"), {mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok),
|
||||
mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts");
|
||||
MmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t)ref::row_bytes(eg.type, N),
|
||||
(uint32_t)ref::row_bytes(ed.type, NFF), N, NFF};
|
||||
ma.gu_stride = (uint64_t)ma.gu_row * NFF;
|
||||
ma.d_stride = (uint64_t)ma.d_row * N;
|
||||
const uint32_t ne = N;
|
||||
mtl::Dispatch g;
|
||||
g.p = P(lib, "mm_gather");
|
||||
g.args = {mtl::Arg::B(mixed), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)};
|
||||
g.tg[0] = 256;
|
||||
g.indirect = &meta;
|
||||
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.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.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.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.indirect = &meta;
|
||||
dn.ind_off = 5 * 4;
|
||||
push(dn, "experts dn");
|
||||
const uint32_t ca[4] = {(uint32_t)T, N, TOPK, 0};
|
||||
D(P(lib, "mm_combine"), {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(yr), mtl::Arg::V(ca)},
|
||||
(T * N + 255) / 256, 1, 256, "experts");
|
||||
gemm(mixh, Wshg, gsh, T, NFF, N, "shared");
|
||||
gemm(mixh, Wshu, ush, T, NFF, N, "shared");
|
||||
const uint32_t nh = T * NFF;
|
||||
D(P(lib, "swiglu_h"), {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256,
|
||||
"shared");
|
||||
gemm(hh, Wshd, ys, T, N, NFF, "shared");
|
||||
gemm(mixh, Wsg, sgv, T, 1, N, "shared");
|
||||
const uint32_t oa[2] = {(uint32_t)T, N};
|
||||
D(P(lib, "moe_out"), {mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(oa)},
|
||||
(T * N + 255) / 256, 1, 256, "shared");
|
||||
}
|
||||
if (head) {
|
||||
// the final read of the last row (decode kernels, 1 token) and the LM head
|
||||
const size_t lr = (size_t)T - 1;
|
||||
const gguf::Tensor &tn = T_("output_hc_norm.weight"), &tdn = T_("output_hc_down.weight"), &tup = T_("output_hc_up.weight");
|
||||
GrArgs a{N, HC, LR, 1, 1, 1e-6f};
|
||||
D(P(lib, "gr_norm"), {mtl::Arg::B(R, lr * HC * N * 4), mtl::Arg::B(bo, lr * N * 4), mtl::Arg::B(inj[cur_inj], lr * HC * 4),
|
||||
mtl::Arg::B(W, tn.offset), mtl::Arg::B(xn1), mtl::Arg::V(a)}, HC, 1, 256, "head");
|
||||
D(P(lib, "gr_down_8_8", mtl::Constants().i(1, 1), "gr_down_8_8#1"),
|
||||
{mtl::Arg::B(W, tdn.offset), mtl::Arg::B(W, tdn.offset), mtl::Arg::B(xn1), mtl::Arg::B(lo1), mtl::Arg::B(inj1),
|
||||
mtl::Arg::V(a)}, LR + HC, 1, 256, "head");
|
||||
D(P(lib, "gr_up", mtl::Constants().i(1, 1), "gr_up#1"),
|
||||
{mtl::Arg::B(W, tup.offset), mtl::Arg::B(lo1), mtl::Arg::B(xn1), mtl::Arg::B(X1), mtl::Arg::V(a)}, N / 16, 1, 256, "head");
|
||||
const gguf::Tensor& to = T_("output.weight");
|
||||
GemvArgs ga{N, NV, 1, (uint32_t)ref::row_bytes(to.type, N), N, NV};
|
||||
const std::string fn = "gemv_" + std::to_string(to.type) + "_16_1_4_1";
|
||||
const int xk = to.type == ref::Q4_K || to.type == ref::Q5_K ? 2 : (to.type == ref::Q6_K || to.type == ref::Q3_K) ? 3 : 0;
|
||||
D(P(lib, fn, mtl::Constants().i(0, 1), fn + "#1"),
|
||||
{mtl::Arg::B(W, to.offset), mtl::Arg::B(X1, (size_t)4 * N * xk), mtl::Arg::B(X1, (size_t)4 * N * 4), mtl::Arg::B(lm),
|
||||
mtl::Arg::V(ga)}, (NV + 7) / 8, 1, 128, "head");
|
||||
ArgArgs aa{NV, 120};
|
||||
D(P(lsmp, "argmax_part"), {mtl::Arg::B(lm), mtl::Arg::B(apv), mtl::Arg::B(api), mtl::Arg::V(aa)}, 120, 1, 256, "head");
|
||||
D(P(lsmp, "argmax_final"), {mtl::Arg::B(apv), mtl::Arg::B(api), mtl::Arg::B(tokout), mtl::Arg::V(aa)}, 1, 1, 256, "head");
|
||||
}
|
||||
};
|
||||
// a chunk: the embedding rows into every stream, then the chunk's command buffer; GDN history ping-pong flipped
|
||||
auto run_chunk = [&](int pos0, int T, bool head) {
|
||||
for (int t = 0; t < T; ++t)
|
||||
for (int c = 0; c < HC; ++c) std::memcpy(R.as<float>() + ((size_t)t * HC + c) * N, E.data() + (size_t)t * N, 4 * N);
|
||||
cur_inj = 0;
|
||||
build(pos0, T, head);
|
||||
const double g = cx.run(ds);
|
||||
for (auto& p : gpar) p ^= 1;
|
||||
return g;
|
||||
};
|
||||
auto prompt = [&](int Pn, double* first_chunk) {
|
||||
double tot = 0;
|
||||
for (int p0 = 0; p0 < Pn; p0 += C) {
|
||||
const int T = std::min(C, Pn - p0);
|
||||
const double g = run_chunk(p0, T, p0 + T >= Pn);
|
||||
if (p0 == 0 && first_chunk) *first_chunk = g;
|
||||
tot += g;
|
||||
}
|
||||
return tot;
|
||||
};
|
||||
std::set<std::string> all_cls;
|
||||
// warm: the first chunk until the clocks are up
|
||||
{
|
||||
const double t0 = mtl::now_s();
|
||||
int n = 0;
|
||||
while (mtl::now_s() - t0 < 1.5) { run_chunk(0, std::min(C, prompts[0]), false); ++n; }
|
||||
std::printf("warm-up: %d chunks; %zu dispatches a chunk\n", n, ds.size());
|
||||
all_cls.insert(cls.begin(), cls.end());
|
||||
}
|
||||
if (breakdown) {
|
||||
// each class alone on the first chunk (pos0 0) and on a chunk at pos0 = PMAX - C (QSA past the budget)
|
||||
for (int p0 : {0, std::max(0, PMAX - C)}) {
|
||||
const std::string keep_only = only;
|
||||
std::map<std::string, double> by;
|
||||
double all = 0;
|
||||
std::vector<std::string> names = {""};
|
||||
for (const auto& c : all_cls) names.push_back(c);
|
||||
for (const std::string& c : names) {
|
||||
only = c;
|
||||
std::vector<double> t;
|
||||
for (int i = 0; i < 3; ++i) t.push_back(run_chunk(p0, C, false));
|
||||
if (c.empty()) all = median(t);
|
||||
else by[c] = median(t);
|
||||
}
|
||||
only = keep_only;
|
||||
double sum = 0;
|
||||
for (auto& kv : by) sum += kv.second;
|
||||
std::printf("chunk of %d at pos %d: all %.2f ms (%.0f tok/s); classes alone (sum %.2f ms):\n", C, p0, all * 1e3, C / all,
|
||||
sum * 1e3);
|
||||
for (auto& kv : by) std::printf(" %-10s %8.2f ms %6.2f us a token\n", kv.first.c_str(), kv.second * 1e3, kv.second / C * 1e6);
|
||||
}
|
||||
}
|
||||
for (int Pn : prompts) {
|
||||
std::vector<double> t;
|
||||
double fc = 0;
|
||||
for (int r = 0; r < reps; ++r) t.push_back(prompt(Pn, &fc));
|
||||
const double m = median(t);
|
||||
std::printf("prompt %6d tokens (chunks of %d): GPU %8.1f ms -> %7.1f tok/s (first chunk %.1f ms)\n", Pn, C, m * 1e3,
|
||||
Pn / m, fc * 1e3);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
// metal/tests/prompt_parity.cpp - the prompt-chunk glue (prompt.metal) against CPU references:
|
||||
// route_topk softmax + top 10 + renormalised weights (moe_route's rule) for 300 tokens, a third of them
|
||||
// with exact ties
|
||||
// the HC read as GEMMs gr_norm -> act_f16 -> [w_down; w_inject] GEMM -> gr_epi_down -> w_up GEMM -> gr_epi_up, on
|
||||
// random BF16 weights, vs the exact formula (fp64) - the fp16 operands set the tolerance
|
||||
// swiglu_h, moe_out elementwise
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
constexpr int N = 2560, HC = 4, LR = 320;
|
||||
struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; };
|
||||
struct EpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; };
|
||||
struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; };
|
||||
struct GemmArgs { int M, N, K, pad; };
|
||||
double sig(double v) { return 1.0 / (1.0 + std::exp(-v)); }
|
||||
} // namespace
|
||||
|
||||
int main() {
|
||||
mtl::Context cx;
|
||||
const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "gr.metal", "prompt.metal"});
|
||||
std::printf("device %s\n", cx.device_name().c_str());
|
||||
int fails = 0;
|
||||
// ---- route_topk
|
||||
{
|
||||
const int T = 300;
|
||||
mtl::Pipeline p = cx.pipeline(lib, "route_topk");
|
||||
mtl::Buffer lg = cx.alloc(4 * T * 512), ids = cx.alloc(4 * T * 10), w = cx.alloc(4 * T * 10);
|
||||
ref::Rng r(9);
|
||||
for (int t = 0; t < T; ++t)
|
||||
for (int e = 0; e < 512; ++e) {
|
||||
float v = r.sym() * 4.0f;
|
||||
if (t % 3 == 0) v = std::round(v * 4.0f) / 4.0f;
|
||||
lg.as<float>()[t * 512 + e] = v;
|
||||
}
|
||||
const uint32_t nt = T;
|
||||
mtl::Dispatch d;
|
||||
d.p = &p;
|
||||
d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::V(nt)};
|
||||
d.grid[0] = (T + 7) / 8;
|
||||
d.tg[0] = 256;
|
||||
cx.run({d});
|
||||
int bad = 0;
|
||||
double worst = 0;
|
||||
for (int t = 0; t < T; ++t) {
|
||||
const float* l = lg.as<float>() + t * 512;
|
||||
float mx = -INFINITY;
|
||||
for (int e = 0; e < 512; ++e) mx = std::max(mx, l[e]);
|
||||
std::vector<double> pr(512);
|
||||
double s = 0;
|
||||
for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx));
|
||||
std::vector<int> order(512);
|
||||
for (int e = 0; e < 512; ++e) order[e] = e;
|
||||
std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; });
|
||||
double ss = 0;
|
||||
for (int k = 0; k < 10; ++k) ss += pr[order[k]] / s;
|
||||
for (int k = 0; k < 10; ++k) {
|
||||
if (ids.as<int>()[t * 10 + k] != order[k]) { ++bad; break; }
|
||||
const double ref = pr[order[k]] / s / std::max(ss, 6.103515625e-5);
|
||||
worst = std::max(worst, std::fabs(w.as<float>()[t * 10 + k] - ref) / ref);
|
||||
}
|
||||
}
|
||||
const bool ok = bad == 0 && worst < 2e-6;
|
||||
fails += !ok;
|
||||
std::printf("route_topk, %d tokens: id mismatches %d, worst weight rel err %.2e %s\n", T, bad, worst, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
// ---- the HC read as GEMMs
|
||||
for (int apply = 0; apply < 2; ++apply) {
|
||||
const int T = 70, NIN = HC * N, NR = LR + HC;
|
||||
mtl::Buffer R = cx.alloc(4 * (size_t)T * NIN), bo = cx.alloc(4 * (size_t)T * N), inj = cx.alloc(4 * T * HC),
|
||||
wn = cx.alloc(4 * NIN), wdn = cx.alloc(2 * (size_t)LR * NIN), win = cx.alloc(2 * (size_t)HC * NIN),
|
||||
wup = cx.alloc(2 * (size_t)NIN * LR), xn = cx.alloc(4 * (size_t)T * NIN), xnh = cx.alloc(2 * (size_t)T * NIN),
|
||||
WD = cx.alloc(2 * (size_t)NR * NIN), WU = cx.alloc(2 * (size_t)NIN * LR), V = cx.alloc(4 * (size_t)T * NR),
|
||||
loh = cx.alloc(2 * (size_t)T * LR), injo = cx.alloc(4 * T * HC), G = cx.alloc(4 * (size_t)T * NIN),
|
||||
mixed = cx.alloc(4 * (size_t)T * N), mixh = cx.alloc(2 * (size_t)T * N);
|
||||
ref::Rng r(11 + apply);
|
||||
std::vector<float> R0((size_t)T * NIN);
|
||||
for (auto& v : R0) v = r.sym();
|
||||
std::copy(R0.begin(), R0.end(), R.as<float>());
|
||||
for (size_t i = 0; i < (size_t)T * N; ++i) bo.as<float>()[i] = r.sym();
|
||||
for (int i = 0; i < T * HC; ++i) inj.as<float>()[i] = r.sym() * 3.0f;
|
||||
for (int i = 0; i < NIN; ++i) wn.as<float>()[i] = 0.5f + r.uni();
|
||||
for (size_t i = 0; i < (size_t)LR * NIN; ++i) wdn.as<uint16_t>()[i] = ref::f32_to_bf16(r.sym() * 0.05f);
|
||||
for (size_t i = 0; i < (size_t)HC * NIN; ++i) win.as<uint16_t>()[i] = ref::f32_to_bf16(r.sym() * 0.05f);
|
||||
for (size_t i = 0; i < (size_t)NIN * LR; ++i) wup.as<uint16_t>()[i] = ref::f32_to_bf16(r.sym() * 0.2f);
|
||||
mtl::Pipeline pn = cx.pipeline(lib, "gr_norm"), pa = cx.pipeline(lib, "act_f16"), pdq = cx.pipeline(lib, "dq_f16_30"),
|
||||
pg = cx.pipeline(lib, "gemm_f16_64_64_4"), ped = cx.pipeline(lib, "gr_epi_down"), peu = cx.pipeline(lib, "gr_epi_up");
|
||||
std::vector<mtl::Dispatch> ds;
|
||||
GrArgs ga{N, HC, LR, (uint32_t)T, (uint32_t)apply, 1e-6f};
|
||||
EpiArgs ea{(uint32_t)T, N, HC, LR};
|
||||
auto D = [&](const mtl::Pipeline& p, std::vector<mtl::Arg> a, uint32_t gx, uint32_t gy, uint32_t tg) {
|
||||
mtl::Dispatch d; d.p = &p; d.args = std::move(a); d.grid[0] = gx; d.grid[1] = gy; d.tg[0] = tg; ds.push_back(d);
|
||||
};
|
||||
auto dq = [&](const mtl::Buffer& w, const mtl::Buffer& out, size_t ooff, int n_in, int n_out) {
|
||||
DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)(2 * n_in)};
|
||||
D(pdq, {mtl::Arg::B(w), mtl::Arg::B(out, ooff), mtl::Arg::V(a)}, (n_in / 16 + 63) / 64, n_out, 64);
|
||||
};
|
||||
auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K) {
|
||||
GemmArgs g{M, Nn, K, 0};
|
||||
D(pg, {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + 63) / 64, (M + 63) / 64, 128);
|
||||
};
|
||||
const uint32_t nx = T * NIN;
|
||||
dq(wdn, WD, 0, NIN, LR);
|
||||
dq(win, WD, (size_t)2 * LR * NIN, NIN, HC);
|
||||
dq(wup, WU, 0, LR, NIN);
|
||||
D(pn, {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xn), mtl::Arg::V(ga)}, HC, T, 256);
|
||||
D(pa, {mtl::Arg::B(xn), mtl::Arg::B(xnh), mtl::Arg::V(nx)}, (nx + 255) / 256, 1, 256);
|
||||
gemm(xnh, WD, V, T, NR, NIN);
|
||||
D(ped, {mtl::Arg::B(V), mtl::Arg::B(loh), mtl::Arg::B(injo), mtl::Arg::V(ea)}, (T * NR + 255) / 256, 1, 256);
|
||||
gemm(loh, WU, G, T, NIN, LR);
|
||||
D(peu, {mtl::Arg::B(G), mtl::Arg::B(xn), mtl::Arg::B(mixed), mtl::Arg::B(mixh), mtl::Arg::V(ea)}, (T * N + 255) / 256, 1, 256);
|
||||
cx.run(ds);
|
||||
// the fused read: gr_norm_h (fp16 xn), the down GEMM, gr_epi_down, the up GEMM with the epilogue (w_up rows
|
||||
// interleaved by dq_f16's perm)
|
||||
mtl::Buffer mixed2 = cx.alloc(4 * (size_t)T * N), mixh2 = cx.alloc(2 * (size_t)T * N), WU2 = cx.alloc(2 * (size_t)NIN * LR),
|
||||
xnh2 = cx.alloc(2 * (size_t)T * NIN), V2 = cx.alloc(4 * (size_t)T * NR), loh2 = cx.alloc(2 * (size_t)T * LR),
|
||||
injo2 = cx.alloc(4 * T * HC);
|
||||
std::copy(R0.begin(), R0.end(), R.as<float>());
|
||||
{
|
||||
mtl::Pipeline pnh = cx.pipeline(lib, "gr_norm_h"), pup = cx.pipeline(lib, "gemm_hc_up_64_64_4");
|
||||
ds.clear();
|
||||
DqArgs da{(uint32_t)LR, (uint32_t)NIN, (uint32_t)(2 * LR), (uint32_t)HC};
|
||||
D(pdq, {mtl::Arg::B(wup), mtl::Arg::B(WU2), mtl::Arg::V(da)}, (LR / 16 + 63) / 64, NIN, 64);
|
||||
D(pnh, {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xnh2), mtl::Arg::V(ga)}, HC, T, 256);
|
||||
gemm(xnh2, WD, V2, T, NR, NIN);
|
||||
D(ped, {mtl::Arg::B(V2), mtl::Arg::B(loh2), mtl::Arg::B(injo2), mtl::Arg::V(ea)}, (T * NR + 255) / 256, 1, 256);
|
||||
GemmArgs g{T, NIN, LR, 0};
|
||||
D(pup, {mtl::Arg::B(loh2), mtl::Arg::B(WU2), mtl::Arg::B(xnh2), mtl::Arg::B(mixed2), mtl::Arg::B(mixh2), mtl::Arg::V(g),
|
||||
mtl::Arg::V(ea)}, NIN / 64, (T + 63) / 64, 128);
|
||||
cx.run(ds);
|
||||
}
|
||||
// reference
|
||||
double se = 0, sr = 0, sei = 0, sri = 0, ser = 0, srr = 0, se2 = 0, sh2 = 0;
|
||||
std::vector<double> x(NIN), lo(LR);
|
||||
for (int t = 0; t < T; ++t) {
|
||||
for (int c = 0; c < HC; ++c) {
|
||||
const double g = apply ? 2.0 * sig(inj.as<float>()[t * HC + c] / (double)HC) : 0.0;
|
||||
double ss = 0;
|
||||
for (int d = 0; d < N; ++d) {
|
||||
const size_t i = (size_t)t * NIN + c * N + d;
|
||||
const double v = R0[i] + (apply ? bo.as<float>()[(size_t)t * N + d] * g : 0.0);
|
||||
x[c * N + d] = v;
|
||||
ss += v * v;
|
||||
ser += (R.as<float>()[i] - v) * (R.as<float>()[i] - v);
|
||||
srr += v * v;
|
||||
}
|
||||
const double rs = 1.0 / std::sqrt(ss / N + 1e-6);
|
||||
for (int d = 0; d < N; ++d) x[c * N + d] *= wn.as<float>()[c * N + d] * rs;
|
||||
}
|
||||
for (int k = 0; k < NR; ++k) {
|
||||
const uint16_t* w = k < LR ? wdn.as<uint16_t>() + (size_t)k * NIN : win.as<uint16_t>() + (size_t)(k - LR) * NIN;
|
||||
double s = 0;
|
||||
for (int i = 0; i < NIN; ++i) s += ref::bf16_to_f32(w[i]) * x[i];
|
||||
if (k < LR) lo[k] = s / HC * sig(s / HC);
|
||||
else {
|
||||
const double gv = injo.as<float>()[t * HC + k - LR];
|
||||
sei += (gv - s) * (gv - s);
|
||||
sri += s * s;
|
||||
}
|
||||
}
|
||||
for (int d = 0; d < N; ++d) {
|
||||
double m = 0;
|
||||
for (int c = 0; c < HC; ++c) {
|
||||
const uint16_t* w = wup.as<uint16_t>() + (size_t)(c * N + d) * LR;
|
||||
double s = 0;
|
||||
for (int k = 0; k < LR; ++k) s += ref::bf16_to_f32(w[k]) * lo[k];
|
||||
m += x[c * N + d] * sig(s);
|
||||
}
|
||||
m /= HC;
|
||||
const double gv = mixed.as<float>()[(size_t)t * N + d], g2 = mixed2.as<float>()[(size_t)t * N + d],
|
||||
h2 = ref::f16_to_f32(mixh2.as<uint16_t>()[(size_t)t * N + d]);
|
||||
se += (gv - m) * (gv - m);
|
||||
se2 += (g2 - m) * (g2 - m);
|
||||
sh2 += (h2 - m) * (h2 - m);
|
||||
sr += m * m;
|
||||
}
|
||||
}
|
||||
const double rm = std::sqrt(se / sr), ri = std::sqrt(sei / sri), rr = std::sqrt(ser / srr);
|
||||
const double r2 = std::sqrt(se2 / sr), rh2 = std::sqrt(sh2 / sr);
|
||||
const bool ok = rm < 2e-3 && ri < 2e-3 && rr < 1e-6 && r2 < 2e-3 && rh2 < 2e-3;
|
||||
fails += !ok;
|
||||
std::printf("HC read as GEMMs, %d tokens, apply %d: mixed rel err %.2e, inject %.2e, residual %.2e; fused read "
|
||||
"(fp16 xn, up GEMM epilogue): mixed %.2e, fp16 %.2e %s\n", T, apply, rm, ri, rr, r2, rh2, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
// ---- swiglu_h, moe_out
|
||||
{
|
||||
const int n = 5000, T = 7, NE = 640;
|
||||
mtl::Pipeline ps = cx.pipeline(lib, "swiglu_h"), pm = cx.pipeline(lib, "moe_out");
|
||||
mtl::Buffer g = cx.alloc(4 * n), u = cx.alloc(4 * n), h = cx.alloc(2 * n), yr = cx.alloc(4 * T * NE),
|
||||
ys = cx.alloc(4 * T * NE), sg = cx.alloc(4 * T), bo = cx.alloc(4 * T * NE);
|
||||
ref::Rng r(5);
|
||||
for (int i = 0; i < n; ++i) { g.as<float>()[i] = r.sym() * 6; u.as<float>()[i] = r.sym(); }
|
||||
for (int i = 0; i < T * NE; ++i) { yr.as<float>()[i] = r.sym(); ys.as<float>()[i] = r.sym(); }
|
||||
for (int i = 0; i < T; ++i) sg.as<float>()[i] = r.sym() * 4;
|
||||
const uint32_t nn = n;
|
||||
const uint32_t a2[2] = {(uint32_t)T, (uint32_t)NE};
|
||||
mtl::Dispatch d1; d1.p = &ps; d1.args = {mtl::Arg::B(g), mtl::Arg::B(u), mtl::Arg::B(h), mtl::Arg::V(nn)};
|
||||
d1.grid[0] = (n + 255) / 256; d1.tg[0] = 256;
|
||||
mtl::Dispatch d2; d2.p = ± d2.args = {mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sg), mtl::Arg::B(bo), mtl::Arg::V(a2)};
|
||||
d2.grid[0] = (T * NE + 255) / 256; d2.tg[0] = 256;
|
||||
cx.run({d1, d2});
|
||||
double w1 = 0, w2 = 0;
|
||||
for (int i = 0; i < n; ++i) {
|
||||
const double gv = g.as<float>()[i], ref = gv * sig(gv) * u.as<float>()[i];
|
||||
w1 = std::max(w1, std::fabs(ref::f16_to_f32(h.as<uint16_t>()[i]) - ref) / (std::fabs(ref) + 1e-3));
|
||||
}
|
||||
for (int i = 0; i < T * NE; ++i) {
|
||||
const double ref = yr.as<float>()[i] + sig(sg.as<float>()[i / NE]) * ys.as<float>()[i];
|
||||
w2 = std::max(w2, std::fabs(bo.as<float>()[i] - ref));
|
||||
}
|
||||
const bool ok = w1 < 1e-3 && w2 < 1e-5;
|
||||
fails += !ok;
|
||||
std::printf("swiglu_h worst rel err %.2e, moe_out worst abs err %.2e %s\n", w1, w2, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS");
|
||||
return fails ? 1 : 0;
|
||||
}
|
||||
@@ -60,7 +60,11 @@ int main(int argc, char** argv) {
|
||||
if (argc < 2) { std::fprintf(stderr, "usage: token_chain <shard-1.gguf> [--ntok N] [--ctx P] [--reps R]\n"); return 2; }
|
||||
int ntok = 1, ctx = 1000, reps = 10;
|
||||
bool skip_qsa = false, graph = false;
|
||||
std::string hc_mode = "3"; // 3: BF16, 3 dispatches; 2: BF16, 2 dispatches; q8: Q8_0 copies, 2 dispatches
|
||||
// the HC read: Q8_0 copies through the 3-dispatch read by default (the owner's decision, as STRATA_HC_Q8 on gfx12);
|
||||
// STRATA_HC_Q8=0 selects the exact BF16 path. --hc 3: BF16, 3 dispatches; 2: BF16, 2 dispatches; q8: Q8_0 copies,
|
||||
// 2 dispatches; q8_3: Q8_0 copies, 3 dispatches
|
||||
const char* hcq8 = std::getenv("STRATA_HC_Q8");
|
||||
std::string hc_mode = hcq8 && std::string(hcq8) == "0" ? "3" : "q8_3";
|
||||
bool fuse = false; // the GDN y prep in the step; no QSA score/top-k within the selection budget
|
||||
bool fold_comb = false; // the MoE combine folded into the next 2-dispatch HC read (measured slower)
|
||||
std::string only, skip;
|
||||
|
||||
Reference in new issue
Block a user