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! |
| } | ||
|
|
||
| // 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 { |
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 |
|
@ggerganov I addressed your issues and updated the PR - unsure if these work - perf is only minutely better and PPL are identical. Also the function is specifically to handle |
|
This commit be79015 is not needed. Better to remove it and just add a TODO in this function to try to optimize it in the future. |
|
Will remove |
02eb201 to
6b2b85c
Compare
|
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
|
@danielhanchen I looked a bit into this and I think that's the compute buffer or more specifically the memory allocated for the QSA mask. That reserves ~9GB for the mask at 128k context for 4k ubatch prompt processing. Later during prompt processing or decoding the graph is different so a realloc happens, that then calculates the size based on the current KV usage, and that's lower (~7GB less in my testcase). AFAIU during decoding this would never require more but if at a deeper context more prompt processing would happen, more memory would be allocated again. Not sure if there's something to fix here (my expectation would be that it just stays allocated at the worst case), or something to optimize with how QSA is implemented (this seems like a lot of memory for the QSA mask alone). |
Ports every upstream qwen4exp (Qwen3.8-Flash-Next) commit from the past 10 days that this fork's manual PR port had not received: - reduce graph splits by hoisting the PLE embedding gather out of the per-layer loop (ggml-org#27880) - sum indexer heads via strided adds instead of transpose+sum_rows (ggml-org#28023) - support recurrent state rollback for MTP speculative decoding (ggml-org#28123) - rewrite QSA sparse-attention block/bias selection: fixes NaN-producing bias rows for short sequences, fixes cross-sequence block pooling in a unified KV cache, adds mrope duplicate-position ranking, and fixes a CUDA rms_norm gridDim.y overflow (ggml-org#27941) - indexer cache seq_cp staleness fix, ext.x/ext.y state-restore fix, PLE-must-be-linear-attention validation, correct -sm tensor disablement (ggml-org#27941) - Hadamard k_rot context-shift crash fix, shared with other archs (ggml-org#27967) Also replaces raw GGML_ASSERT aborts in hparams loading with proper error messages, and adds test coverage: a PLE fixture in test-llama-archs (which required porting the per_layer_token_embd row-count-from-metadata fix to make it loadable) and a state round-trip test in test-save-load-state. Verified against the real Qwen3.8-Flash-Next model: correct generation at short and long (~66k token) context, and test-llama-archs passes qwen4exp on both CUDA and CPU.
- llama-model.cpp: drop unused mtp_on_hybrid_nemotron (from ggml-org#27941 pick; fork's nemotron path diverges, variable never read) — fixes cuda/ubuntu CI - ggml-cpu.c GGML_MOE_LOG: flockfile/funlockfile are POSIX-only; use _lock_file/_unlock_file on _WIN32 — fixes windows CI builds
…-org#28123) build_conv_state_at wrote a single conv-state snapshot plane, but seq-rollback restores state at multiple positions — a rollback restored a convolution history that was never captured. Now writes one snapshot per slot, each ending one token earlier, for both the delta-net QKV convolution and the PLE convolution. Adds QWEN4EXP to the rs_rollback arch list. Upstream 0eadefe, clean 3-way apply. Candidate fix for the sticky 'CUDA error: unspecified launch failure' seen after slot switches with partial KV reuse (rollback path) — crashed on both MTP and non-MTP serving with ggml-org#27941 already applied.
…add tests (ggml-org#27941) * qwen4exp: follow up fixes * -kvu NaN collapse fix Assisted-by: Claude * indexer cache ext.x/ext.y restore fix Assisted-by: Claude * kv-cells: rename seq_set to seq_get_all seq_get is already taken by the single-id getter, so the suggested name cannot be overloaded on return type alone. Assisted-by: Claude * memory-hybrid-idx: implement set_input_qsa on the memory class 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. * tests: check that a sequence state survives a save/restore round-trip 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. * tests: give the synthetic qwen4exp a PLE so the state test bites 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 * llama: disable -sm tensor for qwen4exp 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
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.