opencoti-llamafile / patches /0093-mtp-dualctx-fused-nextn.patch
ManniX-ITA's picture
Upload folder using huggingface_hub
b69d9d8 verified
Raw
History Blame
165 kB
From: opencoti
Subject: [PATCH 0093] bug-858 MTP parity: dual-ctx assistant port + fused NextN + verify-path fixes
Wholesale alignment of speculative MTP to upstream b9859 (#590/#607/#599/#609):
- gemma4-assistant dual-context execution (ctx_dft sharing the target KV via
cparams.ctx_other / mem_other + share selector; OPENCOTI_MTP_DUAL_CTX gate),
draft_mtp impl with pre-norm embedding tap/harvest, defer-gather boot fix.
- fused NextN greedy N-step draft (build_one_step + decode_mtp_fused_nextn on a
dedicated sched_nextn; can_reuse base-drift + ensure_cleared aliased-KV fixes).
- bug-2103: mmvf ne11 cap raised to 8 under TinyBLAS (n3 verify GEMV cliff).
- bug-2094 diag + rolling-KV visibility hunks ride along in llama-kv-cache.cpp
(former 0091-rolling-kv-diag-visibility, folded here during the bug-2108
chain repair).
diff --git a/llama.cpp/common/arg.cpp b/llama.cpp/common/arg.cpp
index 89bd9dd..3fb1518 100644
--- a/llama.cpp/common/arg.cpp
+++ b/llama.cpp/common/arg.cpp
@@ -1455,7 +1455,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
// opencoti F5 M7 Rolling KV — Rung 0 auto residency — see docs/features/rolling_kv.md
add_opt(common_arg(
{"--vram-target"}, "MiB",
- string_format("rolling-kv: VRAM budget cap in MiB for auto KV residency (used with --headinfer-gpu-heads-frac auto). 0 = use all free VRAM minus a compute reserve (maximize). (default: %d)", params.vram_target_mib),
+ string_format("rolling-kv: TOTAL GPU VRAM cap in MiB (weights + KV + compute), NOT a KV-only budget. Auto KV residency keeps as much KV on-device as fits under this cap after weights are loaded; the overflow streams from host. Set it to your card's size to simulate/bound that card. 0 = use all free VRAM minus a compute reserve (maximize). (default: %d)", params.vram_target_mib),
[](common_params & params, int value) {
params.vram_target_mib = value;
}
diff --git a/llama.cpp/common/common.h b/llama.cpp/common/common.h
index b688492..3bb6aa2 100644
--- a/llama.cpp/common/common.h
+++ b/llama.cpp/common/common.h
@@ -303,9 +303,12 @@ struct common_params_speculative_draft {
int32_t n_max = 3; // maximum number of tokens to draft during speculative decoding
int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding
- // opencoti F5 M6-S4 mtp — MTP (Gemma 4 assistant): draft block size B produces
- // B-1 draft tokens per round. Default 3 (= 2 chained MTP draft steps).
- int32_t draft_block_size = 3;
+ // opencoti F5 M6-S4 mtp — MTP (Gemma 4 assistant): OPTIONAL explicit draft-depth ceiling.
+ // 0 (default) = draft depth follows --spec-draft-n-max (n_max); the MTP head re-runs
+ // autoregressively per step (decode_mtp_sync/fused), so depth is NOT limited to a fixed
+ // block size. bug-858: the old default 3 hard-capped ours at 2 drafts/round and blocked
+ // n-max scaling vs upstream b9859. When set >1, caps draft depth at draft_block_size-1.
+ int32_t draft_block_size = 0;
float p_split = 0.1f; // speculative decoding split probability
float p_min = 0.0f; // minimum speculative decoding probability (greedy)
diff --git a/llama.cpp/common/speculative.cpp b/llama.cpp/common/speculative.cpp
index 7faeac7..ee7962f 100644
--- a/llama.cpp/common/speculative.cpp
+++ b/llama.cpp/common/speculative.cpp
@@ -481,6 +481,13 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
int32_t n_embd = 0;
+ // opencoti bug-858 dual-context MTP: true for the gemma4 assistant (ctx_dft shares the target's
+ // KV via ctx_other). Derived in the ctor. When set, process() skips the ctx_dft catch-up decode
+ // (the shared cache is already populated by the target's own decode) and draft() pins all draft
+ // tokens to the same position dp.n_past (KV-less: no per-step cell growth). Mirrors upstream
+ // b9859 common_speculative_impl_draft_mtp::is_mem_shared. false = Qwen NextN (byte-identical).
+ bool is_mem_shared = false;
+
// Per-sequence cross-batch carryover: pair (h_p, x_{p+1}) at MTP pos p+1.
// The last h-row of one process() call needs the first token of the NEXT
// call to pair with, so it's stashed here until that next call fires.
@@ -507,7 +514,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
auto * ctx_dft = this->params.ctx_dft;
GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set");
- n_embd = llama_model_n_embd(llama_get_model(ctx_dft));
+ // opencoti bug-858 dual-context MTP: the gemma4 assistant runs as ctx_dft with
+ // ctx_other=ctx_tgt (shared KV). Detect it here; every branch below keys off this flag.
+ is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt;
+
+ // n_embd = width of an h row (target backbone hidden). For the assistant the ctx_dft model
+ // is the assistant, whose n_embd_out() == the target n_embd (validated at load); for Qwen
+ // NextN the ctx_dft model IS the target, so keep the original accessor (byte-identical).
+ n_embd = is_mem_shared
+ ? llama_model_n_embd_out(llama_get_model(ctx_dft))
+ : llama_model_n_embd(llama_get_model(ctx_dft));
LOG_INF("%s: adding speculative implementation 'draft-mtp'\n", __func__);
LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling);
@@ -550,7 +566,19 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
}
}
- llama_set_embeddings_pre_norm(ctx_tgt, true, /*masked*/ false);
+ // opencoti bug-858 losslessness (2026-07-05): the target tap must not change the target's graph.
+ // The unmasked pre-norm tap engages gemma4.cpp's mtp_defer_out_ids (late row-strip: the last
+ // layer + output_norm run on ALL verify rows instead of the output rows) — structurally equal to
+ // upstream b9859's defer, but ours' GEMM kernels are batch-shape-sensitive (#495), so the deferred
+ // graph flips near-tie argmax vs plain decode (gate4/5: fib DIFF@2, forces DIFF@126; the facade,
+ // which never defers, is clean). Fix: for the shared-KV assistant set the target tap MASKED —
+ // masked keeps the normal early strip (target graph BIT-IDENTICAL to plain decode; t_h_pre_norm
+ // aliases the always-captured t_embd) and captures n_outputs rows with the output_ids-translated
+ // _ith accessor. Every verify token has logits=true, and prefill only ever consumes its last row
+ // (also an output row), so all consumed values equal upstream's all-rows defer capture — see the
+ // logits-rows harvest in process(). Qwen NextN (is_mem_shared==false) keeps the dense unmasked
+ // tap (its catch-up decode reads the whole shifted prefill buffer) — byte-identical.
+ llama_set_embeddings_pre_norm(ctx_tgt, true, /*masked*/ is_mem_shared);
llama_set_embeddings_pre_norm(ctx_dft, true, /*masked*/ true);
pending_h.assign(n_seq, std::vector<float>(n_embd, 0.0f));
@@ -591,7 +619,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
}
auto * ctx_dft = this->params.ctx_dft;
const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);
- if (pos_max < N - 1) {
+ // opencoti bug-858 dual-context MTP: for the shared-KV assistant, ctx_dft's positions are
+ // managed by the TARGET (no catch-up decode), so this warning is spurious — suppress it.
+ if (pos_max < N - 1 && !is_mem_shared) {
LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - "
"process() hook may not have run on every prefill ubatch "
"(need_embd / logits=1 on every prompt position?). "
@@ -634,6 +664,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
const size_t row_bytes = (size_t) n_embd * sizeof(float);
+ // opencoti bug-858 dual-context MTP: with a shared KV (gemma4 assistant) the ctx_dft cache is
+ // already populated by the TARGET's own decode (A2/A3 cell sharing), so skip the catch-up
+ // decode entirely and go straight to harvesting the target's verify hidden rows below.
+ // Mirrors upstream b9859 process() (`if (!is_mem_shared) { ...catch-up... }`). Qwen NextN
+ // (is_mem_shared==false) runs the catch-up exactly as before — byte-identical.
+ if (!is_mem_shared) {
common_batch_clear(batch);
for (int k = 0; k < n_tokens; ++k) {
@@ -677,19 +713,44 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (pos=%d)\n", __func__, (int) rc, (int) batch_in.pos[0]);
return false;
}
+ } // end if (!is_mem_shared): shared-KV assistant skips the ctx_dft catch-up decode
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
if (i_batch_end[seq_id] < 0) {
continue;
}
- const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
+ // opencoti bug-858 losslessness: with the MASKED target tap (is_mem_shared) only rows with
+ // logits set exist in the capture (verify batches: every row; prefill ubatches: the last row).
+ // Harvest exactly those; accept()'s row indexing is unchanged because every verify token has
+ // logits=true, and prefill only ever consumes pending_h (the last row). Qwen NextN (dense
+ // unmasked tap) takes the skip_row==false path on every row — byte-identical to before.
+ auto skip_row = [&](int32_t k) {
+ return is_mem_shared && batch_in.logits && !batch_in.logits[k];
+ };
+
+ int32_t n_rows = 0;
+ for (int32_t k = i_batch_beg[seq_id]; k <= i_batch_end[seq_id]; ++k) {
+ if (skip_row(k)) {
+ continue;
+ }
+ n_rows++;
+ }
+ if (n_rows <= 0) {
+ continue;
+ }
verify_h_rows[seq_id] = n_rows;
verify_h[seq_id].resize((size_t) n_rows * n_embd);
- for (int32_t i = 0; i < n_rows; ++i) {
- const float * h = llama_get_embeddings_pre_norm_ith(ctx_tgt, i_batch_beg[seq_id] + i);
- std::memcpy(verify_h[seq_id].data() + (size_t) i * n_embd, h, row_bytes);
+ int32_t r = 0;
+ for (int32_t k = i_batch_beg[seq_id]; k <= i_batch_end[seq_id]; ++k) {
+ if (skip_row(k)) {
+ continue;
+ }
+ // harvest via the pre-norm tap for BOTH paths (== upstream llama_get_embeddings_nextn_ith)
+ const float * h = llama_get_embeddings_pre_norm_ith(ctx_tgt, k);
+ std::memcpy(verify_h[seq_id].data() + (size_t) r * n_embd, h, row_bytes);
+ r++;
}
std::memcpy(pending_h[seq_id].data(),
@@ -711,6 +772,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
const float * h_row = nullptr;
const size_t row_bytes = (size_t) n_embd * sizeof(float);
+ // opencoti #590/bug-858 MTP profiling — gated, byte-identical when unset.
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
auto & dp = dparams[seq_id];
@@ -734,6 +796,55 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
return;
}
+ // opencoti fused-NextN (#590 / bug-858): the seed decode above put id_last into ctx_dft's
+ // cache at n_past. When OPENCOTI_MTP_FUSED_NEXTN is set, emit all N drafts from ONE fused
+ // graph (single host sync) instead of the per-step AR loop below — the throughput fix.
+ // GREEDY-ONLY: no p_min early-stop (n_steps = n_max). Default OFF until the S5 gate proves
+ // acceptance holds vs the AR path. See docs/features/fused_nextn_mtp.md.
+ static const bool fused_nextn = [] {
+ const char * e = getenv("OPENCOTI_MTP_FUSED_NEXTN");
+ return e && e[0] && e[0] != '0';
+ }();
+ // opencoti bug-858 dual-context MTP: the fused-NextN kernel is Qwen-NextN-specific
+ // (llama_decode_mtp_fused_nextn); never take it for the shared-KV assistant.
+ if (fused_nextn && params.n_max >= 2 && !is_mem_shared) {
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ if (!drafting[seq_id]) {
+ continue;
+ }
+ auto & dp = dparams[seq_id];
+ llama_memory_t mem = llama_get_memory(ctx_dft);
+ llama_pos attn_pos = mem ? llama_memory_seq_pos_max(mem, seq_id) : (llama_pos) 0;
+ if (attn_pos < 0) {
+ attn_pos = 0;
+ }
+ const int32_t n_steps = params.n_max;
+ std::vector<llama_token> out((size_t) n_steps, 0);
+ const int32_t rc = llama_decode_mtp_fused_nextn(
+ ctx_dft, seq_id, attn_pos, dp.id_last, pending_h[seq_id].data(),
+ n_steps, out.data(), /*out_h_prev_last=*/ nullptr);
+ if (rc != 0) {
+ LOG_WRN("%s: fused NextN draft failed rc=%d seq_id=%d\n", __func__, (int) rc, (int) seq_id);
+ continue;
+ }
+ auto & result = *dp.result;
+ for (int32_t k = 0; k < n_steps; ++k) {
+ result.push_back(out[(size_t) k]);
+ }
+ }
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ auto & dp = dparams[seq_id];
+ if (!dp.drafting) {
+ continue;
+ }
+ if (dp.result->size() < (size_t) params.n_min) {
+ dp.result->clear();
+ }
+ last_n_drafted[seq_id] = (uint16_t) dp.result->size();
+ }
+ return;
+ }
+
int i = 0;
while (n_drafting > 0) {
@@ -748,23 +859,39 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
auto * smpl = smpls[seq_id].get();
- common_sampler_sample(smpl, ctx_dft, i_batch, true);
+ // opencoti #590/bug-858: on-device greedy draft-token fast path. When p_min<=0 (the
+ // default — no confidence early-stop is configured), the full-vocab common_sampler_sample
+ // (logit D2H + CPU softmax/sort over ~152k vocab) is MEASURED as 0.98 ms/step = 100% of
+ // the draft-loop host tax (vs 0.14 ms decode-launch, 0.003 ms hidden). The draft sampler
+ // is deterministic-argmax and the graph already published that argmax on-device
+ // (qwen35 t_argmax), so read it back (4 B) and skip the host sampler entirely. p_min>0
+ // falls through to the exact host path (which computes the top-token probability).
+ llama_token id = -1;
+ float id_p = 1.0f; // on-device greedy is unconditionally accepted when p_min<=0
+ // opencoti bug-858 dual-context MTP: the shared-KV assistant graph publishes no
+ // on-device argmax (t_logits + t_h_pre_norm only), so host-sample it (matches
+ // upstream b9859, which always common_sampler_sample's the draft). Byte-identical
+ // for Qwen NextN (is_mem_shared==false).
+ if (params.p_min <= 0.0f && !is_mem_shared) {
+ id = llama_get_draft_greedy_ith(ctx_dft, i_batch); // includes synchronize() = GPU-wait
+ }
+ if (id < 0) {
+ // exact host sampler: p_min>0, or the graph published no on-device argmax
+ common_sampler_sample(smpl, ctx_dft, i_batch, true);
+ const auto * cur_p = common_sampler_get_candidates(smpl, true);
+ for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
+ LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+ seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
+ common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
+ }
+ id = cur_p->data[0].id;
+ id_p = cur_p->data[0].p;
+ }
h_row = llama_get_embeddings_pre_norm_ith(ctx_dft, i_batch);
++i_batch;
- const auto * cur_p = common_sampler_get_candidates(smpl, true);
-
- for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
- LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
- seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
- common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
- }
-
- // add drafted token for each sequence
- const llama_token id = cur_p->data[0].id;
-
// only collect very high-confidence draft tokens
- if (cur_p->data[0].p < params.p_min) {
+ if (id_p < params.p_min) {
drafting[seq_id] = false;
n_drafting--;
@@ -784,7 +911,13 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
continue;
}
- common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true);
+ // opencoti bug-858 dual-context MTP: the shared-KV assistant is KV-less and pins
+ // EVERY draft token to dp.n_past (no per-step cell growth; all drafts attend the
+ // frozen target prefix) — this mirrors upstream b9859 draft()'s is_mem_shared branch
+ // and the HF gemma4_assistant reference. Qwen NextN advances the position per step
+ // (byte-identical to before).
+ const llama_pos draft_pos = is_mem_shared ? dp.n_past : dp.n_past + i + 1;
+ common_batch_add(batch, id, draft_pos, { seq_id }, true);
std::memcpy(batch.embd + n_embd*(batch.n_tokens - 1), h_row, row_bytes);
}
@@ -814,6 +947,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
last_n_drafted[seq_id] = (uint16_t) dp.result->size();
}
+
}
void accept(llama_seq_id seq_id, uint16_t n_accepted, bool /*is_other*/) override {
@@ -832,10 +966,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
}
bool need_embd() const override {
+ // opencoti bug-858 losslessness (2026-07-04): NEVER put the target in plain embeddings mode.
+ // cparams.embeddings=true forces output_all=true (llama-context.cpp:2077), which changes the
+ // target's own decode graph and marks t_embd (lm_head input) as an output → perturbs the target's
+ // verify logits vs plain decode → near-tie argmax flips → spec != plain. Both A4B assistant and
+ // Qwen NextN harvest via the pre-norm tap now (matches upstream b9859, cparams.embeddings=false).
return false;
}
bool need_embd_pre_norm() const override {
+ // opencoti bug-858: both paths harvest the pre-norm hidden from ctx_tgt (== upstream embeddings_nextn).
return true;
}
};
@@ -1409,6 +1549,8 @@ struct common_speculative_impl_draft_assistant : public common_speculative_impl
// can select the last-accepted token's hidden state. The spec framework calls process() after the
// target decode, so embeddings are populated. No decode here (unlike native draft_mtp, which
// mirrors into a separate ctx_dft) — the assistant reads the target's own embeddings directly.
+ // NOTE (bug-858): PRE-norm harvest was TESTED and REFUTED — it collapses accept to ~0.05 at ALL
+ // positions (incl short ctx). This drafter expects the post-norm hidden; do NOT swap to pre_norm.
bool process(const llama_batch & batch_in) override {
if (batch_in.n_tokens <= 0 || batch_in.token == nullptr || batch_in.embd != nullptr) {
return true;
@@ -1463,8 +1605,12 @@ struct common_speculative_impl_draft_assistant : public common_speculative_impl
return;
}
- const int32_t block = params.draft_block_size;
- const int32_t n_steps_raw = block > 1 ? block - 1 : 0;
+ // bug-858: draft depth follows --spec-draft-n-max (n_max) and the room left in context
+ // (dp.n_max = n_draft_max), NOT draft_block_size. The MTP head re-runs autoregressively per
+ // step (decode_mtp_sync/fused build a fresh single-step graph each step: argmax -> next
+ // last_token, h_post -> next h_prev), so depth is bounded by n_max, not a fixed block. The
+ // legacy draft_block_size-1 ceiling (default 3 -> 2) hard-capped ours at 2 drafts/round and
+ // blocked n-max scaling vs upstream b9859 (which drafts to n_max autoregressively) -> REMOVED.
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
auto & dp = dparams[seq_id];
@@ -1488,11 +1634,10 @@ struct common_speculative_impl_draft_assistant : public common_speculative_impl
continue; // empty result -> server falls back to single-token verify
}
- int32_t n_steps = n_steps_raw;
+ int32_t n_steps = params.n_max > 0 ? params.n_max : 1;
if (dp.n_max > 0) {
n_steps = std::min(n_steps, dp.n_max);
}
- n_steps = std::min(n_steps, params.n_max);
if (n_steps <= 0) {
prev_n_acc_at_draft[seq_id] = n_acc_drafts;
continue;
@@ -1789,7 +1934,15 @@ common_speculative * common_speculative_init(common_params_speculative & params,
break;
}
case COMMON_SPECULATIVE_TYPE_MTP: { // opencoti F5 M6-S4 mtp: Gemma-4 assistant drafter
- impls.push_back(std::make_unique<common_speculative_impl_draft_assistant>(config.params, n_seq));
+ // opencoti bug-858 dual-context MTP: when the server created a real ctx_dft FROM the
+ // assistant model (ctx_other=ctx_tgt; --spec-type draft-assistant + OPENCOTI_MTP_DUAL_CTX),
+ // run the shared-KV draft_mtp driver (upstream b9859 is_mem_shared). Otherwise fall
+ // back to the proven single-context in-target assistant (ctx_dft == nullptr).
+ if (config.params.draft.ctx_dft != nullptr) {
+ impls.push_back(std::make_unique<common_speculative_impl_draft_mtp>(config.params, n_seq));
+ } else {
+ impls.push_back(std::make_unique<common_speculative_impl_draft_assistant>(config.params, n_seq));
+ }
break;
}
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
diff --git a/llama.cpp/ggml/src/ggml-backend.cpp b/llama.cpp/ggml/src/ggml-backend.cpp
index c2601b1..752e7ca 100644
--- a/llama.cpp/ggml/src/ggml-backend.cpp
+++ b/llama.cpp/ggml/src/ggml-backend.cpp
@@ -276,6 +276,14 @@ void ggml_backend_tensor_get_async(ggml_backend_t backend, const struct ggml_ten
GGML_ASSERT(backend);
GGML_ASSERT(tensor);
GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+ if (offset + size > ggml_nbytes(tensor)) { // opencoti bug-858 diag (REMOVE after root-cause)
+ fprintf(stderr, "TENSOR_GET_ASYNC_OOB: name='%s' type=%d ne=[%lld,%lld,%lld,%lld] nbytes=%zu offset=%zu size=%zu\n",
+ tensor->name, (int) tensor->type,
+ (long long) tensor->ne[0], (long long) tensor->ne[1],
+ (long long) tensor->ne[2], (long long) tensor->ne[3],
+ ggml_nbytes(tensor), offset, size);
+ fflush(stderr);
+ }
GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds");
if (backend->iface.get_tensor_async == NULL) {
diff --git a/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/llama.cpp/ggml/src/ggml-cuda/fattn.cu
index e53ac9c..e63be74 100644
--- a/llama.cpp/ggml/src/ggml-cuda/fattn.cu
+++ b/llama.cpp/ggml/src/ggml-cuda/fattn.cu
@@ -42,13 +42,6 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const ggml_tensor * Q = dst->src[0];
- if constexpr (ncols2 <= 8) {
- if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) {
- ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 8/ncols2, ncols2>(ctx, dst);
- return;
- }
- }
-
if constexpr (ncols2 <= 16) {
if (Q->ne[1] <= 16/ncols2) {
ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 16/ncols2, ncols2>(ctx, dst);
@@ -597,6 +590,13 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
if (can_use_vector_kernel) {
if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) {
+ // opencoti bug-858 (2026-07-05): REVERTED the earlier <=2 experiment back to upstream's
+ // Q->ne[1]==1. Routing the n_q=2 verify to VEC (cols_per_block=2) did NOT make it bit-equal
+ // to the n_q=1 decode (different column tiling → different accumulation order) and measurably
+ // WORSENED dual-vs-plain losslessness at n_max=1 (gate4b: 0/4 prompts byte-equal vs 1/4 on
+ // upstream ==1 semantics; the earlier accept win attributed to this change was actually the
+ // output_all/masked-tap defect, fixed in common/speculative.cpp). Keep upstream semantics:
+ // VEC for decode only; verify (n_q>=2) takes MMA exactly like b9859.
if (cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && K->ne[1] >= 8192)) {
return BEST_FATTN_KERNEL_VEC;
}
@@ -1328,8 +1328,27 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
// all_f16 fast path (the memcpy lift) but cheap to compute.
const size_t ts_k = ggml_type_size(K->type);
const size_t ts_v = ggml_type_size(V->type);
- ggml_cuda_pool_alloc<char> slotK(pool, (size_t)n_slots * k_slot);
- ggml_cuda_pool_alloc<char> slotV(pool, (size_t)n_slots * v_slot);
+ // opencoti-hook: f5-rolling-kv #586 — PERSISTENT staging ring (was
+ // ggml_cuda_pool_alloc). The pool re-handed the SAME slot addresses every
+ // op/layer, which is the sole reason the coarse per-op cs_sync barrier
+ // existed (#587/bug-255): copy_stream had to wait for ALL of the previous
+ // layer's compute before ANY lift, serialising the tail DMA across every
+ // streaming layer — the S0 overflow-decode cliff (docs/features/rolling_kv.md).
+ // A persistent ring gives STABLE slot addresses across ops, so the existing
+ // per-slot compute_done events gate cross-op WAR at slot granularity and the
+ // barrier is dropped below. SAME n_slots × slot_size footprint → no extra
+ // VRAM (the serving target is VRAM-bound); realloc-on-grow as tile_kv_full /
+ // heads grow during decode. Single-device static (matches the existing
+ // single-static copy_done/compute_done event pool; the overlap path is
+ // single-device in practice).
+ static char * g_slotK = nullptr; static size_t g_capK = 0;
+ static char * g_slotV = nullptr; static size_t g_capV = 0;
+ const size_t need_k = (size_t)n_slots * k_slot;
+ const size_t need_v = (size_t)n_slots * v_slot;
+ if (g_capK < need_k) { if (g_slotK) CUDA_CHECK(cudaFree(g_slotK)); CUDA_CHECK(cudaMalloc((void**)&g_slotK, need_k)); g_capK = need_k; }
+ if (g_capV < need_v) { if (g_slotV) CUDA_CHECK(cudaFree(g_slotV)); CUDA_CHECK(cudaMalloc((void**)&g_slotV, need_v)); g_capV = need_v; }
+ char * const slotK_base = g_slotK;
+ char * const slotV_base = g_slotV;
// R2-b S3b (#312): with >=2 slots the op runs a double-buffer ping-pong — the
// per-tile lift is issued on a dedicated copy_stream while FA/lse for the
@@ -1348,7 +1367,20 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
if (use_overlap) {
const int cs_idx = GGML_CUDA_MAX_STREAMS - 1; // fixed, collision-free
GGML_ASSERT(cs_idx >= 1);
- cudaStream_t cs = ctx.stream(ctx.device, cs_idx);
+ // opencoti-hook: f5-rolling-kv — bug-1841 (#588/#587). At DECODE (n_q==1) collapse
+ // the copy/compute double-buffer onto ctx.stream(): the cross-stream
+ // copy_done/compute_done waits below CANNOT be captured into the decode CUDA graph
+ // ("dependency created on uncaptured work in another stream") — which forced graphs
+ // OFF and made the streaming op launch/sync-bound (per-tile × 48-layer event
+ // serialization, PCIe-independent ~1 tps). On a single stream the tile lift → FA →
+ // lse are stream-ordered (every real dependency implicit; the slot ping-pong's WAR
+ // is satisfied because compute(t) is enqueued before lift(t+2) reuses its slot), so
+ // the whole op is graph-capturable and the 48-layer decode is amortized by graph
+ // replay. n_q>1 (prefill) keeps the true copy/compute overlap (compute-heavy, never
+ // graph-captured). Trades the decode copy/compute overlap — cheap on the fast
+ // serving-target PCIe — for graph amortization, the actual decode-tps lever.
+ const bool single_stream = (n_q == 1);
+ cudaStream_t cs = single_stream ? stream : ctx.stream(ctx.device, cs_idx);
// Pre-created event pool (>= S3a slot cap 4). cudaEventDisableTiming;
// created ONCE during the eager warmup (op dispatch is single-thread,
@@ -1364,20 +1396,19 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
ev_init = true;
}
- // bug-255: the slot pool buffers are reclaimed/re-handed across op
- // invocations (every layer reuses the same pool addresses). In S3a
- // copy+FA shared one stream so layer L+1's copy was implicitly ordered
- // after layer L's FA. With the copy moved to copy_stream, layer L+1's
- // lift could overwrite a slot while layer L's FA (on the compute
- // stream) is still reading it → cross-op WAR. Order copy_stream after
- // all prior compute-stream work so a reused slot is never clobbered
- // mid-read. (One event per op entry; intra-op order is the ping-pong.)
- static cudaEvent_t cs_sync = nullptr;
- if (cs_sync == nullptr) {
- CUDA_CHECK(cudaEventCreateWithFlags(&cs_sync, cudaEventDisableTiming));
- }
- CUDA_CHECK(cudaEventRecord(cs_sync, stream));
- CUDA_CHECK(cudaStreamWaitEvent(cs, cs_sync, 0));
+ // opencoti-hook: f5-rolling-kv #587 — cross-layer run-ahead. The coarse
+ // per-op cs_sync barrier (record on `stream`, wait on `cs`) used to force
+ // copy_stream to wait for ALL of the previous layer's compute before any
+ // lift — because the pool re-handed the same slot addresses each op
+ // (bug-255). With the PERSISTENT ring (#586, above) slot addresses are
+ // stable, so per-slot compute_done events gate cross-op WAR at slot
+ // granularity: slot_live[s] marks a slot holding an un-retired writer from
+ // a prior op/layer, and each lift into slot s waits compute_done[s] iff
+ // slot_live[s]. Dropping the whole-stream barrier lets layer L+1's tail
+ // DMA overlap layer L's compute — inside the replayed cuda-graph, exactly
+ // where S0 located the cliff cost. Byte-identical (only a conservative
+ // barrier is removed; the per-slot edges preserve every real dependency).
+ static bool slot_live[4] = {};
// Lift tile t's K/V into its slot on the copy_stream (cudaMemcpyDefault
// → UVA H2D/D2D). Same packed 1-D-per-(b,h) layout as the S3a path.
@@ -1402,8 +1433,8 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
const int slot = t % n_slots;
const size_t k_nb2 = (size_t)this_kv * k_row;
const size_t v_nb2 = (size_t)this_kv * v_row;
- char * sK = slotK.ptr + (size_t)slot * k_slot;
- char * sV = slotV.ptr + (size_t)slot * v_slot;
+ char * sK = slotK_base + (size_t)slot * k_slot;
+ char * sV = slotV_base + (size_t)slot * v_slot;
if (!all_f16) {
// S3d dequant-on-lift: ONE nc-converter launch per tile reads the
// (head_dim × this_kv × n_head_kv × n_b) source sub-block at its
@@ -1422,17 +1453,45 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
(int64_t)(tV->nb[1]/ts_v), (int64_t)(tV->nb[2]/ts_v), (int64_t)(tV->nb[3]/ts_v), cs);
} else
for (int bb = 0; bb < n_b; ++bb) {
+ // opencoti-hook: f5-rolling-kv #586 — bulk H2D + on-device repack.
+ // The per-head strided 2D copy issues n_head_kv rows of only k_row
+ // (=head_dim·2B = 256B) bytes each; over PCIe those tiny strided
+ // transfers waste bandwidth (measured ~3.1 GB/s vs the 6.5 GB/s
+ // solidPC ceiling — ~2× loss; the loss is worse at higher BW where
+ // small rows can't saturate the link). The host tail is CONTIGUOUS
+ // in cell-major [n_embd_gqa × this_kv] (nb[1]=cell stride), so for a
+ // HOST source bulk-copy the whole block H2D once (full bandwidth)
+ // into a device scratch, then do the cell→head repack ON-DEVICE
+ // (D2D strided ≈ GPU bandwidth, ~free). Byte-identical: the slot ends
+ // up with the exact same packed bytes. A device window source is
+ // already D2D, so keep its in-place strided path (no scratch).
+ const bool bulk = use_2d && tr.host;
+ const size_t kblk = (size_t)this_kv * tK->nb[1]; // contiguous cell-major span
+ const size_t vblk = (size_t)this_kv * tV->nb[1];
+ ggml_cuda_pool_alloc<char> scrK(pool, bulk ? kblk : 1);
+ ggml_cuda_pool_alloc<char> scrV(pool, bulk ? vblk : 1);
+ const char * baseK;
+ const char * baseV;
+ if (bulk) {
+ const char * hK = (const char *)tK->data + (size_t)bb*tK->nb[3] + (size_t)loc*tK->nb[1];
+ const char * hV = (const char *)tV->data + (size_t)bb*tV->nb[3] + (size_t)loc*tV->nb[1];
+ CUDA_CHECK(cudaMemcpyAsync(scrK.ptr, hK, kblk, cudaMemcpyDefault, cs)); // one full-BW H2D
+ CUDA_CHECK(cudaMemcpyAsync(scrV.ptr, hV, vblk, cudaMemcpyDefault, cs));
+ baseK = scrK.ptr; // scratch is cell-major, cell `loc` at offset 0
+ baseV = scrV.ptr;
+ } else {
+ baseK = (const char *)tK->data + (size_t)bb*tK->nb[3] + (size_t)loc*tK->nb[1];
+ baseV = (const char *)tV->data + (size_t)bb*tV->nb[3] + (size_t)loc*tV->nb[1];
+ }
for (int hh = 0; hh < n_head_kv; ++hh) {
const size_t dstoff = (size_t)bb*n_head_kv + hh;
- const char * srcK = (const char *)tK->data + (size_t)bb*tK->nb[3] + (size_t)hh*tK->nb[2] + (size_t)loc*tK->nb[1];
- const char * srcV = (const char *)tV->data + (size_t)bb*tV->nb[3] + (size_t)hh*tV->nb[2] + (size_t)loc*tV->nb[1];
+ const char * srcK = baseK + (size_t)hh*tK->nb[2]; // loc already folded into base
+ const char * srcV = baseV + (size_t)hh*tV->nb[2];
if (use_2d) {
- // S3c: strided source (pinned-host CPU-half, OR the 3c-5
- // permuted device window — all heads per cell). Key rows
- // strided by nb[1]; the 2D copy packs them into the
- // contiguous slot (dst pitch == width == row), byte-
- // equivalent to the device 1-D path. cudaMemcpyDefault
- // (UVA) reads host or device transparently.
+ // repack cell→head-major: rows strided by nb[1], packed into
+ // the contiguous slot (dst pitch == width == row). For a bulk
+ // host tail this is now D2D (scratch→slot, fast); for a device
+ // window it is the original in-place strided D2D.
CUDA_CHECK(cudaMemcpy2DAsync(
sK + dstoff*k_nb2, k_row, srcK, tK->nb[1], k_row, this_kv, cudaMemcpyDefault, cs));
CUDA_CHECK(cudaMemcpy2DAsync(
@@ -1446,8 +1505,10 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
};
// Prologue: stage tile 0 so its compute can start as soon as it lands.
+ // #587: if slot 0 still holds a prior op's writer, wait its compute first.
+ if (!single_stream && slot_live[0]) CUDA_CHECK(cudaStreamWaitEvent(cs, compute_done[0], 0));
lift(0);
- CUDA_CHECK(cudaEventRecord(copy_done[0], cs));
+ if (!single_stream) CUDA_CHECK(cudaEventRecord(copy_done[0], cs));
for (int t = 0; t < n_tiles; ++t) {
const int slot = t % n_slots;
@@ -1457,11 +1518,13 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
// compute to release it (compute_done was recorded n_slots tiles ago).
if (t + 1 < n_tiles) {
const int nslot = (t + 1) % n_slots;
- if (t + 1 >= n_slots) {
+ // #587: wait if the slot is reused within this op (t+1>=n_slots)
+ // OR still holds a prior op/layer's un-retired writer (slot_live).
+ if (!single_stream && (t + 1 >= n_slots || slot_live[nslot])) {
CUDA_CHECK(cudaStreamWaitEvent(cs, compute_done[nslot], 0));
}
lift(t + 1);
- CUDA_CHECK(cudaEventRecord(copy_done[nslot], cs));
+ if (!single_stream) CUDA_CHECK(cudaEventRecord(copy_done[nslot], cs));
}
// Compute tile t on ctx.stream(): wait for its copy, then FA + lse.
@@ -1473,9 +1536,9 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
const int kv0 = trc.global_kv0;
const int this_kv = trc.this_kv;
float * lse_t = lse_tiles.ptr + (size_t)t * n_rows;
- char * sK = slotK.ptr + (size_t)slot * k_slot;
- char * sV = slotV.ptr + (size_t)slot * v_slot;
- CUDA_CHECK(cudaStreamWaitEvent(stream, copy_done[slot], 0));
+ char * sK = slotK_base + (size_t)slot * k_slot;
+ char * sV = slotV_base + (size_t)slot * v_slot;
+ if (!single_stream) CUDA_CHECK(cudaStreamWaitEvent(stream, copy_done[slot], 0));
if (this_kv > 0) {
const size_t k_nb2 = (size_t)this_kv * k_row;
@@ -1538,13 +1601,28 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
ggml_cuda_flash_attn_ext(ctx, &dst_t);
if (!decode_lse) {
- streaming_lse_kernel<<<n_rows, WARP_SIZE, 0, stream>>>(
- (const char *)Q->data, sK,
- mask ? (const char *)mask->data + (size_t)kv0*mask->nb[0] : nullptr,
- lse_t, scale, logit_softcap, this_kv, gqa, head_dim, n_head, n_q,
- Q->nb[1], Q->nb[2], Q->nb[3], k_row, k_nb2, k_nb2 * n_head_kv,
- mask ? mask->nb[0] : 0, mask ? mask->nb[1] : 0);
- CUDA_CHECK(cudaGetLastError());
+ // opencoti F5 bug-1840 fix (b) DIAGNOSTIC — OPENCOTI_LSE_NOOP skips
+ // the per-tile streaming_lse recompute (a low-parallelism
+ // <<<n_rows,32>>> rescan run once per tile because Qwen/head_dim≤256
+ // decode picks the finalize-less VEC FA → no dst_lse). Skipping it
+ // CORRUPTS the online-softmax combine (needle lost) — it exists ONLY
+ // to measure, via real decode tps, how much of the streaming-spill
+ // cost is this recompute vs the per-tile FA itself. If the tps ceiling
+ // justifies it, the real fix is a key-tiled parallel lse kernel (or
+ // emitting lse from the VEC FA). Default-off ⇒ byte-identical.
+ static const bool lse_noop = [] {
+ const char * e = getenv("OPENCOTI_LSE_NOOP");
+ return e != nullptr && atoi(e) == 1;
+ }();
+ if (!lse_noop) {
+ streaming_lse_kernel<<<n_rows, WARP_SIZE, 0, stream>>>(
+ (const char *)Q->data, sK,
+ mask ? (const char *)mask->data + (size_t)kv0*mask->nb[0] : nullptr,
+ lse_t, scale, logit_softcap, this_kv, gqa, head_dim, n_head, n_q,
+ Q->nb[1], Q->nb[2], Q->nb[3], k_row, k_nb2, k_nb2 * n_head_kv,
+ mask ? mask->nb[0] : 0, mask ? mask->nb[1] : 0);
+ CUDA_CHECK(cudaGetLastError());
+ }
}
} else {
// Empty tile (this_kv==0): the kernel writes lse=-inf and never
@@ -1555,7 +1633,8 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
Q->nb[1], Q->nb[2], Q->nb[3], k_row, 0, 0, 0, 0);
CUDA_CHECK(cudaGetLastError());
}
- CUDA_CHECK(cudaEventRecord(compute_done[slot], stream));
+ if (!single_stream) CUDA_CHECK(cudaEventRecord(compute_done[slot], stream));
+ slot_live[slot] = true; // #587: slot holds an un-retired writer for the next op/layer
}
} else
// S3a: single compute stream. The per-tile slot lift, the FA over the
@@ -1569,8 +1648,8 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor
const int kv0 = t * tile_kv_full;
const int this_kv = kv0 >= n_kv ? 0 : min(tile_kv_full, n_kv - kv0);
float * lse_t = lse_tiles.ptr + (size_t)t * n_rows;
- char * sK = slotK.ptr + (size_t)slot * k_slot;
- char * sV = slotV.ptr + (size_t)slot * v_slot;
+ char * sK = slotK_base + (size_t)slot * k_slot;
+ char * sV = slotV_base + (size_t)slot * v_slot;
if (this_kv > 0) {
const size_t k_nb2 = (size_t)this_kv * k_row;
diff --git a/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
index 4db22f3..bad56a5 100644
--- a/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -3344,6 +3344,23 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
continue;
}
+ // opencoti-hook: f5-rolling-kv — bug-1841 (#588/#587). The POSITION_WINDOW streaming FA op
+ // (ggml_cuda_streaming_flash_attn, fattn.cu) uses a copy/compute double-buffer ONLY at
+ // prefill (n_q>1): a SEPARATE copy stream + cross-stream copy_done/compute_done event
+ // waits, which cannot be captured into a CUDA graph ("dependency created on uncaptured
+ // work in another stream", fattn.cu:1528). At DECODE (n_q==1, src[0]->ne[1]==1) the op
+ // now runs single-stream (all lift/FA/lse on ctx.stream(), no cross-stream events) →
+ // fully graph-capturable, which is the actual decode-tps lever (the streaming path was
+ // launch/sync-bound, PCIe-independent, with graphs off). So disable graphs ONLY for the
+ // prefill overlap path (n_q>1, which is never graph-captured anyway — belt-and-suspenders);
+ // decode keeps CUDA graphs. See buglog bug-1841, docs/features/rolling_kv.md.
+ if (node->op == GGML_OP_STREAMING_FLASH_ATTN && node->src[0] && node->src[0]->ne[1] > 1) {
+ use_cuda_graph = false;
+#ifndef NDEBUG
+ GGML_LOG_DEBUG("%s: disabling CUDA graphs due to rolling-KV streaming FA (prefill overlap)\n", __func__);
+#endif
+ }
+
if (node->src[0] && node->src[0]->buffer && ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft)) {
use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture
#ifndef NDEBUG
@@ -4617,7 +4634,41 @@ static enum ggml_status GGML_CALL ggml_backend_cuda_graph_compute(ggml_backend_t
if (graph_compatible) {
const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph);
- if (!graph->warmup_complete) {
+ // opencoti F5 bug-1840 fix (a) — graph-live decode. The stock 0.10.3
+ // warmup machine only engages graphs after TWO consecutive calls with
+ // NO property change, and RESETS to eager on any subsequent change.
+ // During decode the KV cache ne[1] grows one cell/step, so
+ // ggml_cuda_graph_update_required reports "changed" every step →
+ // warmup never completes → graphs never engage (measured:
+ // cudaGraphLaunch=0 across every decode profile). Mainline llama.cpp
+ // instead keeps graphs live across a growing KV via cudaGraphExecUpdate
+ // (already implemented here as ggml_cuda_graph_update_executable, 3440:
+ // cudaGraphExecUpdate + topology-change re-instantiate fallback). Under
+ // OPENCOTI_GRAPH_LIVE, complete warmup after ONE stabilization call, then
+ // stay on graphs and drive recapture-record + exec-update on each property
+ // change instead of dropping to eager. graph_key (nodes[0]) is stable
+ // across steps when llama.cpp reuses the graph result (can_reuse), so the
+ // per-step cost is a CPU recapture-record + cheap exec-update, not a
+ // re-instantiate. Default-off ⇒ byte-identical to stock.
+ // opencoti #590/bug-858: default-ON was TESTED and REVERTED — the
+ // bs2 A/B (graph-live on vs off) showed NO MTP benefit (35B n2
+ // 233≈234, n3 still collapsed 156≈158; spec cells inert) and a
+ // REGRESSION on base decode (35B 201 vs 212) + A4B n3 (189 vs 209).
+ // The n=3 spec collapse is NOT CUDA-graph thrash — it is per-step
+ // spec-loop overhead (host readback/sync in the un-fused draft loop);
+ // graph-live is the wrong lever. Keep the knob env-gated, default-off.
+ static const bool graph_live = [] {
+ const char * e = getenv("OPENCOTI_GRAPH_LIVE");
+ return e != nullptr && atoi(e) == 1;
+ }();
+ if (graph_live) {
+ if (!graph->warmup_complete) {
+ graph->warmup_complete = true; // one stabilization call, then engage
+ } else {
+ use_cuda_graph = true;
+ cuda_graph_update_required = properties_changed || graph->instance == nullptr;
+ }
+ } else if (!graph->warmup_complete) {
// Warmup: need at least 2 calls with no property change on the 2nd call
if (!properties_changed) {
graph->warmup_complete = true;
diff --git a/llama.cpp/ggml/src/ggml-cuda/mmvf.cu b/llama.cpp/ggml/src/ggml-cuda/mmvf.cu
index 09d95f3..4051188 100644
--- a/llama.cpp/ggml/src/ggml-cuda/mmvf.cu
+++ b/llama.cpp/ggml/src/ggml-cuda/mmvf.cu
@@ -800,6 +800,14 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0
switch (type) {
case GGML_TYPE_F32:
if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
+#ifdef GGML_USE_TINYBLAS
+ // opencoti-hook: mtp-verify-mmvf (bug-858/#609) — upstream caps NVIDIA f32 at
+ // ne11<=3 because past that its cuBLAS fallback wins; TinyBLAS builds have no
+ // cuBLAS and the fallback is a single-block tinyblasGE (~174us vs ~1.2us for
+ // the fused-GDN dot GEMVs in a 4-token MTP verify batch => 35% of decode wall
+ // at spec depth 3). Use the kernel's real max (8, the non-NVIDIA default below).
+ return ne11 <= 8;
+#endif // GGML_USE_TINYBLAS
if (ampere_mma_available(cc)) {
return ne11 <= 3;
}
diff --git a/llama.cpp/include/llama.h b/llama.cpp/include/llama.h
index a0a435c..e6517be 100644
--- a/llama.cpp/include/llama.h
+++ b/llama.cpp/include/llama.h
@@ -424,6 +424,9 @@ extern "C" {
ggml_abort_callback abort_callback;
void * abort_callback_data;
+ // opencoti bug-858: dual-context MTP — target context whose KV the draft context shares (nullptr if none)
+ struct llama_context * ctx_other;
+
// Keep the booleans together and at the end of the struct to avoid misalignment during copy-by-value.
bool embeddings; // if true, extract embeddings (together with logits)
bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU
@@ -617,6 +620,8 @@ extern "C" {
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx);
LLAMA_API llama_memory_t llama_get_memory (const struct llama_context * ctx);
+ // opencoti bug-858: dual-context MTP — the target context this context shares KV with (nullptr if none)
+ LLAMA_API struct llama_context * llama_get_ctx_other(struct llama_context * ctx);
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type
LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
@@ -1052,6 +1057,19 @@ extern "C" {
float * out_logits,
float * out_h_prev_last);
+ // opencoti fused-NextN (#590): fuse N per-step NextN draft decodes into one graph on a
+ // DECODER_MTP (Qwen NextN) draft context. GREEDY-ONLY (no p_min early-stop). Returns 0 on
+ // success. last_token must already be resident in ctx's cache at attn_pos (seed decode first).
+ LLAMA_API int32_t llama_decode_mtp_fused_nextn(
+ struct llama_context * ctx,
+ llama_seq_id seq_id,
+ llama_pos attn_pos,
+ llama_token last_token,
+ const float * h_prev,
+ int32_t n_steps,
+ llama_token * out_drafts,
+ float * out_h_prev_last);
+
// Set the number of threads used for decoding
// n_threads is the number of threads used for generation (single token)
// n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens)
diff --git a/llama.cpp/src/llama-context.cpp b/llama.cpp/src/llama-context.cpp
index 9dbc574..7d27212 100644
--- a/llama.cpp/src/llama-context.cpp
+++ b/llama.cpp/src/llama-context.cpp
@@ -155,6 +155,7 @@ llama_context::llama_context(
cparams.cb_eval_user_data = params.cb_eval_user_data;
cparams.ctx_type = params.ctx_type;
+ cparams.ctx_other = params.ctx_other; // opencoti bug-858: dual-context MTP KV sharing
// Initialize backend samplers here so they are part of the sampling graph
// before the reserve passes run later in this function. This avoids a later
@@ -398,6 +399,7 @@ llama_context::llama_context(
/*.type_v =*/ params.type_v,
/*.swa_full =*/ params.swa_full,
/*.ctx_type= */ cparams.ctx_type,
+ /*.mem_other=*/ cparams.ctx_other ? llama_get_memory(cparams.ctx_other) : nullptr, // opencoti bug-858
};
// opencoti bug-1342: pass 1 forces a minimal resident window (measure mode).
@@ -498,6 +500,7 @@ llama_context::llama_context(
/*.type_v =*/ params.type_v,
/*.swa_full =*/ params.swa_full,
/*.ctx_type= */ cparams.ctx_type,
+ /*.mem_other=*/ cparams.ctx_other ? llama_get_memory(cparams.ctx_other) : nullptr, // opencoti bug-858
};
memory.reset(model.create_memory(params_mem2, cparams));
@@ -576,6 +579,8 @@ void llama_context::sched_reserve() {
gf_res_prev.reset(new llm_graph_result(max_nodes));
gf_res_reserve.reset(new llm_graph_result(max_nodes));
+ // opencoti #590/bug-858: sched_nextn + gf_res_prev_nextn are created lazily in
+ // decode_mtp_fused_nextn (sized graph_max_nodes*n_steps) — see header note.
sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
@@ -2012,7 +2017,17 @@ int llama_context::decode(const llama_batch & batch_inp) {
const auto & hparams = model.hparams;
const int64_t n_vocab = vocab.n_tokens();
- const int64_t n_embd = hparams.n_embd_inp();
+ // opencoti bug-858 (dual-context MTP feed width): the gemma4-assistant drafter runs as ctx_dft
+ // (LLAMA_CONTEXT_TYPE_MTP) and consumes the TARGET's backbone hidden — n_embd_out wide (2816) —
+ // as its input embd, NOT its own n_embd (1024). n_embd_inp() returns 1024, so the batch→ubatch
+ // split (llama-batch.cpp: memcpy stride n_embd) would carry 1024-wide rows while
+ // llm_graph_input_embd::set_input writes n_embd_out(2816) → overread of uninitialised staging →
+ // partial NaN in inp_h → 0% accept. Size the batch by n_embd_out for such a context. Guarded to
+ // the genuine mismatch (n_embd_out_impl != n_embd), so normal models stay byte-identical.
+ const int64_t n_embd = (cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_embd_out_impl > 0 &&
+ hparams.n_embd_out_impl != hparams.n_embd)
+ ? (int64_t) hparams.n_embd_out()
+ : (int64_t) hparams.n_embd_inp();
// when computing embeddings, all tokens are output
const bool output_all = cparams.embeddings;
@@ -2836,7 +2851,13 @@ int32_t llama_context::decode_mtp_sync(
// graph; each step's argmax feeds next step's last_token; h_post -> next step's h_prev.
for (int32_t k = 0; k < n_steps; ++k) {
data->token[0] = last_token;
- data->pos[0] = attn_pos + 1 + (llama_pos) k;
+ // bug-858: the gemma4-assistant (is_mem_shared) drafts ALL steps from the SAME query
+ // position — the head is trained to predict multiple future tokens from one position, with
+ // the recurrence carried by the h_prev hidden chain (NOT the RoPE angle). Incrementing the
+ // position per step (attn_pos+1+k) hands the head a mistrained RoPE angle from step 1 on,
+ // degrading step-1+ drafts (aggregate accept ~0.69 vs upstream 0.79). Match b9859 draft_mtp
+ // is_mem_shared: common_batch_add(batch, id, dp.n_past, ...) — constant dp.n_past = attn_pos+1.
+ data->pos[0] = attn_pos + 1;
std::memcpy(data->embd.data(), h_prev, n_bb * sizeof(float));
llama_memory_context_ptr mctx = kv_iswa->init_mtp(seq_id, ub);
@@ -2978,8 +2999,12 @@ int32_t llama_context::decode_mtp_fused(
data->token[0] = last_token;
std::memcpy(data->embd.data(), h_prev, n_bb * sizeof(float));
+ // bug-858: constant query position for ALL draft steps (see decode_mtp_sync note). The
+ // gemma4-assistant (is_mem_shared) is trained to predict every future draft token from the
+ // SAME position (dp.n_past = attn_pos+1); recurrence lives in the h_prev chain, not the RoPE
+ // angle. The old attn_pos+1+k mistrained step-1+ RoPE -> accept 0.69 vs upstream 0.79.
for (int32_t k = 0; k < n_steps; ++k) {
- data->pos[k] = attn_pos + 1 + (llama_pos) k;
+ data->pos[k] = attn_pos + 1;
}
llama_memory_context_ptr mctx = kv_iswa->init_mtp(seq_id, ub);
@@ -3016,6 +3041,192 @@ int32_t llama_context::decode_mtp_fused(
return 0;
}
+// opencoti fused-NextN (#590 / bug-858): fuse the N per-step NextN draft decodes into ONE graph on
+// the draft context (ctx_dft, type LLAMA_CONTEXT_TYPE_MTP → LLM_GRAPH_TYPE_DECODER_MTP). This is the
+// throughput fix: the un-fused path pays a host sync per draft step; here one graph emits N greedy
+// drafts with a single sync. Twin of decode_mtp_fused (Gemma assistant) but for the UNIFIED cache —
+// so no mtp_assistant / sched_mtp / kv_iswa. The N fused steps read-only cross-attend the frozen
+// prefix (build_attn_readonly_nextn, Option A); recurrence is carried by the on-device hidden chain.
+// GREEDY-ONLY: no p_min early-stop (the caller reconciles draft length). See
+// docs/features/fused_nextn_mtp.md. Returns 0 on success, negative on error.
+int32_t llama_context::decode_mtp_fused_nextn(
+ llama_seq_id seq_id,
+ llama_pos attn_pos,
+ llama_token last_token,
+ const float * h_prev,
+ int32_t n_steps,
+ llama_token * out_drafts,
+ float * out_h_prev_last) {
+ if (!memory) {
+ LLAMA_LOG_ERROR("%s: context has no KV memory\n", __func__);
+ return -2;
+ }
+ // NextN self-spec targets are full-attention → unified cache. (iSWA NextN is not a thing;
+ // the Gemma assistant path is decode_mtp_fused.)
+ auto * kv = dynamic_cast<llama_kv_cache *>(memory.get());
+ if (!kv) {
+ LLAMA_LOG_ERROR("%s: fused NextN requires a unified llama_kv_cache\n", __func__);
+ return -3;
+ }
+ if (n_steps <= 1) {
+ // Nothing to fuse — the caller should use its per-step path.
+ return -9;
+ }
+
+ const uint32_t n_embd = model.hparams.n_embd;
+ if (n_embd == 0) {
+ return -4;
+ }
+
+ // n_tokens stays 1 (each fused step processes a single token); pos[] carries the N per-step
+ // query positions consumed by the wrapper's inp_pos_steps (set_input).
+ auto data = std::make_shared<llama_ubatch::data_t>();
+ data->token.resize(1);
+ data->embd.resize(n_embd);
+ data->pos.resize(n_steps);
+ data->n_seq_id.resize(1);
+ data->seq_id.resize(1);
+ data->seq_id_data.resize(1);
+ data->output.resize(1);
+ data->seq_idx.resize(LLAMA_MAX_SEQ, -1);
+ data->seq_id_unq.push_back(seq_id);
+ data->seq_idx[(size_t) seq_id] = 0;
+
+ llama_ubatch ub{};
+ ub.b_equal_seqs = 1;
+ ub.n_tokens = 1;
+ ub.n_seq_tokens = 1;
+ ub.n_seqs = 1;
+ ub.n_seqs_unq = 1;
+ ub.n_pos = (uint32_t) n_steps;
+ ub.token = data->token.data();
+ ub.embd = data->embd.data();
+ ub.pos = data->pos.data();
+ ub.n_seq_id = data->n_seq_id.data();
+ ub.seq_id = data->seq_id.data();
+ ub.seq_id_unq = data->seq_id_unq.data();
+ ub.seq_idx = data->seq_idx.data();
+ ub.output = data->output.data();
+ ub.data = data;
+
+ data->n_seq_id[0] = 1;
+ data->seq_id_data[0] = seq_id;
+ data->seq_id[0] = &data->seq_id_data[0];
+ data->output[0] = 1;
+
+ data->token[0] = last_token;
+ std::memcpy(data->embd.data(), h_prev, n_embd * sizeof(float));
+ for (int32_t k = 0; k < n_steps; ++k) {
+ data->pos[k] = attn_pos + 1 + (llama_pos) k;
+ }
+
+ // Read-only cross-attention context over the frozen prefix (single mtp_slot_info cell; get_n_kv
+ // reports the full used extent so the mask exposes 0..pmax). NO apply() → no draft-KV write.
+ llama_kv_cache::slot_info_vec_t sinfos;
+ sinfos.push_back(kv->mtp_slot_info(seq_id));
+ std::vector<llama_ubatch> ubatches;
+ ubatches.push_back(ub);
+ auto mctx = std::make_unique<llama_kv_cache_context>(kv, std::move(sinfos), std::move(ubatches));
+ if (mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) {
+ LLAMA_LOG_ERROR("%s: failed to build mtp_slot_info context\n", __func__);
+ return -5;
+ }
+
+ // Build/compute the fused N-step graph on the NORMAL sched (graph_params sets n_mtp_steps=N for
+ // DECODER_MTP). Inlined like process_ubatch_mtp — deliberately NOT process_ubatch, whose apply()
+ // would overwrite the pmax prefix cell.
+ mtp_fused_steps = n_steps;
+
+ // opencoti #590/bug-858: lazily create a DEDICATED scheduler + result for the fused NextN draft so
+ // its ggml graph allocation AND CUDA-graph capture survive across draft rounds. Sharing the main
+ // sched with target decode reset the fused graph's compute buffers every round → 100% rebuild +
+ // recapture AND garbage argmax on reuse. Sized graph_max_nodes*n_steps (the graph unrolls N steps).
+ const int32_t nextn_need = std::max(n_steps, 1);
+ if (!sched_nextn || !gf_res_prev_nextn || nextn_reserved_steps < nextn_need) {
+ sched_nextn.reset();
+ gf_res_prev_nextn.reset();
+ const uint32_t n_tok_nextn = std::min(cparams.n_ctx, cparams.n_ubatch);
+ const size_t mn_nextn = (size_t) this->graph_max_nodes(n_tok_nextn) * (size_t) nextn_need;
+ gf_res_prev_nextn.reset(new llm_graph_result(mn_nextn));
+ sched_nextn.reset(ggml_backend_sched_new(
+ backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(),
+ mn_nextn, /*pipeline_parallel*/ false, cparams.op_offload));
+ if (!sched_nextn) {
+ LLAMA_LOG_ERROR("%s: ggml_backend_sched_new failed for sched_nextn\n", __func__);
+ gf_res_prev_nextn.reset();
+ mtp_fused_steps = 1;
+ return -6;
+ }
+ nextn_reserved_steps = nextn_need;
+ }
+
+ auto * res = gf_res_prev_nextn.get();
+ llm_graph_params gparams = graph_params(res, ub, mctx.get(), LLM_GRAPH_TYPE_DECODER_MTP);
+ gparams.n_outputs = 1;
+ gparams.sched = sched_nextn.get();
+
+ ggml_status status = GGML_STATUS_SUCCESS;
+ if (!graph_reuse_disable && res->can_reuse(gparams)) {
+ // reuse — dedicated sched_nextn is NOT clobbered by interleaved target decode
+ } else {
+ res->reset();
+ ggml_backend_sched_reset(sched_nextn.get());
+ ggml_backend_sched_set_eval_callback(sched_nextn.get(), cparams.cb_eval, cparams.cb_eval_user_data);
+
+ ggml_cgraph * gf = model.build_graph(gparams);
+ if (!gf) {
+ LLAMA_LOG_ERROR("%s: failed to build fused NextN graph\n", __func__);
+ mtp_fused_steps = 1;
+ return -6;
+ }
+ if (!ggml_backend_sched_alloc_graph(sched_nextn.get(), gf)) {
+ LLAMA_LOG_ERROR("%s: failed to allocate fused NextN graph\n", __func__);
+ mtp_fused_steps = 1;
+ return -6;
+ }
+ }
+ res->set_inputs(&ub);
+
+ // Compute on the dedicated sched (mirror graph_compute_mtp's CPU-threadpool dance).
+ if (backend_cpu != nullptr) {
+ auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_cpu));
+ auto * set_tp_fn = (decltype(ggml_backend_cpu_set_threadpool) *)
+ ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_set_threadpool");
+ if (set_tp_fn) {
+ set_tp_fn(backend_cpu, threadpool);
+ }
+ }
+ for (const auto & set_n_threads_fn : set_n_threads_fns) {
+ set_n_threads_fn.second(set_n_threads_fn.first, cparams.n_threads);
+ }
+ status = ggml_backend_sched_graph_compute_async(sched_nextn.get(), res->get_gf());
+ mtp_fused_steps = 1;
+ if (status != GGML_STATUS_SUCCESS) {
+ LLAMA_LOG_ERROR("%s: fused NextN graph compute failed (status %d)\n", __func__, (int) status);
+ return -6;
+ }
+
+ ggml_backend_sched_synchronize(sched_nextn.get());
+
+ ggml_tensor * t_arg = res->get_argmax();
+ GGML_ASSERT(t_arg && "fused NextN graph must publish the in-graph argmax block");
+ GGML_ASSERT(t_arg->ne[0] == (int64_t) n_steps && "fused NextN argmax must be I32[n_steps]");
+
+ std::vector<int32_t> drafts((size_t) n_steps);
+ ggml_backend_tensor_get(t_arg, drafts.data(), 0, (size_t) n_steps * sizeof(int32_t));
+ for (int32_t k = 0; k < n_steps; ++k) {
+ out_drafts[k] = (llama_token) drafts[(size_t) k];
+ }
+
+ if (out_h_prev_last) {
+ ggml_tensor * t_post = res->get_embd();
+ GGML_ASSERT(t_post);
+ ggml_backend_tensor_get(t_post, out_h_prev_last, 0, n_embd * sizeof(float));
+ }
+
+ return 0;
+}
+
// opencoti F5 M6-S4 mtp (P4): submit one async MTP draft request to the worker. At most one
// in-flight request per context (returns -7 if a prior request was not yet waited).
int32_t llama_context::decode_mtp_async(
@@ -3255,6 +3466,10 @@ llm_graph_params llama_context::graph_params(
/*.n_outputs =*/ n_outputs,
/*.cb =*/ graph_get_cb(),
/*.res =*/ res,
+ // opencoti fused-NextN (#590): the DECODER_MTP draft graph fuses N greedy steps when the
+ // fused driver sets mtp_fused_steps=N (reset to 1 after). Every other gtype — and the
+ // per-step AR NextN loop — keeps n_mtp_steps=1 (single-step build_one_step).
+ /*.n_mtp_steps =*/ (gtype == LLM_GRAPH_TYPE_DECODER_MTP ? mtp_fused_steps : 1),
};
}
@@ -4330,6 +4545,7 @@ llama_context_params llama_context_default_params() {
/*.type_v =*/ GGML_TYPE_F16,
/*.abort_callback =*/ nullptr,
/*.abort_callback_data =*/ nullptr,
+ /*.ctx_other =*/ nullptr, // opencoti bug-858: dual-context MTP
/*.embeddings =*/ false,
/*.offload_kqv =*/ true,
/*.no_perf =*/ true,
@@ -4415,8 +4631,13 @@ llama_context * llama_init_from_model(
model->hparams.pooling_type, params.pooling_type);
}
+ // opencoti bug-858: a dual-context gemma4-assistant drafter (ctx_other set) legitimately has
+ // nextn_predict_layers == 0 — its "MTP layers" are the TARGET's backbone K/V, shared in-place
+ // via ctx_other (a KV-less assistant reads the target cells). Only reject a STANDALONE MTP
+ // context (ctx_other == nullptr) that lacks nextn layers; Qwen NextN (layers > 0) is unaffected.
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
- model->hparams.nextn_predict_layers == 0) {
+ model->hparams.nextn_predict_layers == 0 &&
+ params.ctx_other == nullptr) {
LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__);
return nullptr;
}
@@ -4572,6 +4793,21 @@ float * llama_get_embeddings_pre_norm_ith(llama_context * ctx, int32_t i) {
return ctx->get_embeddings_pre_norm_ith(i);
}
+// opencoti #590/bug-858: return the on-device backend-sampled draft token (sampling.sampled), produced
+// during the draft decode by the offloaded sampler chain attached via llama_set_sampler (top_k-only →
+// greedy argmax, matching common_sampler_sample). This is the cheap read that replaces the ~152k-vocab
+// common_sampler_sample D2H + CPU softmax/sort (MEASURED 0.98 ms/step). Returns LLAMA_TOKEN_NULL (-1)
+// when backend sampling is inactive → the AR draft loop falls back to the exact CPU sampler.
+llama_token llama_context::get_draft_greedy_ith(int32_t i) {
+ return get_sampled_token_ith(i);
+}
+
+llama_token llama_get_draft_greedy_ith(llama_context * ctx, int32_t i) {
+ ctx->synchronize();
+
+ return ctx->get_draft_greedy_ith(i);
+}
+
bool llama_set_sampler(llama_context * ctx, llama_seq_id seq_id, llama_sampler * smpl) {
return ctx->set_sampler(seq_id, smpl);
}
@@ -4667,6 +4903,10 @@ llama_memory_t llama_get_memory(const struct llama_context * ctx) {
return ctx->get_memory();
}
+llama_context * llama_get_ctx_other(struct llama_context * ctx) {
+ return ctx->get_cparams().ctx_other; // opencoti bug-858: dual-context MTP
+}
+
void llama_memory_clear(llama_memory_t mem, bool data) {
if (!mem) {
return;
@@ -4955,6 +5195,22 @@ int32_t llama_decode_mtp(
return ctx->decode_mtp(seq_id, attn_pos, last_token, h_prev, n_steps, out_drafts, out_logits, out_h_prev_last);
}
+int32_t llama_decode_mtp_fused_nextn(
+ llama_context * ctx,
+ llama_seq_id seq_id,
+ llama_pos attn_pos,
+ llama_token last_token,
+ const float * h_prev,
+ int32_t n_steps,
+ llama_token * out_drafts,
+ float * out_h_prev_last) {
+ if (!ctx) {
+ LLAMA_LOG_ERROR("%s: ctx is NULL\n", __func__);
+ return -1;
+ }
+ return ctx->decode_mtp_fused_nextn(seq_id, attn_pos, last_token, h_prev, n_steps, out_drafts, out_h_prev_last);
+}
+
//
// perf
//
diff --git a/llama.cpp/src/llama-context.h b/llama.cpp/src/llama-context.h
index 7f85a13..b6b0efb 100644
--- a/llama.cpp/src/llama-context.h
+++ b/llama.cpp/src/llama-context.h
@@ -94,6 +94,11 @@ struct llama_context {
float * get_embeddings_pre_norm();
float * get_embeddings_pre_norm_ith(int32_t i);
+ // opencoti #590/bug-858: on-device greedy draft token from the previous MTP decode's t_argmax
+ // (returns -1 if the graph published none). Lets the AR spec draft loop skip the full-vocab
+ // common_sampler_sample host round-trip (~0.98 ms/step) when the draft sampler is argmax-greedy.
+ llama_token get_draft_greedy_ith(int32_t i);
+
llama_token * get_sampled_tokens() const;
llama_token get_sampled_token_ith(int32_t idx);
@@ -182,6 +187,17 @@ struct llama_context {
llama_token * out_drafts,
float * out_h_prev_last);
+ // opencoti fused-NextN (#590): fused N-step greedy draft on a DECODER_MTP (Qwen NextN) draft
+ // context using the unified cache. Twin of decode_mtp_fused minus mtp_assistant/sched_mtp/iSWA.
+ int32_t decode_mtp_fused_nextn(
+ llama_seq_id seq_id,
+ llama_pos attn_pos,
+ llama_token last_token,
+ const float * h_prev,
+ int32_t n_steps,
+ llama_token * out_drafts,
+ float * out_h_prev_last);
+
// opencoti F5 M6-S4 mtp (P4): async MTP draft pipeline. decode_mtp() (above) is a sync
// facade — out_logits!=NULL keeps the in-thread decode_mtp_sync path (the worker streams
// no per-step logits); otherwise it submits via decode_mtp_async and blocks in
@@ -430,6 +446,19 @@ private:
ggml_backend_sched_ptr sched_mtp;
llm_graph_result_ptr gf_res_prev_mtp;
+ // opencoti #590/bug-858: DEDICATED scheduler + result for the fused NextN self-spec draft
+ // (decode_mtp_fused_nextn). It must NOT share the main sched/gf_res_prev with target decode:
+ // every interleaved target decode RESETS the main sched (reallocating compute buffers) and
+ // gf_res_prev, so a fused graph parked there can never satisfy can_reuse AND its tensor memory is
+ // clobbered → 100% rebuild + CUDA-recapture per draft round (measured: build 0.5ms + comp-launch
+ // 1.5ms, rebuilds=N/N) AND garbage argmax when reused. A dedicated sched (like the Gemma assistant's
+ // sched_mtp, but ensure_sched_mtp is assistant-only: needs mtp_assistant + iswa cache) isolates the
+ // fused graph so can_reuse compares fused-vs-fused (stable within a 256-cell n_kv bucket) and the
+ // ggml graph + CUDA capture persist across rounds. NextN and Gemma-assistant are mutually exclusive.
+ ggml_backend_sched_ptr sched_nextn;
+ llm_graph_result_ptr gf_res_prev_nextn;
+ int32_t nextn_reserved_steps = 0;
+
// opencoti F5 M6-S4 mtp (P4 fused): number of draft steps to unroll into the MTP graph for
// the NEXT build (decode_mtp_fused sets it to n_steps; the sequential/reserve paths leave it
// at 1). mtp_reserved_steps records how many steps the current sched_mtp reserve covers, so
diff --git a/llama.cpp/src/llama-cparams.h b/llama.cpp/src/llama-cparams.h
index 934adf5..8f0b36b 100644
--- a/llama.cpp/src/llama-cparams.h
+++ b/llama.cpp/src/llama-cparams.h
@@ -112,4 +112,9 @@ struct llama_cparams {
ggml_backend_sched_eval_callback cb_eval;
void * cb_eval_user_data;
+
+ // opencoti bug-858: dual-context MTP — the target context whose KV cells the
+ // draft (MTP) context shares IN-PLACE (upstream cparams.ctx_other + mem_other).
+ // nullptr for ordinary contexts.
+ llama_context * ctx_other;
};
diff --git a/llama.cpp/src/llama-ext.h b/llama.cpp/src/llama-ext.h
index edfa71c..9d57c33 100644
--- a/llama.cpp/src/llama-ext.h
+++ b/llama.cpp/src/llama-ext.h
@@ -104,3 +104,8 @@ LLAMA_API float * llama_get_embeddings_pre_norm (struct llama_context * ctx);
// LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i);
LLAMA_API float * llama_get_embeddings_pre_norm_ith(struct llama_context * ctx, int32_t i);
+
+// opencoti #590/bug-858: on-device greedy draft token from the previous MTP decode's in-graph argmax
+// (t_argmax). Returns the token id, or -1 if none was published. Lets the AR spec draft loop skip the
+// ~0.98 ms/step full-vocab common_sampler_sample (logit D2H + CPU softmax/sort) when p_min==0.
+LLAMA_API llama_token llama_get_draft_greedy_ith(struct llama_context * ctx, int32_t i);
diff --git a/llama.cpp/src/llama-graph.cpp b/llama.cpp/src/llama-graph.cpp
index e0d5e2a..e0d712e 100644
--- a/llama.cpp/src/llama-graph.cpp
+++ b/llama.cpp/src/llama-graph.cpp
@@ -125,14 +125,28 @@ void llm_graph_input_mtp::set_input(const llama_ubatch * ubatch) {
// opencoti F5 M6-S4 mtp (P4 fused): step k's RoPE query position = ubatch->pos[k]
// (= attn_pos + 1 + k, filled by decode_mtp_fused). Single-step path leaves this empty.
+ // bug-870: mrope archs (qwen35) size each step's pos tensor I32[4] — fill the M-RoPE
+ // text-token layout [p,p,p,0] (mirrors llm_graph_input_pos::set_input); standard rope = [p].
for (size_t k = 0; k < inp_pos_steps.size(); ++k) {
- ggml_backend_tensor_set(inp_pos_steps[k], ubatch->pos + k, 0, sizeof(int32_t));
+ ggml_tensor * pk = inp_pos_steps[k];
+ const int32_t p = ubatch->pos[k];
+ if (pk->ne[0] == 4) {
+ const int32_t pos4[4] = { p, p, p, 0 };
+ ggml_backend_tensor_set(pk, pos4, 0, sizeof(pos4));
+ continue;
+ }
+ ggml_backend_tensor_set(pk, &p, 0, sizeof(int32_t));
}
}
// opencoti F5 M6-S4 mtp
bool llm_graph_input_mtp::can_reuse(const llm_graph_params & params) {
- if (params.gtype != LLM_GRAPH_TYPE_MTP) {
+ // bug-870/#590: the fused NextN driver (decode_mtp_fused_nextn) builds a DECODER_MTP graph, so
+ // gating reuse on MTP-only forced a full N-block MoE graph rebuild+realloc on EVERY draft call —
+ // the per-call build dwarfed the saved host-syncs and made fused SLOWER than the AR loop. Accept
+ // both MTP (Gemma assistant) and DECODER_MTP (Qwen NextN fused); the attn input's own can_reuse
+ // independently vetoes reuse when the prefix mask shape changes, so this stays correct.
+ if (params.gtype != LLM_GRAPH_TYPE_MTP && params.gtype != LLM_GRAPH_TYPE_DECODER_MTP) {
return false;
}
const bool base = inp_last_token && inp_last_token->ne[0] == 1 && inp_h_prev && inp_h_prev->ne[1] == 1;
@@ -607,11 +621,20 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
bool res = true;
+ // opencoti bug-858 (2026-07-05): match upstream b9859 — the kq_mask checks must be
+ // INDEPENDENT of the k_idxs blocks. In the KV-less gemma4-assistant draft context the
+ // k/v_idxs are dead nodes (no cache store) and never get buffers, so nesting the mask
+ // check under them skipped it entirely → the boot-time reserve graph (n_kv = full cache
+ // width) was reused for every draft decode → set_input filled only the first n_kv
+ // columns of a full-width mask, leaving stale pool garbage → accept collapse on any
+ // prompt past one KV pad block (gate10: 0.83@114tok → 0.37@354 → 0.00@1314).
// base tensors may not be allocated if there are no non-SWA attention layers
if (self_k_idxs && self_k_idxs->buffer) {
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
//res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
+ }
+ if (self_kq_mask && self_kq_mask->buffer) {
res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
}
@@ -619,7 +642,9 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
//res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
+ }
+ if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
}
@@ -2894,7 +2919,32 @@ ggml_tensor * llm_graph_context::build_attn(
// null (build_attn for KV asserts this above at :2102), kq_b too;
// the M2 gate ensures non-transposed V. See docs/features/advanced_kv.md.
ggml_tensor * cur;
- if (mctx_cur->get_layer_tactic(il) == llama_kv_cache::headinfer_tactic::GPU_STREAM) {
+ // opencoti-hook: f5-rolling-kv — bug-1840 (#588). POSITION_WINDOW on the FULL-ATTENTION
+ // build_attn overload. The position-window residency tactic was previously wired ONLY into
+ // the iSWA overload (build_attn(llm_graph_input_attn_kv_iswa*), see the mirror ~line 3272);
+ // full-attention models (Qwen etc. — the ONLY class that actually spills, since iSWA KV is
+ // bounded) fell through to the get_k/get_v else-branch below, whose window fallback returns
+ // ggml_concat(window_device, tail_host). That single cross-backend logical tensor makes the
+ // ggml scheduler round-trip the ENTIRE KV cache host<->device EVERY decode step (measured
+ // 12.6 GiB/step both ways -> 0.99 tps at a 4.2 GB tail). Routing spilling layers through
+ // build_attn_mha_position_window keeps the window FA device-resident and streams ONLY the
+ // tail (one H2D copy, no D2H write-back). Guarded to POSITION_WINDOW + FA + no KQ-bias/MLA/
+ // sinks + an active window, so non-spilling layers are byte-identical. See bug-1840,
+ // docs/features/rolling_kv.md.
+ const bool window_layer =
+ cparams.flash_attn && kq_b == nullptr && v_mla == nullptr && sinks == nullptr &&
+ mctx_cur->get_layer_tactic(il) == llama_kv_cache::headinfer_tactic::POSITION_WINDOW &&
+ mctx_cur->headinfer_window_active(il);
+ if (window_layer) {
+ ggml_tensor * k_win = mctx_cur->get_k_window(ctx0, il);
+ ggml_tensor * v_win = mctx_cur->get_v_window(ctx0, il);
+ ggml_tensor * k_tail = mctx_cur->get_k_tail (ctx0, il); // nullptr when n_kv <= wc
+ ggml_tensor * v_tail = mctx_cur->get_v_tail (ctx0, il);
+ const bool cpu_fa_tail =
+ (k_tail != nullptr) && (q->ne[2] == 1) && mctx_cur->headinfer_tail_is_cpu_fa();
+ cur = build_attn_mha_position_window(q, k_win, v_win, k_tail, v_tail,
+ kq_b, kq_mask, sinks, v_mla, kq_scale, il, cpu_fa_tail);
+ } else if (mctx_cur->get_layer_tactic(il) == llama_kv_cache::headinfer_tactic::GPU_STREAM) {
// opencoti F5 M7 Rolling KV — Rung 2 (W4 M7-D #304), S1. GPU_STREAM
// rides the M2 head-split: the device head-group runs a resident FA,
// the host (spilled) head-group runs the streaming-FA op, and the two
@@ -3036,6 +3086,28 @@ ggml_tensor * llm_graph_context::build_attn_mtp(
return cur;
}
+// opencoti fused-NextN (Option A / Gemma-mirror): read-only cross-attention into the frozen prefix
+// KV on the UNIFIED cache. Twin of build_attn_mtp for the non-iSWA llm_graph_input_attn_kv used by
+// Qwen NextN drafters. Reads mctx->get_k/get_v(il) with the input's kq_mask — which exposes the
+// full prefix 0..pmax because get_n_kv reports the cache's used extent, not the mtp_slot_info cell
+// count. Writes NO draft KV; recurrence is carried by the hidden-state chain. wo is applied by the
+// caller (the qwen35* MTP block multiplies by the gate then projects wo). See
+// docs/features/fused_nextn_mtp.md delta #6.
+ggml_tensor * llm_graph_context::build_attn_readonly_nextn(
+ llm_graph_input_attn_kv * inp,
+ ggml_tensor * q_cur,
+ float kq_scale,
+ int il) const {
+ const auto * mctx_cur = inp->mctx;
+
+ ggml_tensor * kq_mask = inp->get_kq_mask();
+
+ ggml_tensor * k = mctx_cur->get_k(ctx0, il);
+ ggml_tensor * v = mctx_cur->get_v(ctx0, il);
+
+ return build_attn_mha(q_cur, k, v, nullptr, kq_mask, nullptr, nullptr, kq_scale, il);
+}
+
static std::unique_ptr<llm_graph_input_attn_k> build_attn_inp_k_impl(
ggml_context * ctx0,
const llama_ubatch & ubatch,
diff --git a/llama.cpp/src/llama-graph.h b/llama.cpp/src/llama-graph.h
index 0c13245..d3a3499 100644
--- a/llama.cpp/src/llama-graph.h
+++ b/llama.cpp/src/llama-graph.h
@@ -1041,6 +1041,17 @@ struct llm_graph_context {
int64_t kv_n_head_v,
bool use_k_as_v) const;
+ // opencoti fused-NextN (Option A): read-only cross-attention into the frozen prefix KV on the
+ // UNIFIED cache. Twin of build_attn_mtp for llm_graph_input_attn_kv (non-iSWA). Reads
+ // mctx->get_k/get_v(il) with the input's kq_mask (which exposes 0..pmax via get_n_kv's used
+ // extent), runs build_attn_mha, writes NO draft KV. wo is applied by the caller.
+ // See docs/features/fused_nextn_mtp.md.
+ ggml_tensor * build_attn_readonly_nextn(
+ llm_graph_input_attn_kv * inp,
+ ggml_tensor * q_cur,
+ float kq_scale,
+ int il) const;
+
llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const;
ggml_tensor * build_attn(
diff --git a/llama.cpp/src/llama-kv-cache-iswa.cpp b/llama.cpp/src/llama-kv-cache-iswa.cpp
index 7611271..9628df1 100644
--- a/llama.cpp/src/llama-kv-cache-iswa.cpp
+++ b/llama.cpp/src/llama-kv-cache-iswa.cpp
@@ -33,8 +33,17 @@ llama_kv_cache_iswa::llama_kv_cache_iswa(
uint32_t n_pad,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse,
+ llama_memory_t mem_other,
+ const layer_share_cb & share,
uint32_t sparse_attn_block_size) : hparams(model.hparams), unified(unified) {
+ // opencoti bug-858 dual-context MTP — split the source iSWA cache into its base/swa halves
+ // so each inner kv_cache shares with the matching half. GUARD: both stay nullptr when
+ // mem_other == nullptr (every existing caller) ⇒ the inner ctors receive nullptr ⇒
+ // byte-identical standalone allocation.
+ llama_memory_t mem_other_base = mem_other ? static_cast<llama_kv_cache_iswa *>(mem_other)->get_base() : nullptr;
+ llama_memory_t mem_other_swa = mem_other ? static_cast<llama_kv_cache_iswa *>(mem_other)->get_swa() : nullptr;
+
// chain filters
const layer_filter_cb filter_base = [&](int32_t il) {
if (filter && !filter(il)) {
@@ -73,7 +82,9 @@ llama_kv_cache_iswa::llama_kv_cache_iswa(
v_trans, offload, unified, size_base, kv_size_initial, slot_shrink_idle_ms,
headinfer_gpu_heads_frac, vram_target_mib, pcie_bw_gbps, kv_residency_mode,
compute_reserve_mib, window_measure_pass,
- n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse, sparse_attn_block_size);
+ n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse,
+ // opencoti bug-858 dual-context MTP — forward the base-half source + share selector.
+ mem_other_base, share, sparse_attn_block_size);
LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa);
@@ -82,7 +93,9 @@ llama_kv_cache_iswa::llama_kv_cache_iswa(
v_trans, offload, unified, size_swa, kv_size_initial, slot_shrink_idle_ms,
headinfer_gpu_heads_frac, vram_target_mib, pcie_bw_gbps, kv_residency_mode,
compute_reserve_mib, window_measure_pass,
- n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, filter_swa, reuse, sparse_attn_block_size);
+ n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, filter_swa, reuse,
+ // opencoti bug-858 dual-context MTP — forward the swa-half source + share selector.
+ mem_other_swa, share, sparse_attn_block_size);
}
void llama_kv_cache_iswa::clear(bool data) {
diff --git a/llama.cpp/src/llama-kv-cache-iswa.h b/llama.cpp/src/llama-kv-cache-iswa.h
index 166b4eb..06fd1ae 100644
--- a/llama.cpp/src/llama-kv-cache-iswa.h
+++ b/llama.cpp/src/llama-kv-cache-iswa.h
@@ -52,6 +52,13 @@ public:
uint32_t n_pad,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse,
+ // opencoti bug-858 dual-context MTP — source iSWA cache; split into
+ // base/swa and forwarded to the two inner kv_base/kv_swa. nullptr =>
+ // ordinary standalone cache (byte-identical). Threaded with `share`.
+ llama_memory_t mem_other,
+ // opencoti bug-858 dual-context MTP — per-layer share selector,
+ // forwarded verbatim to both inner caches. Used only when mem_other != null.
+ const layer_share_cb & share,
// opencoti #551 sparse-attn — forwarded to both kv_base and
// kv_swa. Quest block-selector B_SEL; 0 = off (no side-cache).
// See docs/features/sparse_attn.md.
diff --git a/llama.cpp/src/llama-kv-cache.cpp b/llama.cpp/src/llama-kv-cache.cpp
index 13c8c67..4416d38 100644
--- a/llama.cpp/src/llama-kv-cache.cpp
+++ b/llama.cpp/src/llama-kv-cache.cpp
@@ -15,6 +15,18 @@
#include <map>
#include <stdexcept>
+// opencoti bug-858 diagnostic (REMOVE after root-cause): checked map_layer_ids.at() — logs the
+// missing layer index + enclosing function + map size on a miss, instead of an opaque throw.
+template <typename M>
+static inline int32_t mli_at_checked(const M & m, int32_t il, const char * fn) {
+ auto it = m.find(il);
+ if (it == m.end()) {
+ LLAMA_LOG_ERROR("%s: map_layer_ids MISS il=%d (n_keys=%zu)\n", fn, il, m.size());
+ throw std::out_of_range("map_layer_ids");
+ }
+ return it->second;
+}
+
// opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md
// On Linux, madvise(addr, len, MADV_DONTNEED) is the call that actually
// decommits pages: the kernel zeroes the page table entries, releases the
@@ -184,9 +196,17 @@ static float opencoti_compute_auto_gpu_heads_frac(
size_t free_vram = 0, total_vram = 0;
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
+ // opencoti bug-1839: --vram-target is a TOTAL-device VRAM cap (rolling_kv.md
+ // Decision 4: "leave headroom for siblings"), NOT a KV-only budget. Weights and
+ // the CUDA context are already resident here, so subtract what the device already
+ // holds (total - free) from the target to get the KV share. free_vram still caps
+ // (we can never make resident more than is actually free right now).
size_t budget = free_vram;
if (vram_target_mib > 0) {
- budget = std::min(budget, (size_t) vram_target_mib * 1024 * 1024);
+ const size_t vt_bytes = (size_t) vram_target_mib * 1024 * 1024;
+ const size_t already_used = total_vram > free_vram ? total_vram - free_vram : 0;
+ const size_t vt_kv_share = vt_bytes > already_used ? vt_bytes - already_used : 0;
+ budget = std::min(budget, vt_kv_share);
}
// opencoti bug-1342: hold back the MEASURED compute buffer (two-pass) instead of the
// fixed 1536 MiB guess when the context supplied it; 0 → fall back to the default.
@@ -277,16 +297,24 @@ static uint32_t opencoti_compute_resident_window_cells(
size_t free_vram = 0, total_vram = 0;
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
+ // opencoti bug-1839: --vram-target is a TOTAL-device VRAM cap (rolling_kv.md
+ // Decision 4: "leave headroom for siblings"), NOT a KV-only budget. Weights and
+ // the CUDA context are already resident here, so subtract what the device already
+ // holds (total - free) from the target to get the KV share. free_vram still caps
+ // (we can never make resident more than is actually free right now).
size_t budget = free_vram;
if (vram_target_mib > 0) {
- budget = std::min(budget, (size_t) vram_target_mib * 1024 * 1024);
+ const size_t vt_bytes = (size_t) vram_target_mib * 1024 * 1024;
+ const size_t already_used = total_vram > free_vram ? total_vram - free_vram : 0;
+ const size_t vt_kv_share = vt_bytes > already_used ? vt_bytes - already_used : 0;
+ budget = std::min(budget, vt_kv_share);
}
budget = budget > OPENCOTI_HEADINFER_AUTO_RESERVE_BYTES
? budget - OPENCOTI_HEADINFER_AUTO_RESERVE_BYTES
: 0;
if (per_cell * (size_t) kv_size <= budget) {
- LLAMA_LOG_INFO("%s: position window = FULLY RESIDENT (KV %.0f MiB fits budget "
+ LLAMA_LOG_WARN("%s: position window = FULLY RESIDENT (KV %.0f MiB fits budget "
"%.0f MiB of %.0f MiB free); no host tail\n", __func__,
per_cell * (size_t) kv_size / 1048576.0, budget / 1048576.0, free_vram / 1048576.0);
return kv_size;
@@ -320,9 +348,10 @@ static uint32_t opencoti_compute_resident_window_cells(
cells = walign;
}
}
- LLAMA_LOG_INFO("%s: position window = %u / %u cells resident (%.0f MiB budget, "
- "%.0f MiB/cell-set); host tail = %u cells\n", __func__,
- cells, kv_size, budget / 1048576.0, per_cell / 1048576.0, kv_size - cells);
+ LLAMA_LOG_WARN("%s: position window = %u / %u cells resident (%.0f MiB budget, "
+ "%.0f MiB/cell-set); host tail = %u cells = %.0f MiB\n", __func__,
+ cells, kv_size, budget / 1048576.0, per_cell / 1048576.0, kv_size - cells,
+ per_cell * (size_t) (kv_size - cells) / 1048576.0);
return cells;
}
@@ -348,9 +377,32 @@ llama_kv_cache::llama_kv_cache(
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse,
+ llama_memory_t mem_other,
+ const layer_share_cb & share,
uint32_t sparse_attn_block_size) :
model(model), hparams(model.hparams), v_trans(v_trans),
- n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type) {
+ n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type),
+ // opencoti bug-858 dual-context MTP — `other` is the TARGET cache (nullptr for every
+ // existing caller). The cells vector is shared behind a shared_ptr: a draft cache reuses
+ // the target's, a standalone cache makes its own. Member-init order is other → v_cells_impl
+ // → v_cells (matches the header declaration order), so v_cells_impl safely reads `other`.
+ other(static_cast<llama_kv_cache *>(mem_other)),
+ v_cells_impl(other ? other->v_cells_impl : std::make_shared<std::vector<llama_kv_cells>>()),
+ v_cells(*v_cells_impl) {
+
+ // opencoti bug-858 dual-context MTP — a draft cache (other != nullptr) shares the TARGET's
+ // cells + K/V tensors in place, so its cell count MUST follow the target's allocation (an
+ // oversized view would overflow the source tensors). Mirrors upstream b9859 get_size().
+ // GUARD: skipped entirely when other == nullptr (every existing caller) ⇒ kv_size is
+ // untouched ⇒ byte-identical to the shipped path.
+ if (other) {
+ const uint32_t size_other = other->get_size();
+ if (kv_size != size_other) {
+ LLAMA_LOG_WARN("%s: kv_size = %u overridden to %u to match the shared source cache\n",
+ __func__, kv_size, size_other);
+ kv_size = size_other;
+ }
+ }
GGML_ASSERT(kv_size % n_pad == 0);
@@ -579,6 +631,37 @@ llama_kv_cache::llama_kv_cache(
continue;
}
+ // opencoti bug-858 dual-context MTP — share this draft layer's cells + K/V with a
+ // TARGET-cache layer IN PLACE (a KV-less drafter reads the backbone K/V) instead of
+ // allocating its own tensors. Copies the target kv_layer struct (tensor POINTERS +
+ // residency / block-sel metadata; the draft only READS them), retags il to this
+ // cache's model-layer index, and skips allocation (continue). GUARD: `share && other`
+ // — when either is null (every existing caller passes both null) this block is inert,
+ // control falls straight through to the standalone allocation ⇒ byte-identical path.
+ if (share && other) {
+ const int32_t il_share = share(il);
+
+ if (il_share >= 0) {
+ // opencoti bug-858: match upstream b9859:194 (operator[], NOT .at()) — a missing key
+ // must NOT throw. Diagnose-log if il_share is absent from the target's map (that would
+ // mean our target iSWA inner-cache layer partition differs from upstream's, and []=>0).
+ if (other->map_layer_ids.find(il_share) == other->map_layer_ids.end()) {
+ LLAMA_LOG_WARN("%s: layer %3d: share target %d ABSENT from other map (%zu keys) — []=>0 fallback\n",
+ __func__, il, il_share, other->map_layer_ids.size());
+ }
+ const auto & layer_share = other->layers[other->map_layer_ids[il_share]];
+
+ LLAMA_LOG_WARN("%s: layer %3d: sharing with target layer %d\n", __func__, il, il_share);
+
+ map_layer_ids[il] = layers.size();
+
+ layers.push_back(layer_share);
+ layers.back().il = il;
+
+ continue;
+ }
+ }
+
if (n_embd_head_k_all == 0) {
n_embd_head_k_all = (int32_t) hparams.n_embd_head_k(il);
} else if (n_embd_head_k_all > 0 && n_embd_head_k_all != (int32_t) hparams.n_embd_head_k(il)) {
@@ -754,6 +837,21 @@ llama_kv_cache::llama_kv_cache(
if (!ctx_cpu_s) {
throw std::runtime_error("failed to create CPU ggml context for headinfer kv cache");
}
+ // opencoti bug-1342 diag: the window-tail host K/V is streamed device-ward
+ // per decode tile via strided cudaMemcpy2DAsync. PINNED host memory makes
+ // that async + full-PCIe; a pageable 'CPU' buft makes it SYNCHRONOUS and
+ // slow (per-row). Log the resolved buft ONCE so a boot log settles
+ // pinned-vs-pageable without a debugger (WARN so -lv 3 keeps it).
+ if (layer_window) {
+ static bool opencoti_tail_buft_logged = false;
+ if (!opencoti_tail_buft_logged) {
+ opencoti_tail_buft_logged = true;
+ LLAMA_LOG_WARN("%s: window-tail host buffer buft = '%s' (pinned host = "
+ "async full-PCIe DMA; plain 'CPU' = pageable → SYNCHRONOUS slow "
+ "strided H2D per decode tile)\n", __func__,
+ ggml_backend_buft_name(cpu_buft));
+ }
+ }
// window mode: ALL heads, tail_c cells; M2: CPU head subset, kv_size cells.
const uint32_t host_k_embd = layer_window ? n_embd_k_gqa : n_embd_k_cpu;
const uint32_t host_v_embd = layer_window ? n_embd_v_gqa : n_embd_v_cpu;
@@ -1014,6 +1112,14 @@ llama_kv_cache::llama_kv_cache(
// opencoti F4 M3 Phase 1+2 — see docs/decisions/0001-lazy-slot-context.md
void llama_kv_cache::ensure_cleared(uint32_t up_to_cells) {
+ // opencoti bug-858 (bug-2102): a shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // ALIASES the target's k/v tensors (A3 layer_share) and starts with n_cells_cleared == 0 —
+ // running the lazy clear here MEMSETS the target's LIVE KV (the drafter's first decode zeroed
+ // the head of the target's prompt keys → 4-17-logit trajectory flips scaling with prompt
+ // length). The shared cache is strictly READ-ONLY [TAG_KV_CACHE_SHARE_CELLS]: never clear.
+ if (other) {
+ return;
+ }
// up_to_cells == 0 means "clear up to the current soft cap". Callers
// that need full coverage (state_read) pass kv_size_max.
const uint32_t target = up_to_cells == 0
@@ -1189,6 +1295,11 @@ static void opencoti_decommit_layer_range(ggml_tensor * t,
}
void llama_kv_cache::shrink_if_idle() {
+ // opencoti bug-858 (bug-2102): shared draft cache (other!=null) aliases the TARGET's k/v
+ // tensors — shrinking here would mutate the target's live cache. READ-ONLY: never shrink.
+ if (other) {
+ return;
+ }
if (slot_shrink_idle_ms_ == 0) {
return; // shrink-on-idle disabled
}
@@ -1296,6 +1407,12 @@ void llama_kv_cache::clear(bool data) {
}
bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return true;
+ }
GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
if (p0 < 0) {
@@ -1490,6 +1607,12 @@ std::vector<float> llama_kv_cache::seq_key_scores(llama_seq_id seq_id, int32_t l
}
void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());
GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());
@@ -1577,6 +1700,12 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll
}
void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
auto & cells = v_cells[seq_to_stream[seq_id]];
@@ -1599,6 +1728,12 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
}
void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1");
@@ -1644,6 +1779,12 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll
}
void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1");
@@ -1678,6 +1819,12 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in
}
llama_pos llama_kv_cache::seq_pos_min(llama_seq_id seq_id) const {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return other->seq_pos_min(seq_id);
+ }
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
const auto & cells = v_cells[seq_to_stream[seq_id]];
@@ -1686,6 +1833,12 @@ llama_pos llama_kv_cache::seq_pos_min(llama_seq_id seq_id) const {
}
llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return other->seq_pos_max(seq_id);
+ }
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
const auto & cells = v_cells[seq_to_stream[seq_id]];
@@ -1865,6 +2018,12 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
}
bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return true;
+ }
bool updated = false;
auto * sched = lctx->get_sched();
@@ -2188,6 +2347,12 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
}
void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
// opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md
// Refresh the idle timestamp on every actual write. shrink_if_idle uses
// this to gate the shrink path.
@@ -2340,7 +2505,7 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
}
ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
@@ -2462,7 +2627,7 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k
}
ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
@@ -2923,7 +3088,7 @@ uint32_t llama_kv_cache::headinfer_gpu_heads(int32_t il) const {
}
ggml_tensor * llama_kv_cache::get_k_gpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.gpu_heads > 0 && "get_k_gpu called without an active headinfer split");
@@ -2959,7 +3124,7 @@ ggml_tensor * llama_kv_cache::get_k_gpu(ggml_context * ctx, int32_t il, uint32_t
// views it in-place. n_kv is irrelevant — kbounds spans all selectable blocks.
ggml_tensor * llama_kv_cache::get_kbounds(ggml_context * ctx, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(ctx);
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
if (layer.kbounds_per_stream.empty()) {
return nullptr;
@@ -2973,7 +3138,7 @@ ggml_tensor * llama_kv_cache::get_kbounds(ggml_context * ctx, int32_t il, const
// re-writes it only every cparams.sparse_attn_refresh decode tokens; FA-VEC reads it every step.
ggml_tensor * llama_kv_cache::get_block_sel(ggml_context * ctx, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(ctx);
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
if (layer.block_sel_per_stream.empty()) {
return nullptr;
@@ -2982,7 +3147,7 @@ ggml_tensor * llama_kv_cache::get_block_sel(ggml_context * ctx, int32_t il, cons
}
ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.gpu_heads > 0
&& !layer.k_cpu_per_stream.empty()
@@ -2998,6 +3163,18 @@ ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t
auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * {
ggml_tensor * kc_s = layer.k_cpu_per_stream[s_cache];
+ // opencoti #586 diag: report the ACTUAL allocated buffer buft (the boot WARN
+ // reports the *intended* cpu_buft, masking a silent cudaMallocHost fallback to
+ // pageable CPU — the H2D tail stream then runs at ~half PCIe / non-overlapping).
+ {
+ static bool opencoti_actual_tail_buft_logged = false;
+ if (!opencoti_actual_tail_buft_logged) {
+ opencoti_actual_tail_buft_logged = true;
+ LLAMA_LOG_WARN("%s: ACTUAL tail: data=%p buffer=%p buft='%s'\n", __func__,
+ (void *) kc_s->data, (void *) kc_s->buffer,
+ kc_s->buffer ? ggml_backend_buft_name(ggml_backend_buffer_get_type(kc_s->buffer)) : "NULL-BUFFER");
+ }
+ }
return ggml_view_4d(ctx, kc_s,
n_embd_head_k, n_head_cpu, n_kv, 1,
ggml_row_size(kc_s->type, n_embd_head_k),
@@ -3014,7 +3191,7 @@ ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t
}
ggml_tensor * llama_kv_cache::get_v_gpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.gpu_heads > 0 && "get_v_gpu called without an active headinfer split");
// M2 requires non-transposed V (FA on). The split-active branch in
@@ -3045,7 +3222,7 @@ ggml_tensor * llama_kv_cache::get_v_gpu(ggml_context * ctx, int32_t il, uint32_t
}
ggml_tensor * llama_kv_cache::get_v_cpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.gpu_heads > 0
&& !layer.v_cpu_per_stream.empty()
@@ -3085,7 +3262,7 @@ ggml_tensor * llama_kv_cache::get_v_cpu(ggml_context * ctx, int32_t il, uint32_t
// ne[1] == window_cells, so its dim-3 (stream/position) stride is computed off
// k_s->ne[1] (the audit point — NOT kv_size); the host tail likewise off kc_s->ne[1].
ggml_tensor * llama_kv_cache::get_k_window(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.window_cells > 0 && "get_k_window called without an active position window");
@@ -3113,7 +3290,7 @@ ggml_tensor * llama_kv_cache::get_k_window(ggml_context * ctx, int32_t il, uint3
}
ggml_tensor * llama_kv_cache::get_k_tail(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.window_cells > 0
&& !layer.k_cpu_per_stream.empty()
@@ -3125,12 +3302,35 @@ ggml_tensor * llama_kv_cache::get_k_tail(ggml_context * ctx, int32_t il, uint32_
const uint64_t n_embd_k_gqa2 = (uint64_t) n_head_kv * n_embd_head_k;
const uint32_t wc = layer.window_cells;
const uint32_t n_tail = n_kv > wc ? n_kv - wc : 0;
+ {
+ static bool opencoti_gkt_entry_logged = false;
+ if (!opencoti_gkt_entry_logged) {
+ opencoti_gkt_entry_logged = true;
+ ggml_tensor * kc0 = layer.k_cpu_per_stream.empty() ? nullptr : layer.k_cpu_per_stream[0];
+ LLAMA_LOG_WARN("%s: ENTRY il=%d n_kv=%u wc=%u n_tail=%u kc0=%p kc0.buffer=%p buft=%s\n",
+ __func__, il, n_kv, wc, n_tail, (void *)(kc0 ? kc0->data : nullptr),
+ (void *)(kc0 ? kc0->buffer : nullptr),
+ (kc0 && kc0->buffer) ? ggml_backend_buft_name(ggml_backend_buffer_get_type(kc0->buffer)) : "NULL");
+ }
+ }
if (n_tail == 0) {
return nullptr; // no overflow → fully resident window, single-region FA
}
auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * {
ggml_tensor * kc_s = layer.k_cpu_per_stream[s_cache];
+ // opencoti #586 diag: report the ACTUAL allocated buffer buft (the boot WARN
+ // reports the *intended* cpu_buft, masking a silent cudaMallocHost fallback to
+ // pageable CPU — the H2D tail stream then runs at ~half PCIe / non-overlapping).
+ {
+ static bool opencoti_actual_tail_buft_logged = false;
+ if (!opencoti_actual_tail_buft_logged) {
+ opencoti_actual_tail_buft_logged = true;
+ LLAMA_LOG_WARN("%s: ACTUAL tail: data=%p buffer=%p buft='%s'\n", __func__,
+ (void *) kc_s->data, (void *) kc_s->buffer,
+ kc_s->buffer ? ggml_backend_buft_name(ggml_backend_buffer_get_type(kc_s->buffer)) : "NULL-BUFFER");
+ }
+ }
return ggml_view_4d(ctx, kc_s,
n_embd_head_k, n_head_kv, n_tail, 1,
ggml_row_size(kc_s->type, n_embd_head_k),
@@ -3147,7 +3347,7 @@ ggml_tensor * llama_kv_cache::get_k_tail(ggml_context * ctx, int32_t il, uint32_
}
ggml_tensor * llama_kv_cache::get_v_window(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.window_cells > 0 && "get_v_window called without an active position window");
GGML_ASSERT(!v_trans && "position window requires non-transposed V");
@@ -3176,7 +3376,7 @@ ggml_tensor * llama_kv_cache::get_v_window(ggml_context * ctx, int32_t il, uint3
}
ggml_tensor * llama_kv_cache::get_v_tail(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
GGML_ASSERT(layer.window_cells > 0
&& !layer.v_cpu_per_stream.empty()
@@ -3230,7 +3430,7 @@ uint32_t llama_kv_cache::headinfer_window_cells(int32_t il) const {
}
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
@@ -3413,7 +3613,7 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm
}
ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const {
- const int32_t ikv = map_layer_ids.at(il);
+ const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__);
const auto & layer = layers[ikv];
@@ -4365,6 +4565,10 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) {
}
ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_context * lctx) const {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ GGML_ASSERT(!other && "shift on a shared draft cache is a bug");
auto * ctx = res->get_ctx();
auto * gf = res->get_gf();
@@ -4467,6 +4671,12 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
}
void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
GGML_UNUSED(flags);
io.write(&n_stream, sizeof(n_stream));
@@ -4520,6 +4730,12 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
}
void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
+ // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null)
+ // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS].
+ // Inert on the target/standalone path (other==nullptr).
+ if (other) {
+ return;
+ }
GGML_UNUSED(flags);
GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
diff --git a/llama.cpp/src/llama-kv-cache.h b/llama.cpp/src/llama-kv-cache.h
index 924f7fe..7b3afc3 100644
--- a/llama.cpp/src/llama-kv-cache.h
+++ b/llama.cpp/src/llama-kv-cache.h
@@ -146,6 +146,14 @@ public:
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse,
+ // opencoti bug-858 dual-context MTP — source cache whose cells +
+ // per-layer K/V tensors the draft cache shares IN PLACE. nullptr =>
+ // ordinary standalone cache (byte-identical to pre-bug-858). Threaded
+ // with `share`; both are null for every target / existing caller.
+ llama_memory_t mem_other,
+ // opencoti bug-858 dual-context MTP — per-layer share selector; see
+ // llama_memory_i::layer_share_cb. Consulted only when mem_other != null.
+ const layer_share_cb & share,
// opencoti #551 sparse-attn — Quest block-selector KV-block
// size B_SEL (docs/features/sparse_attn.md). 0 = sparse-attn
// OFF: no kmin/kmax side-cache allocated (byte-identical).
@@ -548,7 +556,20 @@ private:
// note: this is not part of the KV state and it's only used to speed-up the find_slot() method
std::vector<uint32_t> v_heads;
- std::vector<llama_kv_cells> v_cells;
+ // opencoti bug-858 dual-context MTP — the TARGET cache this draft cache shares its
+ // cells + layer tensors with, or nullptr for an ordinary standalone cache. Set once
+ // at construction from the ctor's mem_other. Every existing caller passes nullptr, so
+ // `other == nullptr` is the shipped, byte-identical path. Declared BEFORE v_cells_impl
+ // so the init-list (which reads `other`) runs in the correct member order.
+ llama_kv_cache * other = nullptr;
+
+ // opencoti bug-858 dual-context MTP — the cells vector is held behind a shared_ptr so a
+ // draft cache (other != nullptr) can SHARE the target's cells in place. `v_cells` is a
+ // reference to *v_cells_impl, so EVERY existing v_cells[...] / v_cells.resize(...) access
+ // is unchanged. When other == nullptr this owns a freshly-made vector — identical layout
+ // and lifetime to the previous plain `std::vector<llama_kv_cells> v_cells;` member.
+ std::shared_ptr<std::vector<llama_kv_cells>> v_cells_impl;
+ std::vector<llama_kv_cells> & v_cells;
// maps from a sequence id to a stream id
std::vector<uint32_t> seq_to_stream;
diff --git a/llama.cpp/src/llama-memory-hybrid-iswa.cpp b/llama.cpp/src/llama-memory-hybrid-iswa.cpp
index 7301ad5..3cf8e4a 100644
--- a/llama.cpp/src/llama-memory-hybrid-iswa.cpp
+++ b/llama.cpp/src/llama-memory-hybrid-iswa.cpp
@@ -65,7 +65,10 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
filter_attn == nullptr ?
[&](int32_t il) { return !hparams.is_recurrent(il); }
: filter_attn,
- nullptr
+ nullptr,
+ // opencoti bug-858 dual-context MTP — hybrid-iswa attn cache is standalone: no sharing.
+ /* mem_other */ nullptr,
+ /* share */ nullptr
)),
mem_recr(new llama_memory_recurrent(
model,
diff --git a/llama.cpp/src/llama-memory-hybrid.cpp b/llama.cpp/src/llama-memory-hybrid.cpp
index 86e4134..0fa114d 100644
--- a/llama.cpp/src/llama-memory-hybrid.cpp
+++ b/llama.cpp/src/llama-memory-hybrid.cpp
@@ -65,7 +65,10 @@ llama_memory_hybrid::llama_memory_hybrid(
filter_attn == nullptr ?
[&](int32_t il) { return !hparams.is_recurrent(il); }
: filter_attn,
- nullptr
+ nullptr,
+ // opencoti bug-858 dual-context MTP — hybrid attn cache is standalone: no sharing.
+ /* mem_other */ nullptr,
+ /* share */ nullptr
)),
mem_recr(new llama_memory_recurrent(
model,
diff --git a/llama.cpp/src/llama-memory.h b/llama.cpp/src/llama-memory.h
index 3e79826..98e4f26 100644
--- a/llama.cpp/src/llama-memory.h
+++ b/llama.cpp/src/llama-memory.h
@@ -24,6 +24,10 @@ struct llama_memory_params {
bool swa_full;
llama_context_type ctx_type;
+
+ // opencoti bug-858: dual-context MTP — source memory whose cells the draft cache
+ // shares in-place (upstream mem_other). nullptr => ordinary standalone cache.
+ llama_memory_t mem_other;
};
enum llama_memory_status {
@@ -77,6 +81,13 @@ struct llama_memory_i {
// return negative value to indicate that the layer il should not reuse memory
using layer_reuse_cb = std::function<int32_t(int32_t il)>;
+ // opencoti bug-858 dual-context MTP — maps a draft-cache layer il to the TARGET
+ // cache's layer whose cells + K/V tensors it shares IN PLACE (a KV-less drafter reads
+ // the backbone K/V). Return negative to indicate layer il does NOT share (allocate
+ // normally). Consulted ONLY when the ctor is given a non-null mem_other; a nullptr
+ // share (every existing caller) means no sharing ⇒ byte-identical standalone cache.
+ using layer_share_cb = std::function<int32_t(int32_t il)>;
+
virtual ~llama_memory_i() = default;
// split the input batch into a set of ubatches and verify that they can fit into the cache
diff --git a/llama.cpp/src/llama-model.cpp b/llama.cpp/src/llama-model.cpp
index 3df76f7..d7804be 100644
--- a/llama.cpp/src/llama-model.cpp
+++ b/llama.cpp/src/llama-model.cpp
@@ -2079,6 +2079,24 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
GGML_ASSERT(hparams.is_swa_any());
+ // opencoti bug-858 dual-context MTP (A3) — when THIS model is the
+ // gemma4-assistant drafter created as ctx_dft (cparams.ctx_other = the
+ // target context), share the TARGET's KV cells IN PLACE: the assistant has
+ // no K/V projections, so each of its attention layers reads the target
+ // backbone's last full (n_layer-1) / SWA (n_layer-2) layer. Mirrors upstream
+ // b9859 create_memory LLM_ARCH_GEMMA4_ASSISTANT. Inert for every other arch
+ // and when ctx_other is null → nullptr/nullptr → byte-identical standalone cache.
+ llama_memory_t mtp_mem_other = nullptr;
+ llama_kv_cache::layer_share_cb mtp_share = nullptr;
+ if (arch == LLM_ARCH_GEMMA4_ASSISTANT && cparams.ctx_other != nullptr) {
+ mtp_mem_other = llama_get_memory(cparams.ctx_other);
+ mtp_share = [&](int32_t il) {
+ const llama_model * model_other = llama_get_model(cparams.ctx_other);
+ const int32_t n_layer_other = (int32_t) llama_model_n_layer(model_other);
+ return hparams.is_swa(il) ? (n_layer_other - 2) : (n_layer_other - 1);
+ };
+ }
+
res = new llama_kv_cache_iswa(
*this,
params.type_k,
@@ -2109,6 +2127,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
1,
filter,
reuse,
+ // opencoti bug-858 dual-context MTP (A3) — share the target's KV when
+ // this is the gemma4-assistant ctx_dft (mtp_mem_other/mtp_share set
+ // above); nullptr/nullptr for the target + every other arch → identical.
+ /* mem_other */ mtp_mem_other,
+ /* share */ mtp_share,
// opencoti #551 sparse-attn — B_SEL when enabled, else 0 (off)
cparams.sparse_attn_enabled ? cparams.sparse_attn_block_size : 0);
} else {
@@ -2144,6 +2167,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
hparams.swa_type,
filter,
nullptr,
+ // opencoti bug-858 dual-context MTP — target / standalone cache:
+ // no sharing in A2 (A3 wires the draft cache with a non-null source).
+ /* mem_other */ nullptr,
+ /* share */ nullptr,
// opencoti #551 sparse-attn — B_SEL when enabled, else 0 (off)
cparams.sparse_attn_enabled ? cparams.sparse_attn_block_size : 0);
}
diff --git a/llama.cpp/src/models/gemma4-assistant.cpp b/llama.cpp/src/models/gemma4-assistant.cpp
index 51f935c..cab4fa8 100644
--- a/llama.cpp/src/models/gemma4-assistant.cpp
+++ b/llama.cpp/src/models/gemma4-assistant.cpp
@@ -382,11 +382,17 @@ void llama_model_gemma4_assistant::load_arch_hparams(llama_model_loader & ml) {
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
- uint32_t n_kv_shared_layers = 0;
- ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false);
-
- hparams.n_layer_kv_from_start = hparams.n_layer - (int32_t) n_kv_shared_layers;
- hparams.f_attention_scale = 1.0f;
+ // opencoti bug-858 dual-context MTP — the gemma4-assistant is a KV-LESS drafter: its attention
+ // layers do NOT own K/V. When run as ctx_dft they alias the TARGET's backbone K/V in place via
+ // mem_other/share (llama-model.cpp create_memory A3 + llama-kv-cache.cpp share branch). So
+ // n_layer_kv_from_start MUST stay -1 (has_kv() default => true for every layer), matching
+ // upstream b9859 which does NOT set it for LLM_ARCH_GEMMA4_ASSISTANT. Deriving it from
+ // ATTENTION_SHARED_KV_LAYERS (= n_layer for the drafter) yields 0 → has_kv() false for all
+ // layers → the KV-cache ctor `continue`s past every layer BEFORE the share branch → empty
+ // map_layer_ids → "get_k: map_layer_ids MISS il=0 (n_keys=0)" boot crash. That GGUF key drives
+ // the TARGET's intra-model elastic KV reuse (#417), a different mechanism from the cross-model
+ // share, and is inert for the KV-less drafter.
+ hparams.f_attention_scale = 1.0f;
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
@@ -434,7 +440,16 @@ void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) {
// opencoti F5 M6-S4 mtp: GGUF tensor names upstream-aligned (#23398/#24282) — nextn.* / masked_embd_*.
// The C++ member names stay mtp_* (control-plane; our single-context engine, not upstream's ctx_other).
- mtp_pre_projection = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_PRE, "weight"), {2 * (int64_t) n_bb, n_embd}, 0);
+ //
+ // opencoti bug-858 (dual-context MTP): pre_projection is classified LLM_TENSOR_LAYER_INPUT, so the
+ // plain create_tensor lands it on the CPU buft. The single-context facade tolerates that (it runs on
+ // the target's sched_mtp, which has the CPU backend registered), but a REAL draft context ctx_dft
+ // (llama_init_from_model over the nested assistant) registers only the assistant's *offloaded* (CUDA)
+ // weight buffers — a CPU-resident weight then reads back uninitialized → the whole drafter forward is
+ // NaN → 0% accept. Force it onto the output-head (GPU-when-offloaded) buft, exactly like tok_embd
+ // (#408): device-config-respecting (stays CPU for a CPU-only drafter) and placement-only, so the
+ // facade path stays byte-identical. post_projection/output_norm are LAYER_OUTPUT → already dev_output.
+ mtp_pre_projection = create_tensor_output_head(tn(LLM_TENSOR_NEXTN_PROJ_PRE, "weight"), {2 * (int64_t) n_bb, n_embd}, 0);
mtp_post_projection = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_POST, "weight"), {n_embd, (int64_t) n_bb}, 0);
if (hparams.use_ordered_embeddings) {
@@ -481,8 +496,167 @@ void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) {
}
}
-std::unique_ptr<llm_graph_context> llama_model_gemma4_assistant::build_arch_graph(const llm_graph_params &) const {
- throw std::runtime_error(
- "gemma4_assistant cannot be used as a primary model (-m). "
- "Load the Gemma 4 target with -m, then call llama_model_load_mtp_from_file() with the assistant GGUF.");
+// opencoti bug-858 dual-context MTP (A6): the gemma4-assistant graph built when the assistant is
+// created as its OWN llama_context (ctx_dft) with cparams.ctx_other = the target context. Mirrors
+// upstream b9859 llama_model_gemma4_assistant::graph. The assistant is KV-LESS (wq/wo only, no
+// wk/wv): each attention layer reads the TARGET's shared K/V *in place* via the standard iSWA
+// build_attn (Qcur only, null k_cur/v_cur -> no store). The draft context's create_memory (A3)
+// shares the target's KV cells and aliases the target's last full (n_layer-1) / SWA (n_layer-2)
+// layer per draft layer, so build_attn at layer il transparently reads the aliased target K/V —
+// which is exactly what the single-context path did explicitly via build_attn_mtp(..., il_kv).
+// The input token embedding + backbone hidden come from the target model via ctx_other.
+//
+// This is DISTINCT from the single-context llm_build_gemma4_mtp (built by the gemma4 TARGET's
+// build_arch_graph, gemma4.cpp) which is kept fully intact as the fallback drafter.
+//
+// Faithful to upstream b9859 (and unlike our single-context build): dense LM head over tok_embd
+// (no ordered-embeddings/centroid path — b9859 loads centroids but ignores them in the graph),
+// no final-logit softcapping (monotonic -> argmax-invariant), and no control-vector application
+// (inert without a control vector). See docs/evaluations/mtp.md.
+struct llm_build_gemma4_assistant_dual : public llm_graph_context {
+ llm_build_gemma4_assistant_dual(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ const int64_t n_bb = hparams.n_embd_out_impl; // backbone width (== target n_embd)
+ GGML_ASSERT(n_bb > 0);
+ GGML_ASSERT(model.mtp_pre_projection != nullptr && model.mtp_post_projection != nullptr);
+
+ GGML_ASSERT(cparams.ctx_other != nullptr &&
+ "gemma4_assistant is a KV-less drafter: it can only run as an MTP draft context "
+ "(ctx_dft) created with cparams.ctx_other = the Gemma 4 target context.");
+ const llama_model * model_other = llama_get_model(cparams.ctx_other);
+
+ // Draft inputs: the drafted token ids + the paired target backbone hidden (h_prev).
+ ggml_tensor * inp_tokens;
+ ggml_tensor * inp_h;
+ {
+ auto inp = std::make_unique<llm_graph_input_embd>(n_bb);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);
+ cb(inp->tokens, "inp_tokens", -1);
+ ggml_set_input(inp->tokens);
+ inp_tokens = inp->tokens;
+ res->t_inp_tokens = inp->tokens;
+
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_bb, ubatch.n_tokens);
+ cb(inp->embd, "inp_h", -1);
+ ggml_set_input(inp->embd);
+ inp_h = inp->embd;
+ res->t_inp_embd = inp->embd;
+
+ res->add_input(std::move(inp));
+ }
+
+ // Read the TARGET's token embedding through ctx_other (the assistant reuses the target's
+ // tok_embd for the input pipeline), scaled by sqrt(n_bb) like Gemma 4's input pipeline.
+ ggml_tensor * x = ggml_get_rows(ctx0, model_other->tok_embd, inp_tokens);
+ x = ggml_scale(ctx0, x, sqrtf((float) n_bb));
+ cb(x, "inp_embd_target", -1);
+
+ ggml_tensor * xh = ggml_concat(ctx0, x, inp_h, 0);
+ cb(xh, "inp_xh", -1);
+
+ ggml_tensor * cur = build_lora_mm(model.mtp_pre_projection, xh);
+ cb(cur, "pre_proj", -1);
+
+ auto * inp_attn = build_attn_inp_kv_iswa();
+ ggml_tensor * inp_pos = build_inp_pos();
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ ggml_tensor * inpL = cur;
+
+ for (int il = 0; il < n_layer; ++il) {
+ const bool is_swa = hparams.is_swa(il);
+
+ const int64_t n_embd_head = hparams.n_embd_head_k(il);
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_v(il));
+ const int64_t n_head = hparams.n_head(il);
+
+ const float freq_base_l = model.get_rope_freq_base(cparams, il);
+ const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
+ const int n_rot_l = hparams.n_rot(il);
+
+ ggml_tensor * cur_norm = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur_norm, "attn_norm", il);
+
+ ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur_norm);
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
+ Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Qcur, "Qcur_normed", il);
+
+ ggml_tensor * freq_factors = is_swa ? nullptr : model.layers[il].rope_freqs;
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig,
+ freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(Qcur, "Qcur_pos", il);
+
+ // KV-less cross-attention into the shared/aliased target K/V (null k_cur/v_cur -> no store).
+ cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr,
+ Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);
+
+ if (il == n_layer - 1 && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+ }
+
+ cur = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_post_norm", il);
+
+ ggml_tensor * attn_out = ggml_add(ctx0, cur, inpL);
+ cb(attn_out, "attn_out", il);
+
+ GGML_ASSERT(model.layers[il].ffn_gate_inp == nullptr && "gemma4_assistant MTP does not support MoE FFN");
+
+ cur = build_norm(attn_out, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, nullptr, nullptr,
+ model.layers[il].ffn_gate, nullptr, nullptr,
+ model.layers[il].ffn_down, nullptr, nullptr,
+ nullptr,
+ LLM_FFN_GELU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+
+ cur = build_norm(cur, model.layers[il].ffn_post_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "ffn_post_norm", il);
+
+ cur = ggml_add(ctx0, cur, attn_out);
+
+ if (model.layers[il].out_scale) {
+ cur = ggml_mul(ctx0, cur, model.layers[il].out_scale);
+ cb(cur, "out_scaled", il);
+ }
+
+ inpL = cur;
+ }
+
+ cur = inpL;
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+
+ // LM head (dense, tied to tok_embd — ours' create_tensor_output_head; upstream uses a
+ // DUPLICATED `output`, same weight). Full-vocab logits for greedy verify/draft.
+ ggml_tensor * logits = build_lora_mm(model.tok_embd, cur);
+ cb(logits, "result_output", -1);
+ res->t_logits = logits;
+
+ // Recurrent backbone hidden for the NEXT draft step (post-projected to n_bb), exposed via
+ // t_h_pre_norm so the draft loop harvests it with llama_get_embeddings_pre_norm_ith(ctx_dft).
+ // Mirrors upstream res->t_h_nextn = mul_mat(nextn_proj_post, cur).
+ ggml_tensor * h_next = build_lora_mm(model.mtp_post_projection, cur);
+ cb(h_next, "h_pre_norm", -1);
+ res->t_h_pre_norm = h_next;
+
+ ggml_build_forward_expand(gf, logits);
+ ggml_build_forward_expand(gf, h_next);
+ }
+};
+
+std::unique_ptr<llm_graph_context> llama_model_gemma4_assistant::build_arch_graph(const llm_graph_params & params) const {
+ // opencoti bug-858 dual-context MTP (A6): build the gemma4-assistant drafter graph. This model
+ // is KV-less and only valid as an MTP draft context (ctx_dft) created FROM it with
+ // cparams.ctx_other = the Gemma 4 target — the ctor GGML_ASSERTs that. The single-context
+ // engine (llm_build_gemma4_mtp) is built by the gemma4 TARGET's build_arch_graph instead and
+ // never reaches here; it remains the fallback drafter.
+ return std::make_unique<llm_build_gemma4_assistant_dual>(*this, params);
}
diff --git a/llama.cpp/src/models/gemma4.cpp b/llama.cpp/src/models/gemma4.cpp
index a7b1598..bdf5990 100644
--- a/llama.cpp/src/models/gemma4.cpp
+++ b/llama.cpp/src/models/gemma4.cpp
@@ -171,6 +171,15 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
ggml_tensor * inp_out_ids = build_inp_out_ids();
+ // opencoti-hook: gemma4-mtp-hidden (bug-858 dual-context MTP) — when a context extracts the
+ // pre-norm hidden for ALL positions (embeddings_pre_norm && !masked, i.e. the MTP *target* whose
+ // shifted per-position hidden seeds the drafter), DEFER the inp_out_ids row-strip until after
+ // t_h_pre_norm is captured, so t_h_pre_norm carries all n_tokens rows. Upstream b9859 gates the
+ // in-loop strip on embeddings_nextn_masked and strips late when !masked (gemma4.cpp:278/416).
+ // Normal decode (embeddings_pre_norm==false) and the masked draft keep ours' early strip →
+ // byte-identical. See docs/evaluations/mtp.md and UPSTREAM_SYNC.md.
+ const bool mtp_defer_out_ids = cparams.embeddings_pre_norm && !cparams.embeddings_pre_norm_masked;
+
ggml_tensor * inp_per_layer = nullptr;
if (model.per_layer_tok_embd) {
inp_per_layer = build_inp_per_layer();
@@ -253,7 +262,9 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
}
// TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing
- if (il == n_layer - 1 && inp_out_ids) {
+ // opencoti bug-858: skip the early strip for the MTP target (mtp_defer_out_ids) so the last
+ // layer runs on all n_tokens rows and t_h_pre_norm below stays ungathered.
+ if (il == n_layer - 1 && inp_out_ids && !mtp_defer_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
@@ -353,7 +364,9 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
// TODO @ngxson : improve this
- if (il == n_layer - 1 && inp_out_ids) {
+ // opencoti bug-858: same deferred strip as the main path (keep per-layer input full for
+ // the MTP target so its last-layer residual matches the ungathered hidden).
+ if (il == n_layer - 1 && inp_out_ids && !mtp_defer_out_ids) {
inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
}
@@ -384,6 +397,22 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
model.output_norm, nullptr,
LLM_NORM_RMS, -1);
+ // opencoti-hook: gemma4-mtp-hidden (bug-858 dual-context MTP) — expose the POST-output-norm
+ // hidden state as t_h_pre_norm so a gemma4-assistant draft context (ctx_dft with ctx_other=this)
+ // can read it via llama_get_embeddings_pre_norm_ith() as the recurrent h input. Mirrors upstream
+ // b9859 gemma4.cpp res->t_h_nextn = cur (post-final-norm, captured PRE row-strip). For the MTP
+ // target (mtp_defer_out_ids) this tensor carries all n_tokens rows so the driver can shift the
+ // per-position hidden; for normal decode / masked draft the rows were already stripped in the
+ // last layer, so it equals the n_outputs t_embd below. See docs/evaluations/mtp.md.
+ cb(cur, "h_pre_norm", -1);
+ res->t_h_pre_norm = cur;
+
+ // opencoti bug-858: deferred row-strip — t_h_pre_norm captured all positions above; now strip to
+ // the output rows for the logits / t_embd path (upstream gemma4.cpp:416 `if (!masked) get_rows`).
+ if (mtp_defer_out_ids && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ }
+
cb(cur, "result_norm", -1);
res->t_embd = cur;
diff --git a/llama.cpp/src/models/qwen35.cpp b/llama.cpp/src/models/qwen35.cpp
index 131f1fa..3a21d69 100644
--- a/llama.cpp/src/models/qwen35.cpp
+++ b/llama.cpp/src/models/qwen35.cpp
@@ -215,6 +215,20 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para
}
cur = inpL;
+ // opencoti-hook: qwen-nextn-mtp-hidden (bug-858) — see docs/evaluations/mtp.md.
+ // Upstream (b9859) exposes the MTP/NextN hidden as the POST-output-norm state (t_h_nextn);
+ // we mirror that here (apply output_norm FIRST, then expose it) so t_h_pre_norm points at the
+ // same tensor upstream's nextn.hnorm consumes — a syncability alignment, not a perf fix.
+ // MEASURED: on Qwen3.6-35B-A3B-MTP this reorder is numerically INERT (output_norm ≈ identity,
+ // so RMSNorm-then-hnorm == hnorm; draft acceptance is byte-identical 0.626 either way on GPU).
+ // The real acceptance gap vs upstream (ours 0.626 vs 0.770) is NOT the hidden-capture point —
+ // it is localized to ours' GPU flash-attention kernel: CPU is at parity (ours 0.696 ≈ upstream
+ // 0.692) and `-fa off` on GPU recovers ours to 0.683, so the divergence is FA-kernel precision
+ // in the nextn draft's full-attn block, not this norm. RMSNorm is row-independent, so applying
+ // it before vs after the inp_out_ids reduction leaves t_embd / t_logits byte-identical — base
+ // decode + RULER unaffected.
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+
cb(cur, "h_pre_norm", -1);
res->t_h_pre_norm = cur;
@@ -222,9 +236,6 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
- // Final norm
- cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
-
cb(cur, "result_norm", -1);
res->t_embd = cur;
@@ -539,9 +550,8 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
- ggml_tensor * h_input = inp->embd;
- ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
- cb(tok_embd, "mtp_tok_embd", il);
+ ggml_tensor * h_input = inp->embd;
+ ggml_tensor * inp_tokens = inp->tokens;
res->add_input(std::move(inp));
@@ -549,6 +559,17 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr
ggml_tensor * inp_out_ids = build_inp_out_ids();
auto * inp_attn = build_attn_inp_kv();
+ // opencoti-hook: fused-nextn-mtp — one MTP step (embed → block → head → greedy argmax) as a
+ // reusable lambda the fused draft graph replays N× on-device (in-graph argmax→embed chain,
+ // one launch+readback). Takes the token index (embeds internally) so step k+1 can chain on
+ // step k's argmax. Single-step (n_mtp_steps==1) is byte-identical to the prior inline build
+ // (the argmax node is pruned when unreferenced). See docs/features/fused_nextn_mtp.md.
+ auto build_one_step = [&](ggml_tensor * h_input, ggml_tensor * tok,
+ ggml_tensor * inp_pos, auto * inp_attn,
+ ggml_tensor * inp_out_ids, bool fused)
+ -> std::tuple<ggml_tensor *, ggml_tensor *, ggml_tensor *> {
+ ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, tok);
+ cb(tok_embd, "mtp_tok_embd", il);
ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
@@ -597,16 +618,27 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr
Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr,
n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
- Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,
- n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
- ext_factor, attn_factor, beta_fast, beta_slow);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
- cur = build_attn(inp_attn,
- nullptr, nullptr, nullptr,
- Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ if (fused) {
+ // opencoti fused-NextN (Option A / Gemma-mirror): read-only cross-attention into the
+ // FROZEN prefix KV. The prefix K/V were already written (rope'd + normed) by the draft
+ // context's prior decodes; we do NOT write draft KV here — recurrence is carried by the
+ // hidden-state chain. mtp_slot_info gives one cell (pmax); get_n_kv reports the full used
+ // extent so the causal kq_mask exposes 0..pmax. Kcur/Vcur above go unused here → pruned.
+ // Mirrors build_attn_mtp (llama-graph.cpp). See docs/features/fused_nextn_mtp.md delta #6.
+ cur = build_attn_readonly_nextn(inp_attn, Qcur, kq_scale, il);
+ } else {
+ Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,
+ n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ cur = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ }
cb(cur, "mtp_attn_pregate", il);
cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));
@@ -634,7 +666,7 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr
// Pre-norm hidden state: used by the AR draft loop to seed the next MTP step.
// (In the trunk graph this is `t_h_pre_norm`; the MTP head reuses the same slot.)
cb(cur, "h_pre_norm", -1);
- res->t_h_pre_norm = cur;
+ ggml_tensor * h_pre_norm = cur;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
@@ -651,6 +683,70 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
- res->t_logits = cur;
- ggml_build_forward_expand(gf, cur);
+ ggml_tensor * arg = ggml_argmax(ctx0, cur);
+ cb(arg, "mtp_argmax", -1);
+
+ return { h_pre_norm, cur, arg };
+ };
+
+ const int32_t n_steps = std::max<int32_t>(params.n_mtp_steps, 1);
+
+ if (n_steps <= 1) {
+ auto [h_pre_norm, logits, arg] = build_one_step(h_input, inp_tokens, inp_pos, inp_attn, inp_out_ids, false);
+ (void) arg; // single-step: draft token comes from the backend sampler (t_sampled) or host sampler
+ res->t_h_pre_norm = h_pre_norm;
+ res->t_logits = logits;
+ ggml_build_forward_expand(gf, logits);
+ return;
+ }
+
+ // opencoti-hook: fused-nextn-mtp — fused N-step draft chain (mirrors gemma4-assistant
+ // llm_build_gemma4_mtp). DORMANT until the driver sets n_mtp_steps>1 (S3). One
+ // I32[n_pos_per_embd] position input per step (bug-870: mrope archs like qwen35 need 4
+ // ids/token — ggml_rope_multi asserts a->ne[2]*4==b->ne[0]; the setter fills the [p,p,p,0]
+ // M-RoPE text layout, or [p] for standard rope). Step k+1's token = step k's on-device argmax,
+ // its seed hidden = step k's t_h_pre_norm — no host round-trip. Reuses llm_graph_input_mtp
+ // (width-agnostic inp_h_prev). See docs/features/fused_nextn_mtp.md.
+ const int64_t n_pos_e = hparams.n_pos_per_embd();
+ auto inp_mtp = std::make_unique<llm_graph_input_mtp>();
+ inp_mtp->inp_last_token = inp_tokens;
+ inp_mtp->inp_h_prev = h_input;
+ inp_mtp->inp_pos_steps.reserve(n_steps);
+ std::vector<ggml_tensor *> pos_steps;
+ pos_steps.reserve(n_steps);
+ for (int32_t k = 0; k < n_steps; ++k) {
+ ggml_tensor * p = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos_e);
+ ggml_set_input(p);
+ cb(p, "mtp_inp_pos_step", k);
+ pos_steps.push_back(p);
+ inp_mtp->inp_pos_steps.push_back(p);
+ }
+ res->add_input(std::move(inp_mtp));
+
+ ggml_tensor * tok_k = inp_tokens;
+ ggml_tensor * h_k = h_input;
+ std::vector<ggml_tensor *> step_args;
+ step_args.reserve(n_steps);
+ ggml_tensor * last_h_pre_norm = nullptr;
+ ggml_tensor * last_logits = nullptr;
+ for (int32_t k = 0; k < n_steps; ++k) {
+ auto [h_pre_norm_k, logits_k, arg_k] = build_one_step(h_k, tok_k, pos_steps[k], inp_attn, inp_out_ids, true);
+ step_args.push_back(arg_k);
+ last_h_pre_norm = h_pre_norm_k;
+ last_logits = logits_k;
+ tok_k = arg_k; // next token = this step's greedy argmax (I32[1])
+ h_k = h_pre_norm_k; // next seed hidden = this step's pre-norm hidden
+ }
+
+ ggml_tensor * all_args = step_args[0];
+ for (int32_t k = 1; k < n_steps; ++k) {
+ all_args = ggml_concat(ctx0, all_args, step_args[k], 0);
+ }
+ cb(all_args, "mtp_fused_argmax", -1);
+
+ res->t_argmax = all_args; // I32[n_steps] drafted tokens
+ res->t_h_pre_norm = last_h_pre_norm; // final hidden (seed for next cycle)
+ res->t_logits = last_logits;
+ ggml_build_forward_expand(gf, all_args);
+ ggml_build_forward_expand(gf, last_h_pre_norm);
}
diff --git a/llama.cpp/src/models/qwen35moe.cpp b/llama.cpp/src/models/qwen35moe.cpp
index f3f8e73..0eb2871 100644
--- a/llama.cpp/src/models/qwen35moe.cpp
+++ b/llama.cpp/src/models/qwen35moe.cpp
@@ -602,9 +602,8 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
- ggml_tensor * h_input = inp->embd;
- ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
- cb(tok_embd, "mtp_tok_embd", il);
+ ggml_tensor * h_input = inp->embd;
+ ggml_tensor * inp_tokens = inp->tokens;
res->add_input(std::move(inp));
@@ -612,7 +611,17 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm
ggml_tensor * inp_out_ids = build_inp_out_ids();
auto * inp_attn = build_attn_inp_kv();
-
+ // opencoti-hook: fused-nextn-mtp — one MTP step (embed → block → head → greedy argmax) as a
+ // reusable lambda the fused draft graph replays N× on-device (in-graph argmax→embed chain,
+ // one launch+readback). Takes the token index (embeds internally) so step k+1 can chain on
+ // step k's argmax. Single-step (n_mtp_steps==1) is byte-identical to the prior inline build
+ // (the argmax node is pruned when unreferenced). See docs/features/fused_nextn_mtp.md.
+ auto build_one_step = [&](ggml_tensor * h_input, ggml_tensor * tok,
+ ggml_tensor * inp_pos, auto * inp_attn,
+ ggml_tensor * inp_out_ids, bool fused)
+ -> std::tuple<ggml_tensor *, ggml_tensor *, ggml_tensor *> {
+ ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, tok);
+ cb(tok_embd, "mtp_tok_embd", il);
ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
@@ -661,16 +670,27 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm
Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr,
n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
- Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,
- n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
- ext_factor, attn_factor, beta_fast, beta_slow);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
- cur = build_attn(inp_attn,
- nullptr, nullptr, nullptr,
- Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ if (fused) {
+ // opencoti fused-NextN (Option A / Gemma-mirror): read-only cross-attention into the
+ // FROZEN prefix KV. The prefix K/V were already written (rope'd + normed) by the draft
+ // context's prior decodes; we do NOT write draft KV here — recurrence is carried by the
+ // hidden-state chain. mtp_slot_info gives one cell (pmax); get_n_kv reports the full used
+ // extent so the causal kq_mask exposes 0..pmax. Kcur/Vcur above go unused here → pruned.
+ // Mirrors build_attn_mtp (llama-graph.cpp). See docs/features/fused_nextn_mtp.md delta #6.
+ cur = build_attn_readonly_nextn(inp_attn, Qcur, kq_scale, il);
+ } else {
+ Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,
+ n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ cur = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ }
cb(cur, "mtp_attn_pregate", il);
cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));
@@ -730,7 +750,7 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm
// Pre-norm hidden state: used by the AR draft loop to seed the next MTP step.
cb(cur, "h_pre_norm", -1);
- res->t_h_pre_norm = cur;
+ ggml_tensor * h_pre_norm = cur;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
@@ -747,6 +767,70 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
- res->t_logits = cur;
- ggml_build_forward_expand(gf, cur);
+ ggml_tensor * arg = ggml_argmax(ctx0, cur);
+ cb(arg, "mtp_argmax", -1);
+
+ return { h_pre_norm, cur, arg };
+ };
+
+ const int32_t n_steps = std::max<int32_t>(params.n_mtp_steps, 1);
+
+ if (n_steps <= 1) {
+ auto [h_pre_norm, logits, arg] = build_one_step(h_input, inp_tokens, inp_pos, inp_attn, inp_out_ids, false);
+ (void) arg; // single-step: host samples via common_sampler_sample (argmax node pruned)
+ res->t_h_pre_norm = h_pre_norm;
+ res->t_logits = logits;
+ ggml_build_forward_expand(gf, logits);
+ return;
+ }
+
+ // opencoti-hook: fused-nextn-mtp — fused N-step draft chain (mirrors gemma4-assistant
+ // llm_build_gemma4_mtp). DORMANT until the driver sets n_mtp_steps>1 (S3). One
+ // I32[n_pos_per_embd] position input per step (bug-870: mrope archs like qwen35 need 4
+ // ids/token — ggml_rope_multi asserts a->ne[2]*4==b->ne[0]; the setter fills the [p,p,p,0]
+ // M-RoPE text layout, or [p] for standard rope). Step k+1's token = step k's on-device argmax,
+ // its seed hidden = step k's t_h_pre_norm — no host round-trip. Reuses llm_graph_input_mtp
+ // (width-agnostic inp_h_prev). See docs/features/fused_nextn_mtp.md.
+ const int64_t n_pos_e = hparams.n_pos_per_embd();
+ auto inp_mtp = std::make_unique<llm_graph_input_mtp>();
+ inp_mtp->inp_last_token = inp_tokens;
+ inp_mtp->inp_h_prev = h_input;
+ inp_mtp->inp_pos_steps.reserve(n_steps);
+ std::vector<ggml_tensor *> pos_steps;
+ pos_steps.reserve(n_steps);
+ for (int32_t k = 0; k < n_steps; ++k) {
+ ggml_tensor * p = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos_e);
+ ggml_set_input(p);
+ cb(p, "mtp_inp_pos_step", k);
+ pos_steps.push_back(p);
+ inp_mtp->inp_pos_steps.push_back(p);
+ }
+ res->add_input(std::move(inp_mtp));
+
+ ggml_tensor * tok_k = inp_tokens;
+ ggml_tensor * h_k = h_input;
+ std::vector<ggml_tensor *> step_args;
+ step_args.reserve(n_steps);
+ ggml_tensor * last_h_pre_norm = nullptr;
+ ggml_tensor * last_logits = nullptr;
+ for (int32_t k = 0; k < n_steps; ++k) {
+ auto [h_pre_norm_k, logits_k, arg_k] = build_one_step(h_k, tok_k, pos_steps[k], inp_attn, inp_out_ids, true);
+ step_args.push_back(arg_k);
+ last_h_pre_norm = h_pre_norm_k;
+ last_logits = logits_k;
+ tok_k = arg_k; // next token = this step's greedy argmax (I32[1])
+ h_k = h_pre_norm_k; // next seed hidden = this step's pre-norm hidden
+ }
+
+ ggml_tensor * all_args = step_args[0];
+ for (int32_t k = 1; k < n_steps; ++k) {
+ all_args = ggml_concat(ctx0, all_args, step_args[k], 0);
+ }
+ cb(all_args, "mtp_fused_argmax", -1);
+
+ res->t_argmax = all_args; // I32[n_steps] drafted tokens
+ res->t_h_pre_norm = last_h_pre_norm; // final hidden (seed for next cycle)
+ res->t_logits = last_logits;
+ ggml_build_forward_expand(gf, all_args);
+ ggml_build_forward_expand(gf, last_h_pre_norm);
}
diff --git a/llama.cpp/tools/server/server-context.cpp b/llama.cpp/tools/server/server-context.cpp
index 7bad046..34b50b1 100644
--- a/llama.cpp/tools/server/server-context.cpp
+++ b/llama.cpp/tools/server/server-context.cpp
@@ -965,10 +965,49 @@ private:
return false;
}
- // No separate draft model/context: the assistant decodes on ctx_tgt via llama_decode_mtp.
- params_base.speculative.draft.ctx_tgt = ctx_tgt;
- params_base.speculative.draft.ctx_dft = nullptr;
- SRV_INF("%s", "MTP assistant loaded into target\n");
+ // opencoti bug-858 dual-context MTP (A5): optionally create a REAL draft context FROM the
+ // nested gemma4-assistant model with cparams.ctx_other = ctx_tgt (upstream b9859
+ // dual-context). The assistant is KV-less: its create_memory (A3) shares the target's KV
+ // cells + aliases the last full/SWA layer, and its graph (A6) reads the target's tok_embd
+ // + shared K/V via ctx_other; the shared-KV draft_mtp driver (is_mem_shared) then runs it.
+ // Gated by OPENCOTI_MTP_DUAL_CTX so the default keeps the proven single-context in-target
+ // path (ctx_dft == nullptr, llama_decode_mtp) byte-identical.
+ const char * mtp_dual_env = getenv("OPENCOTI_MTP_DUAL_CTX");
+ const bool mtp_dual_ctx = mtp_dual_env && mtp_dual_env[0] && mtp_dual_env[0] != '0';
+ if (mtp_dual_ctx) {
+ // llama_init_from_model wants a mutable model*; the nested assistant is owned mutably
+ // by the target (unique_ptr<llama_model>), so the const_cast is safe.
+ llama_model * assistant = const_cast<llama_model *>(llama_model_get_mtp_assistant(model_tgt));
+ if (assistant == nullptr) {
+ SRV_ERR("%s", "MTP dual-context requested but the assistant is not loaded into the target\n");
+ return false;
+ }
+
+ auto cparams_mtp = common_context_params_to_llama(params_dft);
+ cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+ cparams_mtp.type_k = params_spec.cache_type_k;
+ cparams_mtp.type_v = params_spec.cache_type_v;
+ cparams_mtp.n_rs_seq = 0;
+ cparams_mtp.ctx_other = ctx_tgt;
+
+ ctx_dft.reset(llama_init_from_model(assistant, cparams_mtp));
+ if (ctx_dft == nullptr) {
+ SRV_ERR("%s", "failed to create dual-context MTP assistant draft context\n");
+ return false;
+ }
+
+ // NOTE: leave ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO (the default) — the
+ // assistant's KV is SHARED with the target, so ctx_dft must never be independently
+ // seq_rm'd / checkpointed (that would corrupt the target's cells).
+ params_base.speculative.draft.ctx_tgt = ctx_tgt;
+ params_base.speculative.draft.ctx_dft = ctx_dft.get();
+ SRV_INF("%s", "MTP assistant dual-context draft created (ctx_other=target, shared KV)\n");
+ } else {
+ // No separate draft context: the assistant decodes on ctx_tgt via llama_decode_mtp.
+ params_base.speculative.draft.ctx_tgt = ctx_tgt;
+ params_base.speculative.draft.ctx_dft = nullptr;
+ SRV_INF("%s", "MTP assistant loaded into target (single-context)\n");
+ }
} else if (params_base.speculative.has_dft()) {
// TODO speculative: move to common/speculative.cpp?
const auto & params_spec = params_base.speculative.draft;
@@ -2641,9 +2710,7 @@ private:
}
// generate the actual drafts (if any)
- {
- common_speculative_draft(spec.get());
- }
+ common_speculative_draft(spec.get());
// make checkpoints if needed
for (auto * slot_ptr : drafting) {
@@ -2674,13 +2741,8 @@ private:
(ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft.get()));
if (use_ckpt_tgt) {
- //const int64_t t_start = ggml_time_us();
-
ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
- //const int64_t t_total = ggml_time_us() - t_start;
- //printf("checkpoint total: %f ms\n", t_total / 1000.0);
-
SLT_DBG(slot, "created speculative checkpoint (pos_min = %d, pos_max = %d, n_tokens = %d, size = %.3f MiB, draft = %.3f MiB)\n",
ckpt.pos_min, ckpt.pos_max, slot.prompt.n_tokens(),
(float) ckpt.size() / 1024 / 1024,