mtfyun/ds4
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Metal: optimize Qwen Q4 decode and verify SSD parity

Browse Source 
Reduce Q4_K gate/up memory traffic by sharing Q8 activation reads across
the shared expert's two projections. Extract the existing IQ2 helper for
reuse, and widen routed activation and scale loads with packed float4 and
half2 operations. Preserve the scalar quantized-code loads and the original
per-lane FP32 accumulation and SIMD reduction order.

Enable MXFP4 down prefetch automatically on M1 Max for the measured
640-to-2560 projection. Preserve existing M5/M6 fast-math policies and the
diagnostic disable switch. Track the compiled library's math mode so
strict math keeps the plain noncontracted kernel even when prefetch is
explicitly requested; the prefetched kernel's explicit FMA otherwise
introduces a one-bit difference in that mode.

Strengthen tests with a plain-kernel oracle independent of automatic
prefetch, explicit/default dispatch checks, and Q4 SSD coverage at the
32/33 and 128-token boundaries. Compare addressed and masked Q4 decode
against the generic resident kernel across cache reuse, fallback and retry.
Add MXFP4 down microbenchmarks and document the full-model measurements.
Keep the existing prefill and row-dispatch policies: the experimental
smaller Q2 sparse-prefill tiles did not show a stable Q4 benefit.

Validation on Apple M1 Max 32 GiB:
- make -j2 all and the affected Metal test targets; git diff --check.
- Q4 ordered geometry and MoE decode exactness with MATH_SAFE=0 and 1.
- IQ2 shared-projection regression: 180 fixtures, 1,119,600 comparisons.
- MoE prefill specialization: 132 passing cases per math mode.
- SSD expert regression: 94 passing cases per math mode, including
  missing-only reads, overlap, output guards, failure and exact retry.
- Q4 SSD, ctx=32768, chunk=1024, temp=0, nothink, 100 generated tokens:
  three alternating pairs improve median decode from 6.07 to 6.42 t/s
  (+5.8%), with identical output in every run. Prefill medians are 6.94
  and 6.91 t/s. SSD/system variation remains substantial; these are local
  measurements and do not establish gains on other devices.
- Q4 MTP with two draft tokens: before/after 64-token output identical.
This commit is contained in:
Giorgio Oppo committed 2026-09-26 15:35:37 +02:00
1 parent 81ee4badb8
commit 4642526317
5 files changed
+150 -49
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/QWEN38_FLASH_NEXT.md
+24
View File
@@ -75,6 +75,30 @@ variation between runs does not establish a 10% improvement. Prefill medians
were 9.98 and 10.05 t/s; its execution path is unchanged. These are local
warm-file measurements, not a guarantee for other Macs or cold SSD reads.
Q4_K gate/up shares activation reads between the shared expert's Q8
projections and uses packed activation/scale loads while preserving the
original FP32 accumulation order. On M1 Max, the 640-to-2560 MXFP4 down
projection also enables prefetch by default. `DS4_QWEN4_MOE_DOWN_PREFETCH=0`
restores the plain down kernel for diagnostics. Safe math always retains that
plain kernel, including when prefetch is explicitly requested, because its
noncontracted arithmetic differs from the prefetched kernel's explicit FMA.
A local Q4 SSD check on M1 Max 32 GiB used `--ctx 32768`,
`--prefill-chunk 1024`, `--temp 0 --nothink`, `-n 100`, and the prompt
`Narrami in italiano la storia di Roma`. Three alternating pairs comparing
the previous kernels with these changes produced:
| Decode | Runs (t/s) | Median (t/s) |
| --- | --- | --- |
| Before | 6.02, 6.07, 6.26 | 6.07 |
| Q4 kernel updates | 6.42, 6.03, 7.27 | 6.42 |
All generated output was identical. The observed median gain was 5.8%; SSD
and system variation remained substantial. Prefill medians were 6.94 and
6.91 t/s. These changes retain the existing prefill and row-dispatch policies;
Q2's smaller sparse-prefill tiles did not show a stable benefit for Q4 in
local comparisons. The measurements do not establish gains on other devices.
Leave the expert-cache budget automatic initially. A plain count passed to
`--ssd-streaming-cache-experts` requests dynamic cache slots; an `NGB` budget
also includes the reserved layer-staging space. Startup accounts for the
ds4_metal.m
+10 -5
View File
@@ -889,6 +889,7 @@ static NSUInteger g_moe_q4_up_slots_bytes;
static NSUInteger g_moe_q4_down_slots_bytes;
static NSUInteger g_attn_out_group_ids_bytes;
static int g_initialized;
static bool g_metal_math_safe;
static int g_quality_mode;
static int g_mpp_invalid_env_reported;
#define DS4_METAL_MAX_ROUTED_EXPERT_USED 8
@@ -6909,6 +6910,7 @@ int ds4_gpu_init(void) {
const int drift_kv_raw_f32 = ds4_gpu_env_bool("DS4_METAL_KV_RAW_F32") > 0; // default OFF
const int drift_rope_exp2_log2 = ds4_gpu_env_bool("DS4_METAL_ROPE_EXP2_LOG2") > 0; // default OFF
const int drift_math_safe = ds4_gpu_env_bool("DS4_METAL_MATH_SAFE") > 0; // default OFF
g_metal_math_safe = drift_math_safe;
if (drift_math_safe) {
// MTLCompileOptions.fastMathEnabled defaults to YES and Apple's
@@ -50104,12 +50106,15 @@ int ds4_gpu_qwen4_moe_down_tensor(
const bool q2 = weight_type == 10u && ff_dim == 640u && rows_per_tg == nsg * 2u &&
ds4_gpu_env_bool("DS4_QWEN4_MOE_MV_SPECIALIZE") != 0 &&
ds4_gpu_device_name_contains("M1 Max");
/* MXFP4 rows with four blocks per lane requested ahead (same lane map and
* chain order, byte-identical); M5/M6 default, DS4_QWEN4_MOE_DOWN_PREFETCH=0/1
* overrides on any device. */
/* Request four MXFP4 blocks per lane ahead. Explicit FMA matches the
* plain kernel in fast math; strict math keeps its noncontracted path,
* even with the diagnostic override. Use the compiled library's mode.
* M1 Max uses the measured 640x2560 shape; M5/M6 retain their policy. */
const int prefetch_override = ds4_gpu_env_bool("DS4_QWEN4_MOE_DOWN_PREFETCH");
const bool prefetch = weight_type == 39u && (ff_dim % 32u) == 0 &&
(prefetch_override >= 0 ? prefetch_override > 0 : ds4_gpu_device_is_m5_or_m6_apple_silicon());
const bool prefetch = !g_metal_math_safe && weight_type == 39u && (ff_dim % 32u) == 0 &&
(prefetch_override >= 0 ? prefetch_override > 0 :
ds4_gpu_device_is_m5_or_m6_apple_silicon() ||
(ff_dim == 640u && out_dim == 2560u && ds4_gpu_device_name_contains("M1 Max")));
const uint32_t dispatch_slots = qwen4_moe_dispatch_slots(&args, n_out);
if (!dispatch_slots) return 1;
return qwen4_dispatch(q2 ? QWEN4_K_MOE_DOWN_Q2K : prefetch ? QWEN4_K_MOE_DOWN_MXFP4_PF : QWEN4_K_MOE_DOWN, &args, sizeof(args), b, 5,
metal/qwen4.metal
+42 -33
View File
@@ -2856,6 +2856,30 @@ static inline float qwen4_row_dot(device const char *row, device const float *x,
return simd_sum(acc);
}
/* Two Q8 projections share only activation loads. Keep each qwen4_row_dot
* lane's block order, four-term expression and final SIMD reduction. */
static inline float2 qwen4_q8_pair_dot(device const char *gr, device const char *ur,
device const float *x, uint in_dim, ushort tiisg) {
const short ix = tiisg / 8, it = tiisg % 8;
const uint nb = in_dim / 32u;
float sumg = 0.0f, sumu = 0.0f;
for (uint ib = (uint)ix; ib < nb; ib += 4) {
device const char *bg = gr + (uint64_t)ib * 34;
device const char *bu = ur + (uint64_t)ib * 34;
device const float *y = x + ib * 32 + (uint)it * 2;
const float dg = (float)(*(device const half *)bg);
const float du = (float)(*(device const half *)bu);
device const char *qg = bg + 2 + it * 2;
device const char *qu = bu + 2 + it * 2;
const float y0 = y[0], y1 = y[1], y16 = y[16], y17 = y[17];
sumg += dg * (y0 * (float)qg[0] + y1 * (float)qg[1] +
y16 * (float)qg[16] + y17 * (float)qg[17]);
sumu += du * (y0 * (float)qu[0] + y1 * (float)qu[1] +
y16 * (float)qu[16] + y17 * (float)qu[17]);
}
return float2(simd_sum(sumg), simd_sum(sumu));
}
/* --- small multi-output GEMV (rows via the generic row dot) ------------- */
struct ds4_metal_args_qwen4_gemv {
@@ -2972,32 +2996,10 @@ kernel void kernel_qwen4_moe_mid_iq2(
const uint64_t mid_base = ((uint64_t)tok * n_out + slot) * args.out_rows;
if (slot == args.n_slots) {
if (args.shared_type == 8u) {
// Same Q8 lane mapping, block order, four-term expression and
// final SIMD sums as two qwen4_row_dot calls; share only x loads.
const short ix = tiisg / 8, it = tiisg % 8;
const uint nb = args.in_dim / 32u;
for (uint r = row0; r < row0 + NR && r < args.out_rows; r++) {
const uint64_t off = (uint64_t)r * args.shared_row_bytes;
device const char *gr = sh_gate + off;
device const char *ur = sh_up + off;
float sumg = 0.0f, sumu = 0.0f;
for (uint ib = (uint)ix; ib < nb; ib += 4) {
device const char *bg = gr + (uint64_t)ib * 34;
device const char *bu = ur + (uint64_t)ib * 34;
device const float *y = xt + ib * 32 + (uint)it * 2;
const float dg = (float)(*(device const half *)bg);
const float du = (float)(*(device const half *)bu);
device const char *qg = bg + 2 + it * 2;
device const char *qu = bu + 2 + it * 2;
const float y0 = y[0], y1 = y[1], y16 = y[16], y17 = y[17];
sumg += dg * (y0 * (float)qg[0] + y1 * (float)qg[1] +
y16 * (float)qg[16] + y17 * (float)qg[17]);
sumu += du * (y0 * (float)qu[0] + y1 * (float)qu[1] +
y16 * (float)qu[16] + y17 * (float)qu[17]);
}
const float g = simd_sum(sumg);
const float u = simd_sum(sumu);
if (tiisg == 0) mid[mid_base + r] = qwen4_silu(g) * u;
const float2 gu = qwen4_q8_pair_dot(sh_gate + off, sh_up + off, xt, args.in_dim, tiisg);
if (tiisg == 0) mid[mid_base + r] = qwen4_silu(gu.x) * gu.y;
}
} else {
for (uint r = row0; r < row0 + NR && r < args.out_rows; r++) {
@@ -3098,9 +3100,11 @@ kernel void kernel_qwen4_moe_mid_q4k(
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;
const float2 gu = args.shared_type == 8u ?
qwen4_q8_pair_dot(sh_gate + off, sh_up + off, xt, args.in_dim, tiisg) :
float2(qwen4_row_dot(sh_gate + off, xt, args.shared_type, args.in_dim, tiisg),
qwen4_row_dot(sh_up + off, xt, args.shared_type, args.in_dim, tiisg));
if (tiisg == 0) mid[mid_base + r] = qwen4_silu(gu.x) * gu.y;
}
return;
}
@@ -3120,16 +3124,19 @@ kernel void kernel_qwen4_moe_mid_q4k(
float sumg[NR] = {0.0f}, sumu[NR] = {0.0f};
for (uint ib = 0; ib < nb; ib++) {
device const float *yp = xt + ib * 256 + group * 32 + l;
float y[8];
for (uint i = 0; i < 8; i++) y[i] = yp[i];
/* Widen memory operations while preserving the scalar FP32 chain.
* Packed types keep the original float/half alignment contract. */
const packed_float4 y0 = *(device const packed_float4 *)yp;
const packed_float4 y1 = *(device const packed_float4 *)(yp + 4u);
const float y[8] = {y0.x, y0.y, y0.z, y0.w, y1.x, y1.y, y1.z, y1.w};
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 * 144;
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 dmg = (float)(*(device const half *)(bg + 2));
const float du = (float)(*(device const half *)bu);
const float dmu = (float)(*(device const half *)(bu + 2));
const packed_half2 scales_g = *(device const packed_half2 *)bg;
const packed_half2 scales_u = *(device const packed_half2 *)bu;
const float dg = (float)scales_g.x, dmg = (float)scales_g.y;
const float du = (float)scales_u.x, dmu = (float)scales_u.y;
device const uchar *scg = bg + 4;
device const uchar *scu = bu + 4;
uint sg, mg, su, mu;
@@ -3144,6 +3151,8 @@ kernel void kernel_qwen4_moe_mid_q4k(
}
const float dsg = dg * (float)sg, dmin_g = dmg * (float)mg;
const float dsu = du * (float)su, dmin_u = dmu * (float)mu;
/* Packed code loads change fast-math contraction on M1 Max.
* Keep scalar bytes to match the original dot bit for bit. */
device const uchar *qg = bg + 16 + (group >> 1) * 32 + l;
device const uchar *qu = bu + 16 + (group >> 1) * 32 + l;
for (uint i = 0; i < 8; i++) {
tests/test_qwen4_kernels.c
+23 -6
View File
@@ -1848,7 +1848,11 @@ static void test_moe_types(arena_t *a, uint32_t NE, uint32_t slots, uint32_t E,
unsetenv("DS4_QWEN4_MOE_MV_SPECIALIZE");
unsetenv("DS4_QWEN4_MOE_MV_NR");
unsetenv("DS4_QWEN4_MOE_MV_NSG");
unsetenv("DS4_QWEN4_MOE_DOWN_PREFETCH");
} else {
/* Mode zero remains the original plain-kernel oracle when
* automatic MXFP4 dispatch starts using prefetch. */
setenv("DS4_QWEN4_MOE_DOWN_PREFETCH", "0", 1);
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);
@@ -1869,20 +1873,23 @@ static void test_moe_types(arena_t *a, uint32_t NE, uint32_t slots, uint32_t E,
free(bm); free(bp); free(am); free(ap);
}
if (dtype == 39u && getenv("DS4_TEST_QWEN4_MV_EXACT")) {
/* The prefetched MXFP4 down rows must match the plain kernel byte for
* bit, shared Q8 slot included, at the default and generic geometries. */
/* Fast math: prefetched MXFP4 must match the plain kernel bit for
* bit, including shared Q8. Safe math: even a forced prefetch request
* must retain the noncontracted plain kernel. Check both geometries. */
const uint64_t np = (uint64_t)T * n_out * E;
float *bp = malloc(np * sizeof(float)), *ap = malloc(np * sizeof(float));
require_ok(bp && ap, "down prefetch allocation");
for (uint32_t spec = 0; spec < 2u; spec++) {
setenv("DS4_QWEN4_MOE_MV_SPECIALIZE", spec ? "1" : "0", 1);
for (uint32_t mode = 0; mode < 2u; mode++) {
setenv("DS4_QWEN4_MOE_DOWN_PREFETCH", mode ? "1" : "0", 1);
for (uint32_t mode = 0; mode < 3u; mode++) {
if (mode == 2u) unsetenv("DS4_QWEN4_MOE_DOWN_PREFETCH");
else setenv("DS4_QWEN4_MOE_DOWN_PREFETCH", mode ? "1" : "0", 1);
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size,
down_off, dtype, NE, T, slots, F, E, sd_off, shared_type), "down prefetch dispatch");
require_ok(ds4_gpu_tensor_read(gpart, 0, mode ? ap : bp, np * sizeof(float)), "down prefetch read");
if (mode) check_exact_f32(spec ? "MXFP4 down prefetch/default, specialized" :
"MXFP4 down prefetch/default, generic", ap, bp, np);
}
check_exact_f32(spec ? "prefetched MXFP4 down, specialized" : "prefetched MXFP4 down, generic", ap, bp, np);
}
unsetenv("DS4_QWEN4_MOE_MV_SPECIALIZE");
unsetenv("DS4_QWEN4_MOE_DOWN_PREFETCH");
@@ -3597,7 +3604,7 @@ static double bench_run(const char *name, bench_fn fn, void *ud, uint32_t reps)
typedef struct {
arena_t *a;
uint64_t off[14];
uint64_t off[15];
ds4_gpu_tensor *t[43];
uint32_t n[3];
} bench_ctx;
@@ -3641,6 +3648,11 @@ static int bench_moe_down(void *ud) {
bench_ctx *c = ud;
return ds4_gpu_qwen4_moe_down_tensor(c->t[17], c->t[1], c->t[16], c->a->base, c->a->size, c->off[10], 8u, 16, 1, 10, 640, 2560, c->off[10], 8u);
}
static int bench_moe_down_mxfp4(void *ud) {
bench_ctx *c = ud;
return ds4_gpu_qwen4_moe_down_tensor(c->t[17], c->t[1], c->t[16], c->a->base, c->a->size,
c->off[14], 39u, 16, c->n[0], 10, 640, 2560, c->off[10], 8u);
}
static int bench_router_gemv(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f32_tensor(c->t[17], c->a->base, c->a->size, c->off[11], 2560, 512, c->t[0], 1); }
static int bench_router_topk(void *ud) {
bench_ctx *c = ud;
@@ -3776,6 +3788,7 @@ static void bench_dispatch(arena_t *a) {
c.off[11] = arena_f32(a, 512ull * 2560, &sh, -0.05f, 0.05f); free(sh);
c.off[12] = arena_q4_K(a, 16ull * 640, 2560, &sh, 0.05f); free(sh);
c.off[13] = arena_q4_K(a, 16ull * 640, 2560, &sh, 0.05f); free(sh);
c.off[14] = arena_mxfp4(a, 16ull * 2560, 640, &sh); free(sh);
c.t[18] = upload(NULL, 256ull * 2560); /* x for 256 tokens */
c.t[19] = upload(NULL, 256ull * 10240); /* wide scratch */
c.t[1] = upload(NULL, 256ull * 10240); /* scratch (also the GDN state) */
@@ -3931,6 +3944,10 @@ static void bench_dispatch(arena_t *a) {
bench_run("moe_mid q4_K 10+1 slots T=1", bench_moe_mid_q4k, &c, 100);
c.n[0] = 2;
bench_run("moe_mid q4_K 10+1 slots T=2", bench_moe_mid_q4k, &c, 100);
c.n[0] = 1;
bench_run("moe_down mxfp4 10+1 slots T=1", bench_moe_down_mxfp4, &c, 100);
c.n[0] = 2;
bench_run("moe_down mxfp4 10+1 slots T=2", bench_moe_down_mxfp4, &c, 100);
}
/* four projections of one input, mixed weight types (q8_0, bf16, q4_0, f16) */
tests/test_qwen4_ssd_experts.c
+51 -5
View File
@@ -192,6 +192,33 @@ static void mm_variant(unsigned nt) {
}
}
static const char *q4_mv_env[] = {"DS4_QWEN4_NO_Q4K_MID", "DS4_QWEN4_Q4K_MID_NR",
"DS4_QWEN4_Q4K_MID_NSG", "DS4_QWEN4_MOE_DOWN_PREFETCH"};
static void q4_mv_save(char *saved[4]) {
for (unsigned i = 0; i < 4; i++) {
const char *v = getenv(q4_mv_env[i]);
saved[i] = v ? strdup(v) : NULL;
need(!v || saved[i], "save Q4 row dispatch environment");
}
}
static void q4_mv_variant(int reference) {
/* Pin the oracle independently of device defaults. The candidate forces
* NR1/NSG1 even at this fixture's small width, exercising addressed and
* masked rows as well as the real model's automatic dispatch elsewhere. */
const char *values[] = {reference ? "1" : NULL, reference ? NULL : "1",
reference ? NULL : "1", reference ? "0" : "1"};
for (unsigned i = 0; i < 4; i++)
need((values[i] ? setenv(q4_mv_env[i], values[i], 1) : unsetenv(q4_mv_env[i])) == 0,
"select original/optimized Q4 row dispatch");
}
static void q4_mv_restore(char *saved[4]) {
for (unsigned i = 0; i < 4; i++) {
need((saved[i] ? setenv(q4_mv_env[i], saved[i], 1) : unsetenv(q4_mv_env[i])) == 0,
"restore Q4 row dispatch environment");
free(saved[i]);
}
}
/* Check the real GPU-built lists against the frozen router IDs, without
* assuming the order in which GPU atomics append a token's slot. */
static void check_mm_lists(const int32_t *ids, uint32_t T, unsigned route,
@@ -264,7 +291,15 @@ static void check_case_impl(const weights *w, uint32_t T, unsigned route, int go
int empty_fd, int fail_first, int warm_probe,
uint32_t reuse_seed_count) {
const int mm = T > 8;
const int compare_mm = mm && T <= 128 && w->gate_type == 16 && w->down_type == 10;
const int compare_mm = mm && T <= 128 &&
((w->gate_type == 16 && w->down_type == 10) ||
(w->gate_type == 12 && w->down_type == 39));
const int compare_q4_mv = !mm && w->gate_type == 12 && w->down_type == 39;
char *saved_q4_mv_env[4] = {NULL, NULL, NULL, NULL};
if (compare_q4_mv) {
q4_mv_save(saved_q4_mv_env);
q4_mv_variant(0); /* Failed-read probes also exercise the candidate. */
}
char *saved_mm_env[5] = {NULL, NULL, NULL, NULL, NULL};
if (compare_mm) {
for (unsigned i = 0; i < 5; i++) {
@@ -369,13 +404,14 @@ static void check_case_impl(const weights *w, uint32_t T, unsigned route, int go
need(ds4_gpu_stream_expert_cache_current_count() == reuse_seed_count,
"reuse cache full before selected batch");
}
/* Preserve the independent generic resident reference. Low-bit MM runs
/* Preserve the independent generic resident reference. Q2/Q4 MM runs
* default SSD cold -> old NT4 SSD on these same buffers, followed by
* NT1/grid4 -> default -> NT1/grid4. This covers both automatic tile
* widths and specialization against forced generic kernels. Warm
* cases remove the readable fd; overflow cases must reread only the
* selected union on every arm. TILES alone does not force old NT4. */
if (compare_mm) mm_variant(variants[run]);
if (compare_q4_mv) q4_mv_variant(run == 0);
const int warm = warm_probe && run >= 2;
if (warm) need(ds4_gpu_set_model_fd(empty_fd), "warm-cache empty fd");
for (unsigned i = 0; i < 4; i++) reset(outputs + i);
@@ -425,6 +461,7 @@ static void check_case_impl(const weights *w, uint32_t T, unsigned route, int go
"restore MM dispatch environment");
free(saved_mm_env[i]);
}
if (compare_q4_mv) q4_mv_restore(saved_q4_mv_env);
free(x); free(mix); free(r); free(inj); free(sg); free(ids); free(read_x); free(read_ids);
}
@@ -445,6 +482,12 @@ static void check_split_one(const weights *w, uint32_t T, uint32_t slots, int sh
need(slots <= MAX_SLOTS && slots >= 4 &&
(mm ? T <= 16384 && slots == S && !shared && !duplicates && (seed_count == 0 || seed_count == 12) :
T >= 1 && T <= 3 && seed_count <= 4), "split fixture shape");
const int compare_q4_mv = !mm && w->gate_type == 12 && w->down_type == 39;
char *saved_q4_mv_env[4] = {NULL, NULL, NULL, NULL};
if (compare_q4_mv) {
q4_mv_save(saved_q4_mv_env);
q4_mv_variant(1);
}
const uint32_t unique = mm ? 24u : slots - (duplicates ? 3u : 0u), stride = slots + shared;
const uint32_t budget = unique > BUDGET ? unique : BUDGET;
int32_t *ids = malloc((size_t)T * slots * sizeof(*ids)), seeds[24], missing[24];
@@ -515,6 +558,7 @@ static void check_split_one(const weights *w, uint32_t T, uint32_t slots, int sh
shared ? inputs[3].view : NULL, NULL, outputs[3].view, inputs[4].view,
T, slots, stride, D, HC) && ds4_gpu_end_commands(), "split reference reduce");
for (unsigned i = 0; i < 4; i++) reference[i] = read_output(outputs + i);
if (compare_q4_mv) q4_mv_variant(0);
ds4_gpu_set_streaming_expert_cache_budget(budget);
ds4_gpu_set_streaming_expert_cache_expert_bytes(2 * w->table.gate_expert_bytes + w->table.down_expert_bytes);
@@ -606,6 +650,7 @@ static void check_split_one(const weights *w, uint32_t T, uint32_t slots, int sh
free(got); free(routing_copy[i]); free_output(routing + i);
}
for (unsigned i = 0; i < 4; i++) { free(reference[i]); free_output(outputs + i); }
if (compare_q4_mv) q4_mv_restore(saved_q4_mv_env);
free(ids); free(x); free(mix); free(r); free(inj); free(shared_gate);
printf("PASS Qwen SSD split T%u %s %u/%u slots=%u shared=%d seeded=%u unique=%u retry=%d: exact all stages, missing-only %llu bytes, warm0\n",
T, mm ? "MM" : "rows", w->gate_type, w->down_type, slots, shared, seed_count, unique, fail_first, (unsigned long long)first.bytes);
@@ -651,7 +696,7 @@ int main(void) {
ds4_gpu_set_quality(false); ds4_gpu_set_ssd_streaming(true);
need(ds4_gpu_set_model_map(map, bytes) && ds4_gpu_set_model_fd(fileno(file)), "model registration");
const uint32_t shapes[][2] = {{0,1}, {0,2}, {0,3}, {0,9}, {0,128},
{1,1}, {1,17}, {2,2}, {2,9}};
{1,1}, {1,17}, {1,128}, {2,2}, {2,9}};
for (unsigned c = 0; c < sizeof(shapes) / sizeof(*shapes); c++) {
const weights *w = formats + shapes[c][0]; const uint32_t T = shapes[c][1];
ds4_gpu_set_streaming_expert_cache_budget(BUDGET);
@@ -664,8 +709,9 @@ int main(void) {
/* Eight experts appear in all 29 rows, sixteen in only 3/4 rows:
* both >8-token counts and partial tiles, with 488 inactive. */
check_case(w, 29, 4, fileno(file), fileno(empty), 0, 0);
if (w->gate_type == 16) {
/* Both sides of the measured M1 SSD batch-size boundary. */
if (w->gate_type == 16 || w->gate_type == 12) {
/* Q2 and Q4 straddle the M1 SSD batch-size boundary with
* more selected experts than fit in the cache. */
check_case(w, 32, 3, fileno(file), fileno(empty), 0, 0);
check_case(w, 33, 3, fileno(file), fileno(empty), 0, 0);
}