From 79da29e2b8070b2b0a31f692c1a015155da04cb7 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Mon, 28 Sep 2026 08:14:05 +1000 Subject: [PATCH 01/24] qwen35 MTP: optional trimmed draft head (top-K frequent tokens, FR-Spec style) blk..nextn.draft_head (rows of output.weight for the K most frequent tokens plus all special tokens) and draft_head_ids (I32). When present, the MTP draft computes logits for those K rows only and scatters them into a full-vocab row of -inf with set_rows, so everything downstream is unchanged and the target model still verifies every draft. The head is listed under prism.hadamard.weight_names (same input-side rotation as output.weight). LLAMA_MTP_DRAFT_HEAD_OFF=1 uses the full head. Loader: GGML_OP_SET_ROWS placement check so the I32 id tensor lands in device memory. Bonsai 27B, K = 32768 (99.6% token coverage of a 4M-token code/prose corpus): the draft head reads 7.6x fewer bytes; fresh code temp 0 90.3 -> 94.8-95.8 t/s with identical output, temp 0.6 neutral within noise. K = 16384 is worse (92.1, 76 vs 81). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit e75a9943ca47e8c6db6fc1781fcfc4de8dc14367 on Torchit1/llama.cpp arc-b580) --- src/llama-arch.cpp | 4 ++++ src/llama-arch.h | 2 ++ src/llama-model-loader.cpp | 7 +++++++ src/llama-model.cpp | 1 + src/llama-model.h | 2 ++ src/models/qwen35.cpp | 23 ++++++++++++++++++++++- 6 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 4c3fcb2575a4..2119098b0420 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -562,6 +562,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, + { LLM_TENSOR_NEXTN_DRAFT_HEAD, "blk.%d.nextn.draft_head" }, + { LLM_TENSOR_NEXTN_DRAFT_HEAD_IDS, "blk.%d.nextn.draft_head_ids" }, { LLM_TENSOR_ATTN_SUB_NORM, "blk.%d.attn_sub_norm" }, { LLM_TENSOR_FFN_SUB_NORM, "blk.%d.ffn_sub_norm" }, { LLM_TENSOR_DEC_OUTPUT_NORM, "dec.output_norm" }, @@ -957,6 +959,8 @@ static const std::map LLM_TENSOR_INFOS = { { LLM_TENSOR_NEXTN_HNORM, { LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL } }, { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, { LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT } }, { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, { LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL } }, + { LLM_TENSOR_NEXTN_DRAFT_HEAD, { LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT } }, + { LLM_TENSOR_NEXTN_DRAFT_HEAD_IDS, { LLM_TENSOR_LAYER_REPEATING, GGML_OP_SET_ROWS } }, // dspark drafter extras. fc / markov factors are matmuls; the two norms are // elementwise muls. All output-side (non-repeating) so the loader does not // require a per-block index. diff --git a/src/llama-arch.h b/src/llama-arch.h index f292e497e5ec..bdd58a7667d1 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -676,6 +676,8 @@ enum llm_tensor { LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + LLM_TENSOR_NEXTN_DRAFT_HEAD, // ARC-LAB: MTP LM head over the top-K frequent tokens only + LLM_TENSOR_NEXTN_DRAFT_HEAD_IDS, // ARC-LAB: vocab ids of its rows LLM_TENSOR_MASKED_EMBD_CENTROIDS, LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 835d172d33f7..23d74c30a396 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -948,6 +948,13 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512); op_tensor = ggml_get_rows(ctx, w, b); } break; + case GGML_OP_SET_ROWS: + { + // ARC-LAB: w = I32 row indices (MTP draft-head vocab ids), scattered into a full logits row + ggml_tensor * dst = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, w->ne[0], 1); + ggml_tensor * src = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, w->ne[0], 1); + op_tensor = ggml_set_rows(ctx, dst, src, w); + } break; case GGML_OP_MUL_MAT: { ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c2ac6a50b316..a31e031fb1c3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1291,6 +1291,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { "ffn_gate_exps", "ffn_up_exps", "ffn_down_exps", "ffn_gate_up_exps", "ffn_gate_shexp", "ffn_up_shexp", "ffn_down_shexp", "ssm_out", + "nextn.draft_head", // ARC-LAB: row subset of output.weight, same input-side rotation }; if (name == "output.weight") { return true; // the output head is built through build_lora_mm in every arch diff --git a/src/llama-model.h b/src/llama-model.h index 8834271a101a..d93326b86458 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -223,6 +223,8 @@ struct llama_layer_nextn { struct ggml_tensor * enorm = nullptr; struct ggml_tensor * hnorm = nullptr; struct ggml_tensor * shared_head_head = nullptr; + struct ggml_tensor * draft_head = nullptr; // ARC-LAB: [n_embd, K] rows of the LM head for top-K tokens + struct ggml_tensor * draft_head_ids = nullptr; // ARC-LAB: I32 [K] vocab id of each row struct ggml_tensor * shared_head_head_s = nullptr; struct ggml_tensor * shared_head_head_in_s = nullptr; struct ggml_tensor * shared_head_norm = nullptr; diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 48a2cd158f35..72fc50a942ce 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -117,6 +117,12 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED); + // ARC-LAB: optional trimmed draft head (FR-Spec style): the draft only proposes the K most frequent tokens + if (const ggml_tensor * dh = ml.get_tensor_meta(tn(LLM_TENSOR_NEXTN_DRAFT_HEAD, "weight", il).str().c_str())) { + const int64_t n_draft_vocab = dh->ne[1]; + layer.nextn.draft_head = create_tensor(tn(LLM_TENSOR_NEXTN_DRAFT_HEAD, "weight", il), { n_embd, n_draft_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); + layer.nextn.draft_head_ids = create_tensor(tn(LLM_TENSOR_NEXTN_DRAFT_HEAD_IDS, "weight", il), { n_draft_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); + } }; for (int i = 0; i < n_layer; ++i) { @@ -777,7 +783,22 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; GGML_ASSERT(head_w && "QWEN35 MTP: missing LM head (nextn.shared_head_head or model.output)"); - cur = build_lora_mm(head_w, cur, head_s); + static const bool draft_head_off = getenv("LLAMA_MTP_DRAFT_HEAD_OFF") != nullptr; + if (layer.nextn.draft_head && layer.nextn.draft_head_ids && !layer.nextn.shared_head_head && !draft_head_off) { + // ARC-LAB: logits for the top-K frequent tokens only, scattered into a full-vocab row of -inf; the target + // model still verifies every draft, so this changes only which tokens can be drafted, not the output + const int64_t n_k = layer.nextn.draft_head->ne[1]; + const int64_t n_out = cur->ne[1]; + const int64_t n_voc = head_w->ne[1]; + ggml_tensor * lk = build_lora_mm(layer.nextn.draft_head, cur, nullptr); // [K, n_out] + ggml_tensor * full = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_voc, n_out); + full = ggml_fill(ctx0, full, -INFINITY); + full = ggml_set_rows(ctx0, full, ggml_reshape_3d(ctx0, lk, 1, n_k, n_out), + ggml_reshape_2d(ctx0, layer.nextn.draft_head_ids, n_k, 1)); + cur = ggml_reshape_2d(ctx0, full, n_voc, n_out); + } else { + cur = build_lora_mm(head_w, cur, head_s); + } cb(cur, "result_output", -1); res->t_logits = cur; From 1a5c67d6191acdac5e4f400dacbf1cbf5d211cf0 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 3 Oct 2026 15:50:28 +1000 Subject: [PATCH 02/24] speculative ngram-mod: adaptive draft length (start 64, double on full accept, halve on rejection) A copy ends where the model stops but the context goes on, so with --spec-ngram-mod-n-max 256 a short copy drafted ~200 tokens past its end; at 48K context verifying those cost ~0.3 s each and made copying slower than writing new code. Drafts are now cut to n_cur (starts at 64, doubles after a fully accepted draft up to n_max, halves after a rejected one). B580, Bonsai 27B, MTP + ngram-mod, same outputs: 48K short copy 38.6 -> 57.3 t/s, 48K 1262-token copy 161.6 -> 172.9, 48K new code 78.5 -> 79.4, short rename 428 -> 411, edit 305 -> 299. LLAMA_ARG_SPEC_NGRAM_ADAPT=0 restores the fixed length; =N sets the starting length. Co-Authored-By: Claude (cherry picked from commit e576a2879f5a1e1c7b1c04e871e880a337824926 on Torchit1/llama.cpp arc-b580) --- common/speculative.cpp | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index 23843abe0139..f8acb3792bfb 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2984,10 +2984,24 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { // consecutive accept rounds with low acceptance fraction (< 0.5) int n_low = 0; + + // ARC-LAB adaptive draft length: drafts are cut to n_cur tokens; n_cur doubles after a fully accepted draft + // (up to n_max) and halves after a rejected one (down to n_start). A copy ends where the model stops, but the + // context continues, so a fixed n_max of 256 threw away ~200 drafted tokens per short copy - at 48K context + // verifying those costs ~0.3 s each. + int n_cur = 0; }; std::vector sinfos; + // LLAMA_ARG_SPEC_NGRAM_ADAPT=0 disables it (fixed n_max as upstream); other values = starting length (default 64). + // B580, n_max 256, t/s fixed vs 16 / 64 / 128: 48K short copy 38.6 / 60.7 / 57.3 / 45.7, 48K long copy 161.6 / 165 / + // 172.9 / 178.9, short rename 428 / 394 / 411 / 416, edit 305 / 296 / 299 / 316 (same outputs). + int n_start = [] { + const char * e = std::getenv("LLAMA_ARG_SPEC_NGRAM_ADAPT"); + return e ? std::atoi(e) : 64; + }(); + common_speculative_impl_ngram_mod( const common_params_speculative & params, uint32_t n_seq) @@ -3016,6 +3030,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { sinfo.i_last = 0; sinfo.n_draft_last = 0; + sinfo.n_cur = n_start > 0 ? std::min(n_start, params.n_max) : params.n_max; const size_t n = mod.get_n(); if (prompt.size() < n) { @@ -3091,6 +3106,11 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { } result.resize(result.size() - n); + // adaptive length: the n_min check above still sees the whole chain; only the returned draft is cut + if (sinfo.n_cur > 0 && result.size() > (size_t) sinfo.n_cur) { + result.resize(sinfo.n_cur); + } + // store length of drafted n-gram for later acceptance analysis sinfo.n_draft_last = result.size(); } @@ -3120,6 +3140,11 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { auto & sinfo = sinfos[seq_id]; + if (n_start > 0 && sinfo.n_draft_last > 0) { + sinfo.n_cur = n_accepted >= sinfo.n_draft_last ? std::min(2 * sinfo.n_cur, (int) params.n_max) + : std::max(std::min(n_start, (int) params.n_max), sinfo.n_cur / 2); + } + // compute acceptance fraction if we have a recorded draft length if (sinfo.n_draft_last > 0) { const double f_acc = (double)n_accepted / (double)sinfo.n_draft_last; From 0b8dd646fb82f44ce9bd8aab6eff432981f0fc68 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:30:58 +1100 Subject: [PATCH 03/24] tools: llama-mtp-export; MTP draft log + lab speculative knobs llama-mtp-export writes, per position, the trunk state the NextN (MTP) head consumes (h_nextn) and the trunk's greedy pick, for training / validating the draft head offline. common/speculative: LLAMA_MTP_DRAFT_LOG (top-3 candidates per draft step), LAB_SPEC_DRAFT_N_MAX / LAB_SPEC_NGRAM_START runtime knobs via GGML_SYCL_KNOBS_FILE. Used for the B580 self-distilled MTP head (+4.8% with MTP, outputs identical). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 7c577265b0e554dee001e7a08db9576fff2b10f4 on Torchit1/llama.cpp arc-b580) --- common/speculative.cpp | 67 +++++++++++++++++- tools/CMakeLists.txt | 1 + tools/mtp-export/CMakeLists.txt | 4 ++ tools/mtp-export/mtp-export.cpp | 119 ++++++++++++++++++++++++++++++++ 4 files changed, 189 insertions(+), 2 deletions(-) create mode 100644 tools/mtp-export/CMakeLists.txt create mode 100644 tools/mtp-export/mtp-export.cpp diff --git a/common/speculative.cpp b/common/speculative.cpp index f8acb3792bfb..626b376a58a2 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -32,6 +32,55 @@ #include #include +// ARC-LAB: lab knobs from GGML_SYCL_KNOBS_FILE (lines NAME=V, re-read when the file changes) so one server can A/B +// speculative settings: LAB_SPEC_DRAFT_N_MAX caps the MTP draft below --spec-draft-n-max, LAB_SPEC_NGRAM_START replaces +// LLAMA_ARG_SPEC_NGRAM_ADAPT's starting length (both from the next request on). Unset = command line / env. +#include +#include +#include +// ARC-LAB: LLAMA_MTP_DRAFT_LOG= appends one JSON line per MTP draft step (position predicted, depth, top 3 candidates +// with probabilities, whether p_min stopped the draft) plus a "begin" line per request, for offline acceptance / tree +// simulation (arctune, 2026-10-05). +static FILE * common_mtp_draft_log() { + static FILE * f = [] { + const char * p = std::getenv("LLAMA_MTP_DRAFT_LOG"); + return p ? fopen(p, "a") : nullptr; + }(); + return f; +} +static int common_lab_knob(const char * name, int def) { + static const char * path = std::getenv("GGML_SYCL_KNOBS_FILE"); + if (!path) { + return def; + } + static std::mutex mtx; + static std::filesystem::file_time_type last{}; + static std::unordered_map m; + std::lock_guard lock(mtx); + std::error_code ec; + const auto t = std::filesystem::last_write_time(path, ec); + if (!ec && t != last) { + last = t; + m.clear(); + if (FILE * fp = fopen(path, "r")) { + char line[256]; + while (fgets(line, sizeof(line), fp)) { + char * eq = strchr(line, '='); + if (!eq || line[0] == '#') { + continue; + } + *eq = 0; + std::string k(line); + k.erase(k.find_last_not_of(" \t") + 1); + m[k] = atoi(eq + 1); + } + fclose(fp); + } + } + const auto it = m.find(name); + return it == m.end() ? def : it->second; +} + #define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) #define SPC_TRC(fmt, ...) LOG_TRC("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) #define SPC_INF(fmt, ...) LOG_INF("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) @@ -2385,6 +2434,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + if (FILE * lf = common_mtp_draft_log()) { + fprintf(lf, "{\"begin\": %zu, \"seq\": %d}\n", prompt.size(), (int) seq_id); + fflush(lf); + } const int32_t N = (int32_t) prompt.size(); if (N <= 0) { return; @@ -2765,6 +2818,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } + if (FILE * lf = common_mtp_draft_log()) { + const auto & dpl = dparams.at(seq_id); + fprintf(lf, "{\"pos\": %d, \"d\": %d, \"c\": [", (int) (dpl.n_past + i + 1), i); + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + fprintf(lf, "%s[%d, %.4f]", k ? ", " : "", cur_p->data[k].id, cur_p->data[k].p); + } + fprintf(lf, "], \"stop\": %d}\n", cur_p->data[0].p < params.p_min ? 1 : 0); + } // add drafted token for each sequence const llama_token id = cur_p->data[0].id; @@ -2791,7 +2852,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + if (std::min((int) params.n_max, common_lab_knob("LAB_SPEC_DRAFT_N_MAX", params.n_max)) <= (int) result.size()) { drafting[seq_id] = false; n_drafting--; continue; @@ -2997,10 +3058,11 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { // LLAMA_ARG_SPEC_NGRAM_ADAPT=0 disables it (fixed n_max as upstream); other values = starting length (default 64). // B580, n_max 256, t/s fixed vs 16 / 64 / 128: 48K short copy 38.6 / 60.7 / 57.3 / 45.7, 48K long copy 161.6 / 165 / // 172.9 / 178.9, short rename 428 / 394 / 411 / 416, edit 305 / 296 / 299 / 316 (same outputs). - int n_start = [] { + const int n_start_env = [] { const char * e = std::getenv("LLAMA_ARG_SPEC_NGRAM_ADAPT"); return e ? std::atoi(e) : 64; }(); + int n_start = n_start_env; common_speculative_impl_ngram_mod( const common_params_speculative & params, @@ -3030,6 +3092,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { sinfo.i_last = 0; sinfo.n_draft_last = 0; + n_start = common_lab_knob("LAB_SPEC_NGRAM_START", n_start_env); sinfo.n_cur = n_start > 0 ? std::min(n_start, params.n_max) : params.n_max; const size_t n = mod.get_n(); diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index da1b2bc707d7..986ad620058b 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -28,6 +28,7 @@ else() add_subdirectory(server) endif() add_subdirectory(tokenize) + add_subdirectory(mtp-export) # ARC-LAB: MTP head training data add_subdirectory(tts) add_subdirectory(mtmd) if (GGML_RPC) diff --git a/tools/mtp-export/CMakeLists.txt b/tools/mtp-export/CMakeLists.txt new file mode 100644 index 000000000000..769f6bb07c92 --- /dev/null +++ b/tools/mtp-export/CMakeLists.txt @@ -0,0 +1,4 @@ +set(TARGET llama-mtp-export) +add_executable(${TARGET} mtp-export.cpp) +target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/tools/mtp-export/mtp-export.cpp b/tools/mtp-export/mtp-export.cpp new file mode 100644 index 000000000000..87e99762c3bf --- /dev/null +++ b/tools/mtp-export/mtp-export.cpp @@ -0,0 +1,119 @@ +// ARC-LAB mtp-export: training / validation data for the NextN (MTP) draft head. +// Runs the trunk over token sequences (teacher forcing) and writes, per position t: +// h[t] the trunk hidden state the MTP head consumes (h_nextn = final-norm output), f16 [n_embd] +// tok[t] the input token at t +// tgt[t] the trunk's own greedy choice for position t + 1 (argmax of the logits at t) - what a draft must match +// The MTP head at position t sees (h[t - 1], tok[t]) and should predict tgt[t]. +// +// llama-mtp-export -m model.gguf -i seqs.i32 -o prefix [-c n_ctx] [-b n_batch] +// seqs.i32: repeated records of int32 n followed by n int32 token ids. Writes prefix.h (f16), prefix.tok (i32), +// prefix.tgt (i32), prefix.idx (int64 start offset of each sequence, plus the total at the end). +#include "llama.h" +#include "ggml.h" +#include "../../src/llama-ext.h" + +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + std::string model_path, in_path, out_prefix; + int n_ctx = 4096, n_batch = 512; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "-m" && i + 1 < argc) model_path = argv[++i]; + else if (a == "-i" && i + 1 < argc) in_path = argv[++i]; + else if (a == "-o" && i + 1 < argc) out_prefix = argv[++i]; + else if (a == "-c" && i + 1 < argc) n_ctx = atoi(argv[++i]); + else if (a == "-b" && i + 1 < argc) n_batch = atoi(argv[++i]); + } + if (model_path.empty() || in_path.empty() || out_prefix.empty()) { + fprintf(stderr, "usage: %s -m model.gguf -i seqs.i32 -o prefix [-c n_ctx] [-b n_batch]\n", argv[0]); + return 1; + } + + // read the sequences + std::vector> seqs; + { + FILE * f = fopen(in_path.c_str(), "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", in_path.c_str()); return 1; } + int32_t n; + while (fread(&n, 4, 1, f) == 1) { + std::vector s(n); + if (fread(s.data(), 4, n, f) != (size_t) n) { fprintf(stderr, "short read\n"); return 1; } + if (n > n_ctx) s.resize(n_ctx); + seqs.push_back(std::move(s)); + } + fclose(f); + } + fprintf(stderr, "mtp-export: %zu sequences\n", seqs.size()); + + llama_backend_init(); + llama_model_params mp = llama_model_default_params(); + mp.n_gpu_layers = 999; + llama_model * model = llama_model_load_from_file(model_path.c_str(), mp); + if (!model) { fprintf(stderr, "load failed\n"); return 1; } + const int n_embd = llama_model_n_embd(model); + const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + + llama_context_params cp = llama_context_default_params(); + cp.n_ctx = n_ctx; + cp.n_batch = n_batch; + cp.n_ubatch = n_batch; + cp.n_seq_max = 1; + llama_context * ctx = llama_init_from_model(model, cp); + if (!ctx) { fprintf(stderr, "context failed\n"); return 1; } + llama_set_embeddings_nextn(ctx, true, /*masked*/ false); + + FILE * fh = fopen((out_prefix + ".h").c_str(), "wb"); + FILE * ft = fopen((out_prefix + ".tok").c_str(), "wb"); + FILE * fg = fopen((out_prefix + ".tgt").c_str(), "wb"); + FILE * fi = fopen((out_prefix + ".idx").c_str(), "wb"); + if (!fh || !ft || !fg || !fi) { fprintf(stderr, "cannot write outputs\n"); return 1; } + + llama_batch batch = llama_batch_init(n_batch, 0, 1); + std::vector hrow(n_embd); + int64_t total = 0; + for (size_t si = 0; si < seqs.size(); ++si) { + const auto & s = seqs[si]; + fwrite(&total, 8, 1, fi); + llama_memory_clear(llama_get_memory(ctx), true); + for (size_t p0 = 0; p0 < s.size(); p0 += n_batch) { + const int n = (int) std::min(n_batch, s.size() - p0); + batch.n_tokens = n; + for (int i = 0; i < n; ++i) { + batch.token[i] = s[p0 + i]; + batch.pos[i] = (llama_pos) (p0 + i); + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; + } + if (llama_decode(ctx, batch) != 0) { fprintf(stderr, "decode failed at seq %zu pos %zu\n", si, p0); return 1; } + for (int i = 0; i < n; ++i) { + const float * h = llama_get_embeddings_nextn_ith(ctx, i); + const float * l = llama_get_logits_ith(ctx, i); + if (!h || !l) { fprintf(stderr, "missing outputs at %d\n", i); return 1; } + int best = 0; + for (int v = 1; v < n_vocab; ++v) if (l[v] > l[best]) best = v; + for (int e = 0; e < n_embd; ++e) hrow[e] = ggml_fp32_to_fp16(h[e]); + fwrite(hrow.data(), sizeof(ggml_fp16_t), n_embd, fh); + const int32_t tok = s[p0 + i], tg = best; + fwrite(&tok, 4, 1, ft); + fwrite(&tg, 4, 1, fg); + } + } + total += (int64_t) s.size(); + if ((si + 1) % 10 == 0 || si + 1 == seqs.size()) { + fprintf(stderr, "mtp-export: %zu / %zu sequences, %lld positions\n", si + 1, seqs.size(), (long long) total); + } + } + fwrite(&total, 8, 1, fi); + fclose(fh); fclose(ft); fclose(fg); fclose(fi); + llama_batch_free(batch); + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + return 0; +} From 15a9adaa55fbad85897aa7af5661de84f4e2ef23 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:50:59 +1100 Subject: [PATCH 04/24] mtp: the draft head's prompt pass computes only its K/V (LLAMA_MTP_KV_ONLY) Independent implementation of the technique described by GRIMOIRE (written from a functional description and this fork's own code; no GRIMOIRE source read). A draft-context decode without output rows (the NextN head's prompt catch-up) only fills the head's KV cache, so graph_mtp now stops after the cache store: embedding (+ Hadamard inverse), hnorm/enorm, eh_proj, attn_norm, K/V projection, K norm, RoPE and cpy_k/cpy_v. No query, attention, output projection, FFN or head. The K/V ops are the full graph's on the same shapes and in the same expand order, so the cache rows are bit-identical. For batches of <= 8 rows the unused Q projection stays in the graph so the SYCL same-input mat-vec grouping (whose tile depends on the largest member) matches the full graph. Only masked nextn extraction qualifies (unmasked readers copy h for every row). Decodes with outputs build the unchanged graph. LLAMA_MTP_KV_ONLY=0 restores the full graph for every decode. Default on. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 595d554c90d5a15eb22e4c66c4ac2564fd71b4cd on Torchit1/llama.cpp arc-b580) --- src/models/qwen35.cpp | 54 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 72fc50a942ce..ecc434f3e31f 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -1,5 +1,6 @@ #include "models.h" #include "llama-memory-recurrent.h" +#include "llama-kv-cache.h" void llama_model_qwen35::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -672,8 +673,19 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr res->add_input(std::move(inp)); + // ARC-LAB: KV-only pass (LLAMA_MTP_KV_ONLY, default on; independent implementation of the technique described by + // GRIMOIRE). A decode with no output rows (the head's prompt catch-up) is only there to fill this layer's K/V, so + // build up to the cache store and stop: no query, attention, output projection, FFN or head. The K/V ops are the + // full graph's, in the same order, so the cache rows are bit-identical. Unmasked nextn readers copy h for every + // token, so they keep the full graph. + static const bool kv_only_on = [] { + const char * e = getenv("LLAMA_MTP_KV_ONLY"); + return e == nullptr || atoi(e) != 0; + }(); + const bool kv_only = kv_only_on && n_outputs == 0 && cparams.embeddings_nextn_masked; + ggml_tensor * inp_pos = build_inp_pos(); - ggml_tensor * inp_out_ids = build_inp_out_ids(); + ggml_tensor * inp_out_ids = kv_only ? nullptr : build_inp_out_ids(); // an unread input would have no buffer auto * inp_attn = build_attn_inp_kv(); @@ -694,6 +706,46 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); + if (kv_only) { + static bool logged = false; + if (!logged) { + logged = true; + LLAMA_LOG_INFO("%s: MTP KV-only pass for decodes without outputs (LLAMA_MTP_KV_ONLY=0 turns it off)\n", __func__); + } + // the SYCL backend runs same-input mat-vecs of <= 8 rows as one launch whose tile choice depends on the + // largest member; keep the Q projection (result unused) so q + k + v group exactly as in the full graph + if (n_tokens <= 8) { + ggml_build_forward_expand(gf, build_lora_mm(layer.wq, cur, layer.wq_s)); + } + + ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "mtp_Kcur_normed", il); + + ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + cb(Vcur, "mtp_Vcur", il); + + 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); + + // the store half of build_attn(): same optional cache rotations, same expand order (v, k, then the copies) + if (inp_attn->self_k_rot) { + Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot); + } + if (inp_attn->self_v_rot) { + Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot); + } + ggml_build_forward_expand(gf, Vcur); + ggml_build_forward_expand(gf, Kcur); + cb(Kcur, "k_cache_in", il); + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il)); + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il)); + return; + } + ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); cb(Qcur_full, "mtp_Qcur_full", il); From c94a7900286d088f88248c1edc9e052bcecf93c1 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:57:24 +1100 Subject: [PATCH 05/24] mtp: one-graph draft chain with device draft-vocab embeddings (LLAMA_MTP_CHAIN) All MTP draft steps in one llama_decode: each step takes the previous step's in-graph argmax and head output row; ids read back once. A 36 MiB device copy of the draft-vocabulary token_embd rows (PTQ1_0) plus SYCL GET_ROWS for PTQ1_0 (same arithmetic as the CPU) keep the lookups on the GPU. On by default only with a trimmed draft head; LLAMA_MTP_CHAIN=1 forces it, =0 off; LLAMA_MTP_CHAIN_GPU_EMBD=0 uses host lookups. Independent implementation of the technique described by llamADALL (clean-room: written from a functional spec, without reading llamADALL's source). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit bcbe5a0ebc48d5541f7425af42a0c40c6defe399 on Torchit1/llama.cpp arc-b580) --- common/speculative.cpp | 116 ++++++++++++ ggml/src/ggml-sycl/dequantize.hpp | 13 ++ ggml/src/ggml-sycl/getrows.cpp | 4 + ggml/src/ggml-sycl/ggml-sycl.cpp | 1 + src/llama-context.cpp | 150 +++++++++++++++ src/llama-context.h | 17 ++ src/llama-cparams.h | 3 + src/llama-ext.h | 10 + src/llama-graph.cpp | 2 + src/llama-graph.h | 11 ++ src/llama-kv-cache.cpp | 42 +++++ src/llama-kv-cache.h | 12 ++ src/models/models.h | 3 + src/models/qwen35.cpp | 295 ++++++++++++++++++++++++++++++ tests/test-backend-ops.cpp | 28 ++- 15 files changed, 704 insertions(+), 3 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 626b376a58a2..b3ec6ff87aa9 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2301,6 +2301,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { }; std::vector deferred; + // ARC-LAB: chained drafts (LLAMA_MTP_CHAIN; independent implementation of the technique described by llamADALL). + // All draft steps run in one llama_decode (llama_set_mtp_chain): each step's argmax token and output row feed the + // next step on the device, and the n draft tokens are read back once, instead of n decodes with a host sample in + // between. Same greedy drafts and the same draft K/V as the loop below. Set in the ctor. + bool chain_mode = false; + bool can_defer() const { return !is_mem_shared && !chain_heads && !getenv("LLAMA_MTP_EAGER_CATCHUP"); } @@ -2411,6 +2417,108 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { verify_h_rows.assign(n_seq, 0); deferred.assign(n_seq, {}); + + // chained drafts: default on when the head has a trimmed draft vocabulary (its token_embd rows fit on the + // device; measured slower without one), LLAMA_MTP_CHAIN=1 forces, =0 off. Greedy single-sequence drafting only. + { + const char * env = getenv("LLAMA_MTP_CHAIN"); + const int32_t sup = llama_mtp_chain_support(ctx_dft); + const bool want = env ? atoi(env) != 0 : (sup & 2) != 0; + const char * why = nullptr; + if (!want) { + why = env ? "LLAMA_MTP_CHAIN=0" : "no trimmed draft head (LLAMA_MTP_CHAIN=1 forces it)"; + } else if (sup == 0) { + why = "not supported by this context (qwen35 single head, flash attention, plain KV cache)"; + } else if (is_mem_shared || chain_heads) { + why = "shared-memory or multi-head MTP"; + } else if (n_seq != 1) { + why = "more than one sequence"; + } else if (this->params.p_min > 0.0f) { + why = "p_min > 0"; + } + chain_mode = why == nullptr; + if (chain_mode) { + SPC_INF("MTP chained drafts on (%s)\n", (sup & 2) ? "draft head" : "full head, host token lookups"); + } else { + SPC_INF("MTP chained drafts off: %s\n", why); + } + } + } + + // ARC-LAB: one chained decode for the whole draft of sequence 0. Returns false when it does not apply (the loop + // runs), true when it ran (dparams[0].result holds the drafts, empty on failure). + bool draft_chain(common_speculative_draft_params_vec & dparams) { + auto * ctx_dft = params.ctx_dft; + auto & dp = dparams[0]; + auto & d = deferred[0]; + if (!dp.drafting || d.catchup_failed) { + return false; + } + + // the loop's step count: it stops once min(n_max, LAB_SPEC_DRAFT_N_MAX) tokens are drafted (at least one) + const int32_t n_draft = std::max(1, std::min((int) params.n_max, common_lab_knob("LAB_SPEC_DRAFT_N_MAX", params.n_max))); + const int32_t n_catch = d.pending ? d.n_valid : 0; + const int32_t n_rows = n_catch + n_draft; + if (n_draft < 2 || n_rows > (int32_t) std::min(llama_n_batch(ctx_dft), llama_n_ubatch(ctx_dft))) { + return false; + } + + const size_t row_bytes = (size_t) n_embd * sizeof(float); + + common_sampler_reset(smpls[0].get()); + common_batch_clear(batch); + + // catch-up rows and the anchor exactly as the loop's first decode, then one placeholder row per later step + // (its cell receives that step's K/V; token and embd are not read) + for (int32_t k = 0; k < n_catch; ++k) { + common_batch_add(batch, d.tokens[k], d.pos[k], { 0 }, false); + std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, d.h.data() + (size_t) k * n_embd, row_bytes); + } + common_batch_add(batch, dp.id_last, dp.n_past, { 0 }, true); + std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[0].data(), row_bytes); + for (int32_t j = 1; j < n_draft; ++j) { + common_batch_add(batch, dp.id_last, dp.n_past + j, { 0 }, false); + std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[0].data(), row_bytes); + } + + llama_set_mtp_chain(ctx_dft, n_draft); + const int32_t rc = llama_decode(ctx_dft, batch); + llama_set_mtp_chain(ctx_dft, 0); + if (rc != 0) { + SPC_ERR("llama_decode(ctx_dft) chained draft returned %d\n", (int) rc); + if (n_catch > 0) { + d.catchup_failed = true; + } + return true; + } + if (n_catch > 0) { + d.pending = false; + } + + const llama_token * ids = llama_get_mtp_chain_ids(ctx_dft); + if (ids == nullptr) { + SPC_ERR("%s", "chained draft produced no tokens\n"); + return true; + } + + auto * smpl = smpls[0].get(); + auto & result = *dp.result; + const int32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx_dft))); + for (int32_t j = 0; j < n_draft; ++j) { + const llama_token id = ids[j]; + if (FILE * lf = common_mtp_draft_log()) { + // the chain reads back token ids only: no candidate list or probabilities (p logged as -1) + fprintf(lf, "{\"pos\": %d, \"d\": %d, \"c\": [[%d, -1.0]], \"stop\": 0, \"chain\": 1}\n", + (int) (dp.n_past + j + 1), j, (int) id); + } + if (id < 0 || id >= n_vocab) { + SPC_ERR("chained draft id %d out of vocab (step %d) - dropping the rest\n", (int) id, j); + break; + } + common_sampler_accept(smpl, id, true); + result.push_back(id); + } + return true; } ~common_speculative_impl_draft_mtp() override { @@ -2716,6 +2824,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } } + if (chain_mode && draft_chain(dparams)) { + auto & dp = dparams[0]; + if (dp.result->size() < (size_t) params.n_min) { + dp.result->clear(); + } + return; + } + common_batch_clear(batch); // keep track of which sequences are still drafting diff --git a/ggml/src/ggml-sycl/dequantize.hpp b/ggml/src/ggml-sycl/dequantize.hpp index 076380eb05a9..8a5c0b58d92e 100644 --- a/ggml/src/ggml-sycl/dequantize.hpp +++ b/ggml/src/ggml-sycl/dequantize.hpp @@ -67,6 +67,19 @@ static __dpct_inline__ int ptq1_0_trit(const uint8_t * qs, const uint8_t * qh, c return (int) ((v * 3) >> 8); } +// ARC-LAB: f32 output whatever GGML_SYCL_F16 says, for GET_ROWS (MTP chain draft-token embeddings). (trit - 1) * d +// with d widened from fp16 is exact in f32 (the product of 0 / +-1 and d), so the rows match the CPU's +// dequantize_row_ptq1_0 bit for bit. +static __dpct_inline__ void dequantize_ptq1_0_f32(const void *vx, const int64_t ib, + const int iqs, float &v0, float &v1) { + const block_ptq1_0 * x = (const block_ptq1_0 *) vx; + + const float d = (float) x[ib].d; + + v0 = (float) (ptq1_0_trit(x[ib].qs, x[ib].qh, iqs + 0) - 1) * d; + v1 = (float) (ptq1_0_trit(x[ib].qs, x[ib].qh, iqs + 1) - 1) * d; +} + static __dpct_inline__ void dequantize_q4_0(const void *vx, const int64_t ib, const int iqs, dfloat2 &v) { const block_q4_0 * x = (const block_q4_0 *) vx; diff --git a/ggml/src/ggml-sycl/getrows.cpp b/ggml/src/ggml-sycl/getrows.cpp index 187640446b48..c6029be7215d 100644 --- a/ggml/src/ggml-sycl/getrows.cpp +++ b/ggml/src/ggml-sycl/getrows.cpp @@ -334,6 +334,10 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { get_rows_sycl(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); break; + case GGML_TYPE_PTQ1_0: + get_rows_sycl_f32(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, + src1_i32, (float *)dst->data, ctx.stream()); + break; case GGML_TYPE_Q2_K: get_rows_sycl_f32(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 963c50806142..354b09854d40 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -7743,6 +7743,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons { switch (op->src[0]->type) { case GGML_TYPE_I32: + case GGML_TYPE_PTQ1_0: // ARC-LAB: dequantize_ptq1_0_f32, bit-exact with the CPU case GGML_TYPE_F16: case GGML_TYPE_BF16: case GGML_TYPE_F32: diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 77eb959cae7e..223273004829 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1338,6 +1338,119 @@ void llama_context::set_nextn_layer_offset(int32_t offset) { cparams.nextn_layer_offset = offset; } +// ARC-LAB: chained MTP drafts (independent implementation of the technique described by llamADALL). The draft loop in +// common/speculative.cpp decodes the NextN head once per drafted token, with a host round trip in between. A chain +// decode runs all steps in one graph (qwen35 graph_mtp::build_chain): step j+1 takes step j's argmax token and output +// state on the device. Bit 0: the chain graph applies to this context; bit 1: the head has a trimmed draft vocabulary, +// so in-graph tokens can be looked up in a small device copy of their token_embd rows. +int32_t llama_context::mtp_chain_support() const { + if (cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP || model.arch != LLM_ARCH_QWEN35 || model.hparams.n_layer_nextn != 1 || + cparams.nextn_layer_offset != 0 || !cparams.flash_attn || !cparams.embeddings_nextn_masked) { + return 0; + } + // per-step K/V views and index slices assume one plain KV cache with a non-transposed V + const auto * kv = dynamic_cast(memory.get()); + if (kv == nullptr || kv->get_v_trans()) { + return 0; + } + const auto & layer = model.layers[model.hparams.n_layer()]; + // same condition as graph_mtp's draft-head branch + const bool draft_head = layer.nextn.draft_head && layer.nextn.draft_head_ids && !layer.nextn.shared_head_head && + getenv("LLAMA_MTP_DRAFT_HEAD_OFF") == nullptr; + return 1 | (draft_head ? 2 : 0); +} + +void llama_context::set_mtp_chain(uint32_t n_steps) { + cparams.mtp_chain_n = n_steps; + if (n_steps == 0) { + return; // keeps the last chain's ids readable + } + mtp_chain_n_out = 0; + + if (mtp_chain_embd_tried) { + return; + } + mtp_chain_embd_tried = true; + + if ((mtp_chain_support() & 2) == 0) { + LLAMA_LOG_INFO("%s: MTP chain: no draft head, in-graph tokens use the host token_embd lookup\n", __func__); + return; + } + + const auto & layer = model.layers[model.hparams.n_layer()]; + ggml_tensor * ids = layer.nextn.draft_head_ids; + GGML_ASSERT(ids->type == GGML_TYPE_I32 && ggml_n_dims(ids) == 1); + mtp_chain_vocab.resize(ids->ne[0]); + ggml_backend_tensor_get(ids, mtp_chain_vocab.data(), 0, ggml_nbytes(ids)); + + const char * env_gpu = getenv("LLAMA_MTP_CHAIN_GPU_EMBD"); + if (env_gpu && atoi(env_gpu) == 0) { + LLAMA_LOG_INFO("%s: MTP chain: LLAMA_MTP_CHAIN_GPU_EMBD=0, in-graph tokens use the host token_embd lookup\n", __func__); + return; + } + + ggml_tensor * src = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + const ggml_backend_dev_t dev = model.dev_layer(model.hparams.n_layer()); + if (dev == nullptr || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU || + (src->type != GGML_TYPE_PTQ1_0 && src->type != GGML_TYPE_F16 && src->type != GGML_TYPE_F32)) { + LLAMA_LOG_INFO("%s: MTP chain: no device copy of token_embd (%s), host lookups\n", __func__, ggml_type_name(src->type)); + return; + } + + const int64_t n_rows = ids->ne[0]; + const size_t row = ggml_row_size(src->type, src->ne[0]); + GGML_ASSERT(src->nb[1] == row); + + ggml_init_params ip = { + /*.mem_size =*/ ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + mtp_chain_ctx.reset(ggml_init(ip)); + ggml_tensor * t = ggml_new_tensor_2d(mtp_chain_ctx.get(), src->type, src->ne[0], n_rows); + ggml_set_name(t, "mtp_chain_embd"); + mtp_chain_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(mtp_chain_ctx.get(), ggml_backend_dev_buffer_type(dev))); + if (!mtp_chain_buf) { + LLAMA_LOG_WARN("%s: MTP chain: could not allocate the device token_embd copy, host lookups\n", __func__); + mtp_chain_ctx.reset(); + return; + } + // a weight to the scheduler: ops reading it run on its device (no copy of the table per graph) + ggml_backend_buffer_set_usage(mtp_chain_buf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + // row k of the copy = token_embd row of draft-head index k, bytes unchanged (same quantization, same latent basis) + std::vector staging((size_t) n_rows * row); + for (int64_t k = 0; k < n_rows; ++k) { + const int32_t id = mtp_chain_vocab[k]; + GGML_ASSERT(id >= 0 && id < src->ne[1]); + ggml_backend_tensor_get(src, staging.data() + (size_t) k * row, (size_t) id * src->nb[1], row); + } + ggml_backend_tensor_set(t, staging.data(), 0, staging.size()); + mtp_chain_embd = t; + + LLAMA_LOG_INFO("%s: MTP chain: %lld draft-vocab token_embd rows (%s, %.1f MiB) on %s\n", __func__, (long long) n_rows, + ggml_type_name(src->type), staging.size() / 1048576.0, ggml_backend_dev_name(dev)); +} + +const int32_t * llama_context::get_mtp_chain_ids() { + synchronize(); + + if (mtp_chain_n_out == 0) { + return nullptr; + } + auto & tok = mtp_chain_tok; + tok.resize(mtp_chain_n_out); + for (uint32_t s = 0; s < mtp_chain_n_out; ++s) { + const int32_t v = mtp_chain_out[s]; + if (mtp_chain_vocab.empty()) { + tok[s] = v; + } else { + tok[s] = v >= 0 && v < (int32_t) mtp_chain_vocab.size() ? mtp_chain_vocab[v] : -1; + } + } + return tok.data(); +} + void llama_context::set_capture_layers(const std::vector & layer_ids) { const int32_t n_layer = (int32_t) model.hparams.n_layer(); if (layer_ids.size() > LLAMA_MAX_LAYERS) { @@ -1940,6 +2053,13 @@ int llama_context::decode(const llama_batch & batch_inp) { const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs(); + // ARC-LAB: a chained MTP draft is one graph over the whole batch (catch-up rows, anchor, one row per later step) + if (cparams.mtp_chain_n > 0 && (n_tokens_all > cparams.n_ubatch || n_tokens_all < cparams.mtp_chain_n)) { + LLAMA_LOG_ERROR("%s: MTP chain of %u steps needs %u..%u tokens in one ubatch, got %u\n", __func__, + cparams.mtp_chain_n, cparams.mtp_chain_n, cparams.n_ubatch, n_tokens_all); + return -1; + } + if (output_all) { // require that all tokens are output if (n_outputs_all != n_tokens_all) { @@ -2196,6 +2316,18 @@ int llama_context::decode(const llama_batch & batch_inp) { extract_layer_inputs(res, n_tokens_prev, ubatch.n_tokens); + // ARC-LAB: chained MTP draft: the per-step argmax rows, mapped to tokens in get_mtp_chain_ids() + if (!res->t_mtp_chain.empty()) { + mtp_chain_out.resize(res->t_mtp_chain.size()); + for (size_t s = 0; s < res->t_mtp_chain.size(); ++s) { + ggml_tensor * t = res->t_mtp_chain[s]; + ggml_backend_t b = ggml_backend_sched_get_tensor_backend(sched.get(), t); + GGML_ASSERT(b != nullptr); + ggml_backend_tensor_get_async(b, t, &mtp_chain_out[s], 0, sizeof(int32_t)); + } + mtp_chain_n_out = (uint32_t) res->t_mtp_chain.size(); + } + // extract nextn embeddings before // only meaningful in LLAMA_POOLING_TYPE_NONE (per-token); other pooling modes are ignored. { @@ -2695,6 +2827,10 @@ ggml_cgraph * llama_context::graph_reserve( this->n_outputs = n_outputs; + // ARC-LAB: reserve with the regular graph, never a chained MTP draft (it needs real tokens and hidden rows) + const uint32_t save_mtp_chain_n = cparams.mtp_chain_n; + cparams.mtp_chain_n = 0; + llama_batch_allocr balloc(model.hparams.n_pos_per_embd()); llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs); @@ -2716,6 +2852,7 @@ ggml_cgraph * llama_context::graph_reserve( } this->n_outputs = save_n_outputs; + cparams.mtp_chain_n = save_mtp_chain_n; // initialize scheduler with the specified graph if (split_only) { @@ -2756,6 +2893,7 @@ llm_graph_params llama_context::graph_params( /*.dspark_ctx_width =*/dspark_ctx.n_embd_cap, /*.hadamard_rotations =*/&hadamard_rotations, /*.hadamard_inverses =*/&hadamard_inverses, + /*.mtp_chain_embd =*/mtp_chain_embd, /*.samplers =*/sampling.samplers, /*.n_outputs =*/n_outputs, /*.cb =*/graph_get_cb(), @@ -4081,6 +4219,18 @@ void llama_set_nextn_layer_offset(llama_context * ctx, int32_t offset) { ctx->set_nextn_layer_offset(offset); } +int32_t llama_mtp_chain_support(const llama_context * ctx) { + return ctx->mtp_chain_support(); +} + +void llama_set_mtp_chain(llama_context * ctx, int32_t n_steps) { + ctx->set_mtp_chain(n_steps > 0 ? (uint32_t) n_steps : 0u); +} + +const llama_token * llama_get_mtp_chain_ids(llama_context * ctx) { + return ctx->get_mtp_chain_ids(); +} + llama_memory_t llama_get_memory(const struct llama_context * ctx) { if (!ctx) { return nullptr; diff --git a/src/llama-context.h b/src/llama-context.h index 6057da7f435f..ef04bf85de9c 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -134,6 +134,11 @@ struct llama_context { void set_embeddings_layer_inp(uint32_t lid, bool enable); void set_nextn_layer_offset(int32_t offset); + + // ARC-LAB: chained MTP drafts (llama-ext.h llama_mtp_chain_support / llama_set_mtp_chain / llama_get_mtp_chain_ids) + int32_t mtp_chain_support() const; + void set_mtp_chain(uint32_t n_steps); + const int32_t * get_mtp_chain_ids(); void set_causal_attn(bool value); void set_warmup(bool value); @@ -403,6 +408,18 @@ struct llama_context { // one-time Hadamard transform-coverage check on the first built graph bool hadamard_verified = false; + // ARC-LAB: chained MTP drafts. mtp_chain_embd: device copy of the token_embd rows the draft head can produce (row k + // = token mtp_chain_vocab[k]), built on first use unless LLAMA_MTP_CHAIN_GPU_EMBD=0. mtp_chain_out: the last chain + // decode's per-step argmax rows, read back asynchronously. + ggml_tensor * mtp_chain_embd = nullptr; + ggml_context_ptr mtp_chain_ctx; + ggml_backend_buffer_ptr mtp_chain_buf; + bool mtp_chain_embd_tried = false; + std::vector mtp_chain_vocab; // draft-head index -> token id (empty: no draft head, indices are tokens) + std::vector mtp_chain_out; + std::vector mtp_chain_tok; // mtp_chain_out mapped to token ids + uint32_t mtp_chain_n_out = 0; + // host buffer for the model output (logits and embeddings) ggml_backend_buffer_ptr buf_output; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 1f013e2f0027..248149614401 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -23,6 +23,9 @@ struct llama_cparams { int32_t nextn_layer_offset = 0; + // ARC-LAB: > 0: the next DECODER_MTP graph drafts this many tokens in one graph (llama_set_mtp_chain) + uint32_t mtp_chain_n = 0; + float rope_freq_base; float rope_freq_scale; diff --git a/src/llama-ext.h b/src/llama-ext.h index c6cfc35bc3cc..55de2f595e27 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -101,6 +101,16 @@ LLAMA_API void llama_set_embeddings_nextn(struct llama_context * ctx, bool value // chain multiple trained NextN heads. Default 0 (first head). LLAMA_API void llama_set_nextn_layer_offset(struct llama_context * ctx, int32_t offset); +// ARC-LAB: chained MTP drafts (single NextN head, qwen35). llama_mtp_chain_support: 0 = not available on this context, +// bit 0 = available, bit 1 = the head has a trimmed draft vocabulary (in-graph token lookups can stay on the device). +// llama_set_mtp_chain(ctx, n): the next llama_decode on this MTP context drafts n tokens in one graph. Its batch must be +// one ubatch: catch-up rows (no output), the anchor row (id_last at n_past, output), then n-1 rows at n_past+1.. whose +// token and embd are placeholders (the graph writes their K/V). Reset with n = 0 after the decode. +// llama_get_mtp_chain_ids: after that decode, the n greedy draft tokens (synchronizes), or nullptr. +LLAMA_API int32_t llama_mtp_chain_support(const struct llama_context * ctx); +LLAMA_API void llama_set_mtp_chain (struct llama_context * ctx, int32_t n_steps); +LLAMA_API const llama_token * llama_get_mtp_chain_ids(struct llama_context * ctx); + // mirrors: // LLAMA_API float * llama_get_embeddings(struct llama_context * ctx); LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9ad7fa925e6b..ceabee25866a 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1354,6 +1354,7 @@ void llm_graph_result::reset() { t_sampled.clear(); t_dspark_greedy.clear(); + t_mtp_chain.clear(); t_sampled_probs.clear(); t_sampled_logits.clear(); t_candidates.clear(); @@ -1521,6 +1522,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) : dspark_ctx_width(params.dspark_ctx_width), hadamard_rotations(params.hadamard_rotations), hadamard_inverses(params.hadamard_inverses), + mtp_chain_embd(params.mtp_chain_embd), samplers(params.samplers), cb_func(params.cb), res(params.res), diff --git a/src/llama-graph.h b/src/llama-graph.h index 6c2db6420d64..ea9952557a47 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -863,6 +863,7 @@ struct llm_graph_params { int64_t dspark_ctx_width; const llama_hadamard_rotations * hadamard_rotations; const llama_hadamard_rotations * hadamard_inverses; + ggml_tensor * mtp_chain_embd; // ARC-LAB: device copy of the draft-vocab token_embd rows, or null std::map samplers; @@ -931,6 +932,11 @@ struct llm_graph_params { return false; } + // ARC-LAB: a chained MTP draft graph bakes per-step KV lengths in; never reuse it or reuse into it + if (cparams.mtp_chain_n > 0 || other.cparams.mtp_chain_n > 0) { + return false; + } + if (!samplers_equal(samplers, other.samplers)) { return false; } @@ -1035,6 +1041,10 @@ class llm_graph_result { std::vector t_sampled; std::vector t_dspark_greedy; + + // ARC-LAB: chained MTP draft (cparams.mtp_chain_n): per step, the argmax row of the draft logits (I32 [1]); an + // index into the draft head's vocabulary when it has one, else a token id + std::vector t_mtp_chain; std::vector t_sampled_probs; std::vector t_sampled_logits; std::vector t_candidates; @@ -1127,6 +1137,7 @@ struct llm_graph_context { int64_t dspark_ctx_width; const llama_hadamard_rotations * hadamard_rotations; const llama_hadamard_rotations * hadamard_inverses; + ggml_tensor * mtp_chain_embd; // Transforms shared by folded weights on the same activation. Key is (input, rotation); // both must match. Valid for one graph build only. diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index d406353d564c..cb806c688166 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1236,6 +1236,32 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const { return v_cells[seq_to_stream[seq_id]]; } +uint32_t llama_kv_cache::get_n_kv_prefix(const slot_info & sinfo, uint32_t n_rows) const { + GGML_ASSERT(sinfo.n_stream() == 1); + + const auto & cells = v_cells[sinfo.strm[0]]; + const auto & idxs = sinfo.idxs[0]; + + // highest used cell, not counting the cells of rows >= n_rows (they were empty before this ubatch) + uint32_t p1 = cells.used_max_p1(); + while (p1 > 0) { + const uint32_t c = p1 - 1; + bool later = false; + for (size_t k = n_rows; k < idxs.size() && !later; ++k) { + later = idxs[k] == c; + } + if (!later && !cells.is_empty(c)) { + break; + } + --p1; + } + + // same padding as get_n_kv + const uint32_t n_pad_cur = std::max(n_pad, 256u); + + return std::min(cells.size(), std::max(n_pad_cur, GGML_PAD(p1, n_pad_cur))); +} + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { uint32_t result = 0; @@ -2783,6 +2809,22 @@ uint32_t llama_kv_cache_context::get_n_kv() const { return n_kv; } +uint32_t llama_kv_cache_context::get_n_kv_prefix(uint32_t n_rows) const { + return kv->get_n_kv_prefix(sinfos[i_cur], n_rows); +} + +ggml_tensor * llama_kv_cache_context::get_k_n(ggml_context * ctx, int32_t il, uint32_t n_kv_view) const { + return kv->get_k(ctx, il, n_kv_view, sinfos[i_cur]); +} + +ggml_tensor * llama_kv_cache_context::get_v_n(ggml_context * ctx, int32_t il, uint32_t n_kv_view) const { + return kv->get_v(ctx, il, n_kv_view, sinfos[i_cur]); +} + +bool llama_kv_cache_context::is_v_trans() const { + return kv->get_v_trans(); +} + ggml_type llama_kv_cache_context::type_k() const { return kv->type_k(); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 4d2229b90274..05add1b853c1 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -172,6 +172,11 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_n_kv(const slot_info & sinfo) const; + // ARC-LAB: get_n_kv as it would be if only the first n_rows of the applied ubatch held their cells (the KV length + // a decode of just those rows would attend over; chained MTP drafts size each step's view with it) + uint32_t get_n_kv_prefix(const slot_info & sinfo, uint32_t n_rows) const; + bool get_v_trans() const { return v_trans; } + // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; @@ -392,6 +397,13 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; + // ARC-LAB: chained MTP drafts. KV length a decode of only the first n_rows of this ubatch would see, and views + // of the cache with an explicit length (the same views get_k / get_v return for that length) + uint32_t get_n_kv_prefix(uint32_t n_rows) const; + ggml_tensor * get_k_n(ggml_context * ctx, int32_t il, uint32_t n_kv_view) const; + ggml_tensor * get_v_n(ggml_context * ctx, int32_t il, uint32_t n_kv_view) const; + bool is_v_trans() const; + ggml_type type_k() const; ggml_type type_v() const; diff --git a/src/models/models.h b/src/models/models.h index a82133099d5c..81993e40ab24 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2302,6 +2302,9 @@ struct llama_model_qwen35 : public llama_model_base { struct graph_mtp : public llm_graph_context { graph_mtp(const llama_model & model, const llm_graph_params & params); + + // ARC-LAB: cparams.mtp_chain_n > 0: all draft steps in one graph (see llama_set_mtp_chain) + void build_chain(const llama_model & model); }; std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index ecc434f3e31f..c5fedfead481 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -609,6 +609,52 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_ffn(ggml_tensor * cur, cons return cur; } +// ARC-LAB: per-step inputs of a chained MTP draft graph (graph_mtp::build_chain): the positions of each step's rows in +// the layout build_inp_pos() gives a decode of just those rows, and the step's output row +class llm_graph_input_mtp_chain : public llm_graph_input_i { +public: + explicit llm_graph_input_mtp_chain(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {} + + struct step { + int64_t r0; // first ubatch row of the step + int64_t nr; // rows + ggml_tensor * pos; // I32 [nr * n_pos_per_embd] + ggml_tensor * out_idx; // I32 [1] = nr - 1 (the step's last row produces its draft) + }; + std::vector steps; + + void set_input(const llama_ubatch * ubatch) override { + const int64_t n_tokens = ubatch->n_tokens; + std::vector v; + for (const auto & st : steps) { + v.assign((size_t) st.nr * n_pos_per_embd, 0); + for (int64_t r = 0; r < st.nr; ++r) { + if (ubatch->token && n_pos_per_embd == 4) { + // text tokens under M-RoPE, as llm_graph_input_pos: three equal dims, the fourth 0 + for (int d = 0; d < 3; ++d) { + v[d * st.nr + r] = ubatch->pos[st.r0 + r]; + } + } else { + for (uint32_t d = 0; d < n_pos_per_embd; ++d) { + v[d * st.nr + r] = ubatch->pos[d * n_tokens + st.r0 + r]; + } + } + } + ggml_backend_tensor_set(st.pos, v.data(), 0, v.size() * sizeof(int32_t)); + const int32_t out_row = (int32_t) st.nr - 1; + ggml_backend_tensor_set(st.out_idx, &out_row, 0, sizeof(out_row)); + } + } + + // never reused (llm_graph_params::allow_reuse refuses chain graphs) + bool can_reuse(const llm_graph_params & params) override { + GGML_UNUSED(params); + return false; + } + + const uint32_t n_pos_per_embd; +}; + // LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3.5/3.6 dense series llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { @@ -618,6 +664,11 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + if (cparams.mtp_chain_n > 0) { + build_chain(model); + return; + } + // hparams.n_layer includes both main model layers and MTP layers. The MTP // layer is stored immediately after the main layers in model.layers[]. const int il = hparams.n_layer(); @@ -856,3 +907,247 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr res->t_logits = cur; ggml_build_forward_expand(gf, cur); } + +// ARC-LAB: chained MTP drafts, n = cparams.mtp_chain_n steps in one graph (independent implementation of the technique +// described by llamADALL). The ubatch holds the deferred catch-up rows and the anchor (step 0, as the first decode of the +// draft loop in common/speculative.cpp) followed by one placeholder row per later step, so the KV cache has given every +// draft position its cell and the mask rows exist. Step j > 0 reads step j-1's argmax token and output state (h_nextn +// row) on the device and writes its K/V into its row's cell, as decode j of the loop does. +// +// Each step is built with the loop's shapes, so the kernels and their rounding are the loop's: its own rows only (step 0 +// the catch-up rows + anchor, later steps one row), its own positions, KV index slices and mask rows, and K/V views and +// mask columns of the KV length a decode of the rows up to and including its own would see (get_n_kv_prefix; flash +// attention splits the KV range by that length). The ops mirror the single-step graph above; keep the two in sync. +void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + const int64_t n_steps = cparams.mtp_chain_n; + GGML_ASSERT(n_steps >= 1 && n_tokens >= n_steps); + GGML_ASSERT(ubatch.token && ubatch.embd && "MTP chain needs tokens and hidden rows"); + const int64_t n_lead = n_tokens - n_steps + 1; // step 0: catch-up rows + the anchor + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + // inputs as the single-step graph creates them + auto inp = std::make_unique(hparams.n_embd); + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + ggml_tensor * tokens_all = inp->tokens; + ggml_tensor * h_all = inp->h; + res->add_input(std::move(inp)); + + auto * inp_attn = build_attn_inp_kv(); + const llama_kv_cache_context * kvc = inp_attn->mctx; + + auto chain_inp = std::make_unique(hparams.n_pos_per_embd()); + for (int64_t s = 0; s < n_steps; ++s) { + llm_graph_input_mtp_chain::step st; + st.r0 = s == 0 ? 0 : n_lead - 1 + s; + st.nr = s == 0 ? n_lead : 1; + st.pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, st.nr * hparams.n_pos_per_embd()); + st.out_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1); + ggml_set_input(st.pos); + ggml_set_input(st.out_idx); + chain_inp->steps.push_back(st); + } + const auto steps = chain_inp->steps; + res->add_input(std::move(chain_inp)); + + // token embeddings: the host table for the ubatch's own tokens (catch-up tokens need not be in the draft + // vocabulary), the draft-vocab device copy (row k = token draft_head_ids[k]) for the argmax tokens + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + const llama_hadamard_transform * inv = nullptr; + if (hadamard_inverses) { + const auto it = hadamard_inverses->find(tok_embd_w); + if (it != hadamard_inverses->end()) { + inv = &it->second; + } + } + auto unrotate = [&](ggml_tensor * e) { + // the single-step graph's inverse transform after the lookup (also applies to the copy: same rows) + if (inv) { + e = llama_mul_mat_hadamard(ctx0, e, inv->rot); + if (inv->signs) { + e = ggml_mul(ctx0, e, inv->signs); + } + } + return e; + }; + + static const bool draft_head_off = getenv("LLAMA_MTP_DRAFT_HEAD_OFF") != nullptr; + const bool use_draft_head = layer.nextn.draft_head && layer.nextn.draft_head_ids && !layer.nextn.shared_head_head && + !draft_head_off; + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm; + GGML_ASSERT(head_w && head_norm_w); + + static bool logged = false; + if (!logged) { + logged = true; + LLAMA_LOG_INFO("%s: MTP chain graph (%lld steps, %s, %s token lookups)\n", __func__, (long long) n_steps, + use_draft_head ? "draft head" : "full head", use_draft_head && mtp_chain_embd ? "device" : "host"); + } + + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + ggml_tensor * tok_embd = unrotate(ggml_get_rows(ctx0, tok_embd_w, ggml_view_1d(ctx0, tokens_all, n_lead, 0))); + ggml_tensor * h_embd = ggml_view_2d(ctx0, h_all, h_all->ne[0], n_lead, h_all->nb[1], 0); + + for (int64_t s = 0; s < n_steps; ++s) { + const auto & st = steps[s]; + const int64_t nr = st.nr; + cb(tok_embd, "mtp_tok_embd", il); + + // this step's slice of the ubatch's KV indices and mask, sized as a decode of rows [0, r0 + nr) would be + const uint32_t n_kv_s = kvc->get_n_kv_prefix((uint32_t) (st.r0 + nr)); + ggml_tensor * k_idxs = inp_attn->get_k_idxs(); + ggml_tensor * v_idxs = inp_attn->get_v_idxs(); + k_idxs = ggml_view_1d(ctx0, k_idxs, nr, st.r0 * ggml_element_size(k_idxs)); + v_idxs = ggml_view_1d(ctx0, v_idxs, nr, st.r0 * ggml_element_size(v_idxs)); + ggml_tensor * kq_mask = inp_attn->get_kq_mask(); + kq_mask = ggml_cont(ctx0, ggml_view_2d(ctx0, kq_mask, n_kv_s, nr, kq_mask->nb[1], st.r0 * kq_mask->nb[1])); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + ggml_tensor * inpSA = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + cb(Qcur_full, "mtp_Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, + n_embd_head, n_head, nr, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + 0); + Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "mtp_Qcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, + n_embd_head, n_head, nr, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, nr); + cb(gate, "mtp_gate", il); + + ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, nr); + Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "mtp_Kcur_normed", il); + + ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, nr); + cb(Vcur, "mtp_Vcur", il); + + Qcur = ggml_rope_multi(ctx0, Qcur, st.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, st.pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + // build_attn() for this step's rows: same rotations, same expand order, store, then attend + if (inp_attn->self_k_rot) { + Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot); + Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot); + } + if (inp_attn->self_v_rot) { + Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot); + } + ggml_build_forward_expand(gf, Qcur); + ggml_build_forward_expand(gf, Vcur); + ggml_build_forward_expand(gf, Kcur); + cb(Kcur, "k_cache_in", il); + ggml_build_forward_expand(gf, kvc->cpy_k(ctx0, Kcur, k_idxs, il)); + ggml_build_forward_expand(gf, kvc->cpy_v(ctx0, Vcur, v_idxs, il)); + + cur = build_attn_mha(Qcur, kvc->get_k_n(ctx0, il, n_kv_s), kvc->get_v_n(ctx0, il, n_kv_s), nullptr, kq_mask, + nullptr, nullptr, kq_scale, il); + cb(cur, "kqv_out", il); + if (inp_attn->self_v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot); + } + cb(cur, "mtp_attn_pregate", il); + + cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate)); + cur = build_lora_mm(layer.wo, cur, layer.wo_s); + cb(cur, "mtp_attn_out", il); + + cur = ggml_add(ctx0, cur, inpSA); + cb(cur, "mtp_attn_residual", il); + + ggml_tensor * ffn_residual = cur; + cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_post_norm", il); + + cur = build_ffn(cur, + layer.ffn_up, nullptr, layer.ffn_up_s, + layer.ffn_gate, nullptr, layer.ffn_gate_s, + layer.ffn_down, nullptr, layer.ffn_down_s, + nullptr, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "mtp_ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_residual); + cb(cur, "mtp_post_ffn", il); + + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_nextn", -1); + + // the step's output row (masked h_nextn: the next step's hidden input) + ggml_tensor * h_out = ggml_get_rows(ctx0, cur, st.out_idx); + cb(h_out, "mtp_shared_head_norm", -1); + + // greedy draft: argmax of the head's logits. With a draft head that is an index into its vocabulary (the + // -inf scatter of the single-step graph only maps it to a token id, which the context does on the host) + ggml_tensor * logits = use_draft_head ? build_lora_mm(layer.nextn.draft_head, h_out, nullptr) + : build_lora_mm(head_w, h_out, head_s); + cb(logits, "result_output", -1); + ggml_tensor * best = ggml_argmax(ctx0, logits); + ggml_set_name(best, "mtp_chain_argmax"); + ggml_set_output(best); + ggml_build_forward_expand(gf, best); + res->t_mtp_chain.push_back(best); + + if (s + 1 == n_steps) { + break; + } + + // next step's inputs: this step's draft token and output row + ggml_tensor * e; + if (use_draft_head && mtp_chain_embd) { + e = ggml_get_rows(ctx0, mtp_chain_embd, best); + } else if (use_draft_head) { + ggml_tensor * ids = layer.nextn.draft_head_ids; + ggml_tensor * tok = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, ids, 1, ids->ne[0]), best); // I32 [1, 1] + e = ggml_get_rows(ctx0, tok_embd_w, ggml_reshape_1d(ctx0, tok, 1)); + } else { + e = ggml_get_rows(ctx0, tok_embd_w, best); + } + tok_embd = unrotate(e); + h_embd = h_out; + } +} diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 327a63de1ed2..3c25cc124f14 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2252,13 +2252,30 @@ struct test_get_rows : public test_case { const int be1; // batch size const int be2; // batch size const bool v; // view (non-contiguous src1) + const bool exact; // ARC-LAB: require bit-identical rows (max NMSE 0), e.g. PTQ1_0 MTP draft embeddings std::string vars() override { - return VARS_TO_STR7(type, n, m, r, be1, be2, v); + return exact ? VARS_TO_STR8(type, n, m, r, be1, be2, v, exact) : VARS_TO_STR7(type, n, m, r, be1, be2, v); } - test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false) - : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v) {} + double max_nmse_err() override { + return exact ? 0.0 : 1e-7; + } + + double err(const float * a, const float * b, size_t n) override { + if (!exact) { + return test_case::err(a, b, n); + } + size_t n_diff = 0; // bit patterns, so -0.0 vs 0.0 counts too + for (size_t i = 0; i < n; ++i) { + n_diff += memcmp(a + i, b + i, sizeof(float)) != 0; + } + return (double) n_diff; + } + + test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false, + bool exact = false) + : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v), exact(exact) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * in = ggml_new_tensor_4d(ctx, type, n, m, be1, be2); @@ -8868,6 +8885,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_get_rows(GGML_TYPE_I32, 256, 5, 4, b, 1, v)); } } + // ARC-LAB: MTP chain draft-token lookups (PTQ1_0 token_embd rows, Bonsai width 5120, 1..8 rows): bit-exact + for (int r : {1, 4, 8}) { + test_cases.emplace_back(new test_get_rows(GGML_TYPE_PTQ1_0, 5120, 300, r, 1, 1, false, true)); + } + test_cases.emplace_back(new test_get_rows(GGML_TYPE_PTQ1_0, 5120, 300, 8, 3, 1, true, true)); test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 8, 2, 1, false)); test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 70000, 4, 1, false)); // row count > CUDA grid-y limit (65535) From 24b56006672cd0b43a49195073b41b72c699b5aa Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 13:15:45 +1100 Subject: [PATCH 06/24] gdn: rollback by replay instead of state snapshots (LLAMA_GDN_REPLAY) Independent implementation of the technique described by llamADALL (clean-room, from the spec and this fork's own code only). With MTP, each verify pass used to write K = n_rs_seq + 1 full DeltaNet states per layer. The new op GGML_OP_GATED_DELTA_NET_REPLAY keeps one base state (the oldest one a rollback can need) plus a small update record (delta, k, g) per later token. The next pass first re-applies the records of the tokens that were kept, then runs the new tokens. The replay repeats the step's own state update, so the state is bit-identical to the snapshot path. - ggml: new op appended at the end of the enum (GGML_OP_COUNT 102, RPC patch version bump); CPU reference. - SYCL: sequential replay kernel. The blocked kernel can emit records (template flag; the existing instantiations are unchanged). A fused cache write (CPY + SET_ROWS skipped) and the gather-skip are used. - llama: recurrent memory keeps the base in plane 0 and two record banks in plane 1; the tensor layout is unchanged. Partial rollback = fewer records to replay. State save/restore carries the pending records (s_trans = 2). The conv state keeps its snapshot planes. - On by default for a single sequence, qwen35, all recurrent layers on SYCL. LLAMA_GDN_REPLAY=0 restores the snapshot path. - test-backend-ops: GATED_DELTA_NET_REPLAY chain tests (replay == snapshot exactly, and vs CPU). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit e37b71a33b4b55853859ca096c156d313732bfc9 on Torchit1/llama.cpp arc-b580) --- ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 30 +++ ggml/src/ggml-cpu/ggml-cpu.c | 11 + ggml/src/ggml-cpu/ops.cpp | 145 ++++++++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml-sycl/gated_delta_net.cpp | 353 ++++++++++++++++++++++++- ggml/src/ggml-sycl/gated_delta_net.hpp | 13 + ggml/src/ggml-sycl/ggml-sycl.cpp | 94 ++++++- ggml/src/ggml.c | 67 ++++- src/llama-graph.cpp | 39 +++ src/llama-graph.h | 7 + src/llama-memory-recurrent.cpp | 165 +++++++++++- src/llama-memory-recurrent.h | 22 +- src/models/delta-net-base.cpp | 40 +++ src/models/qwen35.cpp | 13 +- tests/test-backend-ops.cpp | 213 ++++++++++++++- 16 files changed, 1192 insertions(+), 25 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 059e4496269a..545b892c0eed 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 5 #define RPC_PROTO_MINOR_VERSION 1 -#define RPC_PROTO_PATCH_VERSION 0 +#define RPC_PROTO_PATCH_VERSION 1 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d875a15ed066..c9339461cf8c 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -596,6 +596,8 @@ extern "C" { GGML_OP_GLU, + GGML_OP_GATED_DELTA_NET_REPLAY, // ARC-LAB: appended last so the existing op numbers do not move + GGML_OP_COUNT, }; @@ -2628,6 +2630,34 @@ extern "C" { struct ggml_tensor * dt_bias, struct ggml_tensor * a); + // ARC-LAB: gated delta net with rollback by replay (independent implementation of the technique described by + // llamADALL). Same recurrence as ggml_gated_delta_net (scalar gate only), but instead of K state snapshots the op + // keeps ONE base state plus, for each later token, its update record: the per-token update is + // S <- g * S + k (x) delta, so a record is (delta [S_v], k [S_v], g [1]) per head, R = H*(2*S_v + 1) floats: + // record = [ delta: S_v*H | k: S_v*H (head h holds k of head h % H_k) | g: H ] (g = the decay factor exp(gate)) + // + // q, k, v, g, beta : as ggml_gated_delta_net (g: [1, H, n_tokens, n_seqs]) + // state : [S_v, S_v, H, n_seqs] stored base state + // rec : F32 [R*(K-1), 2, n_seqs] stored records, two banks per seq (a pass reads one bank, the caller + // stores the new records in the other one) + // ctrl : I32 [2, n_seqs] = {n_rec, bank}: records 0..n_rec-1 of `bank` are first re-applied to state + // + // output, packed: [ attn scores S_v*H*n_tokens*n_seqs | base D*n_seqs | records (K-1)*R*n_seqs ], D = S_v*S_v*H + // with m = min(n_tokens - 1, K - 1): base = state after token n_tokens-1-m (the oldest state a rollback of up + // to m tokens can need), record j = token n_tokens-m+j for j < m; records m..K-2 are left unwritten. + // Rolling back r <= m tokens = re-applying m - r records to base on the next call. + GGML_API struct ggml_tensor * ggml_gated_delta_net_replay( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * rec, + struct ggml_tensor * ctrl, + int K); + // DSA lightning indexer // // q: [n_embd_idx, n_head_idx, n_batch, ne3 ] diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 10d4d690c862..666c1576011a 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2112,6 +2112,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_gated_delta_net(params, tensor); } break; + case GGML_OP_GATED_DELTA_NET_REPLAY: + { + ggml_compute_forward_gated_delta_net_replay(params, tensor); + } break; case GGML_OP_LIGHTNING_INDEXER: { ggml_compute_forward_lightning_indexer(params, tensor); @@ -2308,6 +2312,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_COUNT_EQUAL: case GGML_OP_SOLVE_TRI: case GGML_OP_GATED_DELTA_NET: + case GGML_OP_GATED_DELTA_NET_REPLAY: case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: @@ -3045,6 +3050,12 @@ struct ggml_cplan ggml_graph_plan( const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0); cur = per_thread * sizeof(float) * n_tasks; } break; + case GGML_OP_GATED_DELTA_NET_REPLAY: + { + // per thread: working state [S_v, S_v] + delta [S_v] + const int64_t S_v = node->src[2]->ne[0]; + cur = (S_v * S_v + S_v) * sizeof(float) * n_tasks; + } break; case GGML_OP_COUNT: { GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index c7992c78bd14..ccb5ec86b4bb 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11171,6 +11171,151 @@ void ggml_compute_forward_gated_delta_net( } } +// ARC-LAB: ggml_compute_forward_gated_delta_net_replay (rollback by replay, see ggml_gated_delta_net_replay). +// The per-token step is the scalar-gate path of ggml_compute_forward_gated_delta_net_one_chunk, operation for +// operation, and a replayed record repeats exactly its state update (scale by g, then mad k * delta), so the replayed +// state is bit-identical to the state that path computes for the same token. +static void ggml_compute_forward_gated_delta_net_replay_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src_q = dst->src[0]; + const ggml_tensor * src_k = dst->src[1]; + const ggml_tensor * src_v = dst->src[2]; + const ggml_tensor * src_g = dst->src[3]; + const ggml_tensor * src_beta = dst->src[4]; + const ggml_tensor * src_state = dst->src[5]; + const ggml_tensor * src_rec = dst->src[6]; + const ggml_tensor * src_ctrl = dst->src[7]; + + GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne); + GGML_TENSOR_LOCALS(size_t, nbq, src_q, nb); + GGML_TENSOR_LOCALS(int64_t, nek, src_k, ne); + GGML_TENSOR_LOCALS(size_t, nbk, src_k, nb); + GGML_TENSOR_LOCALS(int64_t, nev, src_v, ne); + GGML_TENSOR_LOCALS(size_t, nbv, src_v, nb); + GGML_TENSOR_LOCALS(size_t, nbg, src_g, nb); + GGML_TENSOR_LOCALS(size_t, nbb, src_beta, nb); + + const int64_t S_v = nev0; + const int64_t H = nev1; + const int64_t n_tokens = nev2; + const int64_t n_seqs = nev3; + const int64_t K = ggml_get_op_params_i32(dst, 0); + const int64_t R = H * (2 * S_v + 1); + const int64_t D = S_v * S_v * H; + + // base = state after token t_base; tokens after it leave a record + const int64_t m = std::min(n_tokens - 1, K - 1); + const int64_t t_base = n_tokens - 1 - m; + + float * attn_out_base = (float *) dst->data; + float * base_out = attn_out_base + S_v * H * n_tokens * n_seqs; + float * rec_out = base_out + D * n_seqs; + + const int32_t * ctrl = (const int32_t *) src_ctrl->data; + + const int64_t rq3 = nev3 / neq3; + const int64_t rk3 = nev3 / nek3; + + const float scale = 1.0f / sqrtf((float) S_v); + + const int ith = params->ith; + const int nth = params->nth; + + float * s = (float *) params->wdata + ith * (S_v * S_v + S_v + CACHE_LINE_SIZE_F32); + float * delta = s + S_v * S_v; + + const int64_t nr = H * n_seqs; + const int64_t dr = (nr + nth - 1) / nth; + const int64_t ir0 = dr * ith; + const int64_t ir1 = std::min(ir0 + dr, nr); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t iv1 = ir % H; // head + const int64_t iv3 = ir / H; // sequence + + const int64_t iq1 = iv1 % neq1; + const int64_t ik1 = iv1 % nek1; + const int64_t iq3 = iv3 / rq3; + const int64_t ik3 = iv3 / rk3; + + // state layout: s[j*S_v + i] = S[i][j] (row j = value column j) + memcpy(s, (const char *) src_state->data + iv3 * src_state->nb[3] + iv1 * S_v * S_v * sizeof(float), + S_v * S_v * sizeof(float)); + + // replay: re-apply the stored updates of the tokens kept since the base state + const int32_t n_rec = ctrl[2 * iv3 + 0]; + const int32_t bank = ctrl[2 * iv3 + 1]; + GGML_ASSERT(n_rec >= 0 && n_rec <= K - 1 && (bank == 0 || bank == 1)); + const float * rec_in = (const float *) ((const char *) src_rec->data + iv3 * src_rec->nb[2] + bank * src_rec->nb[1]); + for (int32_t j = 0; j < n_rec; ++j) { + const float * r = rec_in + j * R; + const float * rd = r + iv1 * S_v; + const float * rk = r + H * S_v + iv1 * S_v; + ggml_vec_scale_f32(S_v * S_v, s, r[2 * H * S_v + iv1]); + for (int64_t jj = 0; jj < S_v; ++jj) { + ggml_vec_mad_f32(S_v, &s[jj * S_v], rk, rd[jj]); + } + } + + float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v; + + for (int64_t t = 0; t < n_tokens; t++) { + const float * q_d = (const float *)((const char *)src_q->data + iq3 * nbq3 + t * nbq2 + iq1 * nbq1); + const float * k_d = (const float *)((const char *)src_k->data + ik3 * nbk3 + t * nbk2 + ik1 * nbk1); + const float * v_d = (const float *)((const char *)src_v->data + iv3 * nbv3 + t * nbv2 + iv1 * nbv1); + + const float beta_val = *(const float *)((const char *)src_beta->data + iv3 * nbb3 + t * nbb2 + iv1 * nbb1); + const float g0 = *(const float *)((const char *)src_g->data + iv3 * nbg3 + t * nbg2 + iv1 * nbg1); + const float g_val = expf(g0); + + ggml_vec_scale_f32(S_v * S_v, s, g_val); + + for (int64_t j = 0; j < S_v; ++j) { + float sum = 0.0f; + ggml_vec_dot_f32(S_v, &sum, 0, &s[j * S_v], 0, k_d, 0, 1); + delta[j] = (v_d[j] - sum) * beta_val; + } + + for (int64_t j = 0; j < S_v; ++j) { + ggml_vec_mad_f32(S_v, &s[j * S_v], k_d, delta[j]); + } + + for (int64_t j = 0; j < S_v; ++j) { + float sum = 0.0f; + ggml_vec_dot_f32(S_v, &sum, 0, &s[j * S_v], 0, q_d, 0, 1); + attn_data[j] = sum * scale; + } + + attn_data += S_v * H; + + if (t == t_base) { + memcpy(base_out + (iv3 * H + iv1) * S_v * S_v, s, S_v * S_v * sizeof(float)); + } else if (t > t_base) { + float * r = rec_out + iv3 * (K - 1) * R + (t - t_base - 1) * R; + memcpy(r + iv1 * S_v, delta, S_v * sizeof(float)); + memcpy(r + H * S_v + iv1 * S_v, k_d, S_v * sizeof(float)); + r[2 * H * S_v + iv1] = g_val; + } + } + } +} + +void ggml_compute_forward_gated_delta_net_replay( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_gated_delta_net_replay_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + // ggml_compute_forward_dsv4_hc_comb diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 4c1642a67603..ed645ad7a852 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -105,6 +105,7 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_gated_delta_net_replay(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-sycl/gated_delta_net.cpp b/ggml/src/ggml-sycl/gated_delta_net.cpp index 740a3a1efded..c5e0c7a7eefd 100644 --- a/ggml/src/ggml-sycl/gated_delta_net.cpp +++ b/ggml/src/ggml-sycl/gated_delta_net.cpp @@ -167,19 +167,33 @@ void gated_delta_net_sycl(const float * q, } } +// ARC-LAB replay (Feature A; independent implementation of the technique described by llamADALL): where a +// GATED_DELTA_NET_REPLAY pass puts the state after token t_base (the base) and the update records of the tokens after +// it. Record layout (R = H*(2*S + 1) floats): [delta S*H | k S*H | g H], see ggml_gated_delta_net_replay. +struct gdn_emit { + float * base = nullptr; // [S, S, H] per seq; nullptr = not written + float * rec = nullptr; // record 0 of seq s at rec + (rec_rows ? rec_rows[s] : s) * rec_row_stride + const int64_t * rec_rows = nullptr; // fused cache write: the SET_ROWS row (bank) per seq + int64_t rec_row_stride = 0; // floats + int64_t t_base = 0; // -1: base = the replayed input state (replay-only launch) + int64_t R = 0; +}; + // ARC-LAB token-blocked sequential kernel (non-KDA, S_v = 128). Same per-token recurrence and per-column // arithmetic/reduction order as gated_delta_net_sycl, but the sequential kernel stalls on global-memory latency // every token (~0.4 us/step on the B580 even with the GPU mostly idle) and needs ~5 rounds of threads at 48 heads. // Here each subgroup owns NC columns and shares its k/q registers across them (CUDA's cols_per_warp), and the // work-group stages T tokens of k, q, v, exp(g), beta in SLM; the next block's global loads are issued into // registers before the current block is computed, so their latency overlaps T steps of compute. -template +// emit_t (ARC-LAB replay): also write the base state and the update records described by `em`; the per-token +// arithmetic is untouched, the emission only stores values the step already computed. +template static void gdn_blocked_sycl(const float * q, const float * k, const float * v, const float * g, const float * beta, const float * curr_state, float * dst, float * state, int64_t H, int64_t n_tokens, int64_t n_seqs, int64_t sq1, int64_t sq2, int64_t sq3, int64_t sv1, int64_t sv2, int64_t sv3, int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, float scale, int64_t state_slot_stride, int K, const int32_t * s_ids, - int64_t s_row_stride, dpct::queue_ptr stream) { + int64_t s_row_stride, dpct::queue_ptr stream, const gdn_emit em = gdn_emit()) { constexpr int S = 128, L = 16, RPL = S / L; constexpr int NCW = NC * NW, WGS = NW * L, ncb = S / NCW; constexpr int PK = T * S / WGS, PV = T * NCW / WGS; // per-thread prefetch registers @@ -272,6 +286,7 @@ static void gdn_blocked_sycl(const float * q, const float * k, const float * v, const float g_val = sgb[t]; const float beta_val = sgb[T + t]; float out = 0.0f; + [[maybe_unused]] float dl[NC]; // emit_t: this token's delta per column #pragma unroll for (int c = 0; c < NC; ++c) { float kv_shard = 0.0f; @@ -281,6 +296,9 @@ static void gdn_blocked_sycl(const float * q, const float * k, const float * v, } const float kv_col = warp_reduce_sum(kv_shard); const float delta_col = (sv[t * NCW + w * NC + c] - g_val * kv_col) * beta_val; + if constexpr (emit_t) { + dl[c] = delta_col; + } float attn_partial = 0.0f; #pragma unroll for (int r = 0; r < RPL; ++r) { @@ -295,6 +313,37 @@ static void gdn_blocked_sycl(const float * q, const float * k, const float * v, if (lane < NC) { attn[(t0 + t) * S * H + col0 + lane] = out * scale; } + if constexpr (emit_t) { + const int64_t tt = t0 + t; + if (tt == em.t_base) { + float * bs = em.base + (seq * H + h) * S * S; +#pragma unroll + for (int c = 0; c < NC; ++c) { +#pragma unroll + for (int r = 0; r < RPL; ++r) { + bs[(col0 + c) * S + r * L + lane] = s[c][r]; + } + } + } else if (tt > em.t_base) { + float * rr = em.rec + (em.rec_rows ? em.rec_rows[seq] : seq) * em.rec_row_stride + + (tt - em.t_base - 1) * em.R; +#pragma unroll + for (int c = 0; c < NC; ++c) { + if (lane == c) { + rr[h * S + col0 + c] = dl[c]; + } + } + if (colb == 0 && w == 0) { +#pragma unroll + for (int r = 0; r < RPL; ++r) { + rr[H * S + h * S + r * L + lane] = kr[r]; + } + if (lane == 0) { + rr[2 * H * S + h] = g_val; + } + } + } + } if constexpr (keep_rs_t) { const int target_slot = (int) (n_tokens - 1 - (t0 + t)); if (target_slot >= 0 && target_slot < K) { @@ -324,15 +373,16 @@ static void gdn_blocked_sycl(const float * q, const float * k, const float * v, }); } -template +template static void launch_gdn_blocked(int variant, const float * q, const float * k, const float * v, const float * g, const float * beta, const float * s0, float * dst, float * state, int64_t H, int64_t n_tokens, int64_t n_seqs, int64_t sq1, int64_t sq2, int64_t sq3, int64_t sv1, int64_t sv2, int64_t sv3, int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, float scale, int64_t state_slot_stride, int K, const int32_t * s_ids, - int64_t s_row_stride, dpct::queue_ptr stream) { -#define GDN_BLK(NC, NW, T) gdn_blocked_sycl(q, k, v, g, beta, s0, dst, state, H, n_tokens, n_seqs, \ - sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, s_ids, s_row_stride, stream) + int64_t s_row_stride, dpct::queue_ptr stream, const gdn_emit & em = gdn_emit()) { +#define GDN_BLK(NC, NW, T) gdn_blocked_sycl(q, k, v, g, beta, s0, dst, state, H, n_tokens, \ + n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, s_ids, s_row_stride, \ + stream, em) if (variant < 0) { // auto: widest NC that still leaves ~768 subgroups (~60% of the B580's 1280 thread slots) to hide latency. // B580, 1024 tokens (old kernel -> this): 48 heads 3383 -> 1688 us (NC 8, T 8); 32 heads 2762 -> 1024 (NC 4); @@ -803,3 +853,294 @@ void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, g scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6); ggml_sycl_op_gated_delta_net_impl(ctx, dst, &cache); } + +// ARC-LAB: GATED_DELTA_NET_REPLAY (Feature A: DeltaNet rollback by replay instead of snapshots; independent +// implementation of the technique described by llamADALL). See ggml_gated_delta_net_replay for the data layout. +// +// Sequential kernel: gated_delta_net_sycl's scalar-gate path (same grid, same per-column arithmetic and reduction +// order, lines kept verbatim) with +// - a replay prologue that re-applies n_rec stored records to the loaded state, using the step's own update +// expression (g * S + k * delta) on the stored g, k and delta, so the state matches what the step produced; +// - an epilogue that stores the base state (after token t_base) and the records of the tokens after it. +// With n_tokens == 0 and t_base == -1 it only replays and writes the result to em.base (used before the blocked +// kernel, which has no replay prologue). +template +static void gdn_replay_seq_sycl(const float * q, const float * k, const float * v, const float * g, const float * beta, + const float * curr_state, float * dst, int64_t H, int64_t n_tokens, int64_t sq1, + int64_t sq2, int64_t sq3, int64_t sv1, int64_t sv2, int64_t sv3, int64_t sb1, + int64_t sb2, int64_t sb3, const sycl::uint3 neqk1_magic, const sycl::uint3 rq3_magic, + float scale, const int32_t * s_ids, int64_t s_row_stride, const float * rec_in, + int64_t rec_in_seq_stride, int64_t rec_in_bank_stride, const int32_t * ctrl, + const gdn_emit em) { + auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); + const uint32_t h_idx = item_ct1.get_group(2); + const uint32_t sequence = item_ct1.get_group(1); + const int lane = item_ct1.get_local_id(2); + const int col = item_ct1.get_group(0) * item_ct1.get_local_range(1) + item_ct1.get_local_id(1); + + const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic); + const uint32_t iq3 = fastdiv(sequence, rq3_magic); + + float * attn_data = dst; + + const int64_t state_in_offset = (s_ids ? (int64_t) s_ids[sequence] * s_row_stride : (int64_t) sequence * H * S_v * S_v) + + h_idx * S_v * S_v; + curr_state += state_in_offset + col * S_v; + attn_data += (sequence * n_tokens * H + h_idx) * S_v; + + constexpr int warp_size = ggml_sycl_get_physical_warp_size() < S_v ? ggml_sycl_get_physical_warp_size() : S_v; + static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size"); + constexpr int rows_per_lane = (S_v + warp_size - 1) / warp_size; + float s_shard[rows_per_lane]; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + s_shard[r] = curr_state[i]; + } + + // replay prologue + if (rec_in != nullptr) { + const int32_t n_rec = ctrl[2 * sequence + 0]; + const int32_t bank = ctrl[2 * sequence + 1]; + const float * rb = rec_in + sequence * rec_in_seq_stride + bank * rec_in_bank_stride; + for (int32_t j = 0; j < n_rec; j++) { + const float * rj = rb + j * em.R; + const float g_val = rj[2 * H * S_v + h_idx]; + const float delta_col = rj[h_idx * S_v + col]; + const float * k_t = rj + H * S_v + h_idx * S_v; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + s_shard[r] = g_val * s_shard[r] + k_t[i] * delta_col; + } + } + } + + float * base_out = em.base ? em.base + (sequence * H + h_idx) * S_v * S_v + col * S_v : nullptr; + float * rec_seq = em.rec ? em.rec + (em.rec_rows ? em.rec_rows[sequence] : (int64_t) sequence) * em.rec_row_stride + : nullptr; + + if (em.t_base < 0 && base_out) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + base_out[r * warp_size + lane] = s_shard[r]; + } + } + + for (int t = 0; t < n_tokens; t++) { + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; + + const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; + const float * beta_t = beta + gb_offset; + const float * g_t = g + gb_offset; + + const float beta_val = *beta_t; + + const float g_val = sycl::native::exp(*g_t); + + // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + kv_shard += s_shard[r] * k_t[i]; + } + float kv_col = warp_reduce_sum(kv_shard); + + // delta[col] = (v[col] - g * kv[col]) * beta + float delta_col = (v_t[col] - g_val * kv_col) * beta_val; + + // fused: S[i][col] = g * S[i][col] + k[i] * delta[col] + // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] + float attn_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + s_shard[r] = g_val * s_shard[r] + k_t[i] * delta_col; + attn_partial += s_shard[r] * q_t[i]; + } + + float attn_col = warp_reduce_sum(attn_partial); + + if (lane == 0) { + attn_data[col] = attn_col * scale; + } + + attn_data += S_v * H; + + // epilogue: base state, then one record per later token + if (t == em.t_base) { + if (base_out) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + base_out[r * warp_size + lane] = s_shard[r]; + } + } + } else if (t > em.t_base) { + float * rt = rec_seq + (t - em.t_base - 1) * em.R; + if (lane == 0) { + rt[h_idx * S_v + col] = delta_col; + } + if (col == 0) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + rt[H * S_v + h_idx * S_v + i] = k_t[i]; + } + if (lane == 0) { + rt[2 * H * S_v + h_idx] = g_val; + } + } + } + } +} + +static void launch_gdn_replay_seq(const float * q_d, const float * k_d, const float * v_d, const float * g_d, + const float * b_d, const float * s_d, float * dst_d, int64_t S_v, int64_t H, + int64_t n_tokens, int64_t n_seqs, int64_t sq1, int64_t sq2, int64_t sq3, int64_t sv1, + int64_t sv2, int64_t sv3, int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, + int64_t rq3, float scale, const int32_t * s_ids, int64_t s_row_stride, + const float * rec_in, int64_t rec_in_seq_stride, int64_t rec_in_bank_stride, + const int32_t * ctrl, const gdn_emit & em, dpct::queue_ptr stream) { + GGML_ASSERT(S_v == 128); // supports_op + const int warp_size = ggml_sycl_info().devices[ggml_sycl_get_device()].warp_size; + + const int num_warps = 4; + dpct::dim3 grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps); + dpct::dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1); + + const sycl::uint3 neqk1_magic = init_fastdiv_values(neqk1); + const sycl::uint3 rq3_magic = init_fastdiv_values(rq3); + + constexpr int sv = 128; + const gdn_emit e = em; + stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims), + [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + gdn_replay_seq_sycl(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, sq1, sq2, sq3, + sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, s_ids, + s_row_stride, rec_in, rec_in_seq_stride, rec_in_bank_stride, ctrl, + e); + }); +} + +void ggml_sycl_op_gated_delta_net_replay(ggml_backend_sycl_context & ctx, ggml_tensor * dst, + const ggml_sycl_gdn_replay_fused * fused) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/8); + + const ggml_tensor * src_q = dst->src[0]; + const ggml_tensor * src_k = dst->src[1]; + const ggml_tensor * src_v = dst->src[2]; + const ggml_tensor * src_g = dst->src[3]; + const ggml_tensor * src_beta = dst->src[4]; + const ggml_tensor * src_s = dst->src[5]; + const ggml_tensor * src_rec = dst->src[6]; + const ggml_tensor * src_ctrl = dst->src[7]; + + GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne); + GGML_TENSOR_LOCALS(size_t , nbq, src_q, nb); + GGML_TENSOR_LOCALS(int64_t, nev, src_v, ne); + GGML_TENSOR_LOCALS(size_t, nbv, src_v, nb); + GGML_TENSOR_LOCALS(size_t, nbb, src_beta, nb); + + const int64_t S_v = nev0; + const int64_t H = nev1; + const int64_t n_tokens = nev2; + const int64_t n_seqs = nev3; + const int64_t K = ggml_get_op_params_i32(dst, 0); + const int64_t R = H * (2 * S_v + 1); + const int64_t D = S_v * S_v * H; + const int64_t m = std::min(n_tokens - 1, K - 1); + + GGML_ASSERT(ggml_are_same_stride(src_q, src_k)); + GGML_ASSERT(ggml_is_contiguous(src_g) && ggml_is_contiguous(src_beta) && ggml_is_contiguous(src_s)); + + const int64_t neqk1 = neq1; + const int64_t rq3 = nev3 / neq3; + + const float * q_d = (const float *) src_q->data; + const float * k_d = (const float *) src_k->data; + const float * v_d = (const float *) src_v->data; + const float * g_d = (const float *) src_g->data; + const float * b_d = (const float *) src_beta->data; + float * dst_d = (float *) dst->data; + + // fused recurrent-state gather (GET_ROWS skipped, see ggml_sycl_try_gdn_gather_skip) + const float * s_d = (const float *) src_s->data; + const int32_t * s_ids = nullptr; + int64_t s_row_stride = 0; + { + auto & mg = g_gdn_gathers[&ctx]; + auto it = mg.find(dst); + if (it != mg.end()) { + s_d = it->second.base; + s_ids = it->second.ids; + s_row_stride = it->second.row_stride; + } + } + + const int64_t sq1 = nbq1 / sizeof(float); + const int64_t sq2 = nbq2 / sizeof(float); + const int64_t sq3 = nbq3 / sizeof(float); + const int64_t sv1 = nbv1 / sizeof(float); + const int64_t sv2 = nbv2 / sizeof(float); + const int64_t sv3 = nbv3 / sizeof(float); + const int64_t sb1 = nbb1 / sizeof(float); + const int64_t sb2 = nbb2 / sizeof(float); + const int64_t sb3 = nbb3 / sizeof(float); + + const float scale = 1.0f / sqrtf((float) S_v); + + const float * rec_in = (const float *) src_rec->data; + const int64_t rec_in_bank_stride = src_rec->nb[1] / sizeof(float); + const int64_t rec_in_seq_stride = src_rec->nb[2] / sizeof(float); + const int32_t * ctrl = (const int32_t *) src_ctrl->data; + + gdn_emit em; + em.R = R; + em.t_base = n_tokens - 1 - m; + if (fused) { + em.base = fused->base; + em.rec = fused->rec; + em.rec_rows = fused->rec_rows; + em.rec_row_stride = fused->rec_row_stride; + } else { + em.base = dst_d + S_v * H * n_tokens * n_seqs; + em.rec = em.base + D * n_seqs; + em.rec_row_stride = (K - 1) * R; + } + + static bool logged = false; + if (!logged) { + logged = true; + GGML_LOG_INFO("%s: ARC-LAB GDN rollback by replay active (cache write %s)\n", __func__, + fused ? "fused" : "via copy"); + } + + dpct::queue_ptr stream = ctx.stream(); + + // same kernel choice as the snapshot path (launch_gated_delta_net, scalar gate, S_v == 128) + static const bool blk_off = getenv("GGML_SYCL_GDN_BLOCKED_OFF") != nullptr; + static const int blk_variant = getenv("GGML_SYCL_GDN_BLK_VARIANT") ? atoi(getenv("GGML_SYCL_GDN_BLK_VARIANT")) : -1; + static const int blk_min = getenv("GGML_SYCL_GDN_BLK_MIN") ? atoi(getenv("GGML_SYCL_GDN_BLK_MIN")) : 16; + if (!blk_off && n_tokens >= blk_min) { + // the blocked kernel has no replay prologue: replay into a scratch state first, then run it from there + ggml_sycl_pool_alloc s_rpl(ctx.pool(), (size_t) D * n_seqs); + gdn_emit em0; + em0.base = s_rpl.get(); + em0.t_base = -1; + em0.R = R; + launch_gdn_replay_seq(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, S_v, H, 0, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, + sb1, sb2, sb3, neqk1, rq3, scale, s_ids, s_row_stride, rec_in, rec_in_seq_stride, + rec_in_bank_stride, ctrl, em0, stream); + // K = 0 with keep_rs_t: no snapshot slots and no final-state store, only the emission + launch_gdn_blocked(blk_variant, q_d, k_d, v_d, g_d, b_d, s_rpl.get(), dst_d, em.base, H, n_tokens, + n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, 0, 0, + nullptr, 0, stream, em); + return; + } + launch_gdn_replay_seq(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, + sb1, sb2, sb3, neqk1, rq3, scale, s_ids, s_row_stride, rec_in, rec_in_seq_stride, + rec_in_bank_stride, ctrl, em, stream); +} diff --git a/ggml/src/ggml-sycl/gated_delta_net.hpp b/ggml/src/ggml-sycl/gated_delta_net.hpp index da5dda3f4969..c086986195da 100644 --- a/ggml/src/ggml-sycl/gated_delta_net.hpp +++ b/ggml/src/ggml-sycl/gated_delta_net.hpp @@ -22,3 +22,16 @@ void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, g void ggml_sycl_gdn_gather_reset(const ggml_backend_sycl_context & ctx); void ggml_sycl_gdn_gather_set(const ggml_backend_sycl_context & ctx, const ggml_tensor * gdn, const float * base, const int32_t * ids, int64_t row_stride); + +// ARC-LAB: GATED_DELTA_NET_REPLAY (rollback by replay). fused = nullptr: base state and records go to dst (the graph +// copies them into the cache); otherwise the kernel writes them straight into the cache (see +// ggml_sycl_try_gdn_replay_fusion): base -> the CPY destination, records -> row rec_rows[seq] of the SET_ROWS target. +struct ggml_sycl_gdn_replay_fused { + float * base; + float * rec; + const int64_t * rec_rows; // device I64, one per seq + int64_t rec_row_stride; // floats +}; + +void ggml_sycl_op_gated_delta_net_replay(ggml_backend_sycl_context & ctx, ggml_tensor * dst, + const ggml_sycl_gdn_replay_fused * fused); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 354b09854d40..f7037bf36087 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6024,6 +6024,9 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg case GGML_OP_GATED_DELTA_NET: ggml_sycl_gated_delta_net(ctx, dst); break; + case GGML_OP_GATED_DELTA_NET_REPLAY: + ggml_sycl_op_gated_delta_net_replay(ctx, dst, nullptr); + break; case GGML_OP_OPT_STEP_ADAMW: ggml_sycl_opt_step_adamw(ctx, dst); break; @@ -6359,6 +6362,83 @@ static int ggml_sycl_try_gdn_cache_fusion(const ggml_cgraph * cgraph, int node_i return skip; } +// ARC-LAB: match GATED_DELTA_NET_REPLAY + the two nodes that store its output in the recurrent cache - a CPY of the +// base state and a SET_ROWS of the update records into the bank picked by its row index - so the kernel writes both +// straight into the cache and the two copies are skipped. Single sequence only. Returns the nodes to skip (0 = none). +static int ggml_sycl_try_gdn_replay_fusion(const ggml_cgraph * cgraph, int node_idx, ggml_sycl_gdn_replay_fused & fused) { + if (!g_ggml_sycl_enable_fusion) { + return 0; + } + const ggml_tensor * gdn = cgraph->nodes[node_idx]; + if (gdn->op != GGML_OP_GATED_DELTA_NET_REPLAY || gdn->type != GGML_TYPE_F32 || (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return 0; + } + const ggml_tensor * src_v = gdn->src[2]; + const int64_t S_v = src_v->ne[0]; + const int64_t H = src_v->ne[1]; + const int64_t n_tokens = src_v->ne[2]; + const int64_t n_seqs = src_v->ne[3]; + const int64_t K = ggml_get_op_params_i32(gdn, 0); + const int64_t D = S_v * S_v * H; + const int64_t R = H * (2 * S_v + 1); + if (n_seqs != 1) { + return 0; + } + const size_t base_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs); + const size_t rec_off = base_off + ggml_row_size(GGML_TYPE_F32, D * n_seqs); + + // the next two computed nodes must be the CPY and the SET_ROWS (either order), nothing in between + const ggml_tensor * cpy = nullptr; + const ggml_tensor * sr = nullptr; + int skip = 0; + for (int j = node_idx + 1; j < cgraph->n_nodes && (cpy == nullptr || sr == nullptr); ++j) { + const ggml_tensor * n = cgraph->nodes[j]; + if (ggml_sycl_is_view_or_noop(n) || (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + continue; + } + if ((n->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return 0; + } + if (n->op == GGML_OP_CPY && cpy == nullptr) { + cpy = n; + } else if (n->op == GGML_OP_SET_ROWS && sr == nullptr) { + sr = n; + } else { + return 0; + } + skip = j - node_idx; + } + if (cpy == nullptr || sr == nullptr) { + return 0; + } + + // CPY: base region of this op's output -> a contiguous F32 cache view of D * n_seqs floats + const ggml_tensor * cs = cpy->src[0]; + const ggml_tensor * cd = cpy->src[1]; + if (cs->op != GGML_OP_VIEW || cs->view_src != gdn || cs->view_offs != base_off || !ggml_is_contiguous(cs) || + ggml_nelements(cs) != D * n_seqs || cd->type != GGML_TYPE_F32 || cd->data == nullptr || + !ggml_is_contiguous(cd)) { + return 0; + } + + // SET_ROWS: one row of R * (K - 1) records from the record region -> F32 rows (banks), I64 index on the device + const ggml_tensor * rs = sr->src[0]; + const ggml_tensor * ri = sr->src[1]; + if (rs->op != GGML_OP_VIEW || rs->view_src != gdn || rs->view_offs != rec_off || !ggml_is_contiguous(rs) || + rs->ne[0] != R * (K - 1) || ggml_nrows(rs) != 1 || sr->type != GGML_TYPE_F32 || sr->data == nullptr || + sr->ne[0] != R * (K - 1) || sr->nb[0] != sizeof(float) || sr->nb[1] % sizeof(float) != 0 || + ri->type != GGML_TYPE_I64 || ggml_nelements(ri) != 1 || ri->data == nullptr || + (ri->buffer && ggml_backend_buffer_is_host(ri->buffer))) { + return 0; + } + + fused.base = (float *) cd->data; + fused.rec = (float *) sr->data; + fused.rec_rows = (const int64_t *) ri->data; + fused.rec_row_stride = (int64_t) (sr->nb[1] / sizeof(float)); + return skip; +} + // GET_ROWS that gathers a recurrent state row feeding only GATED_DELTA_NET src[5] (possibly through a RESHAPE): // skip it and let the kernel read the cache row directly. Single sequence only. GGML_SYCL_GDN_GATHER=0 disables. static bool ggml_sycl_try_gdn_gather_skip(ggml_backend_sycl_context & ctx, const ggml_cgraph * cgraph, int node_idx) { @@ -6385,7 +6465,7 @@ static bool ggml_sycl_try_gdn_gather_skip(ggml_backend_sycl_context & ctx, const const ggml_tensor * cur = gr; for (int j = node_idx + 1; j < cgraph->n_nodes; ++j) { const ggml_tensor * n = cgraph->nodes[j]; - if (n->op == GGML_OP_GATED_DELTA_NET && n->src[5] == cur) { + if ((n->op == GGML_OP_GATED_DELTA_NET || n->op == GGML_OP_GATED_DELTA_NET_REPLAY) && n->src[5] == cur) { const ggml_tensor * v = n->src[2]; const int64_t D = v->ne[0] * v->ne[0] * v->ne[1]; if (gr->ne[0] != D || v->ne[3] != 1 || ggml_nelements(cur) != D) { @@ -7220,6 +7300,15 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc continue; } } + if (node->op == GGML_OP_GATED_DELTA_NET_REPLAY) { + ggml_sycl_gdn_replay_fused fused; + const int rpl_nodes_to_skip = ggml_sycl_try_gdn_replay_fusion(cgraph, i, fused); + if (rpl_nodes_to_skip > 0) { + ggml_sycl_op_gated_delta_net_replay(*sycl_ctx, node, &fused); + i += rpl_nodes_to_skip; + continue; + } + } if (node->op == GGML_OP_MUL_MAT && ggml_sycl_try_dn_gates(*sycl_ctx, cgraph, i)) { continue; } @@ -8033,6 +8122,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_GATED_DELTA_NET: // rows-indexed state read (src[6]) and raw gates (op_params[1]) not implemented here return op->src[6] == NULL && ggml_get_op_params_i32(op, 1) == 0; + case GGML_OP_GATED_DELTA_NET_REPLAY: + // ARC-LAB: kernels exist for the scalar gate at head size 128 (Bonsai / Qwen3.5) only + return op->src[2]->ne[0] == 128 && op->src[3]->ne[0] == 1 && ggml_are_same_stride(op->src[0], op->src[1]); case GGML_OP_SSM_CONV: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index fa2a9c44f5cf..762655731dbb 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1114,9 +1114,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "OPT_STEP_SGD", "GLU", + + "GATED_DELTA_NET_REPLAY", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1229,9 +1231,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "sgd(x)", "glu(x)", + + "gated_delta_net_replay(q, k, v, g, beta, s, rec, ctrl)", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6441,6 +6445,65 @@ void ggml_gated_delta_net_set_raw_gates( gdn->src[8] = a; } +// ARC-LAB: ggml_gated_delta_net_replay (see ggml.h) +struct ggml_tensor * ggml_gated_delta_net_replay( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * rec, + struct ggml_tensor * ctrl, + int K) { + GGML_ASSERT(ggml_is_contiguous_rows(q)); + GGML_ASSERT(ggml_is_contiguous_rows(k)); + GGML_ASSERT(ggml_is_contiguous_rows(v)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(beta)); + GGML_ASSERT(ggml_is_contiguous(state)); + GGML_ASSERT(ggml_is_contiguous(ctrl)); + + GGML_ASSERT(q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F32 && v->type == GGML_TYPE_F32); + GGML_ASSERT(g->type == GGML_TYPE_F32 && beta->type == GGML_TYPE_F32 && state->type == GGML_TYPE_F32); + GGML_ASSERT(rec->type == GGML_TYPE_F32); + GGML_ASSERT(ctrl->type == GGML_TYPE_I32); + + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + + GGML_ASSERT(K >= 2 && n_tokens >= 1); + GGML_ASSERT(q->ne[0] == S_v && k->ne[0] == S_v && q->ne[1] == k->ne[1] && H % q->ne[1] == 0); + GGML_ASSERT(g->ne[0] == 1 && g->ne[1] == H && g->ne[2] == n_tokens && g->ne[3] == n_seqs); // scalar gate only + GGML_ASSERT(beta->ne[0] == 1 && beta->ne[1] == H && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs); + GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H && state->ne[3] == n_seqs); + + const int64_t R = H * (2 * S_v + 1); + GGML_ASSERT(rec->ne[0] == R * (K - 1) && rec->ne[1] == 2 && rec->ne[2] == n_seqs); + GGML_ASSERT(rec->nb[0] == sizeof(float) && rec->nb[1] == (size_t) rec->ne[0] * sizeof(float)); + GGML_ASSERT(ggml_nelements(ctrl) == 2 * n_seqs); + + const int64_t ne[4] = { S_v * H * n_tokens * n_seqs + S_v * S_v * H * n_seqs + (K - 1) * R * n_seqs, 1, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, (int32_t) K); + + result->op = GGML_OP_GATED_DELTA_NET_REPLAY; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = g; + result->src[4] = beta; + result->src[5] = state; + result->src[6] = rec; + result->src[7] = ctrl; + + return result; +} + // ggml_lightning_indexer struct ggml_tensor * ggml_lightning_indexer( diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index ceabee25866a..44a31198b323 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -385,6 +385,27 @@ void llm_graph_input_rs::set_input_rs(const llama_memory_recurrent_context * mct } } + if (s_rpl_ctrl) { + GGML_ASSERT(ggml_backend_buffer_is_host(s_rpl_rows->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(s_rpl_ctrl->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(s_rpl_wrow->buffer)); + int32_t * rows = (int32_t *) s_rpl_rows->data; + int32_t * ctrl = (int32_t *) s_rpl_ctrl->data; + int64_t * wrow = (int64_t *) s_rpl_wrow->data; + + for (uint32_t i = 0; i < n_rs; ++i) { + rows[i] = mctx_cur->rpl_row(i); + } + for (uint32_t i = 0; i < ubatch->n_seqs; ++i) { + int32_t n_rec = 0; + int32_t bank = 0; + mctx_cur->rpl_take(i, ubatch->n_seq_tokens, n_rec, bank); + ctrl[2*i + 0] = n_rec; + ctrl[2*i + 1] = bank; + wrow[i] = bank ^ 1; + } + } + if (s_write_rows) { GGML_ASSERT(ggml_backend_buffer_is_host(s_write_rows->buffer)); int64_t * data = (int64_t *) s_write_rows->data; @@ -417,6 +438,12 @@ bool llm_graph_input_rs::can_reuse_rs(const llama_memory_recurrent_context * mct res &= head == mctx_cur->get_head(); res &= rs_z == mctx_cur->get_rs_z(); + res &= (s_rpl_ctrl != nullptr) == mctx_cur->get_rpl(); + if (s_rpl_ctrl) { + res &= s_rpl_rows->ne[0] == mctx_cur->get_n_rs(); + res &= s_rpl_ctrl->ne[0] == 2 * params.ubatch.n_seqs; + } + return res; } @@ -3578,6 +3605,18 @@ static std::unique_ptr build_rs_inp_impl( inp->s_copy_main = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0); inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]); + if (mctx_cur->get_rpl()) { + inp->s_rpl_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs); + ggml_set_input(inp->s_rpl_rows); + inp->s_rpl_main = ggml_view_1d(ctx0, inp->s_rpl_rows, n_seqs, 0); + inp->s_rpl_extra = ggml_view_1d(ctx0, inp->s_rpl_rows, n_rs - n_seqs, n_seqs * inp->s_rpl_rows->nb[0]); + + inp->s_rpl_ctrl = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 2 * n_seqs); + ggml_set_input(inp->s_rpl_ctrl); + inp->s_rpl_wrow = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_seqs); + ggml_set_input(inp->s_rpl_wrow); + } + inp->head = mctx_cur->get_head(); inp->rs_z = mctx_cur->get_rs_z(); diff --git a/src/llama-graph.h b/src/llama-graph.h index ea9952557a47..0a5b378ddc58 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -346,6 +346,13 @@ class llm_graph_input_rs : public llm_graph_input_i { ggml_tensor * s_write_rows = nullptr; int64_t s_write_K = 0; + // ARC-LAB GDN rollback by replay (llama_memory_recurrent::rpl), null when off: + ggml_tensor * s_rpl_rows = nullptr; // I32 [n_rs]: plane-0 S row each state is read from (no rollback plane) + ggml_tensor * s_rpl_main = nullptr; // view [n_seqs] + ggml_tensor * s_rpl_extra = nullptr; // view [n_rs - n_seqs] + ggml_tensor * s_rpl_ctrl = nullptr; // I32 [2, n_seqs]: {records to re-apply, their bank} + ggml_tensor * s_rpl_wrow = nullptr; // I64 [n_seqs]: bank the new records go to + const llama_memory_recurrent_context * mctx; // used in view offsets, need to match for valid graph reuse diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..f5d33f30397b 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -34,6 +35,8 @@ llama_memory_recurrent::llama_memory_recurrent( this->n_rs_seq = n_rs_seq; rs_idx.assign(n_seq_max, 0); + rpl_n.assign(n_seq_max, 0); + rpl_bank.assign(n_seq_max, 0); cells.clear(); cells.resize(mem_size); @@ -105,6 +108,39 @@ llama_memory_recurrent::llama_memory_recurrent( s_l[i] = s; } + // ARC-LAB: GDN rollback by replay (LLAMA_GDN_REPLAY=0 keeps the snapshot planes). Applies to a single sequence + // of a qwen35 model whose recurrent layers all run on SYCL; anything else silently keeps the snapshot path. + // The tensor layout is the same either way (planes 0 and 1 of S are reused), only their meaning changes. + { + const char * env = getenv("LLAMA_GDN_REPLAY"); + bool on = !(env && atoi(env) == 0) && n_rs_seq > 0 && n_seq_max == 1 && mem_size == 1 && offload && + model.arch == LLM_ARCH_QWEN35 && type_s == GGML_TYPE_F32 && + getenv("GGML_SYCL_GDN_CHUNKED") == nullptr; // the opt-in chunked prefill has no replay variant + const uint32_t S = hparams.ssm_d_state; + const uint32_t H = hparams.ssm_dt_rank; + on = on && S == 128 && H > 0 && hparams.ssm_d_inner == S * H && hparams.n_embd_s() == S * S * H; + const uint32_t R = H * (2 * S + 1); + // both banks of n_rs_seq records must fit in one S row + on = on && (uint64_t) 2 * n_rs_seq * R <= (uint64_t) hparams.n_embd_s(); + bool any = false; + for (int i = 0; on && i < n_layer; i++) { + if (s_l[i] == nullptr) { + continue; + } + any = true; + ggml_backend_dev_t dev = model.dev_layer(i); + ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; + on = reg && strcmp(ggml_backend_reg_name(reg), "SYCL") == 0; + } + rpl = on && any; + rpl_R = rpl ? R : 0; + if (rpl) { + LLAMA_LOG_INFO("%s: GDN rollback by replay: on (%u records of %u floats per bank)\n", __func__, n_rs_seq, R); + } else if (n_rs_seq > 0) { + LLAMA_LOG_INFO("%s: GDN rollback by replay: off (snapshot planes)\n", __func__); + } + } + // allocate tensors and initialize the buffers to avoid NaNs in the padding for (auto & [buft, ctx] : ctx_map) { ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); @@ -145,6 +181,7 @@ void llama_memory_recurrent::clear(bool data) { } std::fill(rs_idx.begin(), rs_idx.end(), 0); + rpl_reset(-1); } bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { @@ -166,6 +203,7 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos const bool rm_all = p0 == 0 && p1 == std::numeric_limits::max(); if (rm_all) { set_rs_idx(seq_id, 0); + rpl_reset(seq_id); } // models like Mamba or RWKV can't have a state partially erased at the end @@ -185,6 +223,14 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos // pending rollback is single-use const bool pending = rs_idx[seq_id] != 0; if (!pending && rollback >= 1 && rollback <= (llama_pos) n_rs_seq) { + if (rpl) { + // ARC-LAB replay: only the tokens that left a record can be rolled back + if ((uint32_t) rollback > rpl_n[seq_id]) { + LLAMA_LOG_WARN("%s: cannot roll back %d tokens, only %u kept\n", __func__, rollback, rpl_n[seq_id]); + return false; + } + rpl_n[seq_id] -= (uint32_t) rollback; + } set_rs_idx(seq_id, (uint32_t) rollback); cell.pos = p0 - 1; return true; @@ -194,6 +240,7 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos // invalidate tails which will be cleared if (p0 <= cell.pos && cell.pos < p1) { tail_id = -1; + rpl_reset(seq_id); } } } else { @@ -406,6 +453,18 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) { rs_idx[seq_id] = idx; } +void llama_memory_recurrent::rpl_reset(llama_seq_id seq_id) { + if (seq_id < 0) { + std::fill(rpl_n.begin(), rpl_n.end(), 0); + std::fill(rpl_bank.begin(), rpl_bank.end(), 0); + return; + } + if ((size_t) seq_id < rpl_n.size()) { + rpl_n[seq_id] = 0; + rpl_bank[seq_id] = 0; + } +} + std::map llama_memory_recurrent::memory_breakdown() const { std::map ret; for (const auto & [_, buf] : ctxs_bufs) { @@ -745,7 +804,10 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq std::vector> cell_ranges; // ranges, from inclusive, to exclusive std::vector> cell_ranges_data; // logical source row ranges + std::vector> cell_ranges_data_s; // ARC-LAB replay: S rows (base, plane 0) uint32_t cell_count = 0; + uint32_t rpl_pending = 0; // ARC-LAB replay: records still to apply to the written base state + uint32_t rpl_bank_cur = 0; // Count the number of cells with the specified seq_id // Find all the ranges of cells with this seq id (or all, when -1) @@ -784,6 +846,23 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq cell_ranges_data.back().second++; } + if (rpl) { + // the base state lives in plane 0 whatever rs_idx says (rs_idx only picks the conv snapshot) + const uint32_t cell_id_s = (cell.src >= 0 ? cell.src : (int32_t) i); + if (cell_ranges_data_s.empty() || cell_ranges_data_s.back().second != cell_id_s) { + cell_ranges_data_s.emplace_back(cell_id_s, cell_id_s + 1); + } else { + cell_ranges_data_s.back().second++; + } + // single sequence (rpl requires n_seq_max == 1): the pending records of the cell's sequence + for (const llama_seq_id cell_seq_id : cell.seq_id) { + if (cell_seq_id >= 0 && (size_t) cell_seq_id < rpl_n.size()) { + rpl_pending = rpl_n[cell_seq_id]; + rpl_bank_cur = rpl_bank[cell_seq_id]; + } + } + } + if (cell_range_begin == size) { cell_range_begin = i; } @@ -818,7 +897,12 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq io.write(&cell_count, sizeof(cell_count)); state_write_meta(io, cell_ranges, seq_id); - state_write_data(io, cell_ranges_data); + if (!rpl) { + state_write_data(io, cell_ranges_data, cell_ranges_data, 0, 0); + } else { + GGML_ASSERT(cell_count <= 1); // single sequence + state_write_data(io, cell_ranges_data, cell_ranges_data_s, rpl_pending, rpl_bank_cur); + } } void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { @@ -871,8 +955,11 @@ void llama_memory_recurrent::state_write_meta(llama_io_write_i & io, const std:: } } -void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges) const { - const uint32_t s_trans = 0; +void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges, + const std::vector> & cell_ranges_s, uint32_t rpl_pending, uint32_t rpl_bank_cur) const { + // ARC-LAB replay: with records pending the state is base + records; the header then says 2 (a reader without + // replay rejects it as an incompatible S layout instead of loading a stale state), and the records follow S + const uint32_t s_trans = rpl_pending > 0 ? 2 : 0; const uint32_t n_layer = hparams.n_layer(); io.write(&s_trans, sizeof(s_trans)); @@ -901,7 +988,7 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std:: } } - if (!s_trans) { + if (s_trans != 1) { for (uint32_t il = 0; il < n_layer; ++il) { // skip null layers (read_data will handle this by checking "r_l" and "s_l" for null) if (s_l[il] == nullptr) continue; @@ -916,7 +1003,7 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std:: // Write each logical cell row range. With pending recurrent rollback, // the logical current state may live in a rollback snapshot plane. - for (const auto & range : cell_ranges) { + for (const auto & range : cell_ranges_s) { const size_t range_size = range.second - range.first; const size_t buf_size = range_size * s_size_row; io.write_tensor(s_l[il], range.first * s_size_row, buf_size); @@ -953,6 +1040,19 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std:: } } } + + if (rpl_pending > 0) { + // ARC-LAB replay: the pending update records of the (single) written cell, per layer + GGML_ASSERT(cell_ranges_s.size() == 1 && cell_ranges_s[0].second - cell_ranges_s[0].first == 1); + const uint32_t cell = cell_ranges_s[0].first; + io.write(&rpl_pending, sizeof(rpl_pending)); + io.write(&rpl_R, sizeof(rpl_R)); + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_l[il] == nullptr) continue; + const size_t off = ((size_t) (size + cell) * hparams.n_embd_s() + (size_t) rpl_bank_cur * n_rs_seq * rpl_R) * sizeof(float); + io.write_tensor(s_l[il], off, (size_t) rpl_pending * rpl_R * sizeof(float)); + } + } } bool llama_memory_recurrent::state_read_meta(llama_io_read_i & io, uint32_t cell_count, llama_seq_id dest_seq_id) { @@ -1065,7 +1165,8 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell LLAMA_LOG_ERROR("%s: not enough cells in kv cache to restore state (%u > %u)\n", __func__, cell_count, size); return false; } - if (false != (bool) s_trans) { + // ARC-LAB replay: 2 = plain S layout followed by pending update records (written only when replay is on) + if (s_trans != 0 && !(s_trans == 2 && rpl)) { LLAMA_LOG_ERROR("%s: incompatible s transposition\n", __func__); return false; } @@ -1099,7 +1200,7 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell } } - if (!s_trans) { + if (s_trans != 1) { for (uint32_t il = 0; il < n_layer; ++il) { // skip null layers if (s_l[il] == nullptr) continue; @@ -1172,6 +1273,28 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell } } + if (s_trans == 2) { + // ARC-LAB replay: pending update records -> bank 0 of the restored cell + uint32_t n_rec = 0; + uint32_t R_ref = 0; + io.read(&n_rec, sizeof(n_rec)); + io.read(&R_ref, sizeof(R_ref)); + if (cell_count != 1 || R_ref != rpl_R || n_rec > n_rs_seq) { + LLAMA_LOG_ERROR("%s: mismatched replay records (%u cells, %u records of %u floats)\n", __func__, cell_count, n_rec, R_ref); + return false; + } + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_l[il] == nullptr) continue; + io.read_tensor(s_l[il], (size_t) (size + head) * hparams.n_embd_s() * sizeof(float), (size_t) n_rec * rpl_R * sizeof(float)); + } + for (const llama_seq_id seq : cells[head].seq_id) { + if (seq >= 0 && (size_t) seq < rpl_n.size()) { + rpl_n[seq] = n_rec; + rpl_bank[seq] = 0; + } + } + } + return true; } @@ -1270,3 +1393,31 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const { } return (int32_t)(idx * mem->size) + src0; } + +bool llama_memory_recurrent_context::get_rpl() const { + return mem->rpl; +} + +int32_t llama_memory_recurrent_context::rpl_row(int i) const { + return mem->cells[i + mem->head].src0; +} + +void llama_memory_recurrent_context::rpl_take(int i, uint32_t n_seq_tokens, int32_t & n_rec, int32_t & bank) const { + const uint32_t cell_idx = i + mem->head; + + n_rec = 0; + bank = 0; + if (mem->cells[cell_idx].seq_id.empty()) { + return; + } + const llama_seq_id seq = *mem->cells[cell_idx].seq_id.begin(); + if (seq < 0 || (size_t) seq >= mem->rpl_n.size()) { + return; + } + n_rec = (int32_t) mem->rpl_n[seq]; + bank = (int32_t) mem->rpl_bank[seq]; + + // this ubatch leaves a base state after token n-1-m and records for its last m tokens, in the other bank + mem->rpl_n[seq] = std::min(n_seq_tokens - 1, mem->n_rs_seq); + mem->rpl_bank[seq] = (uint8_t) (bank ^ 1); +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..b44915e1a690 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -78,6 +78,18 @@ class llama_memory_recurrent : public llama_memory_i { void set_rs_idx(llama_seq_id seq_id, uint32_t idx); + // ARC-LAB: rollback by replay for the gated delta net state (LLAMA_GDN_REPLAY, default on where it applies; + // independent implementation of the technique described by llamADALL). The S rows of plane 0 hold the base + // state; plane 1 of each cell holds two banks of per-token update records (ggml_gated_delta_net_replay). + // Rolling back r tokens drops the last r pending records instead of selecting a snapshot plane. The conv + // state (r_l) keeps its snapshot planes and rs_idx. + bool rpl = false; + uint32_t rpl_R = 0; // floats per record: n_v_heads * (2 * head_dim + 1) + std::vector rpl_n; // per seq: records still to be re-applied onto the base state + std::vector rpl_bank; // per seq: bank holding them + + void rpl_reset(llama_seq_id seq_id); // seq_id < 0: all + // computed before each graph build uint32_t n = 0; @@ -127,7 +139,8 @@ class llama_memory_recurrent : public llama_memory_i { size_t size_s_bytes() const; void state_write_meta(llama_io_write_i & io, const std::vector> & cell_ranges, llama_seq_id seq_id = -1) const; - void state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges) const; + void state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges, + const std::vector> & cell_ranges_s, uint32_t rpl_pending, uint32_t rpl_bank_cur) const; bool state_read_meta(llama_io_read_i & io, uint32_t cell_count, llama_seq_id dest_seq_id = -1); bool state_read_data(llama_io_read_i & io, uint32_t cell_count); @@ -173,6 +186,13 @@ class llama_memory_recurrent_context : public llama_memory_context_i { int32_t s_copy(int i) const; + // ARC-LAB GDN replay (see llama_memory_recurrent::rpl) + bool get_rpl() const; + int32_t rpl_row(int i) const; // plane-0 S row the state of rs i is read from + // records to re-apply for seq set i of the ubatch and their bank; also books this ubatch's records (call once + // per ubatch, when setting the inputs - like s_copy() consumes rs_idx) + void rpl_take(int i, uint32_t n_seq_tokens, int32_t & n_rec, int32_t & bank) const; + private: const llama_memory_status status; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ae67400c6131..ca8e69085183 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -573,6 +573,46 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const int64_t D = S_v * S_v * H_v; const int64_t K = cparams.n_rs_seq + 1; + if (inp->s_rpl_ctrl) { + // ARC-LAB: rollback by replay (independent implementation of the technique described by llamADALL). One + // base state in plane 0 plus the update records of the last K-1 tokens in plane 1 (two banks: this pass + // reads one, its records go to the other) instead of K snapshot planes; see ggml_gated_delta_net_replay. + GGML_ASSERT(state_rows == nullptr && n_seqs == 1 && g->ne[0] == 1); + const int64_t R = H_v * (2 * S_v + 1); + const size_t row_size = ssm_states_all->nb[1]; + const size_t rec_size = ggml_row_size(ssm_states_all->type, R * (K - 1)); + + // both banks of this cell's records: [R*(K-1), 2, n_seqs] + ggml_tensor * rec = ggml_view_3d(ctx0, ssm_states_all, R * (K - 1), 2, n_seqs, rec_size, row_size, + (size_t) (mem_size + kv_head) * row_size); + + ggml_tensor * gdn_out = ggml_gated_delta_net_replay(ctx0, q, k, v, g, b, s, rec, inp->s_rpl_ctrl, (int) K); + res->add_fused_node({n_seq_tokens > 1 ? LLM_FUSED_OP_GDN_CH : LLM_FUSED_OP_GDN_AR, gdn_out, il}); + + ggml_tensor * output = ggml_view_4d(ctx0, gdn_out, + S_v, H_v, n_seq_tokens, n_seqs, + ggml_row_size(gdn_out->type, S_v), + ggml_row_size(gdn_out->type, S_v * H_v), + ggml_row_size(gdn_out->type, S_v * H_v * n_seq_tokens), + 0); + cb(output, "attn_output", il); + + const size_t base_off = ggml_row_size(gdn_out->type, S_v * H_v * n_seq_tokens * n_seqs); + const size_t rec_off = base_off + ggml_row_size(gdn_out->type, D * n_seqs); + + // base state -> plane 0; records -> the bank s_rpl_wrow picks (the SYCL backend fuses both into the op) + ggml_build_forward_expand(gf, ggml_cpy(ctx0, + ggml_view_2d(ctx0, gdn_out, D, n_seqs, ggml_row_size(gdn_out->type, D), base_off), + ggml_view_2d(ctx0, ssm_states_all, D, n_seqs, row_size, (size_t) kv_head * row_size))); + + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, + ggml_view_2d(ctx0, ssm_states_all, R * (K - 1), 2, rec_size, (size_t) (mem_size + kv_head) * row_size), + ggml_view_2d(ctx0, gdn_out, R * (K - 1), n_seqs, rec_size, rec_off), + inp->s_rpl_wrow)); + + return output; + } + const bool raw = gdn_raw_beta && gdn_raw_alpha && gdn_raw_dt_bias && gdn_raw_a; ggml_tensor * gg = raw ? gdn_raw_alpha : g; ggml_tensor * bb = raw ? gdn_raw_beta : b; diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index c5fedfead481..fc7841e004b2 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -491,10 +491,19 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // GPU device in the model is Metal. static const bool gdn_state_rows_env = getenv("GGML_GDN_STATE_GATHER") == nullptr; - const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0; + // ARC-LAB: GDN rollback by replay (the memory decides, see llama_memory_recurrent::rpl): the state is the base + // row in plane 0, gathered with its own row indices (s_copy carries the conv rollback plane) + const bool gdn_rpl = inp->s_rpl_ctrl != nullptr; + + const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0 && !gdn_rpl; ggml_tensor * state; - if (gdn_state_rows) { + if (gdn_rpl) { + state = build_rs(ssm_states_all, inp->s_rpl_main, inp->s_rpl_extra, hparams.n_embd_s(), n_seqs, + mctx_cur->get_n_rs(), mctx_cur->get_head(), mctx_cur->get_size(), mctx_cur->get_rs_z()); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + cb(state, "state_predelta", il); + } else if (gdn_state_rows) { state = build_rs_cache_view(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); cb(state, "state_cache_view", il); } else { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 3c25cc124f14..bb39c3de59e5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1218,6 +1218,8 @@ struct test_case { virtual bool run_whole_graph() { return false; } virtual std::vector fusion_test_nodes() { return {}; } virtual bool use_weight_context() { return false; } + // tensor capacity of the test contexts (ARC-LAB: multi-launch chain tests need more than the default) + virtual size_t ctx_tensors() { return 128; } ggml_cgraph * gf = nullptr; ggml_cgraph * gb = nullptr; @@ -1320,7 +1322,7 @@ struct test_case { mode = MODE_TEST; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead(), + /* .mem_size = */ ggml_tensor_overhead()*ctx_tensors() + ggml_graph_overhead(), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -1511,7 +1513,7 @@ struct test_case { static const size_t graph_nodes = 8192; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead_custom(graph_nodes, false), + /* .mem_size = */ ggml_tensor_overhead()*ctx_tensors() + ggml_graph_overhead_custom(graph_nodes, false), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -1658,7 +1660,7 @@ struct test_case { static const size_t graph_nodes = 8192; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead_custom(graph_nodes, false), + /* .mem_size = */ ggml_tensor_overhead()*ctx_tensors() + ggml_graph_overhead_custom(graph_nodes, false), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -1692,7 +1694,7 @@ struct test_case { const std::vector expect = grad_expect(); ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + 2*ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), + /* .mem_size = */ ggml_tensor_overhead()*ctx_tensors() + 2*ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -4350,6 +4352,197 @@ struct test_rwkv_wkv6 : public test_case { } }; +// ARC-LAB: GGML_OP_GATED_DELTA_NET_REPLAY (rollback by replay; independent implementation of the technique described +// by llamADALL). One graph runs a chain of launches twice, the way the recurrent memory does across MTP steps: +// snapshot path: GATED_DELTA_NET(K) -> CPY of the snapshots into a K-plane cache; the next launch reads plane r +// replay path: GATED_DELTA_NET_REPLAY -> CPY of the base into row 0 + SET_ROWS of the records into one bank of +// row 1; the next launch re-applies m - r records ({n_rec, bank} in its ctrl input) +// where r is the rollback (rejected drafts) chosen per step. Each launch reads the cache through the previous CPY / +// SET_ROWS result, which orders the launches and keeps the cache-write fusion of the real graph. +// exact = true : output = snapshot results - replay results (attention outputs of every step + final state); +// must be exactly 0 on each backend (max |x|, threshold 0) +// exact = false: output = replay results, compared with the CPU backend (nmse) +// The last step must have 1 token, so both paths hold the final state in their first row. +struct test_gdn_replay_chain : public test_case { + const int64_t head_count; // q/k heads + const int64_t v_repeat; // v heads = head_count * v_repeat + const int64_t K; // snapshot slots = 1 + records per bank + const std::vector> steps; // (n_tokens, rollback after the step) + const bool exact; + + static constexpr int64_t S = 128; + + test_gdn_replay_chain(int64_t head_count, int64_t v_repeat, int64_t K, + std::vector> steps, bool exact) + : head_count(head_count), v_repeat(v_repeat), K(K), steps(std::move(steps)), exact(exact) {} + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "GATED_DELTA_NET_REPLAY"; + } + + std::string vars() override { + std::string st; + for (const auto & p : steps) { + st += (st.empty() ? "" : " ") + std::to_string(p.first) + "-" + std::to_string(p.second); + } + return VARS_TO_STR4(head_count, v_repeat, K, exact) + ",steps=[" + st + "]"; + } + + bool run_whole_graph() override { return true; } + size_t ctx_tensors() override { return 64 * steps.size() + 128; } + + double max_nmse_err() override { return exact ? 0.0 : 1e-7; } + + double err(const float * a, const float * b, size_t n) override { + if (!exact) { + return nmse(a, b, n); + } + double mx = 0.0; + for (size_t i = 0; i < n; i++) { + mx = std::max(mx, (double) std::fabs(a[i])); + mx = std::max(mx, (double) std::fabs(b[i])); + } + if (mx > 0.0) { + printf("[replay != snapshot: max |diff| %g] ", mx); + } + return mx; + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t H = head_count * v_repeat; + const int64_t D = S * S * H; + const int64_t R = H * (2 * S + 1); + const int64_t RK = R * (K - 1); + GGML_ASSERT(!steps.empty() && steps.back().first == 1 && steps.back().second == 0); + GGML_ASSERT(2 * RK <= D); + + ggml_tensor * snap_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, K); + ggml_tensor * rpl_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, 2); + ggml_set_name(snap_cache, "snap_cache"); + ggml_set_name(rpl_cache, "rpl_cache"); + + ggml_tensor * snap_src = snap_cache; // [D, rows]: the rows the next snapshot launch can read + ggml_tensor * base_src = rpl_cache; // [D, rows]: row 0 = base state + ggml_tensor * rec_src = ggml_view_3d(ctx, rpl_cache, RK, 2, 1, RK * sizeof(float), rpl_cache->nb[1], rpl_cache->nb[1]); + + ggml_tensor * zero_idx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_name(zero_idx, "idx_0"); + + ggml_tensor * out_snap = nullptr; + ggml_tensor * out_rpl = nullptr; + auto append = [&](ggml_tensor * acc, ggml_tensor * t) { + t = ggml_view_1d(ctx, t, ggml_nelements(t), 0); + return acc ? ggml_concat(ctx, acc, t, 0) : t; + }; + + int64_t snap_plane = 0; // rollback plane the next snapshot launch reads + int64_t n_rec = 0, bank = 0; + for (size_t i = 0; i < steps.size(); ++i) { + const int64_t n = steps[i].first; + const int64_t r = steps[i].second; + const int64_t m = std::min(n - 1, K - 1); + GGML_ASSERT(r <= m); + + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, head_count, n, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, head_count, n, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_set_name(v, "v"); + ggml_set_name(g, "g"); + ggml_set_name(b, "beta"); + q = ggml_l2_norm(ctx, q, 1e-6f); + k = ggml_l2_norm(ctx, k, 1e-6f); + + // snapshot path + ggml_tensor * sidx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_name(sidx, ("idx_" + std::to_string(snap_plane)).c_str()); + ggml_tensor * st = ggml_reshape_4d(ctx, ggml_get_rows(ctx, snap_src, sidx), S, S, H, 1); + ggml_tensor * o = ggml_gated_delta_net(ctx, q, k, v, g, b, st, K); + const int64_t n_written = std::min(n, K); + ggml_tensor * cs = ggml_cpy(ctx, + ggml_view_3d(ctx, o, D, 1, n_written, D * sizeof(float), D * sizeof(float), S * H * n * sizeof(float)), + ggml_view_3d(ctx, snap_cache, D, 1, n_written, snap_cache->nb[1], snap_cache->nb[1], 0)); + if (gf) { + ggml_build_forward_expand(gf, cs); + } + snap_src = ggml_reshape_2d(ctx, cs, D, n_written); + snap_plane = r; + out_snap = append(out_snap, ggml_view_1d(ctx, o, S * H * n, 0)); + + // replay path + ggml_tensor * ctrl = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2); + ggml_set_name(ctrl, ("ctrl_" + std::to_string(n_rec) + "_" + std::to_string(bank)).c_str()); + ggml_tensor * wrow = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + ggml_set_name(wrow, ("wrow_" + std::to_string(bank ^ 1)).c_str()); + ggml_tensor * bs = ggml_reshape_4d(ctx, ggml_get_rows(ctx, base_src, zero_idx), S, S, H, 1); + ggml_tensor * o2 = ggml_gated_delta_net_replay(ctx, q, k, v, g, b, bs, rec_src, ctrl, (int) K); + const size_t base_off = S * H * n * sizeof(float); + ggml_tensor * cb = ggml_cpy(ctx, ggml_view_2d(ctx, o2, D, 1, D * sizeof(float), base_off), + ggml_view_2d(ctx, rpl_cache, D, 1, rpl_cache->nb[1], 0)); + ggml_tensor * sr = ggml_set_rows(ctx, + ggml_view_2d(ctx, rpl_cache, RK, 2, RK * sizeof(float), rpl_cache->nb[1]), + ggml_view_2d(ctx, o2, RK, 1, RK * sizeof(float), base_off + D * sizeof(float)), wrow); + if (gf) { + ggml_build_forward_expand(gf, cb); + ggml_build_forward_expand(gf, sr); + } + base_src = ggml_reshape_2d(ctx, cb, D, 1); + rec_src = ggml_reshape_3d(ctx, sr, RK, 2, 1); + n_rec = m - r; + bank ^= 1; + out_rpl = append(out_rpl, ggml_view_1d(ctx, o2, S * H * n, 0)); + + if (i + 1 == steps.size()) { + out_snap = append(out_snap, ggml_view_1d(ctx, cs, D, 0)); + out_rpl = append(out_rpl, ggml_view_1d(ctx, cb, D, 0)); + } + } + + ggml_tensor * out = exact ? ggml_sub(ctx, out_snap, out_rpl) : out_rpl; + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + std::mt19937 rng(1234); + std::uniform_real_distribution u(-1.0f, 1.0f); + std::vector s0(S * S * head_count * v_repeat); + for (auto & x : s0) { + x = u(rng); + } + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (ggml_is_view_op(t->op) || t->op != GGML_OP_NONE) { continue; } + const std::string name = t->name; + if (name == "snap_cache" || name == "rpl_cache") { + // the same initial state in row 0 of both caches, zeros elsewhere (no records to replay yet) + std::vector z(ggml_nelements(t), 0.0f); + std::copy(s0.begin(), s0.end(), z.begin()); + ggml_backend_tensor_set(t, z.data(), 0, ggml_nbytes(t)); + } else if (name.rfind("idx_", 0) == 0) { + const int32_t x = std::stoi(name.substr(4)); + ggml_backend_tensor_set(t, &x, 0, sizeof(x)); + } else if (name.rfind("ctrl_", 0) == 0) { + const size_t us = name.find('_', 5); + const int32_t x[2] = { std::stoi(name.substr(5, us - 5)), std::stoi(name.substr(us + 1)) }; + ggml_backend_tensor_set(t, x, 0, sizeof(x)); + } else if (name.rfind("wrow_", 0) == 0) { + const int64_t x = std::stoll(name.substr(5)); + ggml_backend_tensor_set(t, &x, 0, sizeof(x)); + } else if (name == "g") { + init_tensor_uniform(t, -3.0f, -1e-4f); // keep some of the state alive across the chain + } else if (name == "beta") { + init_tensor_uniform(t, 0.0f, 1.0f); + } else if (name == "v") { + init_tensor_uniform(t, -0.3f, 5.0f); + } else if (t->type == GGML_TYPE_F32) { + init_tensor_uniform(t); + } + } + } +}; + // GGML_OP_GATED_DELTA_NET struct test_gated_delta_net : public test_case { const ggml_type type; @@ -10795,6 +10988,18 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 1, 2, 2, false, false, 1, false, -1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 64, 1, 1, false, false, 1, false, -1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 1, 2, 1, false, false, 2, true, -1, true)); + + // ARC-LAB: GDN rollback by replay - replay chain == snapshot chain (exact), and replay chain vs CPU. + // steps (n_tokens, rollback): prompt-sized launches (>= 16 tokens: blocked kernel on SYCL), MTP verify launches + // with partial acceptance, full rejection (rollback = n - 1) and a prompt after pending records. + for (bool exact : { true, false }) { + const std::vector> chain_a = { {20, 0}, {5, 2}, {5, 0}, {3, 1}, {5, 4}, {2, 0}, {1, 0} }; + const std::vector> chain_b = { {5, 1}, {17, 0}, {4, 3}, {18, 2}, {5, 1}, {1, 0} }; + test_cases.emplace_back(new test_gdn_replay_chain(4, 1, 5, chain_a, exact)); + test_cases.emplace_back(new test_gdn_replay_chain(2, 3, 5, chain_b, exact)); + test_cases.emplace_back(new test_gdn_replay_chain(16, 3, 5, chain_a, exact)); // Bonsai: 16 q/k, 48 v heads + test_cases.emplace_back(new test_gdn_replay_chain(2, 2, 3, { {7, 1}, {3, 2}, {3, 0}, {1, 0} }, exact)); + } test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 16, 1, 1, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 64, 1, 2)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 1)); From 602cd2f0621f03f494b6657650d00c3c529dd358 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 13:50:04 +1100 Subject: [PATCH 07/24] mtp: draft chain off by default until its graph is reused (measured -0.2..-0.8% on the B580) Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 4798af11c85414539b2c8f8b7c3b5b9d049ba4db on Torchit1/llama.cpp arc-b580) --- common/speculative.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index b3ec6ff87aa9..e9ece1006bb1 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2418,15 +2418,15 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { deferred.assign(n_seq, {}); - // chained drafts: default on when the head has a trimmed draft vocabulary (its token_embd rows fit on the - // device; measured slower without one), LLAMA_MTP_CHAIN=1 forces, =0 off. Greedy single-sequence drafting only. + // chained drafts: default OFF for now (B580, 2026-10-10: identical drafts but -0.2..-0.8%, the chain graph is rebuilt + // on every draft; revisit with graph reuse). LLAMA_MTP_CHAIN=1 turns it on. Greedy single-sequence drafting only. { const char * env = getenv("LLAMA_MTP_CHAIN"); const int32_t sup = llama_mtp_chain_support(ctx_dft); - const bool want = env ? atoi(env) != 0 : (sup & 2) != 0; + const bool want = env ? atoi(env) != 0 : false; const char * why = nullptr; if (!want) { - why = env ? "LLAMA_MTP_CHAIN=0" : "no trimmed draft head (LLAMA_MTP_CHAIN=1 forces it)"; + why = env ? "LLAMA_MTP_CHAIN=0" : "off by default (LLAMA_MTP_CHAIN=1 turns it on)"; } else if (sup == 0) { why = "not supported by this context (qwen35 single head, flash attention, plain KV cache)"; } else if (is_mem_shared || chain_heads) { From fdb28c12ec3c49ed0054e33870bc7c8c697395cb Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:04:20 +1100 Subject: [PATCH 08/24] mtp: reuse the draft-chain graph across drafts (cached per batch shape) Independent implementation of the technique described by llamADALL (follow-up). The chain graph was built, split and allocated on every draft because allow_reuse refused chain graphs, which cost more host time than the chain saved (-0.2..-0.8% on the B580, drafts identical). Now: - llama_context keeps a scheduler + graph result per chain batch shape (steps, rows) and swaps them in around process_ubatch, so a draft of a known shape only sets inputs and computes; regular decodes keep their own gf_res_prev. - allow_reuse accepts chain graphs with the same step count; the chain input's can_reuse also requires the same rows and the same per-step KV lengths (get_n_kv_prefix), which are baked into the step views. - per-step position and output-row inputs are packed into two tensors (two uploads per draft instead of 2 x steps). Same graph, same values: outputs unchanged. Still default off. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit d6435c6f00c28bf31a463f6b557f9a0ca4fa0e7d on Torchit1/llama.cpp arc-b580) --- src/llama-context.cpp | 55 +++++++++++++++++++++------- src/llama-context.h | 10 ++++++ src/llama-graph.h | 4 +-- src/models/qwen35.cpp | 83 +++++++++++++++++++++++++++---------------- 4 files changed, 108 insertions(+), 44 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 223273004829..eb4f6ea76edc 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -711,6 +711,8 @@ void llama_context::sched_reserve() { synchronize(); + mtp_chain_cache.clear(); // ARC-LAB: built under the old settings + const int64_t t_start_us = ggml_time_us(); const uint32_t n_seqs = cparams.n_seq_max; @@ -1432,6 +1434,24 @@ void llama_context::set_mtp_chain(uint32_t n_steps) { ggml_type_name(src->type), staging.size() / 1048576.0, ggml_backend_dev_name(dev)); } +llama_context::mtp_chain_entry * llama_context::mtp_chain_entry_for(uint32_t n_tokens) { + const uint32_t key = (cparams.mtp_chain_n << 16) | n_tokens; + auto it = mtp_chain_cache.find(key); + if (it != mtp_chain_cache.end()) { + return &it->second; + } + if (mtp_chain_cache.size() >= 16) { // shapes are catch-up rows (1..n_max+1) + steps: a handful in practice + synchronize(); + mtp_chain_cache.clear(); + } + const size_t max_nodes = graph_max_nodes(std::min(cparams.n_ctx, cparams.n_ubatch)); + mtp_chain_entry e; + e.sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, + cparams.op_offload)); + e.res.reset(new llm_graph_result(max_nodes)); + return &(mtp_chain_cache[key] = std::move(e)); +} + const int32_t * llama_context::get_mtp_chain_ids() { synchronize(); @@ -2174,8 +2194,31 @@ int llama_context::decode(const llama_batch & batch_inp) { ggml_status status; + // ARC-LAB: a chained MTP draft runs on the cached scheduler + graph of its batch shape (mtp_chain_cache) + mtp_chain_entry * chain_e = cparams.mtp_chain_n > 0 ? mtp_chain_entry_for(ubatch.n_tokens) : nullptr; + if (chain_e) { + std::swap(sched, chain_e->sched); + std::swap(gf_res_prev, chain_e->res); + } + const auto * res = process_ubatch(ubatch, ctx_type_to_graph_type(cparams.ctx_type), mctx.get(), status); + if (chain_e) { + // the per-step argmax rows, mapped to tokens in get_mtp_chain_ids() + if (res && !res->t_mtp_chain.empty()) { + mtp_chain_out.resize(res->t_mtp_chain.size()); + for (size_t s = 0; s < res->t_mtp_chain.size(); ++s) { + ggml_tensor * t = res->t_mtp_chain[s]; + ggml_backend_t b = ggml_backend_sched_get_tensor_backend(sched.get(), t); + GGML_ASSERT(b != nullptr); + ggml_backend_tensor_get_async(b, t, &mtp_chain_out[s], 0, sizeof(int32_t)); + } + mtp_chain_n_out = (uint32_t) res->t_mtp_chain.size(); + } + std::swap(sched, chain_e->sched); + std::swap(gf_res_prev, chain_e->res); + } + if (!res) { // the last ubatch failed or was aborted -> remove all positions of that ubatch from the memory module llama_pos pos_min[LLAMA_MAX_SEQ]; @@ -2316,18 +2359,6 @@ int llama_context::decode(const llama_batch & batch_inp) { extract_layer_inputs(res, n_tokens_prev, ubatch.n_tokens); - // ARC-LAB: chained MTP draft: the per-step argmax rows, mapped to tokens in get_mtp_chain_ids() - if (!res->t_mtp_chain.empty()) { - mtp_chain_out.resize(res->t_mtp_chain.size()); - for (size_t s = 0; s < res->t_mtp_chain.size(); ++s) { - ggml_tensor * t = res->t_mtp_chain[s]; - ggml_backend_t b = ggml_backend_sched_get_tensor_backend(sched.get(), t); - GGML_ASSERT(b != nullptr); - ggml_backend_tensor_get_async(b, t, &mtp_chain_out[s], 0, sizeof(int32_t)); - } - mtp_chain_n_out = (uint32_t) res->t_mtp_chain.size(); - } - // extract nextn embeddings before // only meaningful in LLAMA_POOLING_TYPE_NONE (per-token); other pooling modes are ignored. { diff --git a/src/llama-context.h b/src/llama-context.h index ef04bf85de9c..16f308b140b2 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -418,6 +418,16 @@ struct llama_context { std::vector mtp_chain_vocab; // draft-head index -> token id (empty: no draft head, indices are tokens) std::vector mtp_chain_out; std::vector mtp_chain_tok; // mtp_chain_out mapped to token ids + + // ARC-LAB: chain graphs are kept per batch shape, each with its own scheduler, so a draft whose shape (rows, steps) + // and per-step KV lengths match an earlier one reuses its built and allocated graph: no build, split or allocation + // per draft, and the regular decodes keep their own reusable graph in gf_res_prev + struct mtp_chain_entry { + ggml_backend_sched_ptr sched; + llm_graph_result_ptr res; + }; + std::map mtp_chain_cache; + mtp_chain_entry * mtp_chain_entry_for(uint32_t n_tokens); uint32_t mtp_chain_n_out = 0; // host buffer for the model output (logits and embeddings) diff --git a/src/llama-graph.h b/src/llama-graph.h index 0a5b378ddc58..7b1fd48989f8 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -939,8 +939,8 @@ struct llm_graph_params { return false; } - // ARC-LAB: a chained MTP draft graph bakes per-step KV lengths in; never reuse it or reuse into it - if (cparams.mtp_chain_n > 0 || other.cparams.mtp_chain_n > 0) { + // ARC-LAB: a chained MTP draft graph only for the same step count (its input checks the per-step KV lengths) + if (cparams.mtp_chain_n != other.cparams.mtp_chain_n) { return false; } diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index fc7841e004b2..0e81f3dba9a9 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -618,47 +618,63 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_ffn(ggml_tensor * cur, cons return cur; } -// ARC-LAB: per-step inputs of a chained MTP draft graph (graph_mtp::build_chain): the positions of each step's rows in -// the layout build_inp_pos() gives a decode of just those rows, and the step's output row +// ARC-LAB: per-step inputs of a chained MTP draft graph (graph_mtp::build_chain), packed into two tensors so a draft +// uploads them in two writes: pos holds each step's positions in the layout build_inp_pos() gives a decode of just +// that step's rows (step after step), out_idx each step's output row. The graph reads them through per-step views. class llm_graph_input_mtp_chain : public llm_graph_input_i { public: explicit llm_graph_input_mtp_chain(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {} struct step { - int64_t r0; // first ubatch row of the step - int64_t nr; // rows - ggml_tensor * pos; // I32 [nr * n_pos_per_embd] - ggml_tensor * out_idx; // I32 [1] = nr - 1 (the step's last row produces its draft) + int64_t r0; // first ubatch row of the step + int64_t nr; // rows + uint32_t n_kv; // KV length the step attends over (baked into the graph's views) }; std::vector steps; + int64_t n_tokens = 0; + ggml_tensor * pos = nullptr; // I32 [n_tokens * n_pos_per_embd] + ggml_tensor * out_idx = nullptr; // I32 [n_steps], step s: nr - 1 (the step's last row produces its draft) void set_input(const llama_ubatch * ubatch) override { - const int64_t n_tokens = ubatch->n_tokens; - std::vector v; - for (const auto & st : steps) { - v.assign((size_t) st.nr * n_pos_per_embd, 0); + GGML_ASSERT(ubatch->n_tokens == n_tokens); + std::vector v((size_t) n_tokens * n_pos_per_embd, 0); + std::vector o(steps.size()); + size_t off = 0; + for (size_t s = 0; s < steps.size(); ++s) { + const auto & st = steps[s]; + int32_t * vs = v.data() + off; for (int64_t r = 0; r < st.nr; ++r) { if (ubatch->token && n_pos_per_embd == 4) { // text tokens under M-RoPE, as llm_graph_input_pos: three equal dims, the fourth 0 for (int d = 0; d < 3; ++d) { - v[d * st.nr + r] = ubatch->pos[st.r0 + r]; + vs[d * st.nr + r] = ubatch->pos[st.r0 + r]; } } else { for (uint32_t d = 0; d < n_pos_per_embd; ++d) { - v[d * st.nr + r] = ubatch->pos[d * n_tokens + st.r0 + r]; + vs[d * st.nr + r] = ubatch->pos[d * n_tokens + st.r0 + r]; } } } - ggml_backend_tensor_set(st.pos, v.data(), 0, v.size() * sizeof(int32_t)); - const int32_t out_row = (int32_t) st.nr - 1; - ggml_backend_tensor_set(st.out_idx, &out_row, 0, sizeof(out_row)); + off += (size_t) st.nr * n_pos_per_embd; + o[s] = (int32_t) st.nr - 1; } + ggml_backend_tensor_set(pos, v.data(), 0, v.size() * sizeof(int32_t)); + ggml_backend_tensor_set(out_idx, o.data(), 0, o.size() * sizeof(int32_t)); } - // never reused (llm_graph_params::allow_reuse refuses chain graphs) + // reusable when the batch has the same rows and every step would attend over the same KV length (the rest of the + // topology is checked by allow_reuse and the other inputs) bool can_reuse(const llm_graph_params & params) override { - GGML_UNUSED(params); - return false; + if (params.ubatch.n_tokens != n_tokens || params.cparams.mtp_chain_n != steps.size()) { + return false; + } + const auto * kvc = static_cast(params.mctx); + for (const auto & st : steps) { + if (kvc->get_n_kv_prefix((uint32_t) (st.r0 + st.nr)) != st.n_kv) { + return false; + } + } + return true; } const uint32_t n_pos_per_embd; @@ -956,18 +972,23 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { auto * inp_attn = build_attn_inp_kv(); const llama_kv_cache_context * kvc = inp_attn->mctx; - auto chain_inp = std::make_unique(hparams.n_pos_per_embd()); + const uint32_t n_pos = hparams.n_pos_per_embd(); + auto chain_inp = std::make_unique(n_pos); + chain_inp->n_tokens = n_tokens; + chain_inp->pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens * n_pos); + chain_inp->out_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_steps); + ggml_set_input(chain_inp->pos); + ggml_set_input(chain_inp->out_idx); for (int64_t s = 0; s < n_steps; ++s) { llm_graph_input_mtp_chain::step st; - st.r0 = s == 0 ? 0 : n_lead - 1 + s; - st.nr = s == 0 ? n_lead : 1; - st.pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, st.nr * hparams.n_pos_per_embd()); - st.out_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1); - ggml_set_input(st.pos); - ggml_set_input(st.out_idx); + st.r0 = s == 0 ? 0 : n_lead - 1 + s; + st.nr = s == 0 ? n_lead : 1; + st.n_kv = kvc->get_n_kv_prefix((uint32_t) (st.r0 + st.nr)); chain_inp->steps.push_back(st); } - const auto steps = chain_inp->steps; + const auto steps = chain_inp->steps; + ggml_tensor * pos_all = chain_inp->pos; + ggml_tensor * out_all = chain_inp->out_idx; res->add_input(std::move(chain_inp)); // token embeddings: the host table for the ubatch's own tokens (catch-up tokens need not be in the draft @@ -1018,7 +1039,9 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { cb(tok_embd, "mtp_tok_embd", il); // this step's slice of the ubatch's KV indices and mask, sized as a decode of rows [0, r0 + nr) would be - const uint32_t n_kv_s = kvc->get_n_kv_prefix((uint32_t) (st.r0 + nr)); + const uint32_t n_kv_s = st.n_kv; + ggml_tensor * st_pos = ggml_view_1d(ctx0, pos_all, nr * n_pos, st.r0 * n_pos * ggml_element_size(pos_all)); + ggml_tensor * st_out = ggml_view_1d(ctx0, out_all, 1, s * ggml_element_size(out_all)); ggml_tensor * k_idxs = inp_attn->get_k_idxs(); ggml_tensor * v_idxs = inp_attn->get_v_idxs(); k_idxs = ggml_view_1d(ctx0, k_idxs, nr, st.r0 * ggml_element_size(k_idxs)); @@ -1071,10 +1094,10 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, nr); cb(Vcur, "mtp_Vcur", il); - Qcur = ggml_rope_multi(ctx0, Qcur, st.pos, nullptr, + Qcur = ggml_rope_multi(ctx0, Qcur, st_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, st.pos, nullptr, + Kcur = ggml_rope_multi(ctx0, Kcur, st_pos, nullptr, n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); @@ -1127,7 +1150,7 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { cb(cur, "h_nextn", -1); // the step's output row (masked h_nextn: the next step's hidden input) - ggml_tensor * h_out = ggml_get_rows(ctx0, cur, st.out_idx); + ggml_tensor * h_out = ggml_get_rows(ctx0, cur, st_out); cb(h_out, "mtp_shared_head_norm", -1); // greedy draft: argmax of the head's logits. With a draft head that is an index into its vocabulary (the From c18c801257de709f6e055f5897bd8dd0c59a157d Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:20:33 +1100 Subject: [PATCH 09/24] mtp: draft-time instrumentation (LLAMA_MTP_DRAFT_TIMING, GGML_SYCL_OP_TIMING) Both measured runs of the draft chain spend ~6.5 ms per draft call, chain or loop, against a ~3.4 ms weight-read floor. To find where it goes: - LLAMA_MTP_DRAFT_TIMING=1: host phases of every MTP draft() call (pre-draft flushes, batch building, llama_decode submission, GPU wait + readback), averaged and printed every 64 calls, for both the chain and the loop. - GGML_SYCL_OP_TIMING=1: GPU time per op type for draft-sized graphs (< 2000 nodes, every mat-mul <= 8 rows), draining the queue after every kernel group (serialised upper bound incl. one launch latency each), every 64 graphs. Both off by default; no change to what is computed. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 60edb38e7d8833629684c07ac6925f4ce88bc0e3 on Torchit1/llama.cpp arc-b580) --- common/speculative.cpp | 55 +++++++++++++++++++++++++ ggml/src/ggml-sycl/ggml-sycl.cpp | 69 ++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index e9ece1006bb1..ed482143e115 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -6,6 +6,7 @@ #include "ggml.h" #include "llama.h" +#include #include #include #ifdef LLAMA_DSPARK_MARKOV_CUDA @@ -41,6 +42,40 @@ // ARC-LAB: LLAMA_MTP_DRAFT_LOG= appends one JSON line per MTP draft step (position predicted, depth, top 3 candidates // with probabilities, whether p_min stopped the draft) plus a "begin" line per request, for offline acceptance / tree // simulation (arctune, 2026-10-05). +// ARC-LAB: LLAMA_MTP_DRAFT_TIMING=1: host-side phase timing of the MTP draft() calls (chain and loop), averaged and +// printed to stderr every 64 calls. pre = deferred flushes before the draft, batch = building batches, decode = the +// llama_decode calls (batch split, slot search, input uploads, kernel submission; no wait), wait = waiting for the GPU +// and reading the draft back (chain: llama_get_mtp_chain_ids; loop: the per-step sampling, which syncs). +struct common_mtp_draft_timing { + using clk = std::chrono::steady_clock; + const bool on = [] { const char * e = std::getenv("LLAMA_MTP_DRAFT_TIMING"); return e && std::atoi(e) != 0; }(); + clk::time_point t0, t; + int calls = 0, chain_calls = 0, decodes = 0; + double total = 0, pre = 0, batch = 0, decode = 0, wait = 0, rest = 0; + void begin() { if (on) { t0 = t = clk::now(); } } + void lap(double & acc) { + if (!on) { return; } + const auto n = clk::now(); + acc += std::chrono::duration(n - t).count(); + t = n; + } + void end(bool chain) { + if (!on) { return; } + lap(rest); + total += std::chrono::duration(t - t0).count(); + calls++; + chain_calls += chain; + if (calls == 64) { + fprintf(stderr, "mtp-draft-timing: %d calls (%d chained), per call: total %.3f ms = pre %.3f + batch %.3f + " + "decode %.3f (%.2f llama_decode) + wait/readback %.3f + rest %.3f\n", calls, chain_calls, total / calls, + pre / calls, batch / calls, decode / calls, (double) decodes / calls, wait / calls, rest / calls); + calls = chain_calls = decodes = 0; + total = pre = batch = decode = wait = rest = 0; + } + } +}; +static common_mtp_draft_timing g_mtp_dt; + static FILE * common_mtp_draft_log() { static FILE * f = [] { const char * p = std::getenv("LLAMA_MTP_DRAFT_LOG"); @@ -2465,6 +2500,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const size_t row_bytes = (size_t) n_embd * sizeof(float); + g_mtp_dt.lap(g_mtp_dt.pre); common_sampler_reset(smpls[0].get()); common_batch_clear(batch); @@ -2481,9 +2517,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[0].data(), row_bytes); } + g_mtp_dt.lap(g_mtp_dt.batch); llama_set_mtp_chain(ctx_dft, n_draft); const int32_t rc = llama_decode(ctx_dft, batch); llama_set_mtp_chain(ctx_dft, 0); + g_mtp_dt.lap(g_mtp_dt.decode); + g_mtp_dt.decodes++; if (rc != 0) { SPC_ERR("llama_decode(ctx_dft) chained draft returned %d\n", (int) rc); if (n_catch > 0) { @@ -2496,6 +2535,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } const llama_token * ids = llama_get_mtp_chain_ids(ctx_dft); + g_mtp_dt.lap(g_mtp_dt.wait); if (ids == nullptr) { SPC_ERR("%s", "chained draft produced no tokens\n"); return true; @@ -2768,9 +2808,18 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } void draft(common_speculative_draft_params_vec & dparams) override { + g_mtp_dt.begin(); + draft_impl(dparams); + g_mtp_dt.end(last_draft_chained); + } + + bool last_draft_chained = false; + + void draft_impl(common_speculative_draft_params_vec & dparams) { auto & ctx_dft = params.ctx_dft; const size_t row_bytes = (size_t) n_embd * sizeof(float); + last_draft_chained = false; // deferred catch-up rows that do not lead straight into this draft position are decoded // on their own first (uses `batch`, so before it is built) @@ -2825,6 +2874,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } if (chain_mode && draft_chain(dparams)) { + last_draft_chained = true; auto & dp = dparams[0]; if (dp.result->size() < (size_t) params.n_min) { dp.result->clear(); @@ -2832,6 +2882,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { return; } + g_mtp_dt.lap(g_mtp_dt.pre); common_batch_clear(batch); // keep track of which sequences are still drafting @@ -2891,7 +2942,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { llama_set_nextn_layer_offset(ctx_dft, i); } + g_mtp_dt.lap(g_mtp_dt.batch); int ret = llama_decode(ctx_dft, batch); + g_mtp_dt.lap(g_mtp_dt.decode); + g_mtp_dt.decodes++; if (ret != 0) { SPC_ERR("llama_decode[%d] returned %d\n", i, ret); if (i == 0) { @@ -2924,6 +2978,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); common_sampler_sample(smpl, ctx_dft, i_last[seq_id], true); + g_mtp_dt.lap(g_mtp_dt.wait); const float * h_row = llama_get_embeddings_nextn_ith(ctx_dft, i_last[seq_id]); const auto * cur_p = common_sampler_get_candidates(smpl, true); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index f7037bf36087..dfd77750deb8 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6255,6 +6255,45 @@ struct sycl_host_timing { }; static sycl_host_timing g_ht; +// ARC-LAB: GGML_SYCL_OP_TIMING=1 -> GPU time per op type for draft-sized graphs (fewer than 2000 nodes, every mat-mul at +// most 8 rows: the MTP draft decodes and chains, not the target's verify graph). The queue is drained after every kernel +// group, so each entry is its kernel time plus one launch latency (a serialised upper bound); the sum is reported next to +// the kernel-group count. Printed every 64 such graphs. Timing-only: it slows the draft, never use it for speed runs. +struct sycl_op_timing { + bool on = getenv("GGML_SYCL_OP_TIMING") && atoi(getenv("GGML_SYCL_OP_TIMING")); + using clk = std::chrono::steady_clock; + int graphs = 0; + long groups = 0, nodes = 0; + double total = 0; + std::map> per_op; + static bool draft_sized(const ggml_cgraph * g) { + if (g->n_nodes >= 2000) return false; + bool mm = false; + for (int i = 0; i < g->n_nodes; ++i) { + const ggml_tensor * n = g->nodes[i]; + if (n->op == GGML_OP_MUL_MAT) { + mm = true; + if (n->src[1]->ne[1] * n->src[1]->ne[2] > 8) return false; + } + } + return mm; + } + void report() { + fprintf(stderr, "op-timing: %d draft-sized graphs, per graph: %.3f ms serialised over %.1f kernel groups (%.0f nodes)\n", + graphs, total / graphs, (double) groups / graphs, (double) nodes / graphs); + std::vector> v; + for (auto & kv : per_op) v.push_back({ kv.second.first, kv.first }); + std::sort(v.rbegin(), v.rend()); + for (size_t i = 0; i < v.size() && i < 16; ++i) { + const auto & a = per_op[v[i].second]; + fprintf(stderr, "op-timing: %-22s %.3f ms/graph (%.1f groups/graph, %.1f us each)\n", v[i].second.c_str(), + a.first / graphs, (double) a.second / graphs, 1000.0 * a.first / std::max(1L, a.second)); + } + graphs = 0; groups = 0; nodes = 0; total = 0; per_op.clear(); + } +}; +static sycl_op_timing g_ot; + static void ggml_backend_sycl_synchronize(ggml_backend_t backend) try { GGML_SYCL_DEBUG("[SYCL] call %s\n", __func__); ggml_backend_sycl_context * sycl_ctx = (ggml_backend_sycl_context *)backend->context; @@ -7249,6 +7288,27 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc const bool ht_on = g_ht.on && cgraph->n_nodes > 500; sycl_host_timing::clk::time_point ht_prev = sycl_host_timing::clk::now(); const char * ht_op = nullptr; + const bool ot_on = g_ot.on && sycl_op_timing::draft_sized(cgraph); + sycl_op_timing::clk::time_point ot_prev{}; + const char * ot_op = nullptr; + auto ot_mark = [&](const char * next) { // drain the queue, charge the time since the last mark to the op that ran + SYCL_CHECK(CHECK_TRY_ERROR(sycl_ctx->stream()->wait())); + const auto now = sycl_op_timing::clk::now(); + if (ot_op) { + const double ms = std::chrono::duration(now - ot_prev).count(); + auto & a = g_ot.per_op[ot_op]; + a.first += ms; + a.second++; + g_ot.total += ms; + g_ot.groups++; + } + ot_prev = now; + ot_op = next; + }; + if (ot_on) { + ot_mark(nullptr); // inputs and earlier work are not charged + g_ot.nodes += cgraph->n_nodes; + } for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (ht_on) { @@ -7273,6 +7333,9 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { continue; } + if (ot_on) { + ot_mark(ggml_op_desc(node)); + } const int nodes_to_skip = ggml_sycl_fuse(*sycl_ctx, cgraph, i); if (nodes_to_skip != 0) { @@ -7412,6 +7475,12 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc } GGML_ASSERT(ok); } + if (ot_on) { + ot_mark(nullptr); + if (++g_ot.graphs == 64) { + g_ot.report(); + } + } } #ifdef GGML_SYCL_GRAPH From 7ac3c1b9c57f703cf6b5a2bc86ae6e42d6fbd507 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:20:56 +1100 Subject: [PATCH 10/24] gdn replay: allocate only S planes 0 and 1 when rollback by replay is on Replay mode keeps the base state in S plane 0 and the two record banks in plane 1; planes 2..n_rs_seq were allocated but never used (432 MiB on Bonsai at n_rs_seq = 4, 48 GDN layers x 3 planes x 3 MiB, independent of -c). - llama_memory_recurrent: the replay decision now runs before the tensors are created (layer selection via the same filter as the allocation), and S gets 2 * mem_size rows in replay mode. The conv state r_l keeps its 1 + n_rs_seq planes and rs_idx. - LLAMA_GDN_REPLAY=0 and every case where replay does not apply keep the full (1 + n_rs_seq)-plane S allocation, unchanged. - No offset changes: the graph, the SYCL fused cache write and the state save/restore only ever address plane 0 (row head) and plane 1 (row size + head), so the s_trans = 0/2 file layout is the same and files from cr-replay still load. - build_recurrent_attn asserts the 2-plane S shape in replay mode. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit bda090edcc7ebc15416b0a368e99d3f276ce4f48 on Torchit1/llama.cpp arc-b580) --- src/llama-memory-recurrent.cpp | 74 ++++++++++++++++++---------------- src/llama-memory-recurrent.h | 4 +- src/models/delta-net-base.cpp | 2 + 3 files changed, 44 insertions(+), 36 deletions(-) diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index f5d33f30397b..808524545b1f 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -72,6 +72,41 @@ llama_memory_recurrent::llama_memory_recurrent( return it->second.get(); }; + // ARC-LAB: GDN rollback by replay (LLAMA_GDN_REPLAY=0 keeps the snapshot planes). Applies to a single sequence + // of a qwen35 model whose recurrent layers all run on SYCL; anything else silently keeps the snapshot path. + // Decided before the tensors are created: in replay mode S gets only planes 0 (base) and 1 (record banks), so + // planes 2..n_rs_seq are not allocated. The snapshot path keeps the full (1 + n_rs_seq)-plane layout. + { + const char * env = getenv("LLAMA_GDN_REPLAY"); + bool on = !(env && atoi(env) == 0) && n_rs_seq > 0 && n_seq_max == 1 && mem_size == 1 && offload && + model.arch == LLM_ARCH_QWEN35 && type_s == GGML_TYPE_F32 && + getenv("GGML_SYCL_GDN_CHUNKED") == nullptr; // the opt-in chunked prefill has no replay variant + const uint32_t S = hparams.ssm_d_state; + const uint32_t H = hparams.ssm_dt_rank; + on = on && S == 128 && H > 0 && hparams.ssm_d_inner == S * H && hparams.n_embd_s() == S * S * H; + const uint32_t R = H * (2 * S + 1); + // both banks of n_rs_seq records must fit in one S row + on = on && (uint64_t) 2 * n_rs_seq * R <= (uint64_t) hparams.n_embd_s(); + bool any = false; + for (int i = 0; on && i < n_layer; i++) { + if (filter && !filter(i)) { // same layer selection as the allocation below + continue; + } + any = true; + ggml_backend_dev_t dev = model.dev_layer(i); + ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; + on = reg && strcmp(ggml_backend_reg_name(reg), "SYCL") == 0; + } + rpl = on && any; + rpl_R = rpl ? R : 0; + if (rpl) { + LLAMA_LOG_INFO("%s: GDN rollback by replay: on (%u records of %u floats per bank, S planes 2 of %u)\n", __func__, + n_rs_seq, R, 1 + n_rs_seq); + } else if (n_rs_seq > 0) { + LLAMA_LOG_INFO("%s: GDN rollback by replay: off (snapshot planes)\n", __func__); + } + } + r_l.resize(n_layer); s_l.resize(n_layer); @@ -99,48 +134,17 @@ llama_memory_recurrent::llama_memory_recurrent( throw std::runtime_error("failed to create ggml context for rs cache"); } - const uint32_t n_rows = mem_size * (1 + n_rs_seq); + const uint32_t n_rows = mem_size * (1 + n_rs_seq); + // ARC-LAB replay: S needs plane 0 (base) and plane 1 (record banks) only; the conv state keeps all its planes + const uint32_t n_rows_s = rpl ? mem_size * 2 : n_rows; ggml_tensor * r = ggml_new_tensor_2d(ctx, type_r, hparams.n_embd_r(), n_rows); - ggml_tensor * s = ggml_new_tensor_2d(ctx, type_s, hparams.n_embd_s(), n_rows); + ggml_tensor * s = ggml_new_tensor_2d(ctx, type_s, hparams.n_embd_s(), n_rows_s); ggml_format_name(r, "cache_r_l%d", i); ggml_format_name(s, "cache_s_l%d", i); r_l[i] = r; s_l[i] = s; } - // ARC-LAB: GDN rollback by replay (LLAMA_GDN_REPLAY=0 keeps the snapshot planes). Applies to a single sequence - // of a qwen35 model whose recurrent layers all run on SYCL; anything else silently keeps the snapshot path. - // The tensor layout is the same either way (planes 0 and 1 of S are reused), only their meaning changes. - { - const char * env = getenv("LLAMA_GDN_REPLAY"); - bool on = !(env && atoi(env) == 0) && n_rs_seq > 0 && n_seq_max == 1 && mem_size == 1 && offload && - model.arch == LLM_ARCH_QWEN35 && type_s == GGML_TYPE_F32 && - getenv("GGML_SYCL_GDN_CHUNKED") == nullptr; // the opt-in chunked prefill has no replay variant - const uint32_t S = hparams.ssm_d_state; - const uint32_t H = hparams.ssm_dt_rank; - on = on && S == 128 && H > 0 && hparams.ssm_d_inner == S * H && hparams.n_embd_s() == S * S * H; - const uint32_t R = H * (2 * S + 1); - // both banks of n_rs_seq records must fit in one S row - on = on && (uint64_t) 2 * n_rs_seq * R <= (uint64_t) hparams.n_embd_s(); - bool any = false; - for (int i = 0; on && i < n_layer; i++) { - if (s_l[i] == nullptr) { - continue; - } - any = true; - ggml_backend_dev_t dev = model.dev_layer(i); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - on = reg && strcmp(ggml_backend_reg_name(reg), "SYCL") == 0; - } - rpl = on && any; - rpl_R = rpl ? R : 0; - if (rpl) { - LLAMA_LOG_INFO("%s: GDN rollback by replay: on (%u records of %u floats per bank)\n", __func__, n_rs_seq, R); - } else if (n_rs_seq > 0) { - LLAMA_LOG_INFO("%s: GDN rollback by replay: off (snapshot planes)\n", __func__); - } - } - // allocate tensors and initialize the buffers to avoid NaNs in the padding for (auto & [buft, ctx] : ctx_map) { ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b44915e1a690..f83671ad8dc6 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -71,6 +71,7 @@ class llama_memory_recurrent : public llama_memory_i { uint32_t used = 0; // used cells (i.e. at least one seq_id) // number of recurrent-state snapshots per seq for rollback; tensors are widened to (1 + n_rs_seq) groups + // (ARC-LAB: with rpl on, S has 2 groups only - base and record banks; r keeps 1 + n_rs_seq) uint32_t n_rs_seq = 0; // per-seq rollback index @@ -80,7 +81,8 @@ class llama_memory_recurrent : public llama_memory_i { // ARC-LAB: rollback by replay for the gated delta net state (LLAMA_GDN_REPLAY, default on where it applies; // independent implementation of the technique described by llamADALL). The S rows of plane 0 hold the base - // state; plane 1 of each cell holds two banks of per-token update records (ggml_gated_delta_net_replay). + // state; plane 1 of each cell holds two banks of per-token update records (ggml_gated_delta_net_replay). S is + // allocated with these 2 planes only (the snapshot path allocates 1 + n_rs_seq). // Rolling back r tokens drops the last r pending records instead of selecting a snapshot plane. The conv // state (r_l) keeps its snapshot planes and rs_idx. bool rpl = false; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ca8e69085183..93213085790a 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -578,6 +578,8 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // base state in plane 0 plus the update records of the last K-1 tokens in plane 1 (two banks: this pass // reads one, its records go to the other) instead of K snapshot planes; see ggml_gated_delta_net_replay. GGML_ASSERT(state_rows == nullptr && n_seqs == 1 && g->ne[0] == 1); + // the memory allocates S with planes 0 and 1 only in replay mode (llama_memory_recurrent::rpl) + GGML_ASSERT(ssm_states_all->ne[1] == 2 * (int64_t) mem_size); const int64_t R = H_v * (2 * S_v + 1); const size_t row_size = ssm_states_all->nb[1]; const size_t rec_size = ggml_row_size(ssm_states_all->type, R * (K - 1)); From 9eec75a16ab3ab9ac054aa5650ddbf09d95da365 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:28:02 +1100 Subject: [PATCH 11/24] sycl: GGML_SYCL_OP_TIMING matches the MTP draft graphs; mat-muls by weight type The draft-size test counted the Hadamard transform mat-muls, whose src1 holds each row as 1024-wide blocks (5 per 5120 row, 17 per FFN row), so every draft graph failed the <= 8 rows check and nothing was timed. Only weight mat-muls are checked now. A one-time "op-timing: active" line confirms the hook engaged, and mat-muls are reported per weight type (Q6_K / Q8_0 / PTQ1_0 head weights, the Hadamard transforms apart). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 24424cd538a23f99f3758f9dfd56b5fda6fb5ea6 on Torchit1/llama.cpp arc-b580) --- ggml/src/ggml-sycl/ggml-sycl.cpp | 32 +++++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index dfd77750deb8..5ea6f2e7aaea 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6271,13 +6271,31 @@ struct sycl_op_timing { bool mm = false; for (int i = 0; i < g->n_nodes; ++i) { const ggml_tensor * n = g->nodes[i]; - if (n->op == GGML_OP_MUL_MAT) { + // weight mat-muls only: a Hadamard transform reshapes each row into 1024-wide blocks (5..17 per row) + if (n->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(n, 1) != GGML_HINT_SRC0_IS_HADAMARD) { mm = true; if (n->src[1]->ne[1] * n->src[1]->ne[2] > 8) return false; } } + if (mm) { + static bool logged = false; + if (!logged) { + logged = true; + fprintf(stderr, "op-timing: active (first draft-sized graph: %d nodes)\n", g->n_nodes); + } + } return mm; } + // op label: mat-muls by weight type (and the Hadamard transforms apart) + static std::string label(const ggml_tensor * n) { + if (n->op == GGML_OP_MUL_MAT) { + if (ggml_get_op_params_i32(n, 1) == GGML_HINT_SRC0_IS_HADAMARD) { + return "MUL_MAT hadamard"; + } + return std::string("MUL_MAT ") + ggml_type_name(n->src[0]->type); + } + return ggml_op_desc(n); + } void report() { fprintf(stderr, "op-timing: %d draft-sized graphs, per graph: %.3f ms serialised over %.1f kernel groups (%.0f nodes)\n", graphs, total / graphs, (double) groups / graphs, (double) nodes / graphs); @@ -7290,11 +7308,11 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc const char * ht_op = nullptr; const bool ot_on = g_ot.on && sycl_op_timing::draft_sized(cgraph); sycl_op_timing::clk::time_point ot_prev{}; - const char * ot_op = nullptr; - auto ot_mark = [&](const char * next) { // drain the queue, charge the time since the last mark to the op that ran + std::string ot_op; + auto ot_mark = [&](const std::string & next) { // drain the queue, charge the time since the last mark to the op that ran SYCL_CHECK(CHECK_TRY_ERROR(sycl_ctx->stream()->wait())); const auto now = sycl_op_timing::clk::now(); - if (ot_op) { + if (!ot_op.empty()) { const double ms = std::chrono::duration(now - ot_prev).count(); auto & a = g_ot.per_op[ot_op]; a.first += ms; @@ -7306,7 +7324,7 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc ot_op = next; }; if (ot_on) { - ot_mark(nullptr); // inputs and earlier work are not charged + ot_mark(std::string()); // inputs and earlier work are not charged g_ot.nodes += cgraph->n_nodes; } for (int i = 0; i < cgraph->n_nodes; i++) { @@ -7334,7 +7352,7 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc continue; } if (ot_on) { - ot_mark(ggml_op_desc(node)); + ot_mark(sycl_op_timing::label(node)); } const int nodes_to_skip = ggml_sycl_fuse(*sycl_ctx, cgraph, i); @@ -7476,7 +7494,7 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc GGML_ASSERT(ok); } if (ot_on) { - ot_mark(nullptr); + ot_mark(std::string()); if (++g_ot.graphs == 64) { g_ot.report(); } From f142204f155db65b906b101a1407c87ebeba79c8 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:44:47 +1100 Subject: [PATCH 12/24] mtp chain: per-step input slices made once, up front, from device copies (LLAMA_MTP_CHAIN_TIGHT bit 1) Feature B round 2, idea 1 (independent implementation of the technique described by llamADALL). The chain graph sliced its packed host inputs (positions, output rows, K/V indices, mask) with per-step views built lazily inside each step. The scheduler makes every view of a host input its own split input: ~20 blocking host->device copies per draft (5 per step), each behind a backend sync, before the device split starts. And the views landed between the step's compute nodes. Now each packed input gets one whole-tensor copy (a device memcpy) pinned to the head's device, all per-step slices are views of those copies, and the host token lookup, the copies and every slice are added to the graph before the first compute node. Each step's compute nodes then run uninterrupted, as in the single-step draft graph; the device split keeps ~9 inputs instead of ~24. Same values in the same kernels: drafts unchanged. LLAMA_MTP_CHAIN_TIGHT is a bit mask (default 1 here); 0 restores the previous layout exactly. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 0e1ba9c6ce7b46224c7135beb596b98f06b6147c on Torchit1/llama.cpp arc-b580) --- src/models/qwen35.cpp | 88 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 77 insertions(+), 11 deletions(-) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 0e81f3dba9a9..55dba5f19451 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -1030,24 +1030,90 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - ggml_tensor * tok_embd = unrotate(ggml_get_rows(ctx0, tok_embd_w, ggml_view_1d(ctx0, tokens_all, n_lead, 0))); + // ARC-LAB: LLAMA_MTP_CHAIN_TIGHT (bit mask, default all on; 0 = the first chain layout): + // 1: every per-step input slice is made once, up front, from ONE device copy of each packed input. A view of a + // host input is its own scheduler split input (one blocking host->device copy per view and step), and a view + // built lazily lands between the step's compute nodes; here the device split starts with 4-5 copy kernels and + // all the slices, then each step's compute nodes run uninterrupted, as in the single-step draft graph. + static const int tight = [] { + const char * e = getenv("LLAMA_MTP_CHAIN_TIGHT"); + return e ? atoi(e) : 1; + }(); + + ggml_tensor * tok_rows = ggml_get_rows(ctx0, tok_embd_w, ggml_view_1d(ctx0, tokens_all, n_lead, 0)); ggml_tensor * h_embd = ggml_view_2d(ctx0, h_all, h_all->ne[0], n_lead, h_all->nb[1], 0); + ggml_tensor * k_all = inp_attn->get_k_idxs(); + ggml_tensor * v_all = inp_attn->get_v_idxs(); + ggml_tensor * mask_all = inp_attn->get_kq_mask(); + + // the device holding the head's weights (none: CPU-only, the slices stay views of the host inputs) + ggml_backend_t dev = nullptr; + if ((tight & 1) && sched && layer.wq->buffer) { + ggml_backend_buffer_type_t wbuft = ggml_backend_buffer_get_type(layer.wq->buffer); + for (int b = 0; b < ggml_backend_sched_get_n_backends(sched) && !dev; ++b) { + ggml_backend_t be = ggml_backend_sched_get_backend(sched, b); + if (be != backend_cpu && ggml_backend_supports_buft(be, wbuft)) { + dev = be; + } + } + } + if (tight & 1) { + ggml_build_forward_expand(gf, tok_rows); // the host-table lookup first: one host split, then one device split + } + if (dev) { + // whole-tensor copies (a device memcpy each), pinned to the device so their views need no upload + auto on_dev = [&](ggml_tensor * t, const char * name) { + ggml_tensor * c = ggml_cont(ctx0, t); + ggml_set_name(c, name); + if (!ggml_backend_supports_op(dev, c)) { + return t; + } + ggml_backend_sched_set_tensor_backend(sched, c, dev); + ggml_build_forward_expand(gf, c); + return c; + }; + pos_all = on_dev(pos_all, "mtp_chain_pos_dev"); + out_all = on_dev(out_all, "mtp_chain_out_dev"); + k_all = on_dev(k_all, "mtp_chain_k_idxs_dev"); + v_all = on_dev(v_all, "mtp_chain_v_idxs_dev"); + mask_all = on_dev(mask_all, "mtp_chain_kq_mask_dev"); + } + + // each step's slices: positions, output row, KV indices and mask rows, sized as a decode of rows [0, r0 + nr) + // would be (created here; added to the graph here with bit 1, else where the step first reads them) + struct step_inputs { + ggml_tensor * pos, * out, * k_idxs, * v_idxs, * kq_mask; + }; + std::vector sin(n_steps); + for (int64_t s = 0; s < n_steps; ++s) { + const auto & st = steps[s]; + const int64_t nr = st.nr; + auto & si = sin[s]; + si.pos = ggml_view_1d(ctx0, pos_all, nr * n_pos, st.r0 * n_pos * ggml_element_size(pos_all)); + si.out = ggml_view_1d(ctx0, out_all, 1, s * ggml_element_size(out_all)); + si.k_idxs = ggml_view_1d(ctx0, k_all, nr, st.r0 * ggml_element_size(k_all)); + si.v_idxs = ggml_view_1d(ctx0, v_all, nr, st.r0 * ggml_element_size(v_all)); + si.kq_mask = ggml_cont(ctx0, ggml_view_2d(ctx0, mask_all, st.n_kv, nr, mask_all->nb[1], st.r0 * mask_all->nb[1])); + if (tight & 1) { + for (ggml_tensor * t : { si.pos, si.out, si.k_idxs, si.v_idxs, si.kq_mask }) { + ggml_build_forward_expand(gf, t); + } + } + } + + ggml_tensor * tok_embd = unrotate(tok_rows); for (int64_t s = 0; s < n_steps; ++s) { const auto & st = steps[s]; const int64_t nr = st.nr; cb(tok_embd, "mtp_tok_embd", il); - // this step's slice of the ubatch's KV indices and mask, sized as a decode of rows [0, r0 + nr) would be - const uint32_t n_kv_s = st.n_kv; - ggml_tensor * st_pos = ggml_view_1d(ctx0, pos_all, nr * n_pos, st.r0 * n_pos * ggml_element_size(pos_all)); - ggml_tensor * st_out = ggml_view_1d(ctx0, out_all, 1, s * ggml_element_size(out_all)); - ggml_tensor * k_idxs = inp_attn->get_k_idxs(); - ggml_tensor * v_idxs = inp_attn->get_v_idxs(); - k_idxs = ggml_view_1d(ctx0, k_idxs, nr, st.r0 * ggml_element_size(k_idxs)); - v_idxs = ggml_view_1d(ctx0, v_idxs, nr, st.r0 * ggml_element_size(v_idxs)); - ggml_tensor * kq_mask = inp_attn->get_kq_mask(); - kq_mask = ggml_cont(ctx0, ggml_view_2d(ctx0, kq_mask, n_kv_s, nr, kq_mask->nb[1], st.r0 * kq_mask->nb[1])); + const uint32_t n_kv_s = st.n_kv; + ggml_tensor * st_pos = sin[s].pos; + ggml_tensor * st_out = sin[s].out; + ggml_tensor * k_idxs = sin[s].k_idxs; + ggml_tensor * v_idxs = sin[s].v_idxs; + ggml_tensor * kq_mask = sin[s].kq_mask; ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); cb(h_norm, "mtp_hnorm", il); From 075c4b212bc2be14ee0b773f0702f34f473af057 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:45:03 +1100 Subject: [PATCH 13/24] mtp chain: copy a step's mask rows only when the slice is not contiguous (LLAMA_MTP_CHAIN_TIGHT bit 2) Feature B round 2, idea 2 (independent implementation of the technique described by llamADALL). Every chain step ran a CONT on its mask slice. A 1-row step's slice is one row of the mask and already contiguous, so flash attention reads it in place (the decode kernels take the mask row stride from nb[1]; with one row it is never used). Only step 0 with catch-up rows, narrower than the full mask (its KV length is a prefix), still needs the copy. Saves n-1 copy kernels per draft (n = 4: 3). Same mask values: drafts unchanged. Default LLAMA_MTP_CHAIN_TIGHT is now 3. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit c9c5053ac86184909e77787e396a923f090fc640 on Torchit1/llama.cpp arc-b580) --- src/models/qwen35.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 55dba5f19451..fb883e8ade53 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -1035,9 +1035,11 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { // host input is its own scheduler split input (one blocking host->device copy per view and step), and a view // built lazily lands between the step's compute nodes; here the device split starts with 4-5 copy kernels and // all the slices, then each step's compute nodes run uninterrupted, as in the single-step draft graph. + // 2: a step's mask rows are copied only when the slice is not contiguous: a 1-row step's row is already + // contiguous, only a multi-row step 0 narrower than the full mask needs the copy (one copy kernel, not n). static const int tight = [] { const char * e = getenv("LLAMA_MTP_CHAIN_TIGHT"); - return e ? atoi(e) : 1; + return e ? atoi(e) : 3; }(); ggml_tensor * tok_rows = ggml_get_rows(ctx0, tok_embd_w, ggml_view_1d(ctx0, tokens_all, n_lead, 0)); @@ -1093,7 +1095,10 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { si.out = ggml_view_1d(ctx0, out_all, 1, s * ggml_element_size(out_all)); si.k_idxs = ggml_view_1d(ctx0, k_all, nr, st.r0 * ggml_element_size(k_all)); si.v_idxs = ggml_view_1d(ctx0, v_all, nr, st.r0 * ggml_element_size(v_all)); - si.kq_mask = ggml_cont(ctx0, ggml_view_2d(ctx0, mask_all, st.n_kv, nr, mask_all->nb[1], st.r0 * mask_all->nb[1])); + si.kq_mask = ggml_view_2d(ctx0, mask_all, st.n_kv, nr, mask_all->nb[1], st.r0 * mask_all->nb[1]); + if (!(tight & 2) || !ggml_is_contiguous(si.kq_mask)) { + si.kq_mask = ggml_cont(ctx0, si.kq_mask); + } if (tight & 1) { for (ggml_tensor * t : { si.pos, si.out, si.k_idxs, si.v_idxs, si.kq_mask }) { ggml_build_forward_expand(gf, t); From 567ba0b893e1e73ced20f01632326ed3e2e14a5b Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:45:55 +1100 Subject: [PATCH 14/24] mtp chain: gather the output row only for a multi-row step (LLAMA_MTP_CHAIN_TIGHT bit 4) Feature B round 2, idea 3 (independent implementation of the technique described by llamADALL). Every chain step ran GET_ROWS(h_nextn, out_idx) to pick its output row. Only step 0 (catch-up rows + the anchor) has a choice; a 1-row step's row is already the output. Now only steps with more than one row gather, and the output-index input holds only their entries (none when step 0 has no catch-up rows: no out_idx input at all). Side effect, on purpose: without the gather, the draft head's input chain of steps 1.. is ADD -> RMS_NORM -> MUL(norm w) -> MUL(signs) -> RESHAPE -> MUL_MAT(Hadamard), which ggml_sycl_try_add_norm_rot folds into one kernel (residual add + head norm + rotation). The loop never takes that path for the head (its out_ids gather sits in between). The fold is documented as the same arithmetic in the same order (byte-identical in the trunk), but it is new for drafts: first suspect if drafts ever differ (bit 4 off isolates it). Default LLAMA_MTP_CHAIN_TIGHT is now 7. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 9d3313b24d9f98fc715486292d7b1e09a14aee17 on Torchit1/llama.cpp arc-b580) --- src/models/qwen35.cpp | 75 +++++++++++++++++++++++++++---------------- 1 file changed, 47 insertions(+), 28 deletions(-) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index fb883e8ade53..affb021fcd3c 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -620,7 +620,7 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_ffn(ggml_tensor * cur, cons // ARC-LAB: per-step inputs of a chained MTP draft graph (graph_mtp::build_chain), packed into two tensors so a draft // uploads them in two writes: pos holds each step's positions in the layout build_inp_pos() gives a decode of just -// that step's rows (step after step), out_idx each step's output row. The graph reads them through per-step views. +// that step's rows (step after step), out_idx the output row of each step that gathers one. The graph reads them through per-step views. class llm_graph_input_mtp_chain : public llm_graph_input_i { public: explicit llm_graph_input_mtp_chain(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {} @@ -629,16 +629,17 @@ class llm_graph_input_mtp_chain : public llm_graph_input_i { int64_t r0; // first ubatch row of the step int64_t nr; // rows uint32_t n_kv; // KV length the step attends over (baked into the graph's views) + bool gather; // picks its output row with a GET_ROWS (else its single row is the output) }; std::vector steps; int64_t n_tokens = 0; ggml_tensor * pos = nullptr; // I32 [n_tokens * n_pos_per_embd] - ggml_tensor * out_idx = nullptr; // I32 [n_steps], step s: nr - 1 (the step's last row produces its draft) + ggml_tensor * out_idx = nullptr; // I32 [gathering steps]: nr - 1 (the step's last row produces its draft); null: none void set_input(const llama_ubatch * ubatch) override { GGML_ASSERT(ubatch->n_tokens == n_tokens); std::vector v((size_t) n_tokens * n_pos_per_embd, 0); - std::vector o(steps.size()); + std::vector o; size_t off = 0; for (size_t s = 0; s < steps.size(); ++s) { const auto & st = steps[s]; @@ -656,10 +657,14 @@ class llm_graph_input_mtp_chain : public llm_graph_input_i { } } off += (size_t) st.nr * n_pos_per_embd; - o[s] = (int32_t) st.nr - 1; + if (st.gather) { + o.push_back((int32_t) st.nr - 1); + } } ggml_backend_tensor_set(pos, v.data(), 0, v.size() * sizeof(int32_t)); - ggml_backend_tensor_set(out_idx, o.data(), 0, o.size() * sizeof(int32_t)); + if (out_idx) { + ggml_backend_tensor_set(out_idx, o.data(), 0, o.size() * sizeof(int32_t)); + } } // reusable when the batch has the same rows and every step would attend over the same KV length (the rest of the @@ -972,20 +977,39 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { auto * inp_attn = build_attn_inp_kv(); const llama_kv_cache_context * kvc = inp_attn->mctx; + // ARC-LAB: LLAMA_MTP_CHAIN_TIGHT (bit mask, default all on; 0 = the first chain layout): + // 1: every per-step input slice is made once, up front, from ONE device copy of each packed input. A view of a + // host input is its own scheduler split input (one blocking host->device copy per view and step), and a view + // built lazily lands between the step's compute nodes; here the device split starts with 4-5 copy kernels and + // all the slices, then each step's compute nodes run uninterrupted, as in the single-step draft graph. + // 2: a step's mask rows are copied only when the slice is not contiguous: a 1-row step's row is already + // contiguous, only a multi-row step 0 narrower than the full mask needs the copy (one copy kernel, not n). + // 4: a step picks its output row with a GET_ROWS only if it has several rows (step 0 with catch-up rows); a + // 1-row step's row is already the output (no gather, no output-index slice). + static const int tight = [] { + const char * e = getenv("LLAMA_MTP_CHAIN_TIGHT"); + return e ? atoi(e) : 7; + }(); + const uint32_t n_pos = hparams.n_pos_per_embd(); auto chain_inp = std::make_unique(n_pos); chain_inp->n_tokens = n_tokens; chain_inp->pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens * n_pos); - chain_inp->out_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_steps); ggml_set_input(chain_inp->pos); - ggml_set_input(chain_inp->out_idx); + int64_t n_gather = 0; for (int64_t s = 0; s < n_steps; ++s) { llm_graph_input_mtp_chain::step st; - st.r0 = s == 0 ? 0 : n_lead - 1 + s; - st.nr = s == 0 ? n_lead : 1; - st.n_kv = kvc->get_n_kv_prefix((uint32_t) (st.r0 + st.nr)); + st.r0 = s == 0 ? 0 : n_lead - 1 + s; + st.nr = s == 0 ? n_lead : 1; + st.n_kv = kvc->get_n_kv_prefix((uint32_t) (st.r0 + st.nr)); + st.gather = !(tight & 4) || st.nr > 1; + n_gather += st.gather; chain_inp->steps.push_back(st); } + if (n_gather > 0) { + chain_inp->out_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_gather); + ggml_set_input(chain_inp->out_idx); + } const auto steps = chain_inp->steps; ggml_tensor * pos_all = chain_inp->pos; ggml_tensor * out_all = chain_inp->out_idx; @@ -1030,18 +1054,6 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - // ARC-LAB: LLAMA_MTP_CHAIN_TIGHT (bit mask, default all on; 0 = the first chain layout): - // 1: every per-step input slice is made once, up front, from ONE device copy of each packed input. A view of a - // host input is its own scheduler split input (one blocking host->device copy per view and step), and a view - // built lazily lands between the step's compute nodes; here the device split starts with 4-5 copy kernels and - // all the slices, then each step's compute nodes run uninterrupted, as in the single-step draft graph. - // 2: a step's mask rows are copied only when the slice is not contiguous: a 1-row step's row is already - // contiguous, only a multi-row step 0 narrower than the full mask needs the copy (one copy kernel, not n). - static const int tight = [] { - const char * e = getenv("LLAMA_MTP_CHAIN_TIGHT"); - return e ? atoi(e) : 3; - }(); - ggml_tensor * tok_rows = ggml_get_rows(ctx0, tok_embd_w, ggml_view_1d(ctx0, tokens_all, n_lead, 0)); ggml_tensor * h_embd = ggml_view_2d(ctx0, h_all, h_all->ne[0], n_lead, h_all->nb[1], 0); ggml_tensor * k_all = inp_attn->get_k_idxs(); @@ -1075,7 +1087,9 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { return c; }; pos_all = on_dev(pos_all, "mtp_chain_pos_dev"); - out_all = on_dev(out_all, "mtp_chain_out_dev"); + if (n_gather > 1) { + out_all = on_dev(out_all, "mtp_chain_out_dev"); + } k_all = on_dev(k_all, "mtp_chain_k_idxs_dev"); v_all = on_dev(v_all, "mtp_chain_v_idxs_dev"); mask_all = on_dev(mask_all, "mtp_chain_kq_mask_dev"); @@ -1087,12 +1101,12 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { ggml_tensor * pos, * out, * k_idxs, * v_idxs, * kq_mask; }; std::vector sin(n_steps); - for (int64_t s = 0; s < n_steps; ++s) { + for (int64_t s = 0, g = 0; s < n_steps; ++s) { const auto & st = steps[s]; const int64_t nr = st.nr; auto & si = sin[s]; si.pos = ggml_view_1d(ctx0, pos_all, nr * n_pos, st.r0 * n_pos * ggml_element_size(pos_all)); - si.out = ggml_view_1d(ctx0, out_all, 1, s * ggml_element_size(out_all)); + si.out = !st.gather ? nullptr : n_gather == 1 ? out_all : ggml_view_1d(ctx0, out_all, 1, g++ * ggml_element_size(out_all)); si.k_idxs = ggml_view_1d(ctx0, k_all, nr, st.r0 * ggml_element_size(k_all)); si.v_idxs = ggml_view_1d(ctx0, v_all, nr, st.r0 * ggml_element_size(v_all)); si.kq_mask = ggml_view_2d(ctx0, mask_all, st.n_kv, nr, mask_all->nb[1], st.r0 * mask_all->nb[1]); @@ -1101,7 +1115,9 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { } if (tight & 1) { for (ggml_tensor * t : { si.pos, si.out, si.k_idxs, si.v_idxs, si.kq_mask }) { - ggml_build_forward_expand(gf, t); + if (t) { + ggml_build_forward_expand(gf, t); + } } } } @@ -1221,8 +1237,11 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { cb(cur, "h_nextn", -1); // the step's output row (masked h_nextn: the next step's hidden input) - ggml_tensor * h_out = ggml_get_rows(ctx0, cur, st_out); - cb(h_out, "mtp_shared_head_norm", -1); + ggml_tensor * h_out = cur; + if (st_out) { + h_out = ggml_get_rows(ctx0, cur, st_out); + cb(h_out, "mtp_shared_head_norm", -1); + } // greedy draft: argmax of the head's logits. With a draft head that is an index into its vocabulary (the // -inf scatter of the single-step graph only maps it to a token id, which the context does on the host) From 31ffa174c0f432df7baf2f41ba3f2b5d4684d1cb Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:53:28 +1100 Subject: [PATCH 15/24] mtp chain: report graph reuse and process_ubatch host time per chain decode (LLAMA_MTP_DRAFT_TIMING) Diagnostics for the question whether the cached chain graph is reused in steady-state decoding (the chain's llama_decode took 1.9 ms at short context and 3.3-3.5 ms at 48K/115K). With LLAMA_MTP_DRAFT_TIMING=1, every 64 chain decodes print to stderr: mtp-chain-decode: 64 chain decodes, graph reused R, rebuilt M ( n, ...); per decode: reuse check or build+alloc X ms, set_inputs Y ms, compute submit Z ms Miss reasons come from llm_graph_result::reuse_miss_reason(): "new" (shape never built), "params" (allow_reuse), or the type name of the first input that refuses (attn_kv = mask n_kv, mtp_chain = a step's KV length). draft-timing "decode" minus X+Y+Z is the llama_decode time outside process_ubatch (batch split, slot search). No change without the env. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 505f7b8cd1ed8da05faf8fa9794654ffe289eb45 on Torchit1/llama.cpp arc-b580) --- src/llama-context.cpp | 41 +++++++++++++++++++++++++++++++++++++++++ src/llama-context.h | 10 ++++++++++ src/llama-graph.cpp | 17 +++++++++++++++++ src/llama-graph.h | 5 +++++ 4 files changed, 73 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index eb4f6ea76edc..a32df588d17e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1452,6 +1452,19 @@ llama_context::mtp_chain_entry * llama_context::mtp_chain_entry_for(uint32_t n_t return &(mtp_chain_cache[key] = std::move(e)); } +void llama_context::mtp_chain_reuse_stats::report() { + std::string why; + for (const auto & [k, v] : miss) { + why += (why.empty() ? "" : ", ") + k + " " + std::to_string(v); + } + fprintf(stderr, "mtp-chain-decode: %d chain decodes, graph reused %d, rebuilt %d (%s); per decode: reuse check or " + "build+alloc %.3f ms, set_inputs %.3f ms, compute submit %.3f ms\n", n, reused, n - reused, + why.empty() ? "-" : why.c_str(), build_ms / n, set_ms / n, compute_ms / n); + n = reused = 0; + build_ms = set_ms = compute_ms = 0; + miss.clear(); +} + const int32_t * llama_context::get_mtp_chain_ids() { synchronize(); @@ -1695,6 +1708,13 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype); + static const bool chain_stats_on = [] { + const char * e = getenv("LLAMA_MTP_DRAFT_TIMING"); + return e && atoi(e) != 0; + }(); + const bool chain_stats = chain_stats_on && cparams.mtp_chain_n > 0; + const int64_t t_cs0 = chain_stats ? ggml_time_us() : 0; + if (!graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); @@ -1706,7 +1726,14 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } n_reused++; + if (chain_stats) { + mtp_chain_rs.reused++; + } } else { + if (chain_stats) { + const std::string why = graph_reuse_disable ? std::string("disabled") : res->reuse_miss_reason(gparams); + mtp_chain_rs.miss[why.empty() ? std::string("?") : why]++; + } res->reset(); ggml_backend_sched_reset(sched.get()); @@ -1731,6 +1758,8 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } } + const int64_t t_cs1 = chain_stats ? ggml_time_us() : 0; + // set the input data for the input tensors { //const auto t_start_us = ggml_time_us(); @@ -1741,7 +1770,19 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll //LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); } + const int64_t t_cs2 = chain_stats ? ggml_time_us() : 0; + const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); + + if (chain_stats) { + const int64_t t_cs3 = ggml_time_us(); + mtp_chain_rs.build_ms += (t_cs1 - t_cs0) / 1000.0; + mtp_chain_rs.set_ms += (t_cs2 - t_cs1) / 1000.0; + mtp_chain_rs.compute_ms += (t_cs3 - t_cs2) / 1000.0; + if (++mtp_chain_rs.n == 64) { + mtp_chain_rs.report(); + } + } if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; diff --git a/src/llama-context.h b/src/llama-context.h index 16f308b140b2..3fcb20cb1991 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -430,6 +430,16 @@ struct llama_context { mtp_chain_entry * mtp_chain_entry_for(uint32_t n_tokens); uint32_t mtp_chain_n_out = 0; + // ARC-LAB: LLAMA_MTP_DRAFT_TIMING=1: per chain decode, whether its cached graph was reused (else why not) and the + // host time of process_ubatch's parts; printed to stderr every 64 chain decodes + struct mtp_chain_reuse_stats { + int n = 0, reused = 0; + double build_ms = 0, set_ms = 0, compute_ms = 0; + std::map miss; + void report(); + }; + mtp_chain_reuse_stats mtp_chain_rs; + // host buffer for the model output (logits and embeddings) ggml_backend_buffer_ptr buf_output; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 44a31198b323..75d3717667a3 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include // dedup helpers @@ -1489,6 +1490,22 @@ bool llm_graph_result::can_reuse(const llm_graph_params & params) { return res; } +std::string llm_graph_result::reuse_miss_reason(const llm_graph_params & params) { + if (inputs.empty()) { + return "new"; + } + if (!this->params.allow_reuse(params)) { + return "params"; + } + for (auto & input : inputs) { + if (!input->can_reuse(params)) { + const llm_graph_input_i & in = *input; + return typeid(in).name(); + } + } + return ""; +} + llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) { inputs.emplace_back(std::move(input)); return inputs.back().get(); diff --git a/src/llama-graph.h b/src/llama-graph.h index 7b1fd48989f8..e1be43b59627 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -11,6 +11,7 @@ #include #include #include +#include #include struct ggml_cgraph; @@ -1025,6 +1026,10 @@ class llm_graph_result { // return true if the graph was updated and can be reused bool can_reuse(const llm_graph_params & params); + // ARC-LAB: why can_reuse(params) refuses (diagnostics, LLAMA_MTP_DRAFT_TIMING): "new" (never built), "params" + // (allow_reuse), or the type name of the first input that refuses; "" if it would not refuse + std::string reuse_miss_reason(const llm_graph_params & params); + llm_graph_input_i * add_input(llm_graph_input_ptr input); void add_fused_node(llm_graph_fused_node result); From 0687925eadf6a58d7f8d43102b9885170952f605 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:28:56 +1100 Subject: [PATCH 16/24] sycl: GGML_SYCL_SYNC_TRACE=1 attributes every synchronize and blocking upload to its cause Each ggml_backend_sycl_synchronize and blocking buffer_set_tensor is counted under a cause, printed per 64 graphs of >= 2000 nodes (the target's graphs, so ~ per MTP cycle) with the totals per such graph: - the scheduler tags its copies (new ggml-base hook ggml_sync_trace_set_cause): "sched: split input " (computed inputs marked), "sched: backend switch"; - anything else is labelled with the first two named callers outside libggml (backtrace + dladdr), e.g. "llama_get_logits_ith <- common_sampler_sample". Lines: "sync-trace: N big graphs (+M smaller), per big graph: S synchronizes, U blocking uploads (B bytes)", then one line per cause. Off by default; costs a backtrace per sync when on (host timing is then not representative). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit b13c1b921e696e1d459d14df257b0f246d1fe183 on Torchit1/llama.cpp arc-b580) --- ggml/src/ggml-backend-impl.h | 5 ++ ggml/src/ggml-backend.cpp | 33 ++++++++++ ggml/src/ggml-sycl/ggml-sycl.cpp | 100 +++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 9c56ec30c5f1..cd94ce4ad2d3 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -270,6 +270,11 @@ extern "C" { # define GGML_BACKEND_DL_SCORE_IMPL(score_fn) #endif + // ARC-LAB: sync/upload attribution (GGML_SYCL_SYNC_TRACE): the scheduler tags the blocking copies it makes with a + // cause (e.g. "split input "); a backend that counts synchronizes / uploads reads it. Thread-local; null = none. + GGML_API void ggml_sync_trace_set_cause(const char * cause); + GGML_API const char * ggml_sync_trace_get_cause(void); + #ifdef __cplusplus } #endif diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 3d6310f3ffea..8dbe75408971 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1591,6 +1591,26 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { return true; } +// ARC-LAB: see ggml-backend-impl.h +static thread_local char g_sync_trace_cause[128]; +static thread_local bool g_sync_trace_has_cause = false; + +void ggml_sync_trace_set_cause(const char * cause) { + g_sync_trace_has_cause = cause != nullptr; + if (cause) { + snprintf(g_sync_trace_cause, sizeof(g_sync_trace_cause), "%s", cause); + } +} + +const char * ggml_sync_trace_get_cause(void) { + return g_sync_trace_has_cause ? g_sync_trace_cause : nullptr; +} + +static bool ggml_sync_trace_on() { + static const bool on = getenv("GGML_SYCL_SYNC_TRACE") && atoi(getenv("GGML_SYCL_SYNC_TRACE")) != 0; + return on; +} + static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); struct ggml_backend_sched_split * splits = sched->splits; @@ -1609,6 +1629,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // ensure the previous split's async work has completed before we start // this split, the allocator may have reused buffer regions across splits if (split->n_inputs == 0 && prev_backend_id >= 0 && prev_backend_id != split_backend_id) { + if (ggml_sync_trace_on()) { + ggml_sync_trace_set_cause("sched: backend switch"); + } if (sched->events[prev_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[prev_backend_id][sched->cur_copy]); } else { @@ -1621,6 +1644,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); struct ggml_tensor * input = split->inputs[input_id]; struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); + if (ggml_sync_trace_on()) { + char cause[128]; + snprintf(cause, sizeof(cause), "sched: split input %s%s", input->flags & GGML_TENSOR_FLAG_INPUT ? "" : "(computed) ", + input->name); + ggml_sync_trace_set_cause(cause); + } if (input->flags & GGML_TENSOR_FLAG_INPUT) { // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done @@ -1739,6 +1768,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + if (ggml_sync_trace_on()) { + ggml_sync_trace_set_cause(nullptr); + } + if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 5ea6f2e7aaea..403e9dc17ce6 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -36,6 +36,9 @@ #include #include #include +#include +#include +#include #include #include @@ -677,6 +680,94 @@ catch (sycl::exception const &exc) { std::exit(1); } +// ARC-LAB: GGML_SYCL_SYNC_TRACE=1: every ggml_backend_sycl_synchronize and blocking buffer_set_tensor is counted under its +// cause: the scheduler's tag (ggml_sync_trace_set_cause: split input , backend switch), else the first named caller +// outside libggml (e.g. llama_get_logits_ith). Printed per 64 graphs of >= 2000 nodes (target graphs: ~ per MTP cycle). +struct ggml_sycl_sync_trace { + const bool on = getenv("GGML_SYCL_SYNC_TRACE") && atoi(getenv("GGML_SYCL_SYNC_TRACE")) != 0; + std::mutex mtx; + int big = 0, small = 0; + std::map syncs; + std::map> uploads; + std::unordered_map names; // return address -> caller label + + std::string caller() { + const char * c = ggml_sync_trace_get_cause(); + if (c) { + return c; + } + void * fr[24]; + const int n = backtrace(fr, 24); + std::string label, last; + int named = 0; + for (int i = 1; i < n && named < 2; ++i) { + auto it = names.find(fr[i]); + std::string nm; + if (it != names.end()) { + nm = it->second; + } else { + Dl_info di{}; + if (dladdr(fr[i], &di) && di.dli_sname && di.dli_fname && !strstr(di.dli_fname, "libggml")) { + int st = 0; + char * dm = abi::__cxa_demangle(di.dli_sname, nullptr, nullptr, &st); + nm = dm && st == 0 ? std::string(dm) : std::string(di.dli_sname); + free(dm); + const size_t paren = nm.find('('); + if (paren != std::string::npos) { + nm.resize(paren); + } + } + names[fr[i]] = nm; + } + if (!nm.empty() && nm != last) { + label += (label.empty() ? "" : " <- ") + nm; + last = nm; + named++; + } + } + return label.empty() ? std::string("?") : label; + } + void note_sync() { + std::lock_guard lock(mtx); + syncs[caller()]++; + } + void note_upload(size_t bytes) { + std::lock_guard lock(mtx); + auto & u = uploads[caller()]; + u.first++; + u.second += bytes; + } + void note_graph(const ggml_cgraph * g) { + std::lock_guard lock(mtx); + if (g->n_nodes < 2000) { + small++; + return; + } + if (++big < 64) { + return; + } + long ns = 0, nu = 0; + size_t nb = 0; + for (auto & kv : syncs) ns += kv.second; + for (auto & kv : uploads) { nu += kv.second.first; nb += kv.second.second; } + fprintf(stderr, "sync-trace: %d big graphs (+%d smaller), per big graph: %.1f synchronizes, %.1f blocking uploads (%.0f bytes)\n", + big, small, (double) ns / big, (double) nu / big, (double) nb / big); + std::vector> v; + for (auto & kv : syncs) v.push_back({ kv.second, kv.first }); + std::sort(v.rbegin(), v.rend()); + for (auto & e : v) fprintf(stderr, "sync-trace: sync %6.2f %s\n", (double) e.first / big, e.second.c_str()); + std::vector> w; + for (auto & kv : uploads) w.push_back({ kv.second.first, kv.first }); + std::sort(w.rbegin(), w.rend()); + for (auto & e : w) fprintf(stderr, "sync-trace: upload %6.2f %8.0f B %s\n", (double) e.first / big, + (double) uploads[e.second].second / big, e.second.c_str()); + big = small = 0; + syncs.clear(); + uploads.clear(); + } +}; +static ggml_sycl_sync_trace g_sync_trace; + // ARC-LAB: per-token accounting of synchronous buffer_set_tensor calls, printed with the GGML_SYCL_HOST_TIMING report static double g_st_ms = 0; static long g_st_calls = 0; static size_t g_st_bytes = 0; static void ggml_sycl_note_set_tensor(size_t bytes, double ms) { g_st_ms += ms; g_st_calls++; g_st_bytes += bytes; } @@ -696,6 +787,9 @@ static void ggml_backend_sycl_buffer_set_tensor(ggml_backend_buffer_t buffer, size_t bytes; std::chrono::steady_clock::time_point t0 = std::chrono::steady_clock::now(); ~st_timer() { ggml_sycl_note_set_tensor(bytes, std::chrono::duration(std::chrono::steady_clock::now() - t0).count()); } } st_guard{ size }; + if (g_sync_trace.on) { + g_sync_trace.note_upload(size); + } // ARC-LAB: GGML_SYCL_FAST_SET_TENSOR=1: small writes (<= 4 MB, i.e. per-token inputs, not weights) skip the device-wide // wait and the malloc'd bounce buffer (the mmap workaround is for weights loaded from disk): one pinned staging copy and // one wait on this queue. Still synchronous for the caller. @@ -6317,6 +6411,9 @@ static void ggml_backend_sycl_synchronize(ggml_backend_t backend) try { ggml_backend_sycl_context * sycl_ctx = (ggml_backend_sycl_context *)backend->context; const queue_ptr stream = sycl_ctx->stream(sycl_ctx->device, 0); const auto ht_t0 = sycl_host_timing::clk::now(); + if (g_sync_trace.on) { + g_sync_trace.note_sync(); + } SYCL_CHECK(CHECK_TRY_ERROR((stream)->wait())); g_upload_off = 0; if (g_ht.on) { @@ -7323,6 +7420,9 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc ot_prev = now; ot_op = next; }; + if (g_sync_trace.on) { + g_sync_trace.note_graph(cgraph); + } if (ot_on) { ot_mark(std::string()); // inputs and earlier work are not charged g_ot.nodes += cgraph->n_nodes; From 0cc9909ca2505cb594ee97f29a7be96b4ebebf63 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:32:01 +1100 Subject: [PATCH 17/24] llama: skip a context synchronize when nothing was queued since the last one (LLAMA_SYNC_TRIM bit 1) Every result getter (llama_get_logits_ith, the sampled-token getters, llama_get_embeddings_nextn_ith, ...) calls llama_context::synchronize, and an MTP cycle calls several per row: the server samples each verified row, then the MTP drafter reads the target's h_nextn row by row. Each is a full device wait on the one in-order queue every context shares. Only the first after a decode waits for this context's work; a later one can only wait for the other context's work queued in between (e.g. the draft graph), which this context does not need. Now the context tracks whether it queued anything (graph compute and the async output reads, decode/encode) since its last synchronize, and skips the device wait if not. Not with pipeline parallelism (its copy rotation runs in the sync). Results are unchanged: they were complete already. LLAMA_SYNC_TRIM is a new bit mask (default 3; 0 = upstream behaviour). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 60ef575f3b84f421c8935358265bc6046e5192aa on Torchit1/llama.cpp arc-b580) --- src/llama-context.cpp | 22 ++++++++++++++++++++++ src/llama-context.h | 1 + src/llama-impl.h | 3 +++ 3 files changed, 26 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a32df588d17e..6d1ef386012d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -826,11 +826,29 @@ void llama_context::sched_reserve() { __func__, (t_end_us - t_start_us)/1000.0, ggml_backend_sched_get_n_copies(sched.get())); } +// ARC-LAB: LLAMA_SYNC_TRIM bit mask (default all on; 0 = upstream behaviour): 1 = skip a context synchronize when the +// context queued nothing since its last one, 2 = KV-cache rotation matrices kept on the device (llama-kv-cache.cpp) +int llama_sync_trim() { + static const int v = [] { + const char * e = getenv("LLAMA_SYNC_TRIM"); + return e ? atoi(e) : 3; + }(); + return v; +} + void llama_context::synchronize() { if (!sched) { return; } + // ARC-LAB: LLAMA_SYNC_TRIM bit 1 (default on): every getter (logits, sampled tokens, nextn rows, ...) synchronizes, + // and the device's one in-order queue is shared by all contexts. With nothing queued by this context since its + // last synchronize, its results are already complete: skip the device wait. + if ((llama_sync_trim() & 1) && !sync_pending && !cparams.pipeline_parallel) { + return; + } + sync_pending = false; + ggml_backend_sched_synchronize(sched.get()); // FIXME: if multiple single tokens are evaluated without a synchronization, @@ -1839,6 +1857,7 @@ int llama_context::encode(const llama_batch & batch_inp) { sched_reserve(); n_queued_tokens += n_tokens; + sync_pending = true; // reserve output buffer if (output_reserve(n_tokens) < n_tokens) { @@ -2145,6 +2164,7 @@ int llama_context::decode(const llama_batch & batch_inp) { t_compute_start_us = ggml_time_us(); } n_queued_tokens += n_tokens_all; + sync_pending = true; output_swaps.clear(); @@ -2992,6 +3012,7 @@ ggml_status llama_context::graph_compute( set_n_threads_fn.second(set_n_threads_fn.first, n_threads); } + sync_pending = true; auto status = ggml_backend_sched_graph_compute_async(sched.get(), gf); if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: ggml_backend_sched_graph_compute_async failed with error %d\n", __func__, status); @@ -3882,6 +3903,7 @@ void llama_context::opt_epoch_iter( const uint32_t n_tokens_all = balloc->get_n_tokens(); n_queued_tokens += n_tokens_all; + sync_pending = true; embd_seq.clear(); diff --git a/src/llama-context.h b/src/llama-context.h index 3fcb20cb1991..7ad06c4469c2 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -447,6 +447,7 @@ struct llama_context { std::map mem_storage; bool has_evaluated_once = false; + bool sync_pending = true; // ARC-LAB: work queued since the last synchronize (LLAMA_SYNC_TRIM) // env: LLAMA_GRAPH_REUSE_DISABLE bool graph_reuse_disable = false; diff --git a/src/llama-impl.h b/src/llama-impl.h index 4988b06d2ca0..997bb985c610 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -74,6 +74,9 @@ static inline ggml_tensor * llama_mul_mat_hadamard( return res; } +// ARC-LAB: LLAMA_SYNC_TRIM bit mask, see llama-context.cpp +int llama_sync_trim(); + struct time_meas { time_meas(int64_t & t_acc, bool disable = false); ~time_meas(); From 84dafe3c7c68917fb14460235b922e2d6b8ab975 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:32:02 +1100 Subject: [PATCH 18/24] kv-cache: keep the attention rotation matrices on the device (LLAMA_SYNC_TRIM bit 2) With a quantized KV cache every attention graph reads attn_inp_k_rot (256 x 256 f32 here) and attn_inp_v_rot (64 x 64) as graph inputs. They live in host memory, so each graph uploads them again: per graph 2 split inputs, i.e. 2 device synchronizes + 2 blocking copies (~272 KB), in the target's verify graph and in the MTP draft graph alike (upstream has a TODO for this). The cache now also allocates them once in its own device buffer (only when the cache has one non-host buffer type) and build_input_k_rot / v_rot return those tensors; set_input_k_rot / v_rot skip a non-input tensor. Same values (ggml_gen_hadamard, copied once): outputs identical. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 4ec75c3416006e44df935387b2ec514eb56f4e23 on Torchit1/llama.cpp arc-b580) --- src/llama-kv-cache.cpp | 41 +++++++++++++++++++++++++++++++++++++++++ src/llama-kv-cache.h | 6 ++++++ 2 files changed, 47 insertions(+) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index cb806c688166..2ac1dc5a01b9 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -361,6 +361,33 @@ llama_kv_cache::llama_kv_cache( ggml_gen_hadamard(tmp); } + + // ARC-LAB: keep them on the cache's device too (one buffer type only, and not host memory) + if ((llama_sync_trim() & 2) && ctxs_bufs.size() == 1 && !hparams.no_alloc) { + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(ctxs_bufs[0].second.get()); + if (!ggml_backend_buft_is_host(buft)) { + ggml_init_params params = { + /* .mem_size = */ attn_rot_hadamard.size()*ggml_tensor_overhead(), + /* .mem_buffer = */ nullptr, + /* .no_alloc = */ true, + }; + attn_rot_dev_ctx.reset(ggml_init(params)); + for (const auto & [n, m] : attn_rot_hadamard) { + ggml_tensor * t = ggml_new_tensor_2d(attn_rot_dev_ctx.get(), GGML_TYPE_F32, n, n); + ggml_format_name(t, "attn_rot_dev_%d", (int) n); + attn_rot_dev[n] = t; + } + attn_rot_dev_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(attn_rot_dev_ctx.get(), buft)); + if (attn_rot_dev_buf) { + ggml_backend_buffer_set_usage(attn_rot_dev_buf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + for (const auto & [n, t] : attn_rot_dev) { + ggml_backend_tensor_set(t, attn_rot_hadamard.at(n).data(), 0, ggml_nbytes(t)); + } + } else { + attn_rot_dev.clear(); + } + } + } } const char * LLAMA_KV_CACHE_DEBUG = getenv("LLAMA_KV_CACHE_DEBUG"); @@ -1654,6 +1681,10 @@ ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const { } while (n_embd_head_k_all % nrot == 0); nrot /= 2; + if (const auto it = attn_rot_dev.find(nrot); it != attn_rot_dev.end()) { + return it->second; + } + res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot); ggml_set_input(res); ggml_set_name(res, "attn_inp_k_rot"); @@ -1674,6 +1705,10 @@ ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const { //} while (hparams.n_embd_head_v() % nrot == 0); //nrot /= 2; + if (const auto it = attn_rot_dev.find(nrot); it != attn_rot_dev.end()) { + return it->second; + } + res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot); ggml_set_input(res); ggml_set_name(res, "attn_inp_v_rot"); @@ -2012,6 +2047,9 @@ void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch } void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const { + if (!(dst->flags & GGML_TENSOR_FLAG_INPUT)) { + return; // the device copy (attn_rot_dev), already filled + } GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); const auto n_rot = dst->ne[0]; @@ -2021,6 +2059,9 @@ void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const { } void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const { + if (!(dst->flags & GGML_TENSOR_FLAG_INPUT)) { + return; // the device copy (attn_rot_dev), already filled + } GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); const auto n_rot = dst->ne[0]; diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 05add1b853c1..7f3d91bec362 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -280,6 +280,12 @@ class llama_kv_cache : public llama_memory_i { // pre-computed hadamard martrices std::unordered_map> attn_rot_hadamard; + // ARC-LAB: device copies of attn_rot_hadamard (LLAMA_SYNC_TRIM bit 2): graphs read these instead of an input that + // is uploaded (one sync + one blocking copy each) on every decode + ggml_context_ptr attn_rot_dev_ctx; + ggml_backend_buffer_ptr attn_rot_dev_buf; + std::unordered_map attn_rot_dev; + // env: LLAMA_KV_CACHE_DEBUG int debug = 0; From 880cc602a51fb67a846401e379cb0fdc4b4c8d5a Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:32:02 +1100 Subject: [PATCH 19/24] mtp chain: one index upload fewer twice over (LLAMA_MTP_CHAIN_TIGHT bit 8) Feature B round 2 (independent implementation of the technique described by llamADALL). Two chain inputs were uploaded separately though they need not be: - the output-row index (step 0 with catch-up rows) now rides at the end of the position input: one upload, sliced on the device; - V is stored with the K indices: for a non-transposed V cache set_input_v_idxs writes exactly what set_input_k_idxs writes, so the V index input is left out of the graph (llm_graph_input_attn_kv::set_input now skips an unallocated V index tensor). Per chain decode -1..2 split inputs (sync + blocking copy each). Same values in the same kernels: drafts unchanged. Default LLAMA_MTP_CHAIN_TIGHT is now 15. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit f3aa69d6a2c4ac3a3479576bf1d9fcada08eca44 on Torchit1/llama.cpp arc-b580) --- src/llama-graph.cpp | 5 ++++- src/models/qwen35.cpp | 34 +++++++++++++++++++++++----------- 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 75d3717667a3..2020398f89cb 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -569,7 +569,10 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch); - mctx->set_input_v_idxs(self_v_idxs, ubatch); + // ARC-LAB: left unallocated when the graph stores V with the K indices (MTP chain, non-transposed V) + if (self_v_idxs && self_v_idxs->buffer) { + mctx->set_input_v_idxs(self_v_idxs, ubatch); + } // the mask is left unallocated when the graph only stores K/V without attending // (e.g. DFlash's KV-injection pass) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index affb021fcd3c..f2e95cf41b6c 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -635,6 +635,7 @@ class llm_graph_input_mtp_chain : public llm_graph_input_i { int64_t n_tokens = 0; ggml_tensor * pos = nullptr; // I32 [n_tokens * n_pos_per_embd] ggml_tensor * out_idx = nullptr; // I32 [gathering steps]: nr - 1 (the step's last row produces its draft); null: none + bool out_in_pos = false; // the output rows follow the positions in pos (one upload), out_idx unused void set_input(const llama_ubatch * ubatch) override { GGML_ASSERT(ubatch->n_tokens == n_tokens); @@ -661,8 +662,11 @@ class llm_graph_input_mtp_chain : public llm_graph_input_i { o.push_back((int32_t) st.nr - 1); } } + if (out_in_pos) { + v.insert(v.end(), o.begin(), o.end()); + } ggml_backend_tensor_set(pos, v.data(), 0, v.size() * sizeof(int32_t)); - if (out_idx) { + if (out_idx && !out_in_pos) { ggml_backend_tensor_set(out_idx, o.data(), 0, o.size() * sizeof(int32_t)); } } @@ -986,16 +990,16 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { // contiguous, only a multi-row step 0 narrower than the full mask needs the copy (one copy kernel, not n). // 4: a step picks its output row with a GET_ROWS only if it has several rows (step 0 with catch-up rows); a // 1-row step's row is already the output (no gather, no output-index slice). + // 8: fewer uploads: the output rows ride in the position input, and V is stored with the K indices (equal for a + // non-transposed V cache; the V index input is then unused and not uploaded). static const int tight = [] { const char * e = getenv("LLAMA_MTP_CHAIN_TIGHT"); - return e ? atoi(e) : 7; + return e ? atoi(e) : 15; }(); const uint32_t n_pos = hparams.n_pos_per_embd(); auto chain_inp = std::make_unique(n_pos); chain_inp->n_tokens = n_tokens; - chain_inp->pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens * n_pos); - ggml_set_input(chain_inp->pos); int64_t n_gather = 0; for (int64_t s = 0; s < n_steps; ++s) { llm_graph_input_mtp_chain::step st; @@ -1006,13 +1010,17 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { n_gather += st.gather; chain_inp->steps.push_back(st); } - if (n_gather > 0) { + chain_inp->out_in_pos = (tight & 8) && n_gather > 0; + chain_inp->pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens * n_pos + (chain_inp->out_in_pos ? n_gather : 0)); + ggml_set_input(chain_inp->pos); + if (n_gather > 0 && !chain_inp->out_in_pos) { chain_inp->out_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_gather); ggml_set_input(chain_inp->out_idx); } - const auto steps = chain_inp->steps; - ggml_tensor * pos_all = chain_inp->pos; - ggml_tensor * out_all = chain_inp->out_idx; + const auto steps = chain_inp->steps; + const bool out_in_pos = chain_inp->out_in_pos; + ggml_tensor * pos_all = chain_inp->pos; + ggml_tensor * out_all = chain_inp->out_idx; res->add_input(std::move(chain_inp)); // token embeddings: the host table for the ubatch's own tokens (catch-up tokens need not be in the draft @@ -1057,7 +1065,7 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { ggml_tensor * tok_rows = ggml_get_rows(ctx0, tok_embd_w, ggml_view_1d(ctx0, tokens_all, n_lead, 0)); ggml_tensor * h_embd = ggml_view_2d(ctx0, h_all, h_all->ne[0], n_lead, h_all->nb[1], 0); ggml_tensor * k_all = inp_attn->get_k_idxs(); - ggml_tensor * v_all = inp_attn->get_v_idxs(); + ggml_tensor * v_all = (tight & 8) && !kvc->is_v_trans() ? k_all : inp_attn->get_v_idxs(); ggml_tensor * mask_all = inp_attn->get_kq_mask(); // the device holding the head's weights (none: CPU-only, the slices stay views of the host inputs) @@ -1087,11 +1095,12 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { return c; }; pos_all = on_dev(pos_all, "mtp_chain_pos_dev"); - if (n_gather > 1) { + if (n_gather > 1 && !out_in_pos) { out_all = on_dev(out_all, "mtp_chain_out_dev"); } + const bool v_is_k = v_all == k_all; k_all = on_dev(k_all, "mtp_chain_k_idxs_dev"); - v_all = on_dev(v_all, "mtp_chain_v_idxs_dev"); + v_all = v_is_k ? k_all : on_dev(v_all, "mtp_chain_v_idxs_dev"); mask_all = on_dev(mask_all, "mtp_chain_kq_mask_dev"); } @@ -1100,6 +1109,9 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { struct step_inputs { ggml_tensor * pos, * out, * k_idxs, * v_idxs, * kq_mask; }; + if (out_in_pos) { + out_all = ggml_view_1d(ctx0, pos_all, n_gather, n_tokens * n_pos * ggml_element_size(pos_all)); + } std::vector sin(n_steps); for (int64_t s = 0, g = 0; s < n_steps; ++s) { const auto & st = steps[s]; From d8e10e01faa54ea05b9bb68351dd2b15fcb8e0fb Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:13:31 +1100 Subject: [PATCH 20/24] sycl: build on Windows again - the sync trace's caller lookup (cxxabi/dlfcn/execinfo) is Linux-only Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 135a7667b60c966925a0d0bee8de41942cbc481f on Torchit1/llama.cpp arc-b580) --- ggml/src/ggml-sycl/ggml-sycl.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 403e9dc17ce6..e102fe6e40b0 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -36,9 +36,11 @@ #include #include #include +#ifndef _WIN32 #include #include #include +#endif #include #include @@ -696,6 +698,9 @@ struct ggml_sycl_sync_trace { if (c) { return c; } +#ifdef _WIN32 + return std::string("?"); // no caller lookup on Windows +#else void * fr[24]; const int n = backtrace(fr, 24); std::string label, last; @@ -726,6 +731,7 @@ struct ggml_sycl_sync_trace { } } return label.empty() ? std::string("?") : label; +#endif } void note_sync() { std::lock_guard lock(mtx); From 1d7d3b241543471623ad6ef515a71b68366887c7 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:37:20 +1100 Subject: [PATCH 21/24] sycl: rename the sync trace's small/big counters (the Windows SDK defines small as a macro) Co-Authored-By: Claude Opus 5.5 (cherry picked from commit b5a7db2682c295addc928b085b000b516b54aef6 on Torchit1/llama.cpp arc-b580) --- ggml/src/ggml-sycl/ggml-sycl.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index e102fe6e40b0..2cede321a1e7 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -688,7 +688,7 @@ catch (sycl::exception const &exc) { struct ggml_sycl_sync_trace { const bool on = getenv("GGML_SYCL_SYNC_TRACE") && atoi(getenv("GGML_SYCL_SYNC_TRACE")) != 0; std::mutex mtx; - int big = 0, small = 0; + int n_big = 0, n_small = 0; std::map syncs; std::map> uploads; std::unordered_map names; // return address -> caller label @@ -746,10 +746,10 @@ struct ggml_sycl_sync_trace { void note_graph(const ggml_cgraph * g) { std::lock_guard lock(mtx); if (g->n_nodes < 2000) { - small++; + n_small++; return; } - if (++big < 64) { + if (++n_big < 64) { return; } long ns = 0, nu = 0; @@ -757,17 +757,17 @@ struct ggml_sycl_sync_trace { for (auto & kv : syncs) ns += kv.second; for (auto & kv : uploads) { nu += kv.second.first; nb += kv.second.second; } fprintf(stderr, "sync-trace: %d big graphs (+%d smaller), per big graph: %.1f synchronizes, %.1f blocking uploads (%.0f bytes)\n", - big, small, (double) ns / big, (double) nu / big, (double) nb / big); + n_big, n_small, (double) ns / n_big, (double) nu / n_big, (double) nb / n_big); std::vector> v; for (auto & kv : syncs) v.push_back({ kv.second, kv.first }); std::sort(v.rbegin(), v.rend()); - for (auto & e : v) fprintf(stderr, "sync-trace: sync %6.2f %s\n", (double) e.first / big, e.second.c_str()); + for (auto & e : v) fprintf(stderr, "sync-trace: sync %6.2f %s\n", (double) e.first / n_big, e.second.c_str()); std::vector> w; for (auto & kv : uploads) w.push_back({ kv.second.first, kv.first }); std::sort(w.rbegin(), w.rend()); - for (auto & e : w) fprintf(stderr, "sync-trace: upload %6.2f %8.0f B %s\n", (double) e.first / big, - (double) uploads[e.second].second / big, e.second.c_str()); - big = small = 0; + for (auto & e : w) fprintf(stderr, "sync-trace: upload %6.2f %8.0f B %s\n", (double) e.first / n_big, + (double) uploads[e.second].second / n_big, e.second.c_str()); + n_big = n_small = 0; syncs.clear(); uploads.clear(); } From b49a95f779609aa33a619bc14a1a0b7067fdc5e8 Mon Sep 17 00:00:00 2001 From: Jesse Symons <155509533+Torchit1@users.noreply.github.com> Date: Sat, 10 Oct 2026 20:05:53 +1100 Subject: [PATCH 22/24] mtp chain: all draft ids in one tensor, read back with one copy (LLAMA_MTP_CHAIN_TIGHT bit 16) Feature B, idea 4 (independent implementation of the technique described by llamADALL). The chain read its n per-step argmax ids back with n separate 4-byte ggml_backend_tensor_get_async calls. On this SYCL/Level Zero stack each small device->host read costs ~0.2 ms (measured: reading only step 0's id cut the draft call from 6.07 to 5.16 ms). Now, after the last step, the ids are concatenated into one I32 [n] tensor (n-1 tiny CONCAT kernels at the end of the graph, after every step's compute nodes), marked as the only output, and llama_context reads it with one async copy into mtp_chain_out. The per-step argmax tensors are no longer graph outputs. get_mtp_chain_ids, the draft log and the ids are unchanged (the same integers, copied). Default LLAMA_MTP_CHAIN_TIGHT is now 31. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit b0376aec393cfe96e2652e9031bf0e5d7c410dd7 on Torchit1/llama.cpp arc-b580) --- src/llama-context.cpp | 10 +++++++++- src/llama-graph.cpp | 1 + src/llama-graph.h | 1 + src/models/qwen35.cpp | 20 ++++++++++++++++++-- 4 files changed, 29 insertions(+), 3 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 6d1ef386012d..9f5b27eef935 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2266,7 +2266,15 @@ int llama_context::decode(const llama_batch & batch_inp) { if (chain_e) { // the per-step argmax rows, mapped to tokens in get_mtp_chain_ids() - if (res && !res->t_mtp_chain.empty()) { + if (res && res->t_mtp_chain_ids) { + // one copy of all the ids (LLAMA_MTP_CHAIN_TIGHT bit 16) + ggml_tensor * t = res->t_mtp_chain_ids; + ggml_backend_t b = ggml_backend_sched_get_tensor_backend(sched.get(), t); + GGML_ASSERT(b != nullptr && t->ne[0] == (int64_t) res->t_mtp_chain.size()); + mtp_chain_out.resize(res->t_mtp_chain.size()); + ggml_backend_tensor_get_async(b, t, mtp_chain_out.data(), 0, mtp_chain_out.size() * sizeof(int32_t)); + mtp_chain_n_out = (uint32_t) res->t_mtp_chain.size(); + } else if (res && !res->t_mtp_chain.empty()) { mtp_chain_out.resize(res->t_mtp_chain.size()); for (size_t s = 0; s < res->t_mtp_chain.size(); ++s) { ggml_tensor * t = res->t_mtp_chain[s]; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 2020398f89cb..c5e428af0e7b 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1386,6 +1386,7 @@ void llm_graph_result::reset() { t_sampled.clear(); t_dspark_greedy.clear(); t_mtp_chain.clear(); + t_mtp_chain_ids = nullptr; t_sampled_probs.clear(); t_sampled_logits.clear(); t_candidates.clear(); diff --git a/src/llama-graph.h b/src/llama-graph.h index e1be43b59627..19dab8400cf6 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1057,6 +1057,7 @@ class llm_graph_result { // ARC-LAB: chained MTP draft (cparams.mtp_chain_n): per step, the argmax row of the draft logits (I32 [1]); an // index into the draft head's vocabulary when it has one, else a token id std::vector t_mtp_chain; + ggml_tensor * t_mtp_chain_ids = nullptr; // I32 [n_steps]: the same ids in one tensor (one readback), or null std::vector t_sampled_probs; std::vector t_sampled_logits; std::vector t_candidates; diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index f2e95cf41b6c..a8e6d608f7ab 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -992,9 +992,11 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { // 1-row step's row is already the output (no gather, no output-index slice). // 8: fewer uploads: the output rows ride in the position input, and V is stored with the K indices (equal for a // non-transposed V cache; the V index input is then unused and not uploaded). + // 16: the n draft ids are concatenated into one I32 tensor at the end of the graph (n-1 tiny kernels), so the host + // reads them with one copy: each small device->host read costs ~0.2 ms on this stack. static const int tight = [] { const char * e = getenv("LLAMA_MTP_CHAIN_TIGHT"); - return e ? atoi(e) : 15; + return e ? atoi(e) : 31; }(); const uint32_t n_pos = hparams.n_pos_per_embd(); @@ -1262,7 +1264,9 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { cb(logits, "result_output", -1); ggml_tensor * best = ggml_argmax(ctx0, logits); ggml_set_name(best, "mtp_chain_argmax"); - ggml_set_output(best); + if (!(tight & 16)) { + ggml_set_output(best); + } ggml_build_forward_expand(gf, best); res->t_mtp_chain.push_back(best); @@ -1284,4 +1288,16 @@ void llama_model_qwen35::graph_mtp::build_chain(const llama_model & model) { tok_embd = unrotate(e); h_embd = h_out; } + + if (tight & 16) { + // after the last step, so no node lands between a step's compute nodes + ggml_tensor * ids = res->t_mtp_chain[0]; + for (size_t s = 1; s < res->t_mtp_chain.size(); ++s) { + ids = ggml_concat(ctx0, ids, res->t_mtp_chain[s], 0); + } + ggml_set_name(ids, "mtp_chain_ids"); + ggml_set_output(ids); + ggml_build_forward_expand(gf, ids); + res->t_mtp_chain_ids = ids; + } } From 9c4763f57906ab4fca6203116f62b38e78b4a789 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:13:44 -0700 Subject: [PATCH 23/24] llama: default sync trim and the MTP K/V-only pass on only with a SYCL device LLAMA_SYNC_TRIM and LLAMA_MTP_KV_ONLY were on by default for every backend but were only measured on SYCL. They now default on only when a SYCL device is registered; both env vars still force either way. Other backends keep the previous behaviour unless asked. --- src/llama-context.cpp | 14 +++++++++++--- src/llama-impl.h | 3 +++ src/models/qwen35.cpp | 5 +++-- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9f5b27eef935..3a37572122df 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -826,12 +826,20 @@ void llama_context::sched_reserve() { __func__, (t_end_us - t_start_us)/1000.0, ggml_backend_sched_get_n_copies(sched.get())); } -// ARC-LAB: LLAMA_SYNC_TRIM bit mask (default all on; 0 = upstream behaviour): 1 = skip a context synchronize when the +// ARC-LAB: LLAMA_SYNC_TRIM bit mask (default all on with a SYCL device, else 0; 0 = upstream behaviour): 1 = skip a context synchronize when the // context queued nothing since its last one, 2 = KV-cache rotation matrices kept on the device (llama-kv-cache.cpp) +bool llama_sycl_present() { + static const bool v = [] { + ggml_backend_reg_t reg = ggml_backend_reg_by_name("SYCL"); + return reg != nullptr && ggml_backend_reg_dev_count(reg) > 0; + }(); + return v; +} + int llama_sync_trim() { static const int v = [] { const char * e = getenv("LLAMA_SYNC_TRIM"); - return e ? atoi(e) : 3; + return e ? atoi(e) : (llama_sycl_present() ? 3 : 0); }(); return v; } @@ -841,7 +849,7 @@ void llama_context::synchronize() { return; } - // ARC-LAB: LLAMA_SYNC_TRIM bit 1 (default on): every getter (logits, sampled tokens, nextn rows, ...) synchronizes, + // ARC-LAB: LLAMA_SYNC_TRIM bit 1 (default on with SYCL): every getter (logits, sampled tokens, nextn rows, ...) synchronizes, // and the device's one in-order queue is shared by all contexts. With nothing queued by this context since its // last synchronize, its results are already complete: skip the device wait. if ((llama_sync_trim() & 1) && !sync_pending && !cparams.pipeline_parallel) { diff --git a/src/llama-impl.h b/src/llama-impl.h index 997bb985c610..c6bc214ccba0 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -77,6 +77,9 @@ static inline ggml_tensor * llama_mul_mat_hadamard( // ARC-LAB: LLAMA_SYNC_TRIM bit mask, see llama-context.cpp int llama_sync_trim(); +// true when a SYCL device is registered; the ARC-LAB defaults below were only measured there +bool llama_sycl_present(); + struct time_meas { time_meas(int64_t & t_acc, bool disable = false); ~time_meas(); diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index a8e6d608f7ab..914798657224 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -1,4 +1,5 @@ #include "models.h" +#include "llama-impl.h" #include "llama-memory-recurrent.h" #include "llama-kv-cache.h" @@ -758,14 +759,14 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr res->add_input(std::move(inp)); - // ARC-LAB: KV-only pass (LLAMA_MTP_KV_ONLY, default on; independent implementation of the technique described by + // ARC-LAB: KV-only pass (LLAMA_MTP_KV_ONLY, default on with a SYCL device; independent implementation of the technique described by // GRIMOIRE). A decode with no output rows (the head's prompt catch-up) is only there to fill this layer's K/V, so // build up to the cache store and stop: no query, attention, output projection, FFN or head. The K/V ops are the // full graph's, in the same order, so the cache rows are bit-identical. Unmasked nextn readers copy h for every // token, so they keep the full graph. static const bool kv_only_on = [] { const char * e = getenv("LLAMA_MTP_KV_ONLY"); - return e == nullptr || atoi(e) != 0; + return e ? atoi(e) != 0 : llama_sycl_present(); }(); const bool kv_only = kv_only_on && n_outputs == 0 && cparams.embeddings_nextn_masked; From 62f926b418bb5e025c10333eb513be9bfcfbd1be Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:19:21 -0700 Subject: [PATCH 24/24] sycl: drop duplicate PTQ1_0 case labels left by the MTP chain merge The one-graph draft chain commit added PTQ1_0 to the supports_op type list and a bit-exact f32 get_rows case, but prism already had both labels, so the SYCL build failed (duplicate case value). Keep one label each; get_rows keeps prism's XMX-layout assert and uses the chain's f32 dequant. --- ggml/src/ggml-sycl/getrows.cpp | 6 +----- ggml/src/ggml-sycl/ggml-sycl.cpp | 1 - 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/ggml/src/ggml-sycl/getrows.cpp b/ggml/src/ggml-sycl/getrows.cpp index c6029be7215d..4e53476ba8f2 100644 --- a/ggml/src/ggml-sycl/getrows.cpp +++ b/ggml/src/ggml-sycl/getrows.cpp @@ -280,7 +280,7 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { GGML_ASSERT(!(dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.xmx_pq2) && "PTQ1_0 in the XMX layout reached get_rows"); - get_rows_sycl(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, + get_rows_sycl_f32(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); break; case GGML_TYPE_PQ2_0: @@ -334,10 +334,6 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { get_rows_sycl(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); break; - case GGML_TYPE_PTQ1_0: - get_rows_sycl_f32(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, - src1_i32, (float *)dst->data, ctx.stream()); - break; case GGML_TYPE_Q2_K: get_rows_sycl_f32(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 2cede321a1e7..91a4cd2adcdb 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -8030,7 +8030,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_TYPE_BF16: case GGML_TYPE_F32: case GGML_TYPE_Q1_0: - case GGML_TYPE_PTQ1_0: case GGML_TYPE_PQ2_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: