// metal/kernels/moe_mm.metal - routed experts for a PROMPT chunk (after quant/units/mma.metal): a grouped GEMM on
// the matrix units through Metal 4's MetalPerformancePrimitives matmul2d (on the M5 its neural accelerators).
//
// Plan (deterministic, three passes): the chunk's T x topk (token, expert) entries are sorted by expert, then by
// entry index; each expert's run is padded to a multiple of TM rows; tile i = (expert, first padded row).
// mm_hist per block of 1024 entries: a histogram of experts -> bcnt[block][expert]
// mm_scan one threadgroup (thread = expert): block offsets, the padded expert offsets, the tile list
// mm_place per block: an entry's padded row = expert offset + earlier blocks' count + rank in its block;
// srt_tok[row] = token (-1 on padding), srt_w[row] = weight, pos[entry] = row
// GEMMs: a threadgroup takes a tile (TM rows of one expert) and TN output rows; per K chunk of TK it converts the
// tile's activations to fp16 and dequantizes TN x TK weights (ref_quant's arithmetic, then rounded to fp16) into
// threadgroup memory, and matmul2d multiplies them, fp32 accumulation.
// mm_gu H[row][f] = silu(x . gate_f) * (x . up_f), fp32, for the n_ff rows
// mm_down YE[row][d] = h . down_d, fp32
// mm_combine y[t][d] = sum_k w(t, k) * YE[pos(t, k)][d], k in order
//
// The idea of feeding matmul2d tiles of dequantized weights from threadgroup memory follows yibie/strata-mlx's
// tile_matmul (MIT); the code is our own.
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
using namespace mpp::tensor_ops;
struct MmPlanArgs {
uint n_ent; // T * topk
uint topk;
uint n_expert;
uint tm; // the tile height (rows of one expert)
uint n_blocks; // ceil(n_ent / 1024)
};
// counts per block: bcnt[block * n_expert + e]
kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buffer(1)]],
constant MmPlanArgs& a [[buffer(2)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
threadgroup atomic_uint h[1024];
if (tid < a.n_expert) atomic_store_explicit(&h[tid], 0u, memory_order_relaxed);
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint i = tg.x * 1024 + tid;
if (i < a.n_ent && ids[i] >= 0) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid < a.n_expert)
bcnt[tg.x * a.n_expert + tid] = atomic_load_explicit(&h[tid], memory_order_relaxed);
}
// 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;
// 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.
// Padding rows [r0 + tot, r0 + rows) of srt_tok (when bound at buffer 6) are cleared to -1 here so mm_clear is not needed.
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)]], device int* srt_tok [[buffer(6)]],
ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]], ushort nsg [[simdgroups_per_threadgroup]]) {
threadgroup uint s_rows[32], s_tiles[32];
const uint e = tid;
uint tot = 0;
if (e < a.n_expert)
for (uint b = 0; b < a.n_blocks; ++b) {
boff[b * a.n_expert + e] = tot;
tot += bcnt[b * a.n_expert + e];
}
const uint nt = (tot + a.tm - 1) / a.tm, rows = nt * a.tm;
const uint pr = simd_prefix_exclusive_sum(rows), pt = simd_prefix_exclusive_sum(nt);
if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; }
threadgroup_barrier(mem_flags::mem_threadgroup);
uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0;
for (uint i = 0; i < nsg; ++i) {
if (i < sg) { orow += s_rows[i]; otile += s_tiles[i]; }
all_rows += s_rows[i];
all_tiles += s_tiles[i];
}
if (e < a.n_expert) {
const uint r0 = orow + pr, t0 = otile + pt;
eoff[e] = r0;
for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm));
if (srt_tok)
for (uint r = tot; r < rows; ++r) srt_tok[r0 + r] = -1;
}
if (tid == 0) {
meta[0] = all_tiles;
meta[1] = all_rows;
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;
}
}
static inline uint mm_sum_nibbles(uint c) {
c = (c & 0x0F0F0F0Fu) + ((c >> 4) & 0x0F0F0F0Fu);
c = (c + (c >> 8)) & 0x00FF00FFu;
return (c + (c >> 16)) & 0xFFFFu;
}
// grid (n_blocks), 1024 threads. Two-level deterministic intra-block rank: 32-lane sub-block nibble counts in p_cnt
// plus simd_broadcast scan within the caller's 32-lane simdgroup (no s_id[1024] threadgroup array).
kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]],
device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]],
device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]],
constant MmPlanArgs& a [[buffer(7)]],
uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
threadgroup atomic_uint p_cnt[4096];
const uint n4e = 4 * a.n_expert;
if (a.n_expert == 512) {
atomic_store_explicit(&p_cnt[tid], 0u, memory_order_relaxed);
atomic_store_explicit(&p_cnt[tid + 1024], 0u, memory_order_relaxed);
} else {
for (uint k = tid; k < n4e; k += 1024) atomic_store_explicit(&p_cnt[k], 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint i = tg * 1024 + tid;
const int e = i < a.n_ent ? ids[i] : -1;
const uint sub = tid >> 5;
if (e >= 0 && sub < 31) {
const uint idx = (sub >> 3) * a.n_expert + uint(e);
const uint sh = (sub & 7u) * 4u;
atomic_fetch_add_explicit(&p_cnt[idx], 1u << sh, memory_order_relaxed);
}
uint rank = 0;
#pragma clang loop unroll(full)
for (ushort l = 0; l < 32; ++l) {
const int el = simd_broadcast(e, l);
if (l < lane && el == e) ++rank;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (i >= a.n_ent) return;
if (e < 0) { pos[i] = -1; return; }
for (uint g = 0; g < (sub >> 3); ++g) {
const uint c = atomic_load_explicit(&p_cnt[g * a.n_expert + uint(e)], memory_order_relaxed);
rank += mm_sum_nibbles(c);
}
if (sub & 7u) {
const uint m = (1u << ((sub & 7u) * 4u)) - 1u;
const uint c = atomic_load_explicit(&p_cnt[(sub >> 3) * a.n_expert + uint(e)], memory_order_relaxed) & m;
rank += mm_sum_nibbles(c);
}
const uint row = eoff[e] + boff[tg * a.n_expert + e] + rank;
srt_tok[row] = int(i / a.topk);
srt_w[row] = w[i];
pos[i] = int(row);
}
kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) {
if (gid < n) srt_tok[gid] = -1;
}
#ifndef MM_GATHER_ONLY
// ---- the GEMMs -----------------------------------------------------------------------------------------------------
// weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride
struct MmArgs {
ulong gate_off;
ulong up_off;
ulong down_off;
ulong gu_stride;
ulong d_stride;
uint gu_row; // bytes per gate/up row (n_embd values)
uint d_row; // bytes per down row (n_ff values)
uint n_embd;
uint n_ff;
};
// TN x TK weights of `nrow`-strided rows into threadgroup half: thread-strided units of the format
template <int F, int TN, int TK>
static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threadgroup half* dst,
threadgroup const uint* tab, ushort tid, ushort nth) {
constexpr int UNIT = fmt<F>::UNIT, UPC = TK / UNIT, H = UNIT / 2;
static_assert(TK % UNIT == 0, "the K chunk must hold whole units of the format");
for (int it = tid; it < TN * UPC; it += nth) {
const int r = it / UPC, uu = it - r * UPC;
float wv[UNIT];
int p0, p1;
DW<F>::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab);
#pragma clang loop unroll(full)
for (int j = 0; j < H; ++j) {
dst[r * TK + p0 + j] = half(wv[j]);
dst[r * TK + p1 + j] = half(wv[H + j]);
}
}
}
#ifndef MM_V3_ONLY
template <int F, int TM, int TN, int TK, int SG>
kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]],
device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]],
device const int2* tiles [[buffer(4)]], constant MmArgs& a [[buffer(5)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
constexpr int NTH = SG * 32;
threadgroup half xa[TM * TK];
threadgroup half wg[TN * TK];
threadgroup half wu[TN * TK];
threadgroup uint tab[fmt<F>::TAB];
threadgroup int tok[TM];
load_tab<F>(tab, tid, NTH);
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN;
for (int m = tid; m < TM; m += NTH) tok[m] = srt_tok[row0 + m];
device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row;
device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false);
matmul2d<desc, execution_simdgroups<SG>> op;
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tA(xa, dextents<int32_t, 2>(TK, TM));
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tG(wg, dextents<int32_t, 2>(TK, TN));
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tU(wu, dextents<int32_t, 2>(TK, TN));
auto accg = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tG), float>();
auto accu = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tG), float>();
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
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) {
for (int it = tid; it < TM * TK; it += NTH) {
const int m = it / TK, k = it - m * TK;
xa[it] = tok[m] >= 0 ? half(x[ulong(tok[m]) * a.n_embd + k0 + k]) : half(0.0f);
}
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);
auto pg = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tG), float>();
auto pu = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tG), float>();
op.run(tA, tG, pg);
op.run(tA, tU, pu);
for (uint16_t i = 0; i < pg.get_capacity(); ++i)
if (pg.is_valid_element(i)) { accg[i] += pg[i]; accu[i] += pu[i]; }
threadgroup_barrier(mem_flags::mem_threadgroup);
}
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
if (accg.is_valid_element(i)) {
const float g = accg[i];
accg[i] = g / (1.0f + exp(-g)) * accu[i];
}
tensor<device float, dextents<int32_t, 2>, tensor_inline> tH(H, dextents<int32_t, 2>(a.n_ff, 1 << 30));
accg.store(tH.slice(f0, row0));
}
template <int F, int TM, int TN, int TK, int SG>
kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [[buffer(1)]],
device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]],
constant MmArgs& a [[buffer(4)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
constexpr int NTH = SG * 32;
threadgroup half ha[TM * TK];
threadgroup half wd[TN * TK];
threadgroup uint tab[fmt<F>::TAB];
load_tab<F>(tab, tid, NTH);
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN;
device const uchar* w0 = W + a.down_off + ulong(e) * a.d_stride + ulong(d0) * a.d_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false);
matmul2d<desc, execution_simdgroups<SG>> op;
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tA(ha, dextents<int32_t, 2>(TK, TM));
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tW(wd, dextents<int32_t, 2>(TK, TN));
auto acc = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tW), float>();
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
if (acc.is_valid_element(i)) acc[i] = 0.0f;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) {
for (int it = tid; it < TM * TK; it += NTH) {
const int m = it / TK, k = it - m * TK;
ha[it] = half(H[ulong(row0 + m) * a.n_ff + k0 + k]);
}
stage_w<F, TN, TK>(w0, a.d_row, k0, wd, tab, tid, NTH);
threadgroup_barrier(mem_flags::mem_threadgroup);
auto p = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tW), float>();
op.run(tA, tW, p);
for (uint16_t i = 0; i < p.get_capacity(); ++i)
if (p.is_valid_element(i)) acc[i] += p[i];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
tensor<device float, dextents<int32_t, 2>, tensor_inline> tY(YE, dextents<int32_t, 2>(a.n_embd, 1 << 30));
acc.store(tY.slice(d0, row0));
}
// grid (T * n_embd / 256), 256 threads
#endif
kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]],
device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]],
constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -)
uint gid [[thread_position_in_grid]]) {
const uint t = gid / a.y, d = gid % a.y;
if (t >= a.x) return;
float s = 0.0f;
for (uint k = 0; k < a.z; ++k) {
const int r = pos[t * a.z + k];
if (r >= 0) s = fma(srt_w[r], YE[ulong(r) * a.y + d], s);
}
y[gid] = s;
}
#ifndef MM_V3_ONLY
#define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu<F, M, N, K, S>) mm_gu<F, M, N, K, S>;
#define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down<F, M, N, K, S>) mm_down<F, M, N, K, S>;
#define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S)
#define MMDN(F, M, N, K, S) MMDN_(F, M, N, K, S)
#define MM_GU_SET(F) MMGU(F, 32, 64, 64, 4) MMGU(F, 64, 64, 64, 4) MMGU(F, 16, 64, 64, 2) MMGU(F, 32, 32, 64, 2)
#define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2)
MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S)
MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL)
#endif
#endif
// ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory ---------------------
// mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads.
kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]],
device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
const int t = srt_tok[tg.x];
device half* o = XS + ulong(tg.x) * n;
if (n == 2560) {
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
const uint k = tid + i * 256;
o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f);
}
} else {
for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f);
}
}
// mm_gather_h: gathers straight from fp16 rows (mixh) with 128-bit uint4 copies across nth threads per row.
kernel void mm_gather_h(device const half* xh [[buffer(0)]], device const int* srt_tok [[buffer(1)]],
device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]],
uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort nth [[threads_per_threadgroup]]) {
const int t = srt_tok[tg];
device uint4* o8 = (device uint4*)(XS + ulong(tg) * n);
const uint n8 = n >> 3;
if (n == 2560) {
device const uint4* x8 = (device const uint4*)(xh + ulong(max(t, 0)) * n);
if (nth == 32) {
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) o8[tid + i * 32] = t >= 0 ? x8[tid + i * 32] : uint4(0u);
return;
}
if (nth == 64) {
#pragma clang loop unroll(full)
for (uint i = 0; i < 5; ++i) o8[tid + i * 64] = t >= 0 ? x8[tid + i * 64] : uint4(0u);
return;
}
if (nth == 160) {
o8[tid] = t >= 0 ? x8[tid] : uint4(0u);
o8[tid + 160] = t >= 0 ? x8[tid + 160] : uint4(0u);
return;
}
if (nth == 256) {
o8[tid] = t >= 0 ? x8[tid] : uint4(0u);
if (tid < 64) o8[tid + 256] = t >= 0 ? x8[tid + 256] : uint4(0u);
return;
}
}
if (t >= 0) {
device const uint4* x8 = (device const uint4*)(xh + ulong(t) * n);
for (uint k = tid; k < n8; k += nth) o8[k] = x8[k];
} else {
for (uint k = tid; k < n8; k += nth) o8[k] = uint4(0u);
}
}
#ifndef MM_GATHER_ONLY
// The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W,
// W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot.
#ifndef STRATA_ARENA_METAL
#define STRATA_ARENA_METAL
constant bool kArenaFc [[function_constant(20)]];
constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc;
static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2,
device const uchar* W3) {
const uint c = v >> 28;
device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3));
return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul;
}
#endif
// 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;
#ifndef MM_V3_ONLY
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)]], device const uint* smap [[buffer(5)]],
device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
constexpr int NTH = SG * 32;
threadgroup half wg[TN * TK];
threadgroup half wu[TN * TK];
threadgroup float ot[TM * TN];
threadgroup uint tab[fmt<F>::TAB];
load_tab<F>(tab, tid, NTH);
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN;
device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride;
device const uchar* g0 = E + a.gate_off + ulong(f0) * a.gu_row;
device const uchar* u0 = E + a.up_off + ulong(f0) * a.gu_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<desc, execution_simdgroups<SG>> op;
tensor<device half, dextents<int32_t, 2>, tensor_inline> tX((device half*)XS, dextents<int32_t, 2>(a.n_embd, 1 << 30));
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tG(wg, dextents<int32_t, 2>(TK, TN));
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tU(wu, dextents<int32_t, 2>(TK, TN));
auto tA = tX.template slice<TK, TM>(0, row0);
auto accg = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tG), float>();
auto accu = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tG), float>();
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
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) {
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);
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)
if (accg.is_valid_element(i)) {
const float g = accg[i];
accg[i] = g / (1.0f + exp(-g)) * accu[i];
}
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot, dextents<int32_t, 2>(TN, TM));
accg.store(tO);
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int it = tid; it < TM * TN; it += NTH) {
const int m = it / TN, f = it - m * TN;
HS[ulong(row0 + m) * a.n_ff + f0 + f] = half(ot[it]);
}
}
template <int F, int TM, int TN, int TK, int SG>
kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]],
device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]],
constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]],
device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
constexpr int NTH = SG * 32;
threadgroup half wd[TN * TK];
threadgroup uint tab[fmt<F>::TAB];
load_tab<F>(tab, tid, NTH);
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN;
device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(d0) * a.d_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<desc, execution_simdgroups<SG>> op;
tensor<device half, dextents<int32_t, 2>, tensor_inline> tH((device half*)HS, dextents<int32_t, 2>(a.n_ff, 1 << 30));
tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tW(wd, dextents<int32_t, 2>(TK, TN));
auto tA = tH.template slice<TK, TM>(0, row0);
auto acc = op.template get_destination_cooperative_tensor<decltype(tA), decltype(tW), float>();
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
if (acc.is_valid_element(i)) acc[i] = 0.0f;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) {
stage_w<F, TN, TK>(w0, a.d_row, k0, wd, tab, tid, NTH);
threadgroup_barrier(mem_flags::mem_threadgroup);
tA = tH.template slice<TK, TM>(k0, row0);
op.run(tA, tW, acc);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
tensor<device float, dextents<int32_t, 2>, tensor_inline> tY(YE, dextents<int32_t, 2>(a.n_embd, 1 << 30));
acc.store(tY.slice(d0, row0));
}
#define MMGU2_(F, M, N, K, S) template [[host_name("mm_gu2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu2<F, M, N, K, S>) mm_gu2<F, M, N, K, S>;
#define MMDN2_(F, M, N, K, S) template [[host_name("mm_down2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down2<F, M, N, K, S>) mm_down2<F, M, N, K, S>;
#define MMGU2(F, M, N, K, S) MMGU2_(F, M, N, K, S)
#define MMDN2(F, M, N, K, S) MMDN2_(F, M, N, K, S)
#define MM_GU2_SET(F) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 32, 32, 32, 2) MMGU2(F, 32, 64, 16, 2) MMGU2(F, 64, 32, 16, 2) MMGU2(F, 16, 32, 32, 2) MMGU2(F, 16, 64, 16, 2) MMGU2(F, 16, 128, 32, 4) MMGU2(F, 32, 64, 16, 4) MMGU2(F, 16, 32, 32, 1) MMGU2(F, 16, 64, 32, 4) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 64, 4) MMGU2(F, 64, 64, 64, 4) MMGU2(F, 16, 32, 64, 2) MMGU2(F, 32, 64, 32, 4) MMGU2(F, 64, 64, 32, 4) MMGU2(F, 16, 64, 32, 2)
#define MM_DN2_SET(F) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 32, 32, 32, 2) MMDN2(F, 32, 64, 16, 2) MMDN2(F, 64, 32, 16, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 16, 64, 16, 2) MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 16, 4) MMDN2(F, 16, 32, 32, 1) MMDN2(F, 16, 64, 32, 4) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 32, 64, 64, 4) MMDN2(F, 64, 64, 64, 4) MMDN2(F, 16, 32, 64, 2) MMDN2(F, 32, 64, 32, 4) MMDN2(F, 64, 64, 32, 4) MMDN2(F, 16, 64, 32, 2)
#define MM_GU2_U32(F) MMGU2(F, 16, 64, 32, 2) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 16, 32, 32, 2) \
MMGU2(F, 64, 64, 32, 4)
#define MM_DN2_U32(F) MMDN2(F, 16, 64, 32, 2) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 32, 32, 32, 2) \
MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 64, 64, 32, 4)
// K-quants: the staged K chunk must hold whole units as DW places them (128 wide): TK 128
#define MM_GU2_KQ(F) MMGU2(F, 16, 32, 128, 2) MMGU2(F, 16, 16, 128, 2) MMGU2(F, 32, 32, 128, 2) MMGU2(F, 16, 32, 128, 4)
MM_GU2_KQ(T_Q4_K) MM_GU2_KQ(T_Q5_K) MM_GU2_KQ(T_Q6_K) MM_GU2_U32(T_Q8_0) MM_DN2_U32(T_Q5_1) MM_DN2_U32(T_Q8_0)
// IQ4_XS (UNIT 32): the K chunk must be >= 32
MMGU2(T_IQ4_XS, 16, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 64, 64, 2) MMGU2(T_IQ4_XS, 32, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 32, 32, 2)
MMGU2(T_IQ4_XS, 64, 64, 32, 4)
MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_XS) MM_GU2_SET(T_IQ2_S) MM_GU2_SET(T_IQ3_S)
MM_DN2_SET(T_Q2_0)
MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4)
#endif
// ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope),
// no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN.
// The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values
// at a time (W4<F>). Activations: the gathered fp16 rows from device memory.
#ifndef STRATA_COOP_K4
#define STRATA_COOP_K4
// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's
// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own.
constant bool COOP_K4 [[function_constant(7)]];
constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true;
#endif
template <int F> struct W4;
// the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16
// the threadgroup forms of quant.metal's field readers (blocks staged in threadgroup memory)
static inline uint rd32(threadgroup const uchar* p) {
threadgroup const ushort* s = (threadgroup const ushort*)p;
return uint(s[0]) | (uint(s[1]) << 16);
}
static inline float rdh(threadgroup const uchar* p) { return float(*(threadgroup const half*)p); }
static inline half4 w4_signed(uint g, uint s, float sc) {
const uchar4 v = as_type<uchar4>(g);
const float4 q = float4((s & 1) ? -float(v[0]) : float(v[0]), (s & 2) ? -float(v[1]) : float(v[1]),
(s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3]));
return half4(q * sc);
}
template <> struct W4<T_Q2_0> { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4)
// the codes c land in fp16 mantissas (1024 + c), so c - 1 and its product with d are exact in fp16: bit for bit the
// fp32 decode for every d and code byte (checked exhaustively), and 8-11% off mm_down3 at T 1024-2048
static inline half4 dec(uint q, half d) {
const uint2 c = (uint2(q | (q << 14), (q >> 4) | (q << 10)) & 0x00030003u) | 0x64006400u;
return (as_type<half4>(c) - 1025.0h) * d;
}
template <int OFF>
static inline half4 run_q(device const uchar* b, uint kl4) {
return dec((b + kl4)[OFF], *(device const half*)b);
}
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
return run_q<2>(row + (k >> 6) * 18, uint(k & 63) >> 2);
}
};
template <> struct W4<T_IQ3_XXS> {
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 98;
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
const uint aux = rd32(b + 2 + 64 + 4 * ib);
const float sc = rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f);
return w4_signed(g_iq3xxs(tab, b[2 + ((k & 255) >> 2)]), sign4_iq2xs((aux >> (7 * l)) & 127, hh), sc);
}
};
template <> struct W4<T_IQ2_XXS> {
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 66;
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4);
const uint idx = (a0 >> (8 * l)) & 255;
const float sc = rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f);
return w4_signed(g_iq2xxs(tab, idx, hh), sign4_iq2xs((a1 >> (7 * l)) & 127, hh), sc);
}
};
template <> struct W4<T_IQ2_XS> {
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 74;
const int ib = (k & 255) >> 5, hh = (k & 7) >> 2;
const uint v = ((device const ushort*)(b + 2))[(k & 255) >> 3];
const uint scb = b[66 + ib];
const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f);
return w4_signed(g_iq2xs(tab, v & 511, hh), sign4_iq2xs(v >> 9, hh), sc);
}
};
template <> struct W4<T_IQ2_S> {
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 82;
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
const uint qh = b[66 + ib], o = (k & 255) >> 3;
const uint idx = uint(b[2 + o]) | ((qh << (8 - 2 * l)) & 0x300);
const uint s = b[34 + o];
const uint scb = b[74 + ib];
const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f);
return w4_signed(g_iq2s(tab, idx, hh), hh ? s >> 4 : s, sc);
}
};
template <> struct W4<T_IQ3_S> {
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 110;
const int ib = (k & 255) >> 5, hh = (k & 7) >> 2;
const uint qh = b[66 + ib];
const uint i = uint(b[2 + ((k & 255) >> 2)]) | ((qh << (8 - ((k & 31) >> 2))) & 256u);
const uint s = b[74 + ((k & 255) >> 3)];
const uint scb = b[106 + (ib >> 1)];
const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15));
return w4_signed(g_iq3s(tab, i), hh ? s >> 4 : s, sc);
}
};
template <> struct W4<T_IQ4_XS> { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 8) * 136;
const int j = k & 255, ib = j >> 5, o = j & 15, hi = (j >> 4) & 1;
const uint sh = *(device const ushort*)(b + 2);
const int ls = int((b[4 + (ib >> 1)] >> (4 * (ib & 1))) & 15) | int(((sh >> (2 * ib)) & 3) << 4);
const float s = rdh(b) * float(ls - 32);
device const uchar* qs = b + 8 + 16 * ib;
float4 q;
for (int c = 0; c < 4; ++c) {
const half2 w = as_type<half2>(tab[qs[o + c]]);
q[c] = float(hi ? w.y : w.x);
}
return half4(q * s);
}
};
template <> struct W4<T_IQ4_NL> { // 32 a block: d, 16 bytes; low nibbles are elements 0..15, high 16..31
static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) {
device const uchar* b = row + (k >> 5) * 18;
const int j = k & 31, hi = j >> 4, o = j & 15;
const float d = rdh(b);
float4 q;
for (int c = 0; c < 4; ++c) {
const half2 w = as_type<half2>(tab[b[2 + o + c]]);
q[c] = float(hi ? w.y : w.x);
}
return half4(q * d);
}
};
// W8<F>: the runs k .. k + 3 and k + 16 .. k + 19 of a row (k % 32 < 16): one 32-group, so the per-group fields (scale,
// sign word) are read once. The generic one is two W4.
// ---- the K-quant / legacy-quant W4 decoders (the UD-Q4/5/6_K_XL packs' experts): values k .. k + 3 (k % 4 == 0)
// of a row in ref_quant's arithmetic (float), then fp16
template <> struct W4<T_Q8_0> { // 32 a block: d, int8 qs[32]
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
device const uchar* b = row + (k >> 5) * 34;
const float d = rdh(b);
device const char* q = (device const char*)(b + 2 + (k & 31));
return half4(float4(float(q[0]), float(q[1]), float(q[2]), float(q[3])) * d);
}
};
template <> struct W4<T_Q5_1> { // 32 a block: d, m, qh (32 bits), qs[16] (low nibbles 0..15, high 16..31)
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
device const uchar* b = row + (k >> 5) * 24;
const float d = rdh(b), m = rdh(b + 2);
const uint hb = rd32(b + 4);
const int j = k & 31, o = j & 15, hi = j >> 4;
float4 w;
for (int c = 0; c < 4; ++c) {
const uint qv = b[8 + o + c];
const int x = int(hi ? (qv >> 4) : (qv & 15)) | int(((hb >> (j + c)) & 1) << 4);
w[c] = float(x) * d + m;
}
return half4(w);
}
};
template <> struct W4<T_Q4_K> { // 256 a block: d, dmin, scales[12], qs[128] (64-chunks: low nibbles then high)
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
device const uchar* b = row + (k >> 8) * 144;
const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31;
float s, mm;
scale_min_k4(2 * c + hi, b + 4, s, mm);
const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm;
device const uchar* q = b + 16 + 32 * c + l;
float4 w;
for (int i = 0; i < 4; ++i) w[i] = d1 * float(hi ? (q[i] >> 4) : (q[i] & 15)) - m1;
return half4(w);
}
};
template <> struct W4<T_Q5_K> { // 256 a block: d, dmin, scales[12], qh[32], qs[128]
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
device const uchar* b = row + (k >> 8) * 176;
const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31;
float s, mm;
scale_min_k4(2 * c + hi, b + 4, s, mm);
const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm;
device const uchar* qh = b + 16 + l;
device const uchar* ql = b + 48 + 32 * c + l;
const uint bit = 1u << (2 * c + hi);
float4 w;
for (int i = 0; i < 4; ++i)
w[i] = d1 * float(int(hi ? (ql[i] >> 4) : (ql[i] & 15)) + ((qh[i] & bit) ? 16 : 0)) - m1;
return half4(w);
}
};
template <> struct W4<T_Q6_K> { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
device const uchar* b = row + (k >> 8) * 210;
const uint j = uint(k & 255), j4 = j >> 4;
const uint la = rd32(b + (((j >> 1) & 64u) | (j & 63u)));
const uint ha = rd32(b + 128u + (((j >> 2) & 32u) | (j & 31u)));
const float ds = rdh(b + 208) * float(((device const char*)(b + 192))[j4]);
const uint qa = ((la >> (j4 & 4u)) & 0x0F0F0F0Fu) | (((ha >> (j4 & 6u)) & 0x03030303u) << 4);
return half4(ds * (float4(as_type<uchar4>(qa)) - 32.0f));
}
};
template <int F> struct W8 {
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
a = W4<F>::run(row, k, tab);
b = W4<F>::run(row, k + 16, tab);
}
};
// (a W8<T_Q2_0> sharing the block scale was measured slower than the two W4 of the generic one: 5.5 vs 3.5-4.2 ms)
template <> struct W8<T_IQ3_XXS> {
template <typename P>
static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
const P bl = row + (k >> 8) * 98;
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
const uint aux = rd32(bl + 2 + 64 + 4 * ib);
const float sc = rdh(bl) * fma(float(aux >> 28), 0.5f, 0.25f);
const P qs = bl + 2 + ((k & 255) >> 2);
if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values)
a = w4_signed(qs[0] * 0x01010101u, aux, sc);
b = w4_signed(qs[4] * 0x01010101u, aux >> 7, sc);
return;
}
const uint j0 = (aux >> (7 * l)) & 127, j1 = (aux >> (7 * (l + 2))) & 127;
a = w4_signed(g_iq3xxs(tab, qs[0]), sign4_iq2xs(j0, hh), sc);
b = w4_signed(g_iq3xxs(tab, qs[4]), sign4_iq2xs(j1, hh), sc);
}
};
template <> struct W8<T_IQ2_XXS> {
template <typename P>
static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) {
const int ib = (k & 255) >> 5, l = (k & 31) >> 3;
const P p = bl + 2 + 8 * ib;
const uint a1 = rd32(p + 4);
const float sc = d * fma(float(a1 >> 28), 0.25f, 0.125f);
const uint i0 = p[l], i1 = p[l + 2];
const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127;
threadgroup const uint* th = tab + hh;
a = w4_signed(ctab ? uint(iq2xxs_grid[i0] >> (32 * hh)) : th[2 * i0], sign4_iq2xs(j0, hh), sc);
b = w4_signed(ctab ? uint(iq2xxs_grid[i1] >> (32 * hh)) : th[2 * i1], sign4_iq2xs(j1, hh), sc);
}
template <typename P>
static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
const P bl = row + (k >> 8) * 66;
block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b);
}
};
template <> struct W8<T_IQ2_XS> {
static inline void block(device const uchar* bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) {
const int ib = (k & 255) >> 5;
device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3);
const uint v0 = q[0], v1 = q[2];
const uint scb = bl[66 + ib];
const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1
const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f);
threadgroup const uint* th = tab + hh;
a = w4_signed(ctab ? uint(iq2xs_grid[v0 & 511] >> (32 * hh)) : th[2 * (v0 & 511)], sign4_iq2xs(v0 >> 9, hh), sc0);
b = w4_signed(ctab ? uint(iq2xs_grid[v1 & 511] >> (32 * hh)) : th[2 * (v1 & 511)], sign4_iq2xs(v1 >> 9, hh), sc1);
}
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
device const uchar* bl = row + (k >> 8) * 74;
block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b);
}
};
template <> struct W8<T_IQ2_S> {
// grid words read through tab + hh and the sign byte as s >> 4 * hh (not a select): 8% off mm_gu3 at T 2048;
// the same for IQ2_XXS / IQ2_XS, whose sign4_iq2xs select stays (a shift form there measured 10% slower)
template <typename P>
static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) {
const int ib = (k & 255) >> 5, l = (k & 31) >> 3;
threadgroup const uint* th = tab + hh;
const uint sh = uint(hh) << 2;
const uint qh = bl[66 + ib], scb = bl[74 + ib];
const P p = bl + 2 + ((k & 255) >> 3);
const uint i0 = uint(p[0]) | ((qh << (8 - 2 * l)) & 0x300);
const uint i1 = uint(p[2]) | ((qh << (4 - 2 * l)) & 0x300);
const uint s0 = p[32], s1 = p[34];
const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f);
const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f);
a = w4_signed(ctab ? uint(iq2s_grid[i0] >> (32 * hh)) : th[2 * i0], s0 >> sh, sc0);
b = w4_signed(ctab ? uint(iq2s_grid[i1] >> (32 * hh)) : th[2 * i1], s1 >> sh, sc1);
}
template <typename P>
static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
const P bl = row + (k >> 8) * 82;
block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b);
}
};
template <> struct W8<T_IQ3_S> {
template <typename P>
static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) {
const int ib = k >> 5;
const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)];
const float sc = d * fma(float((scb >> (4 * (ib & 1))) & 15), 2.0f, 1.0f);
const P qs = bl + 2 + (k >> 2);
const int sh0 = 8 - ((k & 31) >> 2);
const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u);
const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u);
const P sp = bl + 74 + (k >> 3);
const uint s0 = sp[0], s1 = sp[2];
if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values)
a = half4(half(float(i0 ^ s0) * sc));
b = half4(half(float(i1 ^ s1) * sc));
return;
}
a = w4_signed(g_iq3s(tab, i0), s0 >> (4 * hh), sc);
b = w4_signed(g_iq3s(tab, i1), s1 >> (4 * hh), sc);
}
template <typename P>
static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
if (mm_probe == 7) {
const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2;
const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15));
const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF;
const float scr = float(as_type<half>(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15));
const uint j0 = ((r.x >> (8 * (l & 1))) & 0xFF) | ((qhv << (8 - 2 * l)) & 256);
const uint j1 = ((r.x >> (16 + 8 * (l & 1))) & 0xFF) | ((qhv << (7 - 2 * l)) & 256);
a = w4_signed(tab[j0], (r.y >> (8 * hh)) & 0xFF, scr);
b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr);
return;
}
const P bl = row + (k >> 8) * 110;
block(bl, k & 255, rdh(bl), tab, (k & 7) >> 2, a, b);
}
};
// W8 for the K-quants and Q5_1: the lane's runs k and k + 16 (k % 32 < 16) share the block's scales / mins / high
// bits; the 4-byte code runs read as one uint where the block layout keeps them 4-aligned (Q4_K 144-, Q5_K 176-,
// Q5_1 24-byte blocks: rows of whole blocks start 4-aligned). The same per-value arithmetic as W4.
static inline float4 nib4(uint v, int hi) { const uint s = hi ? (v >> 4) : v; return float4(as_type<uchar4>(s & 0x0F0F0F0Fu)); }
template <> struct W8<T_Q4_K> {
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
device const uchar* bl = row + (k >> 8) * 144;
const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31;
float s, mm;
scale_min_k4(2 * c + hi, bl + 4, s, mm);
const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm;
device const uint* q = (device const uint*)(bl + 16 + 32 * c + l);
const float4 qa = nib4(q[0], hi), qb = nib4(q[4], hi);
a = half4(d1 * qa - m1);
b = half4(d1 * qb - m1);
}
};
template <> struct W8<T_Q5_K> {
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
device const uchar* bl = row + (k >> 8) * 176;
const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31;
float s, mm;
scale_min_k4(2 * c + hi, bl + 4, s, mm);
const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm;
device const uint* qh = (device const uint*)(bl + 16 + l);
device const uint* ql = (device const uint*)(bl + 48 + 32 * c + l);
const uint sh = uint(2 * c + hi);
const float4 ha = float4(as_type<uchar4>((qh[0] >> sh) & 0x01010101u)) * 16.0f;
const float4 hb = float4(as_type<uchar4>((qh[4] >> sh) & 0x01010101u)) * 16.0f;
a = half4(d1 * (nib4(ql[0], hi) + ha) - m1);
b = half4(d1 * (nib4(ql[4], hi) + hb) - m1);
}
};
template <> struct W8<T_Q5_1> {
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
device const uchar* bl = row + (k >> 5) * 24;
const float d = rdh(bl), m = rdh(bl + 2);
const uint hb = rd32(bl + 4);
const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones
const uint v = *(device const uint*)(bl + 8 + j);
float4 wa, wb;
for (int c = 0; c < 4; ++c) {
const uint qv = (v >> (8 * c)) & 0xFFu;
wa[c] = float(int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) * d + m;
wb[c] = float(int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) * d + m;
}
a = half4(wa);
b = half4(wb);
}
};
template <> struct W8<T_Q6_K> {
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
device const uchar* bl = row + (k >> 8) * 210;
const uint j = uint(k & 255), j4 = j >> 4; // l = j & 31 < 16; the pair at l + 16
device const uchar* ql = bl + (((j >> 1) & 64u) | (j & 63u));
device const uchar* qh = bl + 128u + (((j >> 2) & 32u) | (j & 31u));
const float d = rdh(bl + 208);
device const char* sc = (device const char*)(bl + 192u + j4);
const float dsa = d * float(sc[0]), dsb = d * float(sc[1]);
const uint la = rd32(ql), lb = rd32(ql + 16);
const uint ha = rd32(qh), hb = rd32(qh + 16);
const uint sl = j4 & 4u, shh = j4 & 6u;
const uint qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4);
const uint qb = ((lb >> sl) & 0x0F0F0F0Fu) | (((hb >> shh) & 0x03030303u) << 4);
a = half4(dsa * (float4(as_type<uchar4>(qa)) - 32.0f));
b = half4(dsb * (float4(as_type<uchar4>(qb)) - 32.0f));
}
};
template <int F, int TM, int TN, int TK, int SG>
kernel void mm_gu3(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)]], device const uint* smap [[buffer(5)]],
device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) {
threadgroup float ot[SG][TM * TN];
threadgroup uint tab[fmt<F>::TAB];
if (fmt<F>::TAB > 1) {
load_tab<F>(tab, tid, SG * 32);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN;
device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride;
device const uchar* g0 = E + a.gate_off + ulong(n0) * a.gu_row;
device const uchar* u0 = E + a.up_off + ulong(n0) * a.gu_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<desc, execution_simdgroup> op;
tensor<device half, dextents<int32_t, 2>, tensor_inline> tX((device half*)XS, dextents<int32_t, 2>(a.n_embd, 1 << 30));
auto rw = op.template get_right_input_cooperative_tensor<half, half, float>();
auto tA = tX.template slice<TK, TM>(0, row0);
auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rw)>, float>();
auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rw)>, float>();
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; }
if (coop_k4 && TK == 32 && mm_probe != 1) {
const uint16_t half_cap = rw.get_capacity() / 2;
const auto ix0 = rw.get_multidimensional_index(uint16_t(0));
const int kl = ix0[0], hh = (kl >> 2) & 1;
const ulong roff0 = ulong(ix0[1]) * a.gu_row;
device const uchar* g0_l = g0 + roff0;
device const uchar* u0_l = u0 + roff0;
const uint step8 = 8u * a.gu_row;
// IQ2 formats: a 256 super-block at a time, its d read once; the lane's hh is fixed for the whole K loop
constexpr bool sb = F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS;
for (int kb = 0; kb < int(a.n_embd); kb += (sb ? 256 : TK)) {
const uint sboff = uint(kb >> 8) * fmt<F>::BB;
float dg[TN * TK / 256], du[TN * TK / 256];
if (sb) {
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap / 4; ++i) {
dg[i] = rdh(g0_l + i * step8 + sboff);
du[i] = rdh(u0_l + i * step8 + sboff);
}
}
for (int ki = 0; ki < (sb ? 256 : TK); ki += TK) {
const int k0 = kb + ki;
tA = tX.template slice<TK, TM>(k0, row0);
const uint boff = F == T_IQ2_XS ? sboff : (fmt<F>::QK == 256 ? uint(k0 >> 8) * fmt<F>::BB : (fmt<F>::QK == 64 ? uint(k0 >> 6) * fmt<F>::BB : 0u));
const int j_in = fmt<F>::QK == 256 ? ((k0 + kl) & 255) : (fmt<F>::QK == 64 ? ((k0 + kl) & 63) : (k0 + kl));
device const uchar* rgw = fmt<F>::QK >= 64 ? g0_l + boff : g0_l;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4, rgw += step8) {
half4 wg, wg2;
if (F == T_IQ2_S) W8<T_IQ2_S>::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2);
else if (F == T_IQ2_XXS) W8<T_IQ2_XXS>::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2);
else if (F == T_IQ2_XS) W8<T_IQ2_XS>::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2);
else if (F == T_IQ3_S) W8<T_IQ3_S>::block(rgw, j_in, rdh(rgw), tab, hh, wg, wg2);
else W8<F>::run(rgw, j_in, tab, wg, wg2);
rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w;
const uint16_t j = i + half_cap;
rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w;
}
op.run(tA, rw, accg);
device const uchar* ruw = fmt<F>::QK >= 64 ? u0_l + boff : u0_l;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4, ruw += step8) {
half4 wu, wu2;
if (F == T_IQ2_S) W8<T_IQ2_S>::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2);
else if (F == T_IQ2_XXS) W8<T_IQ2_XXS>::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2);
else if (F == T_IQ2_XS) W8<T_IQ2_XS>::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2);
else if (F == T_IQ3_S) W8<T_IQ3_S>::block(ruw, j_in, rdh(ruw), tab, hh, wu, wu2);
else W8<F>::run(ruw, j_in, tab, wu, wu2);
rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w;
const uint16_t j = i + half_cap;
rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w;
}
op.run(tA, rw, accu);
}
}
} else
for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) {
const uint16_t cap = rw.get_capacity(), half_cap = cap / 2;
tA = tX.template slice<TK, TM>(k0, row0);
if (!coop_k4) {
for (uint16_t i = 0; i < cap; ++i) {
const auto ix = rw.get_multidimensional_index(i);
const int k = k0 + ix[0], n = ix[1];
rw[i] = W4<F>::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3];
}
op.run(tA, rw, accg);
for (uint16_t i = 0; i < cap; ++i) {
const auto ix = rw.get_multidimensional_index(i);
const int k = k0 + ix[0], n = ix[1];
rw[i] = W4<F>::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3];
}
op.run(tA, rw, accu);
} else {
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4) {
const auto ix = rw.get_multidimensional_index(i);
const auto iy = rw.get_multidimensional_index(i + half_cap);
const int k = k0 + ix[0], n = ix[1];
device const uchar* rgw = g0 + ulong(n) * a.gu_row;
half4 wg, wg2;
if (mm_probe == 1) {
wg = half4(half(n & 7)); wg2 = half4(half(k & 7));
} else if (TK == 32) {
W8<F>::run(rgw, k, tab, wg, wg2);
} else {
const int k2 = k0 + iy[0];
wg = W4<F>::run(rgw, k, tab);
wg2 = W4<F>::run(g0 + ulong(iy[1]) * a.gu_row, k2, tab);
}
rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w;
const uint16_t j = i + half_cap;
rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w;
}
op.run(tA, rw, accg);
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4) {
const auto ix = rw.get_multidimensional_index(i);
const auto iy = rw.get_multidimensional_index(i + half_cap);
const int k = k0 + ix[0], n = ix[1];
device const uchar* ruw = u0 + ulong(n) * a.gu_row;
half4 wu, wu2;
if (mm_probe == 1) {
wu = half4(half(n & 7)); wu2 = half4(half(k & 7));
} else if (TK == 32) {
W8<F>::run(ruw, k, tab, wu, wu2);
} else {
const int k2 = k0 + iy[0];
wu = W4<F>::run(ruw, k, tab);
wu2 = W4<F>::run(u0 + ulong(iy[1]) * a.gu_row, k2, tab);
}
rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w;
const uint16_t j = i + half_cap;
rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w;
}
op.run(tA, rw, accu);
}
}
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
if (accg.is_valid_element(i)) {
const float g = accg[i];
accg[i] = g / (1.0f + exp(-g)) * accu[i];
}
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot[sg], dextents<int32_t, 2>(TN, TM));
accg.store(tO);
simdgroup_barrier(mem_flags::mem_threadgroup);
if ((F == T_IQ2_XS || F == T_IQ3_XXS || F == T_IQ2_XXS || F == T_Q2_0 || F == T_IQ2_S || F == T_IQ3_S) && 128 % TN == 0 && (TM * TN) % 128 == 0) {
// a lane stores 4 adjacent outputs a pass, 128 a pass over the simdgroup
constexpr int RP = 128 / TN; // rows a pass
const int mr = lane / (TN / 4), f = (lane % (TN / 4)) * 4;
threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]);
device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f;
const uint hs_step = uint(RP) * a.n_ff;
#pragma clang loop unroll(full)
for (int m = 0; m < TM / RP; ++m, hs_p += hs_step)
*(device half4*)hs_p = half4(ot_p[m * (RP * TN / 4)]);
} else
for (int it = lane; it < TM * TN; it += 32) {
const int m = it / TN, f = it - m * TN;
HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]);
}
}
template <int F, int TM, int TN, int TK, int SG>
kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]],
device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]],
device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort tid [[thread_index_in_threadgroup]]) {
threadgroup uint tab[fmt<F>::TAB];
if (fmt<F>::TAB > 1) {
load_tab<F>(tab, tid, SG * 32);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN;
device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<desc, execution_simdgroup> op;
tensor<device half, dextents<int32_t, 2>, tensor_inline> tH((device half*)HS, dextents<int32_t, 2>(a.n_ff, 1 << 30));
auto rw = op.template get_right_input_cooperative_tensor<half, half, float>();
auto tA = tH.template slice<TK, TM>(0, row0);
auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rw)>, float>();
for (uint16_t i = 0; i < acc.get_capacity(); ++i)
if (acc.is_valid_element(i)) acc[i] = 0.0f;
const auto ix0 = rw.get_multidimensional_index(uint16_t(0));
const int kl = ix0[0];
device const uchar* w0_l = w0 + ulong(ix0[1]) * a.d_row;
const uint step8 = 8u * a.d_row;
if (coop_k4 && F == T_Q2_0 && TK == 32) {
const uint16_t half_cap = rw.get_capacity() / 2;
const uint kl4 = uint(kl) >> 2;
for (int k0 = 0; k0 < int(a.n_ff); k0 += 64, w0_l += 18) {
device const uchar* rdw = w0_l;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) {
const half4 w = W4<T_Q2_0>::run_q<2>(rdw, kl4);
const half4 w2 = W4<T_Q2_0>::run_q<6>(rdw, kl4);
rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w;
const uint16_t j = i + half_cap;
rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w;
}
tA = tH.template slice<TK, TM>(k0, row0);
op.run(tA, rw, acc);
rdw = w0_l;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) {
const half4 w = W4<T_Q2_0>::run_q<10>(rdw, kl4);
const half4 w2 = W4<T_Q2_0>::run_q<14>(rdw, kl4);
rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w;
const uint16_t j = i + half_cap;
rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w;
}
tA = tH.template slice<TK, TM>(k0 + 32, row0);
op.run(tA, rw, acc);
}
} else
for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) {
if (!coop_k4) {
for (uint16_t i = 0; i < rw.get_capacity(); ++i) {
const auto ix = rw.get_multidimensional_index(i);
const int k = k0 + ix[0];
rw[i] = W4<F>::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3];
}
} else if (TK == 32) {
// the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair
const uint16_t half_cap = rw.get_capacity() / 2;
const uint boff = fmt<F>::QK == 64 ? uint(k0 >> 6) * fmt<F>::BB : (fmt<F>::QK == 32 ? uint(k0 >> 5) * fmt<F>::BB : 0u);
const int j_in = fmt<F>::QK == 64 ? ((k0 + kl) & 63) : (fmt<F>::QK == 32 ? kl : (k0 + kl));
device const uchar* rdw = w0_l + boff;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) {
half4 w, w2;
W8<F>::run(rdw, j_in, tab, w, w2);
rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w;
const uint16_t j = i + half_cap;
rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w;
}
} else
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < rw.get_capacity(); i += 4) {
const auto ix = rw.get_multidimensional_index(i);
const half4 w = W4<F>::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab);
rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w;
}
tA = tH.template slice<TK, TM>(k0, row0);
op.run(tA, rw, acc);
}
tensor<device float, dextents<int32_t, 2>, tensor_inline> tY(YE, dextents<int32_t, 2>(a.n_embd, 1 << 30));
acc.store(tY.slice(n0, row0));
}
#define MMGU3_(F, M, N, K, S) template [[host_name("mm_gu3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu3<F, M, N, K, S>) mm_gu3<F, M, N, K, S>;
#define MMDN3_(F, M, N, K, S) template [[host_name("mm_down3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down3<F, M, N, K, S>) mm_down3<F, M, N, K, S>;
#define MMGU3(F, M, N, K, S) MMGU3_(F, M, N, K, S)
#define MMDN3(F, M, N, K, S) MMDN3_(F, M, N, K, S)
#define MM3_ONE(F, M, N, K, S) MMGU3(F, M, N, K, S) MMDN3(F, M, N, K, S)
#define MM3_GU(F) MMGU3(F, 16, 64, 32, 2) MMGU3(F, 16, 32, 32, 4) MMGU3(F, 32, 64, 32, 2) MMGU3(F, 16, 64, 64, 2) \
MMGU3(F, 16, 32, 16, 2) MMGU3(F, 32, 32, 32, 2) MMGU3(F, 16, 64, 32, 1) MMGU3(F, 16, 64, 32, 4) MMGU3(F, 16, 32, 16, 4) MMGU3(F, 16, 32, 32, 5) MMGU3(F, 16, 32, 16, 5) MMGU3(F, 16, 16, 32, 4) MMGU3(F, 16, 16, 32, 8) MMGU3(F, 16, 16, 16, 8)
#define MM3_DN(F) MMDN3(F, 16, 64, 32, 2) MMDN3(F, 16, 32, 32, 4) MMDN3(F, 32, 64, 32, 2) MMDN3(F, 16, 64, 64, 2) \
MMDN3(F, 16, 32, 16, 2) MMDN3(F, 32, 32, 32, 2) MMDN3(F, 16, 64, 32, 1) MMDN3(F, 16, 64, 32, 4) MMDN3(F, 16, 64, 16, 2) MMDN3(F, 16, 32, 16, 4) MMDN3(F, 16, 64, 32, 5) MMDN3(F, 16, 64, 16, 4)
MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S)
MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0)
#ifndef MM_V3_ONLY
// mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token
// tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times
// as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same
// per-row arithmetic as mm_gu3 (each 16-row op is mm_gu3's): bitwise for the rows both compute. COOP_K4 layout only.
template <int F, int MS, int TN, int TK, int SG>
kernel void mm_gu4(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)]],
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 TM = 16;
threadgroup float ot[SG][TM * TN];
threadgroup uint tab[fmt<F>::TAB];
load_tab<F>(tab, tid, SG * 32);
threadgroup_barrier(mem_flags::mem_threadgroup);
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN;
device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row;
device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<desc, execution_simdgroup> op;
tensor<device half, dextents<int32_t, 2>, tensor_inline> tX((device half*)XS, dextents<int32_t, 2>(a.n_embd, 1 << 30));
auto rg = op.template get_right_input_cooperative_tensor<half, half, float>();
auto ru = op.template get_right_input_cooperative_tensor<half, half, float>();
auto tA0 = tX.template slice<TK, TM>(0, row0);
auto tA1 = tX.template slice<TK, TM>(0, row0 + TM);
#define MM4_DST op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA0)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rg)>, float>()
auto accg0 = MM4_DST;
auto accu0 = MM4_DST;
auto accg1 = MM4_DST;
auto accu1 = MM4_DST;
#undef MM4_DST
using DT = decltype(accg0);
for (uint16_t i = 0; i < accg0.get_capacity(); ++i)
if (accg0.is_valid_element(i)) { accg0[i] = 0.0f; accu0[i] = 0.0f; accg1[i] = 0.0f; accu1[i] = 0.0f; }
for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) {
const uint16_t cap = rg.get_capacity(), half_cap = cap / 2;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4) {
const auto ix = rg.get_multidimensional_index(i);
const int k = k0 + ix[0], n = ix[1];
device const uchar* rgw = g0 + ulong(n) * a.gu_row;
device const uchar* ruw = u0 + ulong(n) * a.gu_row;
half4 wg, wu, wg2, wu2;
W8<F>::run(rgw, k, tab, wg, wg2);
W8<F>::run(ruw, k, tab, wu, wu2);
rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w;
ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w;
const uint16_t j = i + half_cap;
rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w;
ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w;
}
tA0 = tX.template slice<TK, TM>(k0, row0);
op.run(tA0, rg, accg0);
op.run(tA0, ru, accu0);
if (MS > 1) {
tA1 = tX.template slice<TK, TM>(k0, row0 + TM);
op.run(tA1, rg, accg1);
op.run(tA1, ru, accu1);
}
}
for (int s = 0; s < MS; ++s) {
thread DT& ag = s == 0 ? accg0 : accg1;
thread DT& au = s == 0 ? accu0 : accu1;
for (uint16_t i = 0; i < ag.get_capacity(); ++i)
if (ag.is_valid_element(i)) {
const float g = ag[i];
ag[i] = g / (1.0f + exp(-g)) * au[i];
}
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot[sg], dextents<int32_t, 2>(TN, TM));
simdgroup_barrier(mem_flags::mem_threadgroup);
ag.store(tO);
simdgroup_barrier(mem_flags::mem_threadgroup);
for (int it = lane; it < TM * TN; it += 32) {
const int m = it / TN, f = it - m * TN;
HS[ulong(row0 + s * TM + m) * a.n_ff + n0 + f] = half(ot[sg][it]);
}
}
}
#define MMGU4_(F, M, N, K, S) template [[host_name("mm_gu4_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu4<F, M, N, K, S>) mm_gu4<F, M, N, K, S>;
#define MMGU4(F, M, N, K, S) MMGU4_(F, M, N, K, S)
#define MM4_GU(F) MMGU4(F, 2, 32, 32, 4) MMGU4(F, 2, 32, 32, 2) MMGU4(F, 2, 16, 32, 4) MMGU4(F, 1, 32, 32, 4)
MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2_S) MM4_GU(T_IQ3_S) MM4_GU(T_IQ4_XS)
// mm_gu5: mm_gu3 (COOP_K4 layout, TK 32) with the threadgroup's weight rows staged block by block (QK = 256 values a
// row, gate and up) into threadgroup memory by coalesced ushort loads, the lanes decoding from there: the decode's
// many small scattered device loads (2-6 a run) become threadgroup reads. Same decoded values and the same matrix
// ops as mm_gu3: bitwise. Formats with 256-value blocks and a templated W8 (IQ3_S, IQ3_XXS, IQ2_S).
template <int F, int TN, int SG>
kernel void mm_gu5(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)]],
device const uint* smap [[buffer(5)]], device const uchar* W1 [[buffer(6)]],
device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]],
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 TM = 16, TK = 32, BB = fmt<F>::BB, ROWS = SG * TN, RU = BB / 2;
threadgroup float ot[SG][TM * TN];
threadgroup uint tab[fmt<F>::TAB];
threadgroup ushort stg[2][ROWS * RU];
load_tab<F>(tab, tid, SG * 32);
const int2 tl = tiles[tg.y];
const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN;
device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride;
device const uchar* gt = E + a.gate_off + ulong(nt0) * a.gu_row;
device const uchar* ut = E + a.up_off + ulong(nt0) * a.gu_row;
constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<desc, execution_simdgroup> op;
tensor<device half, dextents<int32_t, 2>, tensor_inline> tX((device half*)XS, dextents<int32_t, 2>(a.n_embd, 1 << 30));
auto rg = op.template get_right_input_cooperative_tensor<half, half, float>();
auto ru = op.template get_right_input_cooperative_tensor<half, half, float>();
auto tA = tX.template slice<TK, TM>(0, row0);
auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rg)>, float>();
auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t<decltype(tA)>, __tensor_ops_detail::__remove_addrspace_t<decltype(rg)>, float>();
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; }
const int NB = int(a.n_embd) / 256;
for (int kb = 0; kb < NB; ++kb) {
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int idx = tid; idx < ROWS * RU; idx += SG * 32) {
const int r = idx / RU, w = idx - r * RU;
stg[0][idx] = ((device const ushort*)(gt + ulong(r) * a.gu_row + kb * BB))[w];
stg[1][idx] = ((device const ushort*)(ut + ulong(r) * a.gu_row + kb * BB))[w];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int ks = 0; ks < 8; ++ks) {
const int k0 = kb * 256 + ks * TK;
const uint16_t half_cap = rg.get_capacity() / 2;
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < half_cap; i += 4) {
const auto ix = rg.get_multidimensional_index(i);
const int kl = ks * TK + ix[0], n = sg * TN + ix[1];
// the block's row in the stage; W8 indexes block (k >> 8) = 0 of it with k = kl
threadgroup const uchar* rgw = (threadgroup const uchar*)(stg[0] + n * RU);
threadgroup const uchar* ruw = (threadgroup const uchar*)(stg[1] + n * RU);
half4 wg, wu, wg2, wu2;
W8<F>::run(rgw, kl, tab, wg, wg2);
W8<F>::run(ruw, kl, tab, wu, wu2);
rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w;
ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w;
const uint16_t j = i + half_cap;
rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w;
ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w;
}
tA = tX.template slice<TK, TM>(k0, row0);
op.run(tA, rg, accg);
op.run(tA, ru, accu);
}
}
for (uint16_t i = 0; i < accg.get_capacity(); ++i)
if (accg.is_valid_element(i)) {
const float g = accg[i];
accg[i] = g / (1.0f + exp(-g)) * accu[i];
}
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot[sg], dextents<int32_t, 2>(TN, TM));
accg.store(tO);
simdgroup_barrier(mem_flags::mem_threadgroup);
for (int it = lane; it < TM * TN; it += 32) {
const int m = it / TN, f = it - m * TN;
HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]);
}
}
#define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5<F, N, S>) mm_gu5<F, N, S>;
#define MMGU5(F, N, S) MMGU5_(F, N, S)
MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1)
MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4)
#endif
#endif