Metal: reuse Qwen IQ2 and HC inputs on M1 Max
Reduce repeated activation work in Qwen decode without changing the model's quantization or per-row arithmetic. Select the new paths automatically on M1 Max and retain the existing diagnostic overrides for direct comparison. Add an IQ2_XXS gate/up kernel with two rows per SIMD group. Load each lane's eight activation values once per quantization block and reuse them across both rows and projections. Preserve the original partial sums, block order, SIMD reduction, shared-expert fallback and output layout. Read 66-byte block headers through aligned ushort loads. Carry the existing SSD address-table function constant and compact slot mapping through the new dispatch, with guards for odd rows and invalid expert IDs. Add an F16 hyper-connection mixer that stores each low-rank input and its sigmoid in threadgroup memory for reuse across streams and output rows. Keep the two values separate to preserve (x * weight) * sigmoid(x) rounding; pin reassociation and contraction at the original arithmetic boundaries. Allocate 2560 bytes of scratch for rank 320 and place the barrier before partial-group exits. Keep the paired two-token mixer and other-device defaults unchanged. No public API or production environment flag is added. Extend numerical tests with an explicit generic oracle versus automatic dispatch, retaining all previous NR/NSG geometry checks. Cover T1/T2/T3, F640/F641, shared and unshared experts, E33/E2560 HC outputs and prefix/suffix guards. Add a focused 96 MiB test/benchmark arena using packed IQ2 weights without dense shadows. Time seven alternating pairs on identical buffers, after 20 warmups per arm, with 100 IQ2 or 1000 HC calls per measured batch. Timings include CPU encoding and the final GPU completion wait. Validation on Apple M1 Max, 32 GiB: - Built all application targets and relevant Metal tests. - Passed the complete Qwen kernel suite and MoE exact comparisons. - Passed Qwen SSD fixtures in default, unretained-command-buffer and Q4_K NR1 configurations, plus existing Metal SSD expert regressions. - Passed real Qwen3.8 Q2 n-gram failure/recovery tests for prefill, decode and MTP. Two paired model runs preserved all logits bit for bit across 410 decode steps, with independent session state and alternating order. Resident microbenchmarks measured IQ2 gate/up at 173.78 -> 163.91 us for T1 (+6.0% throughput), with similar gains at T2/T3. The HC mixer measured 36.57 -> 17.64 us for T1 and 101.22 -> 41.28 us for T3. Full-model gains are smaller and variable. Two paired runs with a shared 7708-expert SSD cache measured +2.51% and -0.08%, or +1.20% pooled over 400 measured steps after warmup. Independent CLI runs did not establish a stable end-to-end speedup; these results do not claim a 10% gain, a prefill improvement, or a benefit on untested devices. Temporary logs, model files and intermediate benchmark harnesses are excluded from this commit.
This commit is contained in:
1 parent
d6cbc77ca6
commit
cdfc0d51e1
3 files changed
+353
-12
No files matched your search
+21
-4
@@ -48141,6 +48141,7 @@ enum {
|
||||
QWEN4_K_HC_NORM_REUSE_Q8,
|
||||
QWEN4_K_HC_GATE_MIX_F16,
|
||||
QWEN4_K_HC_GATE_MIX_F16_PF,
|
||||
QWEN4_K_HC_GATE_MIX_F16_REUSE,
|
||||
QWEN4_K_HC_GATE_MIX_F32,
|
||||
QWEN4_K_HC_GATE_MIX_Q8,
|
||||
QWEN4_K_HC_GATE_MIX_PAIR_F16,
|
||||
@@ -48175,6 +48176,7 @@ enum {
|
||||
QWEN4_K_ATTN_MERGE_WIDE_NPT4,
|
||||
QWEN4_K_ATTN_MM,
|
||||
QWEN4_K_MOE_MID,
|
||||
QWEN4_K_MOE_MID_IQ2,
|
||||
QWEN4_K_MOE_MID_Q4K,
|
||||
QWEN4_K_MOE_MID_Q4K_NR1,
|
||||
QWEN4_K_MOE_DOWN,
|
||||
@@ -48228,6 +48230,7 @@ static const char *const qwen4_kernel_names[QWEN4_K_COUNT] = {
|
||||
"kernel_qwen4_hc_norm_reuse_q8",
|
||||
"kernel_qwen4_hc_gate_mix_f16",
|
||||
"kernel_qwen4_hc_gate_mix_f16_pf",
|
||||
"kernel_qwen4_hc_gate_mix_f16_reuse",
|
||||
"kernel_qwen4_hc_gate_mix_f32",
|
||||
"kernel_qwen4_hc_gate_mix_q8",
|
||||
"kernel_qwen4_hc_gate_mix_pair_f16",
|
||||
@@ -48262,6 +48265,7 @@ static const char *const qwen4_kernel_names[QWEN4_K_COUNT] = {
|
||||
"kernel_qwen4_attn_merge_wide_npt4",
|
||||
"kernel_qwen4_attn_mm",
|
||||
"kernel_qwen4_moe_mid",
|
||||
"kernel_qwen4_moe_mid_iq2",
|
||||
"kernel_qwen4_moe_mid_q4k",
|
||||
"kernel_qwen4_moe_mid_q4k_nr1",
|
||||
"kernel_qwen4_moe_down",
|
||||
@@ -48365,10 +48369,11 @@ static int qwen4_dispatch(int kernel, const void *args, size_t args_len,
|
||||
@autoreleasepool {
|
||||
const bool addresses = g_qwen4_stream_weights && g_qwen4_stream_weights->addresses;
|
||||
const bool q4_mid = kernel == QWEN4_K_MOE_MID_Q4K || kernel == QWEN4_K_MOE_MID_Q4K_NR1;
|
||||
const bool iq2_mid = kernel == QWEN4_K_MOE_MID_IQ2;
|
||||
id<MTLComputePipelineState> pipeline = nil;
|
||||
if (kernel == QWEN4_K_MOE_MID || kernel == QWEN4_K_MOE_DOWN || kernel == QWEN4_K_MOE_DOWN_MXFP4_PF || q4_mid) {
|
||||
if (kernel == QWEN4_K_MOE_MID || kernel == QWEN4_K_MOE_DOWN || kernel == QWEN4_K_MOE_DOWN_MXFP4_PF || q4_mid || iq2_mid) {
|
||||
const qwen4_moe_args *a = args;
|
||||
const bool specialize = !q4_mid && qwen4_moe_mv_specialize(a->weight_type);
|
||||
const bool specialize = !q4_mid && !iq2_mid && qwen4_moe_mv_specialize(a->weight_type);
|
||||
const uint32_t values[] = {a->weight_type, a->shared_type, specialize ? a->in_dim : 0u, specialize ? qwen4_moe_mv_rows() : 0u};
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_shared=%u_dim=%u_rows=%u_addr=%u",
|
||||
qwen4_kernel_names[kernel], values[0], values[1], values[2], values[3], addresses];
|
||||
@@ -48603,11 +48608,17 @@ int ds4_gpu_qwen4_hc_gate_mix_tensor(
|
||||
/* Register-prefetched F16 rows (same lane order and rounding, pinned
|
||||
* against the plain kernel by tests/test_qwen4_kernels.c); M5 default. */
|
||||
const int prefetch_override = ds4_gpu_env_bool("DS4_QWEN4_HC_MIX_PREFETCH");
|
||||
/* M1 Max shares each low-rank activation and sigmoid across the four
|
||||
* output rows in a threadgroup. The existing override retains both
|
||||
* original kernels for numerical and timing comparisons. */
|
||||
const bool reuse = !pair && weight_type == 1u && n_rank == 320u &&
|
||||
prefetch_override < 0 && ds4_gpu_device_name_contains("M1 Max");
|
||||
const bool prefetch = weight_type == 1u &&
|
||||
(prefetch_override >= 0 ? prefetch_override > 0 : ds4_gpu_device_is_m5_apple_silicon());
|
||||
const int kernel = pair ? (prefetch ? QWEN4_K_HC_GATE_MIX_PAIR_F16_PF
|
||||
: qwen4_hc_kernel(weight_type, QWEN4_K_HC_GATE_MIX_PAIR_F16,
|
||||
QWEN4_K_HC_GATE_MIX_PAIR_F32, QWEN4_K_HC_GATE_MIX_PAIR_Q8))
|
||||
: reuse ? QWEN4_K_HC_GATE_MIX_F16_REUSE
|
||||
: prefetch ? QWEN4_K_HC_GATE_MIX_F16_PF
|
||||
: qwen4_hc_kernel(weight_type, QWEN4_K_HC_GATE_MIX_F16, QWEN4_K_HC_GATE_MIX_F32,
|
||||
QWEN4_K_HC_GATE_MIX_Q8);
|
||||
@@ -48619,7 +48630,7 @@ int ds4_gpu_qwen4_hc_gate_mix_tensor(
|
||||
(uint32_t)ds4_gpu_env_u64("DS4_QWEN4_HC_PAIR_NSG", default_nsg, 1u, 16u) : 4u;
|
||||
return qwen4_dispatch(kernel, &args, sizeof(args), b, 4,
|
||||
MTLSizeMake((n_embd + nsg - 1u) / nsg, pair ? 1u : n_tokens, 1), MTLSizeMake(nsg * 32u, 1, 1),
|
||||
pair ? (NSUInteger)n_rank * 2u * sizeof(float) : 0u);
|
||||
(pair || reuse) ? (NSUInteger)n_rank * 2u * sizeof(float) : 0u);
|
||||
}
|
||||
|
||||
int ds4_gpu_qwen4_hc_combine_tensor(
|
||||
@@ -49194,7 +49205,13 @@ int ds4_gpu_qwen4_moe_mid_tensor(
|
||||
(uint32_t)ds4_gpu_env_u64("DS4_QWEN4_Q4K_MID_NSG", default_nsg, 1u, 8u) :
|
||||
(specialize ? qwen4_moe_mv_groups(weight_type) : 4u);
|
||||
const uint32_t rows_per_tg = nr * nsg;
|
||||
const int kernel = !q4k ? QWEN4_K_MOE_MID : nr == 1u ? QWEN4_K_MOE_MID_Q4K_NR1 : QWEN4_K_MOE_MID_Q4K;
|
||||
/* IQ2 gate and up reuse the same eight input values across both rows.
|
||||
* Preserve the NR2 mapping, including the existing SSD slot masks. */
|
||||
const bool iq2 = weight_type == 16u && nr == 2u &&
|
||||
ds4_gpu_env_bool("DS4_QWEN4_MOE_MV_SPECIALIZE") != 0 &&
|
||||
ds4_gpu_device_name_contains("M1 Max");
|
||||
const int kernel = iq2 ? QWEN4_K_MOE_MID_IQ2 : !q4k ? QWEN4_K_MOE_MID :
|
||||
nr == 1u ? QWEN4_K_MOE_MID_Q4K_NR1 : QWEN4_K_MOE_MID_Q4K;
|
||||
/* Masked dispatches compact only the grid. The kernel restores original
|
||||
* slot indices, so shared placement and output strides stay unchanged. */
|
||||
const uint32_t dispatch_slots = args.slot_mask ?
|
||||
|
||||
@@ -319,6 +319,68 @@ kernel void kernel_qwen4_hc_gate_mix_f16_pf(
|
||||
if (tiisg == 0) mixed[(uint64_t)tok * E + d] = g / (float)hc;
|
||||
}
|
||||
|
||||
/* Reuse the low-rank sigmoid across streams and output rows. Keep x and
|
||||
* sigmoid(x) separate: storing silu(x) would change the single-row
|
||||
* accumulation's (x * weight) * sigmoid(x) rounding. */
|
||||
kernel void kernel_qwen4_hc_gate_mix_f16_reuse(
|
||||
constant ds4_metal_args_qwen4_hc_gate_mix & args,
|
||||
device const float *xn,
|
||||
device const float *lo,
|
||||
device const char *w_up,
|
||||
device float *mixed,
|
||||
threadgroup float2 *low_sig [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]]) {
|
||||
const uint hc = args.n_hc, E = args.n_embd, rank = args.n_rank;
|
||||
const uint d = tgpig.x * (ntg.x / 32) + sgitg;
|
||||
const uint tok = tgpig.y;
|
||||
if (tok >= args.n_tokens) return;
|
||||
device const float *l = lo + (uint64_t)tok * rank;
|
||||
const float inv_hc = 1.0f / (float)hc;
|
||||
for (uint r = tid; r < rank; r += ntg.x) {
|
||||
#pragma clang fp reassociate(off)
|
||||
#pragma clang fp contract(off)
|
||||
const float x = l[r] * inv_hc;
|
||||
low_sig[r] = float2(x, qwen4_sigmoid(x));
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (d >= E) return;
|
||||
const uint s = tiisg / 8, lane = tiisg % 8;
|
||||
device const half *wr = (device const half *)w_up + (uint64_t)(s * E + d) * rank;
|
||||
float acc = 0.0f;
|
||||
uint r = lane;
|
||||
for (; r + 56u < rank; r += 64u) {
|
||||
float wv[8];
|
||||
float2 av[8];
|
||||
for (uint i = 0; i < 8; i++) { wv[i] = (float)wr[r + 8u * i]; av[i] = low_sig[r + 8u * i]; }
|
||||
for (uint i = 0; i < 8; i++) {
|
||||
#pragma clang fp reassociate(off)
|
||||
#pragma clang fp contract(off)
|
||||
const float t = av[i].x * wv[i];
|
||||
const float u = t * av[i].y;
|
||||
acc = acc + u;
|
||||
}
|
||||
}
|
||||
for (; r < rank; r += 8u) {
|
||||
#pragma clang fp reassociate(off)
|
||||
#pragma clang fp contract(off)
|
||||
const float2 a = low_sig[r];
|
||||
const float t = a.x * (float)wr[r];
|
||||
const float u = t * a.y;
|
||||
acc = acc + u;
|
||||
}
|
||||
acc += simd_shuffle_xor(acc, 1);
|
||||
acc += simd_shuffle_xor(acc, 2);
|
||||
acc += simd_shuffle_xor(acc, 4);
|
||||
float g = qwen4_sigmoid(acc) * xn[(uint64_t)tok * E * hc + s * E + d];
|
||||
g += simd_shuffle_xor(g, 8);
|
||||
g += simd_shuffle_xor(g, 16);
|
||||
if (tiisg == 0) mixed[(uint64_t)tok * E + d] = g / (float)hc;
|
||||
}
|
||||
|
||||
/* Two-token MTP verification: reuse each up-projection weight for both
|
||||
* rows, and activate the low-rank inputs once per threadgroup. */
|
||||
template <typename W>
|
||||
@@ -2339,6 +2401,90 @@ kernel void kernel_qwen4_moe_mid(
|
||||
}
|
||||
}
|
||||
|
||||
/* IQ2XXS gate/up input reuse, two output rows per SIMD group. Each lane
|
||||
* keeps the generic dot's eight-value part and ordered dl * part additions;
|
||||
* only independent rows and projections share the activation loads. */
|
||||
kernel void kernel_qwen4_moe_mid_iq2(
|
||||
constant ds4_metal_args_qwen4_moe & args,
|
||||
device const char *gate_base,
|
||||
device const char *up_base,
|
||||
device const int32_t *selected,
|
||||
device const float *x,
|
||||
device float *mid,
|
||||
device const char *sh_gate,
|
||||
device const char *sh_up,
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
constexpr uint NR = 2;
|
||||
const uint slot = qwen4_moe_mid_slot(args, tgpig.y), tok = tgpig.z;
|
||||
const uint n_out = args.n_slots + args.has_shared;
|
||||
const uint row0 = (tgpig.x * (ntg.x / 32u) + (uint)sgitg) * NR;
|
||||
if (row0 >= args.out_rows || slot >= n_out || tok >= args.n_tokens) return;
|
||||
device const float *xt = x + (uint64_t)tok * args.in_dim;
|
||||
const uint64_t mid_base = ((uint64_t)tok * n_out + slot) * args.out_rows;
|
||||
if (slot == args.n_slots) {
|
||||
for (uint r = row0; r < row0 + NR && r < args.out_rows; r++) {
|
||||
const uint64_t off = (uint64_t)r * args.shared_row_bytes;
|
||||
const float g = qwen4_row_dot(sh_gate + off, xt, args.shared_type, args.in_dim, tiisg);
|
||||
const float u = qwen4_row_dot(sh_up + off, xt, args.shared_type, args.in_dim, tiisg);
|
||||
if (tiisg == 0) mid[mid_base + r] = qwen4_silu(g) * u;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const int32_t expert = selected[(uint64_t)tok * args.n_slots + slot];
|
||||
if (expert < 0 || (uint)expert >= args.n_total_expert) {
|
||||
if (tiisg == 0) {
|
||||
for (uint r = row0; r < row0 + NR && r < args.out_rows; r++) mid[mid_base + r] = 0.0f;
|
||||
}
|
||||
return;
|
||||
}
|
||||
device const char *gb = qwen4_expert_base(gate_base, (uint)expert, args.expert_bytes);
|
||||
device const char *ub = qwen4_expert_base(up_base, (uint)expert, args.expert_bytes);
|
||||
const uint nb = args.in_dim / 256;
|
||||
const uint ib32 = tiisg / 4, j = tiisg % 4;
|
||||
float sumg[NR] = {0.0f}, sumu[NR] = {0.0f};
|
||||
for (uint ib = 0; ib < nb; ib++) {
|
||||
device const float *yp = xt + ib * 256 + ib32 * 32 + j * 8;
|
||||
float y[8];
|
||||
for (uint i = 0; i < 8; i++) y[i] = yp[i];
|
||||
for (uint r = 0; r < NR && row0 + r < args.out_rows; r++) {
|
||||
const uint64_t off = (uint64_t)(row0 + r) * args.row_bytes + (uint64_t)ib * 66;
|
||||
device const uchar *bg = (device const uchar *)(gb + off);
|
||||
device const uchar *bu = (device const uchar *)(ub + off);
|
||||
const float dg = (float)(*(device const half *)bg);
|
||||
const float du = (float)(*(device const half *)bu);
|
||||
/* IQ2XXS blocks are only two-byte aligned. */
|
||||
device const ushort *qg = (device const ushort *)(bg + 2) + 4 * ib32;
|
||||
device const ushort *qu = (device const ushort *)(bu + 2) + 4 * ib32;
|
||||
const uint ag = (uint)qg[0] | ((uint)qg[1] << 16);
|
||||
const uint sg = (uint)qg[2] | ((uint)qg[3] << 16);
|
||||
const uint au = (uint)qu[0] | ((uint)qu[1] << 16);
|
||||
const uint su = (uint)qu[2] | ((uint)qu[3] << 16);
|
||||
const float dlg = dg * (0.5f + (float)(sg >> 28)) * 0.25f;
|
||||
const float dlu = du * (0.5f + (float)(su >> 28)) * 0.25f;
|
||||
constant const uchar *gridg = (constant const uchar *)(ds4_metal_iq2xxs_grid + ((ag >> (8 * j)) & 0xFFu));
|
||||
constant const uchar *gridu = (constant const uchar *)(ds4_metal_iq2xxs_grid + ((au >> (8 * j)) & 0xFFu));
|
||||
const uint signsg = ds4_metal_ksigns_iq2xs[(sg >> (7 * j)) & 127u];
|
||||
const uint signsu = ds4_metal_ksigns_iq2xs[(su >> (7 * j)) & 127u];
|
||||
float partg = 0.0f, partu = 0.0f;
|
||||
for (uint i = 0; i < 8; i++) {
|
||||
partg += (float)gridg[i] * ((signsg >> i) & 1u ? -y[i] : y[i]);
|
||||
partu += (float)gridu[i] * ((signsu >> i) & 1u ? -y[i] : y[i]);
|
||||
}
|
||||
sumg[r] += dlg * partg;
|
||||
sumu[r] += dlu * partu;
|
||||
}
|
||||
}
|
||||
for (uint r = 0; r < NR && row0 + r < args.out_rows; r++) {
|
||||
const float g = simd_sum(sumg[r]);
|
||||
const float u = simd_sum(sumu[r]);
|
||||
if (tiisg == 0) mid[mid_base + row0 + r] = qwen4_silu(g) * u;
|
||||
}
|
||||
}
|
||||
|
||||
/* Q4_K gate/up input reuse with the original qwen4_row_dot lane mapping
|
||||
* and accumulation order. Each lane still visits every block in order and
|
||||
* adds its eight elements individually; only independent rows/projections
|
||||
|
||||
+186
-8
@@ -414,7 +414,9 @@ static uint64_t arena_iq2_xxs(arena_t *a, uint64_t rows, uint64_t cols, double *
|
||||
const uint64_t blocks = cols / 256;
|
||||
const uint64_t off = arena_alloc(a, rows * blocks * 66u);
|
||||
uint8_t *w = a->base + off;
|
||||
*shadow = malloc(rows * cols * sizeof(double));
|
||||
/* Exact GPU-to-GPU fixtures need packed weights, without a large dense
|
||||
* shadow that would dominate their bounded memory footprint. */
|
||||
if (shadow) *shadow = malloc(rows * cols * sizeof(double));
|
||||
for (uint64_t r = 0; r < rows; r++) {
|
||||
for (uint64_t b = 0; b < blocks; b++) {
|
||||
uint8_t *blk = w + (r * blocks + b) * 66u;
|
||||
@@ -431,7 +433,7 @@ static uint64_t arena_iq2_xxs(arena_t *a, uint64_t rows, uint64_t cols, double *
|
||||
uint16_t q2[4] = { (uint16_t)aux_g, (uint16_t)(aux_g >> 16), (uint16_t)aux_s, (uint16_t)(aux_s >> 16) };
|
||||
memcpy(blk + 2 + ib32 * 8, q2, 8);
|
||||
const float dl = dq * (0.5f + (float)(aux_s >> 28)) * 0.25f;
|
||||
for (int j = 0; j < 4; j++) {
|
||||
for (int j = 0; shadow && j < 4; j++) {
|
||||
const uint8_t *grid = (const uint8_t *)(iq2xxs_grid + ((aux_g >> (8 * j)) & 0xFF));
|
||||
const uint32_t signs = ksigns_iq2xs[(aux_s >> (7 * j)) & 127];
|
||||
for (int i = 0; i < 8; i++) {
|
||||
@@ -441,7 +443,7 @@ static uint64_t arena_iq2_xxs(arena_t *a, uint64_t rows, uint64_t cols, double *
|
||||
}
|
||||
}
|
||||
}
|
||||
dump_quant("iq2_xxs", w, rows * blocks * 66u, *shadow, rows * cols);
|
||||
if (shadow) dump_quant("iq2_xxs", w, rows * blocks * 66u, *shadow, rows * cols);
|
||||
return off;
|
||||
}
|
||||
|
||||
@@ -1452,12 +1454,20 @@ static void test_moe_types(arena_t *a, uint32_t NE, uint32_t slots, uint32_t E,
|
||||
float *bm = malloc(nm * sizeof(float)), *bp = malloc(np * sizeof(float));
|
||||
float *am = malloc(nm * sizeof(float)), *ap = malloc(np * sizeof(float));
|
||||
require_ok(bm && bp && am && ap, "MoE exact allocation");
|
||||
for (uint32_t mode = 0; mode < 9; mode++) {
|
||||
for (uint32_t mode = 0; mode < 10; mode++) {
|
||||
const char *nr[] = {"2", "1", "1", "2", "2", "4", "4", "1", "2"};
|
||||
const char *nsg[] = {"4", "4", "8", "4", "8", "4", "8", "16", "16"};
|
||||
setenv("DS4_QWEN4_MOE_MV_SPECIALIZE", mode ? "1" : "0", 1);
|
||||
setenv("DS4_QWEN4_MOE_MV_NR", nr[mode], 1);
|
||||
setenv("DS4_QWEN4_MOE_MV_NSG", nsg[mode], 1);
|
||||
if (mode == 9u) {
|
||||
/* Automatic dispatch must compare against explicit generic
|
||||
* mode zero, even when a new kernel becomes the default. */
|
||||
unsetenv("DS4_QWEN4_MOE_MV_SPECIALIZE");
|
||||
unsetenv("DS4_QWEN4_MOE_MV_NR");
|
||||
unsetenv("DS4_QWEN4_MOE_MV_NSG");
|
||||
} else {
|
||||
setenv("DS4_QWEN4_MOE_MV_SPECIALIZE", mode ? "1" : "0", 1);
|
||||
setenv("DS4_QWEN4_MOE_MV_NR", nr[mode], 1);
|
||||
setenv("DS4_QWEN4_MOE_MV_NSG", nsg[mode], 1);
|
||||
}
|
||||
require_ok(ds4_gpu_qwen4_moe_mid_tensor(gmid, gx, gsel, a->base, a->size,
|
||||
gate_off, up_off, wtype, NE, T, slots, E, F, sg_off, su_off, shared_type), "exact mid");
|
||||
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size,
|
||||
@@ -2576,6 +2586,161 @@ static double hc_norm_median7(const double *samples) {
|
||||
return sorted[3];
|
||||
}
|
||||
|
||||
typedef struct {
|
||||
arena_t *arena;
|
||||
bool hc;
|
||||
uint32_t E, F, T, NE, slots, shared_type;
|
||||
uint64_t gate, up, shared_gate, shared_up;
|
||||
ds4_gpu_tensor *x, *low, *selected, *out;
|
||||
} m1_reuse_case;
|
||||
|
||||
static int m1_reuse_dispatch(const m1_reuse_case *c) {
|
||||
arena_t *a = c->arena;
|
||||
if (c->hc) return ds4_gpu_qwen4_hc_gate_mix_tensor(c->out, c->x, c->low,
|
||||
a->base, a->size, c->up, 1u, c->T, c->E, 4u, c->F);
|
||||
return ds4_gpu_qwen4_moe_mid_tensor(c->out, c->x, c->selected,
|
||||
a->base, a->size, c->gate, c->up, 16u, c->NE, c->T, c->slots,
|
||||
c->E, c->F, c->shared_gate, c->shared_up, c->shared_type);
|
||||
}
|
||||
|
||||
static void m1_reuse_select(bool hc, uint32_t mode) {
|
||||
const char *name = hc ? "DS4_QWEN4_HC_MIX_PREFETCH" : "DS4_QWEN4_MOE_MV_SPECIALIZE";
|
||||
require_ok(mode == 1u ? unsetenv(name) == 0 : setenv(name, mode ? "1" : "0", 1) == 0,
|
||||
"select generic, automatic or prefetched kernel");
|
||||
}
|
||||
|
||||
static double m1_reuse_batch(const m1_reuse_case *c, uint32_t mode, uint32_t calls) {
|
||||
m1_reuse_select(c->hc, mode);
|
||||
const double start = bench_now();
|
||||
require_ok(ds4_gpu_begin_commands(), "reuse benchmark begin");
|
||||
for (uint32_t i = 0; i < calls; i++) require_ok(m1_reuse_dispatch(c), "reuse benchmark dispatch");
|
||||
require_ok(ds4_gpu_end_commands(), "reuse benchmark end and GPU wait");
|
||||
return 1e6 * (bench_now() - start) / calls;
|
||||
}
|
||||
|
||||
/* Both arms use the same physical buffers. Inputs are immutable and output is
|
||||
* overwritten, so repeating a dispatch does not change the next one's work. */
|
||||
static void m1_reuse_compare(m1_reuse_case *c, bool benchmark) {
|
||||
const uint64_t n = c->hc ? (uint64_t)c->T * c->E :
|
||||
(uint64_t)c->T * (c->slots + (c->shared_type != UINT32_MAX)) * c->F;
|
||||
const uint64_t guard = 17u, total = n + 2u * guard;
|
||||
const float sentinel = 127.25f;
|
||||
ds4_gpu_tensor *storage = upload(NULL, total);
|
||||
c->out = ds4_gpu_tensor_view(storage, guard * sizeof(float), n * sizeof(float));
|
||||
require_ok(c->out != NULL, "reuse guarded output view");
|
||||
float *reference = NULL;
|
||||
char label[128];
|
||||
snprintf(label, sizeof(label), "%s E=%u %s=%u T=%u shared=%u",
|
||||
c->hc ? "HC F16 reuse" : "IQ2 mid reuse", c->E, c->hc ? "rank" : "F", c->F, c->T,
|
||||
!c->hc && c->shared_type != UINT32_MAX);
|
||||
for (uint32_t mode = 0; mode < (c->hc ? 3u : 2u); mode++) {
|
||||
m1_reuse_select(c->hc, mode);
|
||||
require_ok(ds4_gpu_tensor_fill_f32(storage, sentinel, total) &&
|
||||
ds4_gpu_begin_commands() && m1_reuse_dispatch(c) && ds4_gpu_end_commands(),
|
||||
"reuse exact dispatch and wait");
|
||||
float *got = download(storage, total);
|
||||
for (uint64_t i = 0; i < guard; i++)
|
||||
require_ok(got[i] == sentinel && got[guard + n + i] == sentinel, "reuse prefix and suffix guards");
|
||||
if (!mode) {
|
||||
reference = got;
|
||||
check_exact_f32(label, reference, reference, total);
|
||||
uint64_t nonzero = 0;
|
||||
for (uint64_t i = guard; i < guard + n; i++) nonzero += got[i] != 0.0f && got[i] != sentinel;
|
||||
require_ok(nonzero > n / 2u, "reuse reference has nontrivial output");
|
||||
} else {
|
||||
check_exact_f32(label, got, reference, total);
|
||||
free(got);
|
||||
}
|
||||
}
|
||||
printf(" %s: generic vs unset default%s exact, guards intact\n", label,
|
||||
c->hc ? " vs explicit prefetch" : "");
|
||||
if (benchmark) {
|
||||
/* Fixed small warmup, then alternating old/new order. Report every
|
||||
* observation: medians alone can hide device clock drift. HC needs
|
||||
* longer measured batches because each dispatch is much shorter. */
|
||||
for (uint32_t mode = 0; mode < 2u; mode++) (void)m1_reuse_batch(c, mode, 20u);
|
||||
const uint32_t calls = c->hc ? 1000u : 100u;
|
||||
double samples[2][7];
|
||||
for (uint32_t rep = 0; rep < 7u; rep++) for (uint32_t order = 0; order < 2u; order++) {
|
||||
const uint32_t mode = (rep + order) & 1u;
|
||||
samples[mode][rep] = m1_reuse_batch(c, mode, calls);
|
||||
}
|
||||
const double old = hc_norm_median7(samples[0]), automatic = hc_norm_median7(samples[1]);
|
||||
printf(" %s timing calls=%u generic=%.3f us automatic=%.3f us speedup=%+.2f%%\n",
|
||||
label, calls, old, automatic, 100.0 * (old / automatic - 1.0));
|
||||
for (uint32_t mode = 0; mode < 2u; mode++) {
|
||||
printf(" %s samples_us:", mode ? "automatic" : "generic");
|
||||
for (uint32_t rep = 0; rep < 7u; rep++) printf(" %.3f", samples[mode][rep]);
|
||||
printf("\n");
|
||||
}
|
||||
/* Repeated writes must still produce the oracle, including guards. */
|
||||
float *got = download(storage, total);
|
||||
check_exact_f32("reuse benchmark final output", got, reference, total);
|
||||
free(got);
|
||||
}
|
||||
free(reference);
|
||||
ds4_gpu_tensor_free(c->out);
|
||||
ds4_gpu_tensor_free(storage);
|
||||
c->out = NULL;
|
||||
}
|
||||
|
||||
static void test_m1_reuse(arena_t *a, bool benchmark, const char *only) {
|
||||
const char *env[] = {"DS4_QWEN4_MOE_MV_SPECIALIZE", "DS4_QWEN4_MOE_MV_NR",
|
||||
"DS4_QWEN4_MOE_MV_NSG", "DS4_QWEN4_HC_MIX_PREFETCH", "DS4_QWEN4_HC_PAIR_NSG"};
|
||||
char *saved[5];
|
||||
for (uint32_t i = 0; i < 5u; i++) {
|
||||
const char *value = getenv(env[i]);
|
||||
saved[i] = value ? strdup(value) : NULL;
|
||||
require_ok(!value || saved[i] != NULL, "save reuse environment");
|
||||
require_ok(unsetenv(env[i]) == 0, "clear reuse geometry override");
|
||||
}
|
||||
printf("M1 reuse screen: explicit generic vs automatic (device-dependent) kernels%s\n",
|
||||
benchmark ? "; timing includes CPU encoding and final GPU wait" : "");
|
||||
if (!only || strcmp(only, "hc") != 0) {
|
||||
const uint32_t rows[] = {640u, 641u};
|
||||
for (uint32_t shape = 0; shape < 2u; shape++) {
|
||||
m1_reuse_case c = { .arena = a, .E = 2560u, .F = rows[shape], .NE = 12u, .slots = 10u };
|
||||
c.gate = arena_iq2_xxs(a, (uint64_t)c.NE * c.F, c.E, NULL, 0.05f);
|
||||
c.up = arena_iq2_xxs(a, (uint64_t)c.NE * c.F, c.E, NULL, 0.05f);
|
||||
double *shadow;
|
||||
c.shared_gate = arena_q8_0(a, c.F, c.E, &shadow, 0.05f); free(shadow);
|
||||
c.shared_up = arena_q8_0(a, c.F, c.E, &shadow, 0.05f); free(shadow);
|
||||
float *x = rand_vec(3u * c.E, 1.0f);
|
||||
for (uint32_t i = 0; i < 3u * c.E; i++) x[i] *= i % 7u == 0u ? 8.0f : i % 7u == 1u ? 0.125f : 1.0f;
|
||||
c.x = upload(x, 3u * c.E); free(x);
|
||||
int32_t selected[30];
|
||||
for (uint32_t t = 0; t < 3u; t++) for (uint32_t s = 0; s < c.slots; s++)
|
||||
selected[t * c.slots + s] = (int32_t)((t * 7u + s * 5u) % c.NE);
|
||||
c.selected = ds4_gpu_tensor_alloc(sizeof(selected));
|
||||
require_ok(c.selected && ds4_gpu_tensor_write(c.selected, 0, selected, sizeof(selected)), "reuse selected IDs");
|
||||
for (c.T = 1u; c.T <= 3u; c.T++) for (uint32_t shared = 0; shared < 2u; shared++) {
|
||||
c.shared_type = shared ? 8u : UINT32_MAX;
|
||||
m1_reuse_compare(&c, benchmark && shape == 0u && shared != 0u);
|
||||
}
|
||||
ds4_gpu_tensor_free(c.selected); ds4_gpu_tensor_free(c.x);
|
||||
}
|
||||
}
|
||||
if (!only || strcmp(only, "iq2") != 0) {
|
||||
const uint32_t widths[] = {33u, 2560u};
|
||||
for (uint32_t shape = 0; shape < 2u; shape++) {
|
||||
m1_reuse_case c = { .arena = a, .hc = true, .E = widths[shape], .F = 320u };
|
||||
double *shadow;
|
||||
c.up = arena_f16(a, (uint64_t)4u * c.E * c.F, &shadow, 0.3f); free(shadow);
|
||||
float *x = rand_vec(3u * 4u * c.E, 1.0f), *low = rand_vec(3u * c.F, 6.0f);
|
||||
c.x = upload(x, 3u * 4u * c.E); c.low = upload(low, 3u * c.F);
|
||||
free(x); free(low);
|
||||
for (c.T = 1u; c.T <= 3u; c.T++)
|
||||
m1_reuse_compare(&c, benchmark && shape == 1u && c.T != 2u);
|
||||
ds4_gpu_tensor_free(c.low); ds4_gpu_tensor_free(c.x);
|
||||
}
|
||||
}
|
||||
for (uint32_t i = 0; i < 5u; i++) {
|
||||
require_ok(saved[i] ? setenv(env[i], saved[i], 1) == 0 : unsetenv(env[i]) == 0, "restore reuse environment");
|
||||
free(saved[i]);
|
||||
}
|
||||
printf("M1 reuse screen passed; arena used %.2f MiB\n", (double)a->used / (1024.0 * 1024.0));
|
||||
}
|
||||
|
||||
/* Optional scheduling screen; excluded from the default correctness path.
|
||||
* Each type reuses its weight matrices at every batch size. Both variants get
|
||||
* warmups, then seven paired observations with alternating execution order. */
|
||||
@@ -3041,8 +3206,14 @@ static void test_dense_mm(arena_t *a, uint32_t in_dim, uint32_t rows, uint32_t T
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
/* A focused path avoids the full suite's large arena and dispatch bench.
|
||||
* ONLY accepts 1 (both), iq2 or hc; BENCH=1 adds paired resident timings. */
|
||||
const char *reuse_only = getenv("DS4_TEST_QWEN4_M1_REUSE_ONLY");
|
||||
const char *reuse_bench = getenv("DS4_TEST_QWEN4_M1_REUSE_BENCH");
|
||||
const bool run_reuse = reuse_only && reuse_only[0] && strcmp(reuse_only, "0") != 0;
|
||||
const bool bench_reuse = reuse_bench && reuse_bench[0] && strcmp(reuse_bench, "0") != 0;
|
||||
arena_t arena;
|
||||
arena.size = (uint64_t)1536 << 20;
|
||||
arena.size = (uint64_t)(run_reuse || bench_reuse ? 96u : 1536u) << 20;
|
||||
arena.base = mmap(NULL, arena.size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
|
||||
arena.used = 0;
|
||||
if (arena.base == MAP_FAILED) { perror("mmap"); return 1; }
|
||||
@@ -3050,6 +3221,12 @@ int main(void) {
|
||||
require_ok(ds4_gpu_init(), "GPU initialization");
|
||||
require_ok(ds4_gpu_set_model_map(arena.base, arena.size), "model map registration");
|
||||
|
||||
if (run_reuse || bench_reuse) {
|
||||
test_m1_reuse(&arena, bench_reuse, run_reuse ? reuse_only : NULL);
|
||||
ds4_gpu_cleanup();
|
||||
munmap(arena.base, arena.size);
|
||||
return 0;
|
||||
}
|
||||
if (getenv("DS4_TEST_QWEN4_DECODE_FUSIONS")) { test_decode_fusions(&arena); return 0; }
|
||||
if (getenv("DS4_TEST_QWEN4_MV_EXACT")) {
|
||||
test_moe_types(&arena, 8, 6, 2560, 640, 1, 16u, 10u);
|
||||
@@ -3094,6 +3271,7 @@ int main(void) {
|
||||
test_hc_pair_groups(&arena);
|
||||
test_mv_ext_groups(&arena);
|
||||
test_hc_mix_prefetch(&arena);
|
||||
test_m1_reuse(&arena, false, NULL);
|
||||
test_hc(&arena, 2560, 320, 3, 1u);
|
||||
test_hc(&arena, 2560, 320, 2, 1u);
|
||||
test_hc(&arena, 2560, 320, 2, 0u);
|
||||
|
||||
Reference in new issue
Block a user