200 lines
8.0 KiB
Metal
200 lines
8.0 KiB
Metal
// MTPLX 9882703, MIT; see MTPLX-LICENSE.txt.
|
|
|
|
constexpr int BM = 64;
|
|
constexpr int BN = 64;
|
|
constexpr int BK = 64;
|
|
constexpr int WN = 2;
|
|
constexpr short SM = 32;
|
|
constexpr short SN = 32;
|
|
constexpr short SK = 32;
|
|
constexpr short TM = SM / 16;
|
|
constexpr short TN = SN / 16;
|
|
constexpr short TK = SK / 16;
|
|
constexpr int PACK = get_pack_factor<BITS, 8>();
|
|
constexpr int PACK_BYTES = get_bytes_per_pack<BITS>();
|
|
constexpr int BK_PAD = BK + 16 / sizeof(T);
|
|
constexpr int ROW_BYTES = KD * PACK_BYTES / PACK;
|
|
constexpr int ROW_GROUPS = KD / GS;
|
|
constexpr int OUT_COLS = OUTC;
|
|
// SWIGLU: the tile's first 32 weight rows are gate rows, the last 32 the
|
|
// matching up rows, each half loaded by 64 threads from its own row range.
|
|
using WeightLoader = metal::conditional_t<
|
|
SWIGLU,
|
|
QuantizedBlockLoader<T, BN / 2, BK, BK_PAD, 1, 64, GS, BITS>,
|
|
QuantizedBlockLoader<T, BN, BK, BK_PAD, 1, 128, GS, BITS>>;
|
|
|
|
threadgroup T w_tile[BN * BK_PAD];
|
|
|
|
// The tile this threadgroup owns, and the one expert whose rows it holds:
|
|
// the last expert whose first tile is at or before it.
|
|
const int tile = int(threadgroup_position_in_grid.y);
|
|
if (tile >= tile_start[ED]) {
|
|
return;
|
|
}
|
|
int lo = 0;
|
|
int hi = ED - 1;
|
|
while (lo < hi) {
|
|
const int mid = (lo + hi + 1) >> 1;
|
|
if (tile_start[mid] <= tile) {
|
|
lo = mid;
|
|
} else {
|
|
hi = mid - 1;
|
|
}
|
|
}
|
|
const int expert = lo;
|
|
const int row0 = row_start[expert] + (tile - tile_start[expert]) * BM;
|
|
const int tile_rows = min(BM, row_start[expert + 1] - row0);
|
|
const int col0 = int(threadgroup_position_in_grid.x) * BN;
|
|
|
|
const ushort sg = ushort(simdgroup_index_in_threadgroup);
|
|
const ushort lane = ushort(thread_index_in_simdgroup);
|
|
const short tm = SM * short(sg / WN);
|
|
const short tn = SN * short(sg % WN);
|
|
// Rows of this tile that fall in this simdgroup's 32-row half (0 to 32).
|
|
const short rows = short(clamp(tile_rows - int(tm), 0, int(SM)));
|
|
|
|
ushort load_sg = sg;
|
|
const device uint8_t* w_src = (const device uint8_t*)w;
|
|
const device T* s_src = scales;
|
|
const device T* b_src = biases;
|
|
size_t w_row = size_t(expert) * ND + size_t(col0);
|
|
threadgroup T* w_dst = w_tile;
|
|
if (SWIGLU) {
|
|
// Output columns [col0 / 2, col0 / 2 + 32): the matching gate rows
|
|
// of w and up rows of wu (the same array at an offset when the
|
|
// weight is a fused [gate | up]).
|
|
const int up_half = int(load_sg / 2);
|
|
const size_t c = size_t(col0 / 2);
|
|
if (up_half) {
|
|
w_src = (const device uint8_t*)wu;
|
|
s_src = scales_u;
|
|
b_src = biases_u;
|
|
w_row = size_t(expert) * UD + size_t(UOFF) + c;
|
|
} else {
|
|
w_row = size_t(expert) * ND + c;
|
|
}
|
|
w_dst = w_tile + up_half * (BN / 2) * BK_PAD;
|
|
load_sg = load_sg % 2;
|
|
}
|
|
thread WeightLoader loader(
|
|
w_src + w_row * ROW_BYTES,
|
|
s_src + w_row * ROW_GROUPS,
|
|
b_src + w_row * ROW_GROUPS,
|
|
KD,
|
|
w_dst,
|
|
load_sg,
|
|
lane);
|
|
|
|
NAXTile<float, TM, TN> acc;
|
|
acc.clear();
|
|
|
|
// A lane holds rows (y, y + 8) and four consecutive columns from x of every
|
|
// 16x16 fragment (MLX's fragment coordinates). Its four token rows are
|
|
// fixed for the whole K loop; rows past the tile read as zero, as MLX's
|
|
// bounded load does.
|
|
const short2 coord = BaseNAXFrag::get_coord();
|
|
const device T* lane_rows[TM][2];
|
|
bool lane_live[TM][2];
|
|
for (short i = 0; i < TM; i++) {
|
|
for (short r = 0; r < 2; r++) {
|
|
const short local = i * 16 + coord.y + r * 8;
|
|
lane_live[i][r] = local < rows;
|
|
const int token = lane_live[i][r] ? int(row_map[row0 + tm + local]) : 0;
|
|
lane_rows[i][r] = x + size_t(token) * KD + coord.x;
|
|
}
|
|
}
|
|
|
|
// The K loop, compiled twice: for a simdgroup whose 32 rows are all live
|
|
// (no bounds) and for a partial or empty half.
|
|
dispatch_bool(rows == SM, [&](auto full) {
|
|
for (int k = 0; k < KD / BK; k++) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
loader.load_unsafe();
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
STEEL_PRAGMA_NO_UNROLL
|
|
for (short kk = 0; kk < BK; kk += SK) {
|
|
if (decltype(full)::value || rows > 0) {
|
|
NAXTile<T, TM, TK> a;
|
|
NAXTile<T, TN, TK> b;
|
|
volatile int keep_order;
|
|
const int col = k * BK + kk;
|
|
STEEL_PRAGMA_UNROLL
|
|
for (short i = 0; i < TM; i++) {
|
|
STEEL_PRAGMA_UNROLL
|
|
for (short j = 0; j < TK; j++) {
|
|
thread auto& frag = a.frag_at(i, j);
|
|
STEEL_PRAGMA_UNROLL
|
|
for (short r = 0; r < 2; r++) {
|
|
const device T* src = lane_rows[i][r] + col + j * 16;
|
|
STEEL_PRAGMA_UNROLL
|
|
for (short c = 0; c < 4; c++) {
|
|
if constexpr (decltype(full)::value) {
|
|
frag[r * 4 + c] = src[c];
|
|
} else {
|
|
frag[r * 4 + c] = lane_live[i][r] ? src[c] : T(0);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
b.template load<T, BK_PAD, 1>(w_tile + tn * BK_PAD + kk);
|
|
tile_matmad_nax(
|
|
acc,
|
|
a,
|
|
metal::bool_constant<false>{},
|
|
b,
|
|
metal::bool_constant<true>{});
|
|
(void)keep_order;
|
|
}
|
|
}
|
|
loader.next();
|
|
}
|
|
});
|
|
|
|
if constexpr (SWIGLU) {
|
|
// The up simdgroup of each 32-row half hands its values, rounded to T
|
|
// as the stock gather's output is, to the gate simdgroup of the same
|
|
// rows through the (now free) weight tile, lane to lane. The gate
|
|
// simdgroup writes silu(gate) * up with MLX's own functors, rounding to
|
|
// T after each op as the stock chain does.
|
|
constexpr short VALS = TM * TN * 8;
|
|
static_assert(2 * 32 * VALS <= BN * BK_PAD, "the hand-off must fit the weight tile");
|
|
threadgroup T* handoff = w_tile + (sg / 2) * (32 * VALS) + lane * VALS;
|
|
const thread float* vals = acc.elems();
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (tn != 0) {
|
|
for (short e = 0; e < VALS; e++) {
|
|
handoff[e] = static_cast<T>(vals[e]);
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (tn == 0 && rows > 0) {
|
|
device T* out = y + size_t(row0 + tm) * OUT_COLS + size_t(col0 / 2);
|
|
for (short i = 0; i < TM; i++) {
|
|
for (short j = 0; j < TN; j++) {
|
|
for (short r = 0; r < 2; r++) {
|
|
const short row = i * 16 + coord.y + r * 8;
|
|
if (row >= rows) {
|
|
continue;
|
|
}
|
|
for (short c = 0; c < 4; c++) {
|
|
const short e = (i * TN + j) * 8 + r * 4 + c;
|
|
const T gate = static_cast<T>(vals[e]);
|
|
const T act = Multiply()(gate, Sigmoid()(gate));
|
|
out[size_t(row) * OUT_COLS + j * 16 + coord.x + c] =
|
|
Multiply()(act, handoff[e]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} else if (rows > 0) {
|
|
device T* out = y + size_t(row0 + tm) * ND + size_t(col0 + tn);
|
|
if (rows == SM) {
|
|
acc.store(out, ND);
|
|
} else {
|
|
acc.store_safe(out, ND, short2(SN, rows));
|
|
}
|
|
}
|