Metal: stage only selected Qwen experts for MTP and cache overflow

Avoid reading all 512 experts when a Qwen layer has a different size from
the persistent cache or a prefill selection exceeds its capacity. The Q2
pack's Q4/MXFP4 MTP layer previously read about 1.29 GiB on each invocation,
even when the router selected only ten experts.

Allocate compact gate/up/down payloads for the distinct selected experts
and build private GPU-address tables indexed by their original IDs. Keep
the existing row/MM/NAX arithmetic, Q2 physical down-row padding and final
reductions unchanged. Preserve the trunk cache and its address tables.
Retain all six staging buffers through GPU completion, including batches
using unretained command buffers, and declare the indirect payloads as
encoder resources without registering nonexistent cache entries.

Use the original three contiguous reads when every expert is selected or
GPU buffer addresses are unavailable. Choose that fallback before I/O;
actual read failures propagate without rereading the complete layer.
Keep the existing conservative memory admission and staging reservation.

Extend the SSD fixture with a test-only pread interceptor, linked through a
separate Metal backend object. It performs real reads, rejects accesses to
unselected expert ranges, checks per-projection byte totals and verifies
zero reads for resident/cache-hit paths. Cover compact cache overflow,
off-size MTP with two/three tokens, full selection of all 512 experts,
failed reads, retries and retention of the original trunk cache.
Update the Qwen documentation to describe selective temporary staging.

Validation on Apple M1 Max, 32 GiB:
- Full build, host Qwen memory test and existing resident MoE exact tests.
- SSD bitwise mid/down/reduce/HC parity and read-range assertions, including
  DS4_METAL_UNRETAINED_COMMAND_BUFFERS=1 and existing Metal SSD regressions.
- Real Q2 n-gram recovery: exact logits after prefill/decode/MTP read errors
  across two SSD sessions.
- Three alternating before/after runs of the Rome prompt, ctx4096, chunk128,
  50 generated tokens, temperature0 and no thinking: median MTP generation
  4.55 -> 7.86 t/s (+72.7%). All generated output matched, with 27 verifier
  cycles and 21 accepted drafts in every MTP run.
- A separate profiling comparison reduced expert pread bytes from
  48.40 -> 12.02 GiB (-75.2%); these include filesystem-cache reads.

The speed result is specific to this short MTP workload. Expert I/O still
finishes before routed computation; no overlap or next-layer prefetch is
introduced, and CUDA/ROCm or M5 NAX execution was not tested here.
This commit is contained in:
GiorgioOppo committed 2026-09-14 17:16:37 +02:00
1 parent e37f18576c
commit 85d37aec2a
5 files changed
+209 -12

No files matched your search

+4 -1
View File
@@ -179,7 +179,10 @@ test-qwen4-moe-mm-specialize: tests/test_qwen4_moe_mm_specialize
tests/test_qwen4_ssd_experts.o: tests/test_qwen4_ssd_experts.c ds4_gpu.h
$(CC) $(CFLAGS) -fno-fast-math -I. -c -o $@ $<
tests/test_qwen4_ssd_experts: tests/test_qwen4_ssd_experts.o $(CORE_OBJS)
tests/ds4_metal_qwen_ssd.o: ds4_metal.m ds4_gpu.h ds4_gpu_tp.h ds4_deepseek41_gpu.h $(METAL_SRCS) tests/qwen4_ssd_pread_probe.h
$(CC) $(OBJCFLAGS) -include tests/qwen4_ssd_pread_probe.h -c -o $@ ds4_metal.m
tests/test_qwen4_ssd_experts: tests/test_qwen4_ssd_experts.o $(filter-out ds4_metal.o,$(CORE_OBJS)) tests/ds4_metal_qwen_ssd.o
$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
.PHONY: test-qwen4-ssd-experts
+9 -2
View File
@@ -44,11 +44,15 @@ weights, shared experts, context and runtime buffers still need RAM; the
n-gram table continues to use its existing disk reads. Missing routed experts
are read into owned buffers. If a prefill batch needs more experts than the
cache can hold, or an MTP layer uses a different expert size, it uses explicit
layer staging while retaining the normal kernel arithmetic.
temporary staging while retaining the normal kernel arithmetic. On devices
with GPU buffer addresses, this staging contains only the distinct selected
experts. Selecting every expert, or running without buffer-address support,
retains the full-layer fallback.
The current implementation does not overlap these expert reads with GPU
computation. An MTP layer whose expert size differs from the main model is
read in full on every invocation, about 1.29 GiB for the Q2 pack.
staged again on each invocation; its temporary experts do not replace the
main model's cache entries.
Leave the expert-cache budget automatic initially. A plain count passed to
`--ssd-streaming-cache-experts` requests dynamic cache slots; an `NGB` budget
@@ -166,6 +170,9 @@ including indices 383, 384 and 511, with cold and warm selections, eviction,
prefill staging, different-size MTP layers and failed reads followed by exact recovery. It
compares routed intermediates, reduced outputs and hyper-connection residuals
bit for bit against the resident kernels, including the padded Q2 down rows.
Its test-only backend checks actual `pread` requests against selected expert
ranges and counts returned bytes: ten experts for off-size MTP, 24 for cache
overflow, zero reads for cache hits, and the full 512-expert fallback.
The memory-estimate test runs without a GPU or model.
Run `tests/test_qwen4_ngram_state MODEL.gguf` under Metal validation to check
failed disk reads during prefill, decode and MTP, then exact recovery.
+91 -4
View File
@@ -48086,10 +48086,12 @@ static bool qwen4_bind_weight(qwen4_bind *b, const void *map, uint64_t size,
/* Scoped to one routed mid/down call. Indirect cache buffers are registered
* on each encoder and protected by the same completion sequence as DeepSeek.
* Full-layer staging instead supplies ordinary owned, contiguous buffers. */
* Compact staging supplies private address tables backed by owned payloads;
* full-layer staging supplies ordinary owned, contiguous buffers. */
typedef struct {
const ds4_gpu_stream_expert_table *table;
__strong id<MTLBuffer> gate, up, down;
__strong id<MTLBuffer> owned_gate, owned_up, owned_down;
ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
uint32_t n_entries;
bool addresses;
@@ -48423,7 +48425,7 @@ static int qwen4_dispatch(int kernel, const void *args, size_t args_len,
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (addresses && !ds4_gpu_stream_expert_cache_mark_entries_inflight(
if (addresses && g_qwen4_stream_weights->n_entries && !ds4_gpu_stream_expert_cache_mark_entries_inflight(
g_qwen4_stream_weights->entries, g_qwen4_stream_weights->n_entries, 0)) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
if (addresses) {
@@ -48433,6 +48435,11 @@ static int qwen4_dispatch(int kernel, const void *args, size_t args_len,
[enc useResource:e->up_buffer usage:MTLResourceUsageRead];
[enc useResource:e->down_buffer usage:MTLResourceUsageRead];
}
if (g_qwen4_stream_weights->owned_gate) {
[enc useResource:g_qwen4_stream_weights->owned_gate usage:MTLResourceUsageRead];
[enc useResource:g_qwen4_stream_weights->owned_up usage:MTLResourceUsageRead];
[enc useResource:g_qwen4_stream_weights->owned_down usage:MTLResourceUsageRead];
}
}
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:args_len atIndex:0];
@@ -49603,6 +49610,81 @@ static int qwen4_stream_read_layer(qwen4_stream_weights *s) {
return ok;
}
static int qwen4_stream_read_selected(qwen4_stream_weights *s,
const uint32_t *frequency, uint32_t unique) {
const ds4_gpu_stream_expert_table *t = s->table;
if (!frequency || !unique || unique > t->n_total_expert) return 0;
if (unique == t->n_total_expert) return qwen4_stream_read_layer(s);
if (@available(macOS 13.0, *)) {
/* Only the payload is compacted. Original expert IDs still index
* private address tables, preserving routing and every kernel stride. */
const uint64_t gate_bytes = t->gate_expert_bytes * unique;
const uint64_t down_bytes = t->down_expert_bytes * unique;
const NSUInteger addr_bytes = (NSUInteger)t->n_total_expert * sizeof(uint64_t);
if (g_model_fd < 0 || gate_bytes > (uint64_t)NSUIntegerMax || down_bytes > (uint64_t)NSUIntegerMax ||
gate_bytes > (uint64_t)[g_device maxBufferLength] || down_bytes > (uint64_t)[g_device maxBufferLength]) return 0;
s->owned_gate = [g_device newBufferWithLength:(NSUInteger)gate_bytes options:MTLResourceStorageModeShared];
s->owned_up = [g_device newBufferWithLength:(NSUInteger)gate_bytes options:MTLResourceStorageModeShared];
s->owned_down = [g_device newBufferWithLength:(NSUInteger)down_bytes options:MTLResourceStorageModeShared];
if (!s->owned_gate || !s->owned_up || !s->owned_down) return 0;
const uint64_t ga = ds4_gpu_buffer_address(s->owned_gate, 0);
const uint64_t ua = ds4_gpu_buffer_address(s->owned_up, 0);
const uint64_t da = ds4_gpu_buffer_address(s->owned_down, 0);
if (!ga || !ua || !da) {
/* Decide before any I/O; devices without indirect addresses keep
* the ordinary three contiguous reads. */
s->owned_gate = nil; s->owned_up = nil; s->owned_down = nil;
return qwen4_stream_read_layer(s);
}
s->gate = [g_device newBufferWithLength:addr_bytes options:MTLResourceStorageModeShared];
s->up = [g_device newBufferWithLength:addr_bytes options:MTLResourceStorageModeShared];
s->down = [g_device newBufferWithLength:addr_bytes options:MTLResourceStorageModeShared];
if (!s->gate || !s->up || !s->down) return 0;
s->owned_gate.label = @"qwen4_stream_selected_gate";
s->owned_up.label = @"qwen4_stream_selected_up";
s->owned_down.label = @"qwen4_stream_selected_down";
s->gate.label = @"qwen4_stream_selected_gate_addresses";
s->up.label = @"qwen4_stream_selected_up_addresses";
s->down.label = @"qwen4_stream_selected_down_addresses";
uint64_t *gate_addr = [s->gate contents], *up_addr = [s->up contents], *down_addr = [s->down contents];
uint8_t *gate_dst = [s->owned_gate contents], *up_dst = [s->owned_up contents], *down_dst = [s->owned_down contents];
if (!gate_addr || !up_addr || !down_addr || !gate_dst || !up_dst || !down_dst) return 0;
memset(gate_addr, 0, addr_bytes);
memset(up_addr, 0, addr_bytes);
memset(down_addr, 0, addr_bytes);
ds4_gpu_stream_expert_pread_task *tasks = calloc((size_t)unique * 3u, sizeof(*tasks));
if (!tasks) return 0;
uint32_t u = 0;
for (uint32_t e = 0; e < t->n_total_expert; e++) {
if (!frequency[e]) continue;
if (u >= unique) { free(tasks); return 0; }
const uint64_t go = (uint64_t)u * t->gate_expert_bytes;
const uint64_t d = (uint64_t)u * t->down_expert_bytes;
gate_addr[e] = ga + go;
up_addr[e] = ua + go;
down_addr[e] = da + d;
tasks[3u * u] = (ds4_gpu_stream_expert_pread_task) {
.offset = t->gate_offset + (uint64_t)e * t->gate_expert_bytes,
.len = t->gate_expert_bytes, .dst = gate_dst + go };
tasks[3u * u + 1u] = (ds4_gpu_stream_expert_pread_task) {
.offset = t->up_offset + (uint64_t)e * t->gate_expert_bytes,
.len = t->gate_expert_bytes, .dst = up_dst + go };
tasks[3u * u + 2u] = (ds4_gpu_stream_expert_pread_task) {
.offset = t->down_offset + (uint64_t)e * t->down_expert_bytes,
.len = t->down_expert_bytes, .dst = down_dst + d };
u++;
}
uint64_t bytes = 0;
double ms = 0.0;
const int ok = u == unique && ds4_gpu_stream_expert_pread_tasks(tasks, 3u * unique, &bytes, &ms);
free(tasks);
ds4_gpu_stream_expert_cache_note_pread(t->layer, bytes, ms);
s->addresses = ok != 0;
return ok;
}
return qwen4_stream_read_layer(s);
}
int ds4_gpu_qwen4_moe_stream_tensor(
ds4_gpu_tensor *mid, ds4_gpu_tensor *part, const ds4_gpu_tensor *x,
const ds4_gpu_tensor *selected, const ds4_gpu_tensor *lists, const ds4_gpu_tensor *counts,
@@ -49737,16 +49819,21 @@ int ds4_gpu_qwen4_moe_stream_tensor(
stream.addresses = true;
} else {
ds4_gpu_stream_expert_cache_note_frequency_hotness(table->layer, frequency, table->n_total_expert);
ok = qwen4_stream_read_layer(&stream);
ok = qwen4_stream_read_selected(&stream, frequency, unique);
}
if (timing) ds4_gpu_stream_expert_timing_note_selected(sync_ms, copy_ms, ds4_gpu_now_ms() - t0);
if (!ok || !ds4_gpu_begin_commands()) { ok = 0; goto done; }
if (!stream.addresses) {
if (!stream.addresses || stream.owned_gate) {
/* An enclosing batch may outlive this stack frame. Retain staging
* until completion even with unretained Metal command buffers. */
[g_transient_buffers addObject:stream.gate];
[g_transient_buffers addObject:stream.up];
[g_transient_buffers addObject:stream.down];
if (stream.owned_gate) {
[g_transient_buffers addObject:stream.owned_gate];
[g_transient_buffers addObject:stream.owned_up];
[g_transient_buffers addObject:stream.owned_down];
}
}
g_qwen4_stream_weights = &stream;
if (mm) {
+5
View File
@@ -0,0 +1,5 @@
/* Only forced into the SSD test's backend object. Include the system
* declaration before renaming calls, leaving Darwin's symbol aliases intact. */
#include <unistd.h>
ssize_t ds4_test_pread(int fd, void *dst, size_t bytes, off_t offset);
#define pread ds4_test_pread
+100 -5
View File
@@ -1,6 +1,8 @@
#define _DARWIN_C_SOURCE
#include "ds4_gpu.h"
#include <errno.h>
#include <math.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
@@ -22,9 +24,91 @@ typedef struct {
} weights;
typedef struct { ds4_gpu_tensor *base, *view; uint64_t n; } guarded;
/* This binary links a test-only backend object whose pread calls reach this
* wrapper. All accepted requests still perform real file I/O. Sparse holes
* alone cannot prove selected-only reads: reading a hole succeeds with zeros. */
static pthread_mutex_t probe_mutex = PTHREAD_MUTEX_INITIALIZER;
static struct {
int enabled, fd, empty_fd, violation;
uint32_t inflight;
uint64_t offset[3], stride[3], bytes, calls, coverage[3][E];
unsigned char selected[E];
} probe;
typedef struct { uint64_t bytes, calls; } read_stats;
static void need(int ok, const char *what) {
if (!ok) { fprintf(stderr, "Qwen SSD: %s failed\n", what); exit(1); }
}
ssize_t ds4_test_pread(int fd, void *dst, size_t bytes, off_t offset) {
pthread_mutex_lock(&probe_mutex);
if (!probe.enabled || (fd != probe.fd && fd != probe.empty_fd)) {
pthread_mutex_unlock(&probe_mutex);
return pread(fd, dst, bytes, offset);
}
int projection = -1;
const uint64_t start = (uint64_t)offset;
if (offset >= 0 && bytes) for (unsigned p = 0; p < 3; p++) {
if (start < probe.offset[p]) continue;
const uint64_t relative = start - probe.offset[p], total = E * probe.stride[p];
if (relative >= total || bytes > total - relative) continue;
projection = (int)p;
const uint32_t first = (uint32_t)(relative / probe.stride[p]);
const uint32_t last = (uint32_t)((relative + bytes - 1) / probe.stride[p]);
for (uint32_t e = first; e <= last; e++) if (!probe.selected[e]) projection = -1;
break;
}
probe.calls++;
if (projection < 0) {
probe.violation = 1;
pthread_mutex_unlock(&probe_mutex);
errno = EIO; return -1;
}
probe.inflight++;
pthread_mutex_unlock(&probe_mutex);
const ssize_t got = pread(fd, dst, bytes, offset);
const int saved_errno = errno;
pthread_mutex_lock(&probe_mutex);
probe.inflight--;
if (got > 0) {
probe.bytes += (uint64_t)got;
uint64_t relative = start - probe.offset[projection], remaining = (uint64_t)got;
const uint64_t stride = probe.stride[projection];
while (remaining) {
const uint32_t e = (uint32_t)(relative / stride);
const uint64_t room = stride - relative % stride, n = remaining < room ? remaining : room;
probe.coverage[projection][e] += n; remaining -= n; relative += n;
}
}
pthread_mutex_unlock(&probe_mutex);
errno = saved_errno; return got;
}
static void probe_begin(const weights *w, const int32_t *ids, uint64_t n, int fd, int empty_fd) {
pthread_mutex_lock(&probe_mutex);
need(!probe.enabled && !probe.inflight, "read probe idle");
memset(&probe, 0, sizeof(probe));
probe.fd = fd; probe.empty_fd = empty_fd; probe.enabled = 1;
probe.offset[0] = w->table.gate_offset; probe.offset[1] = w->table.up_offset;
probe.offset[2] = w->table.down_offset;
probe.stride[0] = probe.stride[1] = w->table.gate_expert_bytes;
probe.stride[2] = w->table.down_expert_bytes;
for (uint64_t i = 0; i < n; i++) {
need(ids[i] >= 0 && ids[i] < E, "probe selected ID"); probe.selected[ids[i]] = 1;
}
pthread_mutex_unlock(&probe_mutex);
}
static read_stats probe_end(int all_selected, int no_reads) {
pthread_mutex_lock(&probe_mutex);
need(probe.enabled && !probe.inflight, "read workers retired");
need(!probe.violation, "pread never accesses an unselected expert");
if (no_reads) need(!probe.calls && !probe.bytes, "resident/cache-hit path performs no pread");
for (unsigned p = 0; p < 3; p++) for (unsigned e = 0; e < E; e++) {
const uint64_t expected = probe.selected[e] ? probe.stride[p] : 0;
need(probe.coverage[p][e] <= expected, "expert bytes read at most once");
if (all_selected) need(probe.coverage[p][e] == expected, "each selected projection fully read");
}
const read_stats result = {probe.bytes, probe.calls}; probe.enabled = 0;
pthread_mutex_unlock(&probe_mutex); return result;
}
static uint32_t random_u32(void) {
rng ^= rng << 13; rng ^= rng >> 17; rng ^= rng << 5; return rng;
}
@@ -127,7 +211,8 @@ static void check_case(const weights *w, uint32_t T, unsigned route, int good_fd
for (uint32_t t = 0; t < T; t++) {
sg[t] = ((int)(t % 7) - 3) * 0.25f;
for (uint32_t s = 0; s < S; s++) {
ids[t * S + s] = active[route == 2 ? (t + s) % 24 : route * 10 + (t + s) % 10];
ids[t * S + s] = route == 3 ? (int32_t)((t * S + s) % E) :
active[route == 2 ? (t + s) % 24 : route * 10 + (t + s) % 10];
mix[t * S + s] = (s + 1) / 55.f;
}
}
@@ -163,13 +248,16 @@ static void check_case(const weights *w, uint32_t T, unsigned route, int good_fd
if (fail_first) {
need(ds4_gpu_set_model_fd(empty_fd), "empty input fd");
need(ds4_gpu_begin_commands(), "failure batch begin");
probe_begin(w, ids, sn, good_fd, empty_fd);
const int accepted = project(w, T, 1, outputs[0].view, outputs[1].view, gx, gi, lists, counts);
need(!accepted, "uncached read must fail on empty file");
need(ds4_gpu_commands_active(), "failure preserves active batch ownership");
need(ds4_gpu_end_commands(), "failure batch drain");
(void)probe_end(0, 0);
need(ds4_gpu_set_model_fd(good_fd), "restore input fd after failure");
}
const unsigned runs = warm_probe ? 3 : 2;
read_stats reads[3] = {{0,0}, {0,0}, {0,0}};
for (unsigned run = 0; run < runs; run++) {
/* The third pass must use cache hits: the backing fd cannot provide
* any bytes. Restore the real fd before cleanup or the next case. */
@@ -177,11 +265,13 @@ static void check_case(const weights *w, uint32_t T, unsigned route, int good_fd
for (unsigned i = 0; i < 4; i++) reset(outputs + i);
need(ds4_gpu_tensor_write(outputs[3].view, 0, r, rn * 4), "reset HC input");
need(ds4_gpu_begin_commands(), "projection batch begin");
probe_begin(w, ids, sn, good_fd, empty_fd);
need(project(w, T, run != 0, outputs[0].view, outputs[1].view, gx, gi, lists, counts), "resident/SSD projection");
need(ds4_gpu_commands_active(), "projection preserves active batch ownership");
need(ds4_gpu_qwen4_moe_reduce_tensor(outputs[2].view, outputs[1].view, gw,
mm ? NULL : gsg, NULL, outputs[3].view, ginj, T, S, stride, D, HC), "reduce and HC");
need(ds4_gpu_end_commands(), "projection batch drain");
reads[run] = probe_end(run && !warm_probe, run == 0 || run == 2);
for (unsigned i = 0; i < 4; i++) {
float *got = read_output(outputs + i);
if (!run) reference[i] = got;
@@ -196,8 +286,9 @@ static void check_case(const weights *w, uint32_t T, unsigned route, int good_fd
need(read_x && read_ids && ds4_gpu_tensor_read(gx, 0, read_x, xn * 4) &&
ds4_gpu_tensor_read(gi, 0, read_ids, sn * 4), "input integrity read");
need(!memcmp(x, read_x, xn * 4) && !memcmp(ids, read_ids, sn * 4), "input and selected IDs unchanged");
printf("PASS Qwen SSD %u/%u T=%u %s route=%u: exact mid/down/reduce/HC, guards%s%s\n",
printf("PASS Qwen SSD %u/%u T=%u %s route=%u: exact mid/down/reduce/HC, guards; pread=%llu bytes/%llu calls%s%s\n",
w->gate_type, w->down_type, T, mm ? "MM" : "rows", route,
(unsigned long long)reads[1].bytes, (unsigned long long)reads[1].calls,
warm_probe ? ", warm hits without readable fd" : "",
fail_first ? ", failed read and exact retry" : "");
for (unsigned i = 0; i < 4; i++) { free(reference[i]); free_output(outputs + i); }
@@ -256,17 +347,21 @@ int main(void) {
check_case(w, T, 0, fileno(file), fileno(empty), 0, 1);
check_case(w, T, 1, fileno(file), fileno(empty), c == 0, 1);
check_case(w, T, 0, fileno(file), fileno(empty), 0, 1);
if (T == 128) check_case(w, T, 2, fileno(file), fileno(empty), 1, 0);
if (T == 128) {
check_case(w, T, 2, fileno(file), fileno(empty), 1, 0);
check_case(w, T, 3, fileno(file), fileno(empty), 0, 0);
}
}
/* A different-size MTP layer must use owned reads without replacing the
* trunk's slab class. The old trunk selection must still hit with no fd. */
ds4_gpu_set_streaming_expert_cache_budget(BUDGET);
ds4_gpu_set_streaming_expert_cache_expert_bytes(2 * formats[0].table.gate_expert_bytes + formats[0].table.down_expert_bytes);
check_case(formats, 2, 0, fileno(file), fileno(empty), 0, 1);
check_case(formats + 1, 2, 0, fileno(file), fileno(empty), 0, 0);
check_case(formats + 1, 2, 0, fileno(file), fileno(empty), 1, 0);
check_case(formats + 1, 3, 0, fileno(file), fileno(empty), 0, 0);
need(ds4_gpu_set_model_fd(fileno(empty)), "trunk retained after off-size MTP");
check_case(formats, 2, 0, fileno(file), fileno(empty), 0, 1);
ds4_gpu_cleanup(); munmap(map, bytes); fclose(empty); fclose(file);
puts("PASS Qwen SSD expert streaming: selected cache, eviction, owned layer fallback and retry");
puts("PASS Qwen SSD: cache, exact selected-only reads, full-selection fallback and retry");
return 0;
}