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

Metal: remove redundant Qwen MTP projections and expert work

Browse Source 
Preparing historical predictor cache rows previously computed query/gate
and indexer-query projections that were never consumed. Catch-up batches
also ran the predictor MoE for every row even though only the final
residual feeds the vocabulary head or the next recursive draft.

Add a cache-only attention preparation kernel that writes K, V and raw
indexer keys using the original normalization, RoPE and cache arithmetic.
Skip query projections at graph level while retaining fused K/V/IK GEMV
for short batches and the existing pooled-key updates. Ordinary attention
continues through the complete preparation path.

Add a Q8 EH projection for the staged [embedding|0] and [0|hidden] rows.
Separate embedding and hidden dispatches within one encoder, traversing
only the active weight half while preserving full row strides and the
original float4 dot products and SIMD reduction order. Check the inactive
half's scale words to preserve NaN/Inf classification for zero products.
Validate tensor spans, alignment and output aliasing before dispatch.
Select this kernel on M1 Max for E2560, four HC streams and T1..3; keep
the existing path on other devices/layouts, non-Q8 weights and diagnostic
GEMM or Q8 prefill-profile selections. No quantization change, new buffer
allocation or production opt-in is introduced.

Keep batched attention, HC mixing and router arithmetic during MTP
catch-up, then bind last-row views for the independent routed/shared MoE
and residual reduction. Restore graph bindings and release views on all
exit paths. Earlier rows still build the full predictor KV prefix; only
the final residual is consumed. Ordinary trunk MoE remains unchanged.

Extend the focused Metal kernel regression with T1/T2/T3, actual Qwen
projection dimensions, cache boundaries and guards, nonzero multimodal
RoPE positions, signed zero, nonfinite Q8 scales, diagnostic fallbacks,
input immutability and invalid-span/alias rejection. Add a frozen original
MTP input/layer oracle to the model-backed prefix test, including recursive
drafts and exact three-row verification. Preserve diagnostic environment
settings and cover unfused execution without assuming speculation runs.

Validation on M1 Max 32 GiB with Qwen3.8-Flash-Next-Q2.gguf:
- Application and relevant test targets built; git diff --check passed.
- Focused synthetic cache/projection kernel tests and memory checks passed.
- Real SSD predictor comparisons matched EH, KV/pooled keys, final residual,
  all 248320 logits and recursive drafts bit for bit for T1/T2/T3.
- Prefix, chunk/append, payload and unfused tests passed. Session checkpoint,
  restore/rewind, exact-sampling continuation, n-gram failure recovery and
  small-buffer regressions passed. Independent target continuation stayed
  within the established batch tolerance (max absolute error 0.000974178).

Performance versus 92f57fc, frozen binaries/shaders, SSD streaming,
ctx4096/chunk128, temp0, no thinking, MTP enabled via --mtp-timing:
- Four runs per version and prompt in ABBA BAAB order, harmonic-mean rates.
- Rome, 29 input/100 output tokens: prefill 9.85 -> 9.88 t/s (+0.3%),
  decode 9.01 -> 9.57 t/s (+6.2%); mean process time increased 0.8%.
- Long prompt, 575 input/50 output tokens: prefill 41.28 -> 42.12 t/s
  (+2.0%), decode 9.40 -> 9.11 t/s (-3.0%); process time decreased 0.5%.
- Generated outputs and first-draft acceptance counts were identical.
- Isolated EH microbenchmark: 16 samples per arm, 64 dispatches per sample,
  including host encoding and GPU completion. T1: 161.28 -> 127.36 us;
  T2: 309.67 -> 188.75 us; T3: 438.84 -> 277.03 us (+27..64% throughput).
- A separate no-MTP control measured 7.58 -> 7.33 t/s decode (-3.3%).
  End-to-end variation remains material; no stable global 10% improvement
  is established. Final runs had no concurrent model processes, but OS
  file-cache and desktop background activity were not controlled.

CUDA/ROCm were not exercised; graph changes are in the Metal Qwen path.
Temporary benchmarks, logs, notes and model files are excluded.
This commit is contained in:
GiorgioOppo committed 2026-09-14 20:53:23 +02:00
1 parent 92f57fc4d7
commit b86c8ae1db
6 files changed
+787 -37
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

ds4.c
+97 -26
View File
@@ -57593,7 +57593,7 @@ typedef struct {
int32_t mrope_delta;
int32_t snap_mrope_delta;
const ds4_vision_span *vis_spans;
uint32_t mtp_last_rows; /* rows of the last predictor pass resident in mtp_R */
uint32_t mtp_last_rows; /* final predictor output is mtp_R[mtp_last_rows - 1]; zero means none */
size_t vis_span_count;
float *host_row;
float *host_logits;
@@ -58234,28 +58234,40 @@ static bool qwen4_graph_attention_core(ds4_qwen4_gpu_graph *g, uint32_t il,
}
static bool qwen4_graph_attention_cache(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
uint32_t il, uint32_t pos0, uint32_t T) {
uint32_t il, uint32_t pos0, uint32_t T, bool cache_only) {
const uint32_t ratio = 4u;
const uint32_t last = pos0 + T - 1u;
{
bool ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T);
/* Historical predictor rows need keys and values only. Keep their
* projection kernels and K normalization identical to a full step. */
bool ok = cache_only || qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T);
if (ok && qwen4_graph_fused(g, T)) {
ds4_gpu_tensor *outs[4] = { g->kp, g->vp, g->iq, g->ik };
ds4_gpu_tensor *outs[4] = { g->kp, g->vp, cache_only ? g->ik : g->iq, g->ik };
const uint64_t offs[4] = { l->attn_k->abs_offset, l->attn_v->abs_offset,
l->indexer_q_proj->abs_offset, l->indexer_k_proj->abs_offset };
const uint32_t types[4] = { l->attn_k->type, l->attn_v->type, l->indexer_q_proj->type, l->indexer_k_proj->type };
cache_only ? l->indexer_k_proj->abs_offset : l->indexer_q_proj->abs_offset,
l->indexer_k_proj->abs_offset };
const uint32_t types[4] = { l->attn_k->type, l->attn_v->type,
cache_only ? l->indexer_k_proj->type : l->indexer_q_proj->type,
l->indexer_k_proj->type };
const uint32_t rows[4] = { (uint32_t)l->attn_k->dim[1], (uint32_t)l->attn_v->dim[1],
(uint32_t)l->indexer_q_proj->dim[1], (uint32_t)l->indexer_k_proj->dim[1] };
ok = ds4_gpu_qwen4_multi_gemv_tensor(g->mixed, T, DS4_N_EMBD, 4, outs, m->map, m->size, offs, types, rows) != 0;
(uint32_t)(cache_only ? l->indexer_k_proj->dim[1] : l->indexer_q_proj->dim[1]),
(uint32_t)l->indexer_k_proj->dim[1] };
ok = ds4_gpu_qwen4_multi_gemv_tensor(g->mixed, T, DS4_N_EMBD, cache_only ? 3u : 4u,
outs, m->map, m->size, offs, types, rows) != 0;
} else if (ok) {
ok = qwen4_gemv(g->kp, m, l->attn_k, g->mixed, T) &&
qwen4_gemv(g->vp, m, l->attn_v, g->mixed, T) &&
qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, T) &&
(cache_only || qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, T)) &&
qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T);
}
if (!ok) return false;
}
if (!(ds4_gpu_qwen4_attn_prep_tensor(g->q, g->gate, g->layer_k_cache[il], g->layer_v_cache[il], g->iqn,
if (cache_only) {
if (!ds4_gpu_qwen4_attn_cache_prep_tensor(g->layer_k_cache[il], g->layer_v_cache[il],
g->layer_ik_cache[il], g->kp, g->vp, g->ik, g->pos3, m->map, m->size,
l->attn_k_norm->abs_offset, T, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT,
DS4_N_INDEXER_HEAD_DIM, pos0, g->ctx_cap, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS)) return false;
} else if (!(ds4_gpu_qwen4_attn_prep_tensor(g->q, g->gate, g->layer_k_cache[il], g->layer_v_cache[il], g->iqn,
g->layer_ik_cache[il], g->qg, g->kp, g->vp, g->iq, g->ik, g->pos3,
m->map, m->size, l->attn_q_norm->abs_offset, l->attn_k_norm->abs_offset,
l->indexer_q_norm->abs_offset, T, DS4_N_HEAD, DS4_N_HEAD_KV,
@@ -58281,7 +58293,7 @@ static bool qwen4_graph_attention(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
const uint32_t ratio = 4u;
const uint32_t n_blocks_after = (pos0 + T) / ratio;
const uint32_t q_dim = DS4_N_HEAD * DS4_N_HEAD_DIM;
if (!qwen4_graph_attention_cache(g, m, l, il, pos0, T)) return false;
if (!qwen4_graph_attention_cache(g, m, l, il, pos0, T, false)) return false;
if (T == 3u && g->verify_rows_exact) {
/* 3-row speculative verify: run the attention core as 2/1-row
* sub-batches so every dispatch keeps the exact T <= 2 kernel paths
@@ -58328,8 +58340,8 @@ 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 il, uint32_t T) {
static bool qwen4_graph_moe_impl(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
uint32_t il, uint32_t T, bool last_only) {
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;
@@ -58338,6 +58350,35 @@ static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds
l->ffn_gate_inp_shexp->abs_offset, l->ffn_gate_inp_shexp->type,
DS4_N_EMBD, g->sh_gate_logit, T, DS4_N_EXPERT, DS4_N_EXPERT_USED) &&
qwen4_moe_profile_boundary(profile, &last, &elapsed[0]);
ds4_gpu_tensor **fields[] = {
&g->mixed, &g->selected, &g->weights, &g->sh_gate_logit, &g->R, &g->inj,
};
ds4_gpu_tensor *saved[6] = {0}, *views[6] = {0};
bool bound = false;
if (ok && last_only && T > 1u) {
/* MTP only consumes the final residual row. Keep the router's original
* batch geometry, then run its independent expert rows once. The HC
* injection has a chunk matrix per row; mid/part/blk remain scratch. */
const uint64_t row_bytes[] = {
(uint64_t)DS4_N_EMBD * sizeof(float),
(uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t),
(uint64_t)DS4_N_EXPERT_USED * sizeof(float),
sizeof(float),
(uint64_t)DS4_N_HC * DS4_N_EMBD * sizeof(float),
(uint64_t)DS4_N_HC * DS4_QWEN4_HC_CHUNKS * DS4_N_HC * sizeof(float),
};
for (uint32_t i = 0; i < 6u; i++) {
saved[i] = *fields[i];
views[i] = ds4_gpu_tensor_view(saved[i], (uint64_t)(T - 1u) * row_bytes[i], row_bytes[i]);
if (!views[i]) {
ok = false;
goto done;
}
}
for (uint32_t i = 0; i < 6u; i++) *fields[i] = views[i];
bound = true;
T = 1u;
}
/* Prefill-sized batches route each expert's tokens through tiled GEMMs
* (weights read once per 32 tokens) and run the shared expert as dense
* GEMMs; decode keeps the per-(token, slot) row kernels with the shared
@@ -58352,7 +58393,10 @@ static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds
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;
l->ffn_gate_exps->bytes != l->ffn_up_exps->bytes) {
ok = false;
goto done;
}
const ds4_gpu_stream_expert_table table = {
.model_map = m->map, .model_size = m->size, .layer = il,
.n_total_expert = DS4_N_EXPERT,
@@ -58380,10 +58424,11 @@ static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds
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(
ok = 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;
goto done;
}
if (mm) {
if (ok) {
@@ -58419,7 +58464,7 @@ static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds
g->pos, T, (int)l->ffn_gate_exps->name.len, l->ffn_gate_exps->name.ptr, ok ? 1 : 0,
elapsed[0] * 1e3, elapsed[1] * 1e3, elapsed[2] * 1e3, elapsed[3] * 1e3,
elapsed[4] * 1e3, elapsed[5] * 1e3, elapsed[6] * 1e3);
return ok;
goto done;
}
if (ok) {
ok = ds4_gpu_qwen4_moe_mid_tensor(g->mid, g->mixed, g->selected, m->map, m->size, l->ffn_gate_exps->abs_offset,
@@ -58437,9 +58482,26 @@ static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds
ok = ds4_gpu_qwen4_moe_reduce_tensor(g->blk, g->part, g->weights, g->sh_gate_logit, NULL, g->R, g->inj, T,
DS4_N_EXPERT_USED, DS4_N_EXPERT_USED + 1u, DS4_N_EMBD, DS4_N_HC) != 0;
}
done:
if (last_only) {
if (bound) {
for (uint32_t i = 0; i < 6u; i++) *fields[i] = saved[i];
}
for (uint32_t i = 0; i < 6u; i++) ds4_gpu_tensor_free(views[i]);
}
return ok;
}
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) {
return qwen4_graph_moe_impl(g, m, l, il, T, false);
}
static bool qwen4_graph_moe_last(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
uint32_t il, uint32_t T) {
return T >= 1u && T <= 3u && qwen4_graph_moe_impl(g, m, l, il, T, true);
}
/* host side: embedding rows tiled into R and the PLE n-gram gather */
static bool qwen4_graph_stage_inputs(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
const int *tokens, uint32_t T) {
@@ -58867,13 +58929,20 @@ static bool qwen4_graph_state_swap2(ds4_qwen4_gpu_graph *g) {
return true;
}
/* Each staged EH row has one zero half. The Q8 kernel skips its quantized
* products while retaining the full projection's reduction order. */
static bool qwen4_graph_mtp_project(ds4_qwen4_gpu_graph *g, const ds4_model *m,
const ds4_layer_weights *l, uint32_t T) {
if (l->nextn_eh_proj->type == DS4_TENSOR_Q8_0)
return ds4_gpu_qwen4_mtp_project_tensor(g->mtp_proj, g->mtp_cat, m->map, m->size,
l->nextn_eh_proj->abs_offset, DS4_N_EMBD, DS4_N_HC, T) != 0;
return qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, T * (DS4_N_HC + 1u));
}
/* The steering bank contains trunk layers only. The predictor remains
* unsteered; its drafts are verified by the steered target trunk.
* One to three causal predictor steps at idx, using consecutive rows of the
* trunk's pre-mixer streams and the embeddings of their following tokens.
* The nextn layer runs on its own residual; want_logits returns the last
* row only. Batching catches up accepted drafts and predicts again with
* one command submission and one set of weight reads. */
* unsteered; its drafts are verified by the steered target trunk. Each row
* uses a trunk pre-mixer stream and the embedding of its following token,
* producing the predictor's own residual without modifying the trunk. */
static bool qwen4_graph_mtp_input(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
uint32_t row, const int *next_tokens, uint32_t T) {
const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC;
@@ -58901,7 +58970,7 @@ static bool qwen4_graph_mtp_input(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
ds4_gpu_tensor_free(R_row);
ds4_gpu_tensor_free(e_row);
}
if (ok) ok = qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, T * (hc + 1u));
if (ok) ok = qwen4_graph_mtp_project(g, m, l, T);
for (uint32_t t = 0; t < T && ok; t++) {
ds4_gpu_tensor *proj_row = ds4_gpu_tensor_view(g->mtp_proj, t * proj_bytes, proj_bytes);
ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(g->mtp_R, t * hc * emb_bytes, hc * emb_bytes);
@@ -58927,7 +58996,7 @@ static bool qwen4_graph_mtp_cache_steps(ds4_qwen4_gpu_graph *g, const ds4_model
bool ok = qwen4_graph_mtp_input(g, m, w, row, next_tokens, T);
g->R = g->mtp_R;
if (ok) ok = qwen4_graph_attention_cache(g, m, &w->layer[DS4_N_LAYER - 1u],
DS4_N_LAYER - 1u, idx, T);
DS4_N_LAYER - 1u, idx, T, true);
g->R = R;
if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false;
g->mtp_last_rows = 0; /* staged inputs are not recursive predictor outputs */
@@ -58941,6 +59010,8 @@ static bool qwen4_graph_mtp_cache_steps(ds4_qwen4_gpu_graph *g, const ds4_model
return ok;
}
/* Catch up one to three consecutive predictor rows. Only the final row's
* output feeds the head or recursive draft; earlier rows contribute KV. */
static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
uint32_t row, const int *next_tokens, uint32_t T, uint32_t idx, bool want_logits,
float *logits_out, int *draft_out) {
@@ -58956,7 +59027,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, il, T);
if (ok) ok = qwen4_graph_moe_last(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 &&
@@ -59041,7 +59112,7 @@ static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *
ds4_gpu_tensor_free(R_row);
ds4_gpu_tensor_free(e_row);
}
if (ok) ok = qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, hc + 1u);
if (ok) ok = qwen4_graph_mtp_project(g, m, l, 1u);
if (ok) {
ds4_gpu_tensor *proj_row = ds4_gpu_tensor_view(g->mtp_proj, 0, proj_bytes);
ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(g->mtp_R, 0, hc * emb_bytes);
ds4_gpu.h
+16
View File
@@ -3358,6 +3358,15 @@ int ds4_gpu_qwen4_attn_prep_tensor(
uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim, uint32_t n_rot,
uint32_t n_idx_head, uint32_t idx_dim, uint32_t pos0, uint32_t cache_cap,
float rope_base, float eps);
/* Predictor prefix preparation: identical K/V/IK arithmetic to attn_prep,
* without query, gate or indexer-query inputs/outputs. Cache spans and all
* inputs must be disjoint from the three output spans. */
int ds4_gpu_qwen4_attn_cache_prep_tensor(
ds4_gpu_tensor *k_cache, ds4_gpu_tensor *v_cache, ds4_gpu_tensor *ik_cache,
const ds4_gpu_tensor *kproj, const ds4_gpu_tensor *vproj, const ds4_gpu_tensor *ik,
const ds4_gpu_tensor *pos3, const void *model_map, uint64_t model_size, uint64_t g_k_offset,
uint32_t n_tokens, uint32_t n_head_kv, uint32_t head_dim, uint32_t n_rot,
uint32_t idx_dim, uint32_t pos0, uint32_t cache_cap, float rope_base, float eps);
int ds4_gpu_qwen4_idx_block_key_tensor(
ds4_gpu_tensor *block_key, const ds4_gpu_tensor *ik_cache, const ds4_gpu_tensor *pos3,
const void *model_map, uint64_t model_size, uint64_t g_ik_offset,
@@ -3478,6 +3487,13 @@ int ds4_gpu_qwen4_mtp_stage_tensor(
ds4_gpu_tensor *cat, const ds4_gpu_tensor *e, const ds4_gpu_tensor *R,
const void *model_map, uint64_t model_size, uint64_t g_e_offset, uint64_t g_h_offset,
uint32_t n_embd, uint32_t n_hc, float eps);
/* Apply Q8 weights [E][2E] to the exact [T][hc+1][2E] layout written by
* mtp_stage: [embedding|0], then hc rows [0|hidden]. Preserves the existing
* Q8 batch arithmetic and full row strides; out must not overlap cat/weights. */
int ds4_gpu_qwen4_mtp_project_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *cat,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens);
int ds4_gpu_qwen4_mtp_combine_tensor(
ds4_gpu_tensor *R_out, const ds4_gpu_tensor *proj, uint32_t n_embd, uint32_t n_hc);
ds4_metal.m
+107
View File
@@ -48161,6 +48161,7 @@ enum {
QWEN4_K_PLE_CONV,
QWEN4_K_ROUTER,
QWEN4_K_ATTN_PREP,
QWEN4_K_ATTN_CACHE_PREP,
QWEN4_K_IDX_BLOCK_KEY,
QWEN4_K_IDX_SCORE,
QWEN4_K_IDX_SCORE_VEC,
@@ -48250,6 +48251,7 @@ static const char *const qwen4_kernel_names[QWEN4_K_COUNT] = {
"kernel_qwen4_ple_conv",
"kernel_qwen4_router_topk",
"kernel_qwen4_attn_prep",
"kernel_qwen4_attn_cache_prep",
"kernel_qwen4_idx_block_key",
"kernel_qwen4_idx_score",
"kernel_qwen4_idx_score_vec",
@@ -48958,6 +48960,51 @@ int ds4_gpu_qwen4_attn_prep_tensor(
MTLSizeMake(n_head + n_head_kv + n_idx_head + 1u, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
}
int ds4_gpu_qwen4_attn_cache_prep_tensor(
ds4_gpu_tensor *k_cache, ds4_gpu_tensor *v_cache, ds4_gpu_tensor *ik_cache,
const ds4_gpu_tensor *kproj, const ds4_gpu_tensor *vproj, const ds4_gpu_tensor *ik,
const ds4_gpu_tensor *pos3, const void *model_map, uint64_t model_size, uint64_t g_k_offset,
uint32_t n_tokens, uint32_t n_head_kv, uint32_t head_dim, uint32_t n_rot,
uint32_t idx_dim, uint32_t pos0, uint32_t cache_cap, float rope_base, float eps) {
if (!n_tokens || !n_head_kv || head_dim < 32u || head_dim > 256u || head_dim % 32u ||
idx_dim < 32u || idx_dim > 128u || idx_dim % 32u || n_rot > 64u || n_rot % 2u ||
n_rot > head_dim || n_rot > idx_dim || pos0 > cache_cap || n_tokens > cache_cap - pos0 ||
n_head_kv > UINT32_MAX / (head_dim * sizeof(float))) return 0;
struct {
uint32_t n_tokens, n_head, n_head_kv, head_dim, n_rot, n_idx_head, idx_dim, pos0, cache_cap;
float rope_base, eps; uint32_t pad0; float rope_mscale; float rope_freq[32];
} args = { n_tokens, 0u, n_head_kv, head_dim, n_rot, 0u, idx_dim, pos0, cache_cap,
rope_base, eps, 0u, 1.0f, { 0 } };
qwen4_rope_fill(args.rope_freq, &args.rope_mscale, n_rot, rope_base);
const uint64_t row_bytes = (uint64_t)n_head_kv * head_dim * sizeof(float);
const uint64_t sizes[8] = {
(uint64_t)n_tokens * row_bytes, (uint64_t)n_tokens * row_bytes,
(uint64_t)n_tokens * idx_dim * sizeof(float), (uint64_t)head_dim * sizeof(float),
(uint64_t)cache_cap * (row_bytes / 2u), (uint64_t)cache_cap * (row_bytes / 2u),
(uint64_t)cache_cap * idx_dim * sizeof(float), (uint64_t)cache_cap * 16u,
};
qwen4_bind b[8];
if (!qwen4_bind_tensor(&b[0], kproj, sizes[0], "cache k projection") ||
!qwen4_bind_tensor(&b[1], vproj, sizes[1], "cache v projection") ||
!qwen4_bind_tensor(&b[2], ik, sizes[2], "cache indexer k projection") ||
!qwen4_bind_weight(&b[3], model_map, model_size, g_k_offset, sizes[3], "cache k_norm") ||
!qwen4_bind_tensor(&b[4], k_cache, sizes[4], "k cache") ||
!qwen4_bind_tensor(&b[5], v_cache, sizes[5], "v cache") ||
!qwen4_bind_tensor(&b[6], ik_cache, sizes[6], "indexer k cache") ||
!qwen4_bind_tensor(&b[7], pos3, sizes[7], "cache rope positions")) return 0;
for (uint32_t i = 0; i < 8u; i++) {
const NSUInteger alignment = i == 7u ? 16u : i == 4u || i == 5u ? 2u : 4u;
if (b[i].off % alignment) return 0;
for (uint32_t o = 4u; o <= 6u; o++) {
if (i == o || b[i].buf != b[o].buf) continue;
const uint64_t a = b[i].off, z = b[o].off;
if (a <= z ? z - a < sizes[i] : a - z < sizes[o]) return 0;
}
}
return qwen4_dispatch(QWEN4_K_ATTN_CACHE_PREP, &args, sizeof(args), b, 8,
MTLSizeMake(n_head_kv + 1u, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
}
int ds4_gpu_qwen4_idx_block_key_tensor(
ds4_gpu_tensor *block_key, const ds4_gpu_tensor *ik_cache, const ds4_gpu_tensor *pos3,
const void *model_map, uint64_t model_size, uint64_t g_ik_offset,
@@ -50101,6 +50148,66 @@ int ds4_gpu_qwen4_mtp_stage_tensor(
MTLSizeMake(n_hc + 1u, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_mtp_project_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *cat,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!n_embd || n_embd > UINT32_MAX / 2u || (n_embd % 16u) != 0u ||
!n_hc || n_hc == UINT32_MAX || !n_tokens || n_tokens > UINT32_MAX / (n_hc + 1u)) return 0;
const uint64_t in_dim = 2u * (uint64_t)n_embd;
const uint64_t rows = (uint64_t)n_tokens * (n_hc + 1u);
const uint64_t row_bytes = (in_dim / 32u) * 34u;
if (rows > UINT64_MAX / (in_dim * sizeof(float))) return 0;
const uint64_t sizes[3] = {
row_bytes * n_embd, rows * in_dim * sizeof(float), rows * n_embd * sizeof(float),
};
qwen4_bind b[3];
if (!qwen4_bind_weight(&b[0], model_map, model_size, weight_offset, sizes[0], "mtp Q8 projection") ||
!qwen4_bind_tensor(&b[1], cat, sizes[1], "mtp concat") ||
!qwen4_bind_tensor(&b[2], out, sizes[2], "mtp projection")) return 0;
if (b[0].off % 2u || b[1].off % 16u || b[2].off % 4u) return 0;
for (uint32_t i = 0; i < 2u; i++) {
if (b[i].buf != b[2].buf) continue;
const uint64_t a = b[i].off, z = b[2].off;
if (a <= z ? z - a < sizes[i] : a - z < sizes[2]) return 0;
}
const uint64_t mv_max = ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u);
/* Keep other layouts/devices and diagnostic GEMM arithmetic on their
* existing dispatch. The optimized layout is emitted by mtp_stage. */
if (n_embd != 2560u || n_hc != 4u || n_tokens > 3u || rows > mv_max ||
(rows > 8u && ds4_gpu_env_bool("DS4_METAL_Q8_PREFILL_PROFILE") > 0) ||
!ds4_gpu_device_name_contains("M1 Max")) {
return ds4_gpu_qwen4_matmul_q8_0_tensor(out, model_map, model_size, weight_offset,
in_dim, n_embd, cat, rows);
}
const int16_t nsg = ds4_gpu_mv_ext_nsg();
const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, rows);
const char *emb_fn = n_tokens == 1u ? "kernel_qwen4_mtp_project_q8_emb_1" :
n_tokens == 2u ? "kernel_qwen4_mtp_project_q8_emb_2" : "kernel_qwen4_mtp_project_q8_emb_3";
id<MTLComputePipelineState> emb_pipeline = ds4_gpu_get_mul_mv_ext_pipeline(emb_fn, nsg, nxpsg);
id<MTLComputePipelineState> hidden_pipeline =
ds4_gpu_get_mul_mv_ext_pipeline("kernel_qwen4_mtp_project_q8_hidden_4", nsg, nxpsg);
if (!emb_pipeline || !hidden_pipeline) return 0;
ds4_gpu_mul_mv_ext_args args = ds4_gpu_make_mv_ext_args(in_dim, n_embd, rows, 34u, row_bytes);
const uint64_t r0ptg = (uint64_t)(32 / nxpsg) * nsg;
const NSUInteger row_groups = ((NSUInteger)n_embd + r0ptg - 1u) / r0ptg;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:emb_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
for (uint32_t i = 0; i < 3u; i++) [enc setBuffer:b[i].buf offset:b[i].off atIndex:i + 1u];
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
[enc setComputePipelineState:hidden_pipeline];
[enc dispatchThreadgroups:MTLSizeMake(row_groups, n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "Qwen MTP Q8 projection");
}
int ds4_gpu_qwen4_mtp_combine_tensor(
ds4_gpu_tensor *R_out, const ds4_gpu_tensor *proj, uint32_t n_embd, uint32_t n_hc) {
struct { uint32_t n_embd, n_hc; } args = { n_embd, n_hc };
metal/qwen4.metal
+147
View File
@@ -1241,6 +1241,60 @@ kernel void kernel_qwen4_attn_prep(
}
}
/* Predictor prefix rows need only K/V and raw indexer K. Keep the full
* preparation kernel's lane walk, normalization and serial RoPE unchanged. */
kernel void kernel_qwen4_attn_cache_prep(
constant ds4_metal_args_qwen4_attn_prep & args,
device const float *kproj,
device const float *vproj,
device const float *ik,
device const float *g_k,
device half *k_cache,
device half *v_cache,
device float *ik_cache,
device const uint4 *pos3,
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tiisg [[thread_index_in_simdgroup]]) {
const uint slot = tgpig.x, tok = tgpig.y;
if (tok >= args.n_tokens || slot > args.n_head_kv) return;
const uint Hkv = args.n_head_kv, D = args.head_dim, Di = args.idx_dim;
const uint pos = args.pos0 + tok;
float v[8];
threadgroup float row[576];
float tmp[64];
if (slot < Hkv) {
const uint h = slot;
const uint npt = D / 32;
const uint4 p3 = pos3[pos];
device const float *src = kproj + ((uint64_t)tok * Hkv + h) * D;
device const float *vs = vproj + ((uint64_t)tok * Hkv + h) * D;
float ss = 0.0f;
for (uint i = 0; i < npt; i++) { v[i] = src[tiisg * npt + i]; ss += v[i] * v[i]; }
ss = simd_sum(ss);
const float r = rsqrt(ss / (float)D + args.eps);
for (uint i = 0; i < npt; i++) v[i] = v[i] * r * g_k[tiisg * npt + i];
for (uint i = 0; i < npt; i++) row[tiisg * npt + i] = v[i];
simdgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) {
for (uint i = 0; i < args.n_rot; i++) tmp[i] = row[i];
qwen4_rope_neox(tmp, args.n_rot, p3, args.rope_freq, args.rope_mscale);
for (uint i = 0; i < args.n_rot; i++) row[i] = tmp[i];
}
simdgroup_barrier(mem_flags::mem_threadgroup);
device half *dk = k_cache + ((uint64_t)pos * Hkv + h) * D;
device half *dv = v_cache + ((uint64_t)pos * Hkv + h) * D;
for (uint i = 0; i < npt; i++) {
dk[tiisg * npt + i] = (half)row[tiisg * npt + i];
dv[tiisg * npt + i] = (half)vs[tiisg * npt + i];
}
return;
}
const uint npt = Di / 32;
device const float *src = ik + (uint64_t)tok * Di;
device float *dst = ik_cache + (uint64_t)pos * Di;
for (uint i = 0; i < npt; i++) dst[tiisg * npt + i] = src[tiisg * npt + i];
}
struct ds4_metal_args_qwen4_idx_block {
uint32_t block0; /* first block to (re)build */
uint32_t n_blocks;
@@ -3787,6 +3841,99 @@ kernel void kernel_qwen4_mtp_stage(
}
}
/* Separate embedding/hidden dispatches keep every row in the SIMD batch
* active. Each walks only its E-wide half, with the original mul_mv_ext
* float4 dots and shuffle tree; weight/concat strides remain full-width. */
template<short r1ptg, bool emb>
kernel void kernel_qwen4_mtp_project_q8(
constant ds4_metal_args_mul_mv_ext & args,
device const char *src0,
device const char *src1,
device char *dst,
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
const short NSG = FC_mul_mv_nsg;
const short nxpsg = FC_mul_mv_nxpsg;
const short chpt = 4, chpb = 8;
const short nypsg = 32 / nxpsg;
const short tx = tiisg % nxpsg, ty = tiisg / nxpsg;
const int i01 = tgpig.x * (nypsg * NSG) + nypsg * sgitg + ty;
const int E = args.ne00 / 2;
const int k0 = emb ? 0 : E;
int out_row[r1ptg];
device const float4 *y4[r1ptg];
for (short ir1 = 0; ir1 < r1ptg; ir1++) {
out_row[ir1] = emb ? ir1 * 5 : (int)tgpig.y * 5 + 1 + ir1;
y4[ir1] = (device const float4 *)(src1 + (uint64_t)out_row[ir1] * args.nb11 +
(uint64_t)k0 * sizeof(float)) + tx;
}
device const block_q8_0 *wrow = i01 < args.ne01 ?
(device const block_q8_0 *)(src0 + (uint64_t)i01 * args.nb01) :
(device const block_q8_0 *)src0;
device const block_q8_0 *xq = wrow + k0 / 32 + tx / chpb;
float sumf[r1ptg] = { [0 ... r1ptg - 1] = 0.0f };
short cch = tx % chpb;
for (int ich = tx; 4 * ich < E; ich += chpt * nxpsg) {
float4 lx[chpt];
#pragma unroll(chpt)
for (short ch = 0; ch < chpt; ch++) {
dequantize_q8_0_t4(xq, cch, lx[ch]);
cch += nxpsg;
if (cch >= chpb) {
xq += cch / chpb;
cch %= chpb;
}
}
#pragma unroll(chpt)
for (short ch = 0; ch < chpt; ch++) {
#pragma unroll(r1ptg)
for (short ir1 = 0; ir1 < r1ptg; ir1++) {
sumf[ir1] += dot(lx[ch], y4[ir1][ch * nxpsg]);
}
}
#pragma unroll(r1ptg)
for (short ir1 = 0; ir1 < r1ptg; ir1++) y4[ir1] += chpt * nxpsg;
}
for (short ir1 = 0; ir1 < r1ptg; ir1++) {
if (nxpsg >= 32) sumf[ir1] += simd_shuffle_down(sumf[ir1], 16);
if (nxpsg >= 16) sumf[ir1] += simd_shuffle_down(sumf[ir1], 8);
if (nxpsg >= 8) sumf[ir1] += simd_shuffle_down(sumf[ir1], 4);
if (nxpsg >= 4) sumf[ir1] += simd_shuffle_down(sumf[ir1], 2);
if (nxpsg >= 2) sumf[ir1] += simd_shuffle_down(sumf[ir1], 1);
}
/* The zero half contributes nothing for finite scales. Any Inf/NaN
* scale makes its original dot NaN, regardless of the packed quants.
* Read only those scale words, once per row across its eight lanes. */
uint nonfinite = 0u;
device const block_q8_0 *zero_half = wrow + (emb ? E / 32 : 0);
for (int b = tx; b < E / 32; b += nxpsg)
nonfinite |= (as_type<ushort>(zero_half[b].d) & 0x7c00u) == 0x7c00u;
if (nxpsg >= 32) nonfinite |= simd_shuffle_down(nonfinite, 16);
if (nxpsg >= 16) nonfinite |= simd_shuffle_down(nonfinite, 8);
if (nxpsg >= 8) nonfinite |= simd_shuffle_down(nonfinite, 4);
if (nxpsg >= 4) nonfinite |= simd_shuffle_down(nonfinite, 2);
if (nxpsg >= 2) nonfinite |= simd_shuffle_down(nonfinite, 1);
if (tx == 0 && i01 < args.ne01) {
for (short ir1 = 0; ir1 < r1ptg; ir1++) {
device float *out = (device float *)dst + (uint64_t)out_row[ir1] * args.ne0;
out[i01] = nonfinite ? as_type<float>(0x7fc00000u) : sumf[ir1];
}
}
}
template [[host_name("kernel_qwen4_mtp_project_q8_emb_1")]]
kernel void kernel_qwen4_mtp_project_q8<1, true>(constant ds4_metal_args_mul_mv_ext &, device const char *,
device const char *, device char *, uint3, ushort, ushort);
template [[host_name("kernel_qwen4_mtp_project_q8_emb_2")]]
kernel void kernel_qwen4_mtp_project_q8<2, true>(constant ds4_metal_args_mul_mv_ext &, device const char *,
device const char *, device char *, uint3, ushort, ushort);
template [[host_name("kernel_qwen4_mtp_project_q8_emb_3")]]
kernel void kernel_qwen4_mtp_project_q8<3, true>(constant ds4_metal_args_mul_mv_ext &, device const char *,
device const char *, device char *, uint3, ushort, ushort);
template [[host_name("kernel_qwen4_mtp_project_q8_hidden_4")]]
kernel void kernel_qwen4_mtp_project_q8<4, false>(constant ds4_metal_args_mul_mv_ext &, device const char *,
device const char *, device char *, uint3, ushort, ushort);
struct ds4_metal_args_qwen4_mtp_combine {
uint32_t n_embd;
uint32_t n_hc;
tests/test_qwen4_kernels.c
+193 -1
View File
@@ -2458,6 +2458,184 @@ static void test_moe_mm_tiles_exact(arena_t *a, uint32_t down_type) {
ds4_gpu_tensor_free(glists); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gx);
}
/* The cache-only path must reproduce full attention preparation, including
* half stores and three-axis RoPE, without touching neighboring cache rows. */
static void test_mtp_cache_prep(arena_t *a) {
enum { H = 24, Hkv = 2, D = 256, Hi = 4, Di = 128, cap = 12, guard = 32 };
double *shadow = NULL;
const uint64_t gq = arena_f32(a, D, &shadow, .5f, 1.5f); free(shadow);
const uint64_t gk = arena_f32(a, D, &shadow, .5f, 1.5f); free(shadow);
const uint64_t giq = arena_f32(a, Di, &shadow, .5f, 1.5f); free(shadow);
uint32_t positions[cap][4];
for (unsigned p = 0; p < cap; p++) {
positions[p][0] = 101u + p * 3u;
positions[p][1] = 307u + p * 5u;
positions[p][2] = 911u + p * 7u;
positions[p][3] = 0;
}
ds4_gpu_tensor *pos3 = ds4_gpu_tensor_alloc(sizeof(positions));
require_ok(pos3 && ds4_gpu_tensor_write(pos3, 0, positions, sizeof(positions)), "MTP cache mRoPE positions");
float freq[32];
for (unsigned i = 0; i < 32u; i++) freq[i] = powf(10000000.f, -(float)i / 32.f) * .875f;
for (unsigned rope = 0; rope < 2u; rope++) {
ds4_gpu_qwen4_set_rope(rope ? freq : NULL, rope ? 32u : 0u, rope ? 1.125f : 1.f);
for (uint32_t T = 1u; T <= 3u; T++) {
const uint64_t counts[5] = {T * H * 2u * D, T * Hkv * D, T * Hkv * D, T * Hi * Di, T * Di};
float *input[5]; ds4_gpu_tensor *in[5];
for (unsigned i = 0; i < 5u; i++) { input[i] = rand_vec(counts[i], 1.f); in[i] = upload(input[i], counts[i]); }
ds4_gpu_tensor *q = upload(NULL, T * H * D), *gate = upload(NULL, T * H * D);
ds4_gpu_tensor *iq = upload(NULL, T * Hi * Di);
const uint64_t row_bytes[3] = {Hkv * D * 2u, Hkv * D * 2u, Di * 4u};
ds4_gpu_tensor *back[2][3], *cache[2][3];
unsigned char *seed[3], *got[3], *want[3];
for (unsigned i = 0; i < 3u; i++) {
const uint64_t bytes = cap * row_bytes[i] + 2u * guard;
seed[i] = malloc((size_t)bytes); got[i] = malloc((size_t)bytes); want[i] = malloc((size_t)bytes);
require_ok(seed[i] && got[i] && want[i], "MTP cache byte buffers");
memset(seed[i], 0x55 + (int)i, (size_t)bytes);
for (unsigned side = 0; side < 2u; side++) {
back[side][i] = ds4_gpu_tensor_alloc(bytes);
cache[side][i] = ds4_gpu_tensor_view(back[side][i], guard, cap * row_bytes[i]);
require_ok(back[side][i] && cache[side][i], "guarded MTP cache allocation");
}
}
for (unsigned where = 0; where < 2u; where++) {
const uint32_t pos = where ? cap - T : 3u;
for (unsigned side = 0; side < 2u; side++) for (unsigned i = 0; i < 3u; i++)
require_ok(ds4_gpu_tensor_write(back[side][i], 0, seed[i], cap * row_bytes[i] + 2u * guard),
"poison complete MTP cache and guards");
require_ok(ds4_gpu_begin_commands() &&
ds4_gpu_qwen4_attn_prep_tensor(q, gate, cache[0][0], cache[0][1], iq, cache[0][2],
in[0], in[1], in[2], in[3], in[4], pos3, a->base, a->size, gq, gk, giq,
T, H, Hkv, D, 64u, Hi, Di, pos, cap, 10000000.f, 1.e-6f) &&
ds4_gpu_qwen4_attn_cache_prep_tensor(cache[1][0], cache[1][1], cache[1][2],
in[1], in[2], in[4], pos3, a->base, a->size, gk,
T, Hkv, D, 64u, Di, pos, cap, 10000000.f, 1.e-6f) &&
ds4_gpu_end_commands(), "full and cache-only attention preparation");
for (unsigned i = 0; i < 3u; i++) {
const uint64_t bytes = cap * row_bytes[i] + 2u * guard;
const uint64_t begin = guard + pos * row_bytes[i], end = begin + T * row_bytes[i];
require_ok(ds4_gpu_tensor_read(back[0][i], 0, want[i], bytes) &&
ds4_gpu_tensor_read(back[1][i], 0, got[i], bytes), "MTP cache comparison read");
require_ok(!memcmp(got[i], want[i], (size_t)bytes), "MTP K/V/IK bitwise full-prep oracle");
require_ok(!memcmp(got[i], seed[i], (size_t)begin) &&
!memcmp(got[i] + end, seed[i] + end, (size_t)(bytes - end)),
"MTP cache untouched prefix, suffix and allocation guards");
require_ok(memcmp(got[i] + begin, seed[i] + begin, (size_t)(end - begin)) != 0,
"MTP cache rows actually written");
}
}
require_ok(!ds4_gpu_qwen4_attn_cache_prep_tensor(cache[1][0], cache[1][0], cache[1][2],
in[1], in[2], in[4], pos3, a->base, a->size, gk, T, Hkv, D, 64u, Di, 3u, cap, 1.e7f, 1.e-6f),
"MTP cache rejects overlapping K/V outputs");
require_ok(!ds4_gpu_qwen4_attn_cache_prep_tensor(cache[1][0], cache[1][1], cache[1][2],
in[1], in[2], in[4], pos3, a->base, a->size, gk, T, Hkv, D, 64u, Di, cap, cap, 1.e7f, 1.e-6f),
"MTP cache rejects an out-of-range row");
require_ok(!ds4_gpu_qwen4_attn_cache_prep_tensor(cache[1][0], cache[1][1], cache[1][2],
in[1], in[2], in[4], pos3, a->base, a->size, gk, T, Hkv, D, 63u, Di, 3u, cap, 1.e7f, 1.e-6f),
"MTP cache rejects odd rotary width");
for (unsigned i = 0; i < 5u; i++) {
float *unchanged = download(in[i], counts[i]);
check_exact_f32("MTP cache input remains unchanged", unchanged, input[i], counts[i]);
free(unchanged); free(input[i]); ds4_gpu_tensor_free(in[i]);
}
for (unsigned i = 0; i < 3u; i++) {
for (unsigned side = 0; side < 2u; side++) {
ds4_gpu_tensor_free(cache[side][i]); ds4_gpu_tensor_free(back[side][i]);
}
free(want[i]); free(got[i]); free(seed[i]);
}
ds4_gpu_tensor_free(iq); ds4_gpu_tensor_free(gate); ds4_gpu_tensor_free(q);
printf(" MTP cache prep T=%u rope=%u: K/V/IK, mRoPE axes, guards and input immutability exact\n", T, rope);
}
}
uint32_t unchanged[cap][4];
require_ok(ds4_gpu_tensor_read(pos3, 0, unchanged, sizeof(unchanged)) &&
!memcmp(positions, unchanged, sizeof(positions)), "MTP cache positions remain unchanged");
ds4_gpu_qwen4_set_rope(NULL, 0u, 1.f);
ds4_gpu_tensor_free(pos3);
}
static void test_mtp_project(arena_t *a) {
enum { E = 2560, hc = 4, guard = 13 };
const uint32_t width = 2u * E, blocks = width / 32u;
const uint64_t weight_bytes = (uint64_t)E * blocks * 34u;
double *shadow = NULL;
const uint64_t finite_off = arena_q8_0(a, E, width, &shadow, .05f); free(shadow);
const uint64_t special_off = arena_alloc(a, weight_bytes);
memcpy(a->base + special_off, a->base + finite_off, (size_t)weight_bytes);
/* Non-finite scales occur once in each half. Even an inactive half must
* retain zero*Inf/NaN behavior; it cannot blindly skip those blocks. */
const uint16_t scales[] = {0x7c00u, 0xfc00u, 0x7e00u, 0xfe00u};
for (unsigned r = 0; r < 4u; r++) {
const uint64_t block = (uint64_t)r * blocks + (r & 1u ? E / 32u : 0u);
memcpy(a->base + special_off + block * 34u, scales + r, sizeof(scales[r]));
}
const char *env_name = "DS4_METAL_Q8_MV_EXT_MAX_TOKENS";
const char *incoming = getenv(env_name);
char *saved = incoming ? strdup(incoming) : NULL;
require_ok(!incoming || saved, "save EH matvec threshold");
for (uint32_t T = 1u; T <= 3u; T++) {
const uint32_t rows = T * (hc + 1u);
const uint64_t input_n = (uint64_t)rows * width, output_n = (uint64_t)rows * E;
float *cat = calloc((size_t)input_n, sizeof(float));
float *want = malloc((size_t)(output_n + guard) * sizeof(float));
float *got = malloc((size_t)(output_n + guard) * sizeof(float));
require_ok(cat && want && got, "EH projection comparison allocation");
ds4_gpu_tensor *x = upload(NULL, input_n), *ref = upload(NULL, output_n + guard);
ds4_gpu_tensor *out = upload(NULL, output_n + guard);
for (unsigned mode = 0; mode < 4u; mode++) {
for (uint32_t r = 0; r < rows; r++) {
const uint32_t active = r % (hc + 1u) ? E : 0u;
for (uint32_t c = 0; c < width; c++) {
const bool nonzero = c >= active && c < active + E;
cat[(uint64_t)r * width + c] = mode == 1u ? 0.f : mode == 2u ?
((c + r) & 1u ? -0.f : 0.f) : nonzero ? frand() : 0.f;
}
}
require_ok(ds4_gpu_tensor_write(x, 0, cat, input_n * sizeof(float)), "EH zero-half input");
const uint64_t off = mode == 3u ? special_off : finite_off;
for (unsigned fallback = 0; fallback < 2u; fallback++) {
require_ok(setenv(env_name, fallback ? "2" : "16", 1) == 0, "select EH matvec/fallback path");
require_ok(ds4_gpu_tensor_fill_f32(ref, 127.25f, output_n + guard) &&
ds4_gpu_tensor_fill_f32(out, 127.25f, output_n + guard) &&
ds4_gpu_begin_commands() &&
ds4_gpu_qwen4_matmul_q8_0_tensor(ref, a->base, a->size, off, width, E, x, rows) &&
ds4_gpu_qwen4_mtp_project_tensor(out, x, a->base, a->size, off, E, hc, T) &&
ds4_gpu_end_commands() &&
ds4_gpu_tensor_read(ref, 0, want, (output_n + guard) * sizeof(float)) &&
ds4_gpu_tensor_read(out, 0, got, (output_n + guard) * sizeof(float)),
"original and zero-skipping EH projection");
unsigned special = 0;
for (uint64_t i = 0; i < output_n + guard; i++) {
uint32_t a_bits, b_bits;
memcpy(&a_bits, got + i, sizeof(a_bits)); memcpy(&b_bits, want + i, sizeof(b_bits));
if ((b_bits & 0x7fffffffu) > 0x7f800000u) {
require_ok((a_bits & 0x7fffffffu) > 0x7f800000u, "EH preserves NaN classification");
special++;
} else {
require_ok(a_bits == b_bits, "EH finite/Inf output and signed-zero bits");
special += (b_bits & 0x7fffffffu) == 0x7f800000u;
}
}
require_ok(mode == 3u ? special > 0u : special == 0u, "EH exceptional scales actually exercised");
for (uint64_t i = output_n; i < output_n + guard; i++)
require_ok(got[i] == 127.25f, "EH projection allocation guard");
}
float *unchanged = download(x, input_n);
require_ok(!memcmp(unchanged, cat, (size_t)input_n * sizeof(float)), "EH input including zeros remains unchanged");
free(unchanged);
}
require_ok(!ds4_gpu_qwen4_mtp_project_tensor(out, x, a->base, a->size, finite_off, E, hc, 0u),
"EH rejects zero tokens");
ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(ref); ds4_gpu_tensor_free(x);
free(got); free(want); free(cat);
printf(" MTP EH T=%u: full Q8 oracle, zero/signed-zero, non-finite scales and fallback exact\n", T);
}
require_ok((saved ? setenv(env_name, saved, 1) : unsetenv(env_name)) == 0, "restore EH matvec threshold");
free(saved);
}
/* MTP input staging: cat rows [rms(e)*g_e | 0] and [0 | rms(R_s)*g_h_s]
* (full-row or per-stream RMS), then R_out = proj[0] + proj[1+s]. */
static void test_mtp(arena_t *a, uint32_t E, uint32_t hc) {
@@ -3210,10 +3388,12 @@ int main(void) {
* ONLY accepts 1 (both), iq2 or hc; BENCH=1 adds paired resident timings. */
const char *reuse_only = getenv("DS4_TEST_QWEN4_M1_REUSE_ONLY");
const char *reuse_bench = getenv("DS4_TEST_QWEN4_M1_REUSE_BENCH");
const char *mtp_only = getenv("DS4_TEST_QWEN4_MTP_OPT_ONLY");
const bool run_reuse = reuse_only && reuse_only[0] && strcmp(reuse_only, "0") != 0;
const bool bench_reuse = reuse_bench && reuse_bench[0] && strcmp(reuse_bench, "0") != 0;
const bool run_mtp = mtp_only && mtp_only[0] && strcmp(mtp_only, "0") != 0;
arena_t arena;
arena.size = (uint64_t)(run_reuse || bench_reuse ? 96u : 1536u) << 20;
arena.size = (uint64_t)(run_reuse || bench_reuse || run_mtp ? 96u : 1536u) << 20;
arena.base = mmap(NULL, arena.size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
arena.used = 0;
if (arena.base == MAP_FAILED) { perror("mmap"); return 1; }
@@ -3227,6 +3407,16 @@ int main(void) {
munmap(arena.base, arena.size);
return 0;
}
if (run_mtp) {
test_mtp_cache_prep(&arena);
test_mtp_project(&arena);
test_mtp(&arena, 2560u, 4u);
test_mtp(&arena, 64u, 4u);
ds4_gpu_cleanup();
munmap(arena.base, arena.size);
puts("all Qwen MTP cache/projection kernel tests passed");
return 0;
}
if (getenv("DS4_TEST_QWEN4_DECODE_FUSIONS")) { test_decode_fusions(&arena); return 0; }
if (getenv("DS4_TEST_QWEN4_MV_EXACT")) {
test_moe_types(&arena, 8, 6, 2560, 640, 1, 16u, 10u);
@@ -3339,6 +3529,8 @@ int main(void) {
test_multi_gemv(&arena, 2560, 2);
test_multi_gemv(&arena, 64, 3);
printf("mtp\n");
test_mtp_cache_prep(&arena);
test_mtp_project(&arena);
test_mtp(&arena, 2560, 4);
test_mtp(&arena, 64, 4);
test_hc_norm_reuse(&arena);
tests/test_qwen4_mtp_prefill.c
+227 -10
View File
@@ -97,6 +97,112 @@ static void poison_nextn(ds4_session *s, unsigned char value) {
}
}
/* Frozen pre-optimization oracle: keep the full zero-padded EH projection
* and all-row MoE here. Sharing the optimized input/steps helpers would let
* both sides reproduce the same regression. Full attention remains the
* independent oracle for the separate cache-only attention preparation. */
static bool reference_mtp_input(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
uint32_t row, const int *next_tokens, uint32_t T) {
const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC;
const ds4_layer_weights *l = &w->layer[DS4_N_LAYER - 1u];
for (uint32_t t = 0; t < T; t++) {
if (next_tokens[t] < 0 || next_tokens[t] >= (int)DS4_N_VOCAB) return false;
qwen4_ref_row(m, w->token_embd, (uint64_t)next_tokens[t], g->host_row + (uint64_t)t * E);
}
if (!ds4_gpu_tensor_write(g->mtp_e, 0, g->host_row, (uint64_t)T * E * sizeof(float)) ||
!glm_graph_begin_commands_if_needed()) return false;
ds4_gpu_tensor *R_save = g->R;
bool ok = true;
const uint64_t emb_bytes = (uint64_t)E * sizeof(float);
const uint64_t cat_bytes = (hc + 1u) * 2u * emb_bytes;
const uint64_t proj_bytes = (hc + 1u) * emb_bytes;
for (uint32_t t = 0; t < T && ok; t++) {
ds4_gpu_tensor *e_row = ds4_gpu_tensor_view(g->mtp_e, t * emb_bytes, emb_bytes);
ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(R_save, (uint64_t)(row + t) * hc * emb_bytes, hc * emb_bytes);
ds4_gpu_tensor *cat_row = ds4_gpu_tensor_view(g->mtp_cat, t * cat_bytes, cat_bytes);
ok = e_row && R_row && cat_row &&
ds4_gpu_qwen4_mtp_stage_tensor(cat_row, e_row, R_row, m->map, m->size,
l->nextn_enorm->abs_offset, l->nextn_hnorm->abs_offset,
E, hc, DS4_RMS_EPS);
ds4_gpu_tensor_free(cat_row);
ds4_gpu_tensor_free(R_row);
ds4_gpu_tensor_free(e_row);
}
if (ok) ok = qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, T * (hc + 1u));
for (uint32_t t = 0; t < T && ok; t++) {
ds4_gpu_tensor *proj_row = ds4_gpu_tensor_view(g->mtp_proj, t * proj_bytes, proj_bytes);
ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(g->mtp_R, t * hc * emb_bytes, hc * emb_bytes);
ok = proj_row && R_row && ds4_gpu_qwen4_mtp_combine_tensor(R_row, proj_row, E, hc);
ds4_gpu_tensor_free(R_row);
ds4_gpu_tensor_free(proj_row);
}
g->R = g->mtp_R;
if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, T);
g->R = R_save;
return ok;
}
static bool reference_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
uint32_t row, const int *next_tokens, uint32_t T, uint32_t idx, bool want_logits,
float *logits_out, int *draft_out) {
if (!g->mtp_R || T == 0u || T > 3u || idx > g->mtp_pos || idx > g->ctx_cap || T > g->ctx_cap - idx ||
row > g->cap_tokens || T > g->cap_tokens - row) return false;
const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC;
const uint32_t il = DS4_N_LAYER - 1u;
const ds4_layer_weights *l = &w->layer[il];
ds4_gpu_tensor *R_save = g->R;
const uint64_t emb_bytes = (uint64_t)E * sizeof(float);
bool ok = reference_mtp_input(g, m, w, row, next_tokens, T);
g->R = g->mtp_R;
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, 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 &&
(!argmax_env || strcmp(argmax_env, "0") != 0);
/* draft-only rows: host logits consumers always see the full head */
const bool gathered = gpu_argmax && qwen4_mtp_draft_head_load(g, m, w->output);
const uint32_t head_rows = gathered ? g->draft_rows : gpu_argmax ? qwen4_mtp_draft_rows() : DS4_N_VOCAB;
if (ok && want_logits) {
last = ds4_gpu_tensor_view(g->mtp_R, (T - 1u) * hc * emb_bytes, hc * emb_bytes);
g->R = last;
ok = last && qwen4_graph_hc_mix(g, m, l->nextn_hc_head_norm, l->nextn_hc_head_down, l->nextn_hc_head_up, NULL, 1) &&
(gathered ? ds4_gpu_qwen4_matmul_q8_0_weights_tensor(g->logits, g->draft_head, (uint32_t)w->output->dim[0],
head_rows, g->mixed) != 0
: qwen4_gemv_rows(g->logits, m, w->output, g->mixed, 1, head_rows));
}
g->R = R_save;
if (ok && gpu_argmax) ok = ds4_gpu_qwen4_argmax_tensor(g->mtp_argmax, g->mtp_argmax_tmp,
g->logits, head_rows) != 0;
if (!ds4_gpu_end_commands()) ok = false;
ds4_gpu_tensor_free(last);
if (ok && gpu_argmax) {
int32_t token = -1;
ok = ds4_gpu_tensor_read(g->mtp_argmax, 0, &token, sizeof(token)) != 0;
if (ok && gathered) {
if (token < 0 || (uint32_t)token >= g->draft_rows) ok = false;
else token = g->draft_ids[token];
}
if (ok && (token < 0 || token >= (int32_t)DS4_N_VOCAB)) ok = false;
if (ok) *draft_out = token;
} else if (ok && want_logits) {
float *dst = logits_out ? logits_out : g->host_logits;
ok = ds4_gpu_tensor_read(g->logits, 0, dst, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0;
if (ok && draft_out) *draft_out = sample_argmax(dst, DS4_N_VOCAB);
}
if (ok) {
g->mtp_pos = idx + T;
g->mtp_last_rows = T;
if (g->mtp_tail_valid && idx + T - 1u == g->mtp_tail_pos) {
g->mtp_tail_cached = true;
g->mtp_tail_next_token = next_tokens[T - 1u];
}
}
return ok;
}
/* Full MTP is the oracle. With a known lookahead it also fills the last row
* of an intermediate chunk, so graph_forward_tokens never needs to repair
* that boundary through the new cache-only helper. */
@@ -111,8 +217,8 @@ static unsigned reference_sync(ds4_session *s, const ds4_tokens *tokens,
need(g->mtp_pos + 1u == start && g->mtp_tail_valid, "reference append boundary");
ds4_gpu_tensor *saved = g->R;
g->R = g->mtp_tail_R;
const bool ok = qwen4_graph_mtp_step(g, &e->model, &e->weights,
0, tokens->v[start], start - 1u, false, NULL, NULL);
const bool ok = reference_mtp_steps(g, &e->model, &e->weights,
0, tokens->v + start, 1u, start - 1u, false, NULL, NULL);
g->R = saved;
need(ok, "full MTP append boundary oracle");
}
@@ -125,7 +231,7 @@ static unsigned reference_sync(ds4_session *s, const ds4_tokens *tokens,
if (pos + T == (uint32_t)tokens->len) rows--;
for (uint32_t row = 0; row < rows;) {
const uint32_t count = rows - row < 3u ? rows - row : 3u;
need(qwen4_graph_mtp_steps(g, &e->model, &e->weights, row,
need(reference_mtp_steps(g, &e->model, &e->weights, row,
tokens->v + pos + row + 1u, count, pos + row, false, NULL, NULL),
"full MTP prefix oracle");
row += count;
@@ -158,6 +264,109 @@ static void checked_sync(ds4_session *s, const ds4_tokens *tokens, frontier *f,
compare_frontier(s, f + n - 1u, true);
}
static void restore_mtp_case(ds4_session *s, const ds4_session_snapshot *snapshot,
const void *rows, bool verify) {
need(ds4_session_load_snapshot(s, snapshot, error, sizeof(error)) == 0,
"restore optimized-path comparison frontier");
/* Payloads retain the final trunk row, not the whole last batch. Reuse
* the same three actual trunk rows explicitly on both sides. */
need(ds4_gpu_tensor_write(s->qwen4_graph.R, 0, rows,
3ull * DS4_N_HC * DS4_N_EMBD * sizeof(float)), "restore identical trunk rows");
s->qwen4_graph.verify_rows_exact = verify;
ds4_gpu_qwen4_set_verify_rows_exact(verify);
}
static void compare_tensor(const char *what, const ds4_gpu_tensor *got,
const ds4_gpu_tensor *want, uint64_t offset, uint64_t bytes) {
void *a = malloc((size_t)bytes), *b = malloc((size_t)bytes);
need(a && b, "tensor comparison allocation");
need(ds4_gpu_tensor_read(got, offset, a, bytes) &&
ds4_gpu_tensor_read(want, offset, b, bytes), "read compared tensor range");
exact(what, a, b, (size_t)bytes);
free(b); free(a);
}
static void check_mtp_paths(ds4_session *live, ds4_session *reference,
const ds4_session_snapshot *snapshot, const int *next_tokens,
const void *rows) {
ds4_engine *e = live->engine;
ds4_qwen4_gpu_graph *g = &live->qwen4_graph, *r = &reference->qwen4_graph;
const uint64_t residual_bytes = (uint64_t)DS4_N_HC * DS4_N_EMBD * sizeof(float);
const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float);
float *want = malloc((size_t)logits_bytes), *got = malloc((size_t)logits_bytes);
need(want && got, "predictor logits allocation");
/* idx=13 means T3 closes pooled block [12,16), while T1/T2 retain a
* partial block. T3 exercises both prefill and exact-verify reductions. */
for (unsigned mode = 0; mode < 4u; mode++) {
const uint32_t T = mode < 3u ? mode + 1u : 3u;
const bool verify = mode == 3u;
int want_draft = -1, got_draft = -1;
restore_mtp_case(reference, snapshot, rows, verify);
need(reference_mtp_steps(r, &e->model, &e->weights, 0, next_tokens, T,
TEST_PREFIX, true, want, &want_draft), "original complete MTP rows");
frontier expected = capture(reference, 3u);
restore_mtp_case(live, snapshot, rows, verify);
need(qwen4_graph_mtp_cache_steps(g, &e->model, &e->weights, 0, next_tokens,
T, TEST_PREFIX), "optimized cache-only rows");
need(g->mtp_last_rows == 0, "cache-only rows are not recursive outputs");
compare_frontier(live, &expected, true);
compare_tensor("zero-skipping EH projection (cache)", g->mtp_proj, r->mtp_proj, 0,
(uint64_t)T * (DS4_N_HC + 1u) * DS4_N_EMBD * sizeof(float));
restore_mtp_case(live, snapshot, rows, verify);
need(qwen4_graph_mtp_steps(g, &e->model, &e->weights, 0, next_tokens, T,
TEST_PREFIX, true, got, &got_draft), "optimized last-row MTP output");
compare_frontier(live, &expected, true);
compare_tensor("zero-skipping EH projection (full)", g->mtp_proj, r->mtp_proj, 0,
(uint64_t)T * (DS4_N_HC + 1u) * DS4_N_EMBD * sizeof(float));
compare_tensor("last predictor residual", g->mtp_R, r->mtp_R,
(T - 1u) * residual_bytes, residual_bytes);
exact("all predictor logits", got, want, (size_t)logits_bytes);
need(got_draft == want_draft && got_draft == sample_argmax(got, DS4_N_VOCAB),
"last-row draft agrees with original complete MTP");
need(g->mtp_last_rows == T && g->mtp_pos == TEST_PREFIX + T,
"last-row layout and contiguous cache frontier");
free_frontier(&expected);
/* The recursive consumer must read row T-1 after earlier FFN rows
* are skipped. Its oracle uses the frozen input/full-layer helper,
* with the original last residual explicitly selected as input. */
const uint32_t idx = TEST_PREFIX + T;
uint32_t pos3[4] = {0};
int32_t delta = r->mrope_delta;
qwen4_mrope_pos(NULL, 0, idx, &delta, pos3);
need(ds4_gpu_tensor_write(r->pos3, (uint64_t)idx * sizeof(pos3), pos3, sizeof(pos3)),
"reference future position");
ds4_gpu_tensor *saved_R = r->R;
ds4_gpu_tensor *last = ds4_gpu_tensor_view(r->mtp_R, (T - 1u) * residual_bytes, residual_bytes);
need(last != NULL, "reference recursive residual view");
r->R = last;
const int parent = want_draft;
const bool ok = reference_mtp_steps(r, &e->model, &e->weights, 0, &parent,
1u, idx, true, want, &want_draft);
r->R = saved_R;
ds4_gpu_tensor_free(last);
need(ok, "original recursive predictor output");
need(qwen4_graph_mtp_chain_step(g, &e->model, &e->weights, got_draft, idx,
&got_draft), "optimized recursive predictor output");
need(ds4_gpu_tensor_read(g->logits, 0, got, logits_bytes), "recursive predictor logits");
exact("all recursive predictor logits", got, want, (size_t)logits_bytes);
need(got_draft == want_draft && g->mtp_last_rows == 1u,
"recursive draft and output layout");
compare_tensor("recursive predictor residual", g->mtp_R, r->mtp_R, 0, residual_bytes);
expected = capture(reference, 3u);
compare_frontier(live, &expected, true);
free_frontier(&expected);
printf("PASS optimized MTP T%u verify=%u: EH, KV, final residual, full logits and recursive draft\n",
T, verify);
}
g->verify_rows_exact = false;
r->verify_rows_exact = false;
ds4_gpu_qwen4_set_verify_rows_exact(false);
free(got); free(want);
}
static void close_logits(const float *got, const float *want) {
/* The existing 2/3-row verifier and one-row target have distinct
* reductions; use the established MTP batch tolerance, not text alone. */
@@ -179,6 +388,7 @@ static void check_cycles(ds4_session *live, ds4_session *target,
need(ds4_session_load_snapshot(live, snapshot, error, sizeof(error)) == 0 &&
ds4_session_load_snapshot(target, snapshot, error, sizeof(error)) == 0, "restore cycle frontier");
const bool forced = mode == 1u, stochastic = mode == 2u;
const bool speculative = qwen4_graph_fused(&live->qwen4_graph, 2u);
need(setenv("DS4_QWEN4_MTP_DEPTH", mode ? "3" : "2", 1) == 0, "set depth case");
need((forced ? setenv("DS4_QWEN4_SPEC_FORCE_ACCEPT", "1", 1) :
unsetenv("DS4_QWEN4_SPEC_FORCE_ACCEPT")) == 0, "set acceptance case");
@@ -204,9 +414,11 @@ static void check_cycles(ds4_session *live, ds4_session *target,
}
close_logits(live->logits, target->logits);
}
if (forced) need(saw_verify && saw_three, "depth-three verifier actually exercised");
if (forced) need(speculative ? saw_verify && saw_three : !saw_verify && !saw_three,
"depth-three verifier or explicit unfused fallback exercised");
printf("PASS MTP continuation %s, independent target logits\n",
forced ? "forced depth3 state fixture" : stochastic ? "exact sampling depth2" : "greedy depth2");
!speculative ? "unfused ordinary fallback" : forced ? "forced depth3 state fixture" :
stochastic ? "exact sampling depth2" : "greedy depth2");
}
static void check_payload_modes(ds4_session *live, ds4_session *reference,
@@ -272,9 +484,11 @@ int main(int argc, char **argv) {
if (argc != 2 && !streaming) {
fprintf(stderr, "usage: %s QWEN_GGUF [--ssd-streaming]\n", argv[0]); return 2;
}
const char *names[] = {"DS4_QWEN4_SPEC_FORCE_ACCEPT", "DS4_QWEN4_MTP_DEPTH"};
char *saved[2] = {NULL, NULL};
for (unsigned i = 0; i < 2; i++) {
const char *names[] = {"DS4_QWEN4_SPEC_FORCE_ACCEPT", "DS4_QWEN4_MTP_DEPTH",
"DS4_QWEN4_MTP_DRAFT_ROWS", "DS4_QWEN4_MTP_DRAFT_VOCAB",
"DS4_QWEN4_MTP_GPU_ARGMAX"};
char *saved[sizeof(names) / sizeof(names[0])] = {0};
for (unsigned i = 0; i < sizeof(names) / sizeof(names[0]); i++) {
const char *v = getenv(names[i]); saved[i] = v ? strdup(v) : NULL;
need(!v || saved[i], "save test environment"); need(unsetenv(names[i]) == 0, "clear test override");
}
@@ -289,7 +503,7 @@ int main(int argc, char **argv) {
need(live->qwen4_graph.cap_tokens == TEST_CHUNK, "requested chunk size");
ds4_tokens prompt = {0};
ds4_encode_chat_prompt(engine, NULL, "Count from one to ten and explain the pattern.", DS4_THINK_NONE, &prompt);
need(prompt.len >= TEST_APPEND, "enough real prompt token IDs");
need(prompt.len > TEST_APPEND, "enough real prompt and lookahead token IDs");
ds4_tokens prefix = prompt; prefix.len = 1;
frontier frames[MAX_FRAMES] = {{0}};
unsigned n = reference_sync(reference, &prefix, 0, frames);
@@ -321,6 +535,9 @@ int main(int argc, char **argv) {
poison_nextn(live, 0x33);
need(ds4_session_load_snapshot(live, &snapshot, error, sizeof(error)) == 0, "restore over stale cache");
compare_frontier(live, frames + n - 1u, false);
need(frames[n - 1u].residual_rows == 3u, "three real trunk rows for optimized-path oracle");
check_mtp_paths(live, reference, &snapshot, prompt.v + TEST_PREFIX + 1u,
frames[n - 1u].residual);
for (unsigned i = 0; i < n; i++) free_frontier(frames + i);
puts("PASS append boundary and checkpoint restore over reused state");
@@ -329,7 +546,7 @@ int main(int argc, char **argv) {
ds4_session_snapshot_free(&snapshot);
ds4_tokens_free(&other); ds4_tokens_free(&prompt);
ds4_session_free(reference); ds4_session_free(live); ds4_engine_close(engine);
for (unsigned i = 0; i < 2; i++) {
for (unsigned i = 0; i < sizeof(names) / sizeof(names[0]); i++) {
need((saved[i] ? setenv(names[i], saved[i], 1) : unsetenv(names[i])) == 0, "restore test environment");
free(saved[i]);
}