fix(qwen): unify prefill chunk resolution and test generation parity

An explicit prefill chunk at or above the context size previously fell back
to DS4_QWEN4_PREFILL_CHUNK or the default in one-shot and session allocation,
while the memory estimator clamped it to the context. This could change the
actual prefill boundaries and make execution disagree with memory admission.

Use one resolver for one-shot generation, session allocation, memory
estimation and SSD admission. Explicit requests always take precedence over
the environment and clamp to the context size. Preserve existing defaults
when no explicit chunk is supplied.

Add CPU coverage for chunk precedence and context-boundary cases. Add a
real-model Metal regression test that checks public one-shot token emission
and progress against sessions, then compares every vocabulary logit between
an independently replayed direct graph and the session through eight decode
steps. Cover conflicting environment settings, oversized requests and long
prompts with configurable token counts; bound the SSD expert cache.

Document the test and the upstream 8db1d1d one-shot behavior that ignores the
CLI prefill chunk. Cross-revision comparisons must also set the same
DS4_QWEN4_PREFILL_CHUNK value to align effective chunk boundaries.

Validation on Apple M1 Max, 32 GiB:
- Built ds4, ds4-server and both affected tests without warnings.
- CPU memory/chunk-resolution tests passed; git diff --check passed.
- Five Qwen Q2 SSD cases passed with Metal API validation, including 5,760
  prompt tokens at chunks 128 and 2,048 and an oversized explicit request.
- All public greedy tokens matched; 45 complete vocabulary frontiers were
  bit-exact between direct replay and sessions (11,174,400 float comparisons).
- The oversized-request regression test fails against f47db4f with the
  observed first frontier 32 rather than the expected 247, and passes here.

These checks establish same-build parity at identical chunk boundaries.
They do not establish main-versus-branch resident parity on the tester's
64 GiB machine, or token invariance across different chunk sizes.
This commit is contained in:
GiorgioOppo committed 2026-09-18 13:22:40 +02:00
1 parent f47db4fb83
commit 0024ca0803
6 files changed
+306 -12

No files matched your search

+1
View File
@@ -61,6 +61,7 @@
/tests/test_qwen4_moe_mm_specialize
/tests/test_qwen4_ssd_experts
/tests/test_qwen4_memory
/tests/test_qwen4_generation
/tests/test_qwen4_vision
/tests/test_deepseek4_vision_image
/tests/test_deepseek41_metal
+11
View File
@@ -832,6 +832,16 @@ tests/test_qwen4_mtp_prefill.o: tests/test_qwen4_mtp_prefill.c ds4.c ds4.h
tests/test_qwen4_mtp_prefill: tests/test_qwen4_mtp_prefill.o $(filter-out ds4.o,$(CORE_OBJS))
$(CC) $(filter-out -ffast-math,$(CFLAGS)) -o $@ $^ $(METAL_LDLIBS)
tests/test_qwen4_generation.o: tests/test_qwen4_generation.c ds4.c ds4.h ds4_qwen4_vision.h
$(CC) $(filter-out -ffast-math,$(CFLAGS)) -Wno-unused-function -I. -c -o $@ $<
tests/test_qwen4_generation: tests/test_qwen4_generation.o $(filter-out ds4.o,$(CORE_OBJS))
ifeq ($(UNAME_S),Darwin)
$(CC) $(filter-out -ffast-math,$(CFLAGS)) -o $@ $^ $(METAL_LDLIBS)
else
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
endif
tests/test_qwen4_prefill.o: tests/test_qwen4_prefill.c ds4.h
$(CC) $(QUALITY_CFLAGS) -I. -c -o $@ $<
@@ -1150,6 +1160,7 @@ clean:
rm -f tests/test_qwen4_memory tests/test_qwen4_ssd_experts
rm -f tests/test_qwen4_ngram_state
rm -f tests/test_qwen4_mtp_prefill
rm -f tests/test_qwen4_generation
rm -f tests/test_web_recovery
rm -f tests/test_metal_ssd_experts
rm -f tests/test_metal_ssd_reuse
+22
View File
@@ -241,6 +241,28 @@ optional shared experts and all 16 routed slots. A failed SSD read must leave
the completed cached/shared gate/up slots exact, preserve the untouched slots,
and allow an exact retry after draining GPU work.
The memory-estimate test runs without a GPU or model.
To compare ordinary one-shot generation with the session path at identical
prefill boundaries, including every vocabulary logit during continuation:
```sh
make tests/test_qwen4_generation
./tests/test_qwen4_generation gguf/Qwen3.8-Flash-Next-Q2.gguf --ssd-streaming
# Exercise sparse attention and many chunk boundaries with a longer fixture:
./tests/test_qwen4_generation gguf/Qwen3.8-Flash-Next-Q2.gguf \
--ssd-streaming --chunk 128 --ctx 8192 --tokens 5760
```
Omit `--ssd-streaming` on a machine that can keep the model resident. The
test checks explicit chunk precedence, progress frontiers, greedy tokens and
full logits. Different chunk sizes can change reduction order and routed
expert choices; compare the actual chunk boundaries before diagnosing token
drift. In particular, upstream `8db1d1d` ignores `--prefill-chunk` in the
ordinary one-shot Qwen path, although its sessions honor the option. For an
A/B comparison with that revision, also set `DS4_QWEN4_PREFILL_CHUNK` to the
same value in both binaries. This branch uses an explicit chunk in preference
to that environment setting and clamps it to the context size.
On CUDA, use `make test-qwen4-cuda` for the kernel tests. They compare the
active kernels with independent CPU references, without model weights.
Vision and end-to-end checks additionally require the checkpoints above.
+15 -12
View File
@@ -1042,10 +1042,17 @@ static bool qwen4_attention_needs_index_query(uint32_t pos0, uint32_t n_tokens,
(mtp || qwen4_attention_dense_rows(pos0, n_tokens, k_blocks) < n_tokens);
}
static uint32_t qwen4_prefill_chunk_tokens(uint32_t ctx) {
const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK");
const unsigned long v = env && env[0] ? strtoul(env, NULL, 10) : 8192ul;
uint32_t chunk = v == 0 || v > 65536ul ? 8192u : (uint32_t)v;
/* The one-shot graph, sessions and memory admission must use the same shape:
* changing chunk boundaries also changes floating-point reduction order.
* An explicit chunk always wins over the environment, even when it reaches
* or exceeds the context; clamp it rather than falling back to a default. */
static uint32_t qwen4_prefill_chunk_resolve(uint32_t ctx, uint32_t requested) {
uint32_t chunk = requested;
if (!chunk) {
const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK");
const unsigned long v = env && env[0] ? strtoul(env, NULL, 10) : 8192ul;
chunk = v == 0 || v > 65536ul ? 8192u : (uint32_t)v;
}
return chunk > ctx ? ctx : chunk;
}
@@ -1054,8 +1061,7 @@ static ds4_context_memory qwen4_graph_memory_estimate(uint32_t ctx, uint32_t pre
/* Match the bounded graph workspace, including CPU staging. This API
* has no MTP flag, so an embedded predictor reserves its primary and
* both lazy verifier snapshots as well as the live recurrent state. */
uint64_t T = prefill_chunk ? prefill_chunk : qwen4_prefill_chunk_tokens(ctx);
if (T > ctx) T = ctx;
const uint64_t T = qwen4_prefill_chunk_resolve(ctx, prefill_chunk);
const uint64_t blocks = ctx / 4u + 1u, E = DS4_N_EMBD, hc_dim = E * DS4_N_HC;
const uint64_t kv_row = 2ull * DS4_N_HEAD_KV * DS4_N_HEAD_DIM * 2u;
const uint64_t conv_dim = DS4_N_LIN_CONV_DIM;
@@ -59530,8 +59536,7 @@ static int generate_qwen4_metal_argmax(
return 1;
}
ds4_qwen4_gpu_graph *g = xcalloc(1, sizeof(*g));
const uint32_t cap = prefill_chunk && prefill_chunk < (uint32_t)ctx_size ?
prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size);
const uint32_t cap = qwen4_prefill_chunk_resolve((uint32_t)ctx_size, prefill_chunk);
if (!qwen4_graph_alloc(g, weights, (uint32_t)ctx_size, cap, false, NULL, NULL, NULL, 0)) {
free(g);
return 1;
@@ -65458,8 +65463,7 @@ int ds4_engine_generate_argmax(
uint64_t qwen4_reserved = 0;
if (ds4_model_is_qwen4() && e->ssd_streaming) {
if (ctx_size <= 0) return 1;
const uint32_t cap = e->prefill_chunk && e->prefill_chunk < (uint32_t)ctx_size ?
e->prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size);
const uint32_t cap = qwen4_prefill_chunk_resolve((uint32_t)ctx_size, e->prefill_chunk);
const ds4_context_memory memory = ds4_context_memory_estimate_with_prefill_mode(
e->backend, ctx_size, cap, true);
if (!qwen4_streaming_memory_admit(e,
@@ -73346,8 +73350,7 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) {
free(s);
return 1;
}
const uint32_t cap_tokens = e->prefill_chunk && e->prefill_chunk < (uint32_t)ctx_size ?
e->prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size);
const uint32_t cap_tokens = qwen4_prefill_chunk_resolve((uint32_t)ctx_size, e->prefill_chunk);
const ds4_context_memory memory = ds4_context_memory_estimate_with_prefill_mode(
e->backend, ctx_size, cap_tokens, e->ssd_streaming);
#ifdef DS4_HAS_QWEN4_METAL
+232
View File
@@ -0,0 +1,232 @@
/* Real-model one-shot/session parity with identical effective prefill chunks.
* Usage: test_qwen4_generation MODEL [--ssd-streaming] [--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"
#ifdef DS4_HAS_QWEN4_METAL
enum { GENERATION_STEPS = 8 };
static char generation_error[256];
static void generation_need(bool ok, const char *what) {
if (ok) return;
fprintf(stderr, "Qwen generation: %s: %s\n", what, generation_error);
exit(1);
}
typedef struct {
const char *path;
uint32_t cap;
int length, last, calls;
int tokens[GENERATION_STEPS], count, done;
} generation_capture;
static void generation_progress(void *opaque, const char *event, int current, int total) {
if (strcmp(event, "prefill_chunk")) return;
generation_capture *c = opaque;
const int remaining = c->length - c->last;
const int expected = c->last + (remaining > (int)c->cap ? (int)c->cap : remaining);
if (remaining <= 0 || total != c->length || current != expected) {
fprintf(stderr, "Qwen generation: %s chunk %d: frontier=%d/%d, expected=%d/%d cap=%u\n",
c->path, c->calls, current, total, expected, c->length, c->cap);
exit(1);
}
c->last = current;
c->calls++;
if (c->calls == 1 || current == total || current % 2048 == 0)
fprintf(stderr, "Qwen generation: %s prefill %d/%d cap=%u\n",
c->path, current, total, c->cap);
}
static void generation_emit(void *opaque, int token) {
generation_capture *c = opaque;
generation_need(!c->done && c->count < GENERATION_STEPS, "bounded ordered token emission");
c->tokens[c->count++] = token;
}
static void generation_done(void *opaque) {
generation_capture *c = opaque;
generation_need(c->done == 0, "one completion callback");
c->done++;
}
static void generation_logits(const float *direct, const float *session, uint32_t cap, int step) {
const uint32_t vocab = DS4_N_VOCAB;
uint32_t changed = 0, first = 0, direct_bits = 0, session_bits = 0;
double worst = 0;
for (uint32_t i = 0; i < vocab; i++) {
if (!isfinite(direct[i]) || !isfinite(session[i])) {
fprintf(stderr, "Qwen generation: cap=%u step=%d token=%u nonfinite logits direct=%g session=%g\n",
cap, step, i, direct[i], session[i]);
exit(1);
}
uint32_t a, b;
memcpy(&a, direct + i, sizeof(a));
memcpy(&b, session + i, sizeof(b));
if (a != b) {
if (!changed) { first = i; direct_bits = a; session_bits = b; }
changed++;
worst = fmax(worst, fabs((double)direct[i] - session[i]));
}
}
if (!changed) return;
fprintf(stderr, "Qwen generation: cap=%u step=%d logits differ: %u/%u, max_error=%.9g, "
"top1=%d/%d; first token=%u direct=%.9g (0x%08x) session=%.9g (0x%08x)\n",
cap, step, changed, vocab, worst, sample_argmax(direct, vocab),
sample_argmax(session, vocab), first, direct[first], direct_bits, session[first], session_bits);
exit(1);
}
static void generation_case(const char *model, bool streaming, 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;
generation_need(qwen4_prefill_chunk_resolve((uint32_t)ctx, requested) == cap,
"explicit prefill chunk overrides environment and clamps to context");
const int length = tokens ? tokens : (2u * cap + 3u < (uint32_t)(ctx - GENERATION_STEPS - 1) ?
(int)(2u * cap + 3u) : ctx - GENERATION_STEPS - 1);
ds4_engine_options opt = {.model_path = model, .backend = DS4_BACKEND_METAL,
.context_size = ctx, .prefill_chunk = requested, .glm_mtp = false,
.ssd_streaming = streaming, .ssd_streaming_cache_experts = streaming ? 1024u : 0u};
ds4_engine *engine = NULL;
generation_need(ds4_engine_open(&engine, &opt) == 0 && ds4_engine_is_qwen4(engine), "open Qwen model");
ds4_tokens seed = {0}, prompt = {0};
ds4_tokenize_text(engine, "Continue this counting sequence: one, two, three, four, five, six, seven, eight. ", &seed);
generation_need(seed.len > 0, "tokenize fixture");
for (int i = 0; i < length; i++) ds4_tokens_push(&prompt, seed.v[i % seed.len]);
generation_capture one_shot = {.path = "one-shot", .cap = cap, .length = length};
generation_need(ds4_engine_generate_argmax(engine, &prompt, GENERATION_STEPS, ctx,
generation_emit, generation_done, &one_shot, generation_progress, &one_shot) == 0,
"public one-shot generation");
generation_need(one_shot.done == 1 && one_shot.last == length && one_shot.calls > 0,
"complete one-shot callbacks");
ds4_session *session = NULL;
generation_need(ds4_session_create(&session, engine, ctx) == 0, "create comparison session");
generation_need(session->qwen4_graph.cap_tokens == cap && !session->qwen4_graph.mtp_R,
"session effective chunk and disabled MTP");
generation_capture session_trace = {.path = "session", .cap = cap, .length = length};
ds4_session_set_progress(session, generation_progress, &session_trace);
generation_need(ds4_session_sync(session, &prompt, generation_error, sizeof(generation_error)) == 0,
"session prefill");
ds4_session_set_progress(session, NULL, NULL);
generation_need(session_trace.last == length && session_trace.calls == one_shot.calls,
"matching session chunk frontiers");
ds4_qwen4_gpu_graph *direct = calloc(1, sizeof(*direct));
generation_need(direct != NULL, "allocate direct graph descriptor");
const ds4_context_memory memory = ds4_context_memory_estimate_with_prefill_mode(
DS4_BACKEND_METAL, ctx, cap, streaming);
if (streaming) {
generation_need(qwen4_streaming_memory_admit(engine,
ds4_add_sat_u64(engine->qwen4_session_bytes, memory.total_bytes), false),
"admit both bounded comparison graphs");
engine->qwen4_session_bytes += memory.total_bytes;
}
generation_need(qwen4_graph_alloc(direct, &engine->weights, (uint32_t)ctx,
qwen4_prefill_chunk_resolve((uint32_t)ctx, requested), false, NULL, NULL, NULL, 0),
"allocate independent direct graph");
direct->ssd_streaming = streaming;
qwen4_graph_reset(direct);
float *direct_logits = malloc((size_t)DS4_N_VOCAB * sizeof(float));
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. */
for (int pos = 0; pos < length;) {
uint32_t rows = (uint32_t)(length - pos);
if (rows > direct->cap_tokens) rows = direct->cap_tokens;
generation_need(qwen4_graph_forward_tokens(direct, &engine->model, &engine->weights,
prompt.v + pos, rows, pos + (int)rows == length ? direct_logits : NULL, false),
"direct prefill");
pos += (int)rows;
}
int emitted = 0;
bool stopped = false;
for (int step = 0; step <= GENERATION_STEPS; step++) {
generation_need(ds4_session_copy_logits(session, session_logits, (int)DS4_N_VOCAB) == (int)DS4_N_VOCAB,
"read complete session vocabulary");
generation_logits(direct_logits, session_logits, cap, step);
generation_need(direct->pos == (uint32_t)ds4_session_pos(session), "matching decode frontier positions");
if (step == GENERATION_STEPS) break;
const int token = ds4_session_argmax(session);
if (!stopped && ds4_token_is_stop(engine, token)) stopped = true;
if (!stopped) {
if (emitted >= one_shot.count || token != one_shot.tokens[emitted]) {
fprintf(stderr, "Qwen generation: cap=%u step=%d emitted token mismatch: one-shot=%d session=%d\n",
cap, step, emitted < one_shot.count ? one_shot.tokens[emitted] : -1, token);
exit(1);
}
emitted++;
}
/* Keep all nine vocabulary checks even if this fixture emits a stop.
* Subsequent rows are teacher-forced ordinary fixture tokens. */
const int next = stopped ? seed.v[step % seed.len] : token;
generation_need(ds4_session_eval(session, next, generation_error, sizeof(generation_error)) == 0,
"teacher-forced session decode");
generation_need(qwen4_graph_forward_token(direct, &engine->model, &engine->weights, next, direct_logits),
"teacher-forced direct decode");
}
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");
free(session_logits); free(direct_logits);
qwen4_graph_free(direct); free(direct);
if (streaming) engine->qwen4_session_bytes -= memory.total_bytes;
ds4_session_free(session);
ds4_tokens_free(&prompt); ds4_tokens_free(&seed);
ds4_engine_close(engine);
}
static int generation_number(const char *text, int low, int high) {
char *end = NULL;
const long value = strtol(text, &end, 10);
if (!*text || *end || value < low || value > high) return 0;
return (int)value;
}
int main(int argc, char **argv) {
bool streaming = 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], "--chunk") && i + 1 < argc) {
chunk = generation_number(argv[++i], 1, 65536);
if (!chunk) goto usage;
} else if (!strcmp(argv[i], "--ctx") && i + 1 < argc) {
ctx = generation_number(argv[++i], 16, 32768);
if (!ctx) goto usage;
} else if (!strcmp(argv[i], "--tokens") && i + 1 < argc) {
tokens = generation_number(argv[++i], 1, 32768);
if (!tokens) goto usage;
} else goto usage;
}
if ((chunk == 0) != (ctx == 0)) goto usage;
if (tokens && (!ctx || tokens > ctx - GENERATION_STEPS - 1)) goto usage;
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);
else {
generation_case(argv[1], streaming, 32u, 512, 0);
generation_case(argv[1], streaming, 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_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]);
return 2;
}
#else
int main(void) { fputs("Qwen generation test requires Metal\n", stderr); return 2; }
#endif
+25
View File
@@ -13,6 +13,30 @@ static ds4_context_memory estimate(int ctx, uint32_t chunk) {
return m;
}
static void chunk_resolution(void) {
const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK");
char *saved = env ? strdup(env) : NULL;
assert(!env || saved);
assert(unsetenv("DS4_QWEN4_PREFILL_CHUNK") == 0);
assert(qwen4_prefill_chunk_resolve(16384, 0) == 8192);
assert(qwen4_prefill_chunk_resolve(512, 0) == 512);
assert(setenv("DS4_QWEN4_PREFILL_CHUNK", "32", 1) == 0);
assert(qwen4_prefill_chunk_resolve(512, 0) == 32);
const uint32_t requests[] = {1, 128, 511, 512, 513, UINT32_MAX};
const uint32_t expected[] = {1, 128, 511, 512, 512, 512};
for (size_t i = 0; i < sizeof(requests) / sizeof(requests[0]); i++) {
assert(qwen4_prefill_chunk_resolve(512, requests[i]) == expected[i]);
assert(estimate(512, requests[i]).prefill_cap == expected[i]);
}
assert(estimate(512, 0).prefill_cap == 32);
assert(setenv("DS4_QWEN4_PREFILL_CHUNK", "65537", 1) == 0);
assert(qwen4_prefill_chunk_resolve(16384, 0) == 8192);
assert(qwen4_prefill_chunk_resolve(16384, 16384) == 16384);
assert((saved ? setenv("DS4_QWEN4_PREFILL_CHUNK", saved, 1) :
unsetenv("DS4_QWEN4_PREFILL_CHUNK")) == 0);
free(saved);
}
static void monotonic(void) {
const uint32_t chunks[] = {1, 2, 8, 128, 1024, 4096, 8192};
uint64_t previous = 0;
@@ -160,6 +184,7 @@ static void attention_query_planning(void) {
int main(void) {
const ds4_shape saved = g_ds4_shape;
g_ds4_shape = DS4_SHAPE_QWEN4_EXP;
chunk_resolution();
monotonic();
streaming_staging();
attention_query_planning();