mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-08 17:35:22 -10:00
metal: decode T=2 items - router GEMV + moe_route in one dispatch (router_route), gemv_cat, gr_norm_comb, Q3_K packed loads
- router.metal (new): router_route<F> - the router GEMV rows as gemv's (LPR 32, KS 1) with 512-thread threadgroups, the last threadgroup to finish (device counter, device-scope fences) runs moe_route's body and resets the counter; logits / ids / w / plan bitwise gemv + moe_route. token_chain default (CHAIN_ROUTE=0: separate): T2 13.08 -> 12.93 ms - gemv.metal gemv_cat: two same-format matrices over one x in one dispatch (ssm_alpha + ssm_beta), bitwise; no gain in graph mode (CHAIN_ALBE opt-in) - gr2.metal gr_norm_comb: the MoE combine folded into gr_norm, bitwise; +0.6 ms at T2 (CHAIN_COMB opt-in, not kept) - units.metal Dec<Q3_K>: qs / hmask as packed_ushort4 loads (same values): T2 2560->10240 33.8 -> 32.2 us - tests/fuse_parity: the three fusions bitwise for 1..4 tokens
This commit is contained in:
Niko1221 committed
1 parent
56526aedc2
commit
3f9ca99632
6 files changed
+419
-14
No files matched your search
@@ -139,3 +139,53 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float*
|
||||
GP_ROW(T_IQ3_S) GP_ROW(T_IQ4_XS) GP_ROW(T_Q2_0) GP_ROW(T_Q3_K) GP_ROW(T_Q4_K) GP_ROW(T_Q5_K) GP_ROW(T_Q6_K)
|
||||
// Q2_0 with raw activations (the shared expert's down reading the pair's raw h)
|
||||
GV_SET(T_Q2_0R)
|
||||
|
||||
// gemv_cat: two matrices of the same format and n_in over the same x in one dispatch - rows [0, n_split) are W's
|
||||
// (into y), rows [n_split, n_out) are W2's (into y2, from its row 0). Each row's arithmetic is gemv's (KS = 1), so the
|
||||
// outputs are bitwise the two separate GEMVs'. (GDN's ssm_alpha + ssm_beta: 48 + 48 rows.)
|
||||
// Buffers: W, xp, s16, y, args (n_out = both), W2, y2, n_split.
|
||||
template <int F, int LPR, int NR, int SG>
|
||||
kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]],
|
||||
device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]],
|
||||
device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) {
|
||||
constexpr int UNIT = fmt<F>::UNIT, RG = 32 / LPR;
|
||||
threadgroup uint tab[fmt<F>::TAB];
|
||||
load_tab<F>(tab, tid, SG * 32);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const int li = lane % LPR;
|
||||
const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR;
|
||||
const int U = int(a.n_in) / UNIT, sp = int(n_split);
|
||||
float acc[NR][8];
|
||||
for (int r = 0; r < NR; ++r)
|
||||
for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f;
|
||||
device const uchar* rp[NR];
|
||||
for (int r = 0; r < NR; ++r) {
|
||||
const int row = min(row0 + r, int(a.n_out) - 1);
|
||||
rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes;
|
||||
}
|
||||
const int s16n = int(a.n_in) / 16;
|
||||
for (int u = li; u < U; u += LPR) {
|
||||
float q[NR][UNIT];
|
||||
float4 sc[NR];
|
||||
for (int r = 0; r < NR; ++r) Dec<F>::run(rp[r], u, q[r], sc[r], tab);
|
||||
for (int t = 0; t < NTOK; ++t) {
|
||||
float xs[UNIT], xsum[2];
|
||||
X<F>::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum);
|
||||
for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine<UNIT>(q[r], sc[r], xs, xsum);
|
||||
}
|
||||
}
|
||||
for (int r = 0; r < NR; ++r)
|
||||
for (int t = 0; t < NTOK; ++t) {
|
||||
const float v = lane_sum<LPR>(acc[r][t]);
|
||||
const int row = row0 + r;
|
||||
if (li == 0 && row < int(a.n_out)) {
|
||||
if (row < sp) y[t * int(a.y_stride) + row] = v;
|
||||
else y2[t * int(a.y_stride) + row - sp] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
#define GVC_(F, L, R, S) template [[host_name("gemvc_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_cat<F, L, R, S>) gemv_cat<F, L, R, S>;
|
||||
#define GVC(F, L, R, S) GVC_(F, L, R, S)
|
||||
GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4)
|
||||
@@ -301,3 +301,41 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch
|
||||
}
|
||||
#define GD3B(T, R) template [[host_name("gr_down3b_" #T "_" #R)]] kernel decltype(gr_down3b<T, R>) gr_down3b<T, R>;
|
||||
GD3B(1, 1) GD3B(2, 1) GD3B(3, 1) GD3B(1, 2) GD3B(2, 2) GD3B(3, 2)
|
||||
|
||||
// gr_norm_comb: gr_norm (gr.metal) with the sublayer output bo formed on the fly as the MoE combine - bo[t][d] =
|
||||
// sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d], moe_combine_sh's arithmetic - so the combine's
|
||||
// dispatch goes away; R', xn bitwise gr_norm's after moe_combine_sh. One threadgroup per (stream c, token t).
|
||||
kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]],
|
||||
device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]],
|
||||
device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]],
|
||||
device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]],
|
||||
device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]],
|
||||
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;
|
||||
const float gs = 1.0f / (1.0f + exp(-sgv[t]));
|
||||
float ss = 0.0f;
|
||||
for (uint d = tid; d < n; d += 256) {
|
||||
float v = r[d];
|
||||
if (a.apply) {
|
||||
float s = 0.0f;
|
||||
for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d];
|
||||
const float bo = fma(gs, ysh[t * n + d], s);
|
||||
v = fma(bo, 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 float* o = xn + (t * a.hc + c) * n;
|
||||
device const float* w = w_norm + c * n;
|
||||
for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs;
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
// metal/kernels/router.metal - the router GEMV and moe_route in ONE dispatch (after quant/units/moe/gemv.metal): the
|
||||
// logits rows as gemv's (LPR = 32 lanes a row, KS = 1: bitwise gemv_<F>_32_1_*_1's), 16 rows a threadgroup of 512
|
||||
// threads; every threadgroup bumps a device counter after its rows (device-scope fences around it), and the LAST one
|
||||
// runs moe_route's body (the softmax / top-10 / weights, then the plan) with its 512 threads (= n_expert) and resets
|
||||
// the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's.
|
||||
// Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero).
|
||||
|
||||
static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan,
|
||||
constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp,
|
||||
threadgroup int* s_ent) {
|
||||
if (sg < a.n_tok) {
|
||||
device const float* l = logits + sg * 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[sg * 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);
|
||||
const float inv = 1.0f / ssum;
|
||||
if (lane < 10) w[sg * 10 + lane] = selected * inv;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent);
|
||||
}
|
||||
|
||||
template <int F>
|
||||
kernel void router_route(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]],
|
||||
device const float* s16 [[buffer(2)]], device float* logits [[buffer(3)]],
|
||||
constant GemvArgs& ga [[buffer(4)]], device int* ids [[buffer(5)]], device float* w [[buffer(6)]],
|
||||
device int* plan [[buffer(7)]], constant MoeArgs& a [[buffer(8)]],
|
||||
device atomic_uint* counter [[buffer(9)]],
|
||||
uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]],
|
||||
ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
|
||||
constexpr int UNIT = fmt<F>::UNIT;
|
||||
threadgroup uint tab[fmt<F>::TAB];
|
||||
threadgroup int s_grp[32], s_ent[32];
|
||||
threadgroup uint last;
|
||||
load_tab<F>(tab, tid, nth);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// the GEMV: simdgroup sg -> row tg.x * 16 + sg, 32 lanes over its units
|
||||
const int row = int(tgx) * 16 + sg;
|
||||
const int rr = min(row, int(ga.n_out) - 1);
|
||||
device const uchar* rp = W + (ulong)rr * ga.row_bytes;
|
||||
const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16;
|
||||
float acc[8];
|
||||
for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f;
|
||||
for (int u = lane; u < U; u += 32) {
|
||||
float q[UNIT];
|
||||
float4 sc;
|
||||
Dec<F>::run(rp, u, q, sc, tab);
|
||||
for (int t = 0; t < NTOK; ++t) {
|
||||
float xs[UNIT], xsum[2];
|
||||
X<F>::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum);
|
||||
acc[t] += unit_combine<UNIT>(q, sc, xs, xsum);
|
||||
}
|
||||
}
|
||||
for (int t = 0; t < NTOK; ++t) {
|
||||
const float v = lane_sum<32>(acc[t]);
|
||||
if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v;
|
||||
}
|
||||
// the last threadgroup to finish routes
|
||||
atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
if (tid == 0) last = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == ngx - 1 ? 1u : 0u;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (!last) return;
|
||||
atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
|
||||
route_body(logits, ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent);
|
||||
if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed);
|
||||
}
|
||||
|
||||
template [[host_name("router_route_30")]] kernel decltype(router_route<T_BF16>) router_route<T_BF16>;
|
||||
@@ -227,8 +227,15 @@ template <> struct Dec<T_Q3_K> {
|
||||
static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) {
|
||||
device const uchar* b = row + (u >> 3) * 110;
|
||||
const int n = (u >> 2) & 1, j = u & 3;
|
||||
device const ushort* qq = (device const ushort*)(b + 32 + 32 * n);
|
||||
device const ushort* hm = (device const ushort*)b;
|
||||
// qs and hmask as packed (2-byte aligned) vectors: blocks are 110 bytes, so odd blocks are only 2-aligned
|
||||
device const packed_ushort4* qv = (device const packed_ushort4*)(b + 32 + 32 * n);
|
||||
device const packed_ushort4* hv = (device const packed_ushort4*)b;
|
||||
ushort qq[16], hm[16];
|
||||
for (int v = 0; v < 4; ++v) {
|
||||
const ushort4 a4 = ushort4(qv[v]), h4 = ushort4(hv[v]);
|
||||
qq[4 * v] = a4.x; qq[4 * v + 1] = a4.y; qq[4 * v + 2] = a4.z; qq[4 * v + 3] = a4.w;
|
||||
hm[4 * v] = h4.x; hm[4 * v + 1] = h4.y; hm[4 * v + 2] = h4.z; hm[4 * v + 3] = h4.w;
|
||||
}
|
||||
device const uchar* s = b + 96;
|
||||
const int i1 = 8 * n + 2 * j, i2 = i1 + 1;
|
||||
const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4);
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
// metal/tests/fuse_parity.cpp - the decode latency fusions against the kernels they replace, BITWISE, for 1..4
|
||||
// tokens:
|
||||
// gemv_cat (gemv.metal) ssm_alpha + ssm_beta (BF16 2560 -> 48 + 48) in one dispatch == two gemv dispatches
|
||||
// router_route the router GEMV + moe_route in one dispatch == gemv + moe_route (logits, ids, w, plan;
|
||||
// twice in a row: the counter resets)
|
||||
// gr_norm_comb (gr2.metal) the MoE combine folded into the next gr_norm == moe_combine_sh + gr_norm (R', xn)
|
||||
#include "mtl_rt.h"
|
||||
#include "ref_quant.h"
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; };
|
||||
struct MoeArgs {
|
||||
uint64_t gate_off, up_off, down_off, gu_stride, d_stride;
|
||||
uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad;
|
||||
};
|
||||
constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2;
|
||||
struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; };
|
||||
struct GlueArgs { uint32_t n, n_tok, n_embd, topk; };
|
||||
} // namespace
|
||||
|
||||
int main() {
|
||||
mtl::Context cx;
|
||||
const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal",
|
||||
"gr2.metal", "glue.metal"});
|
||||
std::printf("device %s\n", cx.device_name().c_str());
|
||||
int fails = 0;
|
||||
const int N = 2560, NO = 48, MAXT = 8;
|
||||
const size_t rb = ref::row_bytes(ref::BF16, N);
|
||||
mtl::Buffer Wa = cx.alloc(rb * NO), Wb = cx.alloc(rb * NO), xp = cx.alloc(4 * ((size_t)4 * MAXT * N + MAXT * N / 16)),
|
||||
y1 = cx.alloc(4 * MAXT * NO), y2 = cx.alloc(4 * MAXT * NO), z1 = cx.alloc(4 * MAXT * NO), z2 = cx.alloc(4 * MAXT * NO);
|
||||
ref::random_blocks(ref::BF16, Wa.as<uint8_t>(), NO, N, 1, 0.05f);
|
||||
ref::random_blocks(ref::BF16, Wb.as<uint8_t>(), NO, N, 2, 0.05f);
|
||||
ref::Rng r(3);
|
||||
for (size_t i = 0; i < xp.bytes / 4; ++i) xp.as<float>()[i] = r.sym();
|
||||
for (int nt = 1; nt <= 4; ++nt) {
|
||||
mtl::Pipeline pg = cx.pipeline(lib, "gemv_30_32_1_4_1", mtl::Constants().i(0, nt)),
|
||||
pc = cx.pipeline(lib, "gemvc_30_32_1_4", mtl::Constants().i(0, nt));
|
||||
const size_t s16 = (size_t)4 * 4 * MAXT * N;
|
||||
GemvArgs a{(uint32_t)N, (uint32_t)NO, (uint32_t)nt, (uint32_t)rb, (uint32_t)N, (uint32_t)NO};
|
||||
auto g = [&](const mtl::Buffer& W, const mtl::Buffer& y) {
|
||||
mtl::Dispatch d;
|
||||
d.p = &pg;
|
||||
d.args = {mtl::Arg::B(W), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(y), mtl::Arg::V(a)};
|
||||
d.grid[0] = NO / 4;
|
||||
d.tg[0] = 128;
|
||||
return d;
|
||||
};
|
||||
cx.run({g(Wa, y1), g(Wb, y2)});
|
||||
GemvArgs ac = a;
|
||||
ac.n_out = 2 * NO;
|
||||
const uint32_t sp = NO;
|
||||
mtl::Dispatch d;
|
||||
d.p = &pc;
|
||||
d.args = {mtl::Arg::B(Wa), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(z1), mtl::Arg::V(ac), mtl::Arg::B(Wb),
|
||||
mtl::Arg::B(z2), mtl::Arg::V(sp)};
|
||||
d.grid[0] = 2 * NO / 4;
|
||||
d.tg[0] = 128;
|
||||
cx.run({d});
|
||||
const bool ok = std::memcmp(y1.ptr, z1.ptr, 4 * nt * NO) == 0 && std::memcmp(y2.ptr, z2.ptr, 4 * nt * NO) == 0;
|
||||
fails += !ok;
|
||||
std::printf("gemv_cat (alpha + beta) ntok %d: bitwise %s\n", nt, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
// ---- router_route
|
||||
{
|
||||
const int NE = 512;
|
||||
const size_t rbr = ref::row_bytes(ref::BF16, N);
|
||||
mtl::Buffer Wr = cx.alloc(rbr * NE), lg1 = cx.alloc(4 * MAXT * NE), lg2 = cx.alloc(4 * MAXT * NE),
|
||||
id1 = cx.alloc(4 * MAXT * 10), id2 = cx.alloc(4 * MAXT * 10), w1 = cx.alloc(4 * MAXT * 10),
|
||||
w2 = cx.alloc(4 * MAXT * 10), pl1 = cx.alloc(4 * P_INTS), pl2 = cx.alloc(4 * P_INTS), cnt = cx.alloc(16);
|
||||
ref::random_blocks(ref::BF16, Wr.as<uint8_t>(), NE, N, 9, 0.05f);
|
||||
mtl::Pipeline pr = cx.pipeline(lib, "moe_route");
|
||||
for (int nt = 1; nt <= 4; ++nt) {
|
||||
mtl::Pipeline pg = cx.pipeline(lib, "gemv_30_32_1_4_1", mtl::Constants().i(0, nt)),
|
||||
pf = cx.pipeline(lib, "router_route_30", mtl::Constants().i(0, nt));
|
||||
const size_t s16 = (size_t)4 * 4 * MAXT * N;
|
||||
GemvArgs ga{(uint32_t)N, (uint32_t)NE, (uint32_t)nt, (uint32_t)rbr, (uint32_t)N, (uint32_t)NE};
|
||||
MoeArgs ma{};
|
||||
ma.n_tok = nt;
|
||||
ma.topk = 10;
|
||||
ma.n_expert = NE;
|
||||
ma.gx_gu = 5;
|
||||
ma.gx_dn = 20;
|
||||
mtl::Dispatch g;
|
||||
g.p = &pg;
|
||||
g.args = {mtl::Arg::B(Wr), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(lg1), mtl::Arg::V(ga)};
|
||||
g.grid[0] = NE / 4;
|
||||
g.tg[0] = 128;
|
||||
mtl::Dispatch rt;
|
||||
rt.p = ≺
|
||||
rt.args = {mtl::Arg::B(lg1), mtl::Arg::B(id1), mtl::Arg::B(w1), mtl::Arg::B(pl1), mtl::Arg::V(ma)};
|
||||
rt.tg[0] = 512;
|
||||
std::memset(pl1.ptr, 0, pl1.bytes);
|
||||
cx.run({g, rt});
|
||||
bool ok = true;
|
||||
for (int rep_ = 0; rep_ < 2; ++rep_) {
|
||||
std::memset(pl2.ptr, 0, pl2.bytes);
|
||||
mtl::Dispatch f;
|
||||
f.p = &pf;
|
||||
f.args = {mtl::Arg::B(Wr), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(lg2), mtl::Arg::V(ga), mtl::Arg::B(id2),
|
||||
mtl::Arg::B(w2), mtl::Arg::B(pl2), mtl::Arg::V(ma), mtl::Arg::B(cnt)};
|
||||
f.grid[0] = NE / 16;
|
||||
f.tg[0] = 512;
|
||||
cx.run({f});
|
||||
ok = ok && std::memcmp(lg1.ptr, lg2.ptr, 4 * nt * NE) == 0 && std::memcmp(id1.ptr, id2.ptr, 4 * nt * 10) == 0 &&
|
||||
std::memcmp(w1.ptr, w2.ptr, 4 * nt * 10) == 0 && std::memcmp(pl1.ptr, pl2.ptr, 4 * P_INTS) == 0 &&
|
||||
cnt.as<uint32_t>()[0] == 0;
|
||||
}
|
||||
fails += !ok;
|
||||
std::printf("router_route ntok %d: logits, ids, w, plan bitwise (2 runs) %s\n", nt, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
}
|
||||
// ---- gr_norm_comb
|
||||
{
|
||||
const int HC = 4, TOPK = 10;
|
||||
mtl::Buffer R0 = cx.alloc(4 * MAXT * HC * N), R1 = cx.alloc(4 * MAXT * HC * N), R2 = cx.alloc(4 * MAXT * HC * N),
|
||||
ye = cx.alloc(4 * (size_t)MAXT * TOPK * N), ysh = cx.alloc(4 * MAXT * N), sgv = cx.alloc(4 * MAXT),
|
||||
plan = cx.alloc(4 * P_INTS), inj = cx.alloc(4 * MAXT * HC), wn = cx.alloc(4 * HC * N), bo = cx.alloc(4 * MAXT * N),
|
||||
x1 = cx.alloc(4 * MAXT * HC * N), x2 = cx.alloc(4 * MAXT * HC * N);
|
||||
for (size_t i = 0; i < R0.bytes / 4; ++i) R0.as<float>()[i] = r.sym();
|
||||
for (size_t i = 0; i < ye.bytes / 4; ++i) ye.as<float>()[i] = r.sym();
|
||||
for (size_t i = 0; i < ysh.bytes / 4; ++i) ysh.as<float>()[i] = r.sym();
|
||||
for (int i = 0; i < MAXT; ++i) sgv.as<float>()[i] = r.sym() * 3;
|
||||
for (int i = 0; i < MAXT * HC; ++i) inj.as<float>()[i] = r.sym() * 3;
|
||||
for (int i = 0; i < HC * N; ++i) wn.as<float>()[i] = 0.5f + r.uni();
|
||||
for (int i = 0; i < MAXT * TOPK; ++i) plan.as<int>()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK);
|
||||
mtl::Pipeline pcm = cx.pipeline(lib, "moe_combine_sh"), pn = cx.pipeline(lib, "gr_norm"), pnc = cx.pipeline(lib, "gr_norm_comb");
|
||||
for (int nt = 1; nt <= 4; ++nt) {
|
||||
std::memcpy(R1.ptr, R0.ptr, R0.bytes);
|
||||
std::memcpy(R2.ptr, R0.ptr, R0.bytes);
|
||||
GlueArgs ga{(uint32_t)N, (uint32_t)nt, (uint32_t)N, TOPK};
|
||||
GrArgs gr{(uint32_t)N, HC, 320, (uint32_t)nt, 1, 1e-6f};
|
||||
mtl::Dispatch c; c.p = &pcm;
|
||||
c.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(ga)};
|
||||
c.grid[0] = (nt * N + 255) / 256; c.tg[0] = 256;
|
||||
mtl::Dispatch n1; n1.p = &pn;
|
||||
n1.args = {mtl::Arg::B(R1), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x1), mtl::Arg::V(gr)};
|
||||
n1.grid[0] = HC; n1.grid[1] = nt; n1.tg[0] = 256;
|
||||
cx.run({c, n1});
|
||||
const uint32_t tk = TOPK;
|
||||
mtl::Dispatch n2; n2.p = &pnc;
|
||||
n2.args = {mtl::Arg::B(R2), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x2), mtl::Arg::V(gr),
|
||||
mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)};
|
||||
n2.grid[0] = HC; n2.grid[1] = nt; n2.tg[0] = 256;
|
||||
cx.run({n2});
|
||||
const bool ok = std::memcmp(R1.ptr, R2.ptr, 4 * (size_t)nt * HC * N) == 0 && std::memcmp(x1.ptr, x2.ptr, 4 * (size_t)nt * HC * N) == 0;
|
||||
fails += !ok;
|
||||
std::printf("gr_norm_comb ntok %d: R', xn bitwise %s\n", nt, ok ? "PASS" : "FAIL");
|
||||
}
|
||||
}
|
||||
std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS");
|
||||
return fails ? 1 : 0;
|
||||
}
|
||||
+62
-12
@@ -102,7 +102,8 @@ int main(int argc, char** argv) {
|
||||
return *t;
|
||||
};
|
||||
|
||||
const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "gr.metal", "gr2.metal", "glue.metal"});
|
||||
const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "gr.metal", "gr2.metal", "glue.metal",
|
||||
"router.metal"});
|
||||
const int lgdn = cx.compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
const int lqsa = cx.compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
const int lsmp = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
@@ -177,6 +178,12 @@ int main(int argc, char** argv) {
|
||||
// ---- SLC prefetch (--prefetch): the next HC read's weights touched concurrently with the first projection GEMV
|
||||
// after the current read, so the latency-bound read finds them in the SLC
|
||||
const bool prefetch = std::getenv("CHAIN_PREFETCH") != nullptr;
|
||||
const bool albe = std::getenv("CHAIN_ALBE") != nullptr; // ssm_alpha + ssm_beta as one gemv_cat
|
||||
// the router GEMV + moe_route as one router_route (default; CHAIN_ROUTE=0: separate)
|
||||
const bool routef = !(std::getenv("CHAIN_ROUTE") && std::string(std::getenv("CHAIN_ROUTE")) == "0");
|
||||
const bool combf = std::getenv("CHAIN_COMB") != nullptr; // the MoE combine folded into the next gr_norm
|
||||
bool comb_pending = false;
|
||||
mtl::Buffer rcounter = cx.alloc(16);
|
||||
struct Range { const mtl::Buffer* b; size_t off, bytes; };
|
||||
std::vector<Range> pending;
|
||||
mtl::Buffer sink = cx.alloc(16);
|
||||
@@ -219,6 +226,7 @@ int main(int argc, char** argv) {
|
||||
if (!pending.empty()) push_prefetch();
|
||||
};
|
||||
const mtl::Pipeline* pnorm = Pn(lib, "gr_norm", "gr_norm", {});
|
||||
const mtl::Pipeline* pnormc = Pn(lib, "gr_norm_comb", "gr_norm_comb", {});
|
||||
ptouch = Pn(lib, "slc_touch", "slc_touch", {});
|
||||
const mtl::Pipeline* pdown = Pn(lib, "gr_down_8_8", "gr_down_8_8", mtl::Constants().i(1, ntok));
|
||||
const mtl::Pipeline* pup = Pn(lib, "gr_up", "gr_up", mtl::Constants().i(1, ntok));
|
||||
@@ -320,6 +328,14 @@ int main(int argc, char** argv) {
|
||||
d1.p = pnorm;
|
||||
d1.args = {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)};
|
||||
if (comb_pending) {
|
||||
// the MoE combine folded into this norm (gr_norm_comb; bitwise moe_combine_sh + gr_norm)
|
||||
const uint32_t tk = TOPK;
|
||||
d1.p = pnormc;
|
||||
d1.args = {mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset), mtl::Arg::B(xn),
|
||||
mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sg), mtl::Arg::V(tk)};
|
||||
comb_pending = false;
|
||||
}
|
||||
d1.grid[0] = HC;
|
||||
d1.grid[1] = ntok;
|
||||
d1.tg[0] = 256;
|
||||
@@ -383,9 +399,25 @@ int main(int argc, char** argv) {
|
||||
nb = true;
|
||||
gemv2(T(b + "attn_gate.weight"), XA, z, 6144, "gdn proj");
|
||||
nb = true;
|
||||
gemv2(T(b + "ssm_alpha.weight"), XA, al, 48, "gdn proj");
|
||||
nb = true;
|
||||
gemv2(T(b + "ssm_beta.weight"), XA, be, 48, "gdn proj");
|
||||
if (albe) {
|
||||
// ssm_alpha + ssm_beta in one dispatch (gemv_cat; bitwise the two GEMVs)
|
||||
const gguf::Tensor &ta = T(b + "ssm_alpha.weight"), &tb = T(b + "ssm_beta.weight");
|
||||
const std::string fn = "gemvc_" + std::to_string(ta.type) + "_32_1_4";
|
||||
mtl::Dispatch d;
|
||||
d.p = Pn(lib, fn, fn + "#" + std::to_string(ntok), mtl::Constants().i(0, ntok));
|
||||
GemvArgs ga{(uint32_t)N, 96, (uint32_t)ntok, (uint32_t)ref::row_bytes(ta.type, N), (uint32_t)N, 48};
|
||||
const uint32_t sp = 48;
|
||||
d.args = {mtl::Arg::B(W, ta.offset), mtl::Arg::B(XA.b, XA.plane_bytes() * xkind(ta.type)),
|
||||
mtl::Arg::B(XA.b, XA.plane_bytes() * 4), mtl::Arg::B(al), mtl::Arg::V(ga), mtl::Arg::B(W, tb.offset),
|
||||
mtl::Arg::B(be), mtl::Arg::V(sp)};
|
||||
d.grid[0] = 96 / 4;
|
||||
d.tg[0] = 128;
|
||||
push(d, "gdn proj", (double)ta.bytes + tb.bytes);
|
||||
} else {
|
||||
gemv2(T(b + "ssm_alpha.weight"), XA, al, 48, "gdn proj");
|
||||
nb = true;
|
||||
gemv2(T(b + "ssm_beta.weight"), XA, be, 48, "gdn proj");
|
||||
}
|
||||
GdnArgs a{16, 48, 10240, 1, 1e-6f, fuse ? (uint32_t)(MAXT * 6144) : 0u};
|
||||
mtl::Dispatch d;
|
||||
d.p = pgdn;
|
||||
@@ -459,8 +491,10 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
// ---- MoE
|
||||
gr_read(b + "hc_ffn", true, XF, false);
|
||||
gemv2(T(b + "ffn_gate_inp.weight"), XF, logits, 512, "router");
|
||||
nb = true;
|
||||
if (!routef) {
|
||||
gemv2(T(b + "ffn_gate_inp.weight"), XF, logits, 512, "router");
|
||||
nb = true;
|
||||
}
|
||||
const bool pair = fuse && std::getenv("CHAIN_PAIR"); // measured slower (+0.11 ms): opt-in
|
||||
if (pair) {
|
||||
// the shared expert's gate + up + SwiGLU in one dispatch, h raw into XH's raw plane
|
||||
@@ -507,11 +541,26 @@ int main(int argc, char** argv) {
|
||||
if (const char* e = std::getenv(dq2 ? "CHAIN_DN" : "CHAIN_DN4")) std::sscanf(e, "%d,%d,%d", &dl, &dr, &ds_);
|
||||
a.gx_gu = NFF / (gs * (32 / gl) * gr);
|
||||
a.gx_dn = N / (ds_ * (32 / dl) * dr);
|
||||
mtl::Dispatch r;
|
||||
r.p = proute;
|
||||
r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a)};
|
||||
r.tg[0] = 512;
|
||||
push(r, "router", 0);
|
||||
if (routef) {
|
||||
// the router GEMV + moe_route in one dispatch (router_route; bitwise), concurrent with the shared GEMVs
|
||||
const gguf::Tensor& tr = T(b + "ffn_gate_inp.weight");
|
||||
mtl::Dispatch r;
|
||||
r.p = Pn(lib, "router_route_30", "router_route#" + std::to_string(ntok), mtl::Constants().i(0, ntok));
|
||||
GemvArgs ga{(uint32_t)N, 512, (uint32_t)ntok, (uint32_t)ref::row_bytes(tr.type, N), (uint32_t)N, 512};
|
||||
r.args = {mtl::Arg::B(W, tr.offset), mtl::Arg::B(XF.b, 0), mtl::Arg::B(XF.b, XF.plane_bytes() * 4),
|
||||
mtl::Arg::B(logits), mtl::Arg::V(ga), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan),
|
||||
mtl::Arg::V(a), mtl::Arg::B(rcounter)};
|
||||
r.grid[0] = 512 / 16;
|
||||
r.tg[0] = 512;
|
||||
nb = true;
|
||||
push(r, "router", (double)tr.bytes);
|
||||
} else {
|
||||
mtl::Dispatch r;
|
||||
r.p = proute;
|
||||
r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a)};
|
||||
r.tg[0] = 512;
|
||||
push(r, "router", 0);
|
||||
}
|
||||
{
|
||||
GlueArgs ga{NFF, MAXT, N, TOPK};
|
||||
mtl::Dispatch d;
|
||||
@@ -557,7 +606,8 @@ int main(int argc, char** argv) {
|
||||
d.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sg), mtl::Arg::B(bo), mtl::Arg::V(a)};
|
||||
d.grid[0] = (ntok * N + 255) / 256;
|
||||
d.tg[0] = 256;
|
||||
if (!(fold_comb && hc_mode != "3")) push(d, "shared", 0); // else folded into the next HC read
|
||||
if (combf && (hc_mode == "3" || hc_mode == "q8_3")) comb_pending = true; // folded into the next gr_norm
|
||||
else if (!(fold_comb && hc_mode != "3")) push(d, "shared", 0); // else folded into the next HC read
|
||||
}
|
||||
}
|
||||
// ---- the final mixer and the head
|
||||
|
||||
Reference in new issue
Block a user