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

fix(qwen): preserve reference prefill arithmetic with explicit projection phases

Browse Source 
Short prompt chunks and tails were selecting kernels intended for independent
decode batches. Float-operand F16/Q8 tiles and split-K FP32 reductions changed
prefill arithmetic and introduced substantial logit drift against e37f185,
including different results for the same prompt at different chunk sizes.

Track prefill/decode projection policy explicitly on each Qwen graph:
- Route contiguous prompt chunks, including single-token tails, through the
  reference projection paths and an unsplit Metal dense GEMM entry point.
- Preserve existing optimized projection dispatch for ordinary decode, native
  session batches, MTP verification, predictor cache updates and draft chains.
- Scope and restore the phase on success and failure; mark copied batch row
  graphs as decode rather than inferring phase from row count or all_rows.
- Keep CUDA dispatch and the separate vision encoder unchanged. The new CUDA
  prefill entry point delegates to the existing dense function.

Add --prefill-reference to the real-model generation regression test. Compare
normal one-shot/session dispatch against legacy unsplit direct prefill, then
restore automatic decode on both sides. Cover 29-token short prompts,
575-token prompts with a 63-token tail, and two complete 128-token chunks.
Reject inherited reference overrides that could hide a dispatch regression.
Update the private MTP oracle calls for the graph-aware projection helpers.

Validation on Apple M1 Max, 32 GiB, Metal SSD streaming:
- ds4/server and relevant test builds pass; memory planner, Metal kernel,
  generation/reference and MTP cache/continuation suites pass.
- All 27 generation/reference vocabulary frontiers are bit-exact. The same
  test compiled against original 0024ca0 fails at the first frontier, as
  expected, confirming that it detects the prefill regression.
- Five configurations compared with e37f185 cover 85 frontiers and 21,107,200
  FP32 logits, using identical prompts and 16 teacher-forced continuation IDs.
  Worst prefill error falls from 0.48730135 to 1.90734863e-6; worst decode
  error falls from 1.37674665 to 0.000828742981. No argmax disagreements.
  Residual numerical differences remain; this is not bit-exact parity with
  e37f185 or a model-quality evaluation.
- Patched 5,760-token captures are bit-exact across chunk sizes 128 and 2048
  at all 17 frontiers, including the subsequent decode steps.
- Sixteen alternating SSD timing runs against 0024ca0 (four per variant per
  workload) show median prefill changes of -0.71% at 575/chunk128 and +2.54%
  at 5760/chunk2048. Decode medians change +1.59% and +8.12%, respectively;
  short 16-step samples and overlapping ranges do not establish a general
  speedup or isolate decode kernel performance.

CUDA/ROCm GPU execution and resident-mode validation were not performed.
Generated reports, binary captures and temporary logs are excluded.
This commit is contained in:
GiorgioOppo committed 2026-09-18 14:15:14 +02:00
1 parent 0024ca0803
commit 2d6a207bbb
6 files changed
+216 -107
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

ds4.c
+135 -85
View File
@@ -57609,7 +57609,15 @@ static const ds4_vision_span *qwen4_fake_spans(size_t *count) {
X(moe_lists) X(moe_counts) X(sh_gate) X(sh_up) X(sh_mid) X(sh_out) \
X(hc_u) X(hc_lo_act)
/* Matrix row count alone does not distinguish prompt tails from independent
* decode rows. Keep the policy local to the graph and scoped to each forward. */
typedef enum {
QWEN4_PROJECTION_DECODE = 0,
QWEN4_PROJECTION_PREFILL,
} qwen4_projection_phase;
typedef struct ds4_qwen4_gpu_graph {
qwen4_projection_phase projection_phase;
uint32_t ctx_cap;
uint32_t pos;
uint32_t cap_tokens;
@@ -58181,12 +58189,13 @@ static void qwen4_graph_reset(ds4_qwen4_gpu_graph *g) {
}
/* rows > 0 limits the product to a contiguous leading prefix of w. */
static bool qwen4_gemv_rows(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *w,
static bool qwen4_gemv_rows(const ds4_qwen4_gpu_graph *g, ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *w,
const ds4_gpu_tensor *x, uint32_t n_tok, uint64_t rows) {
const uint64_t in_dim = w->dim[0], full_dim = w->ndim >= 2 ? w->dim[1] : 1u;
const uint64_t out_dim = rows && rows < full_dim ? rows : full_dim;
int rc = 0;
#if !defined(__APPLE__)
(void)g;
/* Qwen's recurrent graph keeps activations in FP32. The generic CUDA
* projections can round them to half or quantize them for other models. */
if (in_dim <= UINT32_MAX && out_dim <= UINT32_MAX)
@@ -58194,23 +58203,29 @@ static bool qwen4_gemv_rows(ds4_gpu_tensor *out, const ds4_model *m, const ds4_t
w->type, n_tok, (uint32_t)in_dim, (uint32_t)out_dim);
#else
/* Small F16 batches use float operands. Two/three-row verification and
* large prefills retain their existing kernels. The legacy switch is an
* arithmetic/performance control, not a user tuning option. */
/* Prefill keeps the reference operand precision and K accumulation even
* for short prompt tails. Float-operand F16/Q8 tiles and split-K belong
* to decode batches; selecting them only by row count moves prompt logits.
* The legacy switch remains a diagnostic control for decode batches. */
const bool prefill = g->projection_phase == QWEN4_PROJECTION_PREFILL;
const bool legacy = getenv("DS4_QWEN4_DENSE_MM_LEGACY") != NULL;
const bool f16_batch = !legacy && n_tok > 3u && w->type == DS4_TENSOR_F16 &&
const bool f16_batch = !prefill && !legacy && n_tok > 3u && w->type == DS4_TENSOR_F16 &&
n_tok <= 64u && (out_dim <= 512u || n_tok > 8u);
if (((n_tok > 8u && w->type == DS4_TENSOR_F32) || f16_batch) &&
(in_dim % 32) == 0) {
rc = ds4_gpu_qwen4_dense_mm_tensor(out, x, m->map, m->size, w->abs_offset, w->type, n_tok,
(uint32_t)in_dim, (uint32_t)out_dim);
rc = prefill ? ds4_gpu_qwen4_dense_mm_prefill_tensor(
out, x, m->map, m->size, w->abs_offset, w->type, n_tok,
(uint32_t)in_dim, (uint32_t)out_dim)
: ds4_gpu_qwen4_dense_mm_tensor(
out, x, m->map, m->size, w->abs_offset, w->type, n_tok,
(uint32_t)in_dim, (uint32_t)out_dim);
if (rc) return true;
}
/* Q8 decode batches of 8 or 16 rows on fp32 simdgroup matrices: the
* per-token kernel reads the matrix once per four rows and streams x
* once per four weight rows; this reads the weights once and x once per
* 32 rows. DS4_QWEN4_NO_BATCH_MM=1 keeps the per-token kernel for A/B. */
if (!legacy && w->type == DS4_TENSOR_Q8_0 && n_tok >= 5u && n_tok <= 32u &&
if (!prefill && !legacy && w->type == DS4_TENSOR_Q8_0 && n_tok >= 5u && n_tok <= 32u &&
(in_dim % 32) == 0 && (out_dim % 32) == 0 && getenv("DS4_QWEN4_NO_BATCH_MM") == NULL &&
ds4_gpu_tensor_bytes(x) >= 32u * in_dim * sizeof(float) &&
ds4_gpu_tensor_bytes(out) >= 32u * out_dim * sizeof(float)) {
@@ -58253,9 +58268,9 @@ static bool qwen4_gemv_rows(ds4_gpu_tensor *out, const ds4_model *m, const ds4_t
return rc != 0;
}
static bool qwen4_gemv(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *w,
static bool qwen4_gemv(const ds4_qwen4_gpu_graph *g, ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *w,
const ds4_gpu_tensor *x, uint32_t n_tok) {
return qwen4_gemv_rows(out, m, w, x, n_tok, 0u);
return qwen4_gemv_rows(g, out, m, w, x, n_tok, 0u);
}
/* The MTP draft only needs its argmax, and the tokenizer assigns ids in
@@ -58352,12 +58367,12 @@ static bool qwen4_graph_hc_mix(ds4_qwen4_gpu_graph *g, const ds4_model *m,
bool ok = ds4_gpu_qwen4_hc_norm_tensor(g->xn, g->inj, g->R, m->map, m->size, norm->abs_offset,
inject ? inject->abs_offset : 0, inject ? inject->type : DS4_TENSOR_F32,
T, DS4_N_EMBD, DS4_N_HC, inject ? DS4_N_HC : 0u, DS4_RMS_EPS) &&
qwen4_gemv(g->lo, m, down, g->xn, T);
qwen4_gemv(g, g->lo, m, down, g->xn, T);
if (T > 8u &&
up->type != DS4_TENSOR_Q8_0) {
/* prefill: the up projection as a GEMM over the activated low-rank rows */
return ok && ds4_gpu_qwen4_hc_lo_act_tensor(g->hc_lo_act, g->lo, T, DS4_N_HC, DS4_N_HC_LOWRANK) &&
qwen4_gemv(g->hc_u, m, up, g->hc_lo_act, T) &&
qwen4_gemv(g, g->hc_u, m, up, g->hc_lo_act, T) &&
ds4_gpu_qwen4_hc_mix_rows_tensor(g->mixed, g->hc_u, g->xn, T, DS4_N_EMBD, DS4_N_HC);
}
if (T == 3u && g->verify_rows_exact) {
@@ -58401,8 +58416,8 @@ static bool qwen4_graph_linear(ds4_qwen4_gpu_graph *g, const ds4_model *m, const
l->lin_qkv->abs_offset, l->lin_gate->abs_offset, DS4_N_EMBD,
l->lin_qkv->dim[1], l->lin_gate->dim[1], g->mixed, T) != 0;
}
if (!paired && !(qwen4_gemv(g->qkv, m, l->lin_qkv, g->mixed, T) &&
qwen4_gemv(g->z, m, l->lin_gate, g->mixed, T))) {
if (!paired && !(qwen4_gemv(g, g->qkv, m, l->lin_qkv, g->mixed, T) &&
qwen4_gemv(g, g->z, m, l->lin_gate, g->mixed, T))) {
return false;
}
bool ok;
@@ -58414,8 +58429,8 @@ static bool qwen4_graph_linear(ds4_qwen4_gpu_graph *g, const ds4_model *m, const
DS4_N_EMBD, g->snap_after_first ? g->snap_lin_hist[il] : NULL, 0u,
g->snap_after_second ? g->snap2_lin_hist[il] : NULL, 1u) != 0;
} else {
ok = qwen4_gemv(g->ga, m, l->lin_alpha, g->mixed, T) &&
qwen4_gemv(g->gb, m, l->lin_beta, g->mixed, T) &&
ok = qwen4_gemv(g, g->ga, m, l->lin_alpha, g->mixed, T) &&
qwen4_gemv(g, g->gb, m, l->lin_beta, g->mixed, T) &&
ds4_gpu_qwen4_conv_stream_tensor(g->qkv, g->layer_lin_hist[il], m->map, m->size, l->lin_conv->abs_offset,
T, conv_dim, DS4_N_LIN_CONV, true) &&
ds4_gpu_qwen4_gdn_prep_tensor(g->qkv, g->ga, g->gb, m->map, m->size, l->lin_a->abs_offset,
@@ -58432,7 +58447,7 @@ static bool qwen4_graph_linear(ds4_qwen4_gpu_graph *g, const ds4_model *m, const
ok = ds4_gpu_qwen4_gdn_out_tensor(g->lin_o, g->z, m->map, m->size, l->lin_norm->abs_offset, T,
DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_RMS_EPS) != 0;
}
if (ok) ok = qwen4_gemv(g->blk, m, l->lin_out, g->lin_o, T);
if (ok) ok = qwen4_gemv(g, g->blk, m, l->lin_out, g->lin_o, T);
return ok;
}
@@ -58542,7 +58557,7 @@ static bool qwen4_graph_attention_cache(ds4_qwen4_gpu_graph *g, const ds4_model
{
/* 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);
bool ok = cache_only || qwen4_gemv(g, 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, need_iq ? g->iq : g->ik, g->ik };
const uint64_t offs[4] = { l->attn_k->abs_offset, l->attn_v->abs_offset,
@@ -58557,10 +58572,10 @@ static bool qwen4_graph_attention_cache(ds4_qwen4_gpu_graph *g, const ds4_model
ok = ds4_gpu_qwen4_multi_gemv_tensor(g->mixed, T, DS4_N_EMBD, need_iq ? 4u : 3u,
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) &&
(!need_iq || qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, T)) &&
qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T);
ok = qwen4_gemv(g, g->kp, m, l->attn_k, g->mixed, T) &&
qwen4_gemv(g, g->vp, m, l->attn_v, g->mixed, T) &&
(!need_iq || qwen4_gemv(g, g->iq, m, l->indexer_q_proj, g->mixed, T)) &&
qwen4_gemv(g, g->ik, m, l->indexer_k_proj, g->mixed, T);
}
if (!ok) return false;
}
@@ -58620,7 +58635,7 @@ static bool qwen4_graph_attention(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
uint32_t il, uint32_t pos0, uint32_t T) {
return qwen4_graph_attention_cache(g, m, l, il, pos0, T, false) &&
qwen4_graph_attention_core_rows(g, il, pos0, T) &&
qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T);
qwen4_gemv(g, g->blk, m, l->attn_output, g->attn_o, T);
}
/* router GEMV, top-k (+ shared gate logit), experts with the shared expert as
@@ -58639,7 +58654,7 @@ static bool qwen4_graph_moe_impl(ds4_qwen4_gpu_graph *g, const ds4_model *m, con
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) &&
bool ok = qwen4_gemv(g, g->router, m, l->ffn_gate_inp, g->mixed, T) &&
ds4_gpu_qwen4_router_topk_tensor(g->selected, g->weights, g->router, g->mixed, m->map, m->size,
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) &&
@@ -58729,10 +58744,10 @@ static bool qwen4_graph_moe_impl(ds4_qwen4_gpu_graph *g, const ds4_model *m, con
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) &&
if (ok && mm) ok = qwen4_gemv(g, g->sh_gate, m, l->ffn_gate_shexp, g->mixed, T) &&
qwen4_gemv(g, 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);
qwen4_gemv(g, g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T);
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),
@@ -58750,8 +58765,8 @@ static bool qwen4_graph_moe_impl(ds4_qwen4_gpu_graph *g, const ds4_model *m, con
l->ffn_gate_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED,
DS4_N_EXPERT_USED, DS4_N_EMBD, DS4_N_FF_EXP, g->cap_tokens) &&
qwen4_moe_profile_boundary(profile, &last, &elapsed[2]) &&
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) &&
qwen4_gemv(g, g->sh_gate, m, l->ffn_gate_shexp, g->mixed, T) &&
qwen4_gemv(g, 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_moe_profile_boundary(profile, &last, &elapsed[3]);
}
@@ -58761,7 +58776,7 @@ static bool qwen4_graph_moe_impl(ds4_qwen4_gpu_graph *g, const ds4_model *m, con
DS4_N_EXPERT_USED, DS4_N_EXPERT_USED, DS4_N_FF_EXP, DS4_N_EMBD,
g->cap_tokens) &&
qwen4_moe_profile_boundary(profile, &last, &elapsed[4]) &&
qwen4_gemv(g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T) &&
qwen4_gemv(g, g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T) &&
qwen4_moe_profile_boundary(profile, &last, &elapsed[5]);
}
if (ok) {
@@ -58788,10 +58803,10 @@ static bool qwen4_graph_moe_impl(ds4_qwen4_gpu_graph *g, const ds4_model *m, con
const bool shared_dense = false;
#endif
if (ok && shared_dense) {
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) &&
ok = qwen4_gemv(g, g->sh_gate, m, l->ffn_gate_shexp, g->mixed, T) &&
qwen4_gemv(g, 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);
qwen4_gemv(g, g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T);
}
/* Sixteen rows pick ten experts each but only a third as many distinct
* ones: the grouped kernels read every expert once per four rows that
@@ -58921,8 +58936,9 @@ static bool qwen4_graph_mtp_boundary(ds4_qwen4_gpu_graph *g, const ds4_model *m,
* everything is causal by construction because the recurrent kernels walk
* tokens in order and attention reads the caches written for the same
* chunk. g->R keeps the pre-mixer streams of every row afterwards. */
static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
const int *tokens, uint32_t T, float *logits_out, bool all_rows) {
static bool qwen4_graph_forward_tokens_with_phase(ds4_qwen4_gpu_graph *g,
const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t T,
float *logits_out, bool all_rows, qwen4_projection_phase phase) {
if (!g || T == 0 || T > g->cap_tokens || g->pos + T > g->ctx_cap) return false;
if (all_rows && T > g->n_logit_rows) return false;
for (uint32_t t = 0; t < T; t++) {
@@ -58955,6 +58971,10 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
!qwen4_graph_stage_inputs(g, m, w, tokens, T)) return false;
const double t1 = timing ? now_sec() : 0.0;
if (!glm_graph_begin_commands_if_needed()) return false;
/* The predictor boundary above retains its decode policy. Scope only the
* trunk and output projections, then restore it on both success/failure. */
const qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = phase;
bool ok = true;
/* DS4_QWEN4_TIMING=2 on prefill batches: sync after each stage group and
* report GPU ms per group (adds sync overhead; diagnostics only) */
@@ -58973,8 +58993,8 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
for (uint32_t il = 0; il < n_trunk && ok; il++) {
const ds4_layer_weights *l = &w->layer[il];
if (ds4_qwen4_layer_is_ple(il)) {
ok = qwen4_gemv(g->ple_key, m, l->ple_key, g->ple_emb, T) &&
qwen4_gemv(g->ple_val, m, l->ple_value, g->ple_emb, T) &&
ok = qwen4_gemv(g, g->ple_key, m, l->ple_key, g->ple_emb, T) &&
qwen4_gemv(g, g->ple_val, m, l->ple_value, g->ple_emb, T) &&
ds4_gpu_qwen4_ple_gate_tensor(g->ple_gated, g->ple_normed, g->R, g->ple_key, g->ple_val,
m->map, m->size, l->ple_norm_key->abs_offset,
l->ple_norm_query->abs_offset, l->ple_norm_conv->abs_offset,
@@ -59008,7 +59028,7 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
* buffers have the same capacity, including for later prefill. */
ds4_gpu_tensor *swap = g->R; g->R = g->hc_u; g->hc_u = swap;
swap = g->inj; g->inj = g->inj_alt; g->inj_alt = swap;
ok = qwen4_gemv(g->lo, m, l->hc_ffn_down, g->xn, T) &&
ok = qwen4_gemv(g, g->lo, m, l->hc_ffn_down, g->xn, T) &&
ds4_gpu_qwen4_hc_gate_mix_tensor(g->mixed, g->xn, g->lo, m->map, m->size,
l->hc_ffn_up->abs_offset, l->hc_ffn_up->type, T, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_LOWRANK);
}
@@ -59038,7 +59058,7 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
#undef QWEN4_PROF
if (ok && logits_out && all_rows) {
ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, T) &&
qwen4_gemv(g->logits, m, w->output, g->mixed, T);
qwen4_gemv(g, g->logits, m, w->output, g->mixed, T);
} else if (ok && logits_out) {
/* final mixer on the last token only */
ds4_gpu_tensor *last = ds4_gpu_tensor_view(g->R, (uint64_t)(T - 1u) * hc_dim * sizeof(float),
@@ -59051,7 +59071,7 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
g->R = R_save;
ds4_gpu_tensor_free(last);
}
if (ok) ok = qwen4_gemv(g->logits, m, w->output, g->mixed, 1);
if (ok) ok = qwen4_gemv(g, g->logits, m, w->output, g->mixed, 1);
}
if (ok && g->mtp_R) {
const uint64_t row_bytes = (uint64_t)hc_dim * sizeof(float);
@@ -59094,12 +59114,22 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
g->snap2_valid = ok;
g->snap_after_second = false;
}
g->projection_phase = saved_phase;
return ok;
}
/* A contiguous prompt chunk remains prefill even when its final tail has a
* single row. Decode and speculative verification enter with their own phase. */
static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
const int *tokens, uint32_t T, float *logits_out, bool all_rows) {
return qwen4_graph_forward_tokens_with_phase(g, m, w, tokens, T, logits_out, all_rows,
QWEN4_PROJECTION_PREFILL);
}
static bool qwen4_graph_forward_token(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
int token, float *logits_out) {
return qwen4_graph_forward_tokens(g, m, w, &token, 1, logits_out, false);
return qwen4_graph_forward_tokens_with_phase(g, m, w, &token, 1, logits_out, false,
QWEN4_PROJECTION_DECODE);
}
/* Copy the recurrent state (GDN states and conv histories, PLE history and
@@ -59290,7 +59320,7 @@ static bool qwen4_graph_mtp_project(ds4_qwen4_gpu_graph *g, const ds4_model *m,
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;
#endif
return qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, T * (DS4_N_HC + 1u));
return qwen4_gemv(g, 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
@@ -59346,6 +59376,8 @@ static bool qwen4_graph_mtp_cache_steps(ds4_qwen4_gpu_graph *g, const ds4_model
const int *next_tokens, uint32_t T, uint32_t idx) {
if (!g->mtp_R || !T || 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 qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = QWEN4_PROJECTION_DECODE;
ds4_gpu_tensor *R = g->R;
bool ok = qwen4_graph_mtp_input(g, m, w, row, next_tokens, T);
g->R = g->mtp_R;
@@ -59361,6 +59393,7 @@ static bool qwen4_graph_mtp_cache_steps(ds4_qwen4_gpu_graph *g, const ds4_model
g->mtp_tail_next_token = next_tokens[T - 1u];
}
}
g->projection_phase = saved_phase;
return ok;
}
@@ -59371,6 +59404,8 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
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 qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = QWEN4_PROJECTION_DECODE;
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];
@@ -59395,7 +59430,7 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
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));
: qwen4_gemv_rows(g, 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,
@@ -59426,6 +59461,7 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
g->mtp_tail_next_token = next_tokens[T - 1u];
}
}
g->projection_phase = saved_phase;
return ok;
}
@@ -59452,6 +59488,8 @@ static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *
qwen4_ref_row(m, w->token_embd, (uint64_t)next_token, g->host_row);
if (!ds4_gpu_tensor_write(g->mtp_e, 0, g->host_row, emb_bytes) ||
!glm_graph_begin_commands_if_needed()) return false;
const qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = QWEN4_PROJECTION_DECODE;
ds4_gpu_tensor *R_save = g->R;
ds4_gpu_tensor *last = NULL;
bool ok = true;
@@ -59489,7 +59527,7 @@ static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *
last = ds4_gpu_tensor_view(g->mtp_R, 0, 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) &&
qwen4_gemv_rows(g->logits, m, w->output, g->mixed, 1, chain_head_rows);
qwen4_gemv_rows(g, g->logits, m, w->output, g->mixed, 1, chain_head_rows);
}
g->R = R_save;
if (ok) ok = ds4_gpu_qwen4_argmax_tensor(g->mtp_argmax, g->mtp_argmax_tmp, g->logits, chain_head_rows) != 0;
@@ -59505,6 +59543,7 @@ static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *
g->mtp_pos = idx + 1u;
g->mtp_last_rows = 1u;
}
g->projection_phase = saved_phase;
return ok;
}
@@ -68128,8 +68167,8 @@ static int qwen4_first_token_test(const ds4_model *model, const ds4_vocab *vocab
p += n;
}
for (uint32_t t = 0; bar && t + 1u < ln && fok; t++) {
fok = qwen4_graph_forward_tokens(g, model, weights, lseq + p_len + t, 1u,
out + (uint64_t)(t + 1u) * DS4_N_VOCAB, false);
fok = qwen4_graph_forward_token(g, model, weights, lseq[p_len + t],
out + (uint64_t)(t + 1u) * DS4_N_VOCAB);
}
if (!fok) {
fprintf(stderr, "ds4: FT_LIST forward failed (%u tokens)\n", ln);
@@ -74650,7 +74689,8 @@ static int ds4_session_qwen4_spec_cycle(ds4_session *s, int first_token, float t
g->snap2_valid = false;
g->verify_rows_exact = deep;
ds4_gpu_qwen4_set_verify_rows_exact(deep);
const bool ok = qwen4_graph_forward_tokens(g, m, w, toks, T, rows, true);
const bool ok = qwen4_graph_forward_tokens_with_phase(g, m, w, toks, T, rows, true,
QWEN4_PROJECTION_DECODE);
ds4_gpu_qwen4_set_verify_rows_exact(false);
g->verify_rows_exact = false;
g->snap_after_second = false;
@@ -74772,7 +74812,7 @@ static int ds4_session_qwen4_spec_cycle(ds4_session *s, int first_token, float t
}
if (replacement >= 0) {
/* the residual sample stands in for the draft: run it from the restored state */
if (!qwen4_graph_forward_tokens(g, m, w, &replacement, 1, s->logits, false)) {
if (!qwen4_graph_forward_token(g, m, w, replacement, s->logits)) {
if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: replacement replay failed");
s->checkpoint_valid = false;
return -1;
@@ -78778,6 +78818,7 @@ static bool qwen4_batch_rows_build(qwen4_batch_row *rows, int count,
static void qwen4_batch_row_graph(ds4_qwen4_gpu_graph *out, ds4_session *s,
const qwen4_batch_row *v) {
*out = s->qwen4_graph;
out->projection_phase = QWEN4_PROJECTION_DECODE;
out->R = v->R;
out->ple_gated = v->ple_gated;
out->ple_normed = v->ple_normed;
@@ -78878,10 +78919,10 @@ static bool qwen4_batch_linear(ds4_decode_item *items, int count,
const ds4_qwen4_gpu_graph *rowg,
ds4_qwen4_gpu_graph *g, const ds4_model *m,
const ds4_layer_weights *l, uint32_t il, uint32_t T) {
if (!qwen4_gemv(g->qkv, m, l->lin_qkv, g->mixed, T) ||
!qwen4_gemv(g->z, m, l->lin_gate, g->mixed, T) ||
!qwen4_gemv(g->ga, m, l->lin_alpha, g->mixed, T) ||
!qwen4_gemv(g->gb, m, l->lin_beta, g->mixed, T)) {
if (!qwen4_gemv(g, g->qkv, m, l->lin_qkv, g->mixed, T) ||
!qwen4_gemv(g, g->z, m, l->lin_gate, g->mixed, T) ||
!qwen4_gemv(g, g->ga, m, l->lin_alpha, g->mixed, T) ||
!qwen4_gemv(g, g->gb, m, l->lin_beta, g->mixed, T)) {
return false;
}
/* The causal convolution and the delta-net scan are the only steps that
@@ -78942,7 +78983,7 @@ static bool qwen4_batch_linear(ds4_decode_item *items, int count,
return ds4_gpu_qwen4_gdn_out_tensor(g->lin_o, g->z, m->map, m->size,
l->lin_norm->abs_offset, T, DS4_N_LIN_V_HEAD,
DS4_N_LIN_HEAD_DIM, DS4_RMS_EPS) != 0 &&
qwen4_gemv(g->blk, m, l->lin_out, g->lin_o, T);
qwen4_gemv(g, g->blk, m, l->lin_out, g->lin_o, T);
}
/* Every row's attention in one dispatch per stage: the projections above
@@ -79020,28 +79061,30 @@ static bool qwen4_batch_attention_rows(int count, ds4_qwen4_gpu_graph *rowg,
static bool qwen4_batch_attention(int count, ds4_qwen4_gpu_graph *rowg,
ds4_qwen4_gpu_graph *g, const ds4_model *m,
const ds4_layer_weights *l, uint32_t il, uint32_t T) {
bool ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T) &&
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) &&
qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T);
bool ok = qwen4_gemv(g, g->qg, m, l->attn_q, g->mixed, T) &&
qwen4_gemv(g, g->kp, m, l->attn_k, g->mixed, T) &&
qwen4_gemv(g, g->vp, m, l->attn_v, g->mixed, T) &&
qwen4_gemv(g, g->iq, m, l->indexer_q_proj, g->mixed, T) &&
qwen4_gemv(g, g->ik, m, l->indexer_k_proj, g->mixed, T);
const bool rows_path = getenv("DS4_QWEN4_NO_BATCH_ATTN") == NULL &&
getenv("DS4_QWEN4_NO_IDX_SELECT") == NULL &&
(DS4_N_HEAD_DIM == 128u || DS4_N_HEAD_DIM == 256u);
if (ok && rows_path) return qwen4_batch_attention_rows(count, rowg, g, m, l, il) &&
qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T);
qwen4_gemv(g, g->blk, m, l->attn_output, g->attn_o, T);
/* Each row appends to its own KV and indexer caches at its own position,
* then reads them back through its own block selection. */
for (int i = 0; ok && i < count; i++) {
ds4_qwen4_gpu_graph *r = &rowg[i];
ok = qwen4_graph_attention_tail(r, m, l, il, r->pos, 1u);
}
return ok && qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T);
return ok && qwen4_gemv(g, g->blk, m, l->attn_output, g->attn_o, T);
}
static bool qwen4_graph_encode_native_session_batch(ds4_decode_item *items, int count,
ds4_qwen4_gpu_graph *g,
const ds4_model *m, const ds4_weights *w) {
const qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = QWEN4_PROJECTION_DECODE;
const uint32_t T = (uint32_t)count;
const uint32_t hc_dim = DS4_N_EMBD * DS4_N_HC;
const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT;
@@ -79094,8 +79137,8 @@ static bool qwen4_graph_encode_native_session_batch(ds4_decode_item *items, int
if (!ok) break;
}
if (ds4_qwen4_layer_is_ple(il)) {
ok = qwen4_gemv(g->ple_key, m, l->ple_key, g->ple_emb, T) &&
qwen4_gemv(g->ple_val, m, l->ple_value, g->ple_emb, T) &&
ok = qwen4_gemv(g, g->ple_key, m, l->ple_key, g->ple_emb, T) &&
qwen4_gemv(g, g->ple_val, m, l->ple_value, g->ple_emb, T) &&
ds4_gpu_qwen4_ple_gate_tensor(g->ple_gated, g->ple_normed, g->R, g->ple_key,
g->ple_val, m->map, m->size,
l->ple_norm_key->abs_offset,
@@ -79141,7 +79184,7 @@ static bool qwen4_graph_encode_native_session_batch(ds4_decode_item *items, int
* so it serves every row in one pass. */
if (ok) ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down,
w->output_hc_up, NULL, T) &&
qwen4_gemv(g->batch_logits, m, w->output, g->mixed, T);
qwen4_gemv(g, g->batch_logits, m, w->output, g->mixed, T);
if (timing) {
fprintf(stderr, "ds4: Qwen3.8 batch(T=%u) ms: stage %.1f ple %.1f hc_attn %.1f "
@@ -79153,6 +79196,7 @@ static bool qwen4_graph_encode_native_session_batch(ds4_decode_item *items, int
qwen4_batch_rows_free(views, count);
free(views);
free(rowg);
g->projection_phase = saved_phase;
return ok;
}
/* Each session contributes its next token and optionally one draft. State
@@ -79241,6 +79285,8 @@ static bool qwen4_batch_stage_embeddings_ragged(const qwen4_batch_member *mem, i
static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_member *mem, int count, uint32_t N,
ds4_qwen4_gpu_graph *g, const ds4_model *m,
const ds4_weights *w) {
const qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = QWEN4_PROJECTION_DECODE;
const uint32_t T = N;
const uint32_t hc_dim = DS4_N_EMBD * DS4_N_HC;
const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT;
@@ -79264,8 +79310,8 @@ static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_mem
if (!ok) break;
}
if (ds4_qwen4_layer_is_ple(il)) {
ok = qwen4_gemv(g->ple_key, m, l->ple_key, g->ple_emb, T) &&
qwen4_gemv(g->ple_val, m, l->ple_value, g->ple_emb, T) &&
ok = qwen4_gemv(g, g->ple_key, m, l->ple_key, g->ple_emb, T) &&
qwen4_gemv(g, g->ple_val, m, l->ple_value, g->ple_emb, T) &&
ds4_gpu_qwen4_ple_gate_tensor(g->ple_gated, g->ple_normed, g->R, g->ple_key, g->ple_val,
m->map, m->size, l->ple_norm_key->abs_offset,
l->ple_norm_query->abs_offset, l->ple_norm_conv->abs_offset,
@@ -79285,10 +79331,10 @@ static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_mem
/* projections over every row; the convolution and the scan walk
* each session's one or two rows against its own state through
* the row table, snapshotting after the first as the verify does */
ok = qwen4_gemv(g->qkv, m, l->lin_qkv, g->mixed, T) &&
qwen4_gemv(g->z, m, l->lin_gate, g->mixed, T) &&
qwen4_gemv(g->ga, m, l->lin_alpha, g->mixed, T) &&
qwen4_gemv(g->gb, m, l->lin_beta, g->mixed, T);
ok = qwen4_gemv(g, g->qkv, m, l->lin_qkv, g->mixed, T) &&
qwen4_gemv(g, g->z, m, l->lin_gate, g->mixed, T) &&
qwen4_gemv(g, g->ga, m, l->lin_alpha, g->mixed, T) &&
qwen4_gemv(g, g->gb, m, l->lin_beta, g->mixed, T);
ds4_gpu_qwen4_gdn_row grows[QWEN4_BATCH_MAX_ROWS];
for (int i = 0; i < count; i++) {
const bool snap = rowg[i].snap_after_first && rowg[i].snap_lin_state[il] && mem[i].n == 2u;
@@ -79315,13 +79361,13 @@ static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_mem
(uint32_t)v_dim) &&
ds4_gpu_qwen4_gdn_out_tensor(g->lin_o, g->z, m->map, m->size, l->lin_norm->abs_offset, T,
DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_RMS_EPS) != 0 &&
qwen4_gemv(g->blk, m, l->lin_out, g->lin_o, T);
qwen4_gemv(g, g->blk, m, l->lin_out, g->lin_o, T);
} else if (ok) {
ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T) &&
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) &&
qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T);
ok = qwen4_gemv(g, g->qg, m, l->attn_q, g->mixed, T) &&
qwen4_gemv(g, g->kp, m, l->attn_k, g->mixed, T) &&
qwen4_gemv(g, g->vp, m, l->attn_v, g->mixed, T) &&
qwen4_gemv(g, g->iq, m, l->indexer_q_proj, g->mixed, T) &&
qwen4_gemv(g, g->ik, m, l->indexer_k_proj, g->mixed, T);
/* one entry per token: a draft row's entry sits at pos + 1 */
ds4_gpu_qwen4_attn_row arows[2 * QWEN4_BATCH_MAX_ROWS];
const uint32_t sparse_pos = (g->k_blocks + 1u) * 4u - 1u;
@@ -79340,7 +79386,7 @@ static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_mem
}
}
if (ok) ok = qwen4_batch_attention_entries(arows, T, g, m, l, il) &&
qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T);
qwen4_gemv(g, g->blk, m, l->attn_output, g->attn_o, T);
}
if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, DS4_N_EMBD, DS4_N_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);
@@ -79358,11 +79404,12 @@ static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_mem
}
}
if (ok) ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, T) &&
qwen4_gemv(g->batch_logits, m, w->output, g->mixed, T);
qwen4_gemv(g, g->batch_logits, m, w->output, g->mixed, T);
qwen4_batch_rows_free(views, count);
free(ple_ids);
free(rowg);
free(views);
g->projection_phase = saved_phase;
return ok;
}
@@ -79403,6 +79450,8 @@ static bool qwen4_batch_mtp_drafts(qwen4_batch_member *mem, int count, const uin
qwen4_ref_row(m, w->token_embd, (uint64_t)ids[t], g->host_row + (uint64_t)t * E);
if (!ds4_gpu_tensor_write(g->batch_head_x, 0, g->host_row, (uint64_t)N * E * sizeof(float)) ||
!glm_graph_begin_commands_if_needed()) return false;
const qwen4_projection_phase saved_phase = g->projection_phase;
g->projection_phase = QWEN4_PROJECTION_DECODE;
/* The trunk's MoE scratch is idle here. Reuse it for predictor inputs
* and project all rows together, keeping the existing stage kernels. */
const uint64_t emb_bytes = (uint64_t)E * sizeof(float);
@@ -79419,7 +79468,7 @@ static bool qwen4_batch_mtp_drafts(qwen4_batch_member *mem, int count, const uin
ds4_gpu_tensor_free(R);
ds4_gpu_tensor_free(emb);
}
if (ok) ok = qwen4_gemv(g->mid, m, l->nextn_eh_proj, g->part, N * (hc + 1u));
if (ok) ok = qwen4_gemv(g, g->mid, m, l->nextn_eh_proj, g->part, N * (hc + 1u));
for (uint32_t t = 0; ok && t < N; t++) {
ds4_gpu_tensor *proj = ds4_gpu_tensor_view(g->mid, t * proj_bytes, proj_bytes);
ds4_gpu_tensor *R = ds4_gpu_tensor_view(g->R, t * row_bytes, row_bytes);
@@ -79428,13 +79477,13 @@ static bool qwen4_batch_mtp_drafts(qwen4_batch_member *mem, int count, const uin
ds4_gpu_tensor_free(proj);
}
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, N);
if (ok) ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, N) &&
qwen4_gemv(g->kp, m, l->attn_k, g->mixed, N) &&
qwen4_gemv(g->vp, m, l->attn_v, g->mixed, N) &&
qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, N) &&
qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, N) &&
if (ok) ok = qwen4_gemv(g, g->qg, m, l->attn_q, g->mixed, N) &&
qwen4_gemv(g, g->kp, m, l->attn_k, g->mixed, N) &&
qwen4_gemv(g, g->vp, m, l->attn_v, g->mixed, N) &&
qwen4_gemv(g, g->iq, m, l->indexer_q_proj, g->mixed, N) &&
qwen4_gemv(g, g->ik, m, l->indexer_k_proj, g->mixed, N) &&
qwen4_batch_attention_entries(arows, N, g, m, l, il) &&
qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, N);
qwen4_gemv(g, g->blk, m, l->attn_output, g->attn_o, N);
if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, N, 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, N);
if (ok) ok = qwen4_graph_moe(g, m, l, il, N);
@@ -79451,7 +79500,7 @@ static bool qwen4_batch_mtp_drafts(qwen4_batch_member *mem, int count, const uin
(uint32_t)count);
g->R = R_save;
g->mixed = mixed_save;
if (ok) ok = qwen4_gemv(g->batch_logits, m, w->output, g->batch_head_x, (uint32_t)count);
if (ok) ok = qwen4_gemv(g, g->batch_logits, m, w->output, g->batch_head_x, (uint32_t)count);
for (int i = 0; ok && i < count; i++) {
ds4_qwen4_gpu_graph *sg = &mem[i].session->qwen4_graph;
ds4_gpu_tensor *row = ds4_gpu_tensor_view(g->batch_logits, (uint64_t)i * DS4_N_VOCAB * sizeof(float),
@@ -79471,6 +79520,7 @@ static bool qwen4_batch_mtp_drafts(qwen4_batch_member *mem, int count, const uin
sg->mtp_tail_cached = true;
sg->mtp_tail_next_token = parents[i];
}
g->projection_phase = saved_phase;
return ok;
}
ds4_gpu.h
+8 -2
View File
@@ -3550,12 +3550,18 @@ int ds4_gpu_qwen4_gdn_front_tensor(
int ds4_gpu_qwen4_vision_encode(float *out, const float *patches, const float *pos_embed, uint32_t n_patches,
uint32_t grid_w, const void *model_map, uint64_t model_size,
const ds4_qwen4_vision_weights *w);
/* prefill dense GEMM (f32/f16/q8_0 rows, 32x32 tiles) and the batched hc mix
* pieces */
/* Dense GEMM (f32/f16/q8_0 rows, 32x32 tiles) and the batched hc mix pieces. */
int ds4_gpu_qwen4_dense_mm_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t weight_type,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows);
/* Contiguous prefill uses the same dense kernel without Metal split-K so
* short chunks retain the prefill reduction order. CUDA keeps its existing
* dense dispatch. Decode batches and vision use the entry point above. */
int ds4_gpu_qwen4_dense_mm_prefill_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t weight_type,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows);
int ds4_gpu_qwen4_hc_lo_act_tensor(ds4_gpu_tensor *lo_act, const ds4_gpu_tensor *lo,
uint32_t n_tokens, uint32_t n_hc, uint32_t n_rank);
int ds4_gpu_qwen4_hc_mix_rows_tensor(ds4_gpu_tensor *mixed, const ds4_gpu_tensor *u, const ds4_gpu_tensor *xn,
ds4_metal.m
+23 -7
View File
@@ -51286,10 +51286,10 @@ static bool qwen4_dense_mm_partials_ensure(uint64_t bytes) {
return true;
}
int ds4_gpu_qwen4_dense_mm_tensor(
static int qwen4_dense_mm_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t weight_type,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows) {
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows, bool allow_split_k) {
const uint32_t row_bytes = weight_type == 0u ? in_dim * 4u : weight_type == 1u ? in_dim * 2u :
qwen4_expert_row_bytes(weight_type, in_dim);
struct { uint32_t n_tokens, in_dim, out_rows, weight_type, row_bytes, n_split, pad1, pad2; } args =
@@ -51309,14 +51309,14 @@ int ds4_gpu_qwen4_dense_mm_tensor(
}
/* A projection whose output is narrow gives this grid only a few
* threadgroups when the batch is a decode step; split k until the
* machine has work, then sum the planes. A prefill chunk's token tiles
* already fill it, and its planes would not fit the scratch anyway.
* The split is chosen so the partials stay small and every threadgroup
* still walks enough k to amortize its staging. */
* machine has work, then sum the planes. Contiguous prefill disables
* splitting explicitly, even for short chunks, to retain its reduction
* order. For other callers the split is chosen so the partials stay
* small and every threadgroup still walks enough k to amortize staging. */
const uint32_t tiles = (out_rows + 31u) / 32u;
const uint32_t threadgroups = tiles * ((n_tokens + 31u) / 32u);
uint32_t n_split = 1u;
if (!ds4_gpu_env_u64("DS4_QWEN4_NO_DENSE_MM_KSPLIT", 0u, 0u, 1u)) {
if (allow_split_k && !ds4_gpu_env_u64("DS4_QWEN4_NO_DENSE_MM_KSPLIT", 0u, 0u, 1u)) {
const uint32_t nk = (in_dim + 31u) / 32u;
const uint32_t target = 128u;
const uint32_t want = threadgroups >= target ? 1u : (target + threadgroups - 1u) / threadgroups;
@@ -51349,6 +51349,22 @@ int ds4_gpu_qwen4_dense_mm_tensor(
MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_dense_mm_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t weight_type,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows) {
return qwen4_dense_mm_tensor(out, x, model_map, model_size, weight_offset,
weight_type, n_tokens, in_dim, out_rows, true);
}
int ds4_gpu_qwen4_dense_mm_prefill_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t weight_type,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows) {
return qwen4_dense_mm_tensor(out, x, model_map, model_size, weight_offset,
weight_type, n_tokens, in_dim, out_rows, false);
}
/* Decode batch of 8 or 16 rows against a Q8 matrix (kernel_qwen4_batch_mm_q8):
* the k-split gives narrow projections enough simdgroups, the partial planes
* are summed by kernel_qwen4_dense_mm_reduce. */
ds4_qwen4_cuda.cuh
+5
View File
@@ -2039,6 +2039,11 @@ extern "C" int ds4_gpu_qwen4_dense_mm_tensor(ds4_gpu_tensor *out, const ds4_gpu_
type, 1, T, 1, 1, K, M, 0, false);
}
extern "C" int ds4_gpu_qwen4_dense_mm_prefill_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *map, uint64_t size, uint64_t off, uint32_t type, uint32_t T, uint32_t K, uint32_t M) {
return ds4_gpu_qwen4_dense_mm_tensor(out, x, map, size, off, type, T, K, M);
}
extern "C" int ds4_gpu_qwen4_matmul_q8_0_tensor(ds4_gpu_tensor *out, const void *map,
uint64_t size, uint64_t off, uint64_t K, uint64_t M, const ds4_gpu_tensor *x, uint64_t T) {
if (K > UINT_MAX || M > UINT_MAX || T > UINT_MAX) return 0;
tests/test_qwen4_generation.c
+43 -11
View File
@@ -1,5 +1,8 @@
/* Real-model one-shot/session parity with identical effective prefill chunks.
* Usage: test_qwen4_generation MODEL [--ssd-streaming] [--chunk N --ctx N [--tokens N]]
* --prefill-reference compares automatic dispatch with the diagnostic
* legacy, unsplit prefill path, then uses automatic dispatch for both decodes.
* Usage: test_qwen4_generation MODEL [--ssd-streaming] [--prefill-reference]
* [--chunk N --ctx N [--tokens N]]
* SSD mode bounds the expert cache to 1024 entries. At most two graphs are
* live; the public one-shot graph is released before numerical comparison. */
#include "../ds4.c"
@@ -77,7 +80,8 @@ static void generation_logits(const float *direct, const float *session, uint32_
exit(1);
}
static void generation_case(const char *model, bool streaming, uint32_t requested, int ctx, int tokens) {
static void generation_case(const char *model, bool streaming, bool prefill_reference,
uint32_t requested, int ctx, int tokens) {
/* Compute the expected public behavior independently of the resolver
* under test: an explicit request wins over the environment, then clamps. */
const uint32_t cap = requested < (uint32_t)ctx ? requested : (uint32_t)ctx;
@@ -133,7 +137,14 @@ static void generation_case(const char *model, bool streaming, uint32_t requeste
float *session_logits = malloc((size_t)DS4_N_VOCAB * sizeof(float));
generation_need(direct_logits && session_logits, "allocate vocabulary comparisons");
/* Match one-shot's omission of intermediate output heads. Session sync
* computes each head for its progress/checkpoint contract. */
* computes each head for its progress/checkpoint contract. The optional
* reference disables the two arithmetic changes implicated in prefill
* drift independently of the graph's automatic phase dispatch. */
if (prefill_reference) {
generation_need(setenv("DS4_QWEN4_DENSE_MM_LEGACY", "1", 1) == 0 &&
setenv("DS4_QWEN4_NO_DENSE_MM_KSPLIT", "1", 1) == 0,
"select legacy unsplit reference prefill");
}
for (int pos = 0; pos < length;) {
uint32_t rows = (uint32_t)(length - pos);
if (rows > direct->cap_tokens) rows = direct->cap_tokens;
@@ -142,6 +153,11 @@ static void generation_case(const char *model, bool streaming, uint32_t requeste
"direct prefill");
pos += (int)rows;
}
if (prefill_reference) {
generation_need(unsetenv("DS4_QWEN4_DENSE_MM_LEGACY") == 0 &&
unsetenv("DS4_QWEN4_NO_DENSE_MM_KSPLIT") == 0,
"restore automatic dispatch for both decodes");
}
int emitted = 0;
bool stopped = false;
@@ -171,8 +187,9 @@ static void generation_case(const char *model, bool streaming, uint32_t requeste
}
generation_need(emitted == one_shot.count, "same public and session greedy token count");
printf("PASS Qwen generation requested=%u cap=%u ctx=%d prompt=%d chunks=%d tokens=%d: "
"%d complete vocabulary frontiers bit-exact (%s)\n", requested, cap, ctx, length,
one_shot.calls, emitted, GENERATION_STEPS + 1, streaming ? "SSD" : "resident");
"%d complete vocabulary frontiers bit-exact (%s%s)\n", requested, cap, ctx, length,
one_shot.calls, emitted, GENERATION_STEPS + 1, streaming ? "SSD" : "resident",
prefill_reference ? ", legacy unsplit prefill reference" : "");
free(session_logits); free(direct_logits);
qwen4_graph_free(direct); free(direct);
if (streaming) engine->qwen4_session_bytes -= memory.total_bytes;
@@ -189,11 +206,12 @@ static int generation_number(const char *text, int low, int high) {
}
int main(int argc, char **argv) {
bool streaming = false;
bool streaming = false, prefill_reference = false;
int chunk = 0, ctx = 0, tokens = 0;
if (argc < 2) goto usage;
for (int i = 2; i < argc; i++) {
if (!strcmp(argv[i], "--ssd-streaming")) streaming = true;
else if (!strcmp(argv[i], "--prefill-reference")) prefill_reference = true;
else if (!strcmp(argv[i], "--chunk") && i + 1 < argc) {
chunk = generation_number(argv[++i], 1, 65536);
if (!chunk) goto usage;
@@ -207,24 +225,38 @@ int main(int argc, char **argv) {
}
if ((chunk == 0) != (ctx == 0)) goto usage;
if (tokens && (!ctx || tokens > ctx - GENERATION_STEPS - 1)) goto usage;
/* An inherited diagnostic override could make both sides take the
* reference path and hide an automatic-dispatch regression. */
if (prefill_reference && (getenv("DS4_QWEN4_DENSE_MM_LEGACY") ||
getenv("DS4_QWEN4_NO_DENSE_MM_KSPLIT"))) {
fputs("Qwen generation: --prefill-reference requires unsetting "
"DS4_QWEN4_DENSE_MM_LEGACY and DS4_QWEN4_NO_DENSE_MM_KSPLIT\n", stderr);
return 2;
}
const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK");
char *saved = env ? strdup(env) : NULL;
generation_need(!env || saved, "save prefill environment");
generation_need(setenv("DS4_QWEN4_PREFILL_CHUNK", "32", 1) == 0, "set conflicting chunk fallback");
if (chunk) generation_case(argv[1], streaming, (uint32_t)chunk, ctx, tokens);
if (chunk) generation_case(argv[1], streaming, prefill_reference, (uint32_t)chunk, ctx, tokens);
else if (prefill_reference) {
generation_case(argv[1], streaming, true, 32u, 512, 29);
generation_case(argv[1], streaming, true, 128u, 1024, 575); /* four chunks + a 63-row tail */
generation_case(argv[1], streaming, true, 128u, 512, 256); /* two full chunks */
}
else {
generation_case(argv[1], streaming, 32u, 512, 0);
generation_case(argv[1], streaming, 128u, 512, 0);
generation_case(argv[1], streaming, false, 32u, 512, 0);
generation_case(argv[1], streaming, false, 128u, 512, 0);
/* An explicit chunk above the context must clamp to ctx, never fall
* back to the conflicting environment value of 32. */
generation_case(argv[1], streaming, 512u, 256, 0);
generation_case(argv[1], streaming, false, 512u, 256, 0);
}
generation_need((saved ? setenv("DS4_QWEN4_PREFILL_CHUNK", saved, 1) :
unsetenv("DS4_QWEN4_PREFILL_CHUNK")) == 0, "restore prefill environment");
free(saved);
return 0;
usage:
fprintf(stderr, "usage: %s QWEN_GGUF [--ssd-streaming] [--chunk 1..65536 --ctx 16..32768 [--tokens N]]\n", argv[0]);
fprintf(stderr, "usage: %s QWEN_GGUF [--ssd-streaming] [--prefill-reference] "
"[--chunk 1..65536 --ctx 16..32768 [--tokens N]]\n", argv[0]);
return 2;
}
#else
tests/test_qwen4_mtp_prefill.c
+2 -2
View File
@@ -128,7 +128,7 @@ static bool reference_mtp_input(ds4_qwen4_gpu_graph *g, const ds4_model *m, cons
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_gemv(g, 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);
@@ -171,7 +171,7 @@ static bool reference_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, cons
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));
: qwen4_gemv_rows(g, 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,