diff --git a/common/speculative.cpp b/common/speculative.cpp index 4d015ed0..317f165e 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -147,7 +147,7 @@ static bool common_speculative_are_compatible( static bool common_speculative_target_has_appended_mtp_contract(const llama_model * model) { return llama_model_is_step35(model) || llama_model_is_deepseek4(model) || - llama_model_is_qwen35_family(model); + llama_model_is_qwen35_family(model) || llama_model_is_qwen4exp(model); } static bool common_speculative_has_recognized_mtp_companion( @@ -160,8 +160,13 @@ static bool common_speculative_has_recognized_mtp_companion( return llama_model_arch_string(target) != nullptr && std::strcmp(llama_model_arch_string(target), "gemma4") == 0; } - return common_speculative_target_has_appended_mtp_contract(target) && - llama_model_mtp_package(companion) == LLAMA_MTP_PACKAGE_COMPANION; + if (llama_model_mtp_package(companion) != LLAMA_MTP_PACKAGE_COMPANION) { + return false; + } + if (llama_model_is_qwen4exp(target) || llama_model_is_qwen4exp(companion)) { + return llama_model_is_qwen4exp(target) && llama_model_is_qwen4exp(companion); + } + return common_speculative_target_has_appended_mtp_contract(target); } // state of an implementation of speculative decoding @@ -2221,6 +2226,18 @@ bool common_speculative_finalize_startup( __func__, n_heads); return false; } + } else if (llama_model_is_qwen4exp(model)) { + if (!llama_model_is_qwen4exp(companion)) { + LOG_ERR("%s: Qwen4Exp MTP requires a Qwen4Exp companion\n", __func__); + return false; + } + + const int32_t n_heads = llama_model_n_nextn_layer(companion); + if (n_heads != 1) { + LOG_ERR("%s: Qwen4Exp MTP companion requires exactly one predictor layer, got %d\n", + __func__, n_heads); + return false; + } } if (common_speculative_has_recognized_mtp_companion(model, companion)) { diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index d16efa06..c53e5d15 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -218,8 +218,9 @@ bool server_context::load_model(const gpt_params& params_) { common_speculative_prepare_startup(params_base, false); if (server_speculative_requires_single_slot(params_base.speculative) && params_base.n_parallel > 1) { - LOG_ERROR("Speculative decoding is currently limited to a single server slot (-np 1).\n", { + LOG_ERROR("Speculative decoding is currently limited to a single server slot: restart with -np 1, or drop the --spec-type stage(s) to keep parallel slots.\n", { {"n_parallel", params_base.n_parallel}, + {"stage_chain", common_speculative_stage_chain_to_str(params_base.speculative)}, }); return false; } diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index 1d8b165e..be9343d2 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -3827,6 +3827,12 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_UNARY_OP_GELU_ERF: ggml_cuda_op_gelu_erf(ctx, dst); break; + case GGML_UNARY_OP_ABS: + ggml_cuda_op_abs(ctx, dst); + break; + case GGML_UNARY_OP_SGN: + ggml_cuda_op_sgn(ctx, dst); + break; case GGML_UNARY_OP_SILU: ggml_cuda_op_silu(ctx, dst); break; @@ -4802,6 +4808,8 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons switch (ggml_get_unary_op(op)) { case GGML_UNARY_OP_GELU: case GGML_UNARY_OP_GELU_ERF: + case GGML_UNARY_OP_ABS: + case GGML_UNARY_OP_SGN: case GGML_UNARY_OP_SILU: case GGML_UNARY_OP_SWIGLU: case GGML_UNARY_OP_SWIGLU_OAI: diff --git a/include/llama.h b/include/llama.h index 160e9643..da39924b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -716,6 +716,8 @@ extern "C" { LLAMA_API bool llama_model_is_qwen35_family(const struct llama_model * model); + LLAMA_API bool llama_model_is_qwen4exp(const struct llama_model * model); + LLAMA_API bool llama_is_gemma4_mtp_file(const char * path); LLAMA_API bool llama_model_is_split_mode_graph(const struct llama_model * model); diff --git a/src/graphs/build_qwen4exp.cpp b/src/graphs/build_qwen4exp.cpp index 134723bf..99a38134 100644 --- a/src/graphs/build_qwen4exp.cpp +++ b/src/graphs/build_qwen4exp.cpp @@ -3,6 +3,8 @@ #include "../llama-context.h" #include "../llama-delta-net.h" +#include + // the [hc_dim] gamma is wider than the per-stream reduction, so ggml_fused_rms_norm cannot // express this and the two ops stay separate static ggml_tensor * qwen4exp_grouped_rms( @@ -428,36 +430,78 @@ ggml_cgraph * llm_build_context::build_qwen4exp() { ggml_cgraph * gf = new_graph_custom(); - delta_net delta(lctx, batch); + const bool is_mtp = lctx.cparams.mtp_op_type != MTP_OP_NONE; + + // the MTP pass walks only the QSA tail: no recurrent state, so the draft + // context has zero qnext state slots and the delta-net ctor must not run + std::optional delta_opt; + if (!is_mtp) { + delta_opt.emplace(lctx, batch); + } const int32_t n_embd_head = hparams.n_embd_head_v(0); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); const int32_t hc = hparams.dsv4_hc_mult; - ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); - ggml_tensor * inp_pos = build_inp_pos(); - ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; - ggml_tensor * KQ_mask = build_inp_KQ_mask(); + const int n_layer_begin = is_mtp ? n_layer - hparams.nextn_predict_layers : 0; + const int n_layer_end = is_mtp ? n_layer : n_layer - hparams.nextn_predict_layers; - lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); - cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); - ggml_set_input(lctx.inp_s_seq_qnext); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = (is_mtp || n_tokens > 1) ? build_inp_out_ids() : nullptr; + ggml_tensor * KQ_mask = build_inp_KQ_mask(); float KQ_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - // the wide residual starts as hc identical copies of the embedding - ggml_tensor * res_hc = ggml_repeat_4d(ctx0, - ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), - n_embd, hc, n_tokens, 1); - cb(res_hc, "hc_residual", -1); + ggml_tensor * res_hc = nullptr; + if (is_mtp) { + // the fill sites assert on these buffers; the MTP graph consumes neither + lctx.inp_s_seq_qnext = nullptr; + lctx.inp_ple_rows = nullptr; - if (hparams.ple_n_heads > 0) { - lctx.inp_ple_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, hparams.ple_n_heads * n_tokens); - cb(lctx.inp_ple_rows, "inp_ple_rows", -1); - ggml_set_input(lctx.inp_ple_rows); + const int32_t hc_dim = hc * n_embd; + ggml_tensor * hidden = build_inp_mtp_states(hc_dim); // target's pre-mixer wide stream + ggml_tensor * tok = build_inp_embd_mtp(model.tok_embd); // [n_embd, n_tokens] + const auto & nextn = model.layers[n_layer_begin].nextn; + + // normalize the embedding (n_embd) and the wide hidden stream (hc*n_embd) + ggml_tensor * e = ggml_rms_norm(ctx0, tok, hparams.f_norm_rms_eps); + e = ggml_mul(ctx0, e, nextn.enorm); + + ggml_tensor * h = ggml_rms_norm(ctx0, hidden, hparams.f_norm_rms_eps); + h = ggml_mul(ctx0, h, nextn.hnorm); + + // eh_proj is [fc_embedding | fc_hidden] along its input dimension. + // Concatenating e_norm and h_norm per stream therefore preserves the old + // fc_embedding@e_norm + fc_hidden@h_norm entry result in one matmul. + e = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, e, n_embd, 1, n_tokens), n_embd, hc, n_tokens, 1); + h = ggml_reshape_3d(ctx0, h, n_embd, hc, n_tokens); + ggml_tensor * eh = ggml_concat(ctx0, e, h, 0); + cb(eh, "mtp_eh_concat", n_layer_begin); + + res_hc = llm_build_lora_mm(lctx, ctx0, nextn.eh_proj, eh); + cb(res_hc, "mtp_eh_proj", n_layer_begin); } else { - lctx.inp_ple_rows = nullptr; + ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + + lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); + cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); + ggml_set_input(lctx.inp_s_seq_qnext); + + // the wide residual starts as hc identical copies of the embedding + res_hc = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), + n_embd, hc, n_tokens, 1); + cb(res_hc, "hc_residual", -1); + + if (hparams.ple_n_heads > 0) { + lctx.inp_ple_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, hparams.ple_n_heads * n_tokens); + cb(lctx.inp_ple_rows, "inp_ple_rows", -1); + ggml_set_input(lctx.inp_ple_rows); + } else { + lctx.inp_ple_rows = nullptr; + } } // the same test the delta-net path uses for its own state @@ -467,11 +511,11 @@ ggml_cgraph * llm_build_context::build_qwen4exp() { ple_reset_pos[i] = batch.pos[i] == 0; } - for (int il = 0; il < n_layer; ++il) { + for (int il = n_layer_begin; il < n_layer_end; ++il) { ggml_tensor * inject = nullptr; if (hparams.is_ple(il)) { - res_hc = qwen4exp_ple(*this, ctx0, gf, lctx, model, hparams, delta, res_hc, lctx.inp_ple_rows, + res_hc = qwen4exp_ple(*this, ctx0, gf, lctx, model, hparams, *delta_opt, res_hc, lctx.inp_ple_rows, kv_self.s_l[il], ple_reset_state, ple_reset_pos, n_embd, n_tokens, il, cb); } @@ -481,7 +525,7 @@ ggml_cgraph * llm_build_context::build_qwen4exp() { &inject, n_embd, il, cb); if (hparams.is_recurrent(il)) { - cur = delta.build_layer_attn_linear(ctx0, gf, cur, nullptr, il, cb, /* external_residual */ true, + cur = delta_opt->build_layer_attn_linear(ctx0, gf, cur, nullptr, il, cb, /* external_residual */ true, GGML_UNARY_OP_SIGMOID); } else { // the indexer reads the same block input as q/k/v, and returns the causal mask @@ -522,6 +566,36 @@ ggml_cgraph * llm_build_context::build_qwen4exp() { cb(res_hc, "l_out", il); } + if (is_mtp) { + // wide stream out: the next draft step's hidden input (scheme A) + ggml_tensor * flat = ggml_reshape_2d(ctx0, res_hc, hc * n_embd, n_tokens); + ggml_tensor * h_next = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat; + cb(h_next, "result_mtp_embd", -1); + ggml_set_output(h_next); + ggml_build_forward_expand(gf, h_next); + + // exit mixer is hc_head-shaped: per-stream norm + low-rank gate, no inject + const auto & nextn = model.layers[n_layer_begin].nextn; + ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, + ggml_reshape_3d(ctx0, h_next, n_embd, hc, h_next->ne[1]), + nextn.hc_head_norm, nextn.hc_head_down, nextn.hc_head_up, + nullptr, nullptr, n_embd, -1, cb); + + cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + cb(cur, "result_output", -1); + ggml_build_forward_expand(gf, cur); + return gf; + } + + if (lctx.cparams.mtp && hparams.n_embd_out > hparams.n_embd) { + // hand the draft the pre-mixer wide stream for its first step + ggml_tensor * flat = ggml_reshape_2d(ctx0, res_hc, hc * n_embd, n_tokens); + ggml_tensor * h_nextn = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat; + cb(h_nextn, "result_mtp_embd", -1); + ggml_set_output(h_nextn); + ggml_build_forward_expand(gf, h_nextn); + } + ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, res_hc, model.hc_head_norm, model.hc_head_down, model.hc_head_up, nullptr, nullptr, n_embd, -1, cb); @@ -530,6 +604,9 @@ ggml_cgraph * llm_build_context::build_qwen4exp() { cur = ggml_get_rows(ctx0, cur, inp_out_ids); } + // name the mixed stream for append_pooling (the nextn tail needs a named embd) + cb(cur, "result_embd", -1); + cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); cb(cur, "result_output", -1); diff --git a/src/llama-arch.h b/src/llama-arch.h index b210bd60..513049b7 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -409,6 +409,9 @@ enum llm_tensor { LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, // 95 + LLM_TENSOR_NEXTN_HC_HEAD_NORM, + LLM_TENSOR_NEXTN_HC_HEAD_DOWN, + LLM_TENSOR_NEXTN_HC_HEAD_UP, LLM_TENSOR_INDEXER_K_NORM, LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 7ecaf109..14b0047c 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -629,20 +629,45 @@ void llm_load_hparams( ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); + // the MTP handover is the wide pre-mixer residual, so the width must not depend on an appended NextN block + if (hparams.dsv4_hc_mult > 1) { + const uint32_t wide = hparams.n_embd * hparams.dsv4_hc_mult; + if (hparams.n_embd_out != hparams.n_embd && hparams.n_embd_out != wide) { + throw std::runtime_error("qwen4exp: embedding_length_out must equal n_embd * hyper-connection count"); + } + hparams.n_embd_out = wide; + } + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); + if (hparams.nextn_predict_layers >= hparams.n_layer) { + throw std::runtime_error("qwen4exp: nextn_predict_layers must be smaller than block_count"); + } + { uint32_t full_attn_interval = 4; ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); for (uint32_t i = 0; i < hparams.n_layer; ++i) { hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); } + // the MTP tail is a full-attention (QSA) layer regardless of the interval pattern + for (uint32_t i = hparams.n_layer - hparams.nextn_predict_layers; i < hparams.n_layer; ++i) { + hparams.recurrent_layer_arr[i] = false; + } } { + // block_count includes the nextn tail, but the converter's ratio array covers + // only the main layers; the tail inherits the last main QSA layer's ratio + const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers; uint32_t n_ratios = 0; - if (ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_ratios, false) && n_ratios >= hparams.n_layer) { + if (ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_ratios, false) && n_ratios >= n_main) { std::vector ratios; ml.get_arr(ml.llm_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS), ratios); - std::copy_n(ratios.begin(), hparams.n_layer, hparams.dsv4_compress_ratios.begin()); + const uint32_t n_copy = std::min(n_ratios, hparams.n_layer); + std::copy_n(ratios.begin(), n_copy, hparams.dsv4_compress_ratios.begin()); + for (uint32_t i = n_copy; i < hparams.n_layer; ++i) { + hparams.dsv4_compress_ratios[i] = ratios[n_main - 1]; + } } } @@ -692,7 +717,7 @@ void llm_load_hparams( } } - switch (hparams.n_layer) { + switch (hparams.n_layer - hparams.nextn_predict_layers) { case 48: model.type = e_model::MODEL_125B_A6B; break; default: model.type = e_model::MODEL_UNKNOWN; } diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 45b346be..3a613e4f 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -1681,19 +1681,41 @@ bool create_tensors_helper::create_qwen4exp_tensors(const LLM_TN & tn) { const int32_t hc_dim = hc * n_embd; const int32_t hc_rank = hparams.hc_low_rank; + // a companion target passes the head via -md and legitimately declares zero NextN blocks + if (model.mtp && hparams.nextn_predict_layers > 1) { + throw std::runtime_error(format( + "qwen4exp: MTP supports at most one appended NextN block; nextn_predict_layers is %u", + hparams.nextn_predict_layers)); + } + model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); + // a predictor-only GGUF reports the full block count but ships only the NextN block + const bool mtp_only = hparams.nextn_predict_layers > 0 && + ml.get_tensor_meta(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).c_str()) == nullptr; + const int trunk_flags = mtp_only + ? llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED : 0; + // The wide residual is normalised and collapsed by a hyper-connection mix rather // than by an output_norm, so this architecture carries none. - model.hc_head_norm = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), {hc_dim}); - model.hc_head_down = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), {hc_dim, hc_rank}); - model.hc_head_up = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_UP, "weight"), {hc_rank, hc_dim}); + const int out_mixer_flags = mtp_only ? llama_model_loader::TENSOR_NOT_REQUIRED : 0; + model.hc_head_norm = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), {hc_dim}, out_mixer_flags); + model.hc_head_down = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), {hc_dim, hc_rank}, out_mixer_flags); + model.hc_head_up = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_UP, "weight"), {hc_rank, hc_dim}, out_mixer_flags); model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED); if (model.output == NULL) { model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); } - if (hparams.ple_n_heads > 0) { + bool ple_table_needed = hparams.ple_n_heads > 0; + if (ple_table_needed && mtp_only) { + // only the predictor tail loads from this file; the table is needed only if a tail block uses PLE + ple_table_needed = false; + for (uint32_t il = hparams.n_layer - hparams.nextn_predict_layers; il < hparams.n_layer; ++il) { + ple_table_needed = ple_table_needed || hparams.is_ple(il); + } + } + if (ple_table_needed) { // The row count comes from the table itself. The converter shards the n-gram // table, so the sum of ple_head_vocab_sizes is not the stored row count. const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"); @@ -1738,60 +1760,171 @@ bool create_tensors_helper::create_qwen4exp_tensors(const LLM_TN & tn) { auto & layer = model.layers[i]; - layer.hc_attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_NORM, "weight", i), {hc_dim}); - layer.hc_attn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", i), {hc_dim, hc_rank}); - layer.hc_attn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_UP, "weight", i), {hc_rank, hc_dim}); - layer.hc_attn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", i), {hc_dim, hc}); - layer.hc_ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_NORM, "weight", i), {hc_dim}); - layer.hc_ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_DOWN, "weight", i), {hc_dim, hc_rank}); - layer.hc_ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_UP, "weight", i), {hc_rank, hc_dim}); - layer.hc_ffn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_INJECT, "weight", i), {hc_dim, hc}); + // the nextn (MTP) tail loads only when speculative MTP is requested + const bool is_nextn = hparams.nextn_predict_layers > 0 && + (uint32_t) i >= hparams.n_layer - hparams.nextn_predict_layers; + int lf = is_nextn && !model.mtp ? llama_model_loader::TENSOR_SKIP : 0; + if (mtp_only && !is_nextn) { + lf |= trunk_flags; + } + + layer.hc_attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_NORM, "weight", i), {hc_dim}, lf); + layer.hc_attn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", i), {hc_dim, hc_rank}, lf); + layer.hc_attn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_UP, "weight", i), {hc_rank, hc_dim}, lf); + layer.hc_attn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", i), {hc_dim, hc}, lf); + layer.hc_ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_NORM, "weight", i), {hc_dim}, lf); + layer.hc_ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_DOWN, "weight", i), {hc_dim, hc_rank}, lf); + layer.hc_ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_UP, "weight", i), {hc_rank, hc_dim}, lf); + layer.hc_ffn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_INJECT, "weight", i), {hc_dim, hc}, lf); if (hparams.is_ple(i)) { - layer.ple_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_KEY, "weight", i), {n_embd, hc_dim}); - layer.ple_value = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_VALUE, "weight", i), {n_embd, n_embd}); - layer.ple_norm_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_KEY, "weight", i), {hc_dim}); - layer.ple_norm_query = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", i), {hc_dim}); - layer.ple_norm_conv = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_CONV, "weight", i), {hc_dim}); - layer.ple_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_CONV1D, "weight", i), {hparams.ple_conv_kernel, hc_dim}); + layer.ple_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_KEY, "weight", i), {n_embd, hc_dim}, lf); + layer.ple_value = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_VALUE, "weight", i), {n_embd, n_embd}, lf); + layer.ple_norm_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_KEY, "weight", i), {hc_dim}, lf); + layer.ple_norm_query = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", i), {hc_dim}, lf); + layer.ple_norm_conv = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_CONV, "weight", i), {hc_dim}, lf); + layer.ple_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_CONV1D, "weight", i), {hparams.ple_conv_kernel, hc_dim}, lf); } if (!hparams.is_recurrent(i)) { // wq carries the query and an equal-width gate - layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head * 2}); - layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}); - layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}); - layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}); + layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head * 2}, lf); + layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, lf); + layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, lf); + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, lf); - layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}); - layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}); + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, lf); + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, lf); // one indexer key is shared across the indexer heads - layer.indexer_q_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, idx_head * idx_n_head}); - layer.indexer_k_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, idx_head}); - layer.indexer_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {idx_head}); - layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_head}); + layer.indexer_q_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, idx_head * idx_n_head}, lf); + layer.indexer_k_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, idx_head}, lf); + layer.indexer_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {idx_head}, lf); + layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_head}, lf); } else { - layer.wqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, key_dim * 2 + value_dim}); - layer.wqkv_gate = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, value_dim}); - layer.ssm_conv1d = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {hparams.ssm_d_conv, conv_dim}); - layer.ssm_dt = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_DT, "bias", i), {hparams.ssm_dt_rank}); - layer.ssm_a = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_A_NOSCAN, i), {hparams.ssm_dt_rank}); - layer.ssm_beta = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, num_v_heads}); - layer.ssm_alpha = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), {n_embd, num_v_heads}); - layer.ssm_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_v_dim}); - layer.ssm_out = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_OUT, "weight", i), {value_dim, n_embd}); + layer.wqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, key_dim * 2 + value_dim}, lf); + layer.wqkv_gate = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, value_dim}, lf); + layer.ssm_conv1d = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {hparams.ssm_d_conv, conv_dim}, lf); + layer.ssm_dt = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_DT, "bias", i), {hparams.ssm_dt_rank}, lf); + layer.ssm_a = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_A_NOSCAN, i), {hparams.ssm_dt_rank}, lf); + layer.ssm_beta = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, num_v_heads}, lf); + layer.ssm_alpha = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), {n_embd, num_v_heads}, lf); + layer.ssm_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_v_dim}, lf); + layer.ssm_out = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_OUT, "weight", i), {value_dim, n_embd}, lf); } auto ffn_ctx = ctx_split; - layer.ffn_gate_inp = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}); - use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, 0, n_ff_exp); + layer.ffn_gate_inp = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, lf); + use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, lf, n_ff_exp); - layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), {n_embd}); - layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}); - layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}); - layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}); + layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), {n_embd}, lf); + layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, lf); + layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, lf); + layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, lf); + } + + if (hparams.nextn_predict_layers > 0) { + const int bid = n_layer - 1; + auto & nextn = model.layers[bid].nextn; + auto * head_ctx = ctx_for_layer_split(bid); + + // Head-level tensors are optional at file-load time. A normal load skips any + // that are present, while an MTP load resolves and validates one complete layout below. + int flags = llama_model_loader::TENSOR_NOT_REQUIRED; + if (!model.mtp) { + flags |= llama_model_loader::TENSOR_SKIP; + } + + const std::string eh_proj_name = tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", bid); + const std::string enorm_name = tn(LLM_TENSOR_NEXTN_ENORM, "weight", bid); + const std::string hnorm_name = tn(LLM_TENSOR_NEXTN_HNORM, "weight", bid); + const std::string shared_norm_name = tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", bid); + const std::string hc_norm_name = tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", bid); + const std::string hc_down_name = tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", bid); + const std::string hc_up_name = tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", bid); + + nextn.eh_proj = create_tensor(head_ctx, eh_proj_name, {2 * n_embd, n_embd}, flags); + nextn.enorm = create_tensor(head_ctx, enorm_name, {n_embd}, flags); + nextn.hnorm = create_tensor(head_ctx, hnorm_name, {hc_dim}, flags); + + // Flavor A: the upstream converter's head-local mixer. + nextn.hc_head_norm = create_tensor(head_ctx, hc_norm_name, {hc_dim}, flags); + nextn.hc_head_down = create_tensor(head_ctx, hc_down_name, {hc_dim, hc_rank}, flags); + nextn.hc_head_up = create_tensor(head_ctx, hc_up_name, {hc_rank, hc_dim}, flags); + + // Flavor B: community merge-script files retain only the head norm and + // use the trunk's already-loaded output_hc_down/output_hc_up pair. + nextn.shared_head_norm = create_tensor(head_ctx, shared_norm_name, {hc_dim}, flags); + + // Legacy files carry the head under fork-local mtp.* names: register them + // known-and-skipped so a non-MTP load balances tensor accounting. + { + const int legacy = llama_model_loader::TENSOR_NOT_REQUIRED | llama_model_loader::TENSOR_SKIP; + create_tensor(head_ctx, "mtp.fc_embd.weight", {n_embd, n_embd}, legacy); + create_tensor(head_ctx, "mtp.fc_hidden.weight", {n_embd, n_embd}, legacy); + create_tensor(head_ctx, "mtp.pre_norm_embd.weight", {n_embd}, legacy); + create_tensor(head_ctx, "mtp.pre_norm_hidden.weight", {hc_dim}, legacy); + create_tensor(head_ctx, "mtp.mixer_norm.weight", {hc_dim}, legacy); + create_tensor(head_ctx, "mtp.mixer_down.weight", {hc_dim, hc_rank}, legacy); + create_tensor(head_ctx, "mtp.mixer_up.weight", {hc_rank, hc_dim}, legacy); + } + + if (model.mtp) { + std::string missing; + auto require = [&](const ggml_tensor * tensor, const std::string & name) { + if (tensor == nullptr) { + if (!missing.empty()) { + missing += ", "; + } + missing += name; + } + }; + + require(nextn.eh_proj, eh_proj_name); + require(nextn.enorm, enorm_name); + require(nextn.hnorm, hnorm_name); + + // Any dedicated-mixer tensor selects flavor A. Do not hide a malformed + // flavor A by mixing it with flavor B. + const bool has_dedicated_mixer = nextn.hc_head_norm != nullptr || + nextn.hc_head_down != nullptr || nextn.hc_head_up != nullptr; + if (has_dedicated_mixer) { + require(nextn.hc_head_norm, hc_norm_name); + require(nextn.hc_head_down, hc_down_name); + require(nextn.hc_head_up, hc_up_name); + } else { + // Flavor B': a predictor-only file may omit shared_head_norm and + // carry the head norm only in output_hc_norm (Dzannotti layout). + // Safe only when mtp_only, where output_hc_* is the head's own mixer. + if (nextn.shared_head_norm == nullptr && mtp_only) { + require(model.hc_head_norm, "output_hc_norm.weight"); + } else { + require(nextn.shared_head_norm, shared_norm_name); + } + require(model.hc_head_down, "output_hc_down.weight"); + require(model.hc_head_up, "output_hc_up.weight"); + } + + if (!missing.empty()) { + throw std::runtime_error(format( + "qwen4exp: MTP requested but the NextN head is incomplete; missing tensor(s): %s. " + "Accepted layouts for block %d are flavor A: %s, %s, %s, %s, %s, %s; " + "or flavor B: %s, %s, %s, %s with trunk output_hc_down.weight and output_hc_up.weight", + missing.c_str(), bid, + eh_proj_name.c_str(), enorm_name.c_str(), hnorm_name.c_str(), + hc_norm_name.c_str(), hc_down_name.c_str(), hc_up_name.c_str(), + eh_proj_name.c_str(), enorm_name.c_str(), hnorm_name.c_str(), shared_norm_name.c_str())); + } + + if (!has_dedicated_mixer) { + nextn.hc_head_norm = nextn.shared_head_norm != nullptr + ? nextn.shared_head_norm : model.hc_head_norm; + nextn.hc_head_down = model.hc_head_down; + nextn.hc_head_up = model.hc_head_up; + LLAMA_LOG_WARN("qwen4exp: NextN head has no dedicated hc mixer; using the trunk output_hc mixer\n"); + } + } } return use_mmap_buffer; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 738b13de..95ca9e69 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -998,10 +998,15 @@ struct ggml_tensor * llama_model_loader::create_tensor(struct ggml_context * ctx return NULL; } - // skip unused tensors + // skip unused tensors (per-tensor detail at debug level; one summary at + // done_getting_tensors so ordinary loads are not flooded by an unused + // NextN/MTP head) if (flags & TENSOR_SKIP) { const size_t nbytes = ggml_nbytes(cur); - LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", name.c_str(), nbytes); + LLAMA_LOG_DEBUG("model has unused tensor %s (size = %zu bytes) -- ignoring\n", name.c_str(), nbytes); + + n_skipped++; + size_skipped += nbytes; size_data -= nbytes; n_created++; @@ -1042,6 +1047,10 @@ struct ggml_tensor * llama_model_loader::create_tensor_as_view(struct ggml_conte } void llama_model_loader::done_getting_tensors() const { + if (n_skipped > 0) { + LLAMA_LOG_WARN("%s: skipped %d unused tensors (%.1f MiB) -- e.g. a NextN/MTP head when MTP is not requested; per-tensor detail at debug log level\n", + __func__, n_skipped, size_skipped / 1024.0 / 1024.0); + } if (n_created != n_tensors) { throw std::runtime_error(format("%s: wrong number of tensors; expected %d, got %d", __func__, n_tensors, n_created)); } diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index c4465032..901f78d8 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -38,6 +38,8 @@ struct llama_model_loader { int n_kv = 0; int n_tensors = 0; int n_created = 0; + int n_skipped = 0; + size_t size_skipped = 0; int64_t n_elements = 0; size_t n_bytes = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d1834619..ace09ab3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -490,6 +490,13 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_HC_HEAD_NORM, "output_hc_norm" }, { LLM_TENSOR_HC_HEAD_DOWN, "output_hc_down" }, { LLM_TENSOR_HC_HEAD_UP, "output_hc_up" }, + { LLM_TENSOR_NEXTN_EH_PROJ, "blk.%d.nextn.eh_proj" }, + { LLM_TENSOR_NEXTN_ENORM, "blk.%d.nextn.enorm" }, + { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, + { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, + { LLM_TENSOR_NEXTN_HC_HEAD_NORM, "blk.%d.nextn.hc_head_norm" }, + { LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "blk.%d.nextn.hc_head_down" }, + { LLM_TENSOR_NEXTN_HC_HEAD_UP, "blk.%d.nextn.hc_head_up" }, { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, @@ -2401,6 +2408,15 @@ bool llama_model_is_step35(const llama_model * model) { return model && model->arch == LLM_ARCH_STEP35; } +bool llama_model_is_qwen4exp(const llama_model * model) { + return model && model->arch == LLM_ARCH_QWEN4EXP; +} + +// qwen4exp marks a present trunk block with hc_attn_norm, every other arch with attn_norm +static bool llama_model_trunk_block_present(const llama_layer & layer) { + return layer.attn_norm != nullptr || layer.hc_attn_norm != nullptr; +} + bool llama_model_is_qwen35_family(const llama_model * model) { return model && (model->arch == LLM_ARCH_QWEN35 || model->arch == LLM_ARCH_QWEN35MOE); } @@ -2417,7 +2433,7 @@ enum llama_mtp_package llama_model_mtp_package(const llama_model * model) { const size_t n_nextn = model->hparams.nextn_predict_layers; const bool has_common_package_contract = llama_model_is_step35(model) || llama_model_is_deepseek4(model) || - llama_model_is_qwen35_family(model) || + llama_model_is_qwen35_family(model) || llama_model_is_qwen4exp(model) || llama_model_is_gemma4_mtp_assistant(model); if (!has_common_package_contract) { return n_nextn > 0 ? LLAMA_MTP_PACKAGE_EMBEDDED : LLAMA_MTP_PACKAGE_NONE; @@ -2425,9 +2441,9 @@ enum llama_mtp_package llama_model_mtp_package(const llama_model * model) { if (n_nextn == 0) { if (llama_model_is_step35(model) || llama_model_is_deepseek4(model) || - llama_model_is_qwen35_family(model)) { + llama_model_is_qwen35_family(model) || llama_model_is_qwen4exp(model)) { for (const auto & layer : model->layers) { - if (layer.attn_norm != nullptr) { + if (llama_model_trunk_block_present(layer)) { return LLAMA_MTP_PACKAGE_TARGET_ONLY; } } @@ -2448,7 +2464,7 @@ enum llama_mtp_package llama_model_mtp_package(const llama_model * model) { bool has_trunk = false; for (size_t il = 0; il < first; ++il) { - if (model->layers[il].attn_norm != nullptr) { + if (llama_model_trunk_block_present(model->layers[il])) { has_trunk = true; break; } diff --git a/src/llama-model.h b/src/llama-model.h index 2f3cf18f..4c28e8f9 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -139,6 +139,12 @@ struct llama_layer_nextn { struct ggml_tensor * hnorm = nullptr; struct ggml_tensor * shared_head_head = nullptr; struct ggml_tensor * shared_head_norm = nullptr; + + // qwen4exp: the NextN head's final hyper-connection mixer. For merged + // community files these are resolved to shared_head_norm + the trunk mixer. + struct ggml_tensor * hc_head_norm = nullptr; + struct ggml_tensor * hc_head_down = nullptr; + struct ggml_tensor * hc_head_up = nullptr; }; // TODO: separate into "llama_layer_enc" and "llama_layer_dec" diff --git a/src/llama-spec-features.cpp b/src/llama-spec-features.cpp index 5c857f70..a24ab593 100644 --- a/src/llama-spec-features.cpp +++ b/src/llama-spec-features.cpp @@ -27,7 +27,10 @@ uint32_t llama_model_mtp_feature_width(const struct llama_model * model) { hparams.mtp_backbone_n_embd > 0) { return hparams.mtp_backbone_n_embd; } - if (model->arch == LLM_ARCH_DEEPSEEK4 && hparams.n_embd_out > hparams.n_embd) { + if ((model->arch == LLM_ARCH_DEEPSEEK4 || + model->arch == LLM_ARCH_QWEN4EXP) && + hparams.n_embd_out > hparams.n_embd) { + // the pre-final-mixer wide stream; n_embd_out holds even when the NextN block lives in a companion file return hparams.n_embd_out; } return hparams.n_embd; diff --git a/src/llama.cpp b/src/llama.cpp index d2f33b71..cc665fa3 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1156,6 +1156,7 @@ static bool llama_mtp_tail_uses_layer_cache(const llama_model & model) { return model.hparams.nextn_predict_layers > 0 && (model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_QWEN35MOE || + model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_STEP35); } @@ -2029,6 +2030,15 @@ bool llama_kv_cache::checkpoint_save(ggml_backend_sched_t sched) { ckpt.used_snapshot = used; std::unordered_set backends_to_sync; + // a row with no resolvable backend is copied generically, outside any stream; + // drain queued work once so the copy cannot race an in-flight graph + bool sched_synced = false; + auto sync_sched_once = [&]() { + if (!sched_synced) { + ggml_backend_sched_synchronize(sched); + sched_synced = true; + } + }; for (uint32_t il = 0; il < n_layer; ++il) { if (s_l[il] == nullptr) { @@ -2040,18 +2050,28 @@ bool llama_kv_cache::checkpoint_save(ggml_backend_sched_t sched) { auto & shadow_split = ckpt.split_s_l_shadow[il]; for (int d = 0; d < split_info->n_device; ++d) { if (split_info->splits[d] && shadow_split[d]) { - //ggml_backend_tensor_copy(split_info->splits[d], shadow_split[d]); auto src_backend = ggml_backend_sched_get_tensor_backend(sched, split_info->splits[d]); - ggml_backend_tensor_copy_async(src_backend, src_backend, split_info->splits[d], shadow_split[d]); - backends_to_sync.insert(src_backend); + if (src_backend == nullptr) { + // a host-resident row (e.g. CUDA_Host) matches no backend's default + // buffer type between graphs; the generic copy handles it + sync_sched_once(); + ggml_backend_tensor_copy(split_info->splits[d], shadow_split[d]); + } else { + ggml_backend_tensor_copy_async(src_backend, src_backend, split_info->splits[d], shadow_split[d]); + backends_to_sync.insert(src_backend); + } } } } else { GGML_ASSERT(ckpt.s_l_shadow[il] != nullptr); auto src_backend = ggml_backend_sched_get_tensor_backend(sched, s_l[il]); - GGML_ASSERT(src_backend != nullptr); - ggml_backend_tensor_copy_async(src_backend, src_backend, s_l[il], ckpt.s_l_shadow[il]); - backends_to_sync.insert(src_backend); + if (src_backend == nullptr) { + sync_sched_once(); + ggml_backend_tensor_copy(s_l[il], ckpt.s_l_shadow[il]); + } else { + ggml_backend_tensor_copy_async(src_backend, src_backend, s_l[il], ckpt.s_l_shadow[il]); + backends_to_sync.insert(src_backend); + } } } @@ -2078,6 +2098,15 @@ bool llama_kv_cache::checkpoint_restore(ggml_backend_sched_t sched) { used = ckpt.used_snapshot; std::unordered_set backends_to_sync; + // a row with no resolvable backend is copied generically, outside any stream; + // drain queued work once so the copy cannot race an in-flight graph + bool sched_synced = false; + auto sync_sched_once = [&]() { + if (!sched_synced) { + ggml_backend_sched_synchronize(sched); + sched_synced = true; + } + }; for (uint32_t il = 0; il < n_layer; ++il) { if (s_l[il] == nullptr) { @@ -2090,17 +2119,26 @@ bool llama_kv_cache::checkpoint_restore(ggml_backend_sched_t sched) { for (int d = 0; d < split_info->n_device; ++d) { if (split_info->splits[d] && shadow_split[d]) { auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, split_info->splits[d]); - ggml_backend_tensor_copy_async(dst_backend, dst_backend, shadow_split[d], split_info->splits[d]); - backends_to_sync.insert(dst_backend); + if (dst_backend == nullptr) { + sync_sched_once(); + ggml_backend_tensor_copy(shadow_split[d], split_info->splits[d]); + } else { + ggml_backend_tensor_copy_async(dst_backend, dst_backend, shadow_split[d], split_info->splits[d]); + backends_to_sync.insert(dst_backend); + } } } } else { GGML_ASSERT(ckpt.s_l_shadow[il] != nullptr); GGML_ASSERT(ggml_nbytes(ckpt.s_l_shadow[il]) == ggml_nbytes(s_l[il])); auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, s_l[il]); - GGML_ASSERT(dst_backend != nullptr); - ggml_backend_tensor_copy_async(dst_backend, dst_backend, ckpt.s_l_shadow[il], s_l[il]); - backends_to_sync.insert(dst_backend); + if (dst_backend == nullptr) { + sync_sched_once(); + ggml_backend_tensor_copy(ckpt.s_l_shadow[il], s_l[il]); + } else { + ggml_backend_tensor_copy_async(dst_backend, dst_backend, ckpt.s_l_shadow[il], s_l[il]); + backends_to_sync.insert(dst_backend); + } } } @@ -8632,6 +8670,7 @@ struct llama_context * llama_init_from_model( model->arch != LLM_ARCH_STEP35 && model->arch != LLM_ARCH_GEMMA4_ASSISTANT && model->arch != LLM_ARCH_OPENPANGU && + model->arch != LLM_ARCH_QWEN4EXP && cparams.mtp != 0) { cparams.mtp = 0; } @@ -9650,6 +9689,14 @@ static bool spec_ckpt_try_per_step(llama_kv_cache & kv, const llama_model & mode return false; } + // qwen4exp's recurrent row carries a PLE conv-history tail the per-step + // checkpoint doesn't size, so it would leave that tail advanced by rejected + // drafts. Decline; the whole-slot (gpu-fallback) shadow covers the full row. + if (model.arch == LLM_ARCH_QWEN4EXP) { + kv.save_per_step_ssm = false; + return false; + } + // Split recurrent tensors are supported as long as each layer exposes // concrete backend buffers for the per-step tensors. CPU-only and mixed // CPU/GPU recurrent placement are also allowed.