mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-03 11:35:37 +02:00
qwen4exp: support recurrent state rollback (#28123)
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.
This commit is contained in:
@@ -1100,6 +1100,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
|
||||
switch (arch) {
|
||||
case LLM_ARCH_QWEN35:
|
||||
case LLM_ARCH_QWEN35MOE:
|
||||
case LLM_ARCH_QWEN4EXP:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_NEMOTRON_H:
|
||||
case LLM_ARCH_NEMOTRON_H_MOE:
|
||||
|
||||
+18
-10
@@ -1087,20 +1087,28 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
|
||||
|
||||
ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);
|
||||
|
||||
// keep the last state_cols columns for the next ubatch
|
||||
// [TAG_RECURRENT_ROLLBACK_SPLITS] keep the last state_cols columns once per rollback slot,
|
||||
// slot s ending s tokens earlier so a rollback of s tokens reads a history that never saw them
|
||||
const size_t row_size = ggml_row_size(conv_states_all->type, row_total);
|
||||
const uint32_t mem_size = mctx_cur->get_size();
|
||||
|
||||
ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
|
||||
state_cols, channels, n_seqs,
|
||||
conv_input->nb[1], conv_input->nb[2],
|
||||
ggml_row_size(conv_input->type, conv_input->ne[0] - state_cols));
|
||||
const int64_t n_slots = (int64_t) cparams.n_rs_seq + 1;
|
||||
|
||||
ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
|
||||
state_cols * channels, n_seqs,
|
||||
conv_states_all->nb[1],
|
||||
kv_head * row_size);
|
||||
for (int64_t slot = 0; slot < n_slots; ++slot) {
|
||||
const int64_t s_idx = std::max<int64_t>(0, conv_input->ne[0] - state_cols - slot);
|
||||
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));
|
||||
ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
|
||||
state_cols, channels, n_seqs,
|
||||
conv_input->nb[1], conv_input->nb[2],
|
||||
ggml_row_size(conv_input->type, s_idx));
|
||||
|
||||
ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
|
||||
state_cols * channels, n_seqs,
|
||||
conv_states_all->nb[1],
|
||||
(slot * mem_size + kv_head) * row_size);
|
||||
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));
|
||||
}
|
||||
|
||||
return conv_input;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user