// quantized_matmul_id_mm.metal — MoE-routed GGML-quantized matrix-matrix kernels.
//
// Implements the `kernel_mul_mm_id_<qtype>_f32` kernel shape and its
// preprocessing helper `kernel_mul_mm_id_map0_ne20_<N>` for mlx-native.
//
// Why mm_id? The MoE variant of quantized_matmul needs to run the same
// expert weight tile against many tokens that routed to that expert. The
// existing `kernel_mul_mv_id_<qtype>_f32` re-reads each expert's weight
// blocks once per routed (token, slot) pair. The mm_id variant stages a
// 64x32 tile of the expert's weights into threadgroup shared memory once
// and reuses it across a 32-row block of the expert's routed tokens —
// the same win as the dense mm kernel, but per-expert.
//
// Two-stage dispatch:
//
// 1. `kernel_mul_mm_id_map0_ne20_8` — scans the (token, slot) -> expert
// id table, for each expert builds a contiguous list of the packed
// `(token_idx * ne20 + slot_idx)` values that routed to it, plus a
// per-expert count. Scratch buffers: `hids` `[n_experts, n_tokens]`
// (int32, row-major), `htpe` `[n_experts]` (uint32).
//
// 2. `kernel_mul_mm_id_<qtype>_f32` — dispatched with one Z-per-expert.
// Each threadgroup owns one (N-tile, M-tile, expert_id) and short-
// circuits when the expert's routed-token count is below its M-tile
// start. Inside a tile, ALL 32 M-rows belong to the same expert,
// so the 64x32 weight tile loaded into shmem is valid for every row.
//
// The preprocessing step is the key to making mm profitable at MoE:
// without it, 32 consecutive output rows in a tile could route to 32
// different experts, defeating weight reuse.
//
// Port rules identical to the dense mm port (see quantized_matmul_mm.metal
// for detail). Both kernels are bit-compatible with the reference source;
// output tolerance-level matches with the existing mv_id kernel.
#include <metal_stdlib>
#include <metal_simdgroup>
#include <metal_simdgroup_matrix>
using namespace metal;
#define QK4_0 32
#define QK8_0 32
#define QK_K 256
#define QK_NL 16
// ---- Params for map0 ----
//
// Maps the reference map0 kargs layout. We omit `ne02` (== n_experts,
// used for a
// threadgroup-size assert upstream — we pass it host-side and dispatch
// exactly `n_experts` threads per threadgroup to mirror the upstream
// `ntg = n_experts` launch).
struct GgmlMatmulIdMm_Map0Params {
int32_t ne10; // unused, kept for struct symmetry
int32_t ne11; // n_expert_used (bcast), always == ne20 in our case
uint64_t nb11; // unused
uint64_t nb12; // unused
int32_t ne21; // n_tokens
int32_t ne20; // n_expert_used (== top_k)
uint64_t nb21; // bytes per token in the source ids table
// (= ne20 * sizeof(int32_t) for our layout)
};
// ---- Params for mul_mm_id ----
//
// Maps the reference mm_id kargs layout.
struct GgmlMatmulIdMm_MmParams {
int32_t ne00; // K
int32_t ne02; // n_experts
uint64_t nb01; // bytes per weight row (within one expert)
uint64_t nb02; // bytes per expert weight slab (== nb01 * ne01)
uint64_t nb03; // unused
int32_t ne11; // n_expert_used (bcast, == ne20)
uint64_t nb10; // = sizeof(float) = 4
uint64_t nb11; // bytes per input row = K * sizeof(float)
uint64_t nb12; // bytes per input batch (n_tokens * nb11)
uint64_t nb13; // unused
int32_t ne20; // n_expert_used (== top_k)
int32_t ne21; // n_tokens
int32_t ne0; // N (per-expert output rows)
int32_t ne1; // batch (always == ne21 * ne20 in our layout)
int16_t r2; // 1
int16_t r3; // 1
int16_t sched_flags; // compact kernels only: bit 0 = half-tile skip (P2c)
int16_t _pad1;
};
// ---- block_q struct layouts (byte-for-byte GGUF) ----
typedef struct {
half d;
uint8_t qs[QK4_0 / 2];
} block_q4_0;
typedef struct {
half d;
uint8_t qh[4];
uint8_t qs[QK4_0 / 2];
} block_q5_0;
typedef struct {
half d;
int8_t qs[QK8_0];
} block_q8_0;
typedef struct {
uint8_t scales[QK_K/16];
uint8_t qs[QK_K/4];
half d;
half dmin;
} block_q2_K;
typedef struct {
uint8_t hmask[QK_K/8];
uint8_t qs[QK_K/4];
uint8_t scales[12];
half d;
} block_q3_K;
static_assert(sizeof(block_q3_K) == sizeof(half) + QK_K/4 + QK_K/8 + 12,
"wrong q3_K block size");
typedef struct {
uint8_t ql[QK_K/2];
uint8_t qh[QK_K/4];
int8_t scales[QK_K/16];
half d;
} block_q6_K;
// ADR-013 P16 — Q4_K block (144 bytes) for mm_id port.
// Layout: [half d][half dmin][uint8_t scales[12]][uint8_t qs[128]]
#define K_SCALE_SIZE 12
typedef struct {
half d;
half dmin;
uint8_t scales[K_SCALE_SIZE];
uint8_t qs[QK_K/2];
} block_q4_K;
// ADR-022 Phase 2 — Q5_K block (176 bytes) for mm_id port.
// Layout: [half d][half dmin][uint8_t scales[12]][uint8_t qh[32]][uint8_t qs[128]]
// Same as Q4_K plus a 32-byte qh "high-bit" array. Mirrors
// quantized_matmul_id_ggml.metal:75 lock-step.
typedef struct {
half d;
half dmin;
uint8_t scales[K_SCALE_SIZE];
uint8_t qh[QK_K/8];
uint8_t qs[QK_K/2];
} block_q5_K;
// ADR-022 Phase 1 — Q5_1 / IQ4_NL block typedefs for the mm_id port.
typedef struct {
half d;
half m;
uint qh;
uint8_t qs[QK4_0 / 2];
} block_q5_1;
typedef struct {
half d;
uint8_t qs[QK4_0 / 2];
} block_iq4_nl;
// ADR-033 §Pi Task #20 — IQ4_XS block (136 bytes) for mm_id port.
// Layout: [half d][uint16_t scales_h][uint8_t scales_l[4]][uint8_t qs[128]]
// QK_K = 256 elements per super-block, organized as 8 sub-blocks of 32
// elements each. Each sub-block has a 6-bit signed scale (offset by 32)
// packed across scales_l (low 4 bits each, 2 per byte) + scales_h (top
// 2 bits each, 2 bits per sub-block in a uint16_t). Per-quant nibbles
// in qs are 4-bit codebook indexes into `kvalues_iq4nl` (shared with
// IQ4_NL). Lock-step with quantized_matmul_id_ggml.metal:139.
typedef struct {
half d;
uint16_t scales_h;
uint8_t scales_l[4];
uint8_t qs[QK_K/2];
} block_iq4_xs;
// ADR-066 N-IQ2XXS — IQ2_XXS block (66 bytes per 256 values).
// Lock-step with quantized_matmul_id_ggml.metal `block_iq2_xxs`.
typedef struct {
half d;
uint16_t qs[QK_K/8];
} block_iq2_xxs;
// BEGIN GENERATED iq2xxs
// Generated from src/quant/iq_tables.rs by tests/iq_tables.rs
// (MLX_BLESS_IQ_TABLES=1). Do not edit by hand.
constant uint64_t iq2xxs_grid[256] = {
0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
};
constant uint8_t ksigns_iq2xs[128] = {
0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15,
144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159,
160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175,
48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63,
192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207,
80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95,
96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
};
constant uint8_t kmask_iq2xs[8] = {
1, 2, 4, 8, 16, 32, 64, 128,
};
// END GENERATED iq2xxs
// IQ4_NL non-linear codebook (frozen by ggml-common.h:1109-1112).
// Lock-step with the duplicates in quantized_matmul_id_ggml.metal,
// quantized_matmul_ggml.metal, and the host-side `KVALUES_IQ4_NL`
// in src/gguf/mod.rs.
constant int8_t kvalues_iq4nl[16] = {
-127, -104, -83, -65, -49, -35, -22, -10,
1, 13, 25, 38, 53, 69, 89, 113
};
// Spec source: reference `get_scale_min_k4_just2`.
// Decodes the (sub-block scale, sub-block min) 6-bit pair at index `j`
// (within sub-block group `k`) from the packed K_SCALE_SIZE=12 array.
static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {
return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)}
: uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)),
uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))};
}
// ---- Dequantize helpers (identical to quantized_matmul_mm.metal) ----
template <typename type4x4>
void dequantize_q4_0(device const block_q4_0 * xb, short il, thread type4x4 & reg) {
device const uint16_t * qs = ((device const uint16_t *)xb + 1);
const float d1 = il ? (xb->d / 16.h) : xb->d;
const float d2 = d1 / 256.f;
const float md = -8.h * xb->d;
const ushort mask0 = il ? 0x00F0 : 0x000F;
const ushort mask1 = mask0 << 8;
float4x4 reg_f;
for (int i = 0; i < 8; i++) {
reg_f[i/2][2*(i%2) + 0] = d1 * (qs[i] & mask0) + md;
reg_f[i/2][2*(i%2) + 1] = d2 * (qs[i] & mask1) + md;
}
reg = (type4x4) reg_f;
}
template <typename type4x4>
void dequantize_q5_0(device const block_q5_0 * xb, short il, thread type4x4 & reg) {
device const uint16_t * qs = ((device const uint16_t *)xb + 3);
const float d = xb->d;
const float md = -16.h * xb->d;
const ushort mask = il ? 0x00F0 : 0x000F;
const uint qh = *((device const uint *)xb->qh);
const int x_mv = il ? 4 : 0;
const int gh_mv = il ? 12 : 0;
const int gh_bk = il ? 0 : 4;
float4x4 reg_f;
for (int i = 0; i < 8; i++) {
const uint8_t xh_0 = ((qh >> (gh_mv + 2*i )) << gh_bk) & 0x10;
const uint8_t xh_1 = ((qh >> (gh_mv + 2*i + 1)) << gh_bk) & 0x10;
const int x0 = ((((qs[i] ) & mask) >> x_mv) | xh_0);
const int x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
reg_f[i/2][2*(i%2) + 0] = d * x0 + md;
reg_f[i/2][2*(i%2) + 1] = d * x1 + md;
}
reg = (type4x4)reg_f;
}
template <typename type4x4>
void dequantize_q8_0(device const block_q8_0 * xb, short il, thread type4x4 & reg) {
device const int8_t * qs = ((device const int8_t *)xb->qs);
const float d = xb->d;
float4x4 reg_f;
for (int i = 0; i < 16; i++) {
reg_f[i/4][i%4] = (qs[i + 16*il] * d);
}
reg = (type4x4) reg_f;
}
// Q2_K dequantize, pinned to a fixed upstream revision.
template <typename type4x4>
void dequantize_q2_K(device const block_q2_K * xb, short il, thread type4x4 & reg) {
const float d = xb->d;
const float min = xb->dmin;
device const uint8_t * q = (device const uint8_t *)xb->qs;
const uint8_t sc = xb->scales[il];
q = q + 32*(il/8) + 16*(il&1);
il = (il/2)%4;
const half coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
const uchar mask = il>1 ? (il>2 ? 192 : 48) : (il>0 ? 12 : 3);
const float dl = d * (sc & 0xF) * coef;
const float ml = min * (sc >> 4);
for (int i = 0; i < 16; ++i) {
reg[i/4][i%4] = dl * (q[i] & mask) - ml;
}
}
template <typename type4x4>
void dequantize_q3_K(device const block_q3_K * xb, short il, thread type4x4 & reg) {
const half d_all = xb->d;
device const uint8_t * q = xb->qs;
device const uint8_t * h = xb->hmask;
device const int8_t * scales = (device const int8_t *)xb->scales;
q += 32*(il/8) + 16*(il&1);
h += 16*(il&1);
const uint8_t m = 1 << (il/2);
const uint16_t kmask1 = (il/4)>1 ? ((il/4)>2 ? 192 : 48) : ((il/4)>0 ? 12 : 3);
const uint16_t kmask2 = il/8 ? 0xF0 : 0x0F;
const uint16_t scale_2 = scales[il%8];
const uint16_t scale_1 = scales[8 + il%4];
const int16_t dl_int = (il/4)&1
? (scale_2&kmask2) | ((scale_1&kmask1) << 2)
: (scale_2&kmask2) | ((scale_1&kmask1) << 4);
float dl = il<8 ? d_all*(dl_int - 32.f) : d_all*(dl_int/16.f - 32.f);
const float ml = 4.f*dl;
il = (il/2)&3;
const half coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
const uint8_t mask = il>1 ? (il>2 ? 192 : 48) : (il>0 ? 12 : 3);
dl *= coef;
for (int i = 0; i < 16; ++i) {
reg[i/4][i%4] = dl*(q[i]&mask) - (h[i]&m ? 0.f : ml);
}
}
template <typename type4x4>
void dequantize_q6_K(device const block_q6_K * xb, short il, thread type4x4 & reg) {
const half d_all = xb->d;
device const uint16_t * ql = (device const uint16_t *)xb->ql;
device const uint16_t * qh = (device const uint16_t *)xb->qh;
device const int8_t * scales = (device const int8_t *)xb->scales;
ql = ql + 32*(il/8) + 16*((il/2)&1) + 8*(il&1);
qh = qh + 16*(il/8) + 8*(il&1);
float sc = scales[(il%2) + 2 * ((il/2))];
il = (il/2) & 3;
const uint32_t kmask1 = il>1 ? (il>2 ? 0xC0C0C0C0 : 0x30303030) : (il>0 ? 0x0C0C0C0C : 0x03030303);
const uint32_t kmask2 = il>1 ? 0xF0F0F0F0 : 0x0F0F0F0F;
const float ml = d_all * sc * 32.f;
const float dl0 = d_all * sc;
const float dl1 = dl0 / 256.f;
const float dl2 = dl0 / (256.f * 256.f);
const float dl3 = dl0 / (256.f * 256.f * 256.f);
const uint8_t shr_h = il>2 ? 2 : 0;
const uint8_t shl_h = il>1 ? 0 : (il>0 ? 2 : 4);
const uint8_t shr_l = il>1 ? 4 : 0;
float4x4 reg_f;
for (int i = 0; i < 4; ++i) {
const uint32_t low = (ql[2*i] | (uint32_t)(ql[2*i+1] << 16)) & kmask2;
const uint32_t high = (qh[2*i] | (uint32_t)(qh[2*i+1] << 16)) & kmask1;
const uint32_t q = ((high << shl_h) >> shr_h) | (low >> shr_l);
reg_f[i][0] = dl0 * ((half)(q & 0xFF)) - ml;
reg_f[i][1] = dl1 * ((float)(q & 0xFF00)) - ml;
reg_f[i][2] = dl2 * ((float)(q & 0xFF0000)) - ml;
reg_f[i][3] = dl3 * ((float)(q & 0xFF000000))- ml;
}
reg = (type4x4) reg_f;
}
// ADR-022 Phase 1 — Q5_1 dequant for the mm_id MMA-tile path.
// Spec source: reference `dequantize_q5_1`.
// Fills a 4x4 tile (16 values) with positions [16*il .. 16*il+16) of
// the 32-element block. il=0 fills low half, il=1 fills high half.
template <typename type4x4>
void dequantize_q5_1(device const block_q5_1 * xb, short il, thread type4x4 & reg) {
device const uint16_t * qs = ((device const uint16_t *)xb + 4);
const float d = xb->d;
const float m = xb->m;
const ushort mask = il ? 0x00F0 : 0x000F;
const uint32_t qh = xb->qh;
const int x_mv = il ? 4 : 0;
const int gh_mv = il ? 12 : 0;
const int gh_bk = il ? 0 : 4;
float4x4 reg_f;
for (int i = 0; i < 8; i++) {
const uint8_t xh_0 = ((qh >> (gh_mv + 2*i )) << gh_bk) & 0x10;
const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
const int32_t x0 = ((((qs[i] ) & mask) >> x_mv) | xh_0);
const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
reg_f[i/2][2*(i%2) + 0] = d * x0 + m;
reg_f[i/2][2*(i%2) + 1] = d * x1 + m;
}
reg = (type4x4) reg_f;
}
// ADR-022 Phase 1 — IQ4_NL dequant for the mm_id MMA-tile path.
// Spec source: reference `dequantize_iq4_nl`.
// 16 elements per call addressed by il ∈ {0, 1}; loops over 4 uint16
// chunks of qs, extracting 4 4-bit indices each via shift+mask, looking
// up into kvalues_iq4nl per nibble.
template <typename type4x4>
void dequantize_iq4_nl(device const block_iq4_nl * xb, short il, thread type4x4 & reg) {
device const uint16_t * q4 = (device const uint16_t *)xb->qs;
const float d = xb->d;
uint32_t aux32;
thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
for (int i = 0; i < 4; ++i) {
aux32 = ((q4[2*i] | (q4[2*i+1] << 16)) >> 4*il) & 0x0f0f0f0f;
reg[i][0] = d * (float)kvalues_iq4nl[q8[0]];
reg[i][1] = d * (float)kvalues_iq4nl[q8[1]];
reg[i][2] = d * (float)kvalues_iq4nl[q8[2]];
reg[i][3] = d * (float)kvalues_iq4nl[q8[3]];
}
}
// ADR-033 §Pi Task #20 — IQ4_XS float codebook table.
// Identical values to `kvalues_iq4nl` above but pre-converted to float
// to match the reference `kvalues_iq4nl_f` representation byte-for-byte.
// Used by `dequantize_iq4_xs` below — using the int8 table + cast was
// producing numerical divergence at top_k=8 mm_id batch shapes, so the
// safer mirror of canonical is to emit a float table directly.
constant float kvalues_iq4nl_f[16] = {
-127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f,
1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f
};
// ADR-033 §Pi Task #20 — IQ4_XS dequant for the mm_id MMA-tile path.
// Spec source: reference `dequantize_iq4_xs`
// — reproduced verbatim modulo formatting.
//
// 16 elements per call addressed by `il ∈ [0,16)`:
// - `ib32 = il/2` selects the 32-element sub-block (0..7) in qs;
// - `il%2` selects which 16-elem half (low/high nibble) within that
// sub-block.
//
// Per-sub-block 6-bit signed scale is reconstructed from 4 low bits in
// `scales_l[ib32/2]` (nibble-packed, 2 sub-blocks per byte) plus the top
// 2 bits in `scales_h` (2 bits per sub-block in a uint16_t). Scale value
// is `xb->d * (ls - 32)` — the offset of 32 is the signed-conversion.
template <typename type4x4>
void dequantize_iq4_xs(device const block_iq4_xs * xb, short il, thread type4x4 & reg) {
// il is 0..15 for QK_K = 256 — index of 32-elem sub-block is il/2
const int ib32 = il / 2;
il = il % 2;
// il = 0 or 1: il=0 → first 16 quants of the 32-elem sub-block,
// il=1 → second 16 quants.
device const uint32_t * q4 = (device const uint32_t *)xb->qs + 4 * ib32;
const int ls = ((xb->scales_l[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf)
| (((xb->scales_h >> 2 * ib32) & 3) << 4);
const float d = (float)xb->d * (ls - 32);
uint32_t aux32;
thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
for (int i = 0; i < 4; ++i) {
aux32 = (q4[i] >> 4 * il) & 0x0f0f0f0f;
reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
}
}
// ADR-066 N-IQ2XXS — IQ2_XXS dequant for the mm_id MMA-tile path.
// Ported from the peer's `dequantize_iq2_xxs` (dequantize.h). 16 values
// per call: `il / 2` selects the 32-value sub-block and `il % 2` its half,
// which uses grid indices 2*(il%2) and 2*(il%2)+1 and the matching two
// 7-bit sign fields.
template <typename type4x4>
void dequantize_iq2_xxs(device const block_iq2_xxs * xb, short il, thread type4x4 & reg) {
const float d = xb->d;
const int ib32 = il / 2;
il = il % 2;
device const uint16_t * q2 = xb->qs + 4 * ib32;
const uint32_t aux32_g = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux32_s = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
thread const uint8_t * aux8 = (thread const uint8_t *)&aux32_g;
const float dl = d * (0.5f + (aux32_s >> 28)) * 0.25f;
constant uint8_t * grid = (constant uint8_t *)(iq2xxs_grid + aux8[2 * il + 0]);
uint8_t signs = ksigns_iq2xs[(aux32_s >> 14 * il) & 127];
for (int i = 0; i < 8; ++i) {
reg[i / 4][i % 4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
}
grid = (constant uint8_t *)(iq2xxs_grid + aux8[2 * il + 1]);
signs = ksigns_iq2xs[(aux32_s >> (14 * il + 7)) & 127];
for (int i = 0; i < 8; ++i) {
reg[2 + i / 4][i % 4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
}
}
// ADR-013 P16 — Q4_K dequant for the mm_id MMA-tile path.
// Spec source: reference `dequantize_q4_K`. Fills a 4x4 tile with the
// 16 dequantized values addressed by `il` (which selects sub-block + half
// nibble), using the same scale/min decode as the mat-vec-id Q4_K kernel
// already shipped at quantized_matmul_id_ggml.metal.
template <typename type4x4>
void dequantize_q4_K(device const block_q4_K * xb, short il, thread type4x4 & reg) {
device const uchar * q = xb->qs;
short is = (il/4) * 2;
q = q + (il/4) * 32 + 16 * (il&1);
il = il & 3;
const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
const float d = il < 2 ? xb->d : xb->d / 16.h;
const float min = xb->dmin;
const float dl = d * sc[0];
const float ml = min * sc[1];
const ushort mask = il < 2 ? 0x0F : 0xF0;
for (int i = 0; i < 16; ++i) {
reg[i/4][i%4] = dl * (q[i] & mask) - ml;
}
}
// ADR-022 Phase 2 — Q5_K dequant for mm_id MMA-tile path.
// Spec source: reference `dequantize_q5_K`.
// Q5_K differs from Q4_K by adding a 32-byte qh "high-bit" array; each
// of the 16 elements per call OR's an extra 16 (low half) or 256 (high
// half) into the dequantized value when the corresponding qh bit is set.
template <typename type4x4>
void dequantize_q5_K(device const block_q5_K * xb, short il, thread type4x4 & reg) {
device const uint8_t * q = xb->qs;
device const uint8_t * qh = xb->qh;
short is = (il/4) * 2;
q = q + 32 * (il/4) + 16 * (il&1);
qh = qh + 16 * (il&1);
uint8_t ul = 1 << (il/2);
il = il & 3;
const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
const float d = il < 2 ? xb->d : xb->d / 16.h;
const float min = xb->dmin;
const float dl = d * sc[0];
const float ml = min * sc[1];
const ushort mask = il < 2 ? 0x0F : 0xF0;
const float qh_val = il < 2 ? 16.f : 256.f;
for (int i = 0; i < 16; ++i) {
reg[i/4][i%4] = dl * ((q[i] & mask) + (qh[i] & ul ? qh_val : 0)) - ml;
}
}
// ====================================================================
// map0 — per-expert routed-token list builder
// ====================================================================
//
// Input: src2 — per-token expert ids `[n_tokens, n_expert_used]` int32.
// Output: htpe — per-expert routed count `[n_experts]` uint32. Bit 31 is
// set in every entry when any token has an out-of-range or
// duplicate expert ID; consumers then poison the operation.
// Output: hids — per-expert routed-token list `[n_experts, n_tokens]` int32.
// Each slot holds `(token_idx * ne20 + slot_idx)` packed.
//
// Dispatch geometry:
// threadgroups = (1, 1, 1)
// threads_per_threadgroup = (n_experts, 1, 1) (ntg == ne02)
//
// One thread per expert id (ide = tpitg). Threads scan the token axis
// in blocks of `ntg` tokens, stage each token's `ne20` expert-ids into
// threadgroup shmem, then each thread scans that block looking for its
// own expert-id and appends entries into its slot in `hids`.
template<short ne20>
kernel void hf2q_mul_mm_id_map0_impl(
constant GgmlMatmulIdMm_Map0Params & args [[buffer(0)]],
device const char * src2 [[buffer(1)]],
device char * htpe [[buffer(2)]],
device char * hids [[buffer(3)]],
threadgroup char * shmem [[threadgroup(0)]],
ushort tpitg[[thread_position_in_threadgroup]],
ushort ntg[[threads_per_threadgroup]]) {
const short ide = tpitg; // expert id owned by this thread
uint32_t n_all = 0;
bool routing_invalid = false;
device int32_t * ids_i32 = (device int32_t *) hids + ide * args.ne21;
for (int i21 = 0; i21 < args.ne21; i21 += ntg) {
if (i21 + tpitg < args.ne21) {
device const uint32_t * src2_u32 =
(device const uint32_t *) (src2 + (i21 + tpitg) * args.nb21);
threadgroup uint32_t * sids = (threadgroup uint32_t *) shmem + tpitg * ne20;
for (short i20 = 0; i20 < ne20; i20++) {
sids[i20] = src2_u32[i20];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (short t = 0; t < ntg; t++) {
if (i21 + t >= args.ne21) break;
threadgroup const uint32_t * sids =
(threadgroup const uint32_t *) shmem + t * ne20;
bool token_invalid = false;
for (short i20 = 0; i20 < ne20; i20++) {
token_invalid |= sids[i20] >= ntg;
for (short other = 0; other < i20; other++) {
token_invalid |= sids[i20] == sids[other];
}
}
routing_invalid |= token_invalid;
if (token_invalid) {
continue;
}
short selected_slot = -1;
for (short i20 = 0; i20 < ne20; i20++) {
if (sids[i20] == ide) {
selected_slot = i20;
break;
}
}
if (selected_slot >= 0) {
ids_i32[n_all++] = (i21 + t) * ne20 + selected_slot;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
tpe_u32[ide] = n_all | (routing_invalid ? 0x80000000u : 0u);
}
// Gemma 4 uses top_k = 8 for the MoE gate_up call and top_k = 1 for the
// MoE down call (where the gate_up output is flattened to
// [seq_len*top_k] rows, each one routed to a single expert via the
// expert-id table). Both must hit the mm_id path — without ne20_1, the
// down call falls back to mv_id and re-reads each expert's weights once
// per (seq_len * top_k) row, defeating the whole point of mm.
template [[host_name("kernel_mul_mm_id_map0_ne20_1")]]
kernel void hf2q_mul_mm_id_map0_impl<1>(
constant GgmlMatmulIdMm_Map0Params &,
device const char *, device char *, device char *,
threadgroup char *, ushort, ushort);
template [[host_name("kernel_mul_mm_id_map0_ne20_6")]]
kernel void hf2q_mul_mm_id_map0_impl<6>(
constant GgmlMatmulIdMm_Map0Params &,
device const char *, device char *, device char *,
threadgroup char *, ushort, ushort);
template [[host_name("kernel_mul_mm_id_map0_ne20_8")]]
kernel void hf2q_mul_mm_id_map0_impl<8>(
constant GgmlMatmulIdMm_Map0Params &,
device const char *, device char *, device char *,
threadgroup char *, ushort, ushort);
// ====================================================================
// map0c — compact per-expert tile schedule builder (ADR-063 P2)
// ====================================================================
//
// Output layout:
//
// htpe word 0 — tile_count, or POISON (bit 31) when any token has
// an out-of-range or duplicate expert id, or the
// tile list would exceed `max_tiles`.
// htpe[1 + 3*t + 0..3] — descriptor of tile t: (expert, first list entry,
// valid rows <= 32). Tiles are emitted per expert
// in ascending expert order, each expert's tiles
// in ascending entry order.
// hids — compact routed-row list. Expert e's entries
// start at the exclusive prefix sum of
// (count + 1) over experts < e, and hold
// `token * ne20 + slot` in ascending token order,
// so each expert's 32-row tiles have exactly the
// membership of the uniform-grid schedule. The
// one padding entry per slice is never written.
//
// Launch: one threadgroup of `ntg = 32 * G` threads (G simdgroups, chosen by
// the host from the threadgroup-memory budget; independent of n_experts).
// Token-parallel: in each round of `ntg` tokens, thread i owns token
// round + i, so lane l of simdgroup g owns token round + 32g + l.
//
// Pass 1 the owner validates its token (every id < n_experts, the ne20
// ids pairwise distinct; the map0 predicate, once per token) and
// adds its valid ids to per-expert counts (threadgroup atomics;
// the sum is order-free).
// Scan exclusive prefix sums over experts of (count + 1) and
// ceil(count / 32); poison on an invalid token or more than
// `max_tiles` tiles; each expert then writes its own descriptors.
// Pass 2 per round, each simdgroup ORs bit `lane` into its per-expert
// mask for every id of its tokens. The entry index of (token,
// slot -> expert e) is e's running write position, plus the set
// bits of e's masks in earlier simdgroups (an exclusive simd scan
// across the G masks), plus the bits below `lane` in its own mask:
// ascending token order, without a per-expert sequential scan.
//
// Threadgroup memory (words): next[n_experts] (pass-1 counts, then slice
// starts, then write positions), masks[max(G, 2) * n_experts] (lane masks,
// then write offsets; the scan keeps the counts and first tiles in its first
// 2 * n_experts words), part[2 * G] (scan partials), invalid[1].
struct GgmlMatmulIdMm_Map0cParams {
int32_t ne21; // n_tokens
int32_t ne20; // n_expert_used (== top_k)
uint64_t nb21; // bytes per token in the source ids table
uint32_t max_tiles; // descriptor capacity bound
uint32_t n_experts; // ne02
};
template<short ne20>
kernel void hf2q_mul_mm_id_map0c_impl(
constant GgmlMatmulIdMm_Map0cParams & args [[buffer(0)]],
device const char * src2 [[buffer(1)]],
device char * htpe [[buffer(2)]],
device char * hids [[buffer(3)]],
threadgroup char * shmem [[threadgroup(0)]],
ushort tpitg[[thread_position_in_threadgroup]],
ushort ntg[[threads_per_threadgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]]) {
const uint n_experts = args.n_experts;
const int n_tokens = args.ne21;
const uint G = ntg / 32;
const uint MG = max(G, 2u);
threadgroup atomic_uint * next = (threadgroup atomic_uint *) shmem;
threadgroup atomic_uint * masks = next + n_experts;
threadgroup uint * part = (threadgroup uint *) (masks + MG * n_experts);
threadgroup atomic_uint * invalid = (threadgroup atomic_uint *) (part + 2 * G);
for (uint e = tpitg; e < n_experts; e += ntg) {
atomic_store_explicit(&next[e], 0u, memory_order_relaxed);
}
if (tpitg == 0) {
atomic_store_explicit(invalid, 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// ---- Pass 1: validate owned tokens, count rows per expert ----
bool bad = false;
for (int i0 = 0; i0 < n_tokens; i0 += ntg) {
const int t = i0 + tpitg;
if (t < n_tokens) {
device const uint32_t * row =
(device const uint32_t *) (src2 + (uint64_t) t * args.nb21);
uint32_t own[ne20];
#pragma unroll
for (short i20 = 0; i20 < ne20; i20++) {
own[i20] = row[i20];
bad |= own[i20] >= n_experts;
#pragma unroll
for (short other = 0; other < i20; other++) {
bad |= own[i20] == own[other];
}
}
#pragma unroll
for (short i20 = 0; i20 < ne20; i20++) {
if (own[i20] < n_experts) {
atomic_fetch_add_explicit(&next[own[i20]], 1u, memory_order_relaxed);
}
}
}
}
if (simd_any(bad) && tiisg == 0) {
atomic_store_explicit(invalid, 1u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// ---- Scan: slice starts and first tiles, in chunks of ntg experts ----
uint carry_entry = 0;
uint carry_tile = 0;
for (uint e0 = 0; e0 < n_experts; e0 += ntg) {
const uint e = e0 + tpitg;
const uint count = e < n_experts
? atomic_load_explicit(&next[e], memory_order_relaxed) : 0u;
const uint v_entry = e < n_experts ? count + 1u : 0u;
const uint v_tile = (count + 31u) / 32u;
const uint p_entry = simd_prefix_exclusive_sum(v_entry);
const uint p_tile = simd_prefix_exclusive_sum(v_tile);
if (tiisg == 31) {
part[sgitg] = p_entry + v_entry;
part[G + sgitg] = p_tile + v_tile;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint s_entry = tiisg < G ? part[tiisg] : 0u;
const uint s_tile = tiisg < G ? part[G + tiisg] : 0u;
const uint before_entry = simd_shuffle(simd_prefix_exclusive_sum(s_entry), sgitg);
const uint before_tile = simd_shuffle(simd_prefix_exclusive_sum(s_tile), sgitg);
if (e < n_experts) {
atomic_store_explicit(&next[e], carry_entry + before_entry + p_entry,
memory_order_relaxed);
atomic_store_explicit(&masks[e], count, memory_order_relaxed);
atomic_store_explicit(&masks[n_experts + e], carry_tile + before_tile + p_tile,
memory_order_relaxed);
}
carry_entry += simd_sum(s_entry);
carry_tile += simd_sum(s_tile);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const bool poison =
atomic_load_explicit(invalid, memory_order_relaxed) != 0u || carry_tile > args.max_tiles;
if (tpitg == 0) {
device uint32_t * word0 = (device uint32_t *) htpe;
word0[0] = poison ? 0x80000000u : carry_tile;
}
if (poison) {
return;
}
// ---- Tile descriptors, each expert its own ----
for (uint e = tpitg; e < n_experts; e += ntg) {
const uint count = atomic_load_explicit(&masks[e], memory_order_relaxed);
const uint first_entry = atomic_load_explicit(&next[e], memory_order_relaxed);
device uint32_t * desc = (device uint32_t *) htpe + 1
+ 3 * atomic_load_explicit(&masks[n_experts + e], memory_order_relaxed);
for (uint r = 0; r < count; r += 32) {
desc[0] = e;
desc[1] = first_entry + r;
desc[2] = (count - r < 32u) ? (count - r) : 32u;
desc += 3;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = tpitg; i < G * n_experts; i += ntg) {
atomic_store_explicit(&masks[i], 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// ---- Pass 2: entries in ascending token order ----
// Every id is valid here (an invalid token poisoned the schedule).
threadgroup atomic_uint * own_masks = masks + sgitg * n_experts;
device int32_t * list = (device int32_t *) hids;
const uint below = (1u << tiisg) - 1u;
for (int i0 = 0; i0 < n_tokens; i0 += ntg) {
const int t = i0 + tpitg;
const bool live = t < n_tokens;
uint32_t own[ne20];
uint32_t rank[ne20];
if (live) {
device const uint32_t * row =
(device const uint32_t *) (src2 + (uint64_t) t * args.nb21);
#pragma unroll
for (short i20 = 0; i20 < ne20; i20++) {
own[i20] = row[i20];
atomic_fetch_or_explicit(&own_masks[own[i20]], 1u << tiisg, memory_order_relaxed);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (live) {
#pragma unroll
for (short i20 = 0; i20 < ne20; i20++) {
rank[i20] = popcount(
atomic_load_explicit(&own_masks[own[i20]], memory_order_relaxed) & below);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// Masks -> write offsets: simdgroup sgitg converts experts
// sgitg, sgitg + G, ...; lane g holds simdgroup g's mask and lane 0
// owns the expert's write position.
for (uint e = sgitg; e < n_experts; e += G) {
threadgroup atomic_uint * cell = masks + tiisg * n_experts + e;
const uint bits = tiisg < G
? popcount(atomic_load_explicit(cell, memory_order_relaxed)) : 0u;
const uint base = simd_shuffle(
tiisg == 0 ? atomic_load_explicit(&next[e], memory_order_relaxed) : 0u, 0);
const uint offset = base + simd_prefix_exclusive_sum(bits);
if (tiisg < G) {
atomic_store_explicit(cell, offset, memory_order_relaxed);
}
const uint total = simd_sum(bits);
if (tiisg == 0) {
atomic_store_explicit(&next[e], base + total, memory_order_relaxed);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (live) {
#pragma unroll
for (short i20 = 0; i20 < ne20; i20++) {
list[atomic_load_explicit(&own_masks[own[i20]], memory_order_relaxed) + rank[i20]] =
t * ne20 + i20;
}
}
if (i0 + (int) ntg < n_tokens) {
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = tpitg; i < G * n_experts; i += ntg) {
atomic_store_explicit(&masks[i], 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}
}
template [[host_name("kernel_mul_mm_id_map0c_ne20_1")]]
kernel void hf2q_mul_mm_id_map0c_impl<1>(
constant GgmlMatmulIdMm_Map0cParams &,
device const char *, device char *, device char *,
threadgroup char *, ushort, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_map0c_ne20_6")]]
kernel void hf2q_mul_mm_id_map0c_impl<6>(
constant GgmlMatmulIdMm_Map0cParams &,
device const char *, device char *, device char *,
threadgroup char *, ushort, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_map0c_ne20_8")]]
kernel void hf2q_mul_mm_id_map0c_impl<8>(
constant GgmlMatmulIdMm_Map0cParams &,
device const char *, device char *, device char *,
threadgroup char *, ushort, ushort, ushort, ushort);
// ====================================================================
// mul_mm_id — MoE matmul using the map0-produced expert token lists
// ====================================================================
//
// Dispatch geometry:
// threadgroups = (ceil(neh1/32), ceil(N/64), n_experts)
// threads_per_threadgroup = (128, 1, 1)
//
// where `neh1 = htpe[expert]` (count of tokens routed to this expert).
//
// Each threadgroup owns:
// * im = tgpig.z — expert index
// * r0-block = tgpig.y — output-N tile index
// * r1-block = tgpig.x — input-M (per-expert) tile index
//
// The kernel short-circuits inactive (expert, r1) combinations via the
// `r1 >= neh1` check — some threadgroups may exit early but the grid
// launch is uniform.
//
// Per-iteration of the K loop: stage a 64x32 A-tile (from the expert's
// weight slab) + 32x32 B-tile (from input rows selected via hids) into
// threadgroup memory, then run 8 simdgroup MMA ops per simdgroup. Four
// simdgroups cooperate on one 64x32 output tile.
//
// Write-back staging ALWAYS goes through shmem because consecutive M-rows
// in this tile belong to the same expert but may correspond to different
// (token, slot) output rows (determined by hids[im * ne21 + r1 + j]).
//
// COMPACT (ADR-063 P2, `_c` host names): the grid is `max_tiles` tile
// slots (grid x, fastest) × ⌈N/64⌉ N tiles (grid y) over the map0c
// schedule; ADR-063 M0c measured this order faster than N tile fastest at
// the served shapes. Word 0 of htpe holds tile_count or POISON; each
// live tile loads its (expert, first entry, valid rows) descriptor and runs
// the tile body below unchanged, with `hids[im*ne21 + r1 + j]` replaced by
// `hids[first_entry + j]`. Bit 0 enables the half-tile skip (P2c): when a
// tile holds at most 16 rows, simdgroups 2-3 — which own rows 16-31 — skip
// their MMA and shared-memory store; all four simdgroups still stage the
// tiles and copy rows out.
template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread half4x4 &), bool COMPACT = false>
kernel void hf2q_mul_mm_id_impl(
constant GgmlMatmulIdMm_MmParams & args [[buffer(0)]],
device const char * src0 [[buffer(1)]],
device const char * src1 [[buffer(2)]],
device const char * htpe [[buffer(3)]],
device const char * hids [[buffer(4)]],
device char * dst [[buffer(5)]],
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiitg[[thread_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
threadgroup half * sa = (threadgroup half *)(shmem);
threadgroup float * sb = (threadgroup float *)(shmem + 4096);
constexpr int NR0 = 64;
constexpr int NR1 = 32;
constexpr int NK = 32;
constexpr int NL0 = NK/16; // 2
constexpr int NL1 = NK/8; // 4
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
device const int32_t * ids_i32 = (device const int32_t *) (hids);
int im;
int r0;
int list_base; // hids index of this tile's first routed row
short nr1;
if (COMPACT) {
const uint32_t schedule_state = tpe_u32[0];
const uint32_t tile = tgpig.x;
const uint32_t ntile = tgpig.y;
if ((schedule_state & 0x80000000u) != 0u) {
// Same poison contract as the uniform grid: one threadgroup
// fills the logical output with NaN and nobody forms a weight
// or activation address.
if (tgpig.x == 0 && tgpig.y == 0) {
device float * out = (device float *) dst;
const uint64_t total =
(uint64_t)args.ne21 * (uint64_t)args.ne1 * (uint64_t)args.ne0;
for (uint64_t index = tiitg; index < total; index += 128) {
out[index] = NAN;
}
}
return;
}
if (tile >= schedule_state) return;
device const uint32_t * desc = tpe_u32 + 1 + 3*tile;
im = (int)desc[0];
list_base = (int)desc[1];
nr1 = (short)desc[2];
r0 = ntile * NR0;
} else {
im = tgpig.z;
r0 = tgpig.y * NR0;
const int r1 = tgpig.x * NR1;
const uint32_t route_state = tpe_u32[im];
if ((route_state & 0x80000000u) != 0u) {
// Invalid IDs poison the whole operation. Exactly one threadgroup
// fills the complete logical output with NaN; every threadgroup exits
// before constructing an expert weight or activation address. The
// consuming graph is responsible for rejecting this sentinel at an
// existing host-visible result readback.
if (im == 0 && tgpig.x == 0 && tgpig.y == 0) {
device float * out = (device float *) dst;
const uint64_t total =
(uint64_t)args.ne21 * (uint64_t)args.ne1 * (uint64_t)args.ne0;
for (uint64_t index = tiitg; index < total; index += 128) {
out[index] = NAN;
}
}
return;
}
const int32_t neh1 = (int32_t)route_state;
// Early exit: this expert has fewer routed tokens than our tile's
// M-base. Whole threadgroup returns.
if (r1 >= neh1) return;
nr1 = (neh1 - r1 < NR1) ? (neh1 - r1) : NR1;
list_base = im * args.ne21 + r1;
}
// P2c: simdgroups 2-3 own tile rows 16-31, which hold no routed row.
const bool upper_half_idle =
COMPACT && (args.sched_flags & 1) != 0 && nr1 <= 16 && sgitg >= 2;
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1;
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1;
const short il0 = (tiitg % NL0);
short il = il0;
// (token_idx, slot_idx) of the row this thread owns.
const int id = ids_i32[list_base + lr1];
const short i11 = (id % args.ne20) % args.ne11; // slot index
// §0.19 FIX (2026-06-30): token/row index must be `int`, not `short`.
// The flat [n_tokens, K] input layout makes i12 = id/ne20 range up to
// n_tokens-1 (65535 at 8k, 131071 at 16k for the MoE down call where
// n_tokens = seq_len*top_k). `short` overflows at 32768 → negative
// i12 → `args.nb12*i12` wraps to an OOB device address (read of
// run-varying memory). llama keeps this index small via its
// [K, n_expert_used, n_tokens] multi-dim layout; our flat port must
// widen it. Bit-identical for n_tokens<=32768 (parity gate).
const int i12 = (id / args.ne20); // token index
const short i13 = 0;
// Base of expert `im`'s weight slab.
const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
const short offset1 = il0/nl;
device const block_q * x =
(device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
const short iy = 8*(tiitg % NL1);
device const float * y = (device const float *)(src1
+ args.nb13*i13
+ args.nb12*(uint64_t)i12
+ args.nb11*i11
+ args.nb10*iy);
simdgroup_half8x8 ma[4];
simdgroup_float8x8 mb[2];
simdgroup_float8x8 mc[8];
for (short i = 0; i < 8; i++){
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
}
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
// ---- A tile dequantize + stage ----
{
half4x4 temp_a;
dequantize_func(x, il, temp_a);
threadgroup_barrier(mem_flags::mem_threadgroup);
// ADR-033 §Pi Task #20 iter 8 — added #pragma unroll to match
// the peer's FOR_UNROLL for the 16-iter A-tile staging loop.
#pragma clang loop unroll(full)
for (short i = 0; i < 16; i++) {
const short sx = 2*il0 + i/8;
const short sy = (tiitg/NL0)/8;
const short lx = (tiitg/NL0)%8;
const short ly = i%8;
const short ib = 8*sx + sy;
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
}
}
// ---- B tile stage ----
//
// ADR-033 §Pi Task #20 iter 8 — vectorized B-tile store fast path
// matching the reference `*(threadgroup S1_2x4 *)(sb + ...) = (S1_2x4) *((device T1_2x4 *) y)`.
// Each thread does ONE 8-float store instead of 8 scalar stores
// when the K-loop iteration is fully in-bounds. For Q6_K with
// K%QK_K==0 and NK=32, K is always divisible by NK, so the fast
// path is taken every iteration (no tail). Kept the scalar
// bounds-checked fallback for non-Q6_K shapes.
{
const short sx = (tiitg%NL1);
const short sy = (tiitg/NL1)/8;
const short ly = (tiitg/NL1)%8;
const short ib = 4*sx + sy;
if (loop_k + NK <= args.ne00) {
// Fast path: vectorized 8-float store (= S1_2x4 / float8).
threadgroup float * dst_row = sb + 64*ib + 8*ly;
device const float * src_row = (device const float *) y;
((threadgroup float4 *) dst_row)[0] = ((device const float4 *) src_row)[0];
((threadgroup float4 *) dst_row)[1] = ((device const float4 *) src_row)[1];
} else {
// Tail: scalar with bounds check.
#pragma clang loop unroll(full)
for (short i = 0; i < 8; ++i) {
const short lx = i;
*(sb + 64*ib + 8*ly + lx) =
(loop_k + iy + i < args.ne00) ? *((device float *) y + i) : 0.f;
}
}
}
il = (il + 2 < nl) ? il + 2 : il % 2;
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
y += NK;
threadgroup_barrier(mem_flags::mem_threadgroup);
// ---- MMA accumulate ----
threadgroup const half * lsma = (sa + 4*64*(sgitg%2));
threadgroup const float * lsmb = (sb + 2*64*(sgitg/2));
if (!upper_half_idle) {
for (short ik = 0; ik < NK/8; ik++) {
simdgroup_barrier(mem_flags::mem_none);
for (short i = 0; i < 4; i++) {
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
for (short i = 0; i < 2; i++) {
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
for (short i = 0; i < 8; i++) {
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
}
lsma += 8*64;
lsmb += 4*64;
}
}
}
// ---- Write-back (always through shmem staging: see header note) ----
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup float * temp_str = ((threadgroup float *) shmem)
+ 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
if (!upper_half_idle) {
for (short i = 0; i < 8; i++) {
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// Each simdgroup `sgitg` strides through the tile in steps of 4 rows.
// For each owned row j (within the tile), look up its (token, slot)
// via hids, compute the destination address in dst, and copy.
for (short j = sgitg; j < nr1; j += 4) {
const int id = ids_i32[list_base + j];
const short ide = id % args.ne20; // slot index (output row-major)
// §0.19 FIX (2026-06-30): output token/batch index must be `int` —
// same overflow as i12 above. For the MoE down call idt = id ranges
// to n_tokens-1 (65535 at 8k); `short` overflows at 32768 → negative
// idt → the `dst + idt*ne1*ne0` write-back wraps to an OOB device
// address (corrupting/reading run-varying memory). Bit-identical for
// n_tokens<=32768.
const int idt = id / args.ne20; // token index (output batch)
// §0.19 FIX (2026-06-30, codex SHIP-WITH-CHANGES): force the output
// ELEMENT offset to 64-bit (promote each factor BEFORE multiply) so the
// dst address never overflows int32 at large ne0/ne1/token counts —
// mirrors the int32 dst-offset fix in dense_mm_bf16_tensor.metal.
// Byte-identical for in-range values.
const uint64_t dst_off =
(uint64_t) r0
+ (uint64_t) ide * (uint64_t) args.ne0
+ (uint64_t) idt * (uint64_t) args.ne1 * (uint64_t) args.ne0;
device float * D = (device float *) dst + dst_off;
device float4 * D4 = (device float4 *) D;
threadgroup float * C = (threadgroup float *) shmem + j*NR0;
threadgroup float4 * C4 = (threadgroup float4 *) C;
int i = tiisg;
for (; i < nr0/4; i += 32) {
*(D4 + i) = *(C4 + i);
}
i = (4*(nr0/4)) + tiisg;
for (; i < nr0; i += 32) {
*(D + i) = *(C + i);
}
}
}
// Template instantiations for the three quant types our GGUF uses.
template [[host_name("kernel_mul_mm_id_q4_0_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q4_0, 2, dequantize_q4_0>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q5_0_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q5_0, 2, dequantize_q5_0>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q8_0_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q8_0, 2, dequantize_q8_0>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q2_K_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q2_K, QK_NL, dequantize_q2_K>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q3_K_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q3_K, QK_NL, dequantize_q3_K>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q6_K_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q6_K, QK_NL, dequantize_q6_K>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ADR-013 P16 — Q4_K mm_id template instantiation.
template [[host_name("kernel_mul_mm_id_q4_K_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q4_K, QK_NL, dequantize_q4_K>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ADR-022 Phase 2 — Q5_K mm_id template instantiation.
template [[host_name("kernel_mul_mm_id_q5_K_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q5_K, QK_NL, dequantize_q5_K>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ADR-022 Phase 1 P1.6 — Q5_1 / IQ4_NL mm_id template instantiations.
template [[host_name("kernel_mul_mm_id_q5_1_f32")]]
kernel void hf2q_mul_mm_id_impl<block_q5_1, 2, dequantize_q5_1>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ADR-033 §Pi Task #20 — IQ4_XS mm_id template instantiation.
// Matches the reference `kernel_mul_mm_id_iq4_xs_f32`
// (parameters: block_iq4_xs + QK_NL=16
// + dequantize_iq4_xs).
template [[host_name("kernel_mul_mm_id_iq4_xs_f32")]]
kernel void hf2q_mul_mm_id_impl<block_iq4_xs, QK_NL, dequantize_iq4_xs>(
constant GgmlMatmulIdMm_MmParams & args [[buffer(0)]],
device const char * src0 [[buffer(1)]],
device const char * src1 [[buffer(2)]],
device const char * htpe [[buffer(3)]],
device const char * hids [[buffer(4)]],
device char * dst [[buffer(5)]],
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiitg[[thread_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]);
template [[host_name("kernel_mul_mm_id_iq4_nl_f32")]]
kernel void hf2q_mul_mm_id_impl<block_iq4_nl, 2, dequantize_iq4_nl>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ADR-066 N-IQ2XXS — IQ2_XXS mm_id template instantiation.
template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]]
kernel void hf2q_mul_mm_id_impl<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ADR-063 P2 — compact-schedule (`_c`) instantiations of every mm_id type.
template [[host_name("kernel_mul_mm_id_q4_0_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q4_0, 2, dequantize_q4_0, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q5_0_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q5_0, 2, dequantize_q5_0, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q8_0_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q8_0, 2, dequantize_q8_0, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q2_K_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q2_K, QK_NL, dequantize_q2_K, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q3_K_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q3_K, QK_NL, dequantize_q3_K, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q4_K_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q4_K, QK_NL, dequantize_q4_K, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q5_K_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q5_K, QK_NL, dequantize_q5_K, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q6_K_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q6_K, QK_NL, dequantize_q6_K, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_q5_1_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_q5_1, 2, dequantize_q5_1, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_iq4_nl_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_iq4_nl, 2, dequantize_iq4_nl, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_iq4_xs_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_iq4_xs, QK_NL, dequantize_iq4_xs, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
template [[host_name("kernel_mul_mm_id_iq2_xxs_f32_c")]]
kernel void hf2q_mul_mm_id_impl<block_iq2_xxs, QK_NL, dequantize_iq2_xxs, true>(
constant GgmlMatmulIdMm_MmParams &,
device const char *, device const char *, device const char *, device const char *,
device char *, threadgroup char *, uint3, ushort, ushort, ushort);
// ============================================================================
// ADR-033 §Pi Task #20 debug kernel — dequantize_iq4_xs dump
// ============================================================================
//
// Standalone diagnostic kernel that calls `dequantize_iq4_xs` at all
// 16 il values for ONE block_iq4_xs and writes the dequantized output
// (256 floats) to a buffer. Lets a Rust test compare the kernel's
// dequant output against the canonical CPU reference
// (`test_only_dequantize_iq4_xs`).
//
// Layout: dst[il * 16 + j] for il ∈ [0,16), j ∈ [0,16) = 256 floats.
// Each `il` call produces 16 elements in row-major reg[i/4][i%4] order,
// matching how the mm_id template stages them via the same indexing.
//
// Dispatch: one threadgroup, one thread. We're not benchmarking — just
// reading the dequant output for parity vs CPU.
kernel void hf2q_dequant_iq4_xs_dump(
device const block_iq4_xs * src [[buffer(0)]],
device float * dst [[buffer(1)]],
uint tid [[thread_position_in_grid]]) {
if (tid != 0) return;
for (short il = 0; il < 16; ++il) {
half4x4 reg;
dequantize_iq4_xs(src, il, reg);
for (int i = 0; i < 16; ++i) {
dst[il * 16 + i] = (float)reg[i / 4][i % 4];
}
}
}