| 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). |
|
|
| |
| |
| |
| |
| @@ -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; |
| } |
| |
| |
| |
| |
| @@ -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) |
| |
| |
| |
| |
| @@ -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: { |
| |
| |
| |
| |
| @@ -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) { |
| |
| |
| |
| |
| @@ -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; |
| |
| |
| |
| |
| @@ -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; |
| |
| |
| |
| |
| @@ -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; |
| } |
| |
| |
| |
| |
| @@ -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) |
| |
| |
| |
| |
| @@ -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 |
| // |
| |
| |
| |
| |
| @@ -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 |
| |
| |
| |
| |
| @@ -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; |
| }; |
| |
| |
| |
| |
| @@ -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); |
| |
| |
| |
| |
| @@ -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, |
| |
| |
| |
| |
| @@ -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( |
| |
| |
| |
| |
| @@ -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) { |
| |
| |
| |
| |
| @@ -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. |
| |
| |
| |
| |
| @@ -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())); |
| |
| |
| |
| |
| @@ -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; |
| |
| |
| |
| |
| @@ -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, |
| |
| |
| |
| |
| @@ -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, |
| |
| |
| |
| |
| @@ -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 |
| |
| |
| |
| |
| @@ -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); |
| } |
| |
| |
| |
| |
| @@ -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); |
| } |
| |
| |
| |
| |
| @@ -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; |
| |
| |
| |
| |
| |
| @@ -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); |
| } |
| |
| |
| |
| |
| @@ -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); |
| } |
| |
| |
| |
| |
| @@ -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, |
|
|