// Derived from candle-metal-kernels (Apache-2.0) kernel_mul_mv_id template
// and mlx-native's quantized_matmul_ggml kernels.
// Combines GGML block format dequantization with expert-indexed dispatch.
//
// Original sources:
// candle-metal-kernels/src/metal_src/quantized.metal:7544-7618 (kernel_mul_mv_id)
// candle-metal-kernels/src/metal_src/quantized.metal:90-293 (Q4_0, Q8_0, Q6_K kernels)
// mlx-native/src/shaders/quantized_matmul_ggml.metal (GGML block dequant)
//
// This kernel performs expert-indexed (MoE) quantized matrix-vector multiply:
// For each (token, slot) pair:
// expert_id = ids[token * top_k + slot]
// output[token*top_k + slot, :] = matmul(input[token, :], weight[expert_id])
//
// The key insight: instead of dispatching one kernel per expert, we dispatch once
// for ALL (token, slot) pairs. The kernel uses the ids buffer to route each output
// row to the correct expert's weight slice.
//
// Copyright the candle Authors.
// See LICENSE-APACHE-candle in this directory.
#include <metal_stdlib>
using namespace metal;
// ---- Constants (must match quantized_matmul_ggml.metal) ----
#define QK4_0 32
#define QK8_0 32
#define QK_K 256
#define N_DST 4
#define N_SIMDGROUP 2
#define N_SIMDWIDTH 32
// ---- Parameters for expert-indexed GGML matmul ----
struct GgmlMatvecIdParams {
int64_t ne00; // K: input dimension
int64_t ne01; // N: output dimension per expert
int64_t ne02; // number of experts (on-device ID bound)
int64_t ne10; // K: input dimension (redundant, == ne00)
int64_t ne12; // 1 (unused)
int64_t ne0; // N: output stride
int64_t ne1; // total output rows = n_tokens * top_k
uint r2; // 1
uint r3; // 1
uint top_k; // experts per token
uint n_tokens; // number of input tokens
int64_t expert_stride; // bytes between expert weight slices
};
// Fail closed without a host readback: expert IDs may be produced by an
// earlier kernel in the same command stream. All threads in one threadgroup
// observe the same ID, return before any expert address calculation, and the
// lane(s) that own this output-column tile replace it with NaN.
inline bool poison_invalid_expert_id(
device float * dst,
constant GgmlMatvecIdParams & p,
uint expert_id,
int output_row,
int first_column,
int column_count,
bool write_poison
) {
if ((uint64_t)expert_id < (uint64_t)p.ne02) {
return false;
}
if (write_poison) {
for (int column = 0; column < column_count; column++) {
const int output_column = first_column + column;
if (output_column < p.ne01) {
dst[output_row * p.ne0 + output_column] = NAN;
}
}
}
return true;
}
// K_SCALE_SIZE: bytes used for scales+mins in Q4_K and Q5_K super-blocks.
#define K_SCALE_SIZE 12
// ---- GGML block struct definitions (byte-for-byte with 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;
static_assert(sizeof(block_q2_K) == 2*sizeof(half) + QK_K/16 + QK_K/4,
"wrong q2_K block size");
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");
// Q5_K: 256 values per block, 176 bytes per block.
// Layout: [half d][half dmin][uint8_t scales[12]][uint8_t qh[32]][uint8_t qs[128]]
typedef struct {
half d; // super-block scale for quantized scales
half dmin; // super-block scale for quantized mins
uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits
uint8_t qh[QK_K/8]; // quants, high bit (32 bytes)
uint8_t qs[QK_K/2]; // quants, low 4 bits (128 bytes)
} block_q5_K;
// Q4_K: 256 values per block, 144 bytes per block.
// Layout: [half d][half dmin][uint8_t scales[12]][uint8_t qs[128]]
// Structurally Q5_K minus the 32-byte qh "high-bit" array.
// Same K_SCALE_SIZE=12 layout for packed (sub-scale, sub-min) 6-bit pairs.
//
// GGUF block_q4_K layout.
typedef struct {
half d; // super-block scale for quantized scales
half dmin; // super-block scale for quantized mins
uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits
uint8_t qs[QK_K/2]; // quants, low 4 bits (128 bytes)
} block_q4_K;
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;
// Q5_1: 32 values per block, 24 bytes per block.
// Layout: [half d][half m][uint qh][uint8_t qs[16]]
// Per-element: x[i] = (qs[i] & 0x0F) | (((qh >> i) & 1) << 4)
// x[i+16] = (qs[i] >> 4) | (((qh >> (i+16)) & 1) << 4)
// out[i] = d * x[i] + m
// out[i+16] = d * x[i+16] + m
// ADR-022 Phase 1.
typedef struct {
half d;
half m;
uint qh;
uint8_t qs[QK4_0 / 2]; // 16 bytes — same QK as Q4_0
} block_q5_1;
// IQ4_NL: 32 values per block, 18 bytes per block.
// Layout: [half d][uint8_t qs[16]]
// Each qs byte holds 2 4-bit codebook indices into kvalues_iq4nl[16].
// Per-element: out[i] = d * kvalues_iq4nl[qs[i] & 0x0F]
// out[i+16] = d * kvalues_iq4nl[qs[i] >> 4]
// ADR-022 Phase 1.
typedef struct {
half d;
uint8_t qs[QK4_0 / 2]; // 16 bytes
} block_iq4_nl;
// IQ4_NL non-linear codebook (frozen by the GGUF format).
// Any drift breaks every IQ4_NL GGUF on disk; do not edit without
// updating the host-side `KVALUES_IQ4_NL` in src/gguf/mod.rs in lock-step.
constant int8_t kvalues_iq4nl[16] = {
-127, -104, -83, -65, -49, -35, -22, -10,
1, 13, 25, 38, 53, 69, 89, 113
};
// IQ4_XS (ADR-033 §Pi 2026-05-22): super-block 4-bit codebook quant.
// 256 elements per super-block, 8 × 32-element sub-blocks each with a
// 6-bit signed scale (4-bit scales_l + 2-bit scales_h). 136 bytes.
// Reference: ggml-common.h:444-450 `block_iq4_xs`.
typedef struct {
half d;
uint16_t scales_h;
uint8_t scales_l[QK_K / 64]; // 4
uint8_t qs[QK_K / 2]; // 128
} block_iq4_xs;
// IQ2_XXS (ADR-066 N-IQ2XXS): 2.0625 bpw grid codebook, 66 bytes per
// 256 values. Each 32-value sub-block reads qs[4*ib32 .. 4*ib32+4] as two
// u32: four 8-bit grid indices, then four 7-bit sign indices and a 4-bit
// scale in bits 28..31. Reference: ggml-common.h `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
// ---- Q5_1 dot product helper (ADR-022 Phase 1) ----
//
// Mirrors the reference `block_q_n_dot_y<block_q5_1>`
// helper. Q5_1 differs from Q4_0 by carrying
// (a) an additional `m` (min) term contributing `m * sumy` to the dot,
// and (b) a 5th high-bit per element packed in `qh`.
//
// `il` is 0 or 8 (mlx-native convention: byte offset within the
// 16-byte qs array). The yl[] vector is pre-scaled by the caller
// (yl[i+1] /= 256, yl[i+8] /= 16, yl[i+9] /= 4096) so that masking
// nibbles via 0x000F / 0x0F00 / 0x00F0 / 0xF000 yields the correct
// sub-element-weighted partial sums.
inline float block_q5_1_dot_y(
device const block_q5_1 * qb,
float sumy,
thread float * yl,
int il
) {
float d = qb->d;
float m = qb->m;
float4 acc = 0.f;
// qs is 16 bytes starting at offset 8 in the block.
// Cast block as uint16_t* to skip d (1 uint16) + m (1 uint16) +
// qh (2 uint16) = 4 uint16; then add il/2 to land at the right
// qs sub-region.
device const uint16_t * qs = ((device const uint16_t *)qb + 4 + il/2);
const uint qh = qb->qh;
for (int i = 0; i < 8; i += 2) {
// Low nibbles, sub-positions i + 0 and i + 1 (within block 0..15).
// qh bit (i+0+il) contributes the 5th bit, placed at position 4
// (mask 0x10) for the 0x000F-masked nibble.
// qh bit (i+1+il) contributes the 5th bit, placed at position 12
// (mask 0x1000) for the 0x0F00-masked nibble.
acc[0] += yl[i + 0]
* (float)((qs[i / 2] & 0x000F) | (((qh >> (i + 0 + il )) << 4 ) & 0x0010));
acc[1] += yl[i + 1]
* (float)((qs[i / 2] & 0x0F00) | (((qh >> (i + 1 + il )) << 12) & 0x1000));
// High nibbles, sub-positions i + 8 and i + 9 (within block 16..31).
// qh bit (i+0+il+16) for nibble at mask 0x00F0 → bit at 0x0100.
// qh bit (i+1+il+16) for nibble at mask 0xF000 → bit at 0x10000.
acc[2] += yl[i + 8]
* (float)((qs[i / 2] & 0x00F0) | (((qh >> (i + 0 + il + 16)) << 8 ) & 0x0100));
acc[3] += yl[i + 9]
* (float)((qs[i / 2] & 0xF000) | (((qh >> (i + 1 + il + 16)) << 16) & 0x10000));
}
return d * (acc[0] + acc[1] + acc[2] + acc[3]) + sumy * m;
}
inline float block_q5_0_dot_y(
device const block_q5_0 * qb,
float sumy,
thread float * yl,
int il
) {
const float d = qb->d;
float4 acc = 0.f;
device const uint16_t * qs = ((device const uint16_t *)qb + 3 + il/2);
const uint qh = *((device const uint *)qb->qh);
for (int i = 0; i < 8; i += 2) {
acc[0] += yl[i + 0] * (float)((qs[i / 2] & 0x000F) | (((qh >> (i + 0 + il )) << 4 ) & 0x0010));
acc[1] += yl[i + 1] * (float)((qs[i / 2] & 0x0F00) | (((qh >> (i + 1 + il )) << 12) & 0x1000));
acc[2] += yl[i + 8] * (float)((qs[i / 2] & 0x00F0) | (((qh >> (i + 0 + il + 16)) << 8 ) & 0x0100));
acc[3] += yl[i + 9] * (float)((qs[i / 2] & 0xF000) | (((qh >> (i + 1 + il + 16)) << 16) & 0x10000));
}
return d * (sumy * -16.f + acc[0] + acc[1] + acc[2] + acc[3]);
}
// ---- IQ4_NL dot product helper (ADR-022 Phase 1) ----
//
// IQ4_NL is a 4-bit codebook quant: each qs nibble selects one of
// 16 entries from `kvalues_iq4nl`. There is no zero-point bias and no
// `m` term. The output formula is purely `out[i] = d * kvalues_iq4nl[idx]`.
//
// Caller convention is identical to Q4_0 / Q5_1: yl[] holds the
// pre-scaled input row, but for IQ4_NL the divisors used by Q4_0
// (`/256`, `/16`, `/4096`) do NOT compose because the codebook lookup
// is non-linear. This helper therefore reads the raw input via a
// caller-supplied `yl_raw[]` array of 16 unscaled values, multiplies
// by the codebook entry directly, and returns `d * sum`.
inline float block_iq4_nl_dot_y(
device const block_iq4_nl * qb,
thread float * yl_raw,
int il
) {
float d = qb->d;
float acc = 0.f;
// qs starts at byte 2 in the block; 16 bytes total.
device const uint8_t * qs = qb->qs + il;
for (int i = 0; i < 8; i++) {
const uint8_t b = qs[i];
const int lo = b & 0x0F;
const int hi = (b >> 4) & 0x0F;
// First half: position i within sub-region [il .. il+8).
acc += yl_raw[i] * (float)kvalues_iq4nl[lo];
// Second half: position i + 16 (the high half of the block).
acc += yl_raw[i + 8] * (float)kvalues_iq4nl[hi];
}
return d * acc;
}
// IQ4_XS dot helper — see comment in quantized_matmul_ggml.metal for full
// docs. Processes ONE 32-element sub-block of a super-block.
// `ib32` ∈ [0,7] selects sub-block; `il` ∈ {0, 8} selects which 8-element
// half of the sub-block to process.
inline float block_iq4_xs_dot_y(
device const block_iq4_xs * qb,
thread float * yl_raw,
int ib32,
int il
) {
const float d = qb->d;
const int ls = ((qb->scales_l[ib32 / 2] >> 4*(ib32 % 2)) & 0xf)
| (((qb->scales_h >> 2*ib32) & 0x3) << 4);
const float dl = d * (float)(ls - 32);
float acc = 0.f;
device const uint8_t * qs = qb->qs + 16*ib32 + il;
for (int i = 0; i < 8; i++) {
const uint8_t b = qs[i];
acc += yl_raw[i] * (float)kvalues_iq4nl[b & 0x0F];
acc += yl_raw[i + 8] * (float)kvalues_iq4nl[(b >> 4) & 0x0F];
}
return dl * acc;
}
// ---- Q4_0 dot product helper (identical to quantized_matmul_ggml.metal) ----
inline float block_q4_0_dot_y(
device const block_q4_0 * qb,
float sumy,
thread float * yl,
int il
) {
float d = qb->d;
float2 acc = 0.f;
device const uint16_t * qs = ((device const uint16_t *)qb + 1 + il/2);
for (int i = 0; i < 8; i += 2) {
acc[0] += yl[i + 0] * (qs[i / 2] & 0x000F)
+ yl[i + 1] * (qs[i / 2] & 0x0F00);
acc[1] += yl[i + 8] * (qs[i / 2] & 0x00F0)
+ yl[i + 9] * (qs[i / 2] & 0xF000);
}
return d * (sumy * -8.f + acc[0] + acc[1]);
}
// ====================================================================
// Q4_0 expert-indexed mat-vec kernel
// ====================================================================
//
// For each output row r (where r = token*top_k + slot):
// expert_id = ids[r] (ids is [n_tokens * top_k] flat, pre-expanded)
// src0_cur = src0 + expert_id * expert_stride
// output[r] = matmul(src1[token], src0_cur)
//
// Dispatch geometry: threadgroups=(ceil(N/8), n_tokens*top_k, 1), tg=(8,8,1)
//
// Routing index is in dim Y, NOT dim Z (despite the peer's mul_mv_id grid
// putting ne123 in z). Tested 2026-04-26 on M5 Max
// dwq46 64-token decode: switching kernel to read tgpig.z + dispatcher
// MTLSize::new(N/8, 1, m) regressed throughput from 114 t/s to 90.9 t/s
// (-19%). Apple GPU's threadgroup scheduler distributes this dispatch
// shape better via Y than Z — 7th confirmed static-evidence kernel
// hypothesis falsified per `project_metal_compiler_auto_optimizes_static_levers.md`.
kernel void kernel_mul_mv_id_q4_0_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nw = N_SIMDWIDTH;
const int nb = p.ne00 / QK4_0;
const int r0 = tgpig.x;
const int output_row = tgpig.y; // flat index into output: token*top_k + slot
// Bounds check
if (output_row >= (int)p.ne1) return;
// Determine which token this output row belongs to and which expert
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
// Point to the expert's weight slice
device const block_q4_0 * x = (device const block_q4_0 *)((device const char *)src0 + expert_id * p.expert_stride) + first_row * nb;
// Point to the input row for this token
device const float * y = src1 + token_idx * p.ne10;
float yl[16];
float sumf[nr] = {0.f};
const int ix = tiisg / 2;
const int il = (tiisg % 2) * 8;
device const float * yb = y + ix * QK4_0 + il;
for (int ib = ix; ib < nb; ib += nw/2) {
float sumy = 0;
for (int i = 0; i < 8; i += 2) {
sumy += yb[i] + yb[i+1];
yl[i+0] = yb[i+0];
yl[i+1] = yb[i+1] / 256.f;
sumy += yb[i+16] + yb[i+17];
yl[i+8] = yb[i+16] / 16.f;
yl[i+9] = yb[i+17] / 4096.f;
}
for (int row = 0; row < nr; row++) {
sumf[row] += block_q4_0_dot_y(x + ib + row*nb, sumy, yl, il);
}
yb += QK4_0 * 16;
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
// ====================================================================
// Q4_0 fused-SwiGLU expert-indexed mat-vec kernel
// ====================================================================
//
// Computes: dst[r][n] = sum_k(dequant(W_q4_0[expert_id][n][k])
// * (silu(gate[r][k]) * up[r][k]))
//
// where r = token*top_k + slot, expert_id = ids[r].
//
// Replaces the dispatch sequence:
// silu_mul_f32(gate, up → h_all) # 1 dispatch + memory_barrier
// kernel_mul_mv_id_q4_0_f32(W, h_all → dst) # 1 dispatch
//
// with a single dispatch that reads gate + up directly and computes
// swiglu inline before the dot product. Closes ~5-10µs/layer × 40
// layers ≈ 0.3-0.4ms/token of CPU dispatch overhead in the dwq46
// decode hot path (ADR-012 §Optimize / Task #15).
//
// Buffer layout:
// buffer(0): src0 - Q4_0 packed weight, [n_experts, N, K/QK4_0] blocks
// buffer(1): gate - f32 [n_tokens*top_k, K]
// buffer(2): up - f32 [n_tokens*top_k, K]
// buffer(3): dst - f32 [n_tokens*top_k, N]
// buffer(4): ids - u32 [n_tokens*top_k]
// buffer(5): params - GgmlMatvecIdParams
//
// Dispatch geometry: identical to kernel_mul_mv_id_q4_0_f32 —
// threadgroups=(ceil(N/8), n_tokens*top_k, 1), tg=(8, 8, 1).
kernel void kernel_mul_mv_id_q4_0_f32_swiglu(
device const char * src0 [[buffer(0)]],
device const float * gate [[buffer(1)]],
device const float * up [[buffer(2)]],
device float * dst [[buffer(3)]],
device const uint * ids [[buffer(4)]],
constant GgmlMatvecIdParams & p [[buffer(5)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nw = N_SIMDWIDTH;
const int nb = p.ne00 / QK4_0;
const int r0 = tgpig.x;
const int output_row = tgpig.y; // flat index into output: token*top_k + slot
if (output_row >= (int)p.ne1) return;
// For expert_down in the decode-time MoE pipeline, the input row IS
// the output row index (one h_all row per (token, expert_slot) pair),
// not token_idx. Each (token, expert_slot) has its own gate/up vectors.
const uint expert_id = ids[output_row];
const uint input_row = output_row; // gate/up are pre-routed per (token, slot).
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
// Expert's weight slice.
device const block_q4_0 * x = (device const block_q4_0 *)((device const char *)src0 + expert_id * p.expert_stride) + first_row * nb;
// Per-row gate and up vectors.
device const float * gate_y = gate + input_row * p.ne10;
device const float * up_y = up + input_row * p.ne10;
float yl[16];
float sumf[nr] = {0.f};
const int ix = tiisg / 2;
const int il = (tiisg % 2) * 8;
device const float * gb = gate_y + ix * QK4_0 + il;
device const float * ub = up_y + ix * QK4_0 + il;
for (int ib = ix; ib < nb; ib += nw/2) {
float sumy = 0;
// Compute swiglu = silu(gate) * up = gate * sigmoid(gate) * up
// for each of the 16 active elements per simdthread, reusing the
// same yl[] / sumy aggregation as the unfused kernel.
for (int i = 0; i < 8; i += 2) {
// Lane block 0 (i=0, i+1).
float g0 = gb[i+0];
float g1 = gb[i+1];
float u0 = ub[i+0];
float u1 = ub[i+1];
float s0 = (g0 / (1.0f + metal::exp(-g0))) * u0;
float s1 = (g1 / (1.0f + metal::exp(-g1))) * u1;
sumy += s0 + s1;
yl[i+0] = s0;
yl[i+1] = s1 / 256.f;
// Lane block 1 (i+16, i+17).
float g2 = gb[i+16];
float g3 = gb[i+17];
float u2 = ub[i+16];
float u3 = ub[i+17];
float s2 = (g2 / (1.0f + metal::exp(-g2))) * u2;
float s3 = (g3 / (1.0f + metal::exp(-g3))) * u3;
sumy += s2 + s3;
yl[i+8] = s2 / 16.f;
yl[i+9] = s3 / 4096.f;
}
for (int row = 0; row < nr; row++) {
sumf[row] += block_q4_0_dot_y(x + ib + row*nb, sumy, yl, il);
}
gb += QK4_0 * 16;
ub += QK4_0 * 16;
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
// ====================================================================
// Q8_0 expert-indexed mat-vec kernel
// ====================================================================
#define NB_Q8_0 8
kernel void kernel_mul_mv_id_q8_0_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nw = N_SIMDWIDTH;
const int nb = p.ne00 / QK8_0;
const int r0 = tgpig.x;
const int output_row = tgpig.y;
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
device const block_q8_0 * x = (device const block_q8_0 *)((device const char *)src0 + expert_id * p.expert_stride) + first_row * nb;
device const float * y = src1 + token_idx * p.ne10;
float yl[NB_Q8_0];
float sumf[nr] = {0.f};
const int ix = tiisg / 4;
const int il = tiisg % 4;
device const float * yb = y + ix * QK8_0 + NB_Q8_0 * il;
for (int ib = ix; ib < nb; ib += nw/4) {
for (int i = 0; i < NB_Q8_0; ++i) {
yl[i] = yb[i];
}
for (int row = 0; row < nr; row++) {
device const int8_t * qs = x[ib + row*nb].qs + NB_Q8_0 * il;
float sumq = 0.f;
for (int iq = 0; iq < NB_Q8_0; ++iq) {
sumq += qs[iq] * yl[iq];
}
sumf[row] += sumq * x[ib + row*nb].d;
}
yb += NB_Q8_0 * nw;
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
// Expert-routed sibling of the dense correctness-first Q2_K matvec.
kernel void kernel_mul_mv_id_q2_K_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = 4;
const int output_row = int(tgpig.y);
if (output_row >= p.ne1) return;
const int nb = p.ne00 / QK_K;
const uint token_idx = uint(output_row) / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (int(tgpig.x) * N_SIMDGROUP + int(sgitg)) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
device const block_q2_K * x = (device const block_q2_K *)(src0 + expert_id*p.expert_stride)
+ first_row*nb;
device const float * y = src1 + token_idx*p.ne10;
float yl[32];
float sumf[nr] = {0.f};
const short ix = tiisg / 8;
const short it = tiisg % 8;
const short iq = it / 4;
const short ir = it % 4;
const short is = (8*ir) / 16;
device const float * y4 = y + ix*QK_K + 128*iq + 8*ir;
for (int ib = ix; ib < nb; ib += 4) {
float4 sumy = {0.f, 0.f, 0.f, 0.f};
for (short i = 0; i < 8; ++i) {
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
}
for (int row = 0; row < nr; ++row) {
device const block_q2_K * block = x + ib + row*nb;
device const uint8_t * sc = block->scales + 8*iq + is;
device const uint16_t * qs = (device const uint16_t *)block->qs + 16*iq + 4*ir;
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
for (int i = 0; i < 8; i += 2) {
acc1[0] += yl[i + 0] * (qs[i/2] & 0x0003);
acc2[0] += yl[i + 1] * (qs[i/2] & 0x0300);
acc1[1] += yl[i + 8] * (qs[i/2] & 0x000c);
acc2[1] += yl[i + 9] * (qs[i/2] & 0x0c00);
acc1[2] += yl[i + 16] * (qs[i/2] & 0x0030);
acc2[2] += yl[i + 17] * (qs[i/2] & 0x3000);
acc1[3] += yl[i + 24] * (qs[i/2] & 0x00c0);
acc2[3] += yl[i + 25] * (qs[i/2] & 0xc000);
}
const float d = block->d;
const float dmin = float(block->dmin) / 16.f;
sumf[row] += d * ((acc1[0] + acc2[0]/256.f) * (sc[0] & 0x0f) +
(acc1[1] + acc2[1]/256.f) * (sc[2] & 0x0f) / 4.f +
(acc1[2] + acc2[2]/256.f) * (sc[4] & 0x0f) / 16.f +
(acc1[3] + acc2[3]/256.f) * (sc[6] & 0x0f) / 64.f) -
dmin * (sumy[0] * (sc[0] & 0xf0) + sumy[1] * (sc[2] & 0xf0) +
sumy[2] * (sc[4] & 0xf0) + sumy[3] * (sc[6] & 0xf0));
}
y4 += 4*QK_K;
}
for (int row = 0; row < nr; ++row) {
const float total = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row*p.ne0 + first_row + row] = total;
}
}
}
// Expert-routed variant of the current `kernel_mul_mv_q3_K_f32_impl`
// scheme, preserving its two rows per simdgroup and accumulator order.
kernel void kernel_mul_mv_id_q3_K_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
constexpr short nr = 2;
constexpr short nsg = 2;
const int output_row = int(tgpig.y);
if (output_row >= p.ne1) return;
const int nb = p.ne00/QK_K;
const uint token_idx = uint(output_row)/p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (int(tgpig.x)*nsg + int(sgitg))*nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
device const block_q3_K * x =
(device const block_q3_K *)(src0 + expert_id*p.expert_stride) + first_row*nb;
device const float * yy = src1 + token_idx*p.ne10;
float yl[32];
const short tid = tiisg/4;
const short ix = tiisg%4;
const short ip = tid/4;
const short il = 2*((tid%4)/2);
const short ir = tid%2;
const short l0 = 8*ir;
const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200},
{0x0004, 0x0400, 0x0008, 0x0800},
{0x0010, 0x1000, 0x0020, 0x2000},
{0x0040, 0x4000, 0x0080, 0x8000}};
const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00},
{0x0030, 0x3000, 0x00c0, 0xc000}};
const ushort4 hm = mm[2*ip + il/2];
const short shift = 2*il;
const float v1 = il == 0 ? 4.f : 64.f;
const float v2 = 4.f*v1;
const uint16_t s_shift1 = 4*ip;
const uint16_t s_shift2 = s_shift1 + il;
const short q_offset = 32*ip + l0;
const short y_offset = 128*ip + 32*il + l0;
device const float * y1 = yy + ix*QK_K + y_offset;
float sumf1[nr] = {0.f};
float sumf2[nr] = {0.f};
for (int ib = ix; ib < nb; ib += 4) {
for (short l = 0; l < 8; ++l) {
yl[l+ 0] = y1[l+ 0];
yl[l+ 8] = y1[l+16];
yl[l+16] = y1[l+32];
yl[l+24] = y1[l+48];
}
for (short row = 0; row < nr; ++row) {
device const block_q3_K * block = x + ib + row*nb;
device const uint16_t * q = (device const uint16_t *)(block->qs + q_offset);
device const uint16_t * h = (device const uint16_t *)(block->hmask + l0);
device const uint16_t * a = (device const uint16_t *)block->scales;
uint32_t scales32;
uint32_t aux32;
thread uint16_t * scales16 = (thread uint16_t *)&scales32;
thread const int8_t * scales = (thread const int8_t *)&scales32;
scales16[0] = a[4]; scales16[1] = a[5];
aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030;
scales16[0] = a[il+0]; scales16[1] = a[il+1];
scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32;
float s1=0.f, s2=0.f, s3=0.f, s4=0.f, s5=0.f, s6=0.f;
for (short l = 0; l < 8; l += 2) {
const int32_t qs = q[l/2];
s1 += yl[l+0]*(qs&qm[il/2][0]);
s2 += yl[l+1]*(qs&qm[il/2][1]);
s3 += ((h[l/2]&hm[0]) ? 0.f : yl[l+0]) + ((h[l/2]&hm[1]) ? 0.f : yl[l+1]);
s4 += yl[l+16]*(qs&qm[il/2][2]);
s5 += yl[l+17]*(qs&qm[il/2][3]);
s6 += ((h[l/2]&hm[2]) ? 0.f : yl[l+16]) + ((h[l/2]&hm[3]) ? 0.f : yl[l+17]);
}
float d1 = float(block->d)*(s1 + s2/256.f - s3*v1);
float d2 = float(block->d)*(s4 + s5/256.f - s6*v2);
sumf1[row] += d1*(scales[0]-32);
sumf2[row] += d2*(scales[2]-32);
s1=s2=s3=s4=s5=s6=0.f;
for (short l = 0; l < 8; l += 2) {
const int32_t qs = q[l/2+8];
s1 += yl[l+8]*(qs&qm[il/2][0]);
s2 += yl[l+9]*(qs&qm[il/2][1]);
s3 += ((h[l/2+8]&hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8]&hm[1]) ? 0.f : yl[l+9]);
s4 += yl[l+24]*(qs&qm[il/2][2]);
s5 += yl[l+25]*(qs&qm[il/2][3]);
s6 += ((h[l/2+8]&hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8]&hm[3]) ? 0.f : yl[l+25]);
}
d1 = float(block->d)*(s1 + s2/256.f - s3*v1);
d2 = float(block->d)*(s4 + s5/256.f - s6*v2);
sumf1[row] += d1*(scales[1]-32);
sumf2[row] += d2*(scales[3]-32);
}
y1 += 4*QK_K;
}
for (short row = 0; row < nr; ++row) {
const float total = simd_sum((sumf1[row] + 0.25f*sumf2[row])/(1 << shift));
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row*p.ne0 + first_row + row] = total;
}
}
}
// ====================================================================
// Q8_0 _id expert-indexed mat-vec kernel — NR0=2 NSG=4 variant
// (ADR-029 iter-6 port; peer N_R0_Q8_0=2 + N_SG_Q8_0=4)
// ====================================================================
//
// Same row coverage as `kernel_mul_mv_id_q8_0_f32` (8 rows/TG) but
// distributes the K-dim across 4 simdgroups with shmem cross-SG reduce
// instead of 2 SGs each independently handling 4 rows. Better latency
// hiding when memory-bandwidth-bound, which the gemma4 MoE down_exps
// Q8_0 mat-vec hits at hidden_size=2816, top_k=8.
//
// Dispatch geometry:
// threadgroups = (ceil(N/N_R0_Q8_0=2), n_tokens*top_k, 1)
// threads_per_tg = (32, 4, 1) = 128 threads = 4 simdgroups × 32
// shmem = N_R0_Q8_0 * N_SIMDWIDTH * sizeof(float) = 256 bytes
//
// Env-gated via HF2Q_Q8_0_ID_MV_NR2=1 in dispatch_id_mv.
#define N_R0_Q8_0_ID 2
#define N_SG_Q8_0_ID 4
#define NQ_Q8_0_ID 8
kernel void kernel_mul_mv_id_q8_0_f32_nr2(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
threadgroup float * shmem [[threadgroup(0)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
constexpr int NR0 = N_R0_Q8_0_ID; // 2
constexpr int NSG = N_SG_Q8_0_ID; // 4
constexpr int NW = N_SIMDWIDTH; // 32
constexpr int NQ = NQ_Q8_0_ID; // 8
const int nb = p.ne00 / QK8_0;
const int r0 = tgpig.x;
const int output_row_base = tgpig.y;
if (output_row_base >= (int)p.ne1) return;
const uint token_idx = output_row_base / p.top_k;
const uint expert_id = ids[output_row_base];
const int first_row = r0 * NR0;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row_base, first_row, NR0,
tiisg == 0 && sgitg == 0)) return;
// Per-row src0 pointers (unrolled NR0=2 iterations), with expert offset.
device const block_q8_0 * ax[NR0];
for (int row = 0; row < NR0; ++row) {
ax[row] = (device const block_q8_0 *)((device const char *)src0 + expert_id * p.expert_stride) + (first_row + row) * nb;
}
device const float * y = src1 + token_idx * p.ne10;
float sumf[NR0] = { 0.f };
const int ix = tiisg / (NW / NQ); // 0..3
const int il = tiisg % (NW / NQ); // 0..3
const int ib0 = sgitg * NQ + ix;
float yl[NQ];
device const float * yb = y + ib0 * QK8_0 + il * NQ;
// Each thread covers NQ quants per iteration; SGs interleave across
// ib by stride NSG*NQ. Mirrors kernel_mul_mv_q8_0_f32_nr2 (regular).
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
for (int i = 0; i < NQ; ++i) {
yl[i] = yb[i];
}
for (int row = 0; row < NR0; ++row) {
device const int8_t * qs = ax[row][ib].qs + il * NQ;
float sumq = 0.f;
for (int iq = 0; iq < NQ; ++iq) {
sumq += qs[iq] * yl[iq];
}
sumf[row] += sumq * ax[row][ib].d;
}
yb += NSG * NQ * QK8_0;
}
// Cross-simdgroup reduction (peer's helper_mv_reduce_and_write pattern).
threadgroup float * shmem_rows[NR0];
for (int row = 0; row < NR0; ++row) {
shmem_rows[row] = shmem + NW * row;
if (sgitg == 0) {
shmem_rows[row][tiisg] = 0.0f;
}
sumf[row] = simd_sum(sumf[row]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int row = 0; row < NR0; ++row) {
if (tiisg == 0) {
shmem_rows[row][sgitg] = sumf[row];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int row = 0; row < NR0 && first_row + row < (int)p.ne01; ++row) {
const float tot = simd_sum(shmem_rows[row][tiisg]);
if (tiisg == 0 && sgitg == 0) {
dst[output_row_base * p.ne0 + first_row + row] = tot;
}
}
}
// ====================================================================
// Q5_1 expert-indexed mat-vec kernel (ADR-022 Phase 1)
// ====================================================================
//
// Same dispatch geometry as Q4_0 / Q8_0 (legacy 32-element formats):
// threadgroups = (ceil(N/8), n_tokens*top_k, 1), tg = (8, 8, 1)
// tgpig.x = block-row index
// tgpig.y = flat output row (token*top_k + slot)
//
// Differs from Q4_0 only in (a) the block-q-n typedef walked, and
// (b) the dot helper used (`block_q5_1_dot_y` carries the m*sumy term
// and qh-bit injection; `block_q4_0_dot_y` does not).
//
// Reference: `mul_vec_q_n_f32_impl<block_q5_1, N_R0_Q5_1>`.
// Inlined here in mlx-native
// style (matching the Q4_0 / Q8_0 ports above) rather than templated.
kernel void kernel_mul_mv_id_q5_1_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nw = N_SIMDWIDTH;
// Q5_1 has the same QK as Q4_0 (32 elements per block).
const int nb = p.ne00 / QK4_0;
const int r0 = tgpig.x;
const int output_row = tgpig.y;
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
// Point to the expert's weight slice. Q5_1 block stride is 24 bytes
// (vs 18 for Q4_0); the expert_stride byte offset is honored by the
// host-side dispatcher.
device const block_q5_1 * x = (device const block_q5_1 *)((device const char *)src0
+ expert_id * p.expert_stride) + first_row * nb;
device const float * y = src1 + token_idx * p.ne10;
float yl[16];
float sumf[nr] = {0.f};
const int ix = tiisg / 2;
const int il = (tiisg % 2) * 8;
device const float * yb = y + ix * QK4_0 + il;
for (int ib = ix; ib < nb; ib += nw/2) {
float sumy = 0;
// Pre-scale yl[] for the masked-nibble × scaled-input trick
// (same as Q4_0 — works because Q5_1's qh-bit injection happens
// at bit positions {4, 12, 8, 16} that align with the same
// masks 0x000F / 0x0F00 / 0x00F0 / 0xF000).
for (int i = 0; i < 8; i += 2) {
sumy += yb[i] + yb[i+1];
yl[i+0] = yb[i+0];
yl[i+1] = yb[i+1] / 256.f;
sumy += yb[i+16] + yb[i+17];
yl[i+8] = yb[i+16] / 16.f;
yl[i+9] = yb[i+17] / 4096.f;
}
for (int row = 0; row < nr; row++) {
sumf[row] += block_q5_1_dot_y(x + ib + row*nb, sumy, yl, il);
}
yb += QK4_0 * 16;
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
kernel void kernel_mul_mv_id_q5_0_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nw = N_SIMDWIDTH;
const int nb = p.ne00 / QK4_0;
const int r0 = tgpig.x;
const int output_row = tgpig.y;
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
device const block_q5_0 * x = (device const block_q5_0 *)((device const char *)src0
+ expert_id * p.expert_stride) + first_row * nb;
device const float * y = src1 + token_idx * p.ne10;
float yl[16];
float sumf[nr] = {0.f};
const int ix = tiisg / 2;
const int il = (tiisg % 2) * 8;
device const float * yb = y + ix * QK4_0 + il;
for (int ib = ix; ib < nb; ib += nw/2) {
float sumy = 0.f;
for (int i = 0; i < 8; i += 2) {
sumy += yb[i] + yb[i+1];
yl[i+0] = yb[i+0];
yl[i+1] = yb[i+1] / 256.f;
sumy += yb[i+16] + yb[i+17];
yl[i+8] = yb[i+16] / 16.f;
yl[i+9] = yb[i+17] / 4096.f;
}
for (int row = 0; row < nr; row++) {
sumf[row] += block_q5_0_dot_y(x + ib + row*nb, sumy, yl, il);
}
yb += QK4_0 * 16;
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
// ====================================================================
// IQ4_NL expert-indexed mat-vec kernel (ADR-022 Phase 1)
// ====================================================================
//
// Same dispatch geometry as Q4_0. IQ4_NL's codebook lookup is
// non-linear, so the masked-nibble × pre-scaled-yl trick that Q4_0 /
// Q5_1 use does not compose. We pass the raw input row to the dot
// helper, which multiplies element-wise by the looked-up codebook
// values.
//
// Reference: `kernel_mul_mv_iq4_nl_f32_impl` (instantiated upstream via
// kernel_mul_mv_id<mmv_fn<...>>); inlined here in mlx-native style.
// ====================================================================
// IQ4_XS expert-indexed mat-vec kernel (ADR-033 §Pi Task #16 2026-05-22)
// ====================================================================
//
// MoE-routed variant of `kernel_mul_mv_iq4_xs_f32` (see
// quantized_matmul_ggml.metal). Each output_row picks an expert id
// from `ids[output_row]` and reads that expert's weight slab
// (offset = expert_id * expert_stride). Geometry mirrors IQ4_NL's
// mv_id (N_SIMDGROUP=2, N_DST=4, threadgroup=(8,8,1)).
//
// Reference: `kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_xs_f32_impl>>`.
// Inlined here in mlx-native style (no
// template metaprogramming).
kernel void kernel_mul_mv_id_iq4_xs_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nb = p.ne00 / QK_K; // super-blocks per row
const int r0 = tgpig.x;
const int output_row = tgpig.y;
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
device const block_iq4_xs * x = (device const block_iq4_xs *)((device const char *)src0
+ expert_id * p.expert_stride) + first_row * nb;
device const float * y = src1 + token_idx * p.ne10;
// tiisg partitions super-block × (sub-block, half) work — see
// kernel_mul_mv_iq4_xs_f32 comment for derivation.
const int ix = tiisg / 16; // 0 or 1
const int it = tiisg % 16; // 0..15
const int ib32 = it / 2;
const int il = (it % 2) * 8;
float yl_raw[16];
float sumf[nr] = {0.f};
for (int ibl = ix; ibl < nb; ibl += 2) {
device const float * yb = y + ibl * QK_K + ib32 * 32 + il;
for (int i = 0; i < 8; i++) {
yl_raw[i] = yb[i];
yl_raw[i + 8] = yb[i + 16];
}
for (int row = 0; row < nr; row++) {
sumf[row] += block_iq4_xs_dot_y(x + ibl + row*nb, yl_raw, ib32, il);
}
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
kernel void kernel_mul_mv_id_iq4_nl_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr = N_DST;
const int nsg = N_SIMDGROUP;
const int nw = N_SIMDWIDTH;
const int nb = p.ne00 / QK4_0;
const int r0 = tgpig.x;
const int output_row = tgpig.y;
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr, tiisg == 0)) return;
device const block_iq4_nl * x = (device const block_iq4_nl *)((device const char *)src0
+ expert_id * p.expert_stride) + first_row * nb;
device const float * y = src1 + token_idx * p.ne10;
float yl_raw[16];
float sumf[nr] = {0.f};
const int ix = tiisg / 2;
const int il = (tiisg % 2) * 8;
device const float * yb = y + ix * QK4_0 + il;
for (int ib = ix; ib < nb; ib += nw/2) {
// Raw yl[] for IQ4_NL — no pre-scale, codebook lookup is
// non-linear and would be polluted by the /256, /16, /4096
// divisors used by the linear-quant helpers.
for (int i = 0; i < 8; i++) {
yl_raw[i] = yb[i];
yl_raw[i + 8] = yb[i + 16];
}
for (int row = 0; row < nr; row++) {
sumf[row] += block_iq4_nl_dot_y(x + ib + row*nb, yl_raw, il);
}
yb += QK4_0 * 16;
}
for (int row = 0; row < nr; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot;
}
}
}
// ====================================================================
// Q5_K expert-indexed mat-vec kernel
// ====================================================================
//
// Dispatch geometry (same as Q6_K): threadgroups = (ceil(N/2), n_tokens*top_k, 1)
// tgpig.x = weight-row-pair index (two rows: 2*r0 + sgitg)
// tgpig.y = flat output row (token*top_k + slot)
// sgitg = selects which of the two rows this simdgroup processes
//
// Ported from candle-metal-kernels kernel_mul_mv_q5_K_f32_impl with
// the expert-routing indirection from the Q6_K _id kernel above.
// Copyright the candle Authors (Apache-2.0).
kernel void kernel_mul_mv_id_q5_K_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nb = p.ne00 / QK_K;
const int64_t r0 = tgpig.x;
const int output_row = tgpig.y; // flat: token*top_k + slot
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
// Each threadgroup covers weight-row pair (2*r0, 2*r0+1);
// sgitg selects which row this simdgroup computes.
const int row = 2 * (int)r0 + (int)sgitg;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, row, 1, tiisg == 0)) return;
// Point to the expert's weight slice and the token's input row.
device const block_q5_K * x = (device const block_q5_K *)((device const char *)src0 + expert_id * p.expert_stride) + row * nb;
device const float * yy = src1 + token_idx * p.ne10;
float sumf = 0.f;
const uint16_t kmask1 = 0x3f3f;
const uint16_t kmask2 = 0x0f0f;
const uint16_t kmask3 = 0xc0c0;
const int tid = tiisg / 4;
const int ix = tiisg % 4;
const int iq = tid / 4;
const int ir = tid % 4;
const int n = 8;
const int l0 = n * ir;
const int q_offset = 32 * iq + l0;
const int y_offset = 64 * iq + l0;
const uint8_t hm1 = 1u << (2 * iq);
const uint8_t hm2 = hm1 << 1;
const uint8_t hm3 = hm1 << 4;
const uint8_t hm4 = hm2 << 4;
uint16_t sc16[4];
thread const uint8_t * sc8 = (thread const uint8_t *)sc16;
device const float * y1 = yy + ix * QK_K + y_offset;
for (int i = ix; i < nb; i += 4) {
device const uint8_t * q1 = x[i].qs + q_offset;
device const uint8_t * q2 = q1 + 64;
device const uint8_t * qh = x[i].qh + l0;
device const half * dh = &x[i].d;
// scales array is uint8_t[12]; cast to uint16_t[6] for the
// sc16 decoding identical to the reference candle kernel.
device const uint16_t * a = (device const uint16_t *)x[i].scales + iq;
device const float * y2 = y1 + 128;
float yl[16], yh[16];
float4 sumy = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
yl[l+0] = y1[l + 0]; sumy[0] += yl[l+0];
yl[l+8] = y1[l + 32]; sumy[1] += yl[l+8];
yh[l+0] = y2[l + 0]; sumy[2] += yh[l+0];
yh[l+8] = y2[l + 32]; sumy[3] += yh[l+8];
}
sc16[0] = a[0] & kmask1;
sc16[1] = a[2] & kmask1;
sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2);
sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2);
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
uint8_t h = qh[l];
acc1[0] += yl[l+0] * (float)(q1[l] & 0x0F);
acc1[1] += yl[l+8] * (float)(q1[l] & 0xF0);
acc1[2] += yh[l+0] * (float)(q2[l] & 0x0F);
acc1[3] += yh[l+8] * (float)(q2[l] & 0xF0);
acc2[0] += (h & hm1) ? yl[l+0] : 0.f;
acc2[1] += (h & hm2) ? yl[l+8] : 0.f;
acc2[2] += (h & hm3) ? yh[l+0] : 0.f;
acc2[3] += (h & hm4) ? yh[l+8] : 0.f;
}
const float dall = (float)dh[0];
const float dmin = (float)dh[1];
sumf += dall * ((float)sc8[0] * (acc1[0] + 16.f * acc2[0]) +
(float)sc8[1] * (acc1[1] / 16.f + 16.f * acc2[1]) +
(float)sc8[4] * (acc1[2] + 16.f * acc2[2]) +
(float)sc8[5] * (acc1[3] / 16.f + 16.f * acc2[3])) -
dmin * (sumy[0] * (float)sc8[2] + sumy[1] * (float)sc8[3] +
sumy[2] * (float)sc8[6] + sumy[3] * (float)sc8[7]);
y1 += 4 * QK_K;
}
const float tot = simd_sum(sumf);
if (tiisg == 0 && row < (int)p.ne01) {
dst[output_row * p.ne0 + row] = tot;
}
}
// ====================================================================
// Q6_K expert-indexed mat-vec kernel
// ====================================================================
kernel void kernel_mul_mv_id_q6_K_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const uint8_t kmask1 = 0x03;
const uint8_t kmask2 = 0x0C;
const uint8_t kmask3 = 0x30;
const uint8_t kmask4 = 0xC0;
const int nb = p.ne00 / QK_K;
const int64_t r0 = tgpig.x;
const int output_row_base = tgpig.y;
if (output_row_base >= (int)p.ne1) return;
// For Q6_K, each threadgroup handles 2 rows (one per SIMD group).
// But we need to handle the _id dimension: output_row_base is the
// flat output-row index. We process row = 2*r0 + sgitg within the
// weight matrix, and the output goes to output_row_base.
//
// Actually, the Q6_K dispatch geometry is different from Q4_0/Q8_0.
// In the non-id version: threadgroups = (ceil(N/2), M, 1)
// Each threadgroup handles 2 adjacent weight rows (r0*2 + sgitg).
//
// For the _id version: threadgroups = (ceil(N/2), n_tokens*top_k, 1)
// tgpig.y = output_row (flat: token*top_k + slot)
// tgpig.x = weight row pair index
// sgitg selects within the pair: row = 2*r0 + sgitg
const uint token_idx = output_row_base / p.top_k;
const uint expert_id = ids[output_row_base];
const int row = 2 * r0 + sgitg;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row_base, row, 1, tiisg == 0)) return;
device const block_q6_K * x = (device const block_q6_K *)((device const char *)src0 + expert_id * p.expert_stride) + row * nb;
device const float * yy = src1 + token_idx * p.ne10;
float sumf = 0;
const int tid = tiisg / 2;
const int ix = tiisg % 2;
const int ip = tid / 8;
const int il = tid % 8;
const int n = 4;
const int l0 = n * il;
const int is = 8*ip + l0/16;
const int y_offset = 128*ip + l0;
const int q_offset_l = 64*ip + l0;
const int q_offset_h = 32*ip + l0;
for (int i = ix; i < nb; i += 2) {
device const uint8_t * q1 = x[i].ql + q_offset_l;
device const uint8_t * q2 = q1 + 32;
device const uint8_t * qh = x[i].qh + q_offset_h;
device const int8_t * sc = x[i].scales + is;
device const float * y = yy + i * QK_K + y_offset;
const float dall = x[i].d;
float4 sums = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
sums[0] += y[l+ 0] * ((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32);
sums[1] += y[l+32] * ((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32);
sums[2] += y[l+64] * ((int8_t)((q1[l] >> 4) | ((qh[l] & kmask3) << 0)) - 32);
sums[3] += y[l+96] * ((int8_t)((q2[l] >> 4) | ((qh[l] & kmask4) >> 2)) - 32);
}
sumf += dall * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]);
}
const float tot = simd_sum(sumf);
if (tiisg == 0 && row < (int)p.ne01) {
dst[output_row_base * p.ne0 + row] = tot;
}
}
// ====================================================================
// Q6_K _id expert-indexed mat-vec kernel — nr0=2 variant (ADR-028 iter-321)
// ====================================================================
//
// Same as `kernel_mul_mv_id_q6_K_f32` above but processes nr0=2 rows
// per simdgroup with cached yl[16]. 4 rows/TG (2 SGs × 2 rows) vs
// baseline 2 rows/TG. Mirrors peer's
// `template kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q6_K_f32_impl<N_R0_Q6_K>>>`
// (ggml-metal.metal:10351).
//
// Dispatch: threadgroups=(ceil(N/4), n_tokens*top_k, 1), threads=(2, 32, 1).
// Env-gated via HF2Q_Q6K_ID_MV_NR2=1 in dispatch_id_mv.
kernel void kernel_mul_mv_id_q6_K_f32_nr2(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
constexpr int NSG = 2;
constexpr int nr0 = 2;
constexpr uint8_t kmask1 = 0x03;
constexpr uint8_t kmask2 = 0x0C;
constexpr uint8_t kmask3 = 0x30;
constexpr uint8_t kmask4 = 0xC0;
const int nb = p.ne00 / QK_K;
const int64_t r0 = tgpig.x;
const int output_row_base = tgpig.y;
if (output_row_base >= (int)p.ne1) return;
const uint token_idx = output_row_base / p.top_k;
const uint expert_id = ids[output_row_base];
const int first_row = (int)((r0 * NSG + sgitg) * nr0);
if (poison_invalid_expert_id(
dst, p, expert_id, output_row_base, first_row, nr0, tiisg == 0)) return;
device const block_q6_K * x_base = (device const block_q6_K *)((device const char *)src0 + expert_id * p.expert_stride) + first_row * nb;
device const float * yy = src1 + token_idx * p.ne10;
float sumf[nr0] = {0.f, 0.f};
float yl[16];
// ADR-028 iter-402: short indexing matches peer's pattern (same as
// iter-401 applied to the non-_id variant).
const short tid = tiisg / 2;
const short ix = tiisg % 2;
const short ip = tid / 8;
const short il = tid % 8;
const short l0 = 4 * il;
const short is = 8*ip + l0/16;
const short y_offset = 128*ip + l0;
const short q_offset_l = 64*ip + l0;
const short q_offset_h = 32*ip + l0;
for (int i = ix; i < nb; i += 2) {
// ADR-028 iter-352: explicit FOR_UNROLL pragma test FALSIFIED here too;
// see non-_id variant (kernel_mul_mv_q6_K_f32_nr2) for the bench data.
// Apple Metal's auto-unroll is already optimal; explicit hint is removed.
device const float * y = yy + i * QK_K + y_offset;
for (int l = 0; l < 4; ++l) {
yl[4*l + 0] = y[l + 0];
yl[4*l + 1] = y[l + 32];
yl[4*l + 2] = y[l + 64];
yl[4*l + 3] = y[l + 96];
}
for (int row = 0; row < nr0; ++row) {
device const block_q6_K * xr = x_base + row * nb;
device const uint8_t * q1 = xr[i].ql + q_offset_l;
device const uint8_t * q2 = q1 + 32;
device const uint8_t * qh = xr[i].qh + q_offset_h;
device const int8_t * sc = xr[i].scales + is;
const float dall = xr[i].d;
float4 sums = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < 4; ++l) {
sums[0] += yl[4*l + 0] * ((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32);
sums[1] += yl[4*l + 1] * ((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32);
sums[2] += yl[4*l + 2] * ((int8_t)((q1[l] >> 4) | ((qh[l] & kmask3) << 0)) - 32);
sums[3] += yl[4*l + 3] * ((int8_t)((q2[l] >> 4) | ((qh[l] & kmask4) >> 2)) - 32);
}
sumf[row] += dall * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]);
}
}
for (int row = 0; row < nr0; ++row) {
const int out_row = first_row + row;
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && out_row < (int)p.ne01) {
dst[output_row_base * p.ne0 + out_row] = tot;
}
}
}
// ====================================================================
// Q4_K expert-indexed mat-vec kernel
// ====================================================================
//
// ADR-013 P7 — implements `kernel_mul_mv_id_q4_K_f32` (the mv_id
// thunk wrapping `kernel_mul_mv_q4_K_f32_impl`).
//
// Mirrors the Q5_K mv_id kernel above, differing only in:
// 1. Block struct has no qh field (saves 32 bytes per block).
// 2. The acc2 high-bit accumulators collapse to zero — Q4_K stores
// only the low 4 bits, no high-bit array.
// 3. The dot-product reduces to (q1[l] & 0x0F) and (q1[l] & 0xF0) >> 4
// paired with the pre-summed yl/yh/sumy.
//
// Scale-decode is byte-identical to Q5_K: same kmask1/kmask2/kmask3,
// same sc16[] packing — verified against the reference
// kernel_mul_mv_q4_K_f32_impl.
//
// Geometry (mirrors Q5_K mv_id):
// 2 simdgroups per threadgroup, 1 row per simdgroup → 2 rows per tg.
// tgpig.x = weight-row-pair index (row = 2*r0 + sgitg)
// tgpig.y = flat output row (token*top_k + slot)
//
// Routing index dim Y, NOT Z — see the Q5_K kernel's note at line 124.
kernel void kernel_mul_mv_id_q4_K_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nb = p.ne00 / QK_K;
const int64_t r0 = tgpig.x;
const int output_row = tgpig.y; // flat: token*top_k + slot
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
// Each threadgroup covers weight-row pair (2*r0, 2*r0+1);
// sgitg selects which row this simdgroup computes.
const int row = 2 * (int)r0 + (int)sgitg;
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, row, 1, tiisg == 0)) return;
device const block_q4_K * x = (device const block_q4_K *)((device const char *)src0 + expert_id * p.expert_stride) + row * nb;
device const float * yy = src1 + token_idx * p.ne10;
float sumf = 0.f;
const uint16_t kmask1 = 0x3f3f;
const uint16_t kmask2 = 0x0f0f;
const uint16_t kmask3 = 0xc0c0;
const int tid = tiisg / 4;
const int ix = tiisg % 4;
const int iq = tid / 4;
const int ir = tid % 4;
const int n = 8;
const int l0 = n * ir;
const int q_offset = 32 * iq + l0;
const int y_offset = 64 * iq + l0;
uint16_t sc16[4];
thread const uint8_t * sc8 = (thread const uint8_t *)sc16;
device const float * y1 = yy + ix * QK_K + y_offset;
for (int i = ix; i < nb; i += 4) {
device const uint8_t * q1 = x[i].qs + q_offset;
device const uint8_t * q2 = q1 + 64;
device const half * dh = &x[i].d;
// scales array is uint8_t[12]; cast to uint16_t[6] for the
// sc16 decoding identical to the reference candle kernel.
device const uint16_t * a = (device const uint16_t *)x[i].scales + iq;
device const float * y2 = y1 + 128;
float yl[16], yh[16];
float4 sumy = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
yl[l+0] = y1[l + 0]; sumy[0] += yl[l+0];
yl[l+8] = y1[l + 32]; sumy[1] += yl[l+8];
yh[l+0] = y2[l + 0]; sumy[2] += yh[l+0];
yh[l+8] = y2[l + 32]; sumy[3] += yh[l+8];
}
sc16[0] = a[0] & kmask1;
sc16[1] = a[2] & kmask1;
sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2);
sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2);
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
acc1[0] += yl[l+0] * (float)(q1[l] & 0x0F);
acc1[1] += yl[l+8] * (float)(q1[l] & 0xF0);
acc1[2] += yh[l+0] * (float)(q2[l] & 0x0F);
acc1[3] += yh[l+8] * (float)(q2[l] & 0xF0);
}
const float dall = (float)dh[0];
const float dmin = (float)dh[1];
sumf += dall * ((float)sc8[0] * (acc1[0] ) +
(float)sc8[1] * (acc1[1] / 16.f ) +
(float)sc8[4] * (acc1[2] ) +
(float)sc8[5] * (acc1[3] / 16.f )) -
dmin * (sumy[0] * (float)sc8[2] + sumy[1] * (float)sc8[3] +
sumy[2] * (float)sc8[6] + sumy[3] * (float)sc8[7]);
y1 += 4 * QK_K;
}
const float tot = simd_sum(sumf);
if (tiisg == 0 && row < (int)p.ne01) {
dst[output_row * p.ne0 + row] = tot;
}
}
// ====================================================================
// Q4_K expert-indexed mat-vec, two rows per lane group (peer layout)
// ====================================================================
//
// Port of the peer's `kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K = 2>` under its
// `kernel_mul_mv_id` thunk (N_SG_Q4_K = 2). Within a SIMD group, 8
// consecutive lanes share one super-block (`ix = tiisg / 8`); lane
// `it = tiisg % 8` owns quarter `iq = it / 4` (64 values) and the 8-value
// run `ir = it % 4` in each of its 4 sub-blocks. The 32 activations a lane
// reads are loaded once per block for both of its rows; the quants are read
// as uint16 (two nibble pairs per load) and masked in place, with the
// high-byte terms scaled by 1/256 and the high-nibble sub-blocks by 1/16.
//
// Rows shorter than 4 super-blocks (K = 256 or 512) leave lane groups idle
// in the peer's kernel (`ix >= nb`). Here those groups take further row
// pairs instead: with FOLD = 4 / nb (nb = 1 or 2, else FOLD = 1) the four
// 8-lane groups split into FOLD row groups of 4 / FOLD lanes-of-8 each, so
// a SIMD group covers 2 * FOLD rows. Each lane's per-block arithmetic is the
// peer's; a row's partials are reduced with one `simd_sum` in which lanes of
// other row groups contribute +0. At K >= 1024 (FOLD = 1) this is exactly
// the peer's lane layout and reduction.
//
// One threadgroup computes one (pair, 4 * FOLD-column) tile from that
// pair's own expert slice and token row only, so every pair's output is
// independent of how many pairs share the dispatch and of the other pairs'
// ids; FOLD depends on K alone.
//
// Differences from the peer besides FOLD: the pair is routed along grid Y
// (as every mlx-native `_id` kernel), an out-of-range id writes the stock
// NaN poison for this SIMD group's columns, and rows past `ne01` re-read
// the last row (their sums are not written) instead of reading past the
// tail.
//
// Dispatch: threadgroups (ceil(N / (4 * FOLD)), n_tokens * top_k, 1),
// threads (32, 2, 1). Selected only by `quantized_matmul_id_ggml_mv_q4_k_nr2`.
kernel void kernel_mul_mv_id_q4_K_f32_nr2(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]
) {
constexpr short NSG = 2;
constexpr short NR0 = 2;
constexpr uint16_t kmask1 = 0x3f3f;
constexpr uint16_t kmask2 = 0x0f0f;
constexpr uint16_t kmask3 = 0xc0c0;
const int pair = (int)tgpig.y; // flat: token * top_k + slot
if (pair >= (int)p.ne1) return;
const uint token_idx = (uint)pair / p.top_k;
const uint expert_id = ids[pair];
const int nb = p.ne00 / QK_K;
const short fold = (nb == 1) ? 4 : (nb == 2) ? 2 : 1;
const short block_lanes = 4 / fold; // 8-lane groups per row group
const int sg_first_row = ((int)tgpig.x * NSG + (int)sgitg) * NR0 * fold;
if (poison_invalid_expert_id(
dst, p, expert_id, pair, sg_first_row, NR0 * fold, tiisg == 0)) return;
if (sg_first_row >= (int)p.ne01) return;
const short ix = (tiisg/8) % block_lanes; // first block of this lane
const short rg = (tiisg/8) / block_lanes; // row group of this lane
const short it = tiisg%8; // 0...7
const short iq = it/4; // 0 or 1
const short ir = it%4; // 0...3
// This lane's rows, clamped to the last row (block_q4_K is 144 bytes =
// 72 uint16 per block); the step to the second row is zero past the tail.
const int first_row = sg_first_row + rg * NR0;
const int last_row = (int)p.ne01 - 1;
const int row0 = min(first_row, last_row);
const int row_step = (first_row + 1 <= last_row)
? nb * (int)(sizeof(block_q4_K) / sizeof(uint16_t)) : 0;
device const block_q4_K * x = (device const block_q4_K *)(src0 + expert_id * p.expert_stride) + row0 * nb;
device const float * y = src1 + token_idx * p.ne10;
float yl[16];
float yh[16];
float sumf[NR0] = {0.f};
device const float * y4 = y + ix * QK_K + 64 * iq + 8 * ir;
uint16_t sc16[4];
thread const uint8_t * sc8 = (thread const uint8_t *)sc16;
for (int ib = ix; ib < nb; ib += block_lanes) {
float4 sumy = {0.f, 0.f, 0.f, 0.f};
for (short i = 0; i < 8; ++i) {
yl[i+0] = y4[i+ 0]; sumy[0] += yl[i+0];
yl[i+8] = y4[i+ 32]; sumy[1] += yl[i+8];
yh[i+0] = y4[i+128]; sumy[2] += yh[i+0];
yh[i+8] = y4[i+160]; sumy[3] += yh[i+8];
}
device const uint16_t * sc = (device const uint16_t *)x[ib].scales + iq;
device const uint16_t * q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
device const half * dh = &x[ib].d;
for (short row = 0; row < NR0; row++) {
sc16[0] = sc[0] & kmask1;
sc16[1] = sc[2] & kmask1;
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
device const uint16_t * q2 = q1 + 32;
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
_Pragma("clang loop unroll(full)")
for (short i = 0; i < 4; ++i) {
acc1[0] += yl[2*i + 0] * (q1[i] & 0x000F);
acc1[1] += yl[2*i + 1] * (q1[i] & 0x0F00);
acc1[2] += yl[2*i + 8] * (q1[i] & 0x00F0);
acc1[3] += yl[2*i + 9] * (q1[i] & 0xF000);
acc2[0] += yh[2*i + 0] * (q2[i] & 0x000F);
acc2[1] += yh[2*i + 1] * (q2[i] & 0x0F00);
acc2[2] += yh[2*i + 8] * (q2[i] & 0x00F0);
acc2[3] += yh[2*i + 9] * (q2[i] & 0xF000);
}
sumf[row] += dh[0] * ((acc1[0] + 1.f/256.f * acc1[1]) * sc8[0] +
(acc1[2] + 1.f/256.f * acc1[3]) * sc8[1] * 1.f/16.f +
(acc2[0] + 1.f/256.f * acc2[1]) * sc8[4] +
(acc2[2] + 1.f/256.f * acc2[3]) * sc8[5] * 1.f/16.f) -
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]);
q1 += row_step;
sc += row_step;
dh += row_step;
}
y4 += block_lanes * QK_K;
}
// One simd_sum per row of the SIMD group; lanes of other row groups add
// +0. Every lane runs every reduction (fold is uniform).
for (short g = 0; g < fold; ++g) {
for (short row = 0; row < NR0; ++row) {
const float tot = simd_sum(rg == g ? sumf[row] : 0.f);
const int out_row = sg_first_row + g * NR0 + row;
if (tiisg == 0 && out_row < (int)p.ne01) {
dst[pair * p.ne0 + out_row] = tot;
}
}
}
}
// ====================================================================
// IQ2_XXS expert-indexed mat-vec kernel (ADR-066 N-IQ2XXS)
// ====================================================================
//
// Port of the peer's `kernel_mul_mv_iq2_xxs_f32_impl` layout (N_R0 = 4
// rows per simdgroup, N_SG = 2 simdgroups, 8 rows per threadgroup) onto
// mlx-native's expert ABI. The grid and sign tables are staged in
// threadgroup memory (256 * 8 + 128 = 2,176 B) once per threadgroup.
//
// Each lane owns 32-value sub-blocks with stride 32 along K. It sums
// `y * grid * sign` in F32, scales by `d * (0.5 + s)` per sub-block and
// applies the final `* 0.25` after `simd_sum`. One threadgroup reduces each
// output row of one routed row, so a call with m routed rows equals m
// one-row calls bit for bit. The peer's split variant (`ntx = nb32` for
// rows shorter than 32 sub-blocks) is omitted; every lane geometry here is
// the `ntx = 32` one.
//
// Dispatch: threadgroups (ceil(N / 8), n_tokens * top_k, 1), threads
// (32, 2, 1), threadgroup(0) = 2,176 bytes.
#define IQ2_XXS_N_R0 4
#define IQ2_XXS_N_SG 2
kernel void kernel_mul_mv_id_iq2_xxs_f32(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int nr0 = IQ2_XXS_N_R0;
const int nsg = IQ2_XXS_N_SG;
const int nb = p.ne00 / QK_K; // super-blocks per row
const int nb32 = nb * (QK_K / 32); // 32-value sub-blocks per row
const int r0 = tgpig.x;
const int output_row = tgpig.y;
if (output_row >= (int)p.ne1) return;
const uint token_idx = output_row / p.top_k;
const uint expert_id = ids[output_row];
const int first_row = (r0 * nsg + sgitg) * nr0;
// Uniform per threadgroup: every thread returns before the barrier.
if (poison_invalid_expert_id(
dst, p, expert_id, output_row, first_row, nr0, tiisg == 0)) return;
threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256);
{
int pos = (32 * sgitg + tiisg) * 4;
for (int i = 0; i < 4; ++i) svalues[pos + i] = iq2xxs_grid[pos + i];
pos = (32 * sgitg + tiisg) * 2;
for (int i = 0; i < 2; ++i) ssigns[pos + i] = ksigns_iq2xs[pos + i];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
device const block_iq2_xxs * x = (device const block_iq2_xxs *)(src0
+ expert_id * p.expert_stride) + first_row * nb;
device const float * y = src1 + token_idx * p.ne10;
// Row stride in u16 units (66-byte blocks are 2-byte aligned).
const int row_stride = nb * (int)(sizeof(block_iq2_xxs) / 2);
float yl[32];
float sumf[nr0] = {0.f};
const int ix = tiisg;
device const float * y4 = y + 32 * ix;
for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
for (short i = 0; i < 32; ++i) {
yl[i] = y4[i];
}
const int ibl = ib32 / (QK_K / 32);
const int ib = ib32 % (QK_K / 32);
device const block_iq2_xxs * xr = x + ibl;
device const uint16_t * q2 = xr->qs + 4 * ib;
device const half * dh = &xr->d;
for (short row = 0; row < nr0; row++) {
const float db = dh[0];
device const uint8_t * aux8 = (device const uint8_t *)q2;
const uint32_t aux32 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
const float d = db * (0.5f + (aux32 >> 28));
float sum = 0;
for (short l = 0; l < 4; ++l) {
const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + aux8[l]);
const uint8_t signs = ssigns[(aux32 >> 7 * l) & 127];
for (short j = 0; j < 8; ++j) {
sum += yl[8 * l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
}
}
sumf[row] += d * sum;
dh += row_stride;
q2 += row_stride;
}
y4 += 32 * 32;
}
for (int row = 0; row < nr0; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < p.ne01) {
dst[output_row * p.ne0 + first_row + row] = tot * 0.25f;
}
}
}
// ====================================================================
// Expert-grouped Q4_K gate_up with the moe_swiglu_batch GELU epilogue
// (ADR-063 Q4-GEMMA-DECODE F2)
// ====================================================================
//
// Exact (bitwise) fusion of `kernel_mul_mv_id_q4_K_f32` into `[pairs, n]`
// followed by `moe_swiglu_batch`. The grid is the stock one along Y (one
// threadgroup row per flat (token, slot) pair).
//
// Leader rule: pair p with expert e = ids[p] is a leader when e is valid and
// no q < p has ids[q] == e. A non-leader returns uniformly after reading
// ids[0..p]. A leader serves every member (p and each later q with
// ids[q] == e) in column chunks of at most 4: the weight rows of e are read
// once per K block for the whole chunk, and each member runs the stock
// per-pair arithmetic on its own input row (token_idx = q / top_k) and writes
// exactly the output row the stock pipeline writes for q. Invalid ids take
// the stock poison path for their own pair and never join a group.
//
// Per (row, member) expressions are literal clones of the stock kernel:
// the same yl/yh/sumy operand forms, accumulator lanes, scale decode, sumf
// expression and simd_sum. Only the loop nest is interleaved. Rows past the
// gate/up tail are never written.
// True when an earlier pair already owns `expert_id`; that leader writes
// this pair's output.
inline bool grouped_mv_id_has_earlier_member(
device const uint * ids,
int pair,
uint expert_id
) {
for (int q = 0; q < pair; ++q) {
if (ids[q] == expert_id) {
return true;
}
}
return false;
}
// Collect up to 4 members of `expert_id` starting at `*cursor`.
inline int grouped_mv_id_next_members(
device const uint * ids,
constant GgmlMatvecIdParams & p,
uint expert_id,
thread int * cursor,
thread int * members
) {
int count = 0;
int q = *cursor;
for (; q < (int)p.ne1 && count < 4; ++q) {
if (ids[q] == expert_id) {
members[count++] = q;
}
}
*cursor = q;
return count;
}
// ---- Q4_K per-row sums: stock lane partition ----
//
// Rows run one after another; each row's K loop serves the whole member
// chunk. Interleaving rows inside one K loop (sharing each column's
// yl/yh/sumy across rows) changes Metal's default-math contraction and
// drifts 1 ULP from the stock kernel, so rows are never interleaved;
// members (columns) are.
// One weight row against up to 4 members: tot[c] is the stock per-pair
// `simd_sum(sumf)` for member c.
template <int NC>
inline void q4_K_grouped_row_sums(
device const block_q4_K * x,
device const float * src1,
constant GgmlMatvecIdParams & p,
thread const int * members,
uint tiisg,
thread float * tot
) {
const int nb = p.ne00 / QK_K;
float sumf[NC];
for (int c = 0; c < NC; ++c) {
sumf[c] = 0.f;
}
const uint16_t kmask1 = 0x3f3f;
const uint16_t kmask2 = 0x0f0f;
const uint16_t kmask3 = 0xc0c0;
const int tid = tiisg / 4;
const int ix = tiisg % 4;
const int iq = tid / 4;
const int ir = tid % 4;
const int n = 8;
const int l0 = n * ir;
const int q_offset = 32 * iq + l0;
const int y_offset = 64 * iq + l0;
uint16_t sc16[4];
thread const uint8_t * sc8 = (thread const uint8_t *)sc16;
device const float * y1[NC];
for (int c = 0; c < NC; ++c) {
const uint token_idx = (uint)members[c] / p.top_k;
y1[c] = src1 + token_idx * p.ne10 + ix * QK_K + y_offset;
}
for (int i = ix; i < nb; i += 4) {
device const uint8_t * q1 = x[i].qs + q_offset;
device const uint8_t * q2 = q1 + 64;
device const half * dh = &x[i].d;
device const uint16_t * a = (device const uint16_t *)x[i].scales + iq;
for (int c = 0; c < NC; ++c) {
device const float * y2 = y1[c] + 128;
float yl[16], yh[16];
float4 sumy = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
yl[l+0] = y1[c][l + 0]; sumy[0] += yl[l+0];
yl[l+8] = y1[c][l + 32]; sumy[1] += yl[l+8];
yh[l+0] = y2[l + 0]; sumy[2] += yh[l+0];
yh[l+8] = y2[l + 32]; sumy[3] += yh[l+8];
}
sc16[0] = a[0] & kmask1;
sc16[1] = a[2] & kmask1;
sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2);
sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2);
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
for (int l = 0; l < n; ++l) {
acc1[0] += yl[l+0] * (float)(q1[l] & 0x0F);
acc1[1] += yl[l+8] * (float)(q1[l] & 0xF0);
acc1[2] += yh[l+0] * (float)(q2[l] & 0x0F);
acc1[3] += yh[l+8] * (float)(q2[l] & 0xF0);
}
const float dall = (float)dh[0];
const float dmin = (float)dh[1];
sumf[c] += dall * ((float)sc8[0] * (acc1[0] ) +
(float)sc8[1] * (acc1[1] / 16.f ) +
(float)sc8[4] * (acc1[2] ) +
(float)sc8[5] * (acc1[3] / 16.f )) -
dmin * (sumy[0] * (float)sc8[2] + sumy[1] * (float)sc8[3] +
sumy[2] * (float)sc8[6] + sumy[3] * (float)sc8[7]);
y1[c] += 4 * QK_K;
}
}
for (int c = 0; c < NC; ++c) {
tot[c] = simd_sum(sumf[c]);
}
}
// ---- Q4_K gate_up with the moe_swiglu_batch epilogue ----
//
// The weight holds [gate (n_ff rows); up (n_ff rows)] per expert (ne01 =
// 2 * n_ff). A SIMD group owns R0 gate rows r and their up rows r + n_ff,
// computes each row with the per-row loop above (rows sequential, members
// interleaved), and lane 0 applies the `moe_swiglu_batch` expression to the
// two stock row sums, writing dst[pair * ne0 + r] with ne0 = n_ff: exactly
// what `kernel_mul_mv_id_q4_K_f32` followed by `moe_swiglu_batch` produces.
// An invalid id feeds the stock NaN poison through the same epilogue.
// Instantiated at the frozen R0 = 4 only.
// Literal clone of the `moe_swiglu_batch` expression (moe_dispatch.metal).
inline float grouped_geglu_epilogue(float gate, float up) {
float gelu = gate * 0.5f * (1.0f + precise::tanh(
0.7978845608f * (gate + 0.044715f * gate * gate * gate)));
return gelu * up;
}
template <int R0, int NC>
inline void q4_K_geglu_chunk(
device const block_q4_K * xe,
device const float * src1,
device float * dst,
constant GgmlMatvecIdParams & p,
thread const int * members,
int first_row,
uint tiisg
) {
const int nb = p.ne00 / QK_K;
const int n_ff = (int)p.ne01 / 2;
for (int r = 0; r < R0; ++r) {
const int row = first_row + r;
// Uniform per SIMD group; rows past the tail do no work.
if (row >= n_ff) {
break;
}
float gate[NC];
float up[NC];
q4_K_grouped_row_sums<NC>(xe + row * nb, src1, p, members, tiisg, gate);
q4_K_grouped_row_sums<NC>(xe + (row + n_ff) * nb, src1, p, members, tiisg, up);
if (tiisg == 0) {
for (int c = 0; c < NC; ++c) {
dst[members[c] * p.ne0 + row] = grouped_geglu_epilogue(gate[c], up[c]);
}
}
}
}
// Geometry: threads (2, 32, 1) = 2 SIMD groups of R0 gate/up row pairs, so a
// threadgroup covers 2 * R0 output columns. tgpig.y = flat pair.
template <int R0>
kernel void kernel_mul_mv_id_q4_K_f32_grouped_geglu_impl(
device const char * src0 [[buffer(0)]],
device const float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
device const uint * ids [[buffer(3)]],
constant GgmlMatvecIdParams & p [[buffer(4)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint tiisg [[thread_index_in_simdgroup]],
uint sgitg [[simdgroup_index_in_threadgroup]]
) {
const int pair = tgpig.y;
if (pair >= (int)p.ne1) return;
const uint expert_id = ids[pair];
const int n_ff = (int)p.ne01 / 2;
const int first_row = (2 * (int)tgpig.x + (int)sgitg) * R0;
if ((uint64_t)expert_id >= (uint64_t)p.ne02) {
// Stock: the gate and up rows hold the NAN poison and
// moe_swiglu_batch maps them at run time. Default Metal math assumes
// no NaNs, so a compile-time NAN operand would be folded away; the
// poison is rebuilt from p.r2 (always 1 in this ABI) so the epilogue
// runs on the same bits at run time, exactly as stock does.
const float poison = as_type<float>(as_type<uint>(NAN) + (p.r2 - 1u));
if (tiisg == 0) {
for (int r = 0; r < R0; ++r) {
if (first_row + r < n_ff) {
dst[pair * p.ne0 + first_row + r] = grouped_geglu_epilogue(poison, poison);
}
}
}
return;
}
if (grouped_mv_id_has_earlier_member(ids, pair, expert_id)) return;
device const block_q4_K * xe = (device const block_q4_K *)((device const char *)src0 + expert_id * p.expert_stride);
int cursor = pair;
int members[4];
for (;;) {
const int count = grouped_mv_id_next_members(ids, p, expert_id, &cursor, members);
if (count == 4) {
q4_K_geglu_chunk<R0, 4>(xe, src1, dst, p, members, first_row, tiisg);
} else if (count == 3) {
q4_K_geglu_chunk<R0, 3>(xe, src1, dst, p, members, first_row, tiisg);
} else if (count == 2) {
q4_K_geglu_chunk<R0, 2>(xe, src1, dst, p, members, first_row, tiisg);
} else if (count == 1) {
q4_K_geglu_chunk<R0, 1>(xe, src1, dst, p, members, first_row, tiisg);
}
if (count < 4) break;
}
}
template [[host_name("kernel_mul_mv_id_q4_K_f32_grouped_geglu_rt4")]]
kernel void kernel_mul_mv_id_q4_K_f32_grouped_geglu_impl<4>(
device const char *, device const float *, device float *, device const uint *,
constant GgmlMatvecIdParams &, uint3, uint, uint);