Skip to content

qwen4exp: follow up fixes - #27941

Merged
ggerganov merged 8 commits into
ggml-org:masterfrom
danielhanchen:qwen4exp/followup-fixes-squashed
Sep 1, 2026
Merged

qwen4exp: follow up fixes#27941
ggerganov merged 8 commits into
ggml-org:masterfrom
danielhanchen:qwen4exp/followup-fixes-squashed

Conversation

@danielhanchen

Copy link
Copy Markdown
Contributor

Overview

Qwen3.8-Flash-Next fixes

  • Sequence copies lost their indexer keys. The update context never built ctx_idx, so a copied sequence kept the destination stream's stale keys. Reachable with no flags through the OpenAI n parameter, and it silently produced wrong output.
  • Blocks were keyed on position alone. Correct only for a single sequence, so under --kv-unified a block could be pooled from another sequence's cells. Now keyed on (sequence set, index bucket).
  • Images collapsed into one block slot. Under M-RoPE every token of an image shares a position, so 299 of 300 image cells were scored by a pooled key they are not part of. Blocks are now cut on rank order.
  • Malformed metadata aborted the process. Eight GGML_ASSERT sites reachable from a hand edited GGUF are now throws, plus two cases that were being accepted in silence.
  • Fixes a pre-existing CUDA abort at long context, where the pooled block count reached the 65535 gridDim.y limit at n_kv 262144.

See unslothai#143 for more details - checked PPL before this PR and after for UD-IQ1_S - exact PPL - no need to change GGUFs.

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: Yes - Claude and Local models. PPL tested before and after and exactly same

@yhn112

yhn112 commented Aug 29, 2026

Copy link
Copy Markdown

Hope it helps with #27886

@matteoserva

Copy link
Copy Markdown
Contributor

Not the author of that issue but this fixed #27886 for me

@github-actions github-actions Bot added the model Model specific label Aug 29, 2026
drluoto added a commit to drluoto/llama.cpp that referenced this pull request Aug 29, 2026
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>
drluoto added a commit to drluoto/llama.cpp that referenced this pull request Aug 29, 2026
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>
EchterAgo pushed a commit to EchterAgo/llama.cpp that referenced this pull request Aug 30, 2026
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.
@ggerganov

Copy link
Copy Markdown
Member

@danielhanchen Are you planning more changes to this branch?

@danielhanchen
danielhanchen marked this pull request as ready for review August 30, 2026 10:55
@danielhanchen

Copy link
Copy Markdown
Contributor Author

@ggerganov I think it should be ready - the rest are just perf + MTP which I'll land tomorrow!

Comment thread src/llama-kv-cells.h Outdated
}

// two cells with the same set are visible to exactly the same sequences
const seq_set_t & seq_set(uint32_t i) const {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rename this to seq_get since it gets the seq member:

Suggested change
const seq_set_t & seq_set(uint32_t i) const {
const seq_set_t & seq_get(uint32_t i) const {

ilmmatias pushed a commit to ilmmatias/llama.cpp that referenced this pull request Aug 30, 2026
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.
tvanderka added a commit to tvanderka/llama.cpp that referenced this pull request Aug 30, 2026
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.
Comment thread src/llama-kv-cells.h Outdated
Comment on lines +306 to +309
// 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 {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// 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 {

Comment thread src/llama-memory-hybrid-idx.cpp Outdated
Comment on lines 349 to 356

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/llama-kv-cache.cpp
Comment on lines +2472 to +2478
// 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]);
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, 02eb201 is a good idea.

However, even without this patch, the test still succeeds. So it's missing something.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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:

bool llama_kv_cache::has_cell_ext() const {
// M-RoPE needs the 2D position, the PLE n-gram hash needs the token id
return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0;
}

So the added test is not effective. I think we have to extend the dummy model to have valid PLE data.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Change the test to actually invoke PLE if that helps

@danielhanchen

Copy link
Copy Markdown
Contributor Author

@ggerganov Let me re-check - also FYI #28068

@danielhanchen

Copy link
Copy Markdown
Contributor Author

@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 --kv-unified as QSA for the KV cache interestingly doesn't store tokens like normal KV cache - the positions are random, so the function needs to work out which cell belongs to which bock

@ggerganov

Copy link
Copy Markdown
Member

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.

@danielhanchen

Copy link
Copy Markdown
Contributor Author

Will remove

@danielhanchen
danielhanchen force-pushed the qwen4exp/followup-fixes-squashed branch from 02eb201 to 6b2b85c Compare August 31, 2026 10:44
Comment thread tests/test-llama-archs.cpp Outdated
@ServeurpersoCom

Copy link
Copy Markdown
Contributor

Cool fixes, I'm running tests on my end

@sdroege

sdroege commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

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.

llama/llama-cli -m models/unsloth/Qwen3.8-Flash-Next-GGUF/Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf --jinja -c 131072 -ngl 999 -cmoe --load-mode mmap -ub 4096 -b 8192

Initially this uses 18GB of VRAM. When processing a big prompt (e.g. /read some-8k-token-file), VRAM usage drops to 11GB (and those 7GB don't go to GTT either) and never recovers. Which seems at least wasteful as those 7GB could be used by something more persistent otherwise. I think some graph reallocation is happening but not entirely sure.

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
@danielhanchen

Copy link
Copy Markdown
Contributor Author

@ggerganov ill do a follow up PR on TP - MTP and some other perf items later

@danielhanchen

Copy link
Copy Markdown
Contributor Author

@sdroege checking

@ServeurpersoCom ServeurpersoCom left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@danielhanchen

danielhanchen commented Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

@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

@ggerganov
ggerganov merged commit 36b1015 into ggml-org:master Sep 1, 2026
22 of 26 checks passed
@sdroege

sdroege commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

@sdroege checking

@danielhanchen related to that there's also

[48011] 2.57.057.057 W ~llama_context:    Vulkan1 compute buffer size of 2532.7807 MiB, does not match expectation of 9950.5625 MiB

on exit which approximately matches the amount of memory that disappears.

MarkShark2 added a commit to MarkShark2/llama.cpp that referenced this pull request Sep 2, 2026
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
oobabooga added a commit to unslothai/llama.cpp that referenced this pull request Sep 2, 2026
@sdroege

sdroege commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

@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).

TheTom pushed a commit to TheTom/llama-cpp-turboquant that referenced this pull request Sep 3, 2026
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.
turbo-tan pushed a commit to turbo-tan/llama.cpp-tq3 that referenced this pull request Sep 4, 2026
- 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
turbo-tan pushed a commit to turbo-tan/llama.cpp-tq3 that referenced this pull request Sep 4, 2026
…-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.
fewtarius pushed a commit to fewtarius/CachyLLama that referenced this pull request Sep 5, 2026
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

model Model specific testing Everything test related

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants