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

perf(metal): reuse half operands for large Qwen SSD MoE prefill

Browse Source 
Add half-storage input and intermediate variants of the existing
simdgroup MoE gate/up/activation and down kernels. Preserve FP32 matrix
accumulation, activation arithmetic and the public float intermediate;
move the existing float-to-half operand conversions to a shared input
copy and a half shadow written by the mid epilogue.

Enable the path automatically only for M1 Max SSD streaming with
IQ2_XXS gate/up, Q2_K down and chunks of at least 8192 tokens. Keep the
original path for smaller chunks: a 2048-token experiment did not show
a reliable full-model benefit despite a faster isolated kernel pair.
No production opt-in flag is added.

Scope the half shadow to a single streaming invocation and validate the
producer tensors and dimensions before consuming it. Invalidate state
on entry/failure, separate resident and missing passes, retain scratch
buffers across submitted commands and reconvert reused input contents.
Allocate both optional buffers before encoding and fall back to float
storage on allocation failure. Account explicitly for scratch storage:
140 MiB at the tested 8192-token model shape, already covered by the
existing conservative workspace estimate.

Add a standalone frozen-source oracle/benchmark with generic weight
function constants, NT1/2/4/8, padded Q2_K rows, sparse routing, guarded
buffers and poisoned outputs. Check both FP32 outputs and the actual
half payloads against CPU binary16 conversion. Add Make targets and
extend SSD tests across the 8191/8192 policy boundary, scratch growth to
16384 and reuse, split resident/missing passes, failed reads and retries.

Validation on Apple M1 Max with Qwen3.8-Flash-Next-Q2:
- Build, full Metal kernel suite, CPU memory tests and SSD tests passed.
- Standalone oracle: 31 fixtures, 1015326720 float comparisons and
  100325376 half comparisons, all exact against the independent reference.
- Generic MoE pair at 8192 tokens: 322.620 -> 294.798 ms (+9.44%
  throughput), including conversion and dual writes; nine ABBA/BAAB rounds.
- Complete model, SSD cache 1024 experts, context 16384, MTP off,
  8192-token prompt plus 128 decode evaluations, four runs per arm:
  prefill 181.993 -> 198.252 t/s (+8.93% aggregate; +5.84% comparing
  per-run medians). All four prefill pairs improve; a slow baseline run
  raises the aggregate. A stable 10% overall gain is not demonstrated.
- Decode at chunk 8192 is inconclusive (+0.34% aggregate).
- All timed/validation replays preserve complete vocabulary logits bit
  for bit. Final Metal-debug replays separately validate the retained
  2048 fallback and 8192 optimized path, each comparing 32033280 logits.

The recorded 2048 optimized-path experiment is not enabled by this
commit. Only the automatic threshold/comment changed after measurement;
the retained 8192 execution path is unchanged. Other Apple chips keep
the previous dispatch and are unbenchmarked; CUDA/ROCm are not modified.
This commit is contained in:
GiorgioOppo committed 2026-09-19 11:04:08 +02:00
1 parent 9bff1ca4e1
commit 11811b96b2
7 files changed
+655 -58
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitignore
+1
View File
@@ -97,3 +97,4 @@ __pycache__/
/tests/test_metal_qwen4_moe_down
/tests/test_metal_q8_gemv_reference
/tests/test_metal_qwen4_moe_half
Makefile
+11
View File
@@ -213,6 +213,16 @@ test-metal-q8-gemv-reference: tests/test_metal_q8_gemv_reference
check-metal-q8-gemv-reference: tests/test_metal_q8_gemv_reference
./tests/test_metal_q8_gemv_reference --compile-only
tests/test_metal_qwen4_moe_half: tests/test_metal_qwen4_moe_half.m
$(CC) $(filter-out -ffast-math,$(OBJCFLAGS)) -o $@ $< $(METAL_LDLIBS)
.PHONY: test-metal-qwen4-moe-half check-metal-qwen4-moe-half
test-metal-qwen4-moe-half: tests/test_metal_qwen4_moe_half
./tests/test_metal_qwen4_moe_half
check-metal-qwen4-moe-half: tests/test_metal_qwen4_moe_half
./tests/test_metal_qwen4_moe_half --compile-only
tests/test_metal_qwen4_hc: tests/test_metal_qwen4_hc.m
$(CC) $(filter-out -ffast-math,$(OBJCFLAGS)) -o $@ $< $(METAL_LDLIBS)
@@ -1204,6 +1214,7 @@ clean:
rm -f speed-bench/session_concurrency_bench
rm -f tests/test_metal_q8_reduction
rm -f tests/test_metal_q8_gemv_reference
rm -f tests/test_metal_qwen4_moe_half
rm -f tests/test_metal_qwen4_hc
rm -f tests/test_metal_qwen4_moe_mid
rm -f tests/test_metal_qwen4_moe_down
ds4.c
+5 -1
View File
@@ -1100,7 +1100,11 @@ static ds4_context_memory qwen4_graph_memory_estimate(uint32_t ctx, uint32_t pre
(uint64_t)(mtp ? 12u : 4u) * DS4_N_VOCAB;
if (mtp) fixed += 12u * (hc + 1u) * E + 4u * hc_dim + 1u +
2u * (((uint64_t)DS4_N_VOCAB + 4095u) / 4096u);
m.scratch_bytes = (T * ((gpu_row > old_row ? gpu_row : old_row) + hc_dim + 4u) + fixed) * 4u;
/* Reserve optional half copies of x and the routed MoE intermediate.
* Units here are float words; the legacy conservative estimate already
* covers these copies for the supported Qwen shapes. */
const uint64_t gpu_reserved_row = gpu_row + (E + used * F + 1u) / 2u;
m.scratch_bytes = (T * ((gpu_reserved_row > old_row ? gpu_reserved_row : old_row) + hc_dim + 4u) + fixed) * 4u;
m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes;
return m;
}
ds4_metal.m
+92 -13
View File
@@ -48328,6 +48328,14 @@ static bool qwen4_bind_weight(qwen4_bind *b, const void *map, uint64_t size,
return b->buf != nil;
}
/* A half shadow is valid only inside one SSD wrapper invocation and for the
* exact routed pass that produced it. Public float tensors remain readable
* and writable without any hidden cross-call cache association. */
typedef struct {
const ds4_gpu_tensor *mid, *lists, *counts;
uint32_t tokens, slots, stride, width, list_cap;
} qwen4_moe_half_pass;
/* Scoped to one routed mid/down call. Indirect cache buffers are registered
* on each encoder and protected by the same completion sequence as DeepSeek.
* Compact staging supplies private address tables backed by owned payloads;
@@ -48345,6 +48353,7 @@ typedef struct {
/* CPU-only immutable pass counts. Never bound or dereferenced by a GPU;
* used only to copy active expert IDs into each encoder's setBytes args. */
const int32_t *mm_pass_counts;
qwen4_moe_half_pass *half_pass;
} qwen4_stream_weights;
static qwen4_stream_weights *g_qwen4_stream_weights;
static bool qwen4_moe_mm_m1_ssd(uint32_t n_tokens, uint32_t type);
@@ -48434,6 +48443,14 @@ enum {
QWEN4_K_MOE_MM_DOWN_NT2,
QWEN4_K_MOE_MM_DOWN_NT8,
QWEN4_K_MOE_MM_MID_NT8,
QWEN4_K_MOE_MM_MID_F16_NT1,
QWEN4_K_MOE_MM_MID_F16_NT2,
QWEN4_K_MOE_MM_MID_F16_NT4,
QWEN4_K_MOE_MM_MID_F16_NT8,
QWEN4_K_MOE_MM_DOWN_F16_NT1,
QWEN4_K_MOE_MM_DOWN_F16_NT2,
QWEN4_K_MOE_MM_DOWN_F16_NT4,
QWEN4_K_MOE_MM_DOWN_F16_NT8,
QWEN4_K_MOE_MM_MID_NAX,
QWEN4_K_MOE_MM_DOWN_NAX,
QWEN4_K_MOE_MM_MID_NAX64,
@@ -48544,6 +48561,14 @@ static const char *const qwen4_kernel_names[QWEN4_K_COUNT] = {
"kernel_qwen4_moe_mm_down_nt2",
"kernel_qwen4_moe_mm_down_nt8",
"kernel_qwen4_moe_mm_mid_nt8",
"kernel_qwen4_moe_mm_mid_f16_nt1",
"kernel_qwen4_moe_mm_mid_f16_nt2",
"kernel_qwen4_moe_mm_mid_f16_nt4",
"kernel_qwen4_moe_mm_mid_f16_nt8",
"kernel_qwen4_moe_mm_down_f16_nt1",
"kernel_qwen4_moe_mm_down_f16_nt2",
"kernel_qwen4_moe_mm_down_f16_nt4",
"kernel_qwen4_moe_mm_down_f16_nt8",
"kernel_qwen4_moe_mm_mid_nax",
"kernel_qwen4_moe_mm_down_nax",
"kernel_qwen4_moe_mm_mid_nax64",
@@ -50318,7 +50343,15 @@ static ds4_gpu_tensor *g_qwen4_nax_half_x, *g_qwen4_nax_half_mid, *g_qwen4_nax_h
static uint64_t g_qwen4_nax_half_x_bytes, g_qwen4_nax_half_mid_bytes, g_qwen4_nax_half_midr_bytes;
static const ds4_gpu_tensor *g_qwen4_nax_half_mid_for;
static uint64_t g_qwen4_nax_half_mid_count;
/* Separate storage from tensor-op scratch: the M1 simdgroup path preserves
* its original accumulation and only shares already-half-rounded operands. */
static ds4_gpu_tensor *g_qwen4_half_x, *g_qwen4_half_mid;
static uint64_t g_qwen4_half_x_bytes, g_qwen4_half_mid_bytes;
static void qwen4_nax_release_scratch(void) {
ds4_gpu_tensor_free(g_qwen4_half_x);
ds4_gpu_tensor_free(g_qwen4_half_mid);
g_qwen4_half_x = g_qwen4_half_mid = NULL;
g_qwen4_half_x_bytes = g_qwen4_half_mid_bytes = 0;
if (g_qwen4_nax_half_x) ds4_gpu_tensor_free(g_qwen4_nax_half_x);
if (g_qwen4_nax_half_mid) ds4_gpu_tensor_free(g_qwen4_nax_half_mid);
if (g_qwen4_nax_half_midr) ds4_gpu_tensor_free(g_qwen4_nax_half_midr);
@@ -50365,11 +50398,13 @@ int ds4_gpu_qwen4_moe_mm_mid_tensor(
const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset,
uint32_t weight_type, uint32_t n_expert, uint32_t n_tokens, uint32_t n_slots, uint32_t n_out,
uint32_t in_dim, uint32_t ff_dim, uint32_t list_cap) {
qwen4_moe_half_pass *hp = g_qwen4_stream_weights ? g_qwen4_stream_weights->half_pass : NULL;
if (hp) hp->mid = NULL;
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, in_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * ff_dim;
const uint32_t tiles = qwen4_moe_mm_tiles(n_tokens, weight_type, true);
const uint32_t nt = qwen4_moe_mm_nt(n_tokens, weight_type, "DS4_QWEN4_MOE_MID_NT");
const int kernel = nt == 1u ? QWEN4_K_MOE_MM_MID_NT1 :
int kernel = nt == 1u ? QWEN4_K_MOE_MM_MID_NT1 :
nt == 2u ? QWEN4_K_MOE_MM_MID_NT2 :
nt == 8u ? QWEN4_K_MOE_MM_MID_NT8 : QWEN4_K_MOE_MM_MID;
qwen4_moe_mm_args args = { n_tokens, n_slots, n_out, in_dim, ff_dim, weight_type, row_bytes, list_cap,
@@ -50391,6 +50426,8 @@ int ds4_gpu_qwen4_moe_mm_mid_tensor(
const uint32_t dispatch_experts = qwen4_moe_mm_dispatch_experts(&args);
if (dispatch_experts == UINT32_MAX) return 0;
if (!dispatch_experts) {
if (g_qwen4_stream_weights && g_qwen4_stream_weights->half_pass)
g_qwen4_stream_weights->half_pass->mid = NULL;
/* No mid rows were produced or consumed. Do not let an unrelated
* earlier NAX half-copy association survive this empty pass. */
g_qwen4_nax_half_mid_for = NULL;
@@ -50439,15 +50476,35 @@ int ds4_gpu_qwen4_moe_mm_mid_tensor(
qwen4_moe_mm_grid((ff_dim + 63u) / 64u, dispatch_experts, 1, args.expert_major),
MTLSizeMake(128, 1, 1), nax_fx || nax_comp ? 12288u : 10240u);
}
if (!qwen4_dispatch(kernel, &args, sizeof(args), b, 6,
const uint64_t half_x_bytes = (uint64_t)n_tokens * in_dim * sizeof(uint16_t);
const uint64_t half_mid_bytes = (uint64_t)n_tokens * n_out * ff_dim * sizeof(uint16_t);
/* Allocate both operands before encoding anything. If optional scratch
* does not fit, the original float path can still serve this pass. */
const bool half = hp &&
qwen4_nax_scratch(&g_qwen4_half_x, &g_qwen4_half_x_bytes, half_x_bytes) &&
qwen4_nax_scratch(&g_qwen4_half_mid, &g_qwen4_half_mid_bytes, half_mid_bytes);
if (half) {
[g_transient_buffers addObject:ds4_gpu_tensor_buffer(g_qwen4_half_x)];
[g_transient_buffers addObject:ds4_gpu_tensor_buffer(g_qwen4_half_mid)];
/* Reconvert on every producer call: x is a reused graph buffer and
* its address alone cannot establish that its contents are current. */
ds4_gpu_tensor *xh = qwen4_nax_half_operand(&g_qwen4_half_x, &g_qwen4_half_x_bytes,
x, (uint64_t)n_tokens * in_dim);
if (!xh || !qwen4_bind_tensor(&b[4], xh, half_x_bytes, "moe input half") ||
!qwen4_bind_tensor(&b[6], g_qwen4_half_mid, half_mid_bytes, "moe mid half shadow")) return 0;
kernel = nt == 1u ? QWEN4_K_MOE_MM_MID_F16_NT1 : nt == 2u ? QWEN4_K_MOE_MM_MID_F16_NT2 :
nt == 8u ? QWEN4_K_MOE_MM_MID_F16_NT8 : QWEN4_K_MOE_MM_MID_F16_NT4;
}
if (!qwen4_dispatch(kernel, &args, sizeof(args), b, half ? 7 : 6,
qwen4_moe_mm_grid((ff_dim + 31u) / 32u, dispatch_experts, tiles, args.expert_major),
MTLSizeMake(128, 1, 1), 0)) return 0;
if (tails) {
const MTLSize tail_grid = MTLSizeMake((ff_dim + 31u) / 32u, dispatch_experts, 1);
if (!qwen4_dispatch(QWEN4_K_MOE_MM_MID_NT1, &args, sizeof(args), b, 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 2u && !qwen4_dispatch(QWEN4_K_MOE_MM_MID_NT2, &args, sizeof(args), b, 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 4u && !qwen4_dispatch(QWEN4_K_MOE_MM_MID, &args, sizeof(args), b, 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (!qwen4_dispatch(half ? QWEN4_K_MOE_MM_MID_F16_NT1 : QWEN4_K_MOE_MM_MID_NT1, &args, sizeof(args), b, half ? 7 : 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 2u && !qwen4_dispatch(half ? QWEN4_K_MOE_MM_MID_F16_NT2 : QWEN4_K_MOE_MM_MID_NT2, &args, sizeof(args), b, half ? 7 : 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 4u && !qwen4_dispatch(half ? QWEN4_K_MOE_MM_MID_F16_NT4 : QWEN4_K_MOE_MM_MID, &args, sizeof(args), b, half ? 7 : 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
}
if (half) *hp = (qwen4_moe_half_pass){mid, lists, counts, n_tokens, n_slots, n_out, ff_dim, list_cap};
return 1;
}
@@ -50456,12 +50513,15 @@ int ds4_gpu_qwen4_moe_mm_down_tensor(
const void *model_map, uint64_t model_size, uint64_t down_offset,
uint32_t weight_type, uint32_t n_expert, uint32_t n_tokens, uint32_t n_slots, uint32_t n_out,
uint32_t ff_dim, uint32_t out_dim, uint32_t list_cap) {
qwen4_moe_half_pass *hp = g_qwen4_stream_weights ? g_qwen4_stream_weights->half_pass : NULL;
const qwen4_moe_half_pass producer = hp ? *hp : (qwen4_moe_half_pass){0};
if (hp) hp->mid = NULL; /* consume even if validation or dispatch fails */
const uint32_t weight_dim = weight_type == 10u ? (ff_dim + 255u) / 256u * 256u : ff_dim;
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, weight_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * out_dim;
const uint32_t tiles = qwen4_moe_mm_tiles(n_tokens, weight_type, false);
const uint32_t nt = qwen4_moe_mm_nt(n_tokens, weight_type, "DS4_QWEN4_MOE_DOWN_NT");
const int kernel = nt == 1u ? QWEN4_K_MOE_MM_DOWN_NT1 :
int kernel = nt == 1u ? QWEN4_K_MOE_MM_DOWN_NT1 :
nt == 2u ? QWEN4_K_MOE_MM_DOWN_NT2 :
nt == 8u ? QWEN4_K_MOE_MM_DOWN_NT8 : QWEN4_K_MOE_MM_DOWN;
qwen4_moe_mm_args args = { n_tokens, n_slots, n_out, ff_dim, out_dim, weight_type, row_bytes, list_cap,
@@ -50482,6 +50542,8 @@ int ds4_gpu_qwen4_moe_mm_down_tensor(
const uint32_t dispatch_experts = qwen4_moe_mm_dispatch_experts(&args);
if (dispatch_experts == UINT32_MAX) return 0;
if (!dispatch_experts) {
if (g_qwen4_stream_weights && g_qwen4_stream_weights->half_pass)
g_qwen4_stream_weights->half_pass->mid = NULL;
/* No mid rows were produced or consumed. Do not let an unrelated
* earlier NAX half-copy association survive this empty pass. */
g_qwen4_nax_half_mid_for = NULL;
@@ -50532,14 +50594,24 @@ int ds4_gpu_qwen4_moe_mm_down_tensor(
qwen4_moe_mm_grid((out_dim + 63u) / 64u, dispatch_experts, 1, args.expert_major),
MTLSizeMake(128, 1, 1), 8192u);
}
const bool half = producer.mid == mid && producer.lists == lists && producer.counts == counts &&
producer.tokens == n_tokens && producer.slots == n_slots && producer.stride == n_out &&
producer.width == ff_dim && producer.list_cap == list_cap;
if (half) {
const uint64_t bytes = (uint64_t)n_tokens * n_out * ff_dim * sizeof(uint16_t);
if (!g_qwen4_half_mid || !qwen4_bind_tensor(&b[3], g_qwen4_half_mid, bytes, "moe mid half shadow")) return 0;
[g_transient_buffers addObject:ds4_gpu_tensor_buffer(g_qwen4_half_mid)];
kernel = nt == 1u ? QWEN4_K_MOE_MM_DOWN_F16_NT1 : nt == 2u ? QWEN4_K_MOE_MM_DOWN_F16_NT2 :
nt == 8u ? QWEN4_K_MOE_MM_DOWN_F16_NT8 : QWEN4_K_MOE_MM_DOWN_F16_NT4;
}
if (!qwen4_dispatch(kernel, &args, sizeof(args), b, 5,
qwen4_moe_mm_grid((out_dim + 31u) / 32u, dispatch_experts, tiles, args.expert_major),
MTLSizeMake(128, 1, 1), 0)) return 0;
if (tails) {
const MTLSize tail_grid = MTLSizeMake((out_dim + 31u) / 32u, dispatch_experts, 1);
if (!qwen4_dispatch(QWEN4_K_MOE_MM_DOWN_NT1, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 2u && !qwen4_dispatch(QWEN4_K_MOE_MM_DOWN_NT2, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 4u && !qwen4_dispatch(QWEN4_K_MOE_MM_DOWN, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (!qwen4_dispatch(half ? QWEN4_K_MOE_MM_DOWN_F16_NT1 : QWEN4_K_MOE_MM_DOWN_NT1, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 2u && !qwen4_dispatch(half ? QWEN4_K_MOE_MM_DOWN_F16_NT2 : QWEN4_K_MOE_MM_DOWN_NT2, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 4u && !qwen4_dispatch(half ? QWEN4_K_MOE_MM_DOWN_F16_NT4 : QWEN4_K_MOE_MM_DOWN, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
}
return 1;
}
@@ -50671,6 +50743,7 @@ typedef struct {
uint32_t resident_mask[3], missing_mask[3];
double resident_ms, missing_t0;
bool started, split, missing_mid_submitted;
qwen4_moe_half_pass *half_pass;
} qwen4_stream_mid_overlap;
static int qwen4_stream_mid_before_read(
@@ -50683,7 +50756,7 @@ static int qwen4_stream_mid_before_read(
const bool mm = p->lists != NULL;
const uint32_t routed_mask = (1u << p->n_slots) - 1u;
qwen4_stream_weights snapshot = { .table = t, .addresses = true,
.masked_tokens = mm ? 0u : p->n_tokens, .frequency = p->frequency };
.masked_tokens = mm ? 0u : p->n_tokens, .frequency = p->frequency, .half_pass = p->half_pass };
for (uint32_t u = 0; u < n_unique; u++) {
if (!entries[u]) continue;
snapshot.entries[snapshot.n_entries++] = entries[u];
@@ -50794,7 +50867,7 @@ static int qwen4_stream_missing_gate_up_ready(
if (!p->lists || !n_pairs) return 0;
const NSUInteger addr_bytes = (NSUInteger)t->n_total_expert * sizeof(uint64_t);
qwen4_stream_weights snapshot = { .table = t, .addresses = true,
.frequency = p->frequency };
.frequency = p->frequency, .half_pass = p->half_pass };
snapshot.gate = [g_device newBufferWithLength:addr_bytes options:MTLResourceStorageModeShared];
snapshot.up = [g_device newBufferWithLength:addr_bytes options:MTLResourceStorageModeShared];
NSMutableSet<id<MTLBuffer>> *payloads = [NSMutableSet set];
@@ -50901,11 +50974,17 @@ int ds4_gpu_qwen4_moe_stream_tensor(
int32_t *ids = malloc((size_t)n_ids * sizeof(*ids));
if (!ids) return 0;
qwen4_stream_weights stream = { .table = table };
qwen4_moe_half_pass half_pass = {0};
/* The complete SSD workload benefits at 8192 tokens. The 2048-token
* microbenchmark gain did not translate into a reliable total gain;
* smaller chunks retain the original path without extra scratch. */
const bool half_prefill = mm && gate_type == 16u && down_type == 10u && n_tokens >= 8192u &&
ds4_gpu_device_name_contains("M1 Max");
qwen4_stream_weights stream = { .table = table, .half_pass = half_prefill ? &half_pass : NULL };
uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = {0};
qwen4_stream_mid_overlap overlap = {
.table = table, .mid = mid, .part = part, .x = x, .selected = selected, .ids = ids,
.lists = lists, .frequency = frequency, .list_cap = list_cap,
.lists = lists, .frequency = frequency, .list_cap = list_cap, .half_pass = stream.half_pass,
.gate_type = gate_type, .down_type = down_type, .n_tokens = n_tokens,
.n_slots = n_slots, .in_dim = in_dim, .ff_dim = ff_dim, .shared_down_type = shared_down_type,
.shared_type = shared_type, .shared_gate_offset = shared_gate_offset,
metal/qwen4.metal
+123 -43
View File
@@ -4038,19 +4038,22 @@ static inline void qwen4_dequant_raw16(qwen4_raw16 r, uint b, uint q0, uint type
* to expert e, as 32-row x (8*NT)-token tiles: A (weights, half) and B
* (activations, half) are staged in threadgroup memory per 64-wide K step
* and multiplied with simdgroup matrices into float accumulators, so each
* expert row is read once per token tile. */
template <uint NT>
kernel void kernel_qwen4_moe_mm_mid(
* expert row is read once per token tile. The half-storage variants only
* move the existing B-operand conversion to a shared copy of x and to the
* mid epilogue; all matrix accumulators and activation math remain float. */
template <uint NT, typename X, bool HALF_SHADOW>
static inline void qwen4_moe_mm_mid_impl(
constant ds4_metal_args_qwen4_moe_mm & args,
device const char *gate_base,
device const char *up_base,
device const int32_t *lists,
device const int32_t *counts,
device const float *x, /* [T][in_dim] */
device const X *x, /* [T][in_dim] */
device float *mid, /* [T][n_out][out_rows] */
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
device half *mid_shadow,
threadgroup half *Ag, threadgroup half *Au, threadgroup half *Bs,
threadgroup float *Cs,
uint3 tgpig, ushort tid, ushort sgitg) {
constexpr uint TT = QWEN4_MM_TOKS * NT;
const uint2 block = qwen4_moe_mm_block(args, tgpig);
const uint rb = block.x, e = qwen4_moe_mm_expert(args, tgpig.y);
@@ -4069,10 +4072,6 @@ kernel void kernel_qwen4_moe_mm_mid(
work_count = count - remainder;
}
}
threadgroup half Ag[QWEN4_MM_ROWS * QWEN4_MM_KS];
threadgroup half Au[QWEN4_MM_ROWS * QWEN4_MM_KS];
threadgroup half Bs[QWEN4_MM_KS * TT];
threadgroup float Cs[4][2][64];
device const char *gbase = qwen4_expert_base(gate_base, e, args.expert_bytes);
device const char *ubase = qwen4_expert_base(up_base, e, args.expert_bytes);
device const int32_t *list = lists + (uint64_t)e * args.list_cap;
@@ -4110,13 +4109,21 @@ kernel void kernel_qwen4_moe_mm_mid(
{
const uint tok = tid % TT, kq = tid / TT;
constexpr uint K_PER_THREAD = QWEN4_MM_KS * TT / 128;
device const float *xr = x + (uint64_t)my_t * args.in_dim + kb * QWEN4_MM_KS + kq * K_PER_THREAD;
device const X *xr = x + (uint64_t)my_t * args.in_dim + kb * QWEN4_MM_KS + kq * K_PER_THREAD;
for (uint j = 0; j < K_PER_THREAD; j += 4) {
const float4 v = my_pair >= 0 ? *(device const float4 *)(xr + j) : float4(0.0f);
Bs[(kq * K_PER_THREAD + j + 0) * TT + tok] = (half)v.x;
Bs[(kq * K_PER_THREAD + j + 1) * TT + tok] = (half)v.y;
Bs[(kq * K_PER_THREAD + j + 2) * TT + tok] = (half)v.z;
Bs[(kq * K_PER_THREAD + j + 3) * TT + tok] = (half)v.w;
if constexpr (is_same<X, half>::value) {
const half4 v = my_pair >= 0 ? *(device const half4 *)(xr + j) : half4(0.0h);
Bs[(kq * K_PER_THREAD + j + 0) * TT + tok] = v.x;
Bs[(kq * K_PER_THREAD + j + 1) * TT + tok] = v.y;
Bs[(kq * K_PER_THREAD + j + 2) * TT + tok] = v.z;
Bs[(kq * K_PER_THREAD + j + 3) * TT + tok] = v.w;
} else {
const float4 v = my_pair >= 0 ? *(device const float4 *)(xr + j) : float4(0.0f);
Bs[(kq * K_PER_THREAD + j + 0) * TT + tok] = (half)v.x;
Bs[(kq * K_PER_THREAD + j + 1) * TT + tok] = (half)v.y;
Bs[(kq * K_PER_THREAD + j + 2) * TT + tok] = (half)v.z;
Bs[(kq * K_PER_THREAD + j + 3) * TT + tok] = (half)v.w;
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -4133,8 +4140,8 @@ kernel void kernel_qwen4_moe_mm_mid(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
for (uint nt = 0; nt < NT; nt++) {
simdgroup_store(Cg[nt], Cs[sgitg][0], 8, 0, false);
simdgroup_store(Cu[nt], Cs[sgitg][1], 8, 0, false);
simdgroup_store(Cg[nt], Cs + ((uint)sgitg * 2u) * 64u, 8, 0, false);
simdgroup_store(Cu[nt], Cs + ((uint)sgitg * 2u + 1u) * 64u, 8, 0, false);
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint idx = tid; idx < 4 * 64; idx += 128) {
const uint sg = idx / 64, el = idx % 64, r = el / 8, tok = nt * 8 + el % 8;
@@ -4142,36 +4149,90 @@ kernel void kernel_qwen4_moe_mm_mid(
if (tok >= n_tile || row >= args.out_rows) continue;
const int pair = list[t0 + tok];
const uint t = (uint)pair / args.n_slots, slot = (uint)pair % args.n_slots;
const float g = Cs[sg][0][el], u = Cs[sg][1][el];
mid[((uint64_t)t * args.n_out + slot) * args.out_rows + row] = qwen4_silu(g) * u;
const float g = Cs[(sg * 2u) * 64u + el], u = Cs[(sg * 2u + 1u) * 64u + el];
if constexpr (HALF_SHADOW) {
/* Round the completed FP32 activation, exactly where
* the down kernel would round its float mid operand. */
const float v = qwen4_silu(g) * u;
mid[((uint64_t)t * args.n_out + slot) * args.out_rows + row] = v;
mid_shadow[((uint64_t)t * args.n_out + slot) * args.out_rows + row] = (half)v;
} else {
mid[((uint64_t)t * args.n_out + slot) * args.out_rows + row] = qwen4_silu(g) * u;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}
}
template [[host_name("kernel_qwen4_moe_mm_mid_nt1")]]
kernel void kernel_qwen4_moe_mm_mid<1>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
/* Keep the existing float API and its intermediate observable to diagnostics.
* The half-storage wrapper adds a shadow consumed only by the paired down
* dispatch; it never changes the public mid tensor's element type. */
template <uint NT>
kernel void kernel_qwen4_moe_mm_mid(
constant ds4_metal_args_qwen4_moe_mm & args,
device const char *gate_base, device const char *up_base,
device const int32_t *lists, device const int32_t *counts,
device const float *x, device float *mid,
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
threadgroup half Ag[QWEN4_MM_ROWS * QWEN4_MM_KS];
threadgroup half Au[QWEN4_MM_ROWS * QWEN4_MM_KS];
threadgroup half Bs[QWEN4_MM_KS * QWEN4_MM_TOKS * NT];
threadgroup float Cs[4 * 2 * 64];
qwen4_moe_mm_mid_impl<NT, float, false>(args, gate_base, up_base,
lists, counts, x, mid, nullptr, Ag, Au, Bs, Cs, tgpig, tid, sgitg);
}
template [[host_name("kernel_qwen4_moe_mm_mid_nt2")]]
kernel void kernel_qwen4_moe_mm_mid<2>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
template <uint NT>
kernel void kernel_qwen4_moe_mm_mid_f16(
constant ds4_metal_args_qwen4_moe_mm & args,
device const char *gate_base, device const char *up_base,
device const int32_t *lists, device const int32_t *counts,
device const half *x, device float *mid, device half *mid_shadow,
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
threadgroup half Ag[QWEN4_MM_ROWS * QWEN4_MM_KS];
threadgroup half Au[QWEN4_MM_ROWS * QWEN4_MM_KS];
threadgroup half Bs[QWEN4_MM_KS * QWEN4_MM_TOKS * NT];
threadgroup float Cs[4 * 2 * 64];
qwen4_moe_mm_mid_impl<NT, half, true>(args, gate_base, up_base,
lists, counts, x, mid, mid_shadow, Ag, Au, Bs, Cs, tgpig, tid, sgitg);
}
template [[host_name("kernel_qwen4_moe_mm_mid")]]
kernel void kernel_qwen4_moe_mm_mid<4>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
#define QWEN4_MOE_MM_MID_F32_INSTANCE(NT_, NAME_) \
template [[host_name(NAME_)]] \
kernel void kernel_qwen4_moe_mm_mid<NT_>(constant ds4_metal_args_qwen4_moe_mm &, \
device const char *, device const char *, device const int32_t *, device const int32_t *, \
device const float *, device float *, uint3, ushort, ushort);
QWEN4_MOE_MM_MID_F32_INSTANCE(1, "kernel_qwen4_moe_mm_mid_nt1")
QWEN4_MOE_MM_MID_F32_INSTANCE(2, "kernel_qwen4_moe_mm_mid_nt2")
QWEN4_MOE_MM_MID_F32_INSTANCE(4, "kernel_qwen4_moe_mm_mid")
/* 64-token tiles retain the existing K order and 8/16/32-token remainders. */
QWEN4_MOE_MM_MID_F32_INSTANCE(8, "kernel_qwen4_moe_mm_mid_nt8")
#undef QWEN4_MOE_MM_MID_F32_INSTANCE
/* 64-token tiles: each decoded weight tile serves twice the tokens (8 KB of
* activations staged per K step); remainders take the 8/16/32-token kernels. */
template [[host_name("kernel_qwen4_moe_mm_mid_nt8")]]
kernel void kernel_qwen4_moe_mm_mid<8>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
#define QWEN4_MOE_MM_MID_F16_INSTANCE(NT_) \
template [[host_name("kernel_qwen4_moe_mm_mid_f16_nt" #NT_)]] \
kernel void kernel_qwen4_moe_mm_mid_f16<NT_>(constant ds4_metal_args_qwen4_moe_mm &, \
device const char *, device const char *, device const int32_t *, device const int32_t *, \
device const half *, device float *, device half *, uint3, ushort, ushort);
QWEN4_MOE_MM_MID_F16_INSTANCE(1)
QWEN4_MOE_MM_MID_F16_INSTANCE(2)
QWEN4_MOE_MM_MID_F16_INSTANCE(4)
QWEN4_MOE_MM_MID_F16_INSTANCE(8)
#undef QWEN4_MOE_MM_MID_F16_INSTANCE
/* part[t][slot][r] = down . mid[t][slot], same tiling with mid as B */
template <uint NT>
template <uint NT, typename MID>
kernel void kernel_qwen4_moe_mm_down(
constant ds4_metal_args_qwen4_moe_mm & args,
device const char *down_base,
device const int32_t *lists,
device const int32_t *counts,
device const float *midv, /* [T][n_out][in_dim] */
device const MID *midv, /* [T][n_out][in_dim] */
device float *part, /* [T][n_out][out_rows] */
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
@@ -4226,13 +4287,21 @@ kernel void kernel_qwen4_moe_mm_down(
{
const uint tok = tid % TT, kq = tid / TT;
constexpr uint K_PER_THREAD = QWEN4_MM_KS * TT / 128;
device const float *mr = midv + my_row * args.in_dim + kb * QWEN4_MM_KS + kq * K_PER_THREAD;
device const MID *mr = midv + my_row * args.in_dim + kb * QWEN4_MM_KS + kq * K_PER_THREAD;
for (uint j = 0; j < K_PER_THREAD; j += 4) {
const float4 v = my_pair >= 0 ? *(device const float4 *)(mr + j) : float4(0.0f);
Bs[(kq * K_PER_THREAD + j + 0) * TT + tok] = (half)v.x;
Bs[(kq * K_PER_THREAD + j + 1) * TT + tok] = (half)v.y;
Bs[(kq * K_PER_THREAD + j + 2) * TT + tok] = (half)v.z;
Bs[(kq * K_PER_THREAD + j + 3) * TT + tok] = (half)v.w;
if constexpr (is_same<MID, half>::value) {
const half4 v = my_pair >= 0 ? *(device const half4 *)(mr + j) : half4(0.0h);
Bs[(kq * K_PER_THREAD + j + 0) * TT + tok] = v.x;
Bs[(kq * K_PER_THREAD + j + 1) * TT + tok] = v.y;
Bs[(kq * K_PER_THREAD + j + 2) * TT + tok] = v.z;
Bs[(kq * K_PER_THREAD + j + 3) * TT + tok] = v.w;
} else {
const float4 v = my_pair >= 0 ? *(device const float4 *)(mr + j) : float4(0.0f);
Bs[(kq * K_PER_THREAD + j + 0) * TT + tok] = (half)v.x;
Bs[(kq * K_PER_THREAD + j + 1) * TT + tok] = (half)v.y;
Bs[(kq * K_PER_THREAD + j + 2) * TT + tok] = (half)v.z;
Bs[(kq * K_PER_THREAD + j + 3) * TT + tok] = (half)v.w;
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -4263,16 +4332,27 @@ kernel void kernel_qwen4_moe_mm_down(
}
template [[host_name("kernel_qwen4_moe_mm_down_nt1")]]
kernel void kernel_qwen4_moe_mm_down<1>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
kernel void kernel_qwen4_moe_mm_down<1, float>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
template [[host_name("kernel_qwen4_moe_mm_down_nt2")]]
kernel void kernel_qwen4_moe_mm_down<2>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
kernel void kernel_qwen4_moe_mm_down<2, float>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
template [[host_name("kernel_qwen4_moe_mm_down")]]
kernel void kernel_qwen4_moe_mm_down<4>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
kernel void kernel_qwen4_moe_mm_down<4, float>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
template [[host_name("kernel_qwen4_moe_mm_down_nt8")]]
kernel void kernel_qwen4_moe_mm_down<8>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
kernel void kernel_qwen4_moe_mm_down<8, float>(constant ds4_metal_args_qwen4_moe_mm &, device const char *, device const int32_t *, device const int32_t *, device const float *, device float *, uint3, ushort, ushort);
#define QWEN4_MOE_MM_DOWN_F16_INSTANCE(NT_) \
template [[host_name("kernel_qwen4_moe_mm_down_f16_nt" #NT_)]] \
kernel void kernel_qwen4_moe_mm_down<NT_, half>(constant ds4_metal_args_qwen4_moe_mm &, \
device const char *, device const int32_t *, device const int32_t *, \
device const half *, device float *, uint3, ushort, ushort);
QWEN4_MOE_MM_DOWN_F16_INSTANCE(1)
QWEN4_MOE_MM_DOWN_F16_INSTANCE(2)
QWEN4_MOE_MM_DOWN_F16_INSTANCE(4)
QWEN4_MOE_MM_DOWN_F16_INSTANCE(8)
#undef QWEN4_MOE_MM_DOWN_F16_INSTANCE
/* rows of floats -> halves (round to nearest even), four values per thread;
* the tensor-op tiles read their activation operand from this copy. */
tests/test_metal_qwen4_moe_half.m
+414
View File
@@ -0,0 +1,414 @@
/* Simdgroup MoE half-storage regression and benchmark, independent of host
* dispatch policy. IQ2XXS gate/up and padded Q2_K F640 down use production
* kernels. A separately compiled --baseline-source freezes the float oracle.
* Build: clang -O2 -Wall -Wextra -fobjc-arc -framework Foundation -framework Metal \
* tests/test_metal_qwen4_moe_half.m -o /tmp/test_metal_qwen4_moe_half
* Run from the repository root. --compile-only creates pipelines but no queue
* or GPU dispatch. --bench includes x conversion and the dual mid write in
* every candidate timing; alternating ABBA/BAAB rounds share the same fixture.
* Generic weight types (function constant 900 = 0) match M1 large-prefill
* production dispatch. --specialized-types retains the specialized comparison.
*/
#import <Foundation/Foundation.h>
#import <Metal/Metal.h>
#include <math.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
typedef struct {
uint32_t n_tokens, n_slots, n_out, in_dim, out_rows, weight_type, row_bytes, list_cap;
uint64_t expert_bytes;
uint32_t n_expert, tiles_per_launch, tail_base, expert_major, n_active_expert;
uint32_t active_expert[512];
} mm_args;
_Static_assert(sizeof(mm_args) == 2112, "production MoE MM ABI");
_Static_assert(offsetof(mm_args, active_expert) == 60, "production active expert ABI");
static const NSUInteger kOffset = 256;
static const unsigned kNT[] = {1, 2, 4, 8};
static uint64_t comparisons, half_comparisons, fixtures;
static bool specialized_types;
static void fail(NSString *message) {
fprintf(stderr, "FAIL Qwen MoE half: %s\n", message.UTF8String);
exit(1);
}
static NSString *read_source(NSString *path) {
NSError *error = nil;
NSString *s = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:&error];
if (!s) fail([NSString stringWithFormat:@"read %@: %@", path, error]);
return s;
}
static NSString *between(NSString *s, NSString *begin, NSString *end) {
if ([s componentsSeparatedByString:begin].count != 2 ||
[s componentsSeparatedByString:end].count != 2) fail(@"source anchor missing/nonunique");
NSRange a = [s rangeOfString:begin], b = [s rangeOfString:end];
if (b.location <= a.location) fail(@"invalid source range");
return [s substringWithRange:NSMakeRange(a.location, b.location - a.location)];
}
static NSString *production_source(NSString *repo, NSString *qwenPath) {
NSString *q = read_source(qwenPath);
NSString *m = read_source([repo stringByAppendingPathComponent:@"metal/moe.metal"]);
NSMutableString *s = [NSMutableString stringWithString:
@"#include <metal_stdlib>\nusing namespace metal;\n#define QWEN4_ROUTER_MAX_EXPERT 512\n"];
[s appendString:between(m, @"static constant float ds4_metal_mxfp4_values[16]", @"// BEGIN GENERATED MXFP4 HALF LUT")];
[s appendString:between(m, @"static constant uchar ds4_metal_ksigns_iq2xs[128]", @"#define kmask_iq2xs")];
[s appendString:between(q, @"static inline float qwen4_sigmoid(", @"static inline float qwen4_softplus(")];
[s appendString:between(q, @"static inline float qwen4_silu(", @"/* --- hyper-connections")];
[s appendString:between(q, @"constant bool qwen4_expert_addresses ", @"#define QWEN4_MOE_NSG")];
[s appendString:between(q, @"struct ds4_metal_args_qwen4_moe_mm {", @"#ifdef DS4_METAL_HAS_TENSOR")];
return s;
}
static uint64_t source_hash(NSString *s) {
NSData *data = [s dataUsingEncoding:NSUTF8StringEncoding];
const uint8_t *p = data.bytes;
uint64_t h = UINT64_C(14695981039346656037);
for (NSUInteger i = 0; i < data.length; i++) h = (h ^ p[i]) * UINT64_C(1099511628211);
return h;
}
static uint32_t mix(uint32_t x) {
x ^= x >> 16; x *= UINT32_C(0x7feb352d);
x ^= x >> 15; x *= UINT32_C(0x846ca68b);
return x ^ (x >> 16);
}
static void *data(id<MTLBuffer> b) { return (char *)b.contents + kOffset; }
static id<MTLBuffer> buffer(id<MTLDevice> device, size_t bytes) {
id<MTLBuffer> b = [device newBufferWithLength:bytes + 2 * kOffset options:MTLResourceStorageModeShared];
if (!b) fail([NSString stringWithFormat:@"allocate %zu bytes", bytes]);
memset(b.contents, 0xA5, b.length);
return b;
}
static void guard(id<MTLBuffer> b) {
const uint8_t *p = b.contents;
for (NSUInteger i = 0; i < kOffset; i++)
if (p[i] != 0xA5 || p[b.length - 1 - i] != 0xA5) fail(@"buffer guard overwritten");
}
static void fill_weights(id<MTLBuffer> b, unsigned type, unsigned seed) {
uint8_t *p = data(b);
const size_t bytes = b.length - 2 * kOffset;
for (size_t i = 0; i < bytes; i++) p[i] = (uint8_t)mix((uint32_t)i + seed);
const unsigned block = type == 16 ? 66 : 84;
for (size_t i = 0; i < bytes; i += block) {
/* Small normal scales keep the complete fixture finite. Alternating
* signs and very different magnitudes still exercise cancellation. */
uint16_t h = (uint16_t)(0x1000u | (mix((uint32_t)i + seed) & 0x83ffu));
memcpy(p + i + (type == 16 ? 0 : 80), &h, 2);
if (type == 10) { h ^= 0x0195u; memcpy(p + i + 82, &h, 2); }
}
}
@interface Fixture : NSObject
@property unsigned tokens, dim, ff, experts, slots;
@property mm_args midArgs, downArgs;
@property(strong) id<MTLBuffer> gate, up, down, lists, counts, x, xHalf;
@property(strong) id<MTLBuffer> refMid, gotMid, midHalf, refPart, gotPart;
@end
@implementation Fixture
@end
static Fixture *fixture(id<MTLDevice> device, unsigned tokens, unsigned dim,
unsigned experts, bool sparse, bool expertMajor) {
Fixture *f = [Fixture new];
f.tokens = tokens; f.dim = dim; f.ff = 640; f.experts = experts; f.slots = 10;
mm_args m = {.n_tokens=tokens, .n_slots=f.slots, .n_out=f.slots,
.in_dim=dim, .out_rows=f.ff, .weight_type=16, .row_bytes=dim/256*66,
.list_cap=tokens, .n_expert=experts, .tiles_per_launch=MIN(8u,(tokens+31u)/32u),
.expert_major=expertMajor};
m.expert_bytes = (uint64_t)m.row_bytes * m.out_rows;
mm_args d = m;
d.in_dim = f.ff; d.out_rows = dim; d.weight_type = 10; d.row_bytes = 252;
d.expert_bytes = (uint64_t)d.row_bytes * d.out_rows;
if (sparse) {
m.n_active_expert = d.n_active_expert = experts - 2;
for (unsigned e = 0; e < m.n_active_expert; e++) m.active_expert[e] = d.active_expert[e] = e;
}
f.midArgs = m; f.downArgs = d;
f.gate = buffer(device, m.expert_bytes * experts); f.up = buffer(device, m.expert_bytes * experts);
f.down = buffer(device, d.expert_bytes * experts);
fill_weights(f.gate, 16, 17); fill_weights(f.up, 16, 59); fill_weights(f.down, 10, 117);
f.lists = buffer(device, (size_t)experts * tokens * sizeof(int32_t));
f.counts = buffer(device, experts * sizeof(int32_t));
int32_t *lists = data(f.lists), *counts = data(f.counts);
memset(counts, 0, experts * sizeof(int32_t));
const unsigned active = sparse ? experts - 2 : experts;
for (unsigned t = 0; t < tokens; t++) for (unsigned slot = 0; slot < f.slots; slot++) {
/* Distinct slots per token: consecutive IDs avoid duplicate expert
* selection even in the small 11/13-expert regression fixtures. */
unsigned e = (t * 7 + slot) % active;
lists[(size_t)e * tokens + (unsigned)counts[e]++] = (int32_t)(t * f.slots + slot);
}
f.x = buffer(device, (size_t)tokens * dim * 4); f.xHalf = buffer(device, (size_t)tokens * dim * 2);
float *x = data(f.x);
for (size_t i = 0; i < (size_t)tokens * dim; i++)
x[i] = ((int)(mix((uint32_t)i + 31) % 16385) - 8192) / 8192.f;
const size_t nm = (size_t)tokens * f.slots * f.ff, np = (size_t)tokens * f.slots * dim;
f.refMid = buffer(device, nm * 4); f.gotMid = buffer(device, nm * 4); f.midHalf = buffer(device, nm * 2);
f.refPart = buffer(device, np * 4); f.gotPart = buffer(device, np * 4);
return f;
}
static id<MTLLibrary> library(id<MTLDevice> device, NSString *source, bool safe) {
MTLCompileOptions *options = [MTLCompileOptions new];
if (safe) {
if (@available(macOS 15.0, *)) options.mathMode = MTLMathModeSafe;
else {
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
options.fastMathEnabled = NO;
#pragma clang diagnostic pop
}
}
NSError *error = nil;
id<MTLLibrary> lib = [device newLibraryWithSource:source options:options error:&error];
if (!lib) fail(error.description);
return lib;
}
static id<MTLComputePipelineState> pipeline(id<MTLDevice> device, id<MTLLibrary> lib,
NSString *name, unsigned type) {
NSError *error = nil;
MTLFunctionConstantValues *constants = [MTLFunctionConstantValues new];
unsigned tail = 0; bool addressed = false;
const unsigned bound_type = specialized_types ? type : 0u;
[constants setConstantValue:&bound_type type:MTLDataTypeUInt atIndex:900];
[constants setConstantValue:&tail type:MTLDataTypeUInt atIndex:905];
[constants setConstantValue:&addressed type:MTLDataTypeBool atIndex:906];
id<MTLFunction> fn = [lib newFunctionWithName:name constantValues:constants error:&error];
if (!fn) fail([NSString stringWithFormat:@"%@: %@", name, error]);
id<MTLComputePipelineState> p = [device newComputePipelineStateWithFunction:fn error:&error];
if (!p) fail(error.description);
if (p.threadExecutionWidth != 32 || p.maxTotalThreadsPerThreadgroup < 128) fail(@"unsupported execution geometry");
return p;
}
static NSString *kernel_name(bool mid, bool half, unsigned nt) {
NSString *kind = mid ? @"mid" : @"down";
if (half) return [NSString stringWithFormat:@"kernel_qwen4_moe_mm_%@_f16_nt%u", kind, nt];
return [NSString stringWithFormat:@"kernel_qwen4_moe_mm_%@%@", kind,
nt == 4 ? @"" : [NSString stringWithFormat:@"_nt%u", nt]];
}
static NSArray *pair(id<MTLDevice> device, id<MTLLibrary> lib, bool half, unsigned midNT, unsigned downNT) {
return @[pipeline(device, lib, kernel_name(true, half, midNT), 16),
pipeline(device, lib, kernel_name(false, half, downNT), 10)];
}
static void bind(id<MTLComputeCommandEncoder> e, id<MTLBuffer> b, unsigned index) {
[e setBuffer:b offset:kOffset atIndex:index];
}
static void encode_convert(id<MTLComputeCommandEncoder> en, id<MTLComputePipelineState> p, Fixture *f) {
const uint32_t n4 = f.tokens * f.dim / 4;
[en setComputePipelineState:p]; [en setBytes:&n4 length:sizeof(n4) atIndex:0];
bind(en, f.x, 1); bind(en, f.xHalf, 2);
[en dispatchThreadgroups:MTLSizeMake(((n4+3)/4+255)/256,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)];
}
static void encode_mm(id<MTLComputeCommandEncoder> en, id<MTLComputePipelineState> p,
mm_args args, NSArray<id<MTLBuffer>> *buffers) {
[en setComputePipelineState:p]; [en setBytes:&args length:sizeof(args) atIndex:0];
for (NSUInteger i = 0; i < buffers.count; i++) bind(en, buffers[i], (unsigned)i + 1);
unsigned rb = (args.out_rows + 31) / 32, ne = args.n_active_expert ?: args.n_expert;
MTLSize grid = args.expert_major ? MTLSizeMake((NSUInteger)rb*args.tiles_per_launch,ne,1)
: MTLSizeMake(rb,ne,args.tiles_per_launch);
[en dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(128,1,1)];
}
static double run(id<MTLCommandQueue> queue, NSArray *ps, id<MTLComputePipelineState> convert,
Fixture *f, bool half, bool reference) {
id<MTLBuffer> mid = reference ? f.refMid : f.gotMid, part = reference ? f.refPart : f.gotPart;
id<MTLCommandBuffer> cb = [queue commandBuffer];
id<MTLComputeCommandEncoder> en = [cb computeCommandEncoder];
if (!cb || !en) fail(@"command allocation");
if (half) encode_convert(en, convert, f);
NSMutableArray *in = [@[f.gate,f.up,f.lists,f.counts,half?f.xHalf:f.x,mid] mutableCopy];
if (half) [in addObject:f.midHalf];
encode_mm(en, ps[0], f.midArgs, in);
encode_mm(en, ps[1], f.downArgs, @[f.down,f.lists,f.counts,half?f.midHalf:mid,part]);
[en endEncoding]; [cb commit]; [cb waitUntilCompleted];
if (cb.status != MTLCommandBufferStatusCompleted) fail(cb.error.description);
const double elapsed = (cb.GPUEndTime - cb.GPUStartTime) * 1000.;
if (!(elapsed > 0)) fail(@"GPU timestamps unavailable");
return elapsed;
}
static void check(id<MTLBuffer> reference, id<MTLBuffer> candidate, NSString *label) {
const size_t n = (reference.length - 2 * kOffset) / 4;
const uint32_t *r = data(reference), *c = data(candidate);
for (size_t i = 0; i < n; i++) {
if (r[i] == UINT32_C(0xa5a5a5a5) || c[i] == UINT32_C(0xa5a5a5a5))
fail([NSString stringWithFormat:@"%@ output untouched at %zu", label, i]);
if ((r[i]&0x7f800000u) == 0x7f800000u || (c[i]&0x7f800000u) == 0x7f800000u)
fail([NSString stringWithFormat:@"%@ nonfinite at %zu", label, i]);
if (r[i] != c[i]) fail([NSString stringWithFormat:@"%@ mismatch at %zu: %08x != %08x", label, i, r[i], c[i]]);
}
comparisons += n; guard(reference); guard(candidate);
}
static void poison(id<MTLBuffer> b) {
/* Payload only: pre-existing boundary guards must remain observable. */
memset(data(b), 0xA5, b.length - 2 * kOffset);
}
static uint16_t half_bits(float value) {
const _Float16 rounded = (_Float16)value;
uint16_t bits;
_Static_assert(sizeof(rounded) == sizeof(bits), "CPU IEEE half storage");
memcpy(&bits, &rounded, sizeof(bits));
return bits;
}
static void check_half_payloads(Fixture *f) {
const float *x = data(f.x), *mid = data(f.refMid);
const uint16_t *xh = data(f.xHalf), *mh = data(f.midHalf);
const size_t nx = (size_t)f.tokens * f.dim;
for (size_t i = 0; i < nx; ++i) {
const uint16_t expected = half_bits(x[i]);
if (xh[i] != expected)
fail([NSString stringWithFormat:@"x half payload at %zu: %04x != %04x", i, xh[i], expected]);
}
/* The routing list defines the mid rows actually written by the kernel.
* Do not treat unselected expert/list padding as an activation result. */
const mm_args args = f.midArgs;
const size_t pairs = (size_t)args.n_tokens * args.n_out;
uint8_t *written = calloc(pairs, sizeof(*written));
if (!written) fail(@"half payload routing mask allocation");
const int32_t *lists = data(f.lists), *counts = data(f.counts);
const unsigned ne = args.n_active_expert ?: args.n_expert;
for (unsigned grid_e = 0; grid_e < ne; ++grid_e) {
const unsigned e = args.n_active_expert ? args.active_expert[grid_e] : grid_e;
if (e >= args.n_expert || counts[e] < 0 || (unsigned)counts[e] > args.list_cap)
fail(@"invalid half payload fixture routing");
for (unsigned i = 0; i < (unsigned)counts[e]; ++i) {
const int32_t pair = lists[(size_t)e * args.list_cap + i];
if (pair < 0 || (uint64_t)pair >= (uint64_t)args.n_tokens * args.n_slots)
fail(@"half payload fixture pair outside token slots");
const unsigned token = (unsigned)pair / args.n_slots, slot = (unsigned)pair % args.n_slots;
if (slot >= args.n_out) fail(@"half payload fixture output slot");
const size_t row = (size_t)token * args.n_out + slot;
if (written[row]++) fail(@"half payload fixture duplicate output row");
}
}
for (size_t pair = 0; pair < pairs; ++pair) {
for (unsigned r = 0; r < args.out_rows; ++r) {
const size_t i = pair * args.out_rows + r;
const uint16_t expected = written[pair] ? half_bits(mid[i]) : UINT16_C(0xa5a5);
if (mh[i] != expected)
fail([NSString stringWithFormat:@"mid half payload at %zu: %04x != %04x", i, mh[i], expected]);
}
if (written[pair]) half_comparisons += args.out_rows;
}
half_comparisons += nx;
free(written);
guard(f.xHalf); guard(f.midHalf);
}
static void validate(id<MTLCommandQueue> queue, NSArray *baseline, NSArray *current,
NSArray *half, id<MTLComputePipelineState> convert, Fixture *f) {
poison(f.refMid); poison(f.refPart);
run(queue, baseline, convert, f, false, true);
poison(f.gotMid); poison(f.gotPart);
run(queue, current, convert, f, false, false);
check(f.refMid, f.gotMid, @"current float mid"); check(f.refPart, f.gotPart, @"current float down");
poison(f.gotMid); poison(f.gotPart); poison(f.xHalf); poison(f.midHalf);
run(queue, half, convert, f, true, false);
check(f.refMid, f.gotMid, @"half-storage float mid"); check(f.refPart, f.gotPart, @"half-storage down");
check_half_payloads(f);
for (id<MTLBuffer> b in @[f.gate,f.up,f.down,f.lists,f.counts,f.x,f.xHalf,f.midHalf]) guard(b);
fixtures++;
}
static double median(NSArray<NSNumber *> *a) {
NSArray *s = [a sortedArrayUsingSelector:@selector(compare:)];
NSUInteger n = s.count;
return n&1 ? [s[n/2] doubleValue] : ([s[n/2-1] doubleValue]+[s[n/2] doubleValue])/2.;
}
static NSDictionary *benchmark(id<MTLCommandQueue> queue, NSArray *baseline, NSArray *current,
NSArray *half, id<MTLComputePipelineState> convert, Fixture *f,
unsigned midNT, unsigned downNT, unsigned reps) {
validate(queue, baseline, current, half, convert, f);
for (unsigned i = 0; i < 2; i++) { run(queue, baseline, convert, f, false, true); run(queue, half, convert, f, true, false); }
NSMutableArray *a = [NSMutableArray new], *b = [NSMutableArray new];
/* Each round has two samples of each arm with equal mean position in
* the sequence. Reversing successive rounds also balances adjacency. */
for (unsigned i = 0; i < reps; i++) for (unsigned j = 0; j < 4; j++) {
const bool candidate = ((j == 1 || j == 2) != ((i & 1u) != 0));
if (candidate) [b addObject:@(run(queue, half, convert, f, true, false))];
else [a addObject:@(run(queue, baseline, convert, f, false, true))];
}
check(f.refMid, f.gotMid, @"benchmark mid"); check(f.refPart, f.gotPart, @"benchmark down");
const double ma = median(a), mb = median(b);
printf("T=%u E=%u F=%u experts=%u slots=%u midNT=%u downNT=%u: float %.3f ms half %.3f ms speedup %.2f%%\n",
f.tokens,f.dim,f.ff,f.experts,f.slots,midNT,downNT,ma,mb,(ma/mb-1)*100);
fflush(stdout);
return @{@"tokens":@(f.tokens),@"dim":@(f.dim),@"ff":@(f.ff),@"experts":@(f.experts),
@"slots":@(f.slots),@"mid_nt":@(midNT),@"down_nt":@(downNT),
@"float_gpu_ms":a,@"half_gpu_ms":b,@"float_median_ms":@(ma),@"half_median_ms":@(mb),
@"speedup_percent":@((ma/mb-1)*100),@"sample_order":@"alternating ABBA/BAAB",@"rounds":@(reps),
@"warmup_runs_per_arm":@2,@"includes_x_conversion":@YES,@"includes_dual_mid_write":@YES};
}
static void write_json(NSString *path, NSDictionary *report) {
if (!path) return;
NSError *error = nil;
NSData *json = [NSJSONSerialization dataWithJSONObject:report options:NSJSONWritingPrettyPrinted|NSJSONWritingSortedKeys error:&error];
if (!json || ![json writeToFile:path options:NSDataWritingAtomic error:&error]) fail(error.description);
}
int main(int argc, const char **argv) { @autoreleasepool {
NSString *repo = [[NSFileManager defaultManager] currentDirectoryPath], *candidatePath = nil, *baselinePath = nil, *output = nil;
bool compileOnly = false, bench = false, safe = false;
unsigned reps = 9;
for (int i = 1; i < argc; i++) {
if (!strcmp(argv[i], "--compile-only")) compileOnly = true;
else if (!strcmp(argv[i], "--bench")) bench = true;
else if (!strcmp(argv[i], "--safe-math")) safe = true;
else if (!strcmp(argv[i], "--generic-types")) specialized_types = false;
else if (!strcmp(argv[i], "--specialized-types")) specialized_types = true;
else if (!strcmp(argv[i], "--repo") && i+1 < argc) repo = [NSString stringWithUTF8String:argv[++i]];
else if (!strcmp(argv[i], "--qwen-source") && i+1 < argc) candidatePath = [NSString stringWithUTF8String:argv[++i]];
else if (!strcmp(argv[i], "--baseline-source") && i+1 < argc) baselinePath = [NSString stringWithUTF8String:argv[++i]];
else if (!strcmp(argv[i], "--output") && i+1 < argc) output = [NSString stringWithUTF8String:argv[++i]];
else if (!strcmp(argv[i], "--reps") && i+1 < argc) reps = (unsigned)strtoul(argv[++i],NULL,10);
else { fprintf(stderr,"usage: %s [--compile-only] [--bench] [--safe-math] [--generic-types|--specialized-types] [--reps N] [--repo PATH] [--qwen-source FILE] [--baseline-source FILE] [--output JSON]\n",argv[0]); return 2; }
}
if (reps < 3 || reps > 100) fail(@"--reps must be in [3,100]");
candidatePath = candidatePath ?: [repo stringByAppendingPathComponent:@"metal/qwen4.metal"];
const bool explicit_baseline = baselinePath != nil;
baselinePath = baselinePath ?: candidatePath;
NSString *candidateSource = production_source(repo,candidatePath), *baselineSource = production_source(repo,baselinePath);
const bool independent_baseline = explicit_baseline && ![baselineSource isEqualToString:candidateSource];
id<MTLDevice> device = MTLCreateSystemDefaultDevice();
if (!device) fail(@"Metal device unavailable");
id<MTLLibrary> candidateLib = library(device,candidateSource,safe), baselineLib = library(device,baselineSource,safe);
NSMutableArray *bases = [NSMutableArray new], *currents = [NSMutableArray new], *halves = [NSMutableArray new];
for (unsigned i = 0; i < 4; i++) {
[bases addObject:pair(device,baselineLib,false,kNT[i],kNT[i])];
[currents addObject:pair(device,candidateLib,false,kNT[i],kNT[i])];
[halves addObject:pair(device,candidateLib,true,kNT[i],kNT[i])];
}
id<MTLComputePipelineState> convert = pipeline(device,candidateLib,@"kernel_qwen4_rows_f32_to_f16",0);
NSMutableDictionary *report = [@{@"device":device.name,@"safe_math":@(safe),@"compile_only":@(compileOnly),
@"baseline_source":baselinePath,@"candidate_source":candidatePath,
@"explicit_baseline":@(explicit_baseline),@"independent_baseline":@(independent_baseline),
@"weight_type_mode":specialized_types ? @"specialized" : @"generic",
@"function_constant_900_mid":@(specialized_types ? 16u : 0u),
@"function_constant_900_down":@(specialized_types ? 10u : 0u),
@"poisoned_validation_outputs":@YES,@"half_payload_cpu_reference":@YES,
@"baseline_fnv1a64":[NSString stringWithFormat:@"%016llx",(unsigned long long)source_hash(baselineSource)],
@"candidate_fnv1a64":[NSString stringWithFormat:@"%016llx",(unsigned long long)source_hash(candidateSource)]} mutableCopy];
if (!compileOnly) {
id<MTLCommandQueue> queue = [device newCommandQueue]; if (!queue) fail(@"command queue");
const unsigned ts[] = {1,7,9,31,33,67,131};
for (unsigned s = 0; s < sizeof(ts)/sizeof(ts[0]); s++) @autoreleasepool {
Fixture *f = fixture(device,ts[s],256,13,(s&1)!=0,(s&2)!=0);
for (unsigned n = 0; n < 4; n++) validate(queue,bases[n],currents[n],halves[n],convert,f);
printf("PASS T=%u NT=1/2/4/8 sparse=%u expert-major=%u\n",ts[s],s&1,(s&2)!=0); fflush(stdout);
}
NSMutableArray *times = [NSMutableArray new];
if (bench) {
const unsigned bt[] = {128,2048,8192};
for (unsigned i = 0; i < 3; i++) @autoreleasepool {
const unsigned midNT = i==0 ? 1 : 4, downNT = i==0 ? 1 : 4;
Fixture *f = fixture(device,bt[i],2560,512,false,false);
NSArray *bp = pair(device,baselineLib,false,midNT,downNT), *cp = pair(device,candidateLib,false,midNT,downNT),
*hp = pair(device,candidateLib,true,midNT,downNT);
[times addObject:benchmark(queue,bp,cp,hp,convert,f,midNT,downNT,reps)];
}
}
report[@"benchmarks"] = times;
}
report[@"fixtures"] = @(fixtures); report[@"compared_floats"] = @(comparisons);
report[@"compared_half_values"] = @(half_comparisons); report[@"status"] = @"PASS";
write_json(output,report);
printf("PASS Qwen MoE half: %llu fixtures, %llu float comparisons%s\n",
(unsigned long long)fixtures,(unsigned long long)comparisons,compileOnly?" (compile only)":"");
return 0;
} }
tests/test_qwen4_ssd_experts.c
+9 -1
View File
@@ -453,7 +453,7 @@ static void check_split_one(const weights *w, uint32_t T, uint32_t slots, int sh
enum { MAX_SLOTS = 16, INJ = HC * DS4_QWEN4_HC_CHUNKS * HC };
const int mm = T > 8;
need(slots <= MAX_SLOTS && slots >= 4 &&
(mm ? T <= 128 && slots == S && !shared && !duplicates && (seed_count == 0 || seed_count == 12) :
(mm ? T <= 16384 && slots == S && !shared && !duplicates && (seed_count == 0 || seed_count == 12) :
T >= 1 && T <= 3 && seed_count <= 4), "split fixture shape");
const uint32_t unique = mm ? 24u : slots - (duplicates ? 3u : 0u), stride = slots + shared;
const uint32_t budget = unique > BUDGET ? unique : BUDGET;
@@ -727,6 +727,14 @@ int main(void) {
check_split_one(formats + f, shapes[i], S, 0, 12, 0, 2, fileno(file), fileno(empty));
}
}
/* Half operands are private to each large SSD pass. Cross the policy
* boundary, grow scratch, then shrink/reuse it with new input contents.
* References use the public float path; down I/O failure tests leave an
* early half-producing mid in flight before the retry. */
check_split_one(formats, 8191, S, 0, 12, 0, 0, fileno(file), fileno(empty));
check_split_one(formats, 8192, S, 0, 12, 0, 2, fileno(file), fileno(empty));
check_split_one(formats, 16384, S, 0, 0, 0, 0, fileno(file), fileno(empty));
check_split_one(formats, 8192, S, 0, 12, 0, 0, fileno(file), fileno(empty));
ds4_gpu_cleanup(); munmap(map, bytes); fclose(empty); fclose(file);
puts("PASS Qwen SSD: cache, exact selected-only reads, full-selection fallback and retry");
return 0;