qwen4exp: upstream catch-up, sparse-fa, and NextN/MTP support - #340
Conversation
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.
Prerequisite for the upcoming sparse-fa flash attention path: adds the smem swizzle layout (fattn-swizzle.cuh) that the sparse gather load tiles need. Reconciled against TurboQuant's turbo2/3/4 SRAM tile loaders in fattn-mma-f16.cuh, which are untouched by this change. Cherry-picked from upstream e4b9af0.
Prerequisite for the upcoming sparse-fa flash attention path: the new sparse-fa test cases need kv_view/v_is_view_of_k as named constructor parameters instead of hardcoded booleans. Also wires kv_view through to the K/V tensor creation (previously hardcoded true), matching an earlier upstream commit's intent that our tree never received. Cherry-picked from upstream 5fff128, plus the kv_view wiring fix.
Cherry-picked from upstream 8e93a97 on top of the swizzle and V-is-view-of-K prerequisites. Adds ggml_flash_attn_ext_set_n_kv_max(): the CUDA flash-attention kernel can treat the mask's finite entries as a sparse K/V set and skip the rest, instead of computing dense attention and masking it out. DeepSeek-V4's own top-k sparse-attention path is wired to it. Reconciled against TurboQuant's turbo2/3/4 tile loaders and dedicated fattn-mma-turbo.cuh kernel, which the upstream diff has no knowledge of. Two real bugs were caught and fixed in that reconciliation, not just merge conflicts: - fattn-mma-turbo.cuh instantiated flash_attn_ext_f16<...> with the old positional template argument list; type_K was landing in the new use_sparse slot (a nonzero ggml_type implicitly converts to true), and type_K/type_V were shifted off the end entirely. - Its launch_fattn call had the same problem one level up: the new use_sparse parameter was inserted before warp_size, so the trailing warp_size_host argument would have silently become use_sparse=true and warp_size would have silently fallen back to its default. - qwen35.cpp's MTP draft-head graph calls build_attn_mha directly and wasn't part of upstream's diff at all, so it still had the pre-change 9-argument signature. qwen4exp's own build_attn_mha call is updated for the new signature but left at n_kv_max=0 (disabled): the sparse kernel path is compile-time gated to DeepSeek-V4/GLM's specific MLA head shapes (512/512 or 576/512, GQA 8 or 16). qwen4exp's actual shape is 256/256 with GQA 12, which doesn't match, so enabling it would take no effect today. Extending the shape gate to cover qwen4exp is separate, higher-risk kernel work not attempted here. Verified: full build, test-llama-archs passes on CUDA and CPU with no regressions (qwen4exp, qwen35, qwen35moe, and all other architectures).
Ports upstream PR ggml-org#27836 (open, unmerged) on top of this fork's qwen4exp support: the MTP head folds the next token's embedding into the trunk's wide hyper-connection residual, runs one trunk-shaped block (dense attention + MoE) over it, and collapses the result with its own hyper-connection mixer before reusing the trunk's LM head. Also updates the converter and gguf-py tensor mappings so conversion/qwen4exp.py --mtp can export the head. Also adds mtp_only/trunk_flags handling to load_arch_tensors, which the ported PR didn't include: without it, loading a standalone MTP-only checkpoint (all trunk tensors absent) fails outright instead of tolerating their absence, unlike the equivalent qwen35.cpp path. Detects MTP-only via the absence of blk.0.hc_attn_norm.weight, a tensor every trunk layer carries. Guards the embedding and LM-head fallbacks with a clear assert instead of a null-pointer crash when a checkpoint has neither a dedicated NextN embedding/head nor the trunk's own. Verified: full build, test-llama-archs shows no regressions. Loading a real Qwen3.8-Flash-Next MTP draft checkpoint (mtp-Qwen3.8-Flash- Next-shared-Q4_K_M.gguf) now gets past tensor loading correctly (the mtp_only path works) and fails with the new clear assertion rather than "token_embd.weight not found": that specific checkpoint has neither nextn.embed_tokens nor a trunk token_embd.weight anywhere, confirmed by inspecting its raw safetensors source directly (15 tensors total, no embedding table in any form). That's a model packaging gap in that specific file, not a code gap.
qwen4exp MTP was crashing with GGML_ASSERT(nb10 % sizeof(src1_t) == 0) in ggml_cuda_op_bin_bcast when running real MTP-packaged checkpoints that store some norm weights (e.g. nextn.enorm) as F16. The CUDA broadcast-mul dispatch has a path for F16-activation x F32/F16-weight, but none for F32-activation x F16-weight, which build_norm's grouped RMSNorm+scale pattern hits whenever a norm weight isn't F32. Cast the weight to F32 first when this combination is detected; a no-op for the common case where norm weights are already F32. Fixes qwen4exp NextN/MTP speculative decoding startup crash against real Qwen3.8-Flash-Next MTP GGUFs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…al NextN heads The speculative driver's is_mem_shared check (true whenever no separate -md draft model is given, e.g. qwen4exp's self-contained MTP) was being used to select Gemma4-assistant's same-llama_pos-for-every-draft-token behavior. That behavior is specific to Gemma4-assistant's early-exit self-speculation, not a property of "shares KV memory with target" in general - a real trained NextN head (qwen35, qwen4exp) still needs an incrementing position per draft step even in that mode. Add llama_model_uses_shared_position_draft(), gated on arch == LLM_ARCH_GEMMA4_ASSISTANT specifically, and use it instead of the blanket is_mem_shared check for that one branch. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Ported PR 27836 gave the MTP block its own private per-layer final mixer (layer.nextn.hc_head_norm/down/up, tensor names blk.N.nextn.hc_head_*). The actual, community-validated implementation (PR 27739, reconciled by LaurentZuijdwijk - the one real MTP-head exports like dzannotti/Qwen3.8-Flash-Next-MTP-GGUF are built against) instead trains a single hc mixer shared between the trunk's own final layer and the MTP head (exported as top-level output_hc_norm/down/up, loaded here as model.hc_head_norm/down/up). Files following that convention have no blk.N.nextn.hc_head_* tensors at all, so graph_mtp's GGML_ASSERT on it would abort at graph-build time; conversely, per PR 27836's convention the head ran through a mixer that was never actually trained as a distinct MTP-specific output projection. Confirmed against a real MTP head export off HF that model.hc_head_* is the correct, always-present tensor for this: switch graph_mtp to use it instead of the private per-layer copy. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
graph_mtp gathered attention output and res_hc down to just the requested output rows immediately after attention, then exported that already-gathered res_hc as t_h_nextn. Any decode with fewer requested-output rows than input rows - notably the speculative driver's catch-up/prefill decode into the draft context, which requests logits for none of its rows - would export a zero-row (or otherwise truncated) t_h_nextn. The driver's per-token shift-by-one hidden-state handoff needs one row per input token regardless of which rows have requested logits. Matches PR 27739's (the community-validated implementation) ordering: res->t_h_nextn is assigned the full per-token combine result, and the inp_out_ids gather happens after, scoped to just the final hc_head mix / LM head projection. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
graph::graph gathers cur/res_hc down to the requested output rows exactly once - via the early gather right after the last layer's attention when gather_now is true, or the deferred one inside the embeddings_nextn block when it's false - the two conditions are complements, so exactly one path always fires when inp_out_ids exists. A third, unconditional gather right before the final output norm re-applied inp_out_ids to the already-reduced tensor, indexing with values sized for the original token count against a tensor that now only has output_row_count rows. This is silently harmless whenever n_outputs == n_tokens (every position requested logits, gathering by the identity permutation twice is a no-op), which is why normal generation never hit it. It reliably crashes (GGML_ASSERT(i01 >= 0 && i01 < ne01) in ggml_compute_forward_get_rows) whenever n_outputs < n_tokens on a context with embeddings_nextn set - concretely, the standard llama.cpp warmup decodes 2 tokens and requests logits for 1, and any target context paired with an MTP speculative draft sets embeddings_nextn unconditionally. So this fired on every server startup once a real -md draft-mtp config was used, independent of which draft checkpoint or graph_mtp bugs were involved. Removing the redundant gather - res_hc is already correctly sized by the time this runs - fixes it. Also drop the now-dead requirement on layer.nextn.hc_head_norm/down/up (graph_mtp reuses model.hc_head_* as of the previous commit): make them TENSOR_NOT_REQUIRED so files exporting per-block hyper-connection mixers in the old PR 27836 tensor layout don't fail to load, without requiring the tensor from files that (correctly) don't have it. Verified against a real, cleanly-exported MTP draft head (dzannotti/Qwen3.8-Flash-Next-MTP-GGUF) paired via -md: loads without crashing, and speculative decoding now gets real acceptance (26/26 on a trivial prompt, 71/88 = 80.7% on natural-language generation - in line with the community-reported 0.74-0.90 range for this same head). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The CUDA graph cache keyed captured graphs by the raw memory address of their first node (cgraph->nodes[0]). A captured graph hard-codes its shapes, but speculative decoding constantly alternates between different batch shapes (draft steps, verify batches, catch-up decodes) on the same context - when a new shape happens to reuse the same first- node address as a stale cached graph for a different shape, capture either reuses the wrong graph or thrashes, permanently resetting warmup instead of ever converging to steady-state replay. Hash node count and both endpoint tensors' shapes into the key instead (O(1) - walking all nodes would defeat the point of a CUDA graph), add LRU eviction capped at 64 graphs so the map can't grow unbounded now that distinct shapes get distinct entries. A shape this still fails to separate re-captures exactly as before, so it can't regress anything. Cherry-picked from ggml-org#28243 (open, unmerged), which found this while working on qwen4exp MTP performance - the effect is generic to any speculative-decoding workload on this fork, not qwen4exp-specific, so pulling in just this piece rather than the rest of that PR (which also reworks qwen4exp trunk/draft tensor sharing and doesn't fix the mixer/export-timing bugs already fixed on this branch). Verified: qwen4exp MTP speculative decoding still produces correct, byte-identical (temp 0) output after this change. Speed effect is hard to isolate cleanly from the dominant MoE-cache warmup effect already documented on this branch, but the fix is justified on its own correctness merits regardless of measured delta. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
hi if you have time kindly check this too Thanks, Mark. That completes the DS4 hardware matrix. I pushed 06b66fd with the upstream LIGHTNING_INDEXER Meta handler. It requires all four inputs to be mirrored and returns a mirrored result, matching every occurrence you traced while rejecting any unsupported split input rather than silently accepting it. Local results on the exact head: Your patched-tree results already cover the important real-hardware cases on this exact rule: coherent output on 2 devices, 4 devices, and uneven 3,4,4,1, including fresh prompts, sequence reset, and prefix reuse. No further DS4 retest is needed unless fresh CI finds something. I am tracking the long-context Qwen4Exp allocation report separately. Its log shows a 5.96 GiB compute-buffer request failing with --batch-size 6700 --ubatch-size 6700; the subsequent segfault after allocation failure is not an acceptable failure mode, but it is separate from the DS4 split correctness work above. turbo qunat is the only fork thats support proper split and proper buffer size |
./build/bin/llama-server \
-m /mnt/storage/models/qwen3.8/flash/UD-IQ4_XS/Qwen3.8-Flash-Next-UD-IQ4_XS-00001-of-00003.gguf \
-md /mnt/storage/models/qwen3.8/MTP/Qwen3.8-Flash-Next-MTP-Q4_K_M.gguf \
-ngld auto \
-c 32000 -ngl auto --fit on -fa on --jinja -t 8 -Cr 0-7 -tb 8 -Crb 0-7 \
--cpu-strict 1 -b 512 -ub 512 -ctv q8_0 -ctk q8_0 \
--top-p 0.95 --top-k 20 --temp 1.0 --parallel 1 \
--host 0.0.0.0 --port 8099 --moe-cache auto \
--spec-type draft-mtp --spec-draft-n-max 3 --spec-draft-p-min 0.75 \
--chat-template-kwargs '{"reasoning_effort":"medium"}' \
--chat-template-file /mnt/storage/llama-server/qwen3.6-chat-template/chat_template.jinja \
--reasoning-preservePrompt 1:
Follow up Prompt 2 about improving a python script:
|
|
Thanks for the writeup, it made this reviewable. Most of it holds up: the sparse-fa port is byte-for-byte upstream One blocker, one thing to fix with it, and one bookkeeping item. Blocker: the XOR-swizzle port breaks turbo2/3/4 flash attention on Turing and newer.
The f16 loader was updated to write through That is the default case, not a corner: Reproduced on a GB10 (sm_121, driver 580.173.02), your head All 743 FLASH_ATTN_EXT failures are Same box, same flags, base branch at 80be9a7: Fix: make the three turbo writers go through Bookkeeping: the description lists 11 commits including Smaller items, fine as follow-ups:
Once the turbo writers are swizzle-aware and a turbo FA case runs green on CUDA, I'm happy to merge this. |
TheTom
left a comment
There was a problem hiding this comment.
Requesting changes for the turbo2/3/4 swizzle mismatch detailed in the comment above (743 CUDA FLASH_ATTN_EXT failures on GB10, base passes 7658/7658). Everything else looks good.
…aders The turbo2/turbo3/turbo4 SMEM loaders wrote dequantized K/V elements with plain linear indexing while the XOR-swizzled ldmatrix reads (used on Turing+ once nbatch_K2/V2 is bank-aligned) expected byte offsets from fattn-swizzle.cuh's bytes_rc(). Every read landed on the wrong element, producing garbage turbo attention output whenever swz was active. Add turbo_store_h2<stride_tile, swz>(), mirroring the f16 loader's existing if constexpr (swz) pattern, and route all six turbo write sites through it. Also fix fattn-mma-turbo.cuh's SMEM sizing to use tile_stride() instead of the stale nbatch_K2+4/nbatch_V2+4 formula, matching the actual swizzled stride. Fixes the turbo3/turbo4 FLASH_ATTN_EXT regression reported in PR TheTom#340. All 10955 FLASH_ATTN_EXT test-backend-ops cases pass, including the previously-failing turbo4_vec_q8_0_turbo4_d128_kv256/d256_kv256 cases. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…TN_EXT The general test_flash_attn_ext matrix already exercised turbo3/turbo4 K==V at hsk=128 with batched nb and GQA nr2>=2, but excluded hsk=256 entirely (only 64/72/128 were allowed for non-F16 types) and never included turbo2_0 at all. That's exactly the gap that let the turbo swizzle write/read mismatch ship: the only hsk=256 turbo coverage was the hand-rolled turbo4_vec case, which is nb=1/nr2=1/no-mask only and doesn't exercise the batched MMA path real speculative-decode batches (nb up to 4) or GQA-packed decode go through. Requested in PR TheTom#340 review: hsk 128 and 256, nb<=4, nr2>=2. Verified: 2592/2592 filtered FLASH_ATTN_EXT cases pass (turbo2/3/4, hsk=256, nb in {1,3,32,75}, nr23=[4,1], masked and unmasked). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Four smaller items from TheTom's review, none of them blockers: - deepseek4.cpp: assert n_kv_max (raw SWA window + top-k count) doesn't exceed k_all's actual concat length (raw_k + csa_k), documenting an invariant that held but was previously unchecked. - llama-kv-cache.cpp: llm_graph_input_k_shift::set_input guarded k_rot on both null and ->buffer before use, but only null-checked k_shift. set_input_k_shift asserts on ggml_backend_buffer_is_host(dst->buffer), so a graph-reserve pass (tensors exist, backends not yet allocated) would fault there. Add the same buffer check used everywhere else in llama-graph.cpp's set_input overrides. - llama-graph.cpp build_norm: note why the mw/mb F32 cast is cheap enough to re-insert per graph build (n_embd-sized weight) rather than caching across builds. - fattn.cu: clarify why quantized K/V (including turbo2/3/4) are skipped in the GQA-opt alignment check — their loaders dequantize via swizzled/padded SMEM tiles rather than reading nb[] directly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Corrected three mistargeted fixes from the previous commit against TheTom's actual review text (fetched via gh pr view, not re-derived from memory): - deepseek4.cpp:784: the ask was an assert catching a misconfigured n_swa == 0 model (which would silently zero out the raw SWA window and leave only csa top-k entries), not a bounds check against k_all->ne[2]. Add GGML_ASSERT(hparams.n_swa > 0); keep the bounds assert too, it documents a real, separate invariant. - llama-kv-cache.cpp:2248: the ask was about k_rot's own silent-skip semantics (is it ever legitimate, and when), not k_shift's missing buffer guard. Document that k_rot is null pre-attn_rot_k-setup for non-rotating caches and unallocated only during graph-reserve. (k_shift's missing ->buffer check from the prior commit was a real, separate bug -- set_input_k_shift asserts on it -- and stays fixed.) - fattn.cu: the stale comment was "DEFAULT OFF" above ggml_cuda_turbo_mma_fused(), whose code has defaulted ON for a while; not the alignment-loop comment at the review's line 190 (line numbers had shifted off the review's head commit by the time this landed). Verified: cmake --build --target llama compiles clean, test-llama-archs still green for qwen4exp/deepseek/deepseek2/deepseek32 (the new n_swa assert doesn't trip on any registered synthetic config). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
llama_memory_hybrid_idx's indexer cache (mem_idx) only ever reads K
(the lightning-indexer top-k selection needs keys, never values), but
its hparams were only narrowing n_embd_head_k_full to indexer_head_size
-- n_embd_head_v_full stayed at the full model's real V head dim, and
nothing marked the cache as MLA/K-only. llama_kv_cache always allocates
V storage unless hparams.is_mla() is true (has_v = !is_mla), so this
cache was wasting VRAM on a same-sized V-cache buffer it never touches.
Fix: set n_embd_head_v_full and both n_embd_head_{k,v}_mla_impl to
indexer_head_size, mirroring dsv4_make_k_only's hparams_lid setup in
llama-kv-cache-dsv4.cpp (deepseek4's own lightning-indexer cache uses
the exact same K-only pattern already). Upstream hit the same bug
independently: ggml-org#28330.
Verified: test-llama-archs qwen4exp still OK (NMSE 8.94e-08, roundtrip
OK) on CUDA/CPU.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
Round two on 69499d4. The blocker is fixed properly: all six turbo store sites (zero-fill and main store for turbo2/3/4) go through Fork preservation: I diffed every file from the merge-base and checked each deleted line that does not exist in upstream. Nothing fork-specific goes away. The 195 such lines are signature rewrites for the sparse and swizzle params, the turbo loader fix itself, the FA test sweep gaining turbo2 and hsk=256 cases, the fused-turbo comment corrected to match the code (still default ON), and the fork's older hand-port of qwen4exp / QSA replaced by upstream's newer version. Local writeup kept on my side. GB10 (sm_121), head 69499d4, The no-graphs full sweep is running now as the last gate; I'll merge when it and the remaining CI jobs are green. Follow-ups, none blocking:
|
b5b813d
into
TheTom:feature/turboquant-kv-cache
|
Merged as b5b813d. For the record, the fork-preservation check in detail, since this was a catch-up and the question was whether any fork-only work got dropped. Method: three-dot diff from the PR merge-base fb2cc35, and for every file the PR touches, every deleted line that does not exist in upstream/master and is not re-added in the same file. 34 files, 195 such lines, each group read for meaning.
No Final gates on a GB10 (sm_121), head 69499d4: FLASH_ATTN_EXT 8774/8774 with 2592 turbo cases; test-llama-archs clean; full sweep with |
@TheTom okay i lied more llama.cpp work.
qwen4exp: upstream catch-up, sparse-fa, and NextN/MTP support
Branch:
qwen4-catchup, 11 commits on top offeature/turboquant-kv-cache(merge-base
fb2cc35ab)Scope
Five pieces of work, landed in this order:
in the ~10 days before this branch started, that this fork's earlier
manual port had missed.
reconciled against TurboQuant's own turbo2/3/4 kernels.
--spec-type draft-mtpto qwen4exp for the first time, porting upstream PR 27836 (open,
unmerged) plus a standalone-checkpoint loading fix the PR didn't have.
speculative decoding was completely non-functional (crash, then 0%
acceptance); this is four independent root causes found and fixed
against a real model, detailed below.
open upstream PR (models: Qwen3.8-Flash-Next MTP ggml-org/llama.cpp#28243), a generic (not qwen4exp-specific) fix to
how captured CUDA graphs get cached, relevant to any speculative
decoding workload on this fork.
1. Upstream catch-up (
925792c69)Ports every upstream qwen4exp (Qwen3.8-Flash-Next) commit from the
prior ~10 days that this fork's earlier manual port had not received:
per-layer loop (model: qwen4exp: reduce number of graph splits ggml-org/llama.cpp#27880)
(qwen4exp: sum the indexer heads by slices ggml-org/llama.cpp#28023)
(qwen4exp: support recurrent state rollback ggml-org/llama.cpp#28123)
NaN-producing bias rows for short sequences, fixes cross-sequence
block pooling in a unified KV cache, adds mrope duplicate-position
ranking, fixes a CUDA rms_norm gridDim.y overflow (qwen4exp: follow up fixes ggml-org/llama.cpp#27941)
PLE-must-be-linear-attention validation, correct
-smtensordisablement (qwen4exp: follow up fixes ggml-org/llama.cpp#27941)
2. sparse-fa for DSV4/GLM (
7fd6a7cc0,f3b503950,3ddea31f7)Cherry-picked from upstream
8e93a9773, on top of two prerequisites:7fd6a7cc0— XOR swizzle layout for flash-attn K/V smem fp16 tiles(
fattn-swizzle.cuh), needed by the sparse gather-load tiles.Reconciled against TurboQuant's own turbo2/3/4 SRAM tile loaders in
fattn-mma-f16.cuh, left untouched.f3b503950— makes the FA V-is-view-of-K test case a namedconstructor parameter instead of a hardcoded bool, and wires
kv_viewthrough to K/V tensor creation, matching the new sparse-fatest cases' needs.
3ddea31f7— the feature itself:ggml_flash_attn_ext_set_n_kv_max()lets the CUDA flash-attention kernel treat the mask's finite entries
as a sparse K/V set and skip the rest, instead of computing dense
attention and masking it out. DeepSeek-V4's own top-k
sparse-attention path is wired to it. Reconciled against
TurboQuant's dedicated
fattn-mma-turbo.cuhkernel, which theupstream diff had no knowledge of — two real positional-argument
bugs were caught and fixed here (a
use_sparsebool landing in thewrong template/call-site slot in two places) before they could ship.
Confirmed neutral/no-regression via full test suite + benchmark, as
expected since qwen4exp's attention shape doesn't hit the new
kernel's gate.
3. Initial NextN/MTP port (
209a77e55)Ports upstream PR ggml-org#27836 (open, unmerged) on top of
this fork's qwen4exp support: the MTP head folds the next token's
embedding into the trunk's wide hyper-connection residual, runs one
trunk-shaped block (dense attention + MoE) over it, and collapses the
result with its own hyper-connection mixer before reusing the trunk's
LM head. Also updates the converter and gguf-py tensor mappings so
conversion/qwen4exp.py --mtpcan export the head.Also adds
mtp_only/trunk_flagshandling toload_arch_tensors,which the ported PR didn't include: without it, loading a standalone
MTP-only checkpoint (all trunk tensors absent) fails outright instead
of tolerating their absence, unlike the equivalent
qwen35.cpppath.Detects MTP-only via the absence of
blk.0.hc_attn_norm.weight, atensor every trunk layer carries.
At this point the port loaded but speculative decoding did not
work — first a hard crash, then (once crash-free) 0% draft
acceptance. The next four commits are the root-cause fixes for that,
found and verified against real models rather than assumed correct.
4. NextN/MTP bug fixes (5 commits)
--spec-type draft-mtpfor qwen4exp went from "crashes on load" to"loads but gives 0% draft acceptance" to "works, with real measured
acceptance in line with the community-validated implementation." Five
commits, four independent root causes, each confirmed against a real
model rather than assumed fixed.
Files changed (this section only)
Commits
1.
7533986a2— cast F16 norm weights to F32 before CUDA broadcast-mulThe original bug report:
llama-serveraborted on load withGGML_ASSERT(nb10 % sizeof(src1_t) == 0)inggml_cuda_op_bin_bcast(
binbcast.cu). Root cause: that CUDA kernel has a dispatch path forF16-activation × F32-or-F16-weight, but none for F32-activation ×
F16-weight. Real MTP checkpoints store some norm weights
(
nextn.enorm, etc.) as F16 — unconventional but valid — while trunknorm weights happen to be F32 in every checkpoint tested so far, which
is why this never fired outside the MTP path. Fixed in the three call
sites that mix a norm weight into an activation (
build_norm,build_hc_mix,build_ple): cast the weight to F32 first when thecombination would otherwise hit the missing dispatch branch. No-op
when weights are already F32.
2.
3888e18f0— don't reuse the Gemma4-assistant same-position draft pathThe speculative driver's
is_mem_sharedflag (true whenever noseparate
-mddraft model is given) was being used to selectGemma4-assistant's same-
llama_pos-for-every-draft-token behavior.That behavior is specific to Gemma4-assistant's early-exit
self-speculation, not a general property of "shares KV memory with
the target." Added
llama_model_uses_shared_position_draft(), gatedon
arch == LLM_ARCH_GEMMA4_ASSISTANTspecifically, and used itinstead of the blanket
is_mem_sharedcheck for that one branch. Realfix, though it turned out not to be the cause of the 0%-acceptance
issue investigated afterward — self-contained MTP (no
-md) on thisarchitecture happened to resolve
is_mem_shared=falseanyway.3.
97686af29— MTP head should reuse the trunk's shared hc_head mixerThe ported PR (27836) gave the MTP block its own private per-layer
final mixer (
layer.nextn.hc_head_norm/down/up,blk.N.nextn.hc_head_*on disk). The actual, community-validatedimplementation (PR 27739, reconciled by LaurentZuijdwijk — what real
MTP-head exports like
dzannotti/Qwen3.8-Flash-Next-MTP-GGUFarebuilt against) instead trains one hc mixer shared between the
trunk's own final layer and the MTP head, exported as top-level
output_hc_norm/down/upand loaded asmodel.hc_head_*. Confirmeddirectly against a real MTP head export: it has no
blk.N.nextn.hc_head_*tensors at all. Switchedgraph_mtpto usemodel.hc_head_*.4.
b7f556e80+e5a17b5b9— t_h_nextn export timing and a trunk double-gatherTwo related bugs in the hidden-state handoff:
graph_mtpgatheredres_hcdown to output-only rows beforeexporting
t_h_nextn. Any decode with fewer requested-output rowsthan input rows — concretely, the driver's catch-up/prefill decode,
which requests logits for none of its rows — exported a zero-row
hidden state instead of the full per-token state the driver's
shift-by-one handoff needs. Fixed by exporting before the gather,
matching PR 27739's ordering.
The trunk graph (
graph::graph, notgraph_mtp) has two gatherpaths for
res_hc/cur— an early one and a deferred one — gatedby complementary conditions, so exactly one always fires when
inp_out_idsexists. A third, unconditional gather right beforethe final output norm re-applied
inp_out_idsagain, indexing analready-reduced tensor with values sized for the original token
count. Silently harmless whenever
n_outputs == n_tokens(whynormal generation never hit it — a redundant identity gather is a
no-op), but a hard crash (
GGML_ASSERT(i01 >= 0 && i01 < ne01)inggml_compute_forward_get_rows) whenevern_outputs < n_tokensona context with
embeddings_nextnset — which is unconditionallytrue for any target paired with an MTP draft. This one crashed on
literally every server startup with a real
-md draft-mtpconfig,independent of which draft checkpoint was used or which
graph_mtpbugs remained. Fixed by removing the redundant gather.
Also made
layer.nextn.hc_head_norm/down/upTENSOR_NOT_REQUIRED(dead weight as of commit 3, kept optional only so PR-27836-style
exports that still carry it continue to load).
5. CUDA graph keying fix (
9f8cb2d3d)While investigating whether anything else upstream would help, found
PR #28243 (open,
unmerged, built on top of ggml-org#27836), which claims "1.3 to 2x faster MTP
support for Qwen3.8-Flash-Next." Checked its diff line-by-line against
what's already fixed here: it does not fix the mixer-sharing or
export-timing bugs above (both still present in its
graph_mtp,unchanged from ggml-org#27836) — its qwen4exp-specific parts are a different
feature (
embed_tokens/outputsharing between trunk and draft tosave VRAM/disk) and a clearer error message for a different edge case
(loading a draft-only export standalone instead of via
-md).The actual source of its speed claim is generic CUDA infrastructure,
unrelated to qwen4exp: the CUDA graph cache keyed captured graphs by
the raw memory address of their first node. A captured graph
hard-codes its shapes, but speculative decoding constantly alternates
between different batch shapes (draft steps, verify batches, catch-up
decodes) on the same context — when a new shape happens to reuse the
same first-node address as a stale cached graph for a different shape,
capture thrashes instead of ever reaching steady-state replay. Cherry-
picked just this piece (not the qwen4exp-specific parts of ggml-org#28243):
hash the node count and both endpoint tensors' shapes into the key
instead of the raw pointer, with LRU eviction capped at 64 graphs.
This benefits any speculative-decoding workload on this fork, not just
qwen4exp. Verified no regression (output still byte-identical at
temp 0); its own speed contribution is hard to isolate cleanly from
the MoE-cache warmup effect documented below, which dominates in the
same test runs — kept on correctness grounds (the old keying scheme
had a real cache-key collision problem) independent of the measured
delta.
Validation
Models used
Target:
unsloth/Qwen3.8-Flash-Next-GGUF, quantUD-Q4_K_XL(4shards, ~112 GB), an unmodified quantization of the official
Qwen/Qwen3.8-Flash-Nextrelease — no abliteration/uncensoring.Draft head:
dzannotti/Qwen3.8-Flash-Next-MTP-GGUF, fileQwen3.8-Flash-Next-MTP-Q4_K_M.gguf(2.62 GB), downloaded to/mnt/storage/models/qwen3.8/MTP/. This is the MTP block that shipsinside the official
Qwen/Qwen3.8-Flash-Nextcheckpoint but thatupstream's own converter (PR 27742) drops during conversion — this
repo exports it standalone (
convert_hf_to_gguf.py --mtp) so it canbe attached to any clean target GGUF via
-md, without needing torequantize the ~180B-parameter full model. Chosen specifically because
it's not built from an abliterated/uncensored merge, unlike the
checkpoint in the original bug report (see below).
Both are the models referenced in
run.txtin the same folder as thisdocument, which has the exact launch command used for the benchmarks
below.
Correctness
with any real
-md draft-mtpconfig, independent of checkpoint —see commit
e5a17b5b9).the community-reported 0.74–0.90 range for this exact head.
The original bug-report checkpoint (a third-party "Uncensored"
abliterated merge, since deleted from local storage) still gets 0%
acceptance with all four fixes applied — that confirms it's the
checkpoint's own weight data at fault, not a code issue.
Speed: MTP on vs off, same target, same hardware, same prompt
All runs:
UD-Q4_K_XLtarget,-c 8192,-fa on,temp 0(deterministic — MTP-on and MTP-off produce byte-identical output).
MTP runs add
-md <head> -ngld 999 --spec-type draft-mtp --spec-draft-n-max 3 --spec-draft-p-min 0.75.The cold-cache number is the honest "what you get from a single fresh
request" figure. The warm-cache number is what a sustained session
converges to — see the next section for exactly how and why. Either
way, MTP is a net win on this target/hardware once fixed; it was a net
loss before these fixes (checkpoint-dependent 0% acceptance made the
draft pass pure overhead).
Speed scales with MoE-cache warmth; acceptance does not
Every single-shot benchmark in this session used a freshly-restarted
server, which measures the MoE expert cache cold. Once the cache
warms up under sustained load, the same target+draft pair is
substantially faster — this is a real, measured effect, not a code
change, and worth knowing before trusting any single-run number.
Setup:
UD-Q4_K_XLtarget +dzannottiMTP head via-md,-c 8192,--spec-draft-n-max 3 --spec-draft-p-min 0.75,--moe-cache auto,temp 0(deterministic output),GGML_CUDA_MOE_CACHE_STATS=20forlive hit-rate telemetry.
Cache fill, same server, three consecutive requests:
used=6998/20998(33% full),evictions=0used=20998/20998(100% full),evictions=10442Run 2 vs run 1: +18.7% tok/s from cache warmth alone, identical
prompt, identical everything else, output byte-identical (temp 0).
Run 3 shows the warmth is partly topic-specific: a fresh topic on an
already-warm server still lands close to the cold-start number, not
the same-topic peak — some experts are apparently "generalist" and
stay hot across any topic (giving a baseline lift over a truly cold
cache), but a large share of run 2's peak came from that specific
topic's expert subset being fully resident from repetition.
Root cause, confirmed by reading
moe-cache.cu: the cache fills andevicts based on live usage with no pre-warming —
evictions=0in run1 means it hadn't even reached its capacity limit yet, so the
33%-full 25%-hit-rate snapshot at the end of run 1 is mid-fill, not
steady state. Run 2 shows what steady state (
used== full capacity,evictions actively happening) actually looks like: 83.4%, over 3x the
end-of-run-1 snapshot.
Draft acceptance did not track this. It stayed in the 77–84% band
across all three runs regardless of cache warmth (77.2% earlier in
this session on a cold
-c 8192server, 84.2%/83.3%/80.2% here) — thevariation looks prompt-dependent, not cache-state-dependent, which
makes sense: the MoE cache returns bit-identical expert weights
whether served from cache or a fresh read, so it changes latency,
not the draft head's predictions. Speed and acceptance are two
independent effects; only speed benefits from warmup.
Practical implication: don't judge this setup's real throughput
from a single fresh-restart benchmark — expect a meaningfully higher
number in a long-running session, especially on repeated or related
topics, plateauing once the cache fills and reaches its own steady
per-topic hit rate.
Known follow-ups (not part of this PR)
--spec-chain(chained in-graph multi-token MTP drafting) is notsupported for qwen4exp (
llama_model_supports_mtp_chain()onlyreturns true for
LLM_ARCH_QWEN35). Porting it means re-derivingqwen4exp's hyper-connection block (attention + MoE + hc-combine) in
the same low-level, per-row, in-graph, direct-KV-write form qwen35's
chain path uses — a real feature port, not a quick fix. The
reference implementation's own author also found deeper/chained
drafting didn't help this specific model (its draft head carries a
full MoE, so every drafted token is already a real forward pass).
MoE cache VRAM sizing uses a static reserve (
GGML_CUDA_MOE_CACHE_RESERVE_MB,default 3072 MiB) rather than measuring actual runtime peak usage.
Reducing it crashed with a genuine CUDA OOM on the first real decode
after a speculative-decoding config was added, meaning something in
the live MTP serving path (plausibly the cache's own async fill
queue) grows past the load-time worst-case graph reservation.
Root-caused to a real gap, not chased to a fix.
AI usage: yes