From: opencoti Subject: [PATCH 0093] bug-858 MTP parity: dual-ctx assistant port + fused NextN + verify-path fixes Wholesale alignment of speculative MTP to upstream b9859 (#590/#607/#599/#609): - gemma4-assistant dual-context execution (ctx_dft sharing the target KV via cparams.ctx_other / mem_other + share selector; OPENCOTI_MTP_DUAL_CTX gate), draft_mtp impl with pre-norm embedding tap/harvest, defer-gather boot fix. - fused NextN greedy N-step draft (build_one_step + decode_mtp_fused_nextn on a dedicated sched_nextn; can_reuse base-drift + ensure_cleared aliased-KV fixes). - bug-2103: mmvf ne11 cap raised to 8 under TinyBLAS (n3 verify GEMV cliff). - bug-2094 diag + rolling-KV visibility hunks ride along in llama-kv-cache.cpp (former 0091-rolling-kv-diag-visibility, folded here during the bug-2108 chain repair). diff --git a/llama.cpp/common/arg.cpp b/llama.cpp/common/arg.cpp index 89bd9dd..3fb1518 100644 --- a/llama.cpp/common/arg.cpp +++ b/llama.cpp/common/arg.cpp @@ -1455,7 +1455,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex // opencoti F5 M7 Rolling KV — Rung 0 auto residency — see docs/features/rolling_kv.md add_opt(common_arg( {"--vram-target"}, "MiB", - string_format("rolling-kv: VRAM budget cap in MiB for auto KV residency (used with --headinfer-gpu-heads-frac auto). 0 = use all free VRAM minus a compute reserve (maximize). (default: %d)", params.vram_target_mib), + string_format("rolling-kv: TOTAL GPU VRAM cap in MiB (weights + KV + compute), NOT a KV-only budget. Auto KV residency keeps as much KV on-device as fits under this cap after weights are loaded; the overflow streams from host. Set it to your card's size to simulate/bound that card. 0 = use all free VRAM minus a compute reserve (maximize). (default: %d)", params.vram_target_mib), [](common_params & params, int value) { params.vram_target_mib = value; } diff --git a/llama.cpp/common/common.h b/llama.cpp/common/common.h index b688492..3bb6aa2 100644 --- a/llama.cpp/common/common.h +++ b/llama.cpp/common/common.h @@ -303,9 +303,12 @@ struct common_params_speculative_draft { int32_t n_max = 3; // maximum number of tokens to draft during speculative decoding int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding - // opencoti F5 M6-S4 mtp — MTP (Gemma 4 assistant): draft block size B produces - // B-1 draft tokens per round. Default 3 (= 2 chained MTP draft steps). - int32_t draft_block_size = 3; + // opencoti F5 M6-S4 mtp — MTP (Gemma 4 assistant): OPTIONAL explicit draft-depth ceiling. + // 0 (default) = draft depth follows --spec-draft-n-max (n_max); the MTP head re-runs + // autoregressively per step (decode_mtp_sync/fused), so depth is NOT limited to a fixed + // block size. bug-858: the old default 3 hard-capped ours at 2 drafts/round and blocked + // n-max scaling vs upstream b9859. When set >1, caps draft depth at draft_block_size-1. + int32_t draft_block_size = 0; float p_split = 0.1f; // speculative decoding split probability float p_min = 0.0f; // minimum speculative decoding probability (greedy) diff --git a/llama.cpp/common/speculative.cpp b/llama.cpp/common/speculative.cpp index 7faeac7..ee7962f 100644 --- a/llama.cpp/common/speculative.cpp +++ b/llama.cpp/common/speculative.cpp @@ -481,6 +481,13 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { int32_t n_embd = 0; + // opencoti bug-858 dual-context MTP: true for the gemma4 assistant (ctx_dft shares the target's + // KV via ctx_other). Derived in the ctor. When set, process() skips the ctx_dft catch-up decode + // (the shared cache is already populated by the target's own decode) and draft() pins all draft + // tokens to the same position dp.n_past (KV-less: no per-step cell growth). Mirrors upstream + // b9859 common_speculative_impl_draft_mtp::is_mem_shared. false = Qwen NextN (byte-identical). + bool is_mem_shared = false; + // Per-sequence cross-batch carryover: pair (h_p, x_{p+1}) at MTP pos p+1. // The last h-row of one process() call needs the first token of the NEXT // call to pair with, so it's stashed here until that next call fires. @@ -507,7 +514,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { auto * ctx_dft = this->params.ctx_dft; GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set"); - n_embd = llama_model_n_embd(llama_get_model(ctx_dft)); + // opencoti bug-858 dual-context MTP: the gemma4 assistant runs as ctx_dft with + // ctx_other=ctx_tgt (shared KV). Detect it here; every branch below keys off this flag. + is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt; + + // n_embd = width of an h row (target backbone hidden). For the assistant the ctx_dft model + // is the assistant, whose n_embd_out() == the target n_embd (validated at load); for Qwen + // NextN the ctx_dft model IS the target, so keep the original accessor (byte-identical). + n_embd = is_mem_shared + ? llama_model_n_embd_out(llama_get_model(ctx_dft)) + : llama_model_n_embd(llama_get_model(ctx_dft)); LOG_INF("%s: adding speculative implementation 'draft-mtp'\n", __func__); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); @@ -550,7 +566,19 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } } - llama_set_embeddings_pre_norm(ctx_tgt, true, /*masked*/ false); + // opencoti bug-858 losslessness (2026-07-05): the target tap must not change the target's graph. + // The unmasked pre-norm tap engages gemma4.cpp's mtp_defer_out_ids (late row-strip: the last + // layer + output_norm run on ALL verify rows instead of the output rows) — structurally equal to + // upstream b9859's defer, but ours' GEMM kernels are batch-shape-sensitive (#495), so the deferred + // graph flips near-tie argmax vs plain decode (gate4/5: fib DIFF@2, forces DIFF@126; the facade, + // which never defers, is clean). Fix: for the shared-KV assistant set the target tap MASKED — + // masked keeps the normal early strip (target graph BIT-IDENTICAL to plain decode; t_h_pre_norm + // aliases the always-captured t_embd) and captures n_outputs rows with the output_ids-translated + // _ith accessor. Every verify token has logits=true, and prefill only ever consumes its last row + // (also an output row), so all consumed values equal upstream's all-rows defer capture — see the + // logits-rows harvest in process(). Qwen NextN (is_mem_shared==false) keeps the dense unmasked + // tap (its catch-up decode reads the whole shifted prefill buffer) — byte-identical. + llama_set_embeddings_pre_norm(ctx_tgt, true, /*masked*/ is_mem_shared); llama_set_embeddings_pre_norm(ctx_dft, true, /*masked*/ true); pending_h.assign(n_seq, std::vector(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 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(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(config.params, n_seq)); + } else { + impls.push_back(std::make_unique(config.params, n_seq)); + } break; } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { diff --git a/llama.cpp/ggml/src/ggml-backend.cpp b/llama.cpp/ggml/src/ggml-backend.cpp index c2601b1..752e7ca 100644 --- a/llama.cpp/ggml/src/ggml-backend.cpp +++ b/llama.cpp/ggml/src/ggml-backend.cpp @@ -276,6 +276,14 @@ void ggml_backend_tensor_get_async(ggml_backend_t backend, const struct ggml_ten GGML_ASSERT(backend); GGML_ASSERT(tensor); GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + if (offset + size > ggml_nbytes(tensor)) { // opencoti bug-858 diag (REMOVE after root-cause) + fprintf(stderr, "TENSOR_GET_ASYNC_OOB: name='%s' type=%d ne=[%lld,%lld,%lld,%lld] nbytes=%zu offset=%zu size=%zu\n", + tensor->name, (int) tensor->type, + (long long) tensor->ne[0], (long long) tensor->ne[1], + (long long) tensor->ne[2], (long long) tensor->ne[3], + ggml_nbytes(tensor), offset, size); + fflush(stderr); + } GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); if (backend->iface.get_tensor_async == NULL) { diff --git a/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/llama.cpp/ggml/src/ggml-cuda/fattn.cu index e53ac9c..e63be74 100644 --- a/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -42,13 +42,6 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * Q = dst->src[0]; - if constexpr (ncols2 <= 8) { - if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); - return; - } - } - if constexpr (ncols2 <= 16) { if (Q->ne[1] <= 16/ncols2) { ggml_cuda_flash_attn_ext_mma_f16_case(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 slotK(pool, (size_t)n_slots * k_slot); - ggml_cuda_pool_alloc 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 scrK(pool, bulk ? kblk : 1); + ggml_cuda_pool_alloc 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<<>>( - (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 + // <<>> 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<<>>( + (const char *)Q->data, sK, + mask ? (const char *)mask->data + (size_t)kv0*mask->nb[0] : nullptr, + lse_t, scale, logit_softcap, this_kv, gqa, head_dim, n_head, n_q, + Q->nb[1], Q->nb[2], Q->nb[3], k_row, k_nb2, k_nb2 * n_head_kv, + mask ? mask->nb[0] : 0, mask ? mask->nb[1] : 0); + CUDA_CHECK(cudaGetLastError()); + } } } else { // Empty tile (this_kv==0): the kernel writes lse=-inf and never @@ -1555,7 +1633,8 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor Q->nb[1], Q->nb[2], Q->nb[3], k_row, 0, 0, 0, 0); CUDA_CHECK(cudaGetLastError()); } - CUDA_CHECK(cudaEventRecord(compute_done[slot], stream)); + if (!single_stream) CUDA_CHECK(cudaEventRecord(compute_done[slot], stream)); + slot_live[slot] = true; // #587: slot holds an un-retired writer for the next op/layer } } else // S3a: single compute stream. The per-tile slot lift, the FA over the @@ -1569,8 +1648,8 @@ void ggml_cuda_streaming_flash_attn(ggml_backend_cuda_context & ctx, ggml_tensor const int kv0 = t * tile_kv_full; const int this_kv = kv0 >= n_kv ? 0 : min(tile_kv_full, n_kv - kv0); float * lse_t = lse_tiles.ptr + (size_t)t * n_rows; - char * sK = slotK.ptr + (size_t)slot * k_slot; - char * sV = slotV.ptr + (size_t)slot * v_slot; + char * sK = slotK_base + (size_t)slot * k_slot; + char * sV = slotV_base + (size_t)slot * v_slot; if (this_kv > 0) { const size_t k_nb2 = (size_t)this_kv * k_row; diff --git a/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 4db22f3..bad56a5 100644 --- a/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3344,6 +3344,23 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { continue; } + // opencoti-hook: f5-rolling-kv — bug-1841 (#588/#587). The POSITION_WINDOW streaming FA op + // (ggml_cuda_streaming_flash_attn, fattn.cu) uses a copy/compute double-buffer ONLY at + // prefill (n_q>1): a SEPARATE copy stream + cross-stream copy_done/compute_done event + // waits, which cannot be captured into a CUDA graph ("dependency created on uncaptured + // work in another stream", fattn.cu:1528). At DECODE (n_q==1, src[0]->ne[1]==1) the op + // now runs single-stream (all lift/FA/lse on ctx.stream(), no cross-stream events) → + // fully graph-capturable, which is the actual decode-tps lever (the streaming path was + // launch/sync-bound, PCIe-independent, with graphs off). So disable graphs ONLY for the + // prefill overlap path (n_q>1, which is never graph-captured anyway — belt-and-suspenders); + // decode keeps CUDA graphs. See buglog bug-1841, docs/features/rolling_kv.md. + if (node->op == GGML_OP_STREAMING_FLASH_ATTN && node->src[0] && node->src[0]->ne[1] > 1) { + use_cuda_graph = false; +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to rolling-KV streaming FA (prefill overlap)\n", __func__); +#endif + } + if (node->src[0] && node->src[0]->buffer && ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft)) { use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture #ifndef NDEBUG @@ -4617,7 +4634,41 @@ static enum ggml_status GGML_CALL ggml_backend_cuda_graph_compute(ggml_backend_t if (graph_compatible) { const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); - if (!graph->warmup_complete) { + // opencoti F5 bug-1840 fix (a) — graph-live decode. The stock 0.10.3 + // warmup machine only engages graphs after TWO consecutive calls with + // NO property change, and RESETS to eager on any subsequent change. + // During decode the KV cache ne[1] grows one cell/step, so + // ggml_cuda_graph_update_required reports "changed" every step → + // warmup never completes → graphs never engage (measured: + // cudaGraphLaunch=0 across every decode profile). Mainline llama.cpp + // instead keeps graphs live across a growing KV via cudaGraphExecUpdate + // (already implemented here as ggml_cuda_graph_update_executable, 3440: + // cudaGraphExecUpdate + topology-change re-instantiate fallback). Under + // OPENCOTI_GRAPH_LIVE, complete warmup after ONE stabilization call, then + // stay on graphs and drive recapture-record + exec-update on each property + // change instead of dropping to eager. graph_key (nodes[0]) is stable + // across steps when llama.cpp reuses the graph result (can_reuse), so the + // per-step cost is a CPU recapture-record + cheap exec-update, not a + // re-instantiate. Default-off ⇒ byte-identical to stock. + // opencoti #590/bug-858: default-ON was TESTED and REVERTED — the + // bs2 A/B (graph-live on vs off) showed NO MTP benefit (35B n2 + // 233≈234, n3 still collapsed 156≈158; spec cells inert) and a + // REGRESSION on base decode (35B 201 vs 212) + A4B n3 (189 vs 209). + // The n=3 spec collapse is NOT CUDA-graph thrash — it is per-step + // spec-loop overhead (host readback/sync in the un-fused draft loop); + // graph-live is the wrong lever. Keep the knob env-gated, default-off. + static const bool graph_live = [] { + const char * e = getenv("OPENCOTI_GRAPH_LIVE"); + return e != nullptr && atoi(e) == 1; + }(); + if (graph_live) { + if (!graph->warmup_complete) { + graph->warmup_complete = true; // one stabilization call, then engage + } else { + use_cuda_graph = true; + cuda_graph_update_required = properties_changed || graph->instance == nullptr; + } + } else if (!graph->warmup_complete) { // Warmup: need at least 2 calls with no property change on the 2nd call if (!properties_changed) { graph->warmup_complete = true; diff --git a/llama.cpp/ggml/src/ggml-cuda/mmvf.cu b/llama.cpp/ggml/src/ggml-cuda/mmvf.cu index 09d95f3..4051188 100644 --- a/llama.cpp/ggml/src/ggml-cuda/mmvf.cu +++ b/llama.cpp/ggml/src/ggml-cuda/mmvf.cu @@ -800,6 +800,14 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 switch (type) { case GGML_TYPE_F32: if (GGML_CUDA_CC_IS_NVIDIA(cc)) { +#ifdef GGML_USE_TINYBLAS + // opencoti-hook: mtp-verify-mmvf (bug-858/#609) — upstream caps NVIDIA f32 at + // ne11<=3 because past that its cuBLAS fallback wins; TinyBLAS builds have no + // cuBLAS and the fallback is a single-block tinyblasGE (~174us vs ~1.2us for + // the fused-GDN dot GEMVs in a 4-token MTP verify batch => 35% of decode wall + // at spec depth 3). Use the kernel's real max (8, the non-NVIDIA default below). + return ne11 <= 8; +#endif // GGML_USE_TINYBLAS if (ampere_mma_available(cc)) { return ne11 <= 3; } diff --git a/llama.cpp/include/llama.h b/llama.cpp/include/llama.h index a0a435c..e6517be 100644 --- a/llama.cpp/include/llama.h +++ b/llama.cpp/include/llama.h @@ -424,6 +424,9 @@ extern "C" { ggml_abort_callback abort_callback; void * abort_callback_data; + // opencoti bug-858: dual-context MTP — target context whose KV the draft context shares (nullptr if none) + struct llama_context * ctx_other; + // Keep the booleans together and at the end of the struct to avoid misalignment during copy-by-value. bool embeddings; // if true, extract embeddings (together with logits) bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU @@ -617,6 +620,8 @@ extern "C" { LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx); LLAMA_API llama_memory_t llama_get_memory (const struct llama_context * ctx); + // opencoti bug-858: dual-context MTP — the target context this context shares KV with (nullptr if none) + LLAMA_API struct llama_context * llama_get_ctx_other(struct llama_context * ctx); LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model); @@ -1052,6 +1057,19 @@ extern "C" { float * out_logits, float * out_h_prev_last); + // opencoti fused-NextN (#590): fuse N per-step NextN draft decodes into one graph on a + // DECODER_MTP (Qwen NextN) draft context. GREEDY-ONLY (no p_min early-stop). Returns 0 on + // success. last_token must already be resident in ctx's cache at attn_pos (seed decode first). + LLAMA_API int32_t llama_decode_mtp_fused_nextn( + struct llama_context * ctx, + llama_seq_id seq_id, + llama_pos attn_pos, + llama_token last_token, + const float * h_prev, + int32_t n_steps, + llama_token * out_drafts, + float * out_h_prev_last); + // Set the number of threads used for decoding // n_threads is the number of threads used for generation (single token) // n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens) diff --git a/llama.cpp/src/llama-context.cpp b/llama.cpp/src/llama-context.cpp index 9dbc574..7d27212 100644 --- a/llama.cpp/src/llama-context.cpp +++ b/llama.cpp/src/llama-context.cpp @@ -155,6 +155,7 @@ llama_context::llama_context( cparams.cb_eval_user_data = params.cb_eval_user_data; cparams.ctx_type = params.ctx_type; + cparams.ctx_other = params.ctx_other; // opencoti bug-858: dual-context MTP KV sharing // Initialize backend samplers here so they are part of the sampling graph // before the reserve passes run later in this function. This avoids a later @@ -398,6 +399,7 @@ llama_context::llama_context( /*.type_v =*/ params.type_v, /*.swa_full =*/ params.swa_full, /*.ctx_type= */ cparams.ctx_type, + /*.mem_other=*/ cparams.ctx_other ? llama_get_memory(cparams.ctx_other) : nullptr, // opencoti bug-858 }; // opencoti bug-1342: pass 1 forces a minimal resident window (measure mode). @@ -498,6 +500,7 @@ llama_context::llama_context( /*.type_v =*/ params.type_v, /*.swa_full =*/ params.swa_full, /*.ctx_type= */ cparams.ctx_type, + /*.mem_other=*/ cparams.ctx_other ? llama_get_memory(cparams.ctx_other) : nullptr, // opencoti bug-858 }; memory.reset(model.create_memory(params_mem2, cparams)); @@ -576,6 +579,8 @@ void llama_context::sched_reserve() { gf_res_prev.reset(new llm_graph_result(max_nodes)); gf_res_reserve.reset(new llm_graph_result(max_nodes)); + // opencoti #590/bug-858: sched_nextn + gf_res_prev_nextn are created lazily in + // decode_mtp_fused_nextn (sized graph_max_nodes*n_steps) — see header note. sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); @@ -2012,7 +2017,17 @@ int llama_context::decode(const llama_batch & batch_inp) { const auto & hparams = model.hparams; const int64_t n_vocab = vocab.n_tokens(); - const int64_t n_embd = hparams.n_embd_inp(); + // opencoti bug-858 (dual-context MTP feed width): the gemma4-assistant drafter runs as ctx_dft + // (LLAMA_CONTEXT_TYPE_MTP) and consumes the TARGET's backbone hidden — n_embd_out wide (2816) — + // as its input embd, NOT its own n_embd (1024). n_embd_inp() returns 1024, so the batch→ubatch + // split (llama-batch.cpp: memcpy stride n_embd) would carry 1024-wide rows while + // llm_graph_input_embd::set_input writes n_embd_out(2816) → overread of uninitialised staging → + // partial NaN in inp_h → 0% accept. Size the batch by n_embd_out for such a context. Guarded to + // the genuine mismatch (n_embd_out_impl != n_embd), so normal models stay byte-identical. + const int64_t n_embd = (cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_embd_out_impl > 0 && + hparams.n_embd_out_impl != hparams.n_embd) + ? (int64_t) hparams.n_embd_out() + : (int64_t) hparams.n_embd_inp(); // when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; @@ -2836,7 +2851,13 @@ int32_t llama_context::decode_mtp_sync( // graph; each step's argmax feeds next step's last_token; h_post -> next step's h_prev. for (int32_t k = 0; k < n_steps; ++k) { data->token[0] = last_token; - data->pos[0] = attn_pos + 1 + (llama_pos) k; + // bug-858: the gemma4-assistant (is_mem_shared) drafts ALL steps from the SAME query + // position — the head is trained to predict multiple future tokens from one position, with + // the recurrence carried by the h_prev hidden chain (NOT the RoPE angle). Incrementing the + // position per step (attn_pos+1+k) hands the head a mistrained RoPE angle from step 1 on, + // degrading step-1+ drafts (aggregate accept ~0.69 vs upstream 0.79). Match b9859 draft_mtp + // is_mem_shared: common_batch_add(batch, id, dp.n_past, ...) — constant dp.n_past = attn_pos+1. + data->pos[0] = attn_pos + 1; std::memcpy(data->embd.data(), h_prev, n_bb * sizeof(float)); llama_memory_context_ptr mctx = kv_iswa->init_mtp(seq_id, ub); @@ -2978,8 +2999,12 @@ int32_t llama_context::decode_mtp_fused( data->token[0] = last_token; std::memcpy(data->embd.data(), h_prev, n_bb * sizeof(float)); + // bug-858: constant query position for ALL draft steps (see decode_mtp_sync note). The + // gemma4-assistant (is_mem_shared) is trained to predict every future draft token from the + // SAME position (dp.n_past = attn_pos+1); recurrence lives in the h_prev chain, not the RoPE + // angle. The old attn_pos+1+k mistrained step-1+ RoPE -> accept 0.69 vs upstream 0.79. for (int32_t k = 0; k < n_steps; ++k) { - data->pos[k] = attn_pos + 1 + (llama_pos) k; + data->pos[k] = attn_pos + 1; } llama_memory_context_ptr mctx = kv_iswa->init_mtp(seq_id, ub); @@ -3016,6 +3041,192 @@ int32_t llama_context::decode_mtp_fused( return 0; } +// opencoti fused-NextN (#590 / bug-858): fuse the N per-step NextN draft decodes into ONE graph on +// the draft context (ctx_dft, type LLAMA_CONTEXT_TYPE_MTP → LLM_GRAPH_TYPE_DECODER_MTP). This is the +// throughput fix: the un-fused path pays a host sync per draft step; here one graph emits N greedy +// drafts with a single sync. Twin of decode_mtp_fused (Gemma assistant) but for the UNIFIED cache — +// so no mtp_assistant / sched_mtp / kv_iswa. The N fused steps read-only cross-attend the frozen +// prefix (build_attn_readonly_nextn, Option A); recurrence is carried by the on-device hidden chain. +// GREEDY-ONLY: no p_min early-stop (the caller reconciles draft length). See +// docs/features/fused_nextn_mtp.md. Returns 0 on success, negative on error. +int32_t llama_context::decode_mtp_fused_nextn( + llama_seq_id seq_id, + llama_pos attn_pos, + llama_token last_token, + const float * h_prev, + int32_t n_steps, + llama_token * out_drafts, + float * out_h_prev_last) { + if (!memory) { + LLAMA_LOG_ERROR("%s: context has no KV memory\n", __func__); + return -2; + } + // NextN self-spec targets are full-attention → unified cache. (iSWA NextN is not a thing; + // the Gemma assistant path is decode_mtp_fused.) + auto * kv = dynamic_cast(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(); + 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 ubatches; + ubatches.push_back(ub); + auto mctx = std::make_unique(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 drafts((size_t) n_steps); + ggml_backend_tensor_get(t_arg, drafts.data(), 0, (size_t) n_steps * sizeof(int32_t)); + for (int32_t k = 0; k < n_steps; ++k) { + out_drafts[k] = (llama_token) drafts[(size_t) k]; + } + + if (out_h_prev_last) { + ggml_tensor * t_post = res->get_embd(); + GGML_ASSERT(t_post); + ggml_backend_tensor_get(t_post, out_h_prev_last, 0, n_embd * sizeof(float)); + } + + return 0; +} + // opencoti F5 M6-S4 mtp (P4): submit one async MTP draft request to the worker. At most one // in-flight request per context (returns -7 if a prior request was not yet waited). int32_t llama_context::decode_mtp_async( @@ -3255,6 +3466,10 @@ llm_graph_params llama_context::graph_params( /*.n_outputs =*/ n_outputs, /*.cb =*/ graph_get_cb(), /*.res =*/ res, + // opencoti fused-NextN (#590): the DECODER_MTP draft graph fuses N greedy steps when the + // fused driver sets mtp_fused_steps=N (reset to 1 after). Every other gtype — and the + // per-step AR NextN loop — keeps n_mtp_steps=1 (single-step build_one_step). + /*.n_mtp_steps =*/ (gtype == LLM_GRAPH_TYPE_DECODER_MTP ? mtp_fused_steps : 1), }; } @@ -4330,6 +4545,7 @@ llama_context_params llama_context_default_params() { /*.type_v =*/ GGML_TYPE_F16, /*.abort_callback =*/ nullptr, /*.abort_callback_data =*/ nullptr, + /*.ctx_other =*/ nullptr, // opencoti bug-858: dual-context MTP /*.embeddings =*/ false, /*.offload_kqv =*/ true, /*.no_perf =*/ true, @@ -4415,8 +4631,13 @@ llama_context * llama_init_from_model( model->hparams.pooling_type, params.pooling_type); } + // opencoti bug-858: a dual-context gemma4-assistant drafter (ctx_other set) legitimately has + // nextn_predict_layers == 0 — its "MTP layers" are the TARGET's backbone K/V, shared in-place + // via ctx_other (a KV-less assistant reads the target cells). Only reject a STANDALONE MTP + // context (ctx_other == nullptr) that lacks nextn layers; Qwen NextN (layers > 0) is unaffected. if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && - model->hparams.nextn_predict_layers == 0) { + model->hparams.nextn_predict_layers == 0 && + params.ctx_other == nullptr) { LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__); return nullptr; } @@ -4572,6 +4793,21 @@ float * llama_get_embeddings_pre_norm_ith(llama_context * ctx, int32_t i) { return ctx->get_embeddings_pre_norm_ith(i); } +// opencoti #590/bug-858: return the on-device backend-sampled draft token (sampling.sampled), produced +// during the draft decode by the offloaded sampler chain attached via llama_set_sampler (top_k-only → +// greedy argmax, matching common_sampler_sample). This is the cheap read that replaces the ~152k-vocab +// common_sampler_sample D2H + CPU softmax/sort (MEASURED 0.98 ms/step). Returns LLAMA_TOKEN_NULL (-1) +// when backend sampling is inactive → the AR draft loop falls back to the exact CPU sampler. +llama_token llama_context::get_draft_greedy_ith(int32_t i) { + return get_sampled_token_ith(i); +} + +llama_token llama_get_draft_greedy_ith(llama_context * ctx, int32_t i) { + ctx->synchronize(); + + return ctx->get_draft_greedy_ith(i); +} + bool llama_set_sampler(llama_context * ctx, llama_seq_id seq_id, llama_sampler * smpl) { return ctx->set_sampler(seq_id, smpl); } @@ -4667,6 +4903,10 @@ llama_memory_t llama_get_memory(const struct llama_context * ctx) { return ctx->get_memory(); } +llama_context * llama_get_ctx_other(struct llama_context * ctx) { + return ctx->get_cparams().ctx_other; // opencoti bug-858: dual-context MTP +} + void llama_memory_clear(llama_memory_t mem, bool data) { if (!mem) { return; @@ -4955,6 +5195,22 @@ int32_t llama_decode_mtp( return ctx->decode_mtp(seq_id, attn_pos, last_token, h_prev, n_steps, out_drafts, out_logits, out_h_prev_last); } +int32_t llama_decode_mtp_fused_nextn( + llama_context * ctx, + llama_seq_id seq_id, + llama_pos attn_pos, + llama_token last_token, + const float * h_prev, + int32_t n_steps, + llama_token * out_drafts, + float * out_h_prev_last) { + if (!ctx) { + LLAMA_LOG_ERROR("%s: ctx is NULL\n", __func__); + return -1; + } + return ctx->decode_mtp_fused_nextn(seq_id, attn_pos, last_token, h_prev, n_steps, out_drafts, out_h_prev_last); +} + // // perf // diff --git a/llama.cpp/src/llama-context.h b/llama.cpp/src/llama-context.h index 7f85a13..b6b0efb 100644 --- a/llama.cpp/src/llama-context.h +++ b/llama.cpp/src/llama-context.h @@ -94,6 +94,11 @@ struct llama_context { float * get_embeddings_pre_norm(); float * get_embeddings_pre_norm_ith(int32_t i); + // opencoti #590/bug-858: on-device greedy draft token from the previous MTP decode's t_argmax + // (returns -1 if the graph published none). Lets the AR spec draft loop skip the full-vocab + // common_sampler_sample host round-trip (~0.98 ms/step) when the draft sampler is argmax-greedy. + llama_token get_draft_greedy_ith(int32_t i); + llama_token * get_sampled_tokens() const; llama_token get_sampled_token_ith(int32_t idx); @@ -182,6 +187,17 @@ struct llama_context { llama_token * out_drafts, float * out_h_prev_last); + // opencoti fused-NextN (#590): fused N-step greedy draft on a DECODER_MTP (Qwen NextN) draft + // context using the unified cache. Twin of decode_mtp_fused minus mtp_assistant/sched_mtp/iSWA. + int32_t decode_mtp_fused_nextn( + llama_seq_id seq_id, + llama_pos attn_pos, + llama_token last_token, + const float * h_prev, + int32_t n_steps, + llama_token * out_drafts, + float * out_h_prev_last); + // opencoti F5 M6-S4 mtp (P4): async MTP draft pipeline. decode_mtp() (above) is a sync // facade — out_logits!=NULL keeps the in-thread decode_mtp_sync path (the worker streams // no per-step logits); otherwise it submits via decode_mtp_async and blocks in @@ -430,6 +446,19 @@ private: ggml_backend_sched_ptr sched_mtp; llm_graph_result_ptr gf_res_prev_mtp; + // opencoti #590/bug-858: DEDICATED scheduler + result for the fused NextN self-spec draft + // (decode_mtp_fused_nextn). It must NOT share the main sched/gf_res_prev with target decode: + // every interleaved target decode RESETS the main sched (reallocating compute buffers) and + // gf_res_prev, so a fused graph parked there can never satisfy can_reuse AND its tensor memory is + // clobbered → 100% rebuild + CUDA-recapture per draft round (measured: build 0.5ms + comp-launch + // 1.5ms, rebuilds=N/N) AND garbage argmax when reused. A dedicated sched (like the Gemma assistant's + // sched_mtp, but ensure_sched_mtp is assistant-only: needs mtp_assistant + iswa cache) isolates the + // fused graph so can_reuse compares fused-vs-fused (stable within a 256-cell n_kv bucket) and the + // ggml graph + CUDA capture persist across rounds. NextN and Gemma-assistant are mutually exclusive. + ggml_backend_sched_ptr sched_nextn; + llm_graph_result_ptr gf_res_prev_nextn; + int32_t nextn_reserved_steps = 0; + // opencoti F5 M6-S4 mtp (P4 fused): number of draft steps to unroll into the MTP graph for // the NEXT build (decode_mtp_fused sets it to n_steps; the sequential/reserve paths leave it // at 1). mtp_reserved_steps records how many steps the current sched_mtp reserve covers, so diff --git a/llama.cpp/src/llama-cparams.h b/llama.cpp/src/llama-cparams.h index 934adf5..8f0b36b 100644 --- a/llama.cpp/src/llama-cparams.h +++ b/llama.cpp/src/llama-cparams.h @@ -112,4 +112,9 @@ struct llama_cparams { ggml_backend_sched_eval_callback cb_eval; void * cb_eval_user_data; + + // opencoti bug-858: dual-context MTP — the target context whose KV cells the + // draft (MTP) context shares IN-PLACE (upstream cparams.ctx_other + mem_other). + // nullptr for ordinary contexts. + llama_context * ctx_other; }; diff --git a/llama.cpp/src/llama-ext.h b/llama.cpp/src/llama-ext.h index edfa71c..9d57c33 100644 --- a/llama.cpp/src/llama-ext.h +++ b/llama.cpp/src/llama-ext.h @@ -104,3 +104,8 @@ LLAMA_API float * llama_get_embeddings_pre_norm (struct llama_context * ctx); // LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i); LLAMA_API float * llama_get_embeddings_pre_norm_ith(struct llama_context * ctx, int32_t i); + +// opencoti #590/bug-858: on-device greedy draft token from the previous MTP decode's in-graph argmax +// (t_argmax). Returns the token id, or -1 if none was published. Lets the AR spec draft loop skip the +// ~0.98 ms/step full-vocab common_sampler_sample (logit D2H + CPU softmax/sort) when p_min==0. +LLAMA_API llama_token llama_get_draft_greedy_ith(struct llama_context * ctx, int32_t i); diff --git a/llama.cpp/src/llama-graph.cpp b/llama.cpp/src/llama-graph.cpp index e0d5e2a..e0d712e 100644 --- a/llama.cpp/src/llama-graph.cpp +++ b/llama.cpp/src/llama-graph.cpp @@ -125,14 +125,28 @@ void llm_graph_input_mtp::set_input(const llama_ubatch * ubatch) { // opencoti F5 M6-S4 mtp (P4 fused): step k's RoPE query position = ubatch->pos[k] // (= attn_pos + 1 + k, filled by decode_mtp_fused). Single-step path leaves this empty. + // bug-870: mrope archs (qwen35) size each step's pos tensor I32[4] — fill the M-RoPE + // text-token layout [p,p,p,0] (mirrors llm_graph_input_pos::set_input); standard rope = [p]. for (size_t k = 0; k < inp_pos_steps.size(); ++k) { - ggml_backend_tensor_set(inp_pos_steps[k], ubatch->pos + k, 0, sizeof(int32_t)); + ggml_tensor * pk = inp_pos_steps[k]; + const int32_t p = ubatch->pos[k]; + if (pk->ne[0] == 4) { + const int32_t pos4[4] = { p, p, p, 0 }; + ggml_backend_tensor_set(pk, pos4, 0, sizeof(pos4)); + continue; + } + ggml_backend_tensor_set(pk, &p, 0, sizeof(int32_t)); } } // opencoti F5 M6-S4 mtp bool llm_graph_input_mtp::can_reuse(const llm_graph_params & params) { - if (params.gtype != LLM_GRAPH_TYPE_MTP) { + // bug-870/#590: the fused NextN driver (decode_mtp_fused_nextn) builds a DECODER_MTP graph, so + // gating reuse on MTP-only forced a full N-block MoE graph rebuild+realloc on EVERY draft call — + // the per-call build dwarfed the saved host-syncs and made fused SLOWER than the AR loop. Accept + // both MTP (Gemma assistant) and DECODER_MTP (Qwen NextN fused); the attn input's own can_reuse + // independently vetoes reuse when the prefix mask shape changes, so this stays correct. + if (params.gtype != LLM_GRAPH_TYPE_MTP && params.gtype != LLM_GRAPH_TYPE_DECODER_MTP) { return false; } const bool base = inp_last_token && inp_last_token->ne[0] == 1 && inp_h_prev && inp_h_prev->ne[1] == 1; @@ -607,11 +621,20 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { bool res = true; + // opencoti bug-858 (2026-07-05): match upstream b9859 — the kq_mask checks must be + // INDEPENDENT of the k_idxs blocks. In the KV-less gemma4-assistant draft context the + // k/v_idxs are dead nodes (no cache store) and never get buffers, so nesting the mask + // check under them skipped it entirely → the boot-time reserve graph (n_kv = full cache + // width) was reused for every draft decode → set_input filled only the first n_kv + // columns of a full-width mask, leaving stale pool garbage → accept collapse on any + // prompt past one KV pad block (gate10: 0.83@114tok → 0.37@354 → 0.00@1314). // base tensors may not be allocated if there are no non-SWA attention layers if (self_k_idxs && self_k_idxs->buffer) { res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there + } + if (self_kq_mask && self_kq_mask->buffer) { res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); } @@ -619,7 +642,9 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { if (self_k_idxs_swa && self_k_idxs_swa->buffer) { res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; //res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there + } + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); } @@ -2894,7 +2919,32 @@ ggml_tensor * llm_graph_context::build_attn( // null (build_attn for KV asserts this above at :2102), kq_b too; // the M2 gate ensures non-transposed V. See docs/features/advanced_kv.md. ggml_tensor * cur; - if (mctx_cur->get_layer_tactic(il) == llama_kv_cache::headinfer_tactic::GPU_STREAM) { + // opencoti-hook: f5-rolling-kv — bug-1840 (#588). POSITION_WINDOW on the FULL-ATTENTION + // build_attn overload. The position-window residency tactic was previously wired ONLY into + // the iSWA overload (build_attn(llm_graph_input_attn_kv_iswa*), see the mirror ~line 3272); + // full-attention models (Qwen etc. — the ONLY class that actually spills, since iSWA KV is + // bounded) fell through to the get_k/get_v else-branch below, whose window fallback returns + // ggml_concat(window_device, tail_host). That single cross-backend logical tensor makes the + // ggml scheduler round-trip the ENTIRE KV cache host<->device EVERY decode step (measured + // 12.6 GiB/step both ways -> 0.99 tps at a 4.2 GB tail). Routing spilling layers through + // build_attn_mha_position_window keeps the window FA device-resident and streams ONLY the + // tail (one H2D copy, no D2H write-back). Guarded to POSITION_WINDOW + FA + no KQ-bias/MLA/ + // sinks + an active window, so non-spilling layers are byte-identical. See bug-1840, + // docs/features/rolling_kv.md. + const bool window_layer = + cparams.flash_attn && kq_b == nullptr && v_mla == nullptr && sinks == nullptr && + mctx_cur->get_layer_tactic(il) == llama_kv_cache::headinfer_tactic::POSITION_WINDOW && + mctx_cur->headinfer_window_active(il); + if (window_layer) { + ggml_tensor * k_win = mctx_cur->get_k_window(ctx0, il); + ggml_tensor * v_win = mctx_cur->get_v_window(ctx0, il); + ggml_tensor * k_tail = mctx_cur->get_k_tail (ctx0, il); // nullptr when n_kv <= wc + ggml_tensor * v_tail = mctx_cur->get_v_tail (ctx0, il); + const bool cpu_fa_tail = + (k_tail != nullptr) && (q->ne[2] == 1) && mctx_cur->headinfer_tail_is_cpu_fa(); + cur = build_attn_mha_position_window(q, k_win, v_win, k_tail, v_tail, + kq_b, kq_mask, sinks, v_mla, kq_scale, il, cpu_fa_tail); + } else if (mctx_cur->get_layer_tactic(il) == llama_kv_cache::headinfer_tactic::GPU_STREAM) { // opencoti F5 M7 Rolling KV — Rung 2 (W4 M7-D #304), S1. GPU_STREAM // rides the M2 head-split: the device head-group runs a resident FA, // the host (spilled) head-group runs the streaming-FA op, and the two @@ -3036,6 +3086,28 @@ ggml_tensor * llm_graph_context::build_attn_mtp( return cur; } +// opencoti fused-NextN (Option A / Gemma-mirror): read-only cross-attention into the frozen prefix +// KV on the UNIFIED cache. Twin of build_attn_mtp for the non-iSWA llm_graph_input_attn_kv used by +// Qwen NextN drafters. Reads mctx->get_k/get_v(il) with the input's kq_mask — which exposes the +// full prefix 0..pmax because get_n_kv reports the cache's used extent, not the mtp_slot_info cell +// count. Writes NO draft KV; recurrence is carried by the hidden-state chain. wo is applied by the +// caller (the qwen35* MTP block multiplies by the gate then projects wo). See +// docs/features/fused_nextn_mtp.md delta #6. +ggml_tensor * llm_graph_context::build_attn_readonly_nextn( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + float kq_scale, + int il) const { + const auto * mctx_cur = inp->mctx; + + ggml_tensor * kq_mask = inp->get_kq_mask(); + + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = mctx_cur->get_v(ctx0, il); + + return build_attn_mha(q_cur, k, v, nullptr, kq_mask, nullptr, nullptr, kq_scale, il); +} + static std::unique_ptr build_attn_inp_k_impl( ggml_context * ctx0, const llama_ubatch & ubatch, diff --git a/llama.cpp/src/llama-graph.h b/llama.cpp/src/llama-graph.h index 0c13245..d3a3499 100644 --- a/llama.cpp/src/llama-graph.h +++ b/llama.cpp/src/llama-graph.h @@ -1041,6 +1041,17 @@ struct llm_graph_context { int64_t kv_n_head_v, bool use_k_as_v) const; + // opencoti fused-NextN (Option A): read-only cross-attention into the frozen prefix KV on the + // UNIFIED cache. Twin of build_attn_mtp for llm_graph_input_attn_kv (non-iSWA). Reads + // mctx->get_k/get_v(il) with the input's kq_mask (which exposes 0..pmax via get_n_kv's used + // extent), runs build_attn_mha, writes NO draft KV. wo is applied by the caller. + // See docs/features/fused_nextn_mtp.md. + ggml_tensor * build_attn_readonly_nextn( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + float kq_scale, + int il) const; + llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; ggml_tensor * build_attn( diff --git a/llama.cpp/src/llama-kv-cache-iswa.cpp b/llama.cpp/src/llama-kv-cache-iswa.cpp index 7611271..9628df1 100644 --- a/llama.cpp/src/llama-kv-cache-iswa.cpp +++ b/llama.cpp/src/llama-kv-cache-iswa.cpp @@ -33,8 +33,17 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( uint32_t n_pad, const layer_filter_cb & filter, const layer_reuse_cb & reuse, + llama_memory_t mem_other, + const layer_share_cb & share, uint32_t sparse_attn_block_size) : hparams(model.hparams), unified(unified) { + // opencoti bug-858 dual-context MTP — split the source iSWA cache into its base/swa halves + // so each inner kv_cache shares with the matching half. GUARD: both stay nullptr when + // mem_other == nullptr (every existing caller) ⇒ the inner ctors receive nullptr ⇒ + // byte-identical standalone allocation. + llama_memory_t mem_other_base = mem_other ? static_cast(mem_other)->get_base() : nullptr; + llama_memory_t mem_other_swa = mem_other ? static_cast(mem_other)->get_swa() : nullptr; + // chain filters const layer_filter_cb filter_base = [&](int32_t il) { if (filter && !filter(il)) { @@ -73,7 +82,9 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( v_trans, offload, unified, size_base, kv_size_initial, slot_shrink_idle_ms, headinfer_gpu_heads_frac, vram_target_mib, pcie_bw_gbps, kv_residency_mode, compute_reserve_mib, window_measure_pass, - n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse, sparse_attn_block_size); + n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse, + // opencoti bug-858 dual-context MTP — forward the base-half source + share selector. + mem_other_base, share, sparse_attn_block_size); LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); @@ -82,7 +93,9 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( v_trans, offload, unified, size_swa, kv_size_initial, slot_shrink_idle_ms, headinfer_gpu_heads_frac, vram_target_mib, pcie_bw_gbps, kv_residency_mode, compute_reserve_mib, window_measure_pass, - n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, filter_swa, reuse, sparse_attn_block_size); + n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, filter_swa, reuse, + // opencoti bug-858 dual-context MTP — forward the swa-half source + share selector. + mem_other_swa, share, sparse_attn_block_size); } void llama_kv_cache_iswa::clear(bool data) { diff --git a/llama.cpp/src/llama-kv-cache-iswa.h b/llama.cpp/src/llama-kv-cache-iswa.h index 166b4eb..06fd1ae 100644 --- a/llama.cpp/src/llama-kv-cache-iswa.h +++ b/llama.cpp/src/llama-kv-cache-iswa.h @@ -52,6 +52,13 @@ public: uint32_t n_pad, const layer_filter_cb & filter, const layer_reuse_cb & reuse, + // opencoti bug-858 dual-context MTP — source iSWA cache; split into + // base/swa and forwarded to the two inner kv_base/kv_swa. nullptr => + // ordinary standalone cache (byte-identical). Threaded with `share`. + llama_memory_t mem_other, + // opencoti bug-858 dual-context MTP — per-layer share selector, + // forwarded verbatim to both inner caches. Used only when mem_other != null. + const layer_share_cb & share, // opencoti #551 sparse-attn — forwarded to both kv_base and // kv_swa. Quest block-selector B_SEL; 0 = off (no side-cache). // See docs/features/sparse_attn.md. diff --git a/llama.cpp/src/llama-kv-cache.cpp b/llama.cpp/src/llama-kv-cache.cpp index 13c8c67..4416d38 100644 --- a/llama.cpp/src/llama-kv-cache.cpp +++ b/llama.cpp/src/llama-kv-cache.cpp @@ -15,6 +15,18 @@ #include #include +// 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 +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(mem_other)), + v_cells_impl(other ? other->v_cells_impl : std::make_shared>()), + 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 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::vectorget_sched(); @@ -2188,6 +2347,12 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) { + // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null) + // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS]. + // Inert on the target/standalone path (other==nullptr). + if (other) { + return; + } // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md // Refresh the idle timestamp on every actual write. shrink_if_idle uses // this to gate the shrink path. @@ -2340,7 +2505,7 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { } ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; @@ -2462,7 +2627,7 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k } ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; @@ -2923,7 +3088,7 @@ uint32_t llama_kv_cache::headinfer_gpu_heads(int32_t il) const { } ggml_tensor * llama_kv_cache::get_k_gpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.gpu_heads > 0 && "get_k_gpu called without an active headinfer split"); @@ -2959,7 +3124,7 @@ ggml_tensor * llama_kv_cache::get_k_gpu(ggml_context * ctx, int32_t il, uint32_t // views it in-place. n_kv is irrelevant — kbounds spans all selectable blocks. ggml_tensor * llama_kv_cache::get_kbounds(ggml_context * ctx, int32_t il, const slot_info & sinfo) const { GGML_UNUSED(ctx); - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; if (layer.kbounds_per_stream.empty()) { return nullptr; @@ -2973,7 +3138,7 @@ ggml_tensor * llama_kv_cache::get_kbounds(ggml_context * ctx, int32_t il, const // re-writes it only every cparams.sparse_attn_refresh decode tokens; FA-VEC reads it every step. ggml_tensor * llama_kv_cache::get_block_sel(ggml_context * ctx, int32_t il, const slot_info & sinfo) const { GGML_UNUSED(ctx); - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; if (layer.block_sel_per_stream.empty()) { return nullptr; @@ -2982,7 +3147,7 @@ ggml_tensor * llama_kv_cache::get_block_sel(ggml_context * ctx, int32_t il, cons } ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.gpu_heads > 0 && !layer.k_cpu_per_stream.empty() @@ -2998,6 +3163,18 @@ ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { ggml_tensor * kc_s = layer.k_cpu_per_stream[s_cache]; + // opencoti #586 diag: report the ACTUAL allocated buffer buft (the boot WARN + // reports the *intended* cpu_buft, masking a silent cudaMallocHost fallback to + // pageable CPU — the H2D tail stream then runs at ~half PCIe / non-overlapping). + { + static bool opencoti_actual_tail_buft_logged = false; + if (!opencoti_actual_tail_buft_logged) { + opencoti_actual_tail_buft_logged = true; + LLAMA_LOG_WARN("%s: ACTUAL tail: data=%p buffer=%p buft='%s'\n", __func__, + (void *) kc_s->data, (void *) kc_s->buffer, + kc_s->buffer ? ggml_backend_buft_name(ggml_backend_buffer_get_type(kc_s->buffer)) : "NULL-BUFFER"); + } + } return ggml_view_4d(ctx, kc_s, n_embd_head_k, n_head_cpu, n_kv, 1, ggml_row_size(kc_s->type, n_embd_head_k), @@ -3014,7 +3191,7 @@ ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t } ggml_tensor * llama_kv_cache::get_v_gpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.gpu_heads > 0 && "get_v_gpu called without an active headinfer split"); // M2 requires non-transposed V (FA on). The split-active branch in @@ -3045,7 +3222,7 @@ ggml_tensor * llama_kv_cache::get_v_gpu(ggml_context * ctx, int32_t il, uint32_t } ggml_tensor * llama_kv_cache::get_v_cpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.gpu_heads > 0 && !layer.v_cpu_per_stream.empty() @@ -3085,7 +3262,7 @@ ggml_tensor * llama_kv_cache::get_v_cpu(ggml_context * ctx, int32_t il, uint32_t // ne[1] == window_cells, so its dim-3 (stream/position) stride is computed off // k_s->ne[1] (the audit point — NOT kv_size); the host tail likewise off kc_s->ne[1]. ggml_tensor * llama_kv_cache::get_k_window(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.window_cells > 0 && "get_k_window called without an active position window"); @@ -3113,7 +3290,7 @@ ggml_tensor * llama_kv_cache::get_k_window(ggml_context * ctx, int32_t il, uint3 } ggml_tensor * llama_kv_cache::get_k_tail(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.window_cells > 0 && !layer.k_cpu_per_stream.empty() @@ -3125,12 +3302,35 @@ ggml_tensor * llama_kv_cache::get_k_tail(ggml_context * ctx, int32_t il, uint32_ const uint64_t n_embd_k_gqa2 = (uint64_t) n_head_kv * n_embd_head_k; const uint32_t wc = layer.window_cells; const uint32_t n_tail = n_kv > wc ? n_kv - wc : 0; + { + static bool opencoti_gkt_entry_logged = false; + if (!opencoti_gkt_entry_logged) { + opencoti_gkt_entry_logged = true; + ggml_tensor * kc0 = layer.k_cpu_per_stream.empty() ? nullptr : layer.k_cpu_per_stream[0]; + LLAMA_LOG_WARN("%s: ENTRY il=%d n_kv=%u wc=%u n_tail=%u kc0=%p kc0.buffer=%p buft=%s\n", + __func__, il, n_kv, wc, n_tail, (void *)(kc0 ? kc0->data : nullptr), + (void *)(kc0 ? kc0->buffer : nullptr), + (kc0 && kc0->buffer) ? ggml_backend_buft_name(ggml_backend_buffer_get_type(kc0->buffer)) : "NULL"); + } + } if (n_tail == 0) { return nullptr; // no overflow → fully resident window, single-region FA } auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { ggml_tensor * kc_s = layer.k_cpu_per_stream[s_cache]; + // opencoti #586 diag: report the ACTUAL allocated buffer buft (the boot WARN + // reports the *intended* cpu_buft, masking a silent cudaMallocHost fallback to + // pageable CPU — the H2D tail stream then runs at ~half PCIe / non-overlapping). + { + static bool opencoti_actual_tail_buft_logged = false; + if (!opencoti_actual_tail_buft_logged) { + opencoti_actual_tail_buft_logged = true; + LLAMA_LOG_WARN("%s: ACTUAL tail: data=%p buffer=%p buft='%s'\n", __func__, + (void *) kc_s->data, (void *) kc_s->buffer, + kc_s->buffer ? ggml_backend_buft_name(ggml_backend_buffer_get_type(kc_s->buffer)) : "NULL-BUFFER"); + } + } return ggml_view_4d(ctx, kc_s, n_embd_head_k, n_head_kv, n_tail, 1, ggml_row_size(kc_s->type, n_embd_head_k), @@ -3147,7 +3347,7 @@ ggml_tensor * llama_kv_cache::get_k_tail(ggml_context * ctx, int32_t il, uint32_ } ggml_tensor * llama_kv_cache::get_v_window(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.window_cells > 0 && "get_v_window called without an active position window"); GGML_ASSERT(!v_trans && "position window requires non-transposed V"); @@ -3176,7 +3376,7 @@ ggml_tensor * llama_kv_cache::get_v_window(ggml_context * ctx, int32_t il, uint3 } ggml_tensor * llama_kv_cache::get_v_tail(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; GGML_ASSERT(layer.window_cells > 0 && !layer.v_cpu_per_stream.empty() @@ -3230,7 +3430,7 @@ uint32_t llama_kv_cache::headinfer_window_cells(int32_t il) const { } ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; @@ -3413,7 +3613,7 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm } ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const { - const int32_t ikv = map_layer_ids.at(il); + const int32_t ikv = mli_at_checked(map_layer_ids, il, __func__); const auto & layer = layers[ikv]; @@ -4365,6 +4565,10 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) { } ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_context * lctx) const { + // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null) + // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS]. + // Inert on the target/standalone path (other==nullptr). + GGML_ASSERT(!other && "shift on a shared draft cache is a bug"); auto * ctx = res->get_ctx(); auto * gf = res->get_gf(); @@ -4467,6 +4671,12 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co } void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null) + // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS]. + // Inert on the target/standalone path (other==nullptr). + if (other) { + return; + } GGML_UNUSED(flags); io.write(&n_stream, sizeof(n_stream)); @@ -4520,6 +4730,12 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // opencoti bug-858 (bug-2097): shared draft cache (gemma4-assistant ctx_dft, other!=null) + // is READ-ONLY over the target's cells — mirror upstream b9859 [TAG_KV_CACHE_SHARE_CELLS]. + // Inert on the target/standalone path (other==nullptr). + if (other) { + return; + } GGML_UNUSED(flags); GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); diff --git a/llama.cpp/src/llama-kv-cache.h b/llama.cpp/src/llama-kv-cache.h index 924f7fe..7b3afc3 100644 --- a/llama.cpp/src/llama-kv-cache.h +++ b/llama.cpp/src/llama-kv-cache.h @@ -146,6 +146,14 @@ public: llama_swa_type swa_type, const layer_filter_cb & filter, const layer_reuse_cb & reuse, + // opencoti bug-858 dual-context MTP — source cache whose cells + + // per-layer K/V tensors the draft cache shares IN PLACE. nullptr => + // ordinary standalone cache (byte-identical to pre-bug-858). Threaded + // with `share`; both are null for every target / existing caller. + llama_memory_t mem_other, + // opencoti bug-858 dual-context MTP — per-layer share selector; see + // llama_memory_i::layer_share_cb. Consulted only when mem_other != null. + const layer_share_cb & share, // opencoti #551 sparse-attn — Quest block-selector KV-block // size B_SEL (docs/features/sparse_attn.md). 0 = sparse-attn // OFF: no kmin/kmax side-cache allocated (byte-identical). @@ -548,7 +556,20 @@ private: // note: this is not part of the KV state and it's only used to speed-up the find_slot() method std::vector v_heads; - std::vector 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 v_cells;` member. + std::shared_ptr> v_cells_impl; + std::vector & v_cells; // maps from a sequence id to a stream id std::vector seq_to_stream; diff --git a/llama.cpp/src/llama-memory-hybrid-iswa.cpp b/llama.cpp/src/llama-memory-hybrid-iswa.cpp index 7301ad5..3cf8e4a 100644 --- a/llama.cpp/src/llama-memory-hybrid-iswa.cpp +++ b/llama.cpp/src/llama-memory-hybrid-iswa.cpp @@ -65,7 +65,10 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( filter_attn == nullptr ? [&](int32_t il) { return !hparams.is_recurrent(il); } : filter_attn, - nullptr + nullptr, + // opencoti bug-858 dual-context MTP — hybrid-iswa attn cache is standalone: no sharing. + /* mem_other */ nullptr, + /* share */ nullptr )), mem_recr(new llama_memory_recurrent( model, diff --git a/llama.cpp/src/llama-memory-hybrid.cpp b/llama.cpp/src/llama-memory-hybrid.cpp index 86e4134..0fa114d 100644 --- a/llama.cpp/src/llama-memory-hybrid.cpp +++ b/llama.cpp/src/llama-memory-hybrid.cpp @@ -65,7 +65,10 @@ llama_memory_hybrid::llama_memory_hybrid( filter_attn == nullptr ? [&](int32_t il) { return !hparams.is_recurrent(il); } : filter_attn, - nullptr + nullptr, + // opencoti bug-858 dual-context MTP — hybrid attn cache is standalone: no sharing. + /* mem_other */ nullptr, + /* share */ nullptr )), mem_recr(new llama_memory_recurrent( model, diff --git a/llama.cpp/src/llama-memory.h b/llama.cpp/src/llama-memory.h index 3e79826..98e4f26 100644 --- a/llama.cpp/src/llama-memory.h +++ b/llama.cpp/src/llama-memory.h @@ -24,6 +24,10 @@ struct llama_memory_params { bool swa_full; llama_context_type ctx_type; + + // opencoti bug-858: dual-context MTP — source memory whose cells the draft cache + // shares in-place (upstream mem_other). nullptr => ordinary standalone cache. + llama_memory_t mem_other; }; enum llama_memory_status { @@ -77,6 +81,13 @@ struct llama_memory_i { // return negative value to indicate that the layer il should not reuse memory using layer_reuse_cb = std::function; + // 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; + virtual ~llama_memory_i() = default; // split the input batch into a set of ubatches and verify that they can fit into the cache diff --git a/llama.cpp/src/llama-model.cpp b/llama.cpp/src/llama-model.cpp index 3df76f7..d7804be 100644 --- a/llama.cpp/src/llama-model.cpp +++ b/llama.cpp/src/llama-model.cpp @@ -2079,6 +2079,24 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { GGML_ASSERT(hparams.is_swa_any()); + // opencoti bug-858 dual-context MTP (A3) — when THIS model is the + // gemma4-assistant drafter created as ctx_dft (cparams.ctx_other = the + // target context), share the TARGET's KV cells IN PLACE: the assistant has + // no K/V projections, so each of its attention layers reads the target + // backbone's last full (n_layer-1) / SWA (n_layer-2) layer. Mirrors upstream + // b9859 create_memory LLM_ARCH_GEMMA4_ASSISTANT. Inert for every other arch + // and when ctx_other is null → nullptr/nullptr → byte-identical standalone cache. + llama_memory_t mtp_mem_other = nullptr; + llama_kv_cache::layer_share_cb mtp_share = nullptr; + if (arch == LLM_ARCH_GEMMA4_ASSISTANT && cparams.ctx_other != nullptr) { + mtp_mem_other = llama_get_memory(cparams.ctx_other); + mtp_share = [&](int32_t il) { + const llama_model * model_other = llama_get_model(cparams.ctx_other); + const int32_t n_layer_other = (int32_t) llama_model_n_layer(model_other); + return hparams.is_swa(il) ? (n_layer_other - 2) : (n_layer_other - 1); + }; + } + res = new llama_kv_cache_iswa( *this, params.type_k, @@ -2109,6 +2127,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, 1, filter, reuse, + // opencoti bug-858 dual-context MTP (A3) — share the target's KV when + // this is the gemma4-assistant ctx_dft (mtp_mem_other/mtp_share set + // above); nullptr/nullptr for the target + every other arch → identical. + /* mem_other */ mtp_mem_other, + /* share */ mtp_share, // opencoti #551 sparse-attn — B_SEL when enabled, else 0 (off) cparams.sparse_attn_enabled ? cparams.sparse_attn_block_size : 0); } else { @@ -2144,6 +2167,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, hparams.swa_type, filter, nullptr, + // opencoti bug-858 dual-context MTP — target / standalone cache: + // no sharing in A2 (A3 wires the draft cache with a non-null source). + /* mem_other */ nullptr, + /* share */ nullptr, // opencoti #551 sparse-attn — B_SEL when enabled, else 0 (off) cparams.sparse_attn_enabled ? cparams.sparse_attn_block_size : 0); } diff --git a/llama.cpp/src/models/gemma4-assistant.cpp b/llama.cpp/src/models/gemma4-assistant.cpp index 51f935c..cab4fa8 100644 --- a/llama.cpp/src/models/gemma4-assistant.cpp +++ b/llama.cpp/src/models/gemma4-assistant.cpp @@ -382,11 +382,17 @@ void llama_model_gemma4_assistant::load_arch_hparams(llama_model_loader & ml) { hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer); - uint32_t n_kv_shared_layers = 0; - ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false); - - hparams.n_layer_kv_from_start = hparams.n_layer - (int32_t) n_kv_shared_layers; - hparams.f_attention_scale = 1.0f; + // opencoti bug-858 dual-context MTP — the gemma4-assistant is a KV-LESS drafter: its attention + // layers do NOT own K/V. When run as ctx_dft they alias the TARGET's backbone K/V in place via + // mem_other/share (llama-model.cpp create_memory A3 + llama-kv-cache.cpp share branch). So + // n_layer_kv_from_start MUST stay -1 (has_kv() default => true for every layer), matching + // upstream b9859 which does NOT set it for LLM_ARCH_GEMMA4_ASSISTANT. Deriving it from + // ATTENTION_SHARED_KV_LAYERS (= n_layer for the drafter) yields 0 → has_kv() false for all + // layers → the KV-cache ctor `continue`s past every layer BEFORE the share branch → empty + // map_layer_ids → "get_k: map_layer_ids MISS il=0 (n_keys=0)" boot crash. That GGUF key drives + // the TARGET's intra-model elastic KV reuse (#417), a different mechanism from the cross-model + // share, and is inert for the KV-less drafter. + hparams.f_attention_scale = 1.0f; ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); @@ -434,7 +440,16 @@ void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) { // opencoti F5 M6-S4 mtp: GGUF tensor names upstream-aligned (#23398/#24282) — nextn.* / masked_embd_*. // The C++ member names stay mtp_* (control-plane; our single-context engine, not upstream's ctx_other). - mtp_pre_projection = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_PRE, "weight"), {2 * (int64_t) n_bb, n_embd}, 0); + // + // opencoti bug-858 (dual-context MTP): pre_projection is classified LLM_TENSOR_LAYER_INPUT, so the + // plain create_tensor lands it on the CPU buft. The single-context facade tolerates that (it runs on + // the target's sched_mtp, which has the CPU backend registered), but a REAL draft context ctx_dft + // (llama_init_from_model over the nested assistant) registers only the assistant's *offloaded* (CUDA) + // weight buffers — a CPU-resident weight then reads back uninitialized → the whole drafter forward is + // NaN → 0% accept. Force it onto the output-head (GPU-when-offloaded) buft, exactly like tok_embd + // (#408): device-config-respecting (stays CPU for a CPU-only drafter) and placement-only, so the + // facade path stays byte-identical. post_projection/output_norm are LAYER_OUTPUT → already dev_output. + mtp_pre_projection = create_tensor_output_head(tn(LLM_TENSOR_NEXTN_PROJ_PRE, "weight"), {2 * (int64_t) n_bb, n_embd}, 0); mtp_post_projection = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_POST, "weight"), {n_embd, (int64_t) n_bb}, 0); if (hparams.use_ordered_embeddings) { @@ -481,8 +496,167 @@ void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) { } } -std::unique_ptr 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(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 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(*this, params); } diff --git a/llama.cpp/src/models/gemma4.cpp b/llama.cpp/src/models/gemma4.cpp index a7b1598..bdf5990 100644 --- a/llama.cpp/src/models/gemma4.cpp +++ b/llama.cpp/src/models/gemma4.cpp @@ -171,6 +171,15 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para ggml_tensor * inp_out_ids = build_inp_out_ids(); + // opencoti-hook: gemma4-mtp-hidden (bug-858 dual-context MTP) — when a context extracts the + // pre-norm hidden for ALL positions (embeddings_pre_norm && !masked, i.e. the MTP *target* whose + // shifted per-position hidden seeds the drafter), DEFER the inp_out_ids row-strip until after + // t_h_pre_norm is captured, so t_h_pre_norm carries all n_tokens rows. Upstream b9859 gates the + // in-loop strip on embeddings_nextn_masked and strips late when !masked (gemma4.cpp:278/416). + // Normal decode (embeddings_pre_norm==false) and the masked draft keep ours' early strip → + // byte-identical. See docs/evaluations/mtp.md and UPSTREAM_SYNC.md. + const bool mtp_defer_out_ids = cparams.embeddings_pre_norm && !cparams.embeddings_pre_norm_masked; + ggml_tensor * inp_per_layer = nullptr; if (model.per_layer_tok_embd) { inp_per_layer = build_inp_per_layer(); @@ -253,7 +262,9 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para } // TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing - if (il == n_layer - 1 && inp_out_ids) { + // opencoti bug-858: skip the early strip for the MTP target (mtp_defer_out_ids) so the last + // layer runs on all n_tokens rows and t_h_pre_norm below stays ungathered. + if (il == n_layer - 1 && inp_out_ids && !mtp_defer_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); } @@ -353,7 +364,9 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens] // TODO @ngxson : improve this - if (il == n_layer - 1 && inp_out_ids) { + // opencoti bug-858: same deferred strip as the main path (keep per-layer input full for + // the MTP target so its last-layer residual matches the ungathered hidden). + if (il == n_layer - 1 && inp_out_ids && !mtp_defer_out_ids) { inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids); } @@ -384,6 +397,22 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para model.output_norm, nullptr, LLM_NORM_RMS, -1); + // opencoti-hook: gemma4-mtp-hidden (bug-858 dual-context MTP) — expose the POST-output-norm + // hidden state as t_h_pre_norm so a gemma4-assistant draft context (ctx_dft with ctx_other=this) + // can read it via llama_get_embeddings_pre_norm_ith() as the recurrent h input. Mirrors upstream + // b9859 gemma4.cpp res->t_h_nextn = cur (post-final-norm, captured PRE row-strip). For the MTP + // target (mtp_defer_out_ids) this tensor carries all n_tokens rows so the driver can shift the + // per-position hidden; for normal decode / masked draft the rows were already stripped in the + // last layer, so it equals the n_outputs t_embd below. See docs/evaluations/mtp.md. + cb(cur, "h_pre_norm", -1); + res->t_h_pre_norm = cur; + + // opencoti bug-858: deferred row-strip — t_h_pre_norm captured all positions above; now strip to + // the output rows for the logits / t_embd path (upstream gemma4.cpp:416 `if (!masked) get_rows`). + if (mtp_defer_out_ids && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; diff --git a/llama.cpp/src/models/qwen35.cpp b/llama.cpp/src/models/qwen35.cpp index 131f1fa..3a21d69 100644 --- a/llama.cpp/src/models/qwen35.cpp +++ b/llama.cpp/src/models/qwen35.cpp @@ -215,6 +215,20 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para } cur = inpL; + // opencoti-hook: qwen-nextn-mtp-hidden (bug-858) — see docs/evaluations/mtp.md. + // Upstream (b9859) exposes the MTP/NextN hidden as the POST-output-norm state (t_h_nextn); + // we mirror that here (apply output_norm FIRST, then expose it) so t_h_pre_norm points at the + // same tensor upstream's nextn.hnorm consumes — a syncability alignment, not a perf fix. + // MEASURED: on Qwen3.6-35B-A3B-MTP this reorder is numerically INERT (output_norm ≈ identity, + // so RMSNorm-then-hnorm == hnorm; draft acceptance is byte-identical 0.626 either way on GPU). + // The real acceptance gap vs upstream (ours 0.626 vs 0.770) is NOT the hidden-capture point — + // it is localized to ours' GPU flash-attention kernel: CPU is at parity (ours 0.696 ≈ upstream + // 0.692) and `-fa off` on GPU recovers ours to 0.683, so the divergence is FA-kernel precision + // in the nextn draft's full-attn block, not this norm. RMSNorm is row-independent, so applying + // it before vs after the inp_out_ids reduction leaves t_embd / t_logits byte-identical — base + // decode + RULER unaffected. + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_pre_norm", -1); res->t_h_pre_norm = cur; @@ -222,9 +236,6 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para cur = ggml_get_rows(ctx0, cur, inp_out_ids); } - // Final norm - cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); - cb(cur, "result_norm", -1); res->t_embd = cur; @@ -539,9 +550,8 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; - ggml_tensor * h_input = inp->embd; - ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); - cb(tok_embd, "mtp_tok_embd", il); + ggml_tensor * h_input = inp->embd; + ggml_tensor * inp_tokens = inp->tokens; res->add_input(std::move(inp)); @@ -549,6 +559,17 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr ggml_tensor * inp_out_ids = build_inp_out_ids(); auto * inp_attn = build_attn_inp_kv(); + // opencoti-hook: fused-nextn-mtp — one MTP step (embed → block → head → greedy argmax) as a + // reusable lambda the fused draft graph replays N× on-device (in-graph argmax→embed chain, + // one launch+readback). Takes the token index (embeds internally) so step k+1 can chain on + // step k's argmax. Single-step (n_mtp_steps==1) is byte-identical to the prior inline build + // (the argmax node is pruned when unreferenced). See docs/features/fused_nextn_mtp.md. + auto build_one_step = [&](ggml_tensor * h_input, ggml_tensor * tok, + ggml_tensor * inp_pos, auto * inp_attn, + ggml_tensor * inp_out_ids, bool fused) + -> std::tuple { + ggml_tensor * 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(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(); + inp_mtp->inp_last_token = inp_tokens; + inp_mtp->inp_h_prev = h_input; + inp_mtp->inp_pos_steps.reserve(n_steps); + std::vector 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 step_args; + step_args.reserve(n_steps); + ggml_tensor * last_h_pre_norm = nullptr; + ggml_tensor * last_logits = nullptr; + for (int32_t k = 0; k < n_steps; ++k) { + auto [h_pre_norm_k, logits_k, arg_k] = build_one_step(h_k, tok_k, pos_steps[k], inp_attn, inp_out_ids, true); + step_args.push_back(arg_k); + last_h_pre_norm = h_pre_norm_k; + last_logits = logits_k; + tok_k = arg_k; // next token = this step's greedy argmax (I32[1]) + h_k = h_pre_norm_k; // next seed hidden = this step's pre-norm hidden + } + + ggml_tensor * all_args = step_args[0]; + for (int32_t k = 1; k < n_steps; ++k) { + all_args = ggml_concat(ctx0, all_args, step_args[k], 0); + } + cb(all_args, "mtp_fused_argmax", -1); + + res->t_argmax = all_args; // I32[n_steps] drafted tokens + res->t_h_pre_norm = last_h_pre_norm; // final hidden (seed for next cycle) + res->t_logits = last_logits; + ggml_build_forward_expand(gf, all_args); + ggml_build_forward_expand(gf, last_h_pre_norm); } diff --git a/llama.cpp/src/models/qwen35moe.cpp b/llama.cpp/src/models/qwen35moe.cpp index f3f8e73..0eb2871 100644 --- a/llama.cpp/src/models/qwen35moe.cpp +++ b/llama.cpp/src/models/qwen35moe.cpp @@ -602,9 +602,8 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; - ggml_tensor * h_input = inp->embd; - ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); - cb(tok_embd, "mtp_tok_embd", il); + ggml_tensor * h_input = inp->embd; + ggml_tensor * inp_tokens = inp->tokens; res->add_input(std::move(inp)); @@ -612,7 +611,17 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm ggml_tensor * inp_out_ids = build_inp_out_ids(); auto * inp_attn = build_attn_inp_kv(); - + // opencoti-hook: fused-nextn-mtp — one MTP step (embed → block → head → greedy argmax) as a + // reusable lambda the fused draft graph replays N× on-device (in-graph argmax→embed chain, + // one launch+readback). Takes the token index (embeds internally) so step k+1 can chain on + // step k's argmax. Single-step (n_mtp_steps==1) is byte-identical to the prior inline build + // (the argmax node is pruned when unreferenced). See docs/features/fused_nextn_mtp.md. + auto build_one_step = [&](ggml_tensor * h_input, ggml_tensor * tok, + ggml_tensor * inp_pos, auto * inp_attn, + ggml_tensor * inp_out_ids, bool fused) + -> std::tuple { + ggml_tensor * 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(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(); + inp_mtp->inp_last_token = inp_tokens; + inp_mtp->inp_h_prev = h_input; + inp_mtp->inp_pos_steps.reserve(n_steps); + std::vector 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 step_args; + step_args.reserve(n_steps); + ggml_tensor * last_h_pre_norm = nullptr; + ggml_tensor * last_logits = nullptr; + for (int32_t k = 0; k < n_steps; ++k) { + auto [h_pre_norm_k, logits_k, arg_k] = build_one_step(h_k, tok_k, pos_steps[k], inp_attn, inp_out_ids, true); + step_args.push_back(arg_k); + last_h_pre_norm = h_pre_norm_k; + last_logits = logits_k; + tok_k = arg_k; // next token = this step's greedy argmax (I32[1]) + h_k = h_pre_norm_k; // next seed hidden = this step's pre-norm hidden + } + + ggml_tensor * all_args = step_args[0]; + for (int32_t k = 1; k < n_steps; ++k) { + all_args = ggml_concat(ctx0, all_args, step_args[k], 0); + } + cb(all_args, "mtp_fused_argmax", -1); + + res->t_argmax = all_args; // I32[n_steps] drafted tokens + res->t_h_pre_norm = last_h_pre_norm; // final hidden (seed for next cycle) + res->t_logits = last_logits; + ggml_build_forward_expand(gf, all_args); + ggml_build_forward_expand(gf, last_h_pre_norm); } diff --git a/llama.cpp/tools/server/server-context.cpp b/llama.cpp/tools/server/server-context.cpp index 7bad046..34b50b1 100644 --- a/llama.cpp/tools/server/server-context.cpp +++ b/llama.cpp/tools/server/server-context.cpp @@ -965,10 +965,49 @@ private: return false; } - // No separate draft model/context: the assistant decodes on ctx_tgt via llama_decode_mtp. - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = nullptr; - SRV_INF("%s", "MTP assistant loaded into target\n"); + // opencoti bug-858 dual-context MTP (A5): optionally create a REAL draft context FROM the + // nested gemma4-assistant model with cparams.ctx_other = ctx_tgt (upstream b9859 + // dual-context). The assistant is KV-less: its create_memory (A3) shares the target's KV + // cells + aliases the last full/SWA layer, and its graph (A6) reads the target's tok_embd + // + shared K/V via ctx_other; the shared-KV draft_mtp driver (is_mem_shared) then runs it. + // Gated by OPENCOTI_MTP_DUAL_CTX so the default keeps the proven single-context in-target + // path (ctx_dft == nullptr, llama_decode_mtp) byte-identical. + const char * mtp_dual_env = getenv("OPENCOTI_MTP_DUAL_CTX"); + const bool mtp_dual_ctx = mtp_dual_env && mtp_dual_env[0] && mtp_dual_env[0] != '0'; + if (mtp_dual_ctx) { + // llama_init_from_model wants a mutable model*; the nested assistant is owned mutably + // by the target (unique_ptr), so the const_cast is safe. + llama_model * assistant = const_cast(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,