qwen4exp: follow up fixes - #27941
Conversation
|
Hope it helps with #27886 |
|
Not the author of that issue but this fixed #27886 for me |
Sequence-copy indexer keys, block keying on (sequence set, index bucket), M-RoPE image blocks cut on rank order, metadata asserts -> throws, and the gridDim.y 65535 abort at n_kv 262144. PPL-identical per the PR; validated on gfx1151 (arithmetic @24K, vision, 41.7 tok/s file-rewrite). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sequence-copy indexer keys, block keying on (sequence set, index bucket), M-RoPE image blocks cut on rank order, metadata asserts -> throws, and the gridDim.y 65535 abort at n_kv 262144. PPL-identical per the PR; validated on gfx1151 (arithmetic @24K, vision, 41.7 tok/s file-rewrite). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_prev_tokens() scans all used cells once per ubatch to resolve the n-gram predecessor tokens (qwen4exp PLE). the scan is O(used) per call and dominates decode at long context: measured ~46 ms per token at ~170k ctx on 2xL40S (out of ~89 ms/token total, ~65% of wall time). llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set<rows>), updated by the existing seq_pos_inc/dec funnels, so add/remove/defrag/copy (memory_seq_cp) paths update it by construction. prev_token(p) resolves to the token of the largest existing position <= p, with the 'last cell wins' tie-break of the general scan preserved. ubatches with shared temporal positions (multimodal) are detected as not-applicable and keep using the general scan: within a shared position, cells resolve by ubatch order, not by row. LLAMA_KV_PREV_TOKENS env: fast (default) | verify | off. verify runs both paths for every call and logs mismatches + a cost heartbeat; used to validate the index against production traffic (455k lookups across prompt-cache loads and checkpoint restores, 0 mismatches; scan 45.9 ms vs index 6.9 us avg). measured (qwen4exp, 2xL40S, unified KV, 256k ctx): tg @155k ctx: ~11.7 -> ~31 t/s tg @200k ctx: ~11.5 -> ~27 t/s related upstream work: ggml-org#27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), ggml-org#27977 (shrinks the general scan constants + qsa gather windows). complementary layers; can be combined.
|
@danielhanchen Are you planning more changes to this branch? |
|
@ggerganov I think it should be ready - the rest are just perf + MTP which I'll land tomorrow! |
Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.
Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.
Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.
Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.
| } | ||
|
|
||
| // two cells with the same set are visible to exactly the same sequences | ||
| const seq_set_t & seq_set(uint32_t i) const { |
There was a problem hiding this comment.
Rename this to seq_get since it gets the seq member:
| const seq_set_t & seq_set(uint32_t i) const { | |
| const seq_set_t & seq_get(uint32_t i) const { |
Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.
Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.
get_prev_tokens() scans all used cells once per ubatch to resolve the n-gram predecessor tokens (qwen4exp PLE). the scan is O(used) per call and dominates decode at long context: measured ~46 ms per token at ~170k ctx on 2xL40S (out of ~89 ms/token total, ~65% of wall time). llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set<rows>), updated by the existing seq_pos_inc/dec funnels, so add/remove/defrag/copy (memory_seq_cp) paths update it by construction. prev_token(p) resolves to the token of the largest existing position <= p, with the 'last cell wins' tie-break of the general scan preserved. ubatches with shared temporal positions (multimodal) are detected as not-applicable and keep using the general scan: within a shared position, cells resolve by ubatch order, not by row. LLAMA_KV_PREV_TOKENS env: fast (default) | verify | off. verify runs both paths for every call and logs mismatches + a cost heartbeat; used to validate the index against production traffic (455k lookups across prompt-cache loads and checkpoint restores, 0 mismatches; scan 45.9 ms vs index 6.9 us avg). measured (qwen4exp, 2xL40S, unified KV, 256k ctx): tg @155k ctx: ~11.7 -> ~31 t/s tg @200k ctx: ~11.5 -> ~27 t/s related upstream work: ggml-org#27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), ggml-org#27977 (shrinks the general scan constants + qsa gather windows). complementary layers; can be combined.
get_prev_tokens() scans all used cells once per ubatch to resolve the n-gram predecessor tokens (qwen4exp PLE). the scan is O(used) per call and dominates decode at long context: measured ~46 ms per token at ~170k ctx on 2xL40S (out of ~89 ms/token total, ~65% of wall time). llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set<rows>), updated by the existing seq_pos_inc/dec funnels, so add/remove/defrag/copy (memory_seq_cp) paths update it by construction. prev_token(p) resolves to the token of the largest existing position <= p, with the 'last cell wins' tie-break of the general scan preserved. ubatches with shared temporal positions (multimodal) are detected as not-applicable and keep using the general scan: within a shared position, cells resolve by ubatch order, not by row. LLAMA_KV_PREV_TOKENS env: fast (default) | verify | off. verify runs both paths for every call and logs mismatches + a cost heartbeat; used to validate the index against production traffic (455k lookups across prompt-cache loads and checkpoint restores, 0 mismatches; scan 45.9 ms vs index 6.9 us avg). measured (qwen4exp, 2xL40S, unified KV, 256k ctx): tg @155k ctx: ~11.7 -> ~31 t/s tg @200k ctx: ~11.5 -> ~27 t/s related upstream work: ggml-org#27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), ggml-org#27977 (shrinks the general scan constants + qsa gather windows). complementary layers; can be combined.
| // the full set of sequences this cell is visible to. two cells with the same | ||
| // set are visible to exactly the same sequences. | ||
| // note: seq_get() above returns the single id and requires a one-sequence cell | ||
| const seq_set_t & seq_get_all(uint32_t i) const { |
There was a problem hiding this comment.
| // the full set of sequences this cell is visible to. two cells with the same | |
| // set are visible to exactly the same sequences. | |
| // note: seq_get() above returns the single id and requires a one-sequence cell | |
| const seq_set_t & seq_get_all(uint32_t i) const { | |
| // the full set of sequences this cell is visible to | |
| const seq_set_t & seq_get_all(uint32_t i) const { |
There was a problem hiding this comment.
I haven't understood yet the logic of this function. Could you explain how it works in more details? It's OK to accept for now since it's impact is limited to Qwen4, but it would be useful for me to know what it does.
Also, I suspect that this implementation will become a CPU bottleneck for larger contexts, so look for optimizations (e.g. reserve vectors, template params, etc.). See for example set_input_kq_mask as a well optimized example. Can optimize in a follow-up PR.
Btw, the existing pattern is to keep these "set input" implementations in the main memory class and to have the memory context just forward the call to it. For example, in llama_kv_cache_context we do like this:
void llama_kv_cache_context::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
kv->set_input_kq_mask(dst, ubatch, causal_attn);
}And the actual logic is implemented in the llama_kv_cache::set_input_kq_mask. You should mirror this here for consistency. In general, the memory contexts should be very "thin" in terms of logic.
| // apply_ubatch() takes the 2D position from the ubatch, and that ubatch is built with this | ||
| // cache's own n_pos_per_embd. a cache that does not use M-RoPE itself but mirrors one that | ||
| // does (the qwen4exp QSA indexer) would drop x and y. put the saved ext back instead, which | ||
| // is what the whole-context path below already does. | ||
| for (uint32_t i = 0; i < (uint32_t) exts.size(); ++i) { | ||
| cells.ext_set(sinfo.idxs[0][i], exts[i]); | ||
| } |
There was a problem hiding this comment.
I wonder if we can extend the tests-save-load-state with a test that would demonstrate this issue on master. Though from the description in unslothai#143, it might be difficult to figure out such a test. Low prio, but mentioning just in case you can think of something simple.
There was a problem hiding this comment.
Yes, 02eb201 is a good idea.
However, even without this patch, the test still succeeds. So it's missing something.
There was a problem hiding this comment.
@danielhanchen In case you missed this comment. The problem seems to be that the generated qwen4 dummy model by test-llama-archs does not have PLE. This makes has_cell_ext() return false:
llama.cpp/src/llama-kv-cache.cpp
Lines 1824 to 1827 in 02eb201
So the added test is not effective. I think we have to extend the dummy model to have valid PLE data.
There was a problem hiding this comment.
Change the test to actually invoke PLE if that helps
|
@ggerganov Let me re-check - also FYI #28068 |
Assisted-by: Claude
seq_get is already taken by the single-id getter, so the suggested name cannot be overloaded on return type alone. Assisted-by: Claude
The context held the whole implementation, where the pattern elsewhere is a thin context forwarding to the memory class, as llama_kv_cache_context does for set_input_kq_mask. The body reads no context state, so it moves unchanged and the context keeps a forwarder. Also shortens the seq_get_all comment as suggested.
Saves seq 0, erases it, restores the blob and saves again, requiring the two blobs to match. Compares blobs rather than generated text, which cannot see a field dropped on the way back in. Note this passes on master for qwen4exp, so it does not demonstrate the ext.x/ext.y drop this PR fixes; reaching that needs 2D mrope content.
6b2b85c to
52bc9ff
Compare
|
@ggerganov Fixed merge conflicts and had to force push |
|
@danielhanchen PTAL at my comments earlier: #27941 (comment) |
|
Oops will check |
ggml-org#27941, squashed) danielhanchen's follow-up series, applied as one diff (d7bd3bf..pr27941) because the per-commit picks conflict with code the series itself later relocates (set_input_qsa moves onto the memory class). Carries: indexer keys follow sequence copies; blocks keyed on (sequence set, index bucket) instead of position alone; image cells cut on rank order under M-RoPE; loader GGML_ASSERTs on malformed GGUFs become throws; state blob round-trip test. Conflict resolution: took the series' side everywhere (validation additions and the set_input_qsa delegation); our out_ids trim, QSA top_k-to-FA pass and per-block bias all survive (blk_bias is a pass-through parameter in the new interface). Squashed from: 132832d 515effd 8161d11 868e2f5 e1f5324 6b2b85c Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
has_cell_ext() is n_pos_per_embd() > 1 || ple_n_heads > 0, and the indexer cache sets rope_type = NONE, so without a PLE it serializes no cell ext at all and the round-trip test cannot see a dropped ext.x/ext.y. With one, removing the ext_set restore in state_read_meta fails the test: 198 of 335692 bytes differ, first at offset 282092. Loading such a model needed two fixes: - the row count of per_layer_token_embd came from require_weight(), which a model synthesised from metadata alone has no file to answer. Derive it from the head ranges and prefer the file's padded count where there is one. - the PLE conv history is a row of the recurrent cache, so a PLE on a full attention layer dereferenced a null p_l. Reject it at load time instead. The meta mirror is skipped for qwen4exp. It returned NaN logits before this fixture carried a PLE, which the nmse check passes since a NaN comparison is false, and aborts with one. -sm tensor on real devices works. Assisted-by: Claude
|
@ggerganov Ok done sorry on the delay |
|
Cool fixes, I'm running tests on my end |
|
I'm not sure if this is relevant here (I can create a new issue), but there's something funny happening with this model (before and after these fixes) at least on the Vulkan backend. Initially this uses 18GB of VRAM. When processing a big prompt (e.g. |
test-llama-archs skipped the tensor split for this arch from inside the test, so the arch still advertised support it does not have. Declare it in llm_arch_supports_sm_tensor instead and drop the test-side exception; the existing llm_arch_supports_sm_tensor branch then does the skipping. Assisted-by: Claude
|
@ggerganov ill do a follow up PR on TP - MTP and some other perf items later |
|
@sdroege checking |
There was a problem hiding this comment.
Went through the C++ against the transformers reference block by block: n-gram hash, PLE layer, QSA indexer, gated residual and the gamma folding all match. Greedy decoding on unsloth's UD-Q4_K_XL, single sequence, 256 tokens, KV cache in f16 and in q8_0: identical output before and after this PR.
Also ran --parallel 4 on unsloth's UD-Q4_K_XL, four concurrent greedy requests, before and after: all four slots answer correctly in both.
|
@ServeurpersoCom @sdroege if ur interested also check unslothai#152 and unslothai#137 and unslothai#144 and unslothai#142 MTP modules at https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF/tree/main/MTP |
@danielhanchen related to that there's also on exit which approximately matches the amount of memory that disappears. |
68 upstream commits, 9 conflicted files. Brings qwen4exp fixes needed for Qwen3.8-Flash-Next: recurrent state rollback (ggml-org#28123), seq_cp/block position keying/mtmd input (ggml-org#27941), indexer head slicing (ggml-org#28023), and the central n_layer_nextn load (ggml-org#28159). Resolutions: - ggml-rpc.cpp: resolved entirely to the fork side again. Upstream's rpc_dispatcher/proto-6.0 port stays deferred; the fork keeps 5.1.3. Note that upstream ggml-org#26500 now implements the same foreign-endpoint skip in serialize_graph that the fork carries, so that one is a candidate to drop when the dispatcher port finally happens. - llama-graph.cpp: adopted upstream's fused ggml_swiglu_clamp for both the dense and MoE clamp paths and kept the fork's LLM_ARCH_GLM5NEXT in the condition. The fused kernel is min(gate,limit) / clamp(up,+-limit) / swiglu_split, i.e. bit-identical to the three-op sequence it replaces, and Vulkan implements it, so the boards get it too. Costs the ffn_*_clamped cb() trace names, same tradeoff upstream took for DSV4. - speculative.cpp: upstream ggml-org#27310 folded the DFlash encoder into the injection decode and deleted features_buf. The fork's non-finite feature scrub moves onto batch_inject.embd in place; the last_tap_nextn tap automerged. - dflash.cpp: kept the fork's Laguna kv_inp norm and deferred gated o_proj, layered on upstream's newly-populated wk_s/wv_s/wo_s scale arguments (previously NULL). The gated path latches wo_s alongside wo. - nemotron-h.cpp: dropped the now-duplicate NEXTN_PREDICT_LAYERS get_key, kept the fork's stricter assert and n_layer_nextn_per_head derivation. - llama-model-loader.cpp: kept the fork's rpc_preloaded skip ahead of upstream's use_mmap || lazy.has(cur) predicate. - tests: kept both sets of flash-attn cases; ported the fork's GLM5NEXT test_dsa_kpool/test_mtp onto upstream's int-verbosity logging convention. Windows CUDA build green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01U3H5motr51eTWujztSXykc
Overview
Qwen3.8-Flash-Next fixes
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.See unslothai#143 for more details - checked PPL before this PR and after for UD-IQ1_S - exact PPL - no need to change GGUFs.