qwen4exp: support recurrent state rollback - #28123
Merged
ServeurpersoCom merged 1 commit intoSep 1, 2026
Merged
Conversation
MTP speculative decoding needs the target state to move back by the number of rejected draft tokens. Without rollback support the context is classified as SEQ_RM_TYPE_FULL and the server serializes the whole recurrent state to host memory on every round, which costs more than the drafting saves. The recurrent cache already holds n_rs_seq + 1 snapshot planes and the delta net writes its SSM state into them, but build_conv_state_at wrote a single plane, so a rollback restored a convolution history that was never captured. It now writes one snapshot per slot, each ending one token earlier, for the delta net QKV convolution and for the PLE convolution alike. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL with the standalone MTP draft, n-max 3 and a single slot: decoding reaches 183 tok/s on code and 144 tok/s on prose. The same branch before this change, where the server falls back to checkpointing the state to host memory, reaches 123 and 83 tok/s, for 108 tok/s without a draft.
CISC
approved these changes
Aug 31, 2026
am17an
approved these changes
Sep 1, 2026
This was referenced Sep 1, 2026
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
danielhanchen
added a commit
to unslothai/llama.cpp
that referenced
this pull request
Sep 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Qwen3.8-Flash-Next currently cannot roll back its recurrent state, so with an MTP draft the server falls back to serializing the whole state to host memory on every speculative round, which costs more than the drafting saves.
Adding the arch to llm_arch_supports_rs_rollback() is not enough on its own: qwen4exp overrides the shared convolution state write with its own build_conv_state_at(), which only ever wrote the current plane, so a rollback would restore a correct SSM state next to a convolution history that was never captured. It now writes one snapshot per rollback slot, like llm_build_delta_net_base::build_conv_state already does, covering both the delta net QKV convolution and the PLE one.
Tested on an RTX PRO 6000 with the MTP head and the draft patch from https://huggingface.co/dzannotti/Qwen3.8-Flash-Next-MTP-GGUF, n-max 3, one slot:
no draft: 108 tok/s
before: 123 tok/s code, 83 tok/s prose
after: 183 tok/s code, 144 tok/s prose
Note the 83: before this change MTP was slower than not drafting at all.
Additional information
Alternative to #28118
Requirements