qwen4exp: fix speculative decoding and add native MTP (NextN) support

Speculative decoding on Qwen3.8-Flash-Next previously ran ~6x SLOWER than
plain decode. This branch fixes seven issues: the recurrent rollback ring was
not enabled for this arch (every spec step serialized the full recurrent state
through host RAM), the ring was granted only to MTP/EAGLE3-style types, the
qwen4exp conv-state write ignored the ring banks, the PLE n-gram history could
not rewind losslessly, the rejection path never restored a draft context that
cannot roll back, prompt-cache reuse desynced such drafts, and
--spec-draft-model with draft-mtp loaded the main model instead of the given
sidecar file.

It also implements the model's native MTP head end to end: an --mtp converter
export (mtp.* tensors -> an mtp- sidecar GGUF) and the draft-mtp graph. The
combiner keeps the hyper-connection residual streams distinct, which is what
the head was trained on (draft acceptance 0.87 vs 0.47 when collapsed).

Measured on Strix Halo (Ryzen AI Max+ 395, Vulkan/RADV), UD-Q3_K_XL target,
greedy: decode 23.6 -> 49.1 t/s (2.08x) at shallow context and 14.0 -> 31.5
t/s (2.24x) at 7.3k tokens, acceptance 0.785 at --spec-draft-n-max 6.
Correctness verified with a greedy identity oracle (spec output equals plain
decode, modulo replay FP-reduction-order noise).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: apepojken

common/common.cpp

@@ -1647,7 +1647,14 @@ void common_memory::init(llama_context * ctx_tgt, llama_context * ctx_dft) {
 void common_memory::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) const {
     common_context_seq_rm(ctx_tgt, seq_id, p0, p1);
     if (ctx_dft) {
-        common_context_seq_rm(ctx_dft, seq_id, p0, p1);
+        // mixed-capability pair: the target may support partial removal (e.g. via the
+        // recurrent rollback ring) while a recurrent draft model does not. the draft cache
+        // is a pure optimization, so instead of aborting, drop the whole draft sequence and
+        // let the driver re-prime it from the prompt on the next draft call.
+        auto * mem = llama_get_memory(ctx_dft);
+        if (!llama_memory_seq_rm(mem, seq_id, p0, p1)) {
+            llama_memory_seq_rm(mem, seq_id, -1, -1);
+        }
     }
 }
 

common/common.h

@@ -384,11 +384,20 @@ struct common_params_speculative {
     }
 
     uint32_t need_n_rs_seq() const {
+        // every speculative type rolls back the target's rejected draft suffix, so any of
+        // them benefits from the recurrent-state snapshot ring on rollback-capable archs;
+        // without it, recurrent/hybrid models fall back to full per-step state
+        // serialization through host memory (SEQ_RM_TYPE_FULL), which is catastrophically
+        // slow (measured ~750 ms/step on qwen4exp)
         bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
-            return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
+            return t != COMMON_SPECULATIVE_TYPE_NONE;
         });
 
-        return needs_rs_seq ? draft.n_max : 0u;
+        // n_max + 1: a verify batch holds the previously sampled token plus up to n_max
+        // drafts, and the checkpoint+replay path (used when the draft context cannot roll
+        // back) rewinds the target across the whole batch, one deeper than the rejected
+        // draft suffix alone
+        return needs_rs_seq ? draft.n_max + 1 : 0u;
     }
 };
 

common/speculative.cpp

@@ -2391,8 +2391,14 @@ common_speculative_init_result::common_speculative_init_result(
     // the draft context holds as many tokens per sequence as the target context
     cparams.n_ctx = llama_n_ctx(ctx_tgt);
 
-    // note: for small models maybe we can set this to the maximum possible draft from all speculative types
-    //       the extra memory for small models is likely negligible?
+    // note: do NOT give the draft context a rollback ring (n_rs_seq > 0). the draft decodes
+    // one token per ubatch while drafting, and the delta-net snapshot store clamps banks
+    // older than the current ubatch to the batch-start state, so a rollback deeper than one
+    // token restores a state that is too new. the draft then drafts from a desynced prefix:
+    // measured on a qwen35 4B target + 2B draft pair, acceptance fell to 0.343 with the ring
+    // vs 0.84 with the (slower, serialize-per-step) FULL path, while outputs stayed exact
+    // because verification is target-authoritative. the target is unaffected: its rollbacks
+    // always land inside its own n_draft+1-token verify batch, whose banks are all fresh.
     cparams.n_rs_seq  = 0;
     cparams.ctx_other = ctx_tgt;
 
@@ -2401,7 +2407,10 @@ common_speculative_init_result::common_speculative_init_result(
         model_path = params.speculative.draft.mparams.path;
         LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str());
 
-        llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams);
+        // load the draft path that was just logged - loading params.model.path here loaded
+        // the main model a second time, which only worked by accident for models whose
+        // nextn tensors live inside the main GGUF, and never for mtp- sidecar files
+        llama_model * model_dft = llama_model_load_from_file(model_path.c_str(), mparams);
         if (model_dft == NULL) {
             LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str());
             return;

conversion/qwen4exp.py

@@ -26,9 +26,9 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
 
     model_arch = gguf.MODEL_ARCH.QWEN4EXP
 
-    # the MTP block is a separate draft head; vLLM drops it too
-    supports_mtp_export = False
-    no_mtp = True
+    # the MTP block loads through the shared _QwenMtpMixin remap; the qwen4exp-specific
+    # glue (fc_embedding/fc_hidden and the head's hyper-connection mixer) is rewritten in
+    # filter_tensors below
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -40,6 +40,25 @@ def __init__(self, *args, **kwargs):
         self._ple_map: np.memmap | None = None
         self._ple_path: Path | None = None
 
+    @classmethod
+    def filter_tensors(cls, item):
+        name = item[0]
+        if name.startswith("model.mtp."):
+            name = name.replace("model.", "", 1)
+            item = (name, item[1])
+        if name.startswith("mtp.") and not cls.no_mtp:
+            obc = cls._original_block_count
+            parts = name.split(".")
+            # separate embedding/hidden projections in place of qwen35's fused eh_proj
+            if len(parts) == 3 and parts[1] == "fc_embedding":
+                return f"model.layers.{obc}.fc_embd_mtp.weight", item[1]
+            if len(parts) == 3 and parts[1] == "fc_hidden":
+                return f"model.layers.{obc}.fc_hidden_mtp.weight", item[1]
+            # the head's own hyper-connection mixer in place of shared_head.norm
+            if len(parts) == 4 and parts[1] == "hyper_connection_mixer":
+                return f"model.layers.{obc}.hc_mixer_mtp.{parts[2]}.{parts[3]}", item[1]
+        return super().filter_tensors(item)
+
     def _read_hash_constants(self, suffix: str) -> list[int]:
         """Read an int64 PLE constant straight from the checkpoint.
 
@@ -68,14 +87,17 @@ def set_gguf_parameters(self):
         self.gguf_writer.add_indexer_top_k(hp["indexer_budget"])
         ratio = hp["indexer_compress_ratio"]
         layer_types = hp["layer_types"]
-        self.gguf_writer.add_attention_compress_ratios(
-            [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
-        )
+        # the MTP block(s) beyond the trunk run dense: pad the per-layer array to block_count
+        ratios = [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
+        ratios += [0] * (self.block_count - n_layer)
+        self.gguf_writer.add_attention_compress_ratios(ratios)
 
         # ple_layer_ids is 1-based in the HF config; empty means no n-gram table,
         # so emit no PLE keys rather than optional ones
         ple_layers = [i - 1 for i in hp["ple_layer_ids"]]
-        if not ple_layers:
+        # an mtp- sidecar ships no n-gram table and its filter drops the PLE constants,
+        # so emit no PLE keys at all: the loader then treats the file as PLE-free
+        if not ple_layers or self.mtp_only:
             return
         self.gguf_writer.add_ple_layers(ple_layers)
         self.gguf_writer.add_ple_ngram_size(hp["ngram_size"])

gguf-py/gguf/constants.py

@@ -1170,6 +1170,11 @@ class MODEL_TENSOR(IntEnum):
     NEXTN_HNORM            = auto()
     NEXTN_SHARED_HEAD_HEAD = auto()
     NEXTN_SHARED_HEAD_NORM = auto()
+    NEXTN_FC_EMBD          = auto()
+    NEXTN_FC_HIDDEN        = auto()
+    NEXTN_HC_NORM          = auto()
+    NEXTN_HC_DOWN          = auto()
+    NEXTN_HC_UP            = auto()
     # eagle3
     FC                     = auto()  # feature fusion layer
     D2T                    = auto()  # draft to target vocabulary mapping
@@ -1940,6 +1945,11 @@ class MODEL_TENSOR(IntEnum):
     MODEL_TENSOR.NEXTN_HNORM:               "blk.{bid}.nextn.hnorm",
     MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD:    "blk.{bid}.nextn.shared_head_head",
     MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM:    "blk.{bid}.nextn.shared_head_norm",
+    MODEL_TENSOR.NEXTN_FC_EMBD:             "blk.{bid}.nextn.fc_embd",
+    MODEL_TENSOR.NEXTN_FC_HIDDEN:           "blk.{bid}.nextn.fc_hidden",
+    MODEL_TENSOR.NEXTN_HC_NORM:             "blk.{bid}.nextn.hc_norm",
+    MODEL_TENSOR.NEXTN_HC_DOWN:             "blk.{bid}.nextn.hc_down",
+    MODEL_TENSOR.NEXTN_HC_UP:               "blk.{bid}.nextn.hc_up",
     MODEL_TENSOR.FC:                        "fc",
     MODEL_TENSOR.DSPARK_MARKOV_W1:          "markov_w1",
     MODEL_TENSOR.DSPARK_MARKOV_W2:          "markov_w2",
@@ -2895,6 +2905,13 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.PLE_NORM_QUERY,
         MODEL_TENSOR.PLE_NORM_CONV,
         MODEL_TENSOR.PLE_CONV1D,
+        MODEL_TENSOR.NEXTN_FC_EMBD,
+        MODEL_TENSOR.NEXTN_FC_HIDDEN,
+        MODEL_TENSOR.NEXTN_HC_NORM,
+        MODEL_TENSOR.NEXTN_HC_DOWN,
+        MODEL_TENSOR.NEXTN_HC_UP,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
     ],
     MODEL_ARCH.PLAMO: [
         MODEL_TENSOR.TOKEN_EMBD,

gguf-py/gguf/tensor_mapping.py

@@ -2661,6 +2661,26 @@ class TensorNameMap:
             "model.layers.{bid}.hnorm",
         ),
 
+        MODEL_TENSOR.NEXTN_FC_EMBD: (
+            "model.layers.{bid}.fc_embd_mtp",
+        ),
+
+        MODEL_TENSOR.NEXTN_FC_HIDDEN: (
+            "model.layers.{bid}.fc_hidden_mtp",
+        ),
+
+        MODEL_TENSOR.NEXTN_HC_NORM: (
+            "model.layers.{bid}.hc_mixer_mtp.hc_norm",
+        ),
+
+        MODEL_TENSOR.NEXTN_HC_DOWN: (
+            "model.layers.{bid}.hc_mixer_mtp.input_mix_weight_down",
+        ),
+
+        MODEL_TENSOR.NEXTN_HC_UP: (
+            "model.layers.{bid}.hc_mixer_mtp.input_mix_weight_up",
+        ),
+
         MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: (
             "model.layers.{bid}.shared_head.head",
         ),

src/llama-arch.cpp

@@ -568,6 +568,11 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
     { LLM_TENSOR_NEXTN_HNORM,                            "blk.%d.nextn.hnorm" },
     { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,                 "blk.%d.nextn.shared_head_head" },
     { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,                 "blk.%d.nextn.shared_head_norm" },
+    { LLM_TENSOR_NEXTN_FC_EMBD,                          "blk.%d.nextn.fc_embd" },
+    { LLM_TENSOR_NEXTN_FC_HIDDEN,                        "blk.%d.nextn.fc_hidden" },
+    { LLM_TENSOR_NEXTN_HC_NORM,                          "blk.%d.nextn.hc_norm" },
+    { LLM_TENSOR_NEXTN_HC_DOWN,                          "blk.%d.nextn.hc_down" },
+    { LLM_TENSOR_NEXTN_HC_UP,                            "blk.%d.nextn.hc_up" },
     { LLM_TENSOR_ATTN_SUB_NORM,                          "blk.%d.attn_sub_norm" },
     { LLM_TENSOR_FFN_SUB_NORM,                           "blk.%d.ffn_sub_norm" },
     { LLM_TENSOR_DEC_OUTPUT_NORM,                        "dec.output_norm" },
@@ -949,6 +954,11 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_NEXTN_HNORM,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+    {LLM_TENSOR_NEXTN_FC_EMBD,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_NEXTN_FC_HIDDEN,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_NEXTN_HC_NORM,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+    {LLM_TENSOR_NEXTN_HC_DOWN,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_NEXTN_HC_UP,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     // Nemotron 3 Super
     // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU
     {LLM_TENSOR_FFN_LATENT_DOWN,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
@@ -1080,6 +1090,10 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
     switch (arch) {
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
+        // qwen4exp shares the delta-net base + build_rs path with qwen35, and its extra
+        // caches already mirror rollback: the indexer cache is addressed by the attention
+        // cells and the PLE history/conv state go through seq_rm/build_rs like the SSM state.
+        case LLM_ARCH_QWEN4EXP:
         case LLM_ARCH_DEEPSEEK4:
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:

src/llama-arch.h

@@ -681,6 +681,13 @@ enum llm_tensor {
     LLM_TENSOR_NEXTN_HNORM,
     LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,
     LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,
+    // qwen4exp MTP glue: separate embedding/hidden projections instead of eh_proj,
+    // and the head's own hyper-connection mixer in place of shared_head_norm
+    LLM_TENSOR_NEXTN_FC_EMBD,
+    LLM_TENSOR_NEXTN_FC_HIDDEN,
+    LLM_TENSOR_NEXTN_HC_NORM,
+    LLM_TENSOR_NEXTN_HC_DOWN,
+    LLM_TENSOR_NEXTN_HC_UP,
     LLM_TENSOR_MASKED_EMBD_CENTROIDS,
     LLM_TENSOR_MASKED_EMBD_ORDERING,
     LLM_TENSOR_FC,

src/llama-memory-hybrid-idx.cpp

@@ -56,7 +56,13 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
             model, hparams_idx, type_k, type_v, v_trans, offload, unified,
             kv_size, n_seq_max, n_pad, n_swa, swa_type,
             nullptr, filter_idx, nullptr, nullptr, "idx_");
-    }()) {}
+    }()) {
+    // keep enough PLE history that a rollback of up to n_rs_seq tokens still leaves the
+    // full (ple_ngram_size - 1) hash window of true predecessors; see ple_hist_keep
+    const uint32_t n_win = model.hparams.ple_ngram_size > 0 ? model.hparams.ple_ngram_size - 1 : 0;
+
+    ple_hist_keep_n = n_win + n_rs_seq;
+}
 
 llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
     // note: this repeats llama_memory_hybrid::init_batch because the indexer cache needs the
@@ -437,9 +443,11 @@ void llama_memory_hybrid_idx::ple_hist_state_read(llama_io_read_i & io, llama_se
         io.read(&next_pos, sizeof(next_pos));
         io.read(&n_toks,   sizeof(n_toks));
 
-        // the window is never longer than ple_ngram_size - 1; anything else is a corrupt or
-        // mismatched blob, and reading it would size an allocation from the file
-        if (n_toks > LLAMA_MAX_PLE_NGRAM - 1) {
+        // the history holds the hash window plus up to ple_hist_keep_n rollback slack;
+        // anything far beyond that is a corrupt or mismatched blob, and reading it would
+        // size an allocation from the file. allow generous slack for blobs written with a
+        // larger ring than the reader's.
+        if (n_toks > (uint32_t) LLAMA_MAX_PLE_NGRAM - 1 + std::max<uint32_t>(ple_hist_keep_n, 64)) {
             throw std::runtime_error("qwen4exp PLE history: implausible token count in state blob");
         }
 
@@ -573,6 +581,12 @@ llama_memory_hybrid_idx::ple_history & llama_memory_hybrid_idx_context::get_ple_
     return mem->ple_hist_get(seq_id);
 }
 
+uint32_t llama_memory_hybrid_idx_context::get_ple_hist_keep() const {
+    GGML_ASSERT(mem != nullptr);
+
+    return mem->ple_hist_keep_n;
+}
+
 void llama_memory_hybrid_idx_context::set_input_qsa(
         ggml_tensor * cell_blk,
         ggml_tensor * blk_cells,

src/llama-memory-hybrid-idx.h

@@ -95,6 +95,12 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid {
     // const because set_input updates it through a const memory context
     ple_history & ple_hist_get(llama_seq_id seq_id) const;
 
+    // how many history tokens to keep per sequence: (ple_ngram_size - 1) for the hash
+    // window, plus n_rs_seq so a ring rollback can truncate without losing the true
+    // predecessors (a shorter history EOS-pads the window after a rewind, which corrupts
+    // the n-gram rows for the first tokens decoded after every speculative rollback)
+    uint32_t ple_hist_keep_n = 0;
+
 private:
     // the indexer cache holds one key head per layer, so it needs its own hparams:
     // llama_kv_cache keeps a reference to what it is given
@@ -161,6 +167,8 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
     // [TAG_PLE_HISTORY] the per-sequence n-gram history of the owning memory, for set_input
     llama_memory_hybrid_idx::ple_history & get_ple_hist(llama_seq_id seq_id) const;
 
+    uint32_t get_ple_hist_keep() const;
+
     // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache.
     // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout:
     //   cell_blk  I32 [n_kv, ns]           block each cell belongs to