Metal: add bounded SSD expert streaming for Qwen3.8 Flash Next
Allow the existing --ssd-streaming mode to run the Qwen Metal graph without making every routed expert resident. Keep dense, shared, GDN, attention and hyper-connection weights mapped; retain the separate disk-only n-gram path. Reuse the explicit-pread expert cache with 512 expert IDs and top-10 routing. Dispatch the existing row, grouped-MM and NAX kernels through specialized expert-address tables, preserving their accumulation and reduction order. Keep Q2_K physical down-row padding separate from logical activations, and pass the shared gate/up and down quantization types independently. Stage complete layers in owned buffers when a batch exceeds cache capacity or the layer has a different expert size, including the Q2 pack's Q4/MXFP4 MTP layer. Never fall back to mapping routed weights for Qwen. Retain these buffers through GPU completion, including unretained command buffers. Validate router IDs, lists, tensor ranges and aliases before dispatch, and preserve command-batch ownership across read failures and retries. Size the cache after accounting for static weights, context, recurrent and speculative state, prefill workspace, staging and runtime headroom. Share the Qwen estimator between GPU and CPU-only builds, fix its large-context integer overflow, and check additional sessions and direct-generation graphs against the same budget. Reject unsupported preload/full-layer options instead of silently applying another model's policy. Add sparse-file GPU parity tests for cold/warm caches, eviction, overflow, mixed expert sizes, invalid input and failed-read recovery. Add host memory regressions and enable the real-model n-gram recovery test in SSD mode. Document cache sizing, commands and current I/O/MTP limitations. Validation: full macOS build; Qwen SSD parity with ordinary and unretained command buffers; existing Qwen resident decode/Q4 geometry tests; existing Metal SSD IQ2/Q4/MXFP4 tests; host cache and Qwen memory tests. A real Q2 Metal SSD run completed the Rome prompt with a 50-token generation limit. The real-model n-gram test also passed prefill, decode and MTP read-failure recovery with exact logits across two SSD sessions.
This commit is contained in:
GiorgioOppo committed
1 parent
9139e2ae58
commit
e37f18576c
10 files changed
+1001
-83
No files matched your search
@@ -57,6 +57,8 @@
|
||||
/tests/test_rocm_memory
|
||||
/tests/test_qwen4_kernels
|
||||
/tests/test_qwen4_moe_mm_specialize
|
||||
/tests/test_qwen4_ssd_experts
|
||||
/tests/test_qwen4_memory
|
||||
/tests/test_qwen4_vision
|
||||
/tests/test_deepseek4_vision_image
|
||||
/tests/test_deepseek41_metal
|
||||
|
||||
@@ -176,6 +176,16 @@ tests/test_qwen4_moe_mm_specialize: tests/test_qwen4_moe_mm_specialize.o $(CORE_
|
||||
test-qwen4-moe-mm-specialize: tests/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)
|
||||
$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
|
||||
|
||||
.PHONY: test-qwen4-ssd-experts
|
||||
test-qwen4-ssd-experts: tests/test_qwen4_ssd_experts
|
||||
./tests/test_qwen4_ssd_experts
|
||||
|
||||
tests/test_qwen4_conv_parallel.o: tests/test_qwen4_conv_parallel.c ds4_gpu.h
|
||||
$(CC) $(CFLAGS) -fno-fast-math -I. -c -o $@ $<
|
||||
|
||||
@@ -710,6 +720,16 @@ tests/test_qwen4_ngrams: tests/test_qwen4_ngrams.o $(filter-out ds4_cpu.o,$(CPU_
|
||||
test-qwen4-ngrams: tests/test_qwen4_ngrams
|
||||
./tests/test_qwen4_ngrams
|
||||
|
||||
tests/test_qwen4_memory.o: tests/test_qwen4_memory.c ds4.c ds4.h
|
||||
$(CC) $(filter-out -ffast-math,$(CFLAGS)) -Wno-unused-function -I. -c -o $@ $<
|
||||
|
||||
tests/test_qwen4_memory: tests/test_qwen4_memory.o $(filter-out ds4_cpu.o,$(CPU_CORE_OBJS))
|
||||
$(CC) $(filter-out -ffast-math,$(CFLAGS)) -o $@ $^ $(LDLIBS)
|
||||
|
||||
.PHONY: test-qwen4-memory
|
||||
test-qwen4-memory: tests/test_qwen4_memory
|
||||
./tests/test_qwen4_memory
|
||||
|
||||
tests/test_qwen4_ngram_state.o: tests/test_qwen4_ngram_state.c ds4.c ds4.h
|
||||
$(CC) $(filter-out -ffast-math,$(CFLAGS)) -Wno-unused-function -I. -c -o $@ $<
|
||||
|
||||
@@ -1019,6 +1039,7 @@ tests/test_session_state_gpu.o: ds4_tool_text.h
|
||||
|
||||
clean:
|
||||
rm -f tests/test_qwen4_ngrams
|
||||
rm -f tests/test_qwen4_memory tests/test_qwen4_ssd_experts
|
||||
rm -f tests/test_qwen4_ngram_state
|
||||
rm -f tests/test_web_recovery
|
||||
rm -f tests/test_metal_ssd_experts
|
||||
|
||||
@@ -32,6 +32,35 @@ Q4_K gate/up and MXFP4 down experts. Context and runtime buffers still need
|
||||
RAM beyond the resident weights; start Q2 with the context and chunk above
|
||||
on a 64 GB machine. Q4 needs a larger machine.
|
||||
|
||||
On Metal, SSD streaming keeps a bounded cache of routed experts instead of
|
||||
requiring all of them in RAM:
|
||||
|
||||
```sh
|
||||
./ds4 --ssd-streaming --ctx 4096 --prefill-chunk 128
|
||||
```
|
||||
|
||||
The Q2 and Q4 packs support streaming during prefill, decode and MTP. Dense
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
context, prefill chunk, static weights, vision weights when present, and
|
||||
staging before choosing the cache size. Reduce the context or prefill chunk
|
||||
if those fixed requirements leave insufficient memory.
|
||||
|
||||
Qwen fills its expert cache on demand, so `--ssd-streaming-cold` does not
|
||||
change its preload behavior. Popularity preloading and the GLM-specific
|
||||
`--ssd-streaming-full-layers` option are not supported for Qwen.
|
||||
|
||||
Use the same model options with `ds4-agent`, `ds4-server`, or
|
||||
`ds4-bench`. Add `--mtp` for speculative decoding using the built-in MTP
|
||||
weights. `--nothink` disables thinking. The server exposes
|
||||
@@ -113,15 +142,15 @@ the responses and diagnostics for inspection; it does not grade image content.
|
||||
Metal only for now. The Metal graph accepts Q8_0, Q4_0, F16, BF16 and F32
|
||||
dense weights, Q8_0/MXFP4/Q4_0/Q4_K/Q2_K/IQ2_XXS experts, F16/F32/Q8_0
|
||||
hyper-connection mixers and the original BF16 n-gram table.
|
||||
Tensor parallelism, pipeline execution, and SSD expert streaming are not
|
||||
implemented for this model yet. CPU code is a correctness reference, not a
|
||||
general inference backend.
|
||||
Tensor parallelism and pipeline execution are not implemented for this model
|
||||
yet. CPU code is a correctness reference, not a general inference backend.
|
||||
|
||||
|
||||
## Validation
|
||||
|
||||
```sh
|
||||
make test-qwen4-kernels test-qwen4-q2 test-qwen4-prefill-reuse test-q8-prefill-variants
|
||||
make test-qwen4-ssd-experts test-qwen4-memory
|
||||
make test-frontends
|
||||
make test-qwen4-ngrams
|
||||
make tests/test_qwen4_ngram_state
|
||||
@@ -132,8 +161,18 @@ python3 -m unittest discover -s gguf-tools/tests -p test_qwen4_native_ngrams.py
|
||||
|
||||
The kernel tests exercise the active Metal API. Vision and end-to-end model
|
||||
checks additionally require the checkpoints described above.
|
||||
The SSD test needs no GGUF: a sparse temporary file uses a 512-expert table,
|
||||
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.
|
||||
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.
|
||||
For a smaller memory footprint, run
|
||||
`tests/test_qwen4_ngram_state MODEL.gguf --ssd-streaming`. This uses a
|
||||
256-token context, an eight-token prefill chunk and a 1024-expert cache request
|
||||
while keeping both comparison sessions and MTP enabled.
|
||||
|
||||
[Checkpoint-fix benchmark charts and measurements](../speed-bench/qwen38-checkpoints/README.md)
|
||||
compare prefill, ordinary decode, and MTP decode against the preceding PR head.
|
||||
|
||||
@@ -1023,6 +1023,63 @@ static bool ds4_qwen4_layer_is_nextn(uint32_t il) {
|
||||
il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER;
|
||||
}
|
||||
|
||||
static uint32_t qwen4_prefill_chunk_tokens(uint32_t ctx) {
|
||||
const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK");
|
||||
const unsigned long v = env && env[0] ? strtoul(env, NULL, 10) : 8192ul;
|
||||
uint32_t chunk = v == 0 || v > 65536ul ? 8192u : (uint32_t)v;
|
||||
return chunk > ctx ? ctx : chunk;
|
||||
}
|
||||
|
||||
static ds4_context_memory qwen4_graph_memory_estimate(uint32_t ctx, uint32_t prefill_chunk) {
|
||||
ds4_context_memory m = {0};
|
||||
/* Match the bounded graph workspace, including CPU staging. This API
|
||||
* has no MTP flag, so an embedded predictor reserves its primary and
|
||||
* both lazy verifier snapshots as well as the live recurrent state. */
|
||||
uint64_t T = prefill_chunk ? prefill_chunk : qwen4_prefill_chunk_tokens(ctx);
|
||||
if (T > ctx) T = ctx;
|
||||
const uint64_t blocks = ctx / 4u + 1u, E = DS4_N_EMBD, hc_dim = E * DS4_N_HC;
|
||||
const uint64_t kv_row = 2ull * DS4_N_HEAD_KV * DS4_N_HEAD_DIM * 2u;
|
||||
const uint64_t conv_dim = DS4_N_LIN_CONV_DIM;
|
||||
const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM;
|
||||
const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM;
|
||||
const uint64_t kv_dim = (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM;
|
||||
const uint64_t iq_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM;
|
||||
const uint64_t k_blocks = DS4_N_INDEXER_TOP_K / 4u;
|
||||
const uint64_t hc = DS4_N_HC, F = DS4_N_FF_EXP, used = DS4_N_EXPERT_USED;
|
||||
const bool mtp = DS4_N_NEXTN_PREDICT != 0;
|
||||
uint32_t n_attn = 0, n_linear = 0;
|
||||
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
|
||||
if (ds4_qwen4_layer_is_linear(il)) n_linear++;
|
||||
else n_attn++;
|
||||
}
|
||||
m.prefill_cap = (uint32_t)T;
|
||||
m.raw_cap = ctx;
|
||||
m.comp_cap = (uint32_t)blocks;
|
||||
m.raw_bytes = (uint64_t)n_attn * ctx * kv_row + (uint64_t)ctx * 16u;
|
||||
m.compressed_bytes = (uint64_t)n_attn * ((uint64_t)ctx * DS4_N_INDEXER_HEAD_DIM * 4u + blocks * DS4_N_INDEXER_HEAD_DIM * 2u);
|
||||
const uint64_t old_row = 12u * hc_dim + 24u * E + (used + 1u) * (E + 2u * F) + 2u * blocks;
|
||||
/* qwen4_graph_alloc: eight HC chunks; decode attention has 64 splits.
|
||||
* Keep these constants available to the CPU-only memory estimator. */
|
||||
const uint64_t gpu_row = 6u * hc_dim + 2u * DS4_N_HC_LOWRANK + 16u * hc * hc +
|
||||
5u * E + conv_dim + 2u * v_dim + 2u * DS4_N_LIN_V_HEAD + 5u * q_dim +
|
||||
2u * kv_dim + 2u * iq_dim + DS4_N_INDEXER_HEAD_DIM + blocks + blocks / 8u +
|
||||
5u * k_blocks + 7u + 2u * DS4_N_EXPERT + 2u * used + (used + 1u) * (F + E) + 3u * F;
|
||||
const uint64_t state = (uint64_t)n_linear *
|
||||
(v_dim * DS4_N_LIN_HEAD_DIM + (DS4_N_LIN_CONV - 1u) * conv_dim) +
|
||||
(uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * hc_dim;
|
||||
const uint64_t attn_part = 3u * (uint64_t)DS4_N_HEAD * 64u * (2u + DS4_N_HEAD_DIM);
|
||||
/* GPU + host logits, session logits/probabilities, and four verify
|
||||
* rows. Predictor vocabulary subsets and optional steering are extra
|
||||
* model allocations, not graph workspace. */
|
||||
uint64_t fixed = (mtp ? 4u : 1u) * state + attn_part + DS4_N_EXPERT +
|
||||
(uint64_t)(mtp ? 12u : 4u) * DS4_N_VOCAB;
|
||||
if (mtp) fixed += 12u * (hc + 1u) * E + 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;
|
||||
m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes;
|
||||
return m;
|
||||
}
|
||||
|
||||
static int g_ds4_lock_fd = -1;
|
||||
|
||||
#if defined(__GNUC__) || defined(__clang__)
|
||||
@@ -39623,12 +39680,6 @@ static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap(
|
||||
#ifdef DS4_HAS_DEEPSEEK41_GPU
|
||||
static ds4_context_memory ds41_graph_memory(uint32_t ctx);
|
||||
#endif
|
||||
static uint32_t qwen4_prefill_chunk_tokens(uint32_t ctx) {
|
||||
const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK");
|
||||
const unsigned long v = env && env[0] ? strtoul(env, NULL, 10) : 8192ul;
|
||||
uint32_t chunk = v == 0 || v > 65536ul ? 8192u : (uint32_t)v;
|
||||
return chunk > ctx ? ctx : chunk;
|
||||
}
|
||||
|
||||
ds4_context_memory ds4_context_memory_estimate_with_prefill_mode(
|
||||
ds4_backend backend,
|
||||
@@ -39637,24 +39688,8 @@ ds4_context_memory ds4_context_memory_estimate_with_prefill_mode(
|
||||
bool ssd_streaming) {
|
||||
ds4_context_memory m = {0};
|
||||
uint32_t ctx = ctx_size > 0 ? (uint32_t)ctx_size : 1u;
|
||||
if (ds4_backend_uses_graph(backend) && ds4_model_is_qwen4()) {
|
||||
/* per-token f16 K/V and raw indexer keys per attention layer, pooled
|
||||
* block keys, rope positions; transients scale with the prefill chunk */
|
||||
const uint64_t T = prefill_chunk ? prefill_chunk : qwen4_prefill_chunk_tokens(ctx);
|
||||
const uint64_t blocks = ctx / 4u + 1u, E = DS4_N_EMBD, hc_dim = E * DS4_N_HC;
|
||||
const uint64_t kv_row = 2ull * DS4_N_HEAD_KV * DS4_N_HEAD_DIM * 2u;
|
||||
uint32_t n_attn = 0;
|
||||
for (uint32_t il = 0; il < DS4_N_LAYER; il++) n_attn += !ds4_qwen4_layer_is_linear(il);
|
||||
m.prefill_cap = (uint32_t)T;
|
||||
m.raw_cap = ctx;
|
||||
m.comp_cap = (uint32_t)blocks;
|
||||
m.raw_bytes = (uint64_t)n_attn * ctx * kv_row + (uint64_t)ctx * 16u;
|
||||
m.compressed_bytes = (uint64_t)n_attn * (ctx * DS4_N_INDEXER_HEAD_DIM * 4u + blocks * DS4_N_INDEXER_HEAD_DIM * 2u);
|
||||
m.scratch_bytes = T * (12u * hc_dim + 24u * E + (uint64_t)(DS4_N_EXPERT_USED + 1u) * (E + 2u * DS4_N_FF_EXP) +
|
||||
2u * blocks) * 4u;
|
||||
m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes;
|
||||
return m;
|
||||
}
|
||||
if (ds4_backend_uses_graph(backend) && ds4_model_is_qwen4())
|
||||
return qwen4_graph_memory_estimate(ctx, prefill_chunk);
|
||||
|
||||
if (ds4_backend_uses_graph(backend)) {
|
||||
#ifdef DS4_HAS_DEEPSEEK41_GPU
|
||||
@@ -42217,6 +42252,7 @@ typedef enum {
|
||||
struct ds4_engine {
|
||||
char *model_path;
|
||||
uint64_t ds41_session_bytes;
|
||||
uint64_t qwen4_session_bytes;
|
||||
ds4_model model;
|
||||
ds4_model mtp_model;
|
||||
ds4_model vision_model;
|
||||
@@ -57485,6 +57521,7 @@ typedef struct {
|
||||
uint32_t k_blocks;
|
||||
uint32_t n_block_cap;
|
||||
uint32_t sel_stride;
|
||||
bool ssd_streaming;
|
||||
ds4_gpu_tensor *R, *xn, *lo, *inj, *mixed, *blk;
|
||||
ds4_gpu_tensor *qkv, *z, *ga, *gb, *lin_o;
|
||||
ds4_gpu_tensor *ple_emb, *ple_key, *ple_val, *ple_gated, *ple_normed, *ple_hist;
|
||||
@@ -58273,8 +58310,9 @@ static bool qwen4_moe_profile_boundary(bool enabled, double *last, double *elaps
|
||||
return glm_graph_begin_commands_if_needed();
|
||||
}
|
||||
|
||||
static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t T) {
|
||||
const bool profile = T > 8u && getenv("DS4_QWEN4_MOE_PROFILE") != NULL;
|
||||
static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
|
||||
uint32_t il, uint32_t T) {
|
||||
const bool profile = !g->ssd_streaming && T > 8u && getenv("DS4_QWEN4_MOE_PROFILE") != NULL;
|
||||
double elapsed[7] = {0}, last = 0;
|
||||
if (!qwen4_moe_profile_boundary(profile, &last, &elapsed[0])) return false;
|
||||
bool ok = qwen4_gemv(g->router, m, l->ffn_gate_inp, g->mixed, T) &&
|
||||
@@ -58292,6 +58330,43 @@ static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds
|
||||
qwen4_expert_type_has_mm(l->ffn_down_exps->type) &&
|
||||
qwen4_graph_dense_ok(l->ffn_gate_shexp) && qwen4_graph_dense_ok(l->ffn_up_shexp) &&
|
||||
qwen4_graph_dense_ok(l->ffn_down_shexp);
|
||||
if (g->ssd_streaming) {
|
||||
uint64_t gate_bytes = 0, down_bytes = 0;
|
||||
if (!streaming_layer_gate_down_expert_bytes(l, &gate_bytes, &down_bytes) ||
|
||||
l->ffn_gate_exps->type != l->ffn_up_exps->type ||
|
||||
l->ffn_gate_exps->bytes != l->ffn_up_exps->bytes) return false;
|
||||
const ds4_gpu_stream_expert_table table = {
|
||||
.model_map = m->map, .model_size = m->size, .layer = il,
|
||||
.n_total_expert = DS4_N_EXPERT,
|
||||
.gate_offset = l->ffn_gate_exps->abs_offset,
|
||||
.up_offset = l->ffn_up_exps->abs_offset,
|
||||
.down_offset = l->ffn_down_exps->abs_offset,
|
||||
.gate_expert_bytes = gate_bytes, .down_expert_bytes = down_bytes,
|
||||
};
|
||||
/* Keep prefill's expert lists and shared GEMMs, and decode's shared
|
||||
* slot, exactly as in the resident path. Only routed weight storage
|
||||
* changes; the final slot reduction and HC combine remain unchanged. */
|
||||
if (ok && mm) ok = ds4_gpu_qwen4_moe_build_lists_tensor(
|
||||
g->moe_lists, g->moe_counts, g->selected, T, DS4_N_EXPERT_USED,
|
||||
DS4_N_EXPERT, g->cap_tokens) != 0;
|
||||
if (ok) ok = ds4_gpu_qwen4_moe_stream_tensor(
|
||||
g->mid, g->part, g->mixed, g->selected,
|
||||
mm ? g->moe_lists : NULL, mm ? g->moe_counts : NULL, &table,
|
||||
l->ffn_gate_exps->type, l->ffn_down_exps->type, T, DS4_N_EXPERT_USED,
|
||||
DS4_N_EMBD, DS4_N_FF_EXP, g->cap_tokens,
|
||||
l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset,
|
||||
l->ffn_down_shexp->abs_offset,
|
||||
mm ? UINT32_MAX : l->ffn_gate_shexp->type,
|
||||
mm ? UINT32_MAX : l->ffn_down_shexp->type) != 0;
|
||||
if (ok && mm) ok = qwen4_gemv(g->sh_gate, m, l->ffn_gate_shexp, g->mixed, T) &&
|
||||
qwen4_gemv(g->sh_up, m, l->ffn_up_shexp, g->mixed, T) &&
|
||||
ds4_gpu_swiglu_tensor(g->sh_mid, g->sh_gate, g->sh_up, T * DS4_N_FF_EXP, 0.0f, 1.0f) &&
|
||||
qwen4_gemv(g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T);
|
||||
return ok && ds4_gpu_qwen4_moe_reduce_tensor(
|
||||
g->blk, g->part, g->weights, g->sh_gate_logit, mm ? g->sh_out : NULL,
|
||||
g->R, g->inj, T, DS4_N_EXPERT_USED, DS4_N_EXPERT_USED + (mm ? 0u : 1u),
|
||||
DS4_N_EMBD, DS4_N_HC) != 0;
|
||||
}
|
||||
if (mm) {
|
||||
if (ok) {
|
||||
ok = ds4_gpu_qwen4_moe_build_lists_tensor(g->moe_lists, g->moe_counts, g->selected, T, DS4_N_EXPERT_USED,
|
||||
@@ -58492,7 +58567,7 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
|
||||
if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T);
|
||||
}
|
||||
QWEN4_PROF(4);
|
||||
if (ok) ok = qwen4_graph_moe(g, m, l, T); /* the reduce folds the combine in */
|
||||
if (ok) ok = qwen4_graph_moe(g, m, l, il, T); /* the reduce folds the combine in */
|
||||
if (ok && g->dump_prompt_rows) qwen4_graph_dump_last_ffn(g, il, T);
|
||||
if (ok) ok = qwen4_graph_apply_steering_ffn(g, il, T);
|
||||
QWEN4_PROF(5);
|
||||
@@ -58764,7 +58839,7 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
|
||||
if (ok) ok = qwen4_graph_attention(g, m, l, il, idx, T);
|
||||
if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, E, hc) != 0;
|
||||
if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T);
|
||||
if (ok) ok = qwen4_graph_moe(g, m, l, T);
|
||||
if (ok) ok = qwen4_graph_moe(g, m, l, il, T);
|
||||
ds4_gpu_tensor *last = NULL;
|
||||
const char *argmax_env = getenv("DS4_QWEN4_MTP_GPU_ARGMAX");
|
||||
const bool gpu_argmax = want_logits && draft_out && !logits_out &&
|
||||
@@ -58854,7 +58929,7 @@ static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *
|
||||
if (ok) ok = qwen4_graph_attention(g, m, l, il, idx, 1);
|
||||
if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, 1, E, hc) != 0;
|
||||
if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, 1);
|
||||
if (ok) ok = qwen4_graph_moe(g, m, l, 1);
|
||||
if (ok) ok = qwen4_graph_moe(g, m, l, il, 1);
|
||||
/* the chained draft only needs its argmax; DS4_QWEN4_MTP_DRAFT_ROWS can
|
||||
* score the frequent-token prefix exactly like the first draft's head */
|
||||
const uint32_t chain_head_rows = qwen4_mtp_draft_rows();
|
||||
@@ -58894,6 +58969,8 @@ static int generate_qwen4_metal_argmax(
|
||||
const token_vec * prompt,
|
||||
int n_predict,
|
||||
int ctx_size,
|
||||
bool ssd_streaming,
|
||||
uint32_t prefill_chunk,
|
||||
const char * directional_steering_file,
|
||||
float directional_steering_attn,
|
||||
float directional_steering_ffn,
|
||||
@@ -58907,10 +58984,13 @@ static int generate_qwen4_metal_argmax(
|
||||
return 1;
|
||||
}
|
||||
ds4_qwen4_gpu_graph *g = xcalloc(1, sizeof(*g));
|
||||
if (!qwen4_graph_alloc(g, weights, (uint32_t)ctx_size, qwen4_prefill_chunk_tokens((uint32_t)ctx_size), false)) {
|
||||
const uint32_t cap = prefill_chunk && prefill_chunk < (uint32_t)ctx_size ?
|
||||
prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size);
|
||||
if (!qwen4_graph_alloc(g, weights, (uint32_t)ctx_size, cap, false)) {
|
||||
free(g);
|
||||
return 1;
|
||||
}
|
||||
g->ssd_streaming = ssd_streaming;
|
||||
if (!qwen4_graph_load_steering(g, directional_steering_file,
|
||||
directional_steering_attn, directional_steering_ffn)) {
|
||||
qwen4_graph_free(g);
|
||||
@@ -59003,7 +59083,8 @@ static int generate_metal_graph_raw_swa(
|
||||
}
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
if (ds4_model_is_qwen4()) {
|
||||
return generate_qwen4_metal_argmax(model, vocab, weights, prompt, n_predict, ctx_size,
|
||||
return generate_qwen4_metal_argmax(model, vocab, weights, prompt, n_predict, ctx_size, ssd_streaming,
|
||||
prefill_chunk,
|
||||
directional_steering_file,
|
||||
directional_steering_attn, directional_steering_ffn,
|
||||
emit, done, emit_ud, progress, progress_ud);
|
||||
@@ -59176,6 +59257,8 @@ ds4_context_memory ds4_context_memory_estimate_with_prefill_mode(
|
||||
(void)ssd_streaming;
|
||||
ds4_context_memory m = {0};
|
||||
uint32_t ctx = ctx_size > 0 ? (uint32_t)ctx_size : 1u;
|
||||
if (ds4_backend_uses_graph(backend) && ds4_model_is_qwen4())
|
||||
return qwen4_graph_memory_estimate(ctx, prefill_chunk);
|
||||
|
||||
m.raw_cap = ds4_default_raw_cap(ctx);
|
||||
m.raw_bytes = (uint64_t)DS4_N_LAYER *
|
||||
@@ -59900,6 +59983,7 @@ struct ds4_session {
|
||||
bool glm_graph_ready;
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
ds4_qwen4_gpu_graph qwen4_graph;
|
||||
uint64_t qwen4_reserved_graph_bytes;
|
||||
bool qwen4_graph_ready;
|
||||
bool qwen4_rewound;
|
||||
float *qwen4_verify_logits;
|
||||
@@ -64708,6 +64792,10 @@ static int ds4_session_eval_splitkv_spec_after_first(
|
||||
* 4. fall back to ordinary one-token decode if the fast verifier cannot prove
|
||||
* the target stream. */
|
||||
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
static bool qwen4_streaming_memory_admit(ds4_engine *e, uint64_t graph_bytes, bool fit_cache);
|
||||
#endif
|
||||
|
||||
int ds4_engine_generate_argmax(
|
||||
ds4_engine *e,
|
||||
const ds4_tokens *prompt,
|
||||
@@ -64785,7 +64873,23 @@ int ds4_engine_generate_argmax(
|
||||
ds4_session_free(s);
|
||||
return rc;
|
||||
}
|
||||
return generate_metal_graph_raw_swa(model, vocab, weights, prompt,
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
uint64_t qwen4_reserved = 0;
|
||||
if (ds4_model_is_qwen4() && e->ssd_streaming) {
|
||||
if (ctx_size <= 0) return 1;
|
||||
const uint32_t cap = e->prefill_chunk && e->prefill_chunk < (uint32_t)ctx_size ?
|
||||
e->prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size);
|
||||
const ds4_context_memory memory = ds4_context_memory_estimate_with_prefill_mode(
|
||||
e->backend, ctx_size, cap, true);
|
||||
if (!qwen4_streaming_memory_admit(e,
|
||||
ds4_add_sat_u64(e->qwen4_session_bytes, memory.total_bytes), false)) return 1;
|
||||
/* This API owns a temporary graph; reserve it across callbacks
|
||||
* as well, so sessions created there include the live graph. */
|
||||
qwen4_reserved = memory.total_bytes;
|
||||
e->qwen4_session_bytes += qwen4_reserved;
|
||||
}
|
||||
#endif
|
||||
const int result = generate_metal_graph_raw_swa(model, vocab, weights, prompt,
|
||||
n_predict, ctx_size, e->quality,
|
||||
e->ssd_streaming,
|
||||
e->ssd_streaming_cold,
|
||||
@@ -64799,6 +64903,10 @@ int ds4_engine_generate_argmax(
|
||||
e->directional_steering_ffn_scale,
|
||||
emit, done, emit_ud,
|
||||
progress, progress_ud);
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
e->qwen4_session_bytes -= qwen4_reserved;
|
||||
#endif
|
||||
return result;
|
||||
#else
|
||||
fprintf(stderr, "ds4: %s generation requested but this build has no graph backend support\n",
|
||||
ds4_backend_name(e->backend));
|
||||
@@ -70002,6 +70110,77 @@ static bool engine_warm_full_model(const ds4_engine_options *opt) {
|
||||
opt->distributed.role == DS4_DISTRIBUTED_NONE;
|
||||
}
|
||||
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
/* Qwen keeps the dense/GDN/HC weights mapped and stages routed weights in
|
||||
* owned buffers. Reserve the largest active layer, including the differently
|
||||
* quantized MTP layer when enabled, before sizing the persistent expert cache.
|
||||
* Two windows cover a previous in-flight layer and its successor. */
|
||||
static bool qwen4_streaming_memory_admit(ds4_engine *e, uint64_t graph_bytes, bool fit_cache) {
|
||||
if (!e->ssd_streaming) return true;
|
||||
const uint64_t gib = UINT64_C(1073741824);
|
||||
const uint64_t host = glm_graph_host_memory_bytes();
|
||||
const uint64_t recommended = ds4_gpu_recommended_working_set_size();
|
||||
if (!host || !recommended) {
|
||||
fprintf(stderr, "ds4: cannot determine Qwen SSD memory budget\n");
|
||||
return false;
|
||||
}
|
||||
uint64_t budget = host / 8u * 7u;
|
||||
if (recommended / 8u * 7u < budget) budget = recommended / 8u * 7u;
|
||||
uint64_t statics = 0, expert_bytes = 0, max_experts = 0, largest_layer = 0;
|
||||
if (!weights_streaming_non_routed_bytes(&e->weights, &statics) ||
|
||||
!ds4_streaming_routed_expert_bytes(&e->weights, &expert_bytes) ||
|
||||
!ds4_streaming_cacheable_expert_count(&e->weights, &max_experts, NULL)) return false;
|
||||
const uint32_t layers = DS4_N_LAYER - (e->glm_mtp ? 0u : DS4_N_NEXTN_PREDICT);
|
||||
for (uint32_t il = 0; il < layers; il++) {
|
||||
uint64_t bytes = 0;
|
||||
if (!streaming_layer_routed_expert_bytes(&e->weights.layer[il], &bytes)) return false;
|
||||
bytes = ds4_mul_sat_u64(bytes, DS4_N_EXPERT);
|
||||
if (bytes > largest_layer) largest_layer = bytes;
|
||||
}
|
||||
const uint64_t staging = ds4_mul_sat_u64(largest_layer, 2u);
|
||||
const uint64_t fixed = ds4_add_sat_u64(ds4_add_sat_u64(statics, e->vision_model.size),
|
||||
ds4_add_sat_u64(graph_bytes, ds4_add_sat_u64(staging, 2u * gib)));
|
||||
if (fixed >= budget || budget - fixed < expert_bytes) {
|
||||
fprintf(stderr, "ds4: Qwen SSD needs %.2f GiB before the expert cache; safe budget %.2f GiB. "
|
||||
"Reduce context or prefill chunk.\n",
|
||||
ds4_bytes_to_gib(fixed), ds4_bytes_to_gib(budget));
|
||||
return false;
|
||||
}
|
||||
uint64_t count = (budget - fixed) / expert_bytes;
|
||||
if (count > max_experts) count = max_experts;
|
||||
if (!fit_cache) {
|
||||
if (e->ssd_streaming_cache_experts > count) {
|
||||
fprintf(stderr, "ds4: no room for another Qwen SSD session; reduce context or expert cache\n");
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
uint64_t requested = e->ssd_streaming_cache_experts;
|
||||
if (e->ssd_streaming_cache_bytes) {
|
||||
if (e->ssd_streaming_cache_bytes <= staging ||
|
||||
e->ssd_streaming_cache_bytes - staging < expert_bytes) {
|
||||
fprintf(stderr, "ds4: Qwen SSD cache target must exceed %.2f GiB for prefill staging\n",
|
||||
ds4_bytes_to_gib(staging + expert_bytes));
|
||||
return false;
|
||||
}
|
||||
requested = (e->ssd_streaming_cache_bytes - staging) / expert_bytes;
|
||||
}
|
||||
if (requested && requested < count) count = requested;
|
||||
if (requested > count) fprintf(stderr, "ds4: Qwen SSD cache capped from %llu to %llu experts\n",
|
||||
(unsigned long long)requested, (unsigned long long)count);
|
||||
e->ssd_streaming_cache_experts = (uint32_t)count;
|
||||
e->ssd_streaming_cache_bytes = count * expert_bytes;
|
||||
e->ssd_streaming_prefill_headroom_bytes = staging;
|
||||
e->ssd_streaming_static_decode_map = true;
|
||||
e->ssd_streaming_decode_map_bytes = statics;
|
||||
fprintf(stderr, "ds4: Qwen SSD: %.2f GiB static weights, %.2f GiB graph reserve, "
|
||||
"%.2f GiB staging, %u cached experts (%.2f GiB)\n",
|
||||
ds4_bytes_to_gib(statics), ds4_bytes_to_gib(graph_bytes), ds4_bytes_to_gib(staging),
|
||||
e->ssd_streaming_cache_experts, ds4_bytes_to_gib(e->ssd_streaming_cache_bytes));
|
||||
return true;
|
||||
}
|
||||
#endif
|
||||
|
||||
#ifdef DS4_HAS_DEEPSEEK41_GPU
|
||||
static bool ds41_memory_admit_for_host(ds4_engine *e, uint64_t graph_bytes,
|
||||
bool fit_cache, uint64_t host,
|
||||
@@ -70203,6 +70382,13 @@ static int ds4_engine_open_internal(ds4_engine **out,
|
||||
}
|
||||
config_validate_model(&e->model);
|
||||
if (ds4_model_is_qwen4() && !opt->inspect_only) {
|
||||
if (e->ssd_streaming && (e->ssd_streaming_full_layers || e->ssd_streaming_preload_experts)) {
|
||||
fprintf(stderr, "ds4: Qwen SSD fills its expert cache on demand; full resident layers "
|
||||
"and popularity preload are not supported\n");
|
||||
ds4_engine_close(e);
|
||||
*out = NULL;
|
||||
return 1;
|
||||
}
|
||||
const bool backend_ok =
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
e->backend == DS4_BACKEND_METAL ||
|
||||
@@ -70210,10 +70396,10 @@ static int ds4_engine_open_internal(ds4_engine **out,
|
||||
(opt->first_token_test && e->backend == DS4_BACKEND_CPU);
|
||||
if (!backend_ok || opt->tp.role != DS4_TP_NONE || opt->cuda_tensor_parallel ||
|
||||
opt->distributed.role != DS4_DISTRIBUTED_NONE || load_slice ||
|
||||
e->ssd_streaming || opt->dspark || e->power_percent != 100 ||
|
||||
(e->ssd_streaming && opt->first_token_test) || opt->dspark || e->power_percent != 100 ||
|
||||
(opt->mtp_path && opt->mtp_path[0])) {
|
||||
fprintf(stderr, "ds4: Qwen3.8 requires single-host Metal (or --cpu --first-token-test); "
|
||||
"tensor parallelism, pipeline execution, SSD streaming, DSpark, "
|
||||
"tensor parallelism, pipeline execution, SSD first-token tests, DSpark, "
|
||||
"external MTP models and power throttling are not supported\n");
|
||||
ds4_engine_close(e);
|
||||
*out = NULL;
|
||||
@@ -70869,12 +71055,21 @@ static int ds4_engine_open_internal(ds4_engine **out,
|
||||
ds4_gpu_set_quality(e->quality);
|
||||
ds4_gpu_set_glm_model(DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA);
|
||||
ds4_gpu_set_ssd_streaming(e->ssd_streaming);
|
||||
if (!ds4_engine_configure_streaming_auto_cache(e, opt->context_size)) {
|
||||
ds4_engine_close(e);
|
||||
*out = NULL;
|
||||
return 1;
|
||||
}
|
||||
if (!ds4_engine_configure_streaming_cache_budget(e)) {
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
if (ds4_model_is_qwen4() && e->ssd_streaming) {
|
||||
const ds4_context_memory memory = ds4_context_memory_estimate_with_prefill_mode(
|
||||
e->backend, opt->context_size > 0 ? opt->context_size : 4096, e->prefill_chunk, true);
|
||||
const uint32_t sessions = e->placement_session_count_hint > 0 ?
|
||||
(uint32_t)e->placement_session_count_hint : 1u;
|
||||
if (!qwen4_streaming_memory_admit(e, ds4_mul_sat_u64(memory.total_bytes, sessions), true)) {
|
||||
ds4_engine_close(e);
|
||||
*out = NULL;
|
||||
return 1;
|
||||
}
|
||||
} else
|
||||
#endif
|
||||
if (!ds4_engine_configure_streaming_auto_cache(e, opt->context_size) ||
|
||||
!ds4_engine_configure_streaming_cache_budget(e)) {
|
||||
ds4_engine_close(e);
|
||||
*out = NULL;
|
||||
return 1;
|
||||
@@ -70982,7 +71177,7 @@ static int ds4_engine_open_internal(ds4_engine **out,
|
||||
routed++;
|
||||
if (!weights_streaming_layer_experts_uniform(&e->weights, il)) boosted++;
|
||||
}
|
||||
if (boosted > 0) {
|
||||
if (boosted > 0 && !ds4_model_is_qwen4()) {
|
||||
fprintf(stderr,
|
||||
"ds4: SSD streaming mixed-precision model: %u/%u routed layers "
|
||||
"off the slab size class will bypass the expert cache and read "
|
||||
@@ -71046,7 +71241,7 @@ static int ds4_engine_open_internal(ds4_engine **out,
|
||||
weights_have_output_head(&e->weights)));
|
||||
ds4_model_map_span_vec spans;
|
||||
bool spans_ok = false;
|
||||
if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) {
|
||||
if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 || ds4_model_is_qwen4()) {
|
||||
spans_ok = weights_model_map_decode_static_spans(&e->weights, true, true, &spans);
|
||||
} else if (load_slice) {
|
||||
if (e->ssd_streaming_static_decode_map) {
|
||||
@@ -72481,10 +72676,18 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) {
|
||||
}
|
||||
const uint32_t cap_tokens = e->prefill_chunk && e->prefill_chunk < (uint32_t)ctx_size ?
|
||||
e->prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size);
|
||||
const ds4_context_memory memory = ds4_context_memory_estimate_with_prefill_mode(
|
||||
e->backend, ctx_size, cap_tokens, e->ssd_streaming);
|
||||
if (!qwen4_streaming_memory_admit(e,
|
||||
ds4_add_sat_u64(e->qwen4_session_bytes, memory.total_bytes), false)) {
|
||||
free(s);
|
||||
return 1;
|
||||
}
|
||||
if (!qwen4_graph_alloc(&s->qwen4_graph, &e->weights, (uint32_t)ctx_size, cap_tokens, e->glm_mtp)) {
|
||||
free(s);
|
||||
return 1;
|
||||
}
|
||||
s->qwen4_graph.ssd_streaming = e->ssd_streaming;
|
||||
qwen4_graph_reset(&s->qwen4_graph);
|
||||
if (!qwen4_graph_load_steering(&s->qwen4_graph,
|
||||
e->directional_steering_file,
|
||||
@@ -72495,6 +72698,8 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) {
|
||||
return 1;
|
||||
}
|
||||
s->qwen4_graph_ready = true;
|
||||
s->qwen4_reserved_graph_bytes = memory.total_bytes;
|
||||
e->qwen4_session_bytes += memory.total_bytes;
|
||||
s->prefill_cap = (uint32_t)ctx_size;
|
||||
s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->logits[0]));
|
||||
s->sample_probs = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->sample_probs[0]));
|
||||
@@ -72850,6 +73055,7 @@ void ds4_session_free(ds4_session *s) {
|
||||
#endif
|
||||
#ifdef DS4_HAS_QWEN4_METAL
|
||||
if (ds4_session_is_qwen4(s)) {
|
||||
s->engine->qwen4_session_bytes -= s->qwen4_reserved_graph_bytes;
|
||||
if (s->engine && s->engine->glm_mtp_timing && s->qwen4_spec_cycles) {
|
||||
fprintf(stderr, "ds4: Qwen3.8 mtp: %" PRIu64 " verify cycles, %" PRIu64 " drafts accepted (%.1f%%)\n",
|
||||
s->qwen4_spec_cycles, s->qwen4_spec_accepted,
|
||||
|
||||
@@ -3419,6 +3419,24 @@ 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);
|
||||
/* Metal SSD routed experts. Table sizes are the stored bytes per expert,
|
||||
* including Q2_K down-row padding; ff_dim remains the logical activation width.
|
||||
* NULL lists/counts use the ordinary row kernels, including the shared slot.
|
||||
* Non-NULL lists/counts use the ordinary expert-grouped MM kernels with routed
|
||||
* slots only; the caller computes the shared expert and performs the reduction.
|
||||
* The call completes the router before reading IDs and preserves command-batch
|
||||
* ownership. Cache overflow and off-size layers use explicit owned SSD reads,
|
||||
* never mapped routed weights. Returns zero on invalid input or I/O failure. */
|
||||
#ifdef __APPLE__
|
||||
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,
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
uint32_t gate_type, uint32_t down_type, uint32_t n_tokens, uint32_t n_slots,
|
||||
uint32_t in_dim, uint32_t ff_dim, uint32_t list_cap,
|
||||
uint64_t shared_gate_offset, uint64_t shared_up_offset, uint64_t shared_down_offset,
|
||||
uint32_t shared_type, uint32_t shared_down_type);
|
||||
#endif
|
||||
/* weight_type covers both the alpha and the beta projection */
|
||||
int ds4_gpu_qwen4_gdn_front_tensor(
|
||||
ds4_gpu_tensor *qkv, ds4_gpu_tensor *state, const ds4_gpu_tensor *mixed,
|
||||
|
||||
+269
-17
@@ -939,8 +939,9 @@ static uint32_t g_model_view_count;
|
||||
|
||||
enum {
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER = 80,
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384,
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = DS4_METAL_MAX_ROUTED_EXPERT_USED,
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 512,
|
||||
/* Qwen routes ten experts; other models retain their kernel-specific limits. */
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = 16,
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES =
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER *
|
||||
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT,
|
||||
@@ -17367,7 +17368,8 @@ static int ds4_gpu_stream_expert_cache_prepare_selected_batch(
|
||||
uint32_t *unique_out,
|
||||
id<MTLBuffer> *overflow_gate,
|
||||
id<MTLBuffer> *overflow_up,
|
||||
id<MTLBuffer> *overflow_down) {
|
||||
id<MTLBuffer> *overflow_down,
|
||||
bool allow_mapped_overflow) {
|
||||
if (overflow_gate) *overflow_gate = nil;
|
||||
if (overflow_up) *overflow_up = nil;
|
||||
if (overflow_down) *overflow_down = nil;
|
||||
@@ -17684,6 +17686,7 @@ static int ds4_gpu_stream_expert_cache_prepare_selected_batch(
|
||||
ds4_gpu_stream_expert_cache_entry *entry = unique_entries[u];
|
||||
const uint32_t expert = (uint32_t)unique_ids[u];
|
||||
if (!entry) {
|
||||
if (!allow_mapped_overflow) { ok = 0; break; }
|
||||
const uint64_t gate_rel = unique_gate_offsets[u] - gate_offset;
|
||||
const uint64_t down_rel = unique_down_offsets[u] - down_offset;
|
||||
if (!*overflow_gate) {
|
||||
@@ -39699,7 +39702,7 @@ static int ds4_gpu_glm_routed_moe_batch_tensor_impl(
|
||||
&stream_unique,
|
||||
&stream_overflow_gate,
|
||||
&stream_overflow_up,
|
||||
&stream_overflow_down)) {
|
||||
&stream_overflow_down, true)) {
|
||||
return 0;
|
||||
}
|
||||
if (stream_unique == 0) {
|
||||
@@ -43733,7 +43736,7 @@ int ds4_gpu_routed_moe_batch_tensor(
|
||||
&stream_unique,
|
||||
&stream_overflow_gate,
|
||||
&stream_overflow_up,
|
||||
&stream_overflow_down)) {
|
||||
&stream_overflow_down, true)) {
|
||||
g_stream_prefill_batch_selected_addr_building--;
|
||||
return 0;
|
||||
}
|
||||
@@ -48081,6 +48084,34 @@ static bool qwen4_bind_weight(qwen4_bind *b, const void *map, uint64_t size,
|
||||
return b->buf != nil;
|
||||
}
|
||||
|
||||
/* 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. */
|
||||
typedef struct {
|
||||
const ds4_gpu_stream_expert_table *table;
|
||||
__strong id<MTLBuffer> gate, up, down;
|
||||
ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
|
||||
uint32_t n_entries;
|
||||
bool addresses;
|
||||
} qwen4_stream_weights;
|
||||
static qwen4_stream_weights *g_qwen4_stream_weights;
|
||||
|
||||
static bool qwen4_bind_experts(qwen4_bind *b, const void *map, uint64_t size,
|
||||
uint64_t offset, uint64_t bytes, const char *what) {
|
||||
const qwen4_stream_weights *s = g_qwen4_stream_weights;
|
||||
if (!s) return qwen4_bind_weight(b, map, size, offset, bytes, what);
|
||||
const ds4_gpu_stream_expert_table *t = s->table;
|
||||
if (map != t->model_map || size != t->model_size) return false;
|
||||
const uint64_t gate_bytes = t->gate_expert_bytes * t->n_total_expert;
|
||||
const uint64_t down_bytes = t->down_expert_bytes * t->n_total_expert;
|
||||
if (offset == t->gate_offset && bytes == gate_bytes) b->buf = s->gate;
|
||||
else if (offset == t->up_offset && bytes == gate_bytes) b->buf = s->up;
|
||||
else if (offset == t->down_offset && bytes == down_bytes) b->buf = s->down;
|
||||
else return false;
|
||||
b->off = 0;
|
||||
return b->buf != nil;
|
||||
}
|
||||
|
||||
enum {
|
||||
QWEN4_K_HC_NORM_F16 = 0,
|
||||
QWEN4_K_HC_NORM_F32,
|
||||
@@ -48311,16 +48342,19 @@ static int qwen4_dispatch(int kernel, const void *args, size_t args_len,
|
||||
MTLSize grid, MTLSize tg, NSUInteger tg_mem) {
|
||||
if (!g_initialized && !ds4_gpu_init()) return 0;
|
||||
@autoreleasepool {
|
||||
const bool addresses = g_qwen4_stream_weights && g_qwen4_stream_weights->addresses;
|
||||
const bool q4_mid = kernel == QWEN4_K_MOE_MID_Q4K || kernel == QWEN4_K_MOE_MID_Q4K_NR1;
|
||||
id<MTLComputePipelineState> pipeline = nil;
|
||||
if (kernel == QWEN4_K_MOE_MID || kernel == QWEN4_K_MOE_DOWN || kernel == QWEN4_K_MOE_DOWN_MXFP4_PF) {
|
||||
if (kernel == QWEN4_K_MOE_MID || kernel == QWEN4_K_MOE_DOWN || kernel == QWEN4_K_MOE_DOWN_MXFP4_PF || q4_mid) {
|
||||
const qwen4_moe_args *a = args;
|
||||
const bool specialize = qwen4_moe_mv_specialize(a->weight_type);
|
||||
const bool specialize = !q4_mid && qwen4_moe_mv_specialize(a->weight_type);
|
||||
const uint32_t values[] = {a->weight_type, a->shared_type, specialize ? a->in_dim : 0u, specialize ? qwen4_moe_mv_rows() : 0u};
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_shared=%u_dim=%u_rows=%u",
|
||||
qwen4_kernel_names[kernel], values[0], values[1], values[2], values[3]];
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_shared=%u_dim=%u_rows=%u_addr=%u",
|
||||
qwen4_kernel_names[kernel], values[0], values[1], values[2], values[3], addresses];
|
||||
pipeline = [g_pipeline_cache objectForKey:key];
|
||||
if (!pipeline) {
|
||||
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
|
||||
[constants setConstantValue:&addresses type:MTLDataTypeBool atIndex:906];
|
||||
for (uint32_t i = 0; specialize && i < 4u; i++)
|
||||
[constants setConstantValue:&values[i] type:MTLDataTypeUInt atIndex:901u+i];
|
||||
NSError *error = nil;
|
||||
@@ -48346,13 +48380,14 @@ static int qwen4_dispatch(int kernel, const void *args, size_t args_len,
|
||||
((const qwen4_moe_mm_args *)args)->weight_type : 0u;
|
||||
if (type >= 40u) return 0;
|
||||
const uint32_t tail_base = ((const qwen4_moe_mm_args *)args)->tail_base;
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_tail=%u",
|
||||
qwen4_kernel_names[kernel], type, tail_base];
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_tail=%u_addr=%u",
|
||||
qwen4_kernel_names[kernel], type, tail_base, addresses];
|
||||
pipeline = [g_pipeline_cache objectForKey:key];
|
||||
if (!pipeline) {
|
||||
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
|
||||
[constants setConstantValue:&type type:MTLDataTypeUInt atIndex:900];
|
||||
[constants setConstantValue:&tail_base type:MTLDataTypeUInt atIndex:905];
|
||||
[constants setConstantValue:&addresses type:MTLDataTypeBool atIndex:906];
|
||||
NSError *error = nil;
|
||||
id<MTLFunction> fn = [g_library newFunctionWithName:
|
||||
[NSString stringWithUTF8String:qwen4_kernel_names[kernel]]
|
||||
@@ -48388,7 +48423,17 @@ 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(
|
||||
g_qwen4_stream_weights->entries, g_qwen4_stream_weights->n_entries, 0)) return 0;
|
||||
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
|
||||
if (addresses) {
|
||||
for (uint32_t i = 0; i < g_qwen4_stream_weights->n_entries; i++) {
|
||||
const ds4_gpu_stream_expert_cache_entry *e = g_qwen4_stream_weights->entries[i];
|
||||
[enc useResource:e->gate_buffer usage:MTLResourceUsageRead];
|
||||
[enc useResource:e->up_buffer usage:MTLResourceUsageRead];
|
||||
[enc useResource:e->down_buffer usage:MTLResourceUsageRead];
|
||||
}
|
||||
}
|
||||
[enc setComputePipelineState:pipeline];
|
||||
[enc setBytes:args length:args_len atIndex:0];
|
||||
for (int i = 0; i < n_binds; i++) {
|
||||
@@ -49086,8 +49131,8 @@ int ds4_gpu_qwen4_moe_mid_tensor(
|
||||
has_shared ? 1u : 0u, has_shared ? shared_type : 0u, sh_row_bytes, n_total_expert };
|
||||
qwen4_bind b[7];
|
||||
if (n_tokens == 0 || n_slots == 0 || row_bytes == 0 || ff_dim == 0 || (has_shared && sh_row_bytes == 0) ||
|
||||
!qwen4_bind_weight(&b[0], model_map, model_size, gate_offset, experts_bytes, "moe gate experts") ||
|
||||
!qwen4_bind_weight(&b[1], model_map, model_size, up_offset, experts_bytes, "moe up experts") ||
|
||||
!qwen4_bind_experts(&b[0], model_map, model_size, gate_offset, experts_bytes, "moe gate experts") ||
|
||||
!qwen4_bind_experts(&b[1], model_map, model_size, up_offset, experts_bytes, "moe up experts") ||
|
||||
!qwen4_bind_tensor(&b[2], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
|
||||
!qwen4_bind_tensor(&b[3], x, (uint64_t)n_tokens * in_dim * sizeof(float), "moe input") ||
|
||||
!qwen4_bind_tensor(&b[4], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid")) {
|
||||
@@ -49147,7 +49192,7 @@ int ds4_gpu_qwen4_moe_down_tensor(
|
||||
qwen4_bind b[5];
|
||||
if (n_tokens == 0 || n_slots == 0 || row_bytes == 0 || (ff_dim % 32u) != 0 ||
|
||||
out_dim == 0 || (has_shared && sh_row_bytes == 0) ||
|
||||
!qwen4_bind_weight(&b[0], model_map, model_size, down_offset, experts_bytes, "moe down experts") ||
|
||||
!qwen4_bind_experts(&b[0], model_map, model_size, down_offset, experts_bytes, "moe down experts") ||
|
||||
!qwen4_bind_tensor(&b[1], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
|
||||
!qwen4_bind_tensor(&b[2], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid") ||
|
||||
!qwen4_bind_tensor(&b[3], part, (uint64_t)n_tokens * n_out * out_dim * sizeof(float), "moe partial")) {
|
||||
@@ -49374,8 +49419,8 @@ int ds4_gpu_qwen4_moe_mm_mid_tensor(
|
||||
if (n_tokens == 0 || n_slots == 0 || n_out < n_slots || row_bytes == 0 ||
|
||||
(weight_type != 8u && weight_type != 39u && weight_type != 12u && weight_type != 10u && weight_type != 16u && weight_type != 2u) ||
|
||||
(in_dim % 64) != 0 || ff_dim == 0 || n_expert == 0 || n_expert > 512 ||
|
||||
!qwen4_bind_weight(&b[0], model_map, model_size, gate_offset, expert_bytes * n_expert, "moe gate experts") ||
|
||||
!qwen4_bind_weight(&b[1], model_map, model_size, up_offset, expert_bytes * n_expert, "moe up experts") ||
|
||||
!qwen4_bind_experts(&b[0], model_map, model_size, gate_offset, expert_bytes * n_expert, "moe gate experts") ||
|
||||
!qwen4_bind_experts(&b[1], model_map, model_size, up_offset, expert_bytes * n_expert, "moe up experts") ||
|
||||
!qwen4_bind_tensor(&b[2], lists, (uint64_t)n_expert * list_cap * sizeof(int32_t), "moe lists") ||
|
||||
!qwen4_bind_tensor(&b[3], counts, (uint64_t)n_expert * sizeof(int32_t), "moe counts") ||
|
||||
!qwen4_bind_tensor(&b[4], x, (uint64_t)n_tokens * in_dim * sizeof(float), "moe input") ||
|
||||
@@ -49457,7 +49502,7 @@ int ds4_gpu_qwen4_moe_mm_down_tensor(
|
||||
if (n_tokens == 0 || n_slots == 0 || n_out < n_slots || row_bytes == 0 ||
|
||||
(weight_type != 8u && weight_type != 39u && weight_type != 12u && weight_type != 10u && weight_type != 16u && weight_type != 2u) ||
|
||||
(ff_dim % 64) != 0 || out_dim == 0 || n_expert == 0 || n_expert > 512 ||
|
||||
!qwen4_bind_weight(&b[0], model_map, model_size, down_offset, expert_bytes * n_expert, "moe down experts") ||
|
||||
!qwen4_bind_experts(&b[0], model_map, model_size, down_offset, expert_bytes * n_expert, "moe down experts") ||
|
||||
!qwen4_bind_tensor(&b[1], lists, (uint64_t)n_expert * list_cap * sizeof(int32_t), "moe lists") ||
|
||||
!qwen4_bind_tensor(&b[2], counts, (uint64_t)n_expert * sizeof(int32_t), "moe counts") ||
|
||||
!qwen4_bind_tensor(&b[3], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid") ||
|
||||
@@ -49520,6 +49565,213 @@ int ds4_gpu_qwen4_moe_mm_down_tensor(
|
||||
return 1;
|
||||
}
|
||||
|
||||
static bool qwen4_stream_range(const ds4_gpu_stream_expert_table *t, uint64_t offset, uint64_t bytes) {
|
||||
return offset <= t->model_size && bytes <= t->model_size - offset;
|
||||
}
|
||||
|
||||
static bool qwen4_stream_overlap(const ds4_gpu_tensor *a, uint64_t an,
|
||||
const ds4_gpu_tensor *b, uint64_t bn) {
|
||||
if (!a || !b || ds4_gpu_tensor_buffer(a) != ds4_gpu_tensor_buffer(b)) return false;
|
||||
const uint64_t ao = ds4_gpu_tensor_offset(a), bo = ds4_gpu_tensor_offset(b);
|
||||
return ao <= bo ? bo - ao < an : ao - bo < bn;
|
||||
}
|
||||
|
||||
static int qwen4_stream_read_layer(qwen4_stream_weights *s) {
|
||||
const ds4_gpu_stream_expert_table *t = s->table;
|
||||
const uint64_t gate_bytes = t->gate_expert_bytes * t->n_total_expert;
|
||||
const uint64_t down_bytes = t->down_expert_bytes * t->n_total_expert;
|
||||
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;
|
||||
/* The graph reserves staging headroom independently of the cache. An
|
||||
* off-size MTP layer never changes the trunk slab's size class. */
|
||||
s->gate = [g_device newBufferWithLength:(NSUInteger)gate_bytes options:MTLResourceStorageModeShared];
|
||||
s->up = [g_device newBufferWithLength:(NSUInteger)gate_bytes options:MTLResourceStorageModeShared];
|
||||
s->down = [g_device newBufferWithLength:(NSUInteger)down_bytes options:MTLResourceStorageModeShared];
|
||||
if (!s->gate || !s->up || !s->down) return 0;
|
||||
s->gate.label = @"qwen4_stream_layer_gate";
|
||||
s->up.label = @"qwen4_stream_layer_up";
|
||||
s->down.label = @"qwen4_stream_layer_down";
|
||||
ds4_gpu_stream_expert_pread_task tasks[3] = {
|
||||
{ .offset = t->gate_offset, .len = gate_bytes, .dst = [s->gate contents] },
|
||||
{ .offset = t->up_offset, .len = gate_bytes, .dst = [s->up contents] },
|
||||
{ .offset = t->down_offset, .len = down_bytes, .dst = [s->down contents] },
|
||||
};
|
||||
uint64_t bytes = 0;
|
||||
double ms = 0.0;
|
||||
const int ok = ds4_gpu_stream_expert_pread_tasks(tasks, 3, &bytes, &ms);
|
||||
ds4_gpu_stream_expert_cache_note_pread(t->layer, bytes, ms);
|
||||
return ok;
|
||||
}
|
||||
|
||||
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,
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
uint32_t gate_type, uint32_t down_type, uint32_t n_tokens, uint32_t n_slots,
|
||||
uint32_t in_dim, uint32_t ff_dim, uint32_t list_cap,
|
||||
uint64_t shared_gate_offset, uint64_t shared_up_offset, uint64_t shared_down_offset,
|
||||
uint32_t shared_type, uint32_t shared_down_type) {
|
||||
if (!g_initialized && !ds4_gpu_init()) return 0;
|
||||
const bool mm = lists != NULL;
|
||||
const bool shared = shared_type != UINT32_MAX;
|
||||
if (!g_ssd_streaming_mode || g_qwen4_stream_weights || !table || !table->model_map ||
|
||||
table->layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
|
||||
table->n_total_expert == 0 || table->n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
|
||||
n_tokens == 0 || n_slots == 0 || n_slots > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
|
||||
n_slots > table->n_total_expert || in_dim == 0 || in_dim > UINT32_MAX / 4u ||
|
||||
ff_dim == 0 || ff_dim > UINT32_MAX / 4u ||
|
||||
mm != (counts != NULL) || (mm && (shared || shared_down_type != UINT32_MAX || list_cap == 0)) ||
|
||||
shared != (shared_down_type != UINT32_MAX)) return 0;
|
||||
const uint32_t n_out = n_slots + (shared ? 1u : 0u);
|
||||
const uint64_t n_pairs = (uint64_t)n_tokens * n_out;
|
||||
if (n_pairs > INT32_MAX) return 0;
|
||||
const uint32_t gate_row = qwen4_expert_row_bytes(gate_type, in_dim);
|
||||
const uint32_t down_dim = down_type == 10u ? (ff_dim + 255u) / 256u * 256u : ff_dim;
|
||||
const uint32_t down_row = qwen4_expert_row_bytes(down_type, down_dim);
|
||||
const uint32_t shared_gate_row = shared ? qwen4_expert_row_bytes(shared_type, in_dim) : 0;
|
||||
const uint32_t shared_down_row = shared ? qwen4_expert_row_bytes(shared_down_type, ff_dim) : 0;
|
||||
if (!gate_row || !down_row || (ff_dim % 32u) != 0 ||
|
||||
table->gate_expert_bytes != (uint64_t)gate_row * ff_dim ||
|
||||
table->down_expert_bytes != (uint64_t)down_row * in_dim ||
|
||||
table->gate_expert_bytes > table->model_size / table->n_total_expert ||
|
||||
table->down_expert_bytes > table->model_size / table->n_total_expert ||
|
||||
(mm && ((in_dim % 64u) != 0 || (ff_dim % 64u) != 0 ||
|
||||
(gate_type != 8u && gate_type != 39u && gate_type != 12u && gate_type != 10u && gate_type != 16u && gate_type != 2u) ||
|
||||
(down_type != 8u && down_type != 39u && down_type != 12u && down_type != 10u && down_type != 16u && down_type != 2u)))) return 0;
|
||||
const uint64_t gate_bytes = table->gate_expert_bytes * table->n_total_expert;
|
||||
const uint64_t down_bytes = table->down_expert_bytes * table->n_total_expert;
|
||||
if (!qwen4_stream_range(table, table->gate_offset, gate_bytes) ||
|
||||
!qwen4_stream_range(table, table->up_offset, gate_bytes) ||
|
||||
!qwen4_stream_range(table, table->down_offset, down_bytes) ||
|
||||
(shared && (!shared_gate_row || !shared_down_row ||
|
||||
!qwen4_stream_range(table, shared_gate_offset, (uint64_t)shared_gate_row * ff_dim) ||
|
||||
!qwen4_stream_range(table, shared_up_offset, (uint64_t)shared_gate_row * ff_dim) ||
|
||||
!qwen4_stream_range(table, shared_down_offset, (uint64_t)shared_down_row * in_dim)))) return 0;
|
||||
const uint64_t n_ids = (uint64_t)n_tokens * n_slots;
|
||||
const uint64_t sizes[] = {
|
||||
n_pairs * ff_dim * sizeof(float), n_pairs * in_dim * sizeof(float),
|
||||
(uint64_t)n_tokens * in_dim * sizeof(float), n_ids * sizeof(int32_t),
|
||||
(uint64_t)table->n_total_expert * list_cap * sizeof(int32_t),
|
||||
(uint64_t)table->n_total_expert * sizeof(int32_t),
|
||||
};
|
||||
const ds4_gpu_tensor *tensors[] = {mid, part, x, selected, lists, counts};
|
||||
qwen4_bind checked;
|
||||
for (uint32_t i = 0; i < (mm ? 6u : 4u); i++) {
|
||||
if (!qwen4_bind_tensor(&checked, tensors[i], sizes[i], "streamed MoE") || (checked.off % 4u) != 0) return 0;
|
||||
for (uint32_t o = 0; o < 2u; o++) {
|
||||
if (i != o && qwen4_stream_overlap(tensors[o], sizes[o], tensors[i], sizes[i])) return 0;
|
||||
}
|
||||
}
|
||||
|
||||
int32_t *ids = malloc((size_t)n_ids * sizeof(*ids));
|
||||
if (!ids) return 0;
|
||||
qwen4_stream_weights stream = { .table = table };
|
||||
uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = {0};
|
||||
uint32_t unique = 0;
|
||||
const int had_batch = g_batch_cb != nil;
|
||||
const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
|
||||
double t0 = timing ? ds4_gpu_now_ms() : 0.0, sync_ms = 0.0, copy_ms = 0.0;
|
||||
int ok = ds4_gpu_synchronize();
|
||||
if (timing) { sync_ms = ds4_gpu_now_ms() - t0; t0 = ds4_gpu_now_ms(); }
|
||||
if (!ok) goto done;
|
||||
/* No CPU access to router results or cache mutation before all prior GPU
|
||||
* work completes. Also join a previous explicit-load job before reusing
|
||||
* its worker pool or replacing cache entries. */
|
||||
ds4_gpu_stream_expert_pending_load_clear();
|
||||
ok = ds4_gpu_tensor_read(selected, 0, ids, n_ids * sizeof(*ids));
|
||||
for (uint64_t i = 0; ok && i < n_ids; i++) {
|
||||
if (ids[i] < 0 || (uint32_t)ids[i] >= table->n_total_expert) { ok = 0; break; }
|
||||
if (frequency[ids[i]]++ == 0) unique++;
|
||||
}
|
||||
if (ok && mm) {
|
||||
/* Lists supplied by the router must address exactly the selected
|
||||
* pairs; stale or truncated lists must not dereference missing cache
|
||||
* addresses or leave partial outputs uninitialized. */
|
||||
int32_t count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
|
||||
int32_t *pairs = malloc((size_t)n_ids * sizeof(*pairs));
|
||||
uint8_t *seen = calloc((size_t)n_ids, 1);
|
||||
ok = pairs && seen && ds4_gpu_tensor_read(counts, 0, count, sizes[5]);
|
||||
for (uint32_t e = 0; ok && e < table->n_total_expert; e++) {
|
||||
if (count[e] < 0 || (uint32_t)count[e] != frequency[e] || frequency[e] > list_cap) { ok = 0; break; }
|
||||
if (!frequency[e]) continue;
|
||||
ok = ds4_gpu_tensor_read(lists, (uint64_t)e * list_cap * sizeof(int32_t),
|
||||
pairs, (uint64_t)frequency[e] * sizeof(*pairs));
|
||||
for (uint32_t i = 0; ok && i < frequency[e]; i++) {
|
||||
const int32_t pair = pairs[i];
|
||||
if (pair < 0 || (uint64_t)pair >= n_ids || ids[pair] != (int32_t)e || seen[pair]) ok = 0;
|
||||
else seen[pair] = 1;
|
||||
}
|
||||
}
|
||||
free(seen);
|
||||
free(pairs);
|
||||
}
|
||||
if (timing) { copy_ms = ds4_gpu_now_ms() - t0; t0 = ds4_gpu_now_ms(); }
|
||||
if (!ok) goto done;
|
||||
ds4_gpu_stream_expert_cache_note_token(table->layer);
|
||||
const bool cache_class = ds4_gpu_stream_expert_cache_note_expert_size(table->gate_expert_bytes, table->down_expert_bytes);
|
||||
const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
|
||||
if (cache_class && budget != 0 && unique <= budget) {
|
||||
id<MTLBuffer> gate = nil, up = nil, down = nil;
|
||||
id<MTLBuffer> overflow_gate = nil, overflow_up = nil, overflow_down = nil;
|
||||
uint32_t prepared_unique = 0;
|
||||
g_stream_prefill_batch_selected_addr_building++;
|
||||
ok = ds4_gpu_stream_expert_cache_prepare_selected_batch(
|
||||
table->model_map, table->model_size, table->layer, selected, n_tokens,
|
||||
table->n_total_expert, n_slots, table->gate_offset, table->up_offset, table->down_offset,
|
||||
table->gate_expert_bytes, table->down_expert_bytes,
|
||||
&gate, &up, &down, stream.entries, &stream.n_entries,
|
||||
&prepared_unique, &overflow_gate, &overflow_up, &overflow_down, false);
|
||||
g_stream_prefill_batch_selected_addr_building--;
|
||||
stream.gate = gate;
|
||||
stream.up = up;
|
||||
stream.down = down;
|
||||
ok = ok && !overflow_gate && !overflow_up && !overflow_down && prepared_unique == unique;
|
||||
/* Entries seeded by another cache API may not yet have GPU addresses. */
|
||||
for (uint32_t e = 0; ok && e < table->n_total_expert; e++) {
|
||||
if (!frequency[e]) continue;
|
||||
const ds4_gpu_stream_expert_cache_entry *entry = &g_stream_expert_cache[table->layer][e];
|
||||
ok = entry->valid && ds4_gpu_stream_expert_cache_set_addr_slot_raw(table->layer, e,
|
||||
entry->gate_buffer, entry->gate_inner, entry->up_buffer, entry->up_inner,
|
||||
entry->down_buffer, entry->down_inner);
|
||||
}
|
||||
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);
|
||||
}
|
||||
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) {
|
||||
/* 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];
|
||||
}
|
||||
g_qwen4_stream_weights = &stream;
|
||||
if (mm) {
|
||||
ok = ds4_gpu_qwen4_moe_mm_mid_tensor(mid, x, lists, counts, table->model_map, table->model_size,
|
||||
table->gate_offset, table->up_offset, gate_type, table->n_total_expert,
|
||||
n_tokens, n_slots, n_slots, in_dim, ff_dim, list_cap) &&
|
||||
ds4_gpu_qwen4_moe_mm_down_tensor(part, mid, lists, counts, table->model_map, table->model_size,
|
||||
table->down_offset, down_type, table->n_total_expert,
|
||||
n_tokens, n_slots, n_slots, ff_dim, in_dim, list_cap);
|
||||
} else {
|
||||
ok = ds4_gpu_qwen4_moe_mid_tensor(mid, x, selected, table->model_map, table->model_size,
|
||||
table->gate_offset, table->up_offset, gate_type, table->n_total_expert,
|
||||
n_tokens, n_slots, in_dim, ff_dim, shared_gate_offset, shared_up_offset, shared_type) &&
|
||||
ds4_gpu_qwen4_moe_down_tensor(part, mid, selected, table->model_map, table->model_size,
|
||||
table->down_offset, down_type, table->n_total_expert,
|
||||
n_tokens, n_slots, ff_dim, in_dim, shared_down_offset, shared_down_type);
|
||||
}
|
||||
g_qwen4_stream_weights = NULL;
|
||||
done:
|
||||
free(ids);
|
||||
if (!had_batch && g_batch_cb && !ds4_gpu_end_commands()) ok = 0;
|
||||
if (had_batch && !g_batch_cb && !ds4_gpu_begin_commands()) ok = 0;
|
||||
return ok;
|
||||
}
|
||||
|
||||
int ds4_gpu_qwen4_matmul_q8_0_weights_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *w,
|
||||
uint32_t in_dim, uint32_t out_dim, const ds4_gpu_tensor *x) {
|
||||
if (!g_initialized && !ds4_gpu_init()) return 0;
|
||||
|
||||
+32
-19
@@ -2089,6 +2089,17 @@ kernel void kernel_qwen4_attn_mm(
|
||||
|
||||
/* --- routed experts ----------------------------------------------------- */
|
||||
|
||||
/* Only SSD cache launches bind address tables. The resident specialization
|
||||
* retains its original contiguous expert addressing and arithmetic. */
|
||||
constant bool qwen4_expert_addresses [[function_constant(906)]];
|
||||
static inline device const char *qwen4_expert_base(device const char *base,
|
||||
uint expert, uint64_t bytes) {
|
||||
if (is_function_constant_defined(qwen4_expert_addresses) && qwen4_expert_addresses) {
|
||||
return (device const char *)((device const uint64_t *)base)[expert];
|
||||
}
|
||||
return base + (uint64_t)expert * bytes;
|
||||
}
|
||||
|
||||
#define QWEN4_MOE_NSG 4
|
||||
#define QWEN4_MOE_NR0 2
|
||||
|
||||
@@ -2299,12 +2310,12 @@ kernel void kernel_qwen4_moe_mid(
|
||||
const bool shared = slot == args.n_slots;
|
||||
const uint type = shared ? st : wt;
|
||||
const uint row_bytes = shared ? args.shared_row_bytes : args.row_bytes;
|
||||
device const char *gb = shared ? sh_gate : gate_base;
|
||||
device const char *ub = shared ? sh_up : up_base;
|
||||
const uint64_t ebase = shared ? 0 : (uint64_t)(uint)selected[(uint64_t)tok * args.n_slots + slot] * args.expert_bytes;
|
||||
const uint expert = shared ? 0u : (uint)selected[(uint64_t)tok * args.n_slots + slot];
|
||||
device const char *gb = shared ? sh_gate : qwen4_expert_base(gate_base, expert, args.expert_bytes);
|
||||
device const char *ub = shared ? sh_up : qwen4_expert_base(up_base, expert, args.expert_bytes);
|
||||
device const float *xt = x + (uint64_t)tok * args.in_dim;
|
||||
for (uint r = row0; r < row0 + nr && r < args.out_rows; r++) {
|
||||
const uint64_t off = ebase + (uint64_t)r * row_bytes;
|
||||
const uint64_t off = (uint64_t)r * row_bytes;
|
||||
const float g = qwen4_row_dot(gb + off, xt, type, dim, tiisg);
|
||||
const float u = qwen4_row_dot(ub + off, xt, type, dim, tiisg);
|
||||
if (tiisg == 0) {
|
||||
@@ -2354,7 +2365,8 @@ kernel void kernel_qwen4_moe_mid_q4k(
|
||||
}
|
||||
return;
|
||||
}
|
||||
const uint64_t ebase = (uint64_t)(uint)expert * args.expert_bytes;
|
||||
device const char *gb = qwen4_expert_base(gate_base, (uint)expert, args.expert_bytes);
|
||||
device const char *ub = qwen4_expert_base(up_base, (uint)expert, args.expert_bytes);
|
||||
const uint nb = args.in_dim / 256;
|
||||
const uint group = tiisg / 4, l = (tiisg % 4) * 8;
|
||||
const uint shift = (group & 1u) * 4u;
|
||||
@@ -2364,9 +2376,9 @@ kernel void kernel_qwen4_moe_mid_q4k(
|
||||
float y[8];
|
||||
for (uint i = 0; i < 8; i++) y[i] = yp[i];
|
||||
for (uint r = 0; r < NR && row0 + r < args.out_rows; r++) {
|
||||
const uint64_t off = ebase + (uint64_t)(row0 + r) * args.row_bytes + (uint64_t)ib * 144;
|
||||
device const uchar *bg = (device const uchar *)(gate_base + off);
|
||||
device const uchar *bu = (device const uchar *)(up_base + off);
|
||||
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);
|
||||
@@ -2432,12 +2444,12 @@ kernel void kernel_qwen4_moe_down(
|
||||
const bool shared = slot == args.n_slots;
|
||||
const uint type = shared ? st : wt;
|
||||
const uint row_bytes = shared ? args.shared_row_bytes : args.row_bytes;
|
||||
device const char *db = shared ? sh_down : down_base;
|
||||
const uint64_t pair = (uint64_t)tok * n_out + slot;
|
||||
const uint64_t ebase = shared ? 0 : (uint64_t)(uint)selected[(uint64_t)tok * args.n_slots + slot] * args.expert_bytes;
|
||||
const uint expert = shared ? 0u : (uint)selected[(uint64_t)tok * args.n_slots + slot];
|
||||
device const char *db = shared ? sh_down : qwen4_expert_base(down_base, expert, args.expert_bytes);
|
||||
device const float *m = mid + pair * args.in_dim;
|
||||
for (uint r = row0; r < row0 + nr && r < args.out_rows; r++) {
|
||||
const float v = qwen4_row_dot(db + ebase + (uint64_t)r * row_bytes, m, type, dim, tiisg);
|
||||
const float v = qwen4_row_dot(db + (uint64_t)r * row_bytes, m, type, dim, tiisg);
|
||||
if (tiisg == 0) part[pair * args.out_rows + r] = v;
|
||||
}
|
||||
}
|
||||
@@ -2487,13 +2499,14 @@ kernel void kernel_qwen4_moe_down_mxfp4_pf(
|
||||
}
|
||||
return;
|
||||
}
|
||||
const uint64_t ebase = (uint64_t)(uint)selected[(uint64_t)tok * args.n_slots + slot] * args.expert_bytes;
|
||||
device const char *db = qwen4_expert_base(down_base,
|
||||
(uint)selected[(uint64_t)tok * args.n_slots + slot], args.expert_bytes);
|
||||
const uint ix = tiisg / 8, it = tiisg % 8;
|
||||
const uint nb = dim / 32;
|
||||
for (uint r = row0; r < row0 + nr && r < args.out_rows; r++) {
|
||||
#pragma clang fp reassociate(off)
|
||||
#pragma clang fp contract(off)
|
||||
device const uchar *row = (device const uchar *)(down_base + ebase + (uint64_t)r * args.row_bytes);
|
||||
device const uchar *row = (device const uchar *)(db + (uint64_t)r * args.row_bytes);
|
||||
float acc = 0.0f;
|
||||
uint ib = ix;
|
||||
for (; ib + 12u < nb; ib += 16u) {
|
||||
@@ -2877,8 +2890,8 @@ kernel void kernel_qwen4_moe_mm_mid(
|
||||
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 = gate_base + (uint64_t)e * args.expert_bytes;
|
||||
device const char *ubase = up_base + (uint64_t)e * args.expert_bytes;
|
||||
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;
|
||||
const uint row0 = rb * QWEN4_MM_ROWS;
|
||||
const uint nk = args.in_dim / QWEN4_MM_KS;
|
||||
@@ -3000,7 +3013,7 @@ kernel void kernel_qwen4_moe_mm_down(
|
||||
threadgroup half As[QWEN4_MM_ROWS * QWEN4_MM_KS];
|
||||
threadgroup half Bs[QWEN4_MM_KS * TT];
|
||||
threadgroup float Cs[4][64];
|
||||
device const char *dbase = down_base + (uint64_t)e * args.expert_bytes;
|
||||
device const char *dbase = qwen4_expert_base(down_base, e, args.expert_bytes);
|
||||
device const int32_t *list = lists + (uint64_t)e * args.list_cap;
|
||||
const uint row0 = rb * QWEN4_MM_ROWS;
|
||||
const uint nk = args.in_dim / QWEN4_MM_KS;
|
||||
@@ -3180,8 +3193,8 @@ kernel void kernel_qwen4_moe_mm_mid_nax_t(
|
||||
threadgroup BT *Bs = (threadgroup BT *)(shmem + 8192); /* [NR1][32] */
|
||||
threadgroup half *Br = (threadgroup half *)(shmem + 8192 + NR1 * 64); /* [NR1][32] residual (COMP) */
|
||||
threadgroup float *Cs = (threadgroup float *)shmem; /* [NR1 tok][64 row] after the K loop */
|
||||
device const char *gbase = gate_base + (uint64_t)e * args.expert_bytes;
|
||||
device const char *ubase = up_base + (uint64_t)e * args.expert_bytes;
|
||||
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;
|
||||
const uint row0 = rb * NR0;
|
||||
const uint nk = args.in_dim / NK;
|
||||
@@ -3337,7 +3350,7 @@ kernel void kernel_qwen4_moe_mm_down_nax_t(
|
||||
threadgroup XT *Bs = (threadgroup XT *)(shmem + 4096); /* [NR1][32] */
|
||||
threadgroup half *Br = (threadgroup half *)(shmem + 4096 + NR1 * 64); /* [NR1][32] residual (COMP) */
|
||||
threadgroup float *Cs = (threadgroup float *)shmem; /* [NR1 tok][64 row] after the K loop */
|
||||
device const char *dbase = down_base + (uint64_t)e * args.expert_bytes;
|
||||
device const char *dbase = qwen4_expert_base(down_base, e, args.expert_bytes);
|
||||
device const int32_t *list = lists + (uint64_t)e * args.list_cap;
|
||||
const uint row0 = rb * NR0;
|
||||
const uint nk = args.in_dim / NK;
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
#define DS4_NO_GPU
|
||||
#include "../ds4.c"
|
||||
#include <assert.h>
|
||||
|
||||
static ds4_context_memory estimate(int ctx, uint32_t chunk) {
|
||||
ds4_context_memory m = ds4_context_memory_estimate_with_prefill_mode(
|
||||
DS4_BACKEND_METAL, ctx, chunk, true);
|
||||
assert(m.total_bytes == m.raw_bytes + m.compressed_bytes + m.scratch_bytes);
|
||||
if (!m.prefill_cap || m.prefill_cap > m.raw_cap)
|
||||
fprintf(stderr, "invalid Qwen caps: ctx=%d chunk=%u family=%u prefill=%u raw=%u\n",
|
||||
ctx, chunk, (unsigned)DS4_MODEL_FAMILY, m.prefill_cap, m.raw_cap);
|
||||
assert(m.prefill_cap > 0 && m.prefill_cap <= m.raw_cap);
|
||||
return m;
|
||||
}
|
||||
|
||||
static void monotonic(void) {
|
||||
const uint32_t chunks[] = {1, 2, 8, 128, 1024, 4096, 8192};
|
||||
uint64_t previous = 0;
|
||||
for (size_t i = 0; i < sizeof(chunks) / sizeof(chunks[0]); i++) {
|
||||
ds4_context_memory m = estimate(8192, chunks[i]);
|
||||
assert(m.prefill_cap == chunks[i]);
|
||||
assert(m.total_bytes > previous);
|
||||
previous = m.total_bytes;
|
||||
}
|
||||
const int contexts[] = {1, 16, 128, 512, 4096, 8192, 262144};
|
||||
previous = 0;
|
||||
for (size_t i = 0; i < sizeof(contexts) / sizeof(contexts[0]); i++) {
|
||||
ds4_context_memory m = estimate(contexts[i], 128);
|
||||
assert(m.total_bytes > previous);
|
||||
previous = m.total_bytes;
|
||||
}
|
||||
ds4_context_memory bounded = estimate(128, 128);
|
||||
ds4_context_memory oversized = estimate(128, UINT32_MAX);
|
||||
assert(oversized.prefill_cap == 128 && oversized.total_bytes == bounded.total_bytes);
|
||||
ds4_context_memory smallest = estimate(0, 128);
|
||||
assert(smallest.raw_cap == 1 && smallest.prefill_cap == 1);
|
||||
ds4_context_memory largest = estimate(INT_MAX, UINT32_MAX);
|
||||
assert(largest.prefill_cap == INT_MAX && largest.total_bytes > previous);
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
const ds4_shape saved = g_ds4_shape;
|
||||
g_ds4_shape = DS4_SHAPE_QWEN4_EXP;
|
||||
monotonic();
|
||||
|
||||
/* Independent allocation lower bounds for the shipped 49-layer shape:
|
||||
* 36 GDN layers, 48 value heads, 128x128 state, three convolution rows;
|
||||
* one live state and three MTP snapshots. No tensor data is needed. */
|
||||
const uint64_t recurrent = UINT64_C(36) *
|
||||
(48u * 128u * 128u + 3u * 10240u) * 4u + 9u * 10240u * 4u;
|
||||
assert(recurrent == UINT64_C(118038528));
|
||||
const uint64_t split_attention = UINT64_C(3) * 24u * 64u * 258u * 4u;
|
||||
const uint64_t logit_buffers = UINT64_C(12) * 248320u * 4u;
|
||||
ds4_context_memory m = estimate(4096, 128);
|
||||
const ds4_context_memory large_ctx = estimate(INT_MAX, 1);
|
||||
assert(large_ctx.compressed_bytes == UINT64_C(13) *
|
||||
((uint64_t)INT_MAX * 128u * 4u + ((uint64_t)INT_MAX / 4u + 1u) * 128u * 2u));
|
||||
assert(m.scratch_bytes >= 4u * recurrent + split_attention + logit_buffers +
|
||||
UINT64_C(128) * (10240u + 4u) * 4u);
|
||||
/* The conservative estimate remains bounded; no complete weight tensor or
|
||||
* n-gram table belongs to the graph workspace. */
|
||||
assert(m.total_bytes < UINT64_C(1024) * 1024u * 1024u);
|
||||
ds4_context_memory two = estimate(4096, 2);
|
||||
ds4_context_memory one = estimate(4096, 1);
|
||||
assert(two.raw_bytes == one.raw_bytes && two.compressed_bytes == one.compressed_bytes);
|
||||
assert(two.scratch_bytes - one.scratch_bytes >= (10240u + 4u) * 4u);
|
||||
|
||||
ds4_context_memory resident = ds4_context_memory_estimate_with_prefill_mode(
|
||||
DS4_BACKEND_METAL, 4096, 128, false);
|
||||
assert(resident.total_bytes == m.total_bytes);
|
||||
assert(setenv("DS4_QWEN4_PREFILL_CHUNK", "128", 1) == 0);
|
||||
ds4_context_memory from_env = estimate(4096, 0);
|
||||
assert(from_env.prefill_cap == 128 && from_env.total_bytes == m.total_bytes);
|
||||
assert(unsetenv("DS4_QWEN4_PREFILL_CHUNK") == 0);
|
||||
|
||||
/* Main-only Qwen files have no predictor or verifier snapshots. */
|
||||
g_ds4_shape.n_layer--;
|
||||
g_ds4_shape.n_nextn_predict = 0;
|
||||
ds4_context_memory main_only = estimate(4096, 128);
|
||||
assert(main_only.total_bytes < m.total_bytes);
|
||||
assert(m.total_bytes - main_only.total_bytes >= 3u * recurrent);
|
||||
g_ds4_shape = DS4_SHAPE_QWEN4_MINI;
|
||||
monotonic();
|
||||
g_ds4_shape = saved;
|
||||
puts("Qwen memory: bounded workspace, recurrent/MTP reserves and monotonic budgets OK");
|
||||
return 0;
|
||||
}
|
||||
@@ -4,13 +4,21 @@
|
||||
/* A failed disk lookup must invalidate the recurrent frontier, including
|
||||
* speculative snapshots. Retrying must rebuild from the retained tokens. */
|
||||
int main(int argc, char **argv) {
|
||||
if (argc != 2) {
|
||||
fprintf(stderr, "usage: %s QWEN_GGUF\n", argv[0]);
|
||||
const bool streaming = argc == 3 && strcmp(argv[2], "--ssd-streaming") == 0;
|
||||
if (argc != 2 && !streaming) {
|
||||
fprintf(stderr, "usage: %s QWEN_GGUF [--ssd-streaming]\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
ds4_engine *engine = NULL;
|
||||
ds4_engine_options opt = {.model_path = argv[1], .backend = DS4_BACKEND_METAL,
|
||||
.glm_mtp = true, .prefill_chunk = 8};
|
||||
if (streaming) {
|
||||
/* Account for the actual context before opening the model and leave
|
||||
* room for the live/control sessions and their speculative snapshots. */
|
||||
opt.context_size = 256;
|
||||
opt.ssd_streaming = true;
|
||||
opt.ssd_streaming_cache_experts = 1024;
|
||||
}
|
||||
assert(ds4_engine_open(&engine, &opt) == 0);
|
||||
assert(ds4_engine_is_qwen4(engine) && engine->model.ngram_tensor);
|
||||
ds4_tokens prompt = {0};
|
||||
|
||||
@@ -0,0 +1,272 @@
|
||||
#define _DARWIN_C_SOURCE
|
||||
#include "ds4_gpu.h"
|
||||
#include <math.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <sys/mman.h>
|
||||
#include <unistd.h>
|
||||
|
||||
/* Real SSD reads with a sparse 512-expert file. Only 24 experts contain data;
|
||||
* indices above the older 384-expert cache bound participate in every route. */
|
||||
enum { D = 256, F = 640, E = 512, S = 10, BUDGET = 12, GUARD = 4, HC = 4 };
|
||||
static const int32_t active[] = {511, 384, 383, 0, 1, 2, 3, 4, 5, 6,
|
||||
7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20};
|
||||
static const float poison = -1234.5f;
|
||||
static uint32_t rng = 123;
|
||||
typedef struct {
|
||||
uint32_t gate_type, down_type, shared_down_type;
|
||||
ds4_gpu_stream_expert_table table;
|
||||
uint64_t shared_gate, shared_up, shared_down;
|
||||
} weights;
|
||||
typedef struct { ds4_gpu_tensor *base, *view; uint64_t n; } guarded;
|
||||
|
||||
static void need(int ok, const char *what) {
|
||||
if (!ok) { fprintf(stderr, "Qwen SSD: %s failed\n", what); exit(1); }
|
||||
}
|
||||
static uint32_t random_u32(void) {
|
||||
rng ^= rng << 13; rng ^= rng >> 17; rng ^= rng << 5; return rng;
|
||||
}
|
||||
static uint64_t align_up(uint64_t x, uint64_t a) { return (x + a - 1) / a * a; }
|
||||
static uint32_t block_bytes(uint32_t type) {
|
||||
switch (type) { case 8: return 34; case 10: return 84; case 12: return 144;
|
||||
case 16: return 66; case 39: return 17; default: abort(); }
|
||||
}
|
||||
static uint64_t row_bytes(uint32_t type, uint32_t n) {
|
||||
const uint32_t block = type == 8 || type == 39 ? 32 : 256;
|
||||
return ((uint64_t)n + block - 1) / block * block_bytes(type);
|
||||
}
|
||||
static void half_at(uint8_t *p, uint16_t h) { memcpy(p, &h, sizeof(h)); }
|
||||
static void fill_weights(uint8_t *p, uint64_t bytes, uint32_t type) {
|
||||
const uint32_t block = block_bytes(type);
|
||||
for (uint64_t off = 0; off < bytes; off += block) {
|
||||
uint8_t *b = p + off;
|
||||
for (uint32_t j = 0; j < block; j++) b[j] = random_u32();
|
||||
const uint16_t scale = 0x1800u + (random_u32() % 5u) * 0x100u;
|
||||
if (type == 39) b[0] = 118u + random_u32() % 5u;
|
||||
else if (type == 10) { half_at(b + 80, scale); half_at(b + 82, 0x1400); }
|
||||
else { half_at(b, scale); if (type == 12) half_at(b + 2, 0x1400); }
|
||||
}
|
||||
}
|
||||
static ds4_gpu_tensor *upload(const void *data, uint64_t bytes) {
|
||||
ds4_gpu_tensor *t = ds4_gpu_tensor_alloc(bytes);
|
||||
need(t != NULL, "tensor allocation");
|
||||
if (data) need(ds4_gpu_tensor_write(t, 0, data, bytes), "tensor upload");
|
||||
return t;
|
||||
}
|
||||
static guarded output(uint64_t n) {
|
||||
guarded t = {upload(NULL, (n + 2 * GUARD) * sizeof(float)), NULL, n};
|
||||
t.view = ds4_gpu_tensor_view(t.base, GUARD * sizeof(float), n * sizeof(float));
|
||||
need(t.view != NULL, "guarded view"); return t;
|
||||
}
|
||||
static void reset(guarded *t) {
|
||||
need(ds4_gpu_tensor_fill_f32(t->base, poison, t->n + 2 * GUARD), "output poison");
|
||||
}
|
||||
static float *read_output(const guarded *t) {
|
||||
const uint64_t n = t->n + 2 * GUARD;
|
||||
float *v = malloc(n * sizeof(float)); need(v != NULL, "host output");
|
||||
need(ds4_gpu_tensor_read(t->base, 0, v, n * sizeof(float)), "read output");
|
||||
for (uint64_t i = 0; i < n; i++) {
|
||||
if (i < GUARD || i >= GUARD + t->n)
|
||||
need(memcmp(v + i, &poison, sizeof(float)) == 0, "output guards");
|
||||
else need(isfinite(v[i]), "finite output");
|
||||
}
|
||||
return v;
|
||||
}
|
||||
static void exact(const char *stage, const float *ref, const float *got, uint64_t n) {
|
||||
if (!memcmp(ref, got, n * sizeof(float))) return;
|
||||
for (uint64_t i = 0; i < n; i++) if (memcmp(ref + i, got + i, sizeof(float))) {
|
||||
uint32_t a, b; memcpy(&a, ref + i, 4); memcpy(&b, got + i, 4);
|
||||
fprintf(stderr, "Qwen SSD %s mismatch at %llu: %08x != %08x\n",
|
||||
stage, (unsigned long long)i, a, b); exit(1);
|
||||
}
|
||||
}
|
||||
static void free_output(guarded *t) {
|
||||
ds4_gpu_tensor_free(t->view); ds4_gpu_tensor_free(t->base);
|
||||
}
|
||||
|
||||
static int project(const weights *w, uint32_t T, int streamed,
|
||||
ds4_gpu_tensor *mid, ds4_gpu_tensor *part, ds4_gpu_tensor *x,
|
||||
ds4_gpu_tensor *ids, ds4_gpu_tensor *lists, ds4_gpu_tensor *counts) {
|
||||
const int mm = lists != NULL;
|
||||
const uint32_t shared = mm ? UINT32_MAX : 8u;
|
||||
const ds4_gpu_stream_expert_table *a = &w->table;
|
||||
if (streamed) return ds4_gpu_qwen4_moe_stream_tensor(mid, part, x, ids, lists, counts,
|
||||
a, w->gate_type, w->down_type, T, S, D, F, T,
|
||||
w->shared_gate, w->shared_up, w->shared_down, shared, mm ? UINT32_MAX : w->shared_down_type);
|
||||
if (mm) return ds4_gpu_qwen4_moe_mm_mid_tensor(mid, x, lists, counts,
|
||||
a->model_map, a->model_size, a->gate_offset, a->up_offset, w->gate_type,
|
||||
E, T, S, S, D, F, T) &&
|
||||
ds4_gpu_qwen4_moe_mm_down_tensor(part, mid, lists, counts,
|
||||
a->model_map, a->model_size, a->down_offset, w->down_type, E, T, S, S, F, D, T);
|
||||
return ds4_gpu_qwen4_moe_mid_tensor(mid, x, ids, a->model_map, a->model_size,
|
||||
a->gate_offset, a->up_offset, w->gate_type, E, T, S, D, F,
|
||||
w->shared_gate, w->shared_up, shared) &&
|
||||
ds4_gpu_qwen4_moe_down_tensor(part, mid, ids, a->model_map, a->model_size,
|
||||
a->down_offset, w->down_type, E, T, S, F, D, w->shared_down, w->shared_down_type);
|
||||
}
|
||||
|
||||
/* Both arms use the same row/MM arithmetic and the same final reduction/HC.
|
||||
* Prefill's shared dense expert is outside this routed-only API. Decode also
|
||||
* checks the shared slot, including mixed Q8/MXFP4 weights and its reduction. */
|
||||
static void check_case(const weights *w, uint32_t T, unsigned route, int good_fd,
|
||||
int empty_fd, int fail_first, int warm_probe) {
|
||||
const int mm = T > 8;
|
||||
const uint32_t stride = S + !mm;
|
||||
const uint64_t xn = (uint64_t)T * D, sn = (uint64_t)T * S;
|
||||
const uint64_t rn = xn * HC, in = (uint64_t)T * HC * DS4_QWEN4_HC_CHUNKS * HC;
|
||||
rng = 123u + T * 17u + route;
|
||||
float *x = malloc(xn * 4), *mix = malloc(sn * 4), *r = malloc(rn * 4);
|
||||
float *inj = malloc(in * 4), *sg = malloc(T * 4);
|
||||
int32_t *ids = malloc(sn * 4);
|
||||
need(x && mix && r && inj && sg && ids, "host inputs");
|
||||
for (uint64_t i = 0; i < xn; i++) x[i] = ((int)(random_u32() % 257) - 128) / 1024.f;
|
||||
for (uint64_t i = 0; i < rn; i++) r[i] = ((int)(random_u32() % 257) - 128) / 256.f;
|
||||
for (uint64_t i = 0; i < in; i++) inj[i] = ((int)(random_u32() % 257) - 128) / 1024.f;
|
||||
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];
|
||||
mix[t * S + s] = (s + 1) / 55.f;
|
||||
}
|
||||
}
|
||||
ds4_gpu_tensor *gx = upload(x, xn * 4), *gi = upload(ids, sn * 4);
|
||||
ds4_gpu_tensor *gw = upload(mix, sn * 4), *ginj = upload(inj, in * 4), *gsg = upload(sg, T * 4);
|
||||
ds4_gpu_tensor *lists = mm ? upload(NULL, (uint64_t)E * T * 4) : NULL;
|
||||
ds4_gpu_tensor *counts = mm ? upload(NULL, E * 4) : NULL;
|
||||
guarded outputs[] = {output((uint64_t)T * stride * F), output((uint64_t)T * stride * D),
|
||||
output(xn), output(rn)};
|
||||
const char *stages[] = {"mid", "down", "reduce", "HC residual"};
|
||||
float *reference[4] = {NULL, NULL, NULL, NULL};
|
||||
if (mm) need(ds4_gpu_qwen4_moe_build_lists_tensor(lists, counts, gi, T, S, E, T), "expert lists");
|
||||
ds4_gpu_qwen4_set_verify_rows_exact(T == 3);
|
||||
if (w->gate_type == 16 && T == 1 && route == 0) {
|
||||
weights invalid = *w; invalid.table.down_expert_bytes--;
|
||||
for (unsigned i = 0; i < 4; i++) reset(outputs + i);
|
||||
need(ds4_gpu_begin_commands(), "invalid-size batch begin");
|
||||
need(!project(&invalid, T, 1, outputs[0].view, outputs[1].view, gx, gi, NULL, NULL),
|
||||
"reject wrong physical Q2 row size");
|
||||
need(ds4_gpu_commands_active() && ds4_gpu_end_commands(), "invalid-size batch ownership");
|
||||
const int32_t invalid_id = E;
|
||||
need(ds4_gpu_tensor_write(gi, 0, &invalid_id, 4) && ds4_gpu_begin_commands(), "invalid-ID batch begin");
|
||||
need(!project(w, T, 1, outputs[0].view, outputs[1].view, gx, gi, NULL, NULL), "reject expert ID512");
|
||||
need(ds4_gpu_commands_active() && ds4_gpu_end_commands(), "invalid-ID batch ownership");
|
||||
need(ds4_gpu_tensor_write(gi, 0, ids, sn * 4), "restore valid IDs");
|
||||
for (unsigned i = 0; i < 4; i++) {
|
||||
float *got = read_output(outputs + i);
|
||||
for (uint64_t j = 0; j < outputs[i].n + 2 * GUARD; j++)
|
||||
need(!memcmp(got + j, &poison, 4), "invalid request leaves output untouched");
|
||||
free(got);
|
||||
}
|
||||
}
|
||||
if (fail_first) {
|
||||
need(ds4_gpu_set_model_fd(empty_fd), "empty input fd");
|
||||
need(ds4_gpu_begin_commands(), "failure batch begin");
|
||||
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");
|
||||
need(ds4_gpu_set_model_fd(good_fd), "restore input fd after failure");
|
||||
}
|
||||
const unsigned runs = warm_probe ? 3 : 2;
|
||||
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. */
|
||||
if (run == 2) need(ds4_gpu_set_model_fd(empty_fd), "warm-cache empty 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");
|
||||
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");
|
||||
for (unsigned i = 0; i < 4; i++) {
|
||||
float *got = read_output(outputs + i);
|
||||
if (!run) reference[i] = got;
|
||||
else { exact(stages[i], reference[i], got, outputs[i].n + 2 * GUARD); free(got); }
|
||||
}
|
||||
const uint32_t cached = ds4_gpu_stream_expert_cache_current_count();
|
||||
need(cached <= BUDGET, "bounded cache count");
|
||||
if (run && warm_probe) need(cached >= S, "selected experts cached");
|
||||
}
|
||||
need(ds4_gpu_set_model_fd(good_fd), "restore input fd");
|
||||
float *read_x = malloc(xn * 4); int32_t *read_ids = malloc(sn * 4);
|
||||
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",
|
||||
w->gate_type, w->down_type, T, mm ? "MM" : "rows", route,
|
||||
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); }
|
||||
ds4_gpu_qwen4_set_verify_rows_exact(false);
|
||||
ds4_gpu_tensor_free(gx); ds4_gpu_tensor_free(gi); ds4_gpu_tensor_free(gw);
|
||||
ds4_gpu_tensor_free(ginj); ds4_gpu_tensor_free(gsg); ds4_gpu_tensor_free(lists); ds4_gpu_tensor_free(counts);
|
||||
free(x); free(mix); free(r); free(inj); free(sg); free(ids); free(read_x); free(read_ids);
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
weights formats[] = {{.gate_type=16, .down_type=10, .shared_down_type=8},
|
||||
{.gate_type=12, .down_type=39, .shared_down_type=39},
|
||||
{.gate_type=8, .down_type=8, .shared_down_type=8}};
|
||||
const uint64_t page = (uint64_t)sysconf(_SC_PAGESIZE);
|
||||
uint64_t bytes = page;
|
||||
for (unsigned i = 0; i < 3; i++) {
|
||||
weights *w = formats + i; ds4_gpu_stream_expert_table *t = &w->table;
|
||||
t->layer = 3 + i; t->n_total_expert = E;
|
||||
t->gate_expert_bytes = row_bytes(w->gate_type, D) * F;
|
||||
t->down_expert_bytes = row_bytes(w->down_type, F) * D;
|
||||
t->gate_offset = bytes; bytes = align_up(bytes + E * t->gate_expert_bytes, page);
|
||||
t->up_offset = bytes; bytes = align_up(bytes + E * t->gate_expert_bytes, page);
|
||||
t->down_offset = bytes; bytes = align_up(bytes + E * t->down_expert_bytes, page);
|
||||
w->shared_gate = bytes; bytes = align_up(bytes + row_bytes(8, D) * F, page);
|
||||
w->shared_up = bytes; bytes = align_up(bytes + row_bytes(8, D) * F, page);
|
||||
w->shared_down = bytes; bytes = align_up(bytes + row_bytes(w->shared_down_type, F) * D, page);
|
||||
}
|
||||
FILE *file = tmpfile(), *empty = tmpfile();
|
||||
need(file && empty && ftruncate(fileno(file), (off_t)bytes) == 0, "temporary sparse expert file");
|
||||
uint8_t *map = mmap(NULL, bytes, PROT_READ | PROT_WRITE, MAP_SHARED, fileno(file), 0);
|
||||
need(map != MAP_FAILED, "expert file mapping");
|
||||
for (unsigned i = 0; i < 3; i++) {
|
||||
weights *w = formats + i; ds4_gpu_stream_expert_table *t = &w->table;
|
||||
t->model_map = map; t->model_size = bytes;
|
||||
for (unsigned j = 0; j < sizeof(active) / sizeof(*active); j++) {
|
||||
const uint32_t e = (uint32_t)active[j];
|
||||
fill_weights(map + t->gate_offset + e * t->gate_expert_bytes, t->gate_expert_bytes, w->gate_type);
|
||||
fill_weights(map + t->up_offset + e * t->gate_expert_bytes, t->gate_expert_bytes, w->gate_type);
|
||||
/* Q2's physical tail has nonzero data and must never enter the dot product. */
|
||||
fill_weights(map + t->down_offset + e * t->down_expert_bytes, t->down_expert_bytes, w->down_type);
|
||||
}
|
||||
fill_weights(map + w->shared_gate, row_bytes(8, D) * F, 8);
|
||||
fill_weights(map + w->shared_up, row_bytes(8, D) * F, 8);
|
||||
fill_weights(map + w->shared_down, row_bytes(w->shared_down_type, F) * D, w->shared_down_type);
|
||||
}
|
||||
need(msync(map, bytes, MS_SYNC) == 0 && ds4_gpu_init(), "fixture initialization");
|
||||
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}};
|
||||
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);
|
||||
ds4_gpu_set_streaming_expert_cache_expert_bytes(2 * w->table.gate_expert_bytes + w->table.down_expert_bytes);
|
||||
need(ds4_gpu_stream_expert_cache_current_count() == 0, "cold cache reset");
|
||||
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);
|
||||
}
|
||||
/* 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);
|
||||
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");
|
||||
return 0;
|
||||
}
|
||||
Reference in new issue
Block a user