qwen4exp: fix QSA correctness defects and harden metadata loading - #143
qwen4exp: fix QSA correctness defects and harden metadata loading#143danielhanchen wants to merge 2 commits into
Conversation
Four runtime side fixes, no GGUF format change: - copy the indexer cache on a cross stream seq_cp. The update constructor never built ctx_idx, so a copied sequence kept the destination stream's old indexer keys while its cell metadata claimed the parent's positions. Reachable through the plain OpenAI n parameter with no flags. - key QSA blocks on the pair of sequence set and position bucket. Keying on position alone is only correct when the cells array holds one sequence, so under a unified cache the last writer won and a block's pooled key was built from another sequence's cells. - cut blocks on each cell's rank in sequence order under M-RoPE. Every token of an image shares one position, so all of them mapped to the same block slot and overwrote each other. - throw instead of aborting on malformed metadata. Eight GGML_ASSERT sites in load_arch_hparams were reachable from a hand edited GGUF, and a short fixed size array was accepted in silence. Also counts pooled blocks along ne1 for the norm, which avoids a pre existing CUDA abort at n_kv 262144 where the block count reached the 65535 gridDim.y limit. Perplexity is identical before and after over 1 and 16 chunks, and test-llama-archs -a qwen4exp passes on CPU and CUDA with roundtrip.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: efaa1ed738
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // a block is r cells one sequence holds at consecutive indices, so it is keyed on | ||
| // (sequence set, index bucket). A unified cache shares one cells array between all | ||
| // sequences and every sequence counts positions from zero, so the bucket alone pools | ||
| // different sequences into one block. Cells with the same sequence set are visible to | ||
| // the same sequences, so that set is the coarsest key that never mixes two of them. |
There was a problem hiding this comment.
This newly added comment is ten lines long and splits sentences across multiple // lines; the same pattern recurs throughout set_input_qsa. The repository explicitly requires comments to be usually 1-2 lines and forbids mid-sentence hard wrapping, so condense these explanations rather than landing the repeated long-form commentary.
AGENTS.md reference: AGENTS.md:L74-L81
Useful? React with 👍 / 👎.
| // cell. Rank the cells in sequence order and key on the rank: that is the token | ||
| // index the reference groups on. A rank only orders cells inside one sequence, so | ||
| // this needs a cache holding one; otherwise the short groups are dropped as above. | ||
| if (dup && ubatch->is_pos_2d() && one_seq) { |
There was a problem hiding this comment.
When a multimodal prompt is copied to another sequence in a unified cache, the shared prompt cells contain both sequence IDs, so n_seq_present is greater than one and one_seq is false even though all those cells have the same sequence set. This condition therefore skips ranking for the repeated M-RoPE image positions; group_cells() leaves the image cells competing for one slot, causing their block to be discarded as incomplete or pooled from only one image cell once QSA is active. The repeated-position ranking needs to handle copied/shared sequence sets, not only caches containing one sequence ID.
Useful? React with 👍 / 👎.
| GGML_ASSERT(r <= 64); | ||
| const uint64_t slots_full = r == 64 ? ~uint64_t(0) : ((uint64_t(1) << r) - 1); |
There was a problem hiding this comment.
For a Qwen4Exp GGUF whose nonzero attention.compress_ratios entry exceeds 64, metadata loading still accepts the value, but the first sparse-attention graph input terminates the process at this GGML_ASSERT. This also undermines the new catchable-error handling for malformed metadata. Validate the upper bound in load_arch_hparams and throw a load error, or represent block occupancy without the 64-bit limit.
Useful? React with 👍 / 👎.
Cut the added comment lines from 64 to 19. Keeps the invariants that the code does not state, the reference citations, and the two traps, and drops restatements, spacers and measured numbers that the pull request description already carries. Also corrects two comments in llama-memory-hybrid-idx.h that became wrong when the update context started building ctx_idx.
|
Superseded by ggml-org#27941, which is the same work with nine more lines in Closing so there is one copy to maintain. |
2 participants
In short
ctx_idx, so a copied sequence kept the destination stream's stale keys. Reachable with no flags through the OpenAInparameter, and it silently produced wrong output.--kv-unifieda block could be pooled from another sequence's cells. Now keyed on (sequence set, index bucket).GGML_ASSERTsites reachable from a hand edited GGUF are now throws, plus two cases that were being accepted in silence.gridDim.ylimit atn_kv262144.No GGUF format change. Every published artifact loads unmodified, and perplexity is identical before and after.
What this fixes
Four correctness defects in the qwen4exp implementation, plus one pre-existing crash that two of the investigations turned up independently. All four were found by auditing ggml-org#27879, a draft PR of machine-generated fixes whose author said plainly that they could not judge whether the fixes were correct. Two of its six claims were real, one was real but already documented, and three were wrong. This PR takes only what survived checking, written fresh rather than cherry-picked, because that PR's own version of the largest fix crashes on CUDA.
Everything here is runtime side. No GGUF changes, no new required metadata, and every published artifact keeps loading byte for byte unmodified.
1. The indexer cache is never copied when a sequence is copied
Of the four constructors of
llama_memory_hybrid_idx_context, only the update one never buildsctx_idx.llama_kv_cache::update()is what applies a pending cross streamseq_cp, so withctx_idxnull the indexer key buffer is never copied while its cell metadata already claims the parent's positions.Instrumented,
n_cmpl=2: the attention cache gotupdate()withn_stream_copies=1, the indexer cache got none, andctx_idxwas null in all 53 update contexts.This is reachable with no flags at all, through the plain OpenAI
nparameter (server-context.cpp:3748tocopy_state_totoseq_cp), and it produces silently wrong output. Two children of one prompt:The contrast is the point. Below the budget attention is dense and the indexer is unused, so the children must agree; above it they must not diverge either. A reproducer that shows only one half of that is not sensitive enough, and the first version of this one was not: with synthetic filler the children matched even on broken code, because the continuation was decided by the always visible tail.
The companion change setting the indexer cache to
LLAMA_ROPE_TYPE_NONEis deliberate. Indexer keys are raw, rotation happens after pooling at read time, andllama-kv-cache.cpp:866guards only the shift graph onrope_type, so stream copies in the sameupdate()still run.2. QSA blocks are keyed on position alone, which is unsafe under a unified cache
A QSA block is
ratiocells one sequence holds at consecutive positions. The old code keyed itb = pos/ratio,slot = pos%ratio. That is only correct when the cells array holds one sequence. Under--kv-unifiedevery sequence sharesv_cells[0]and every sequence counts positions from zero, so the last writer wins and a block's pooled key is built from another sequence's cells.Two sequences of 6387 tokens, both keyings from one binary via a toggle:
Cells were actually placed into blocks in the fixed arms, so the zeros are not the result of nothing running.
A note on the magnitude, because it is easy to over-read. These counters are cumulative over
set_input_qsacalls, so the "before" figure tracks the context size and the number of calls rather than being a property of the defect: the same test at a different-creports 74,448 instead of 272,779. The portable claim is the shape, not the number, and it holds in every configuration measured: before is massively non zero under a unified cache, after is exactly zero, and the non unified path is unaffected either way.This is a selection quality defect and not a content leak.
build_attn_mask_top_kreturnsggml_add(kq_mask_top_k, kq_mask)(llama-graph.cpp:3016), so the ordinary causal mask is re-applied over the top-k result and a foreign cell stays at negative infinity however wrong the blocks are. What breaks is which of a sequence's own cells it selects. The description in ggml-org#27879 overstates this, and so did an earlier draft of these notes.Blocks are now keyed on the pair of sequence set and position bucket. Cells carrying the same sequence set are visible to exactly the same sequences, so that set is the coarsest key that can never mix two of them. A bucket whose cells disagree splits into short groups, and a short group is dropped exactly as a partial tail block already was, which matches all three reference implementations: partial blocks are excluded, not pooled over fewer elements.
No tensor changes shape, so graph input bytes are identical before and after: 24.04 GiB at
n_kv262144, ratio 4,n_ubatch2048 unified, in both cases.Known cost. Generation throughput at unified batch 16 drops 0.44 percent, with a standard deviation of 0.22 and a 95 percent confidence interval of 0.26 to 0.62 percent. Six repetitions per arm, arm order rotated every repetition, page cache pre-warmed. It is the per cell sequence set comparison, in exactly the path that was previously incorrect. Prompt throughput is unchanged. A first version also cost 1.49 percent on the non unified path; a single sequence fast path returns that to within noise, since a stream holding one sequence needs no set test at all.
An earlier draft of this description reported 2.33 percent, and that figure is withdrawn. It did not reproduce. Re-measured on identical hardware with a third arm carrying the block keying change alone, so the merge could be ruled in or out:
The merged branch differs from the keying change alone by +0.15 percent +/- 0.77, which is indistinguishable from zero, so merging the two fixes neither introduced nor removed a cost. Since 2.33 percent fails to reproduce for the unmerged change either, the discrepancy is in the original measurement configuration rather than in the code, and no claim is made that anything here improved it.
The pitfall that produced the original number is worth naming, because it is easy to repeat: a cold page cache. Run first and unwarmed, the base arm scored 139.21 t/s against a warm 391, so a single-sample comparison in the wrong order reports a fictitious result of well over 100 percent in either direction. Every figure above is warmed and order rotated.
3. Under M-RoPE many cells share a position, and they all collapse into one block slot
mtmd.cpp:2393-2399gives every token of an image the samepos.t, varying onlyxandy, andllama-kv-cache.cpp:1127stores section 0 as the cell position. So an image contributes hundreds of cells at one position. All of them compute the sameband the sameslot, they overwrite each other, the remaining slots keep a fill value that is a real cell index, and the block is nevertheless declared complete because the fill count reachedratio.Measured with
llama-mtmd-clion a real image above the budget, reading back the arrays that had just been emitted:299 of the 300 image cells were being scored by a pooled key they are not part of.
An honest limit on the impact. This does not visibly break image understanding. The same image is described correctly before and after, at 2816 and at 15360 cells. Three quarters of the corrupted pooled key is cell 0, which behaves like an attention sink and scores high, so the image blocks tend to get selected anyway. The corruption is real and it does change the selection, but no prompt was found where it produces a wrong answer.
Blocks are now cut on each cell's rank in sequence order, ordering by position and then by the two spatial components, which is the same total order the M-RoPE causal mask already uses. The pooled key is roped with the full four section position of the block's first cell rather than a scalar. The path is gated on actually observing a duplicate position in a 2D batch, so with unique positions nothing changes and the text path is unaffected.
All three references agree on both points: block membership is by token index and never by position, and the group start indexes the full interleaved cosine and sine row.
4. Malformed metadata aborts instead of returning an error
load_arch_hparamshad eightGGML_ASSERTsites that a hand edited or truncated GGUF can reach, which aborts the process rather than failing the load. These become throws naming the offending key and its value:ssm.conv_kernel = 0qwen4exp.ssm.conv_kernel must be greater than zero, got 0hyper_connection.low_rank = 0qwen4exp.hyper_connection.low_rank must be greater than zero, got 0ple.layers = [1,5]qwen4exp.ple.layers lists 2 layers, but only one PLE layer is supportedple.layer_multipliershas 1 entries, but at least 3 are requiredThe last row is a genuinely new class rather than a tidier message.
get_arr()into a fixed array bounds checks the upper end, but a short array is copied as is and leaves the destination tail alone. Sincehparamsis value initialised that tail reads as zero, and the n-gram hash quietly drops those positions. A short offsets array made heads 4 through 15 all read from offset 0. There were two such silent cases, not one.Checks were adopted only where the reference implementations agree the value is impossible, not merely unused today.
hc_count == 1is rejected because all three raise on it. An absentcompress_ratiosis accepted because all three treat it as optional, so ggml-org#27879's check on that is not taken. Multiple PLE layers are permitted by all three, so the single layer limit here is a llama.cpp representational limit rather than an architectural rule, and the message says so instead of claiming the architecture forbids it.5. A pre-existing CUDA abort at long context
Found independently by two of the investigations above, and present on master with none of this PR's code.
build_qsa_top_kreshapes the pooled keys so that the block count lands onne2, andrms_normlaunchesgridDim.y = ne2, which is capped at 65535. Atn_kv262144 with ratio 4 that is 65536 blocks, over the limit by exactly one:Blocks are now counted along
ne1for the norm, which is bounded bygridDim.xat 2^31, and only reshaped for the rope, which flattens rows intogridDim.xanyway. Every other operation in the path stays inside its limits forn_kvup to 262144 andn_ubatchup to 2048.This is the same kernel and the same limit that ggml-org#27879 hits, reached through
n_kvrather than throughn_ubatch. That PR did not introduce the failure class; it made an existing hazard easy to trigger.Verification
Every fix was developed against its own build of the same base, and each was checked against pristine master rather than against a sibling fix.
4.8154 +/- 0.04186test-llama-archs -a qwen4expblk_biason and offNote on the budget boundary: it is
indexer_top_k + ratio - 1 = 2051cells, and the cache pads to a multiple of 256. A 2051 token prompt therefore allocates 2304 cells and is already above the budget, so the bit-identical check belongs at 2040 tokens and 2048 cells.Every new test was verified to fail on deliberately broken code. Two did not, and are reported rather than counted:
What is not covered