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:
Pascal
2026-09-01 06:24:49 +02:00
committed by GitHub
parent 09412af38a
commit 0eadefebd3
2 changed files with 19 additions and 10 deletions
+1
View File
@@ -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
View File
@@ -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;
}