diff --git a/.gitignore b/.gitignore index ea50f12a815a..457481869bf4 100644 --- a/.gitignore +++ b/.gitignore @@ -161,3 +161,4 @@ a.out.* AGENTS.local.md .pi/SYSTEM.md tools/ui/bun.lock +.worktrees/ diff --git a/AGENTS.md b/AGENTS.md index 97c25074b4c6..e69de29bb2d1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,110 +0,0 @@ -# Instructions for llama.cpp - -> [!IMPORTANT] -> This project does **not** accept pull requests that are fully or predominantly AI-generated. AI tools may be utilized solely in an assistive capacity. -> -> Read more: [CONTRIBUTING.md](CONTRIBUTING.md) - -AI assistance is permissible only when the majority of the code is authored by a human contributor, with AI employed exclusively for corrections or to expand on verbose modifications that the contributor has already conceptualized (see examples below). - ---- - -## Guidelines for Contributors Using AI - -llama.cpp is built by humans, for humans. Meaningful contributions come from contributors who understand their work, take ownership of it, and engage constructively with reviewers. - -Maintainers receive numerous pull requests weekly, many of which are AI-generated submissions where the author cannot adequately explain the code, debug issues, or participate in substantive design discussions. Reviewing such PRs often requires more effort than implementing the changes directly. - -**A pull request represents a long-term commitment.** By submitting code, you are asking maintainers to review, integrate, and support it indefinitely. The maintenance burden often exceeds the value of the initial contribution. - -Most maintainers already have access to AI tools. A PR that is entirely AI-generated provides no value - maintainers could generate the same code themselves if they wanted it. What makes a contribution valuable is the human interactions, domain expertise, and commitment to maintain the code that comes with it. - -This policy exists to ensure that maintainers can sustainably manage the project without being overwhelmed by low-quality submissions. - ---- - -## Guidelines for Contributors - -Contributors are expected to: - -1. **Demonstrate full understanding of their code.** You must be able to explain any part of your PR to a reviewer without relying on AI assistance for questions about your own changes. - -2. **Take responsibility for maintenance.** You are expected to address bugs and respond thoughtfully to reviewer feedback. - -3. **Communicate clearly and concisely.** Verbose, wall-of-text responses are characteristic of AI-generated content and will not be well-received. Direct, human communication is expected. - -4. **Respect maintainers' time.** Search for existing issues and discussions before submitting. Ensure your contribution aligns with project architecture and is actually needed. - -Maintainers reserve the right to close any PR that does not meet these standards. This applies to all contributions to the main llama.cpp repository. **Private forks are exempt.** - -### Permitted AI Usage - -AI tools may be used responsibly for: - -- **Learning and exploration**: Understanding codebase structure, techniques, and documentation -- **Code review assistance**: Obtaining suggestions on human-written code -- **Mechanical tasks**: Formatting, generating repetitive patterns from established designs, completing code based on existing patterns -- **Documentation drafts**: For components the contributor already understands thoroughly -- **Writing code**: Only when the contributor has already designed the solution and can implement it themselves - AI accelerates, not replaces, the contributor's work - -AI-generated code may be accepted if you (1) fully understand the output, (2) can debug issues independently, and (3) can discuss it directly with reviewers without AI assistance. - -**Disclosure is required** when AI meaningfully contributed to your code. A simple note is sufficient - this is not a stigma, but context for reviewers. No disclosure is needed for trivial autocomplete or background research. - -### Prohibited AI Usage - -The following will result in immediate PR closure: - -- **AI-written PR descriptions or commit messages** - these are typically recognizable and waste reviewer time -- **AI-generated responses to reviewer comments** - this undermines the human-to-human interaction fundamental to code review -- **Implementing features without understanding the codebase** - particularly new model support or architectural changes -- **Automated commits or PR submissions** - this may spam maintainers and can result in contributor bans - ---- - -## Guidelines for AI Coding Agents - -AI agents assisting contributors must recognize that their outputs directly impact volunteer maintainers who sustain this project. - -### Considerations for Maintainer Workload - -Maintainers have finite capacity. Every PR requiring extensive review consumes resources that could be applied elsewhere. Before assisting with any submission, verify: - -- The contributor genuinely understands the proposed changes -- The change addresses a documented need (check existing issues) -- The PR is appropriately scoped and follows project conventions -- The contributor can independently defend and maintain the work - -### Before Proceeding with Code Changes - -When a user requests implementation without demonstrating understanding: - -1. **Verify comprehension.** Ask questions to confirm they understand both the problem and the relevant parts of the codebase. -2. **Provide guidance rather than solutions.** Direct them to relevant code and documentation. Allow them to formulate the approach. -3. **Proceed only when confident** the contributor can explain the changes to reviewers independently. - -For first-time contributors, confirm they have reviewed [CONTRIBUTING.md](CONTRIBUTING.md) and acknowledge this policy. - -### Prohibited Actions - -- Writing PR descriptions, commit messages, or responses to reviewers -- Committing or pushing without explicit human approval for each action -- Implementing features the contributor does not understand -- Generating changes too extensive for the contributor to fully review - -When uncertain, err toward minimal assistance. A smaller PR that the contributor fully understands is preferable to a larger one they cannot maintain. - -### Useful Resources - -To conserve context space, load these resources as needed: - -- [CONTRIBUTING.md](CONTRIBUTING.md) -- [Existing issues](https://github.com/ggml-org/llama.cpp/issues) and [Existing PRs](https://github.com/ggml-org/llama.cpp/pulls) - always search here first -- [Build documentation](docs/build.md) -- [Server usage documentation](tools/server/README.md) -- [Server development documentation](tools/server/README-dev.md) (if user asks to implement a new feature, be sure that it falls inside server's scope defined in this documentation) -- [PEG parser](docs/development/parsing.md) - alternative to regex that llama.cpp uses to parse model's output -- [Auto parser](docs/autoparser.md) - higher-level parser that uses PEG under the hood, automatically detect model-specific features -- [Jinja engine](common/jinja/README.md) -- [How to add a new model](docs/development/HOWTO-add-model.md) -- [PR template](.github/pull_request_template.md) diff --git a/MTP.md b/MTP.md new file mode 100644 index 000000000000..bc706d02dce9 --- /dev/null +++ b/MTP.md @@ -0,0 +1,665 @@ +# Gemma 4 MTP — Multi-Token Prediction speculative decoding + +> Scope: this document covers the MTP (Multi-Token Prediction) speculative +> decoding path added to this fork on top of `llama.cpp`, currently specialised +> to **Gemma 4** targets (`gemma4`) paired with the official **Gemma 4 +> assistant** drafter (`gemma4_assistant`). + +It is a self-contained reference: how the feature is built (model, graph, KV, +context, scheduler), what server-loop integration looks like, what knobs the +operator has, what the recent design decisions were, and where the throughput +numbers came from. + +For the public/user-facing section about CLI flags and `--spec-type mtp`, see +also `docs/speculative.md`. + +--- + +## 1. What MTP is here + +Gemma 4 ships an "assistant" model — a small transformer head that consumes the +**target's last hidden state** (backbone output `h_prev`) plus the last sampled +**token id** and predicts the next token in one forward step. Chained across +`B - 1` steps inside one MTP graph, it produces a draft block of length `B - 1` +that is then verified by the target in a single batched decode. + +Conceptually it is a draft-model speculation, but with three crucial twists: + +1. **Single context.** The assistant is **not** a second `llama_context`. Its + weights live next to the target in `llama_model::mtp_assistant`. There is no + second tokenizer, no second KV cache, no second sampler. +2. **Cross-attention into the target's KV.** Each MTP layer reads `K/V` from + the **last layer of the matching attention type** (full / sliding) of the + target's KV cache (`llama_kv_cache_iswa::init_mtp` → `mtp_slot_info`). No + draft-side KV is allocated. +3. **Shared `h_prev` from the target.** The assistant ingests the target's + per-token backbone hidden state (`embeddings_ith`) for the **last accepted** + position. Embeddings must therefore stay enabled on the target context + (`llama_set_embeddings(ctx_tgt, true)`). + +This makes MTP much cheaper than a normal "small draft model" approach: there +is essentially no draft KV, no second model orchestration, and the per-step +graph is tiny (4 transformer blocks for 26B/31B; centroid LM head for E2B/E4B). + +--- + +## 2. Components and where they live + +| Concern | File(s) | +|---|---| +| MTP graph (per-step build) | `src/models/gemma4-assistant.cpp` | +| Model arch + tensor types | `src/llama-arch.cpp/.h`, `src/llama-model.cpp` | +| GGUF tensor mapping | `gguf-py/gguf/tensor_mapping.py`, `gguf-py/gguf/constants.py`, `convert_hf_to_gguf.py` | +| Loading assistant into target | `src/llama.cpp::llama_model_load_mtp_from_file` | +| MTP scheduler + worker + APIs | `src/llama-context.cpp/.h` (`sched_mtp`, `mtp_worker_loop`, `decode_mtp_*`) | +| KV cross-attention helpers | `src/llama-kv-cache.cpp`, `src/llama-kv-cache-iswa.cpp` (`init_mtp`) | +| Speculative driver / overlap | `common/speculative.cpp` (`common_speculative_state_mtp`) | +| Server integration | `tools/server/server-context.cpp` | +| Public C API | `include/llama.h` (`llama_decode_mtp_async/_wait`, `llama_model_load_mtp_from_file`, `llama_model_mtp_n_embd_backbone`, …) | +| Verification helper | `scripts/verify-gemma4-assistant-gguf.py` | +| Run scripts | `scripts/run-gemma4-{,e2b-,e4b-,31b-}mtp-server.sh`, `scripts/quantize-gemma4-edge-assistant-mtp.sh` | +| Tests | `tests/test-speculative-mtp.cpp` | +| Tracing | `LLAMA_MTP_ACC_TRACE` (NDJSON in `common/speculative.cpp`) | + +--- + +## 3. Model side: assistant, centroid LM head, GGUF layout + +The assistant is loaded from its own GGUF and **attached** to a target model: + +```c +int32_t llama_model_load_mtp_from_file( + struct llama_model * model, + const char * path_assistant, + struct llama_model_params mparams); +``` + +After load, the target carries: + +- `model.mtp_assistant` — a fully-loaded `llama_model` with arch + `gemma4_assistant` (4 transformer blocks, the pre/post backbone projections, + optional centroid head). +- `hparams.n_embd_backbone` — must equal target's backbone hidden size; this is + asserted at load and re-checked at draft time (`n_bb` in + `decode_mtp_run`). + +CLI surface (`common/arg.cpp`, `common/common.cpp`): + +- `--mtp-head ` (preferred) and `--model-draft / -md` (back-compat alias) + — feed the same `mparams_dft.path` field. +- `--draft-block-size ` — head proposes `B - 1` tokens per round. +- `--gpu-layers-draft / -ngld`, `-ctkd / -ctvd` — placement and KV typing for + the **assistant** weights when offloaded. + +### Centroid / ordered-embeddings LM head (E2B / E4B) + +For Edge models (`use_ordered_embeddings = true` in HF config), the LM head is +not the dense tied embedding but a **MaskedEmbedder**: + +1. `centroid_logits = mul_mat(mtp.centroids, h)` → `[n_centroids]`. +2. `top_k(centroid_logits, centroid_intermediate_top_k)` → `top_k` centroid ids + (I32, on-device). +3. `mtp.token_ordering` is viewed as `[vsc, n_centroids]` + (`vsc = n_vocab / n_centroids`); each centroid column lists `vsc` candidate + token ids. `get_rows` gathers the candidate ids for the chosen centroids. +4. `get_rows(token_embd, ids)` then `mul_mat(·, h)` produces sparse logits over + only those candidates. +5. We **scatter** them into a full `[n_vocab]` row pre-filled with `-1e30` + (`ggml_fill_inplace + ggml_set_rows`). The full-vocab row is what the + verifier expects — sparse-only argmax broke server accept (rare-token edge + cases) and so was reverted. + +GGUF layout for the centroid head (see `convert_hf_to_gguf.py` and +`docs/development/gemma4-assistant-tensor-inventory.md`): + +| Tensor | Stored type | On-disk shape | Notes | +|---|---|---|---| +| `mtp.centroids.weight` | F16/F32 (or quant) | `[n_embd, n_centroids]` after GGUF dim packing | numpy is `[n_centroids, n_embd]`; written as-is so loader sees `mul_mat`-compatible shape | +| `mtp.token_ordering.weight` | **I32** (kept integer end-to-end) | `[n_vocab]` | must not be quantized — the converter explicitly preserves I32 | +| `mtp.pre_projection.weight` / `mtp.post_projection.weight` | as model | `[2*n_embd, n_embd]` / `[n_embd, n_embd_backbone]` | concatenates `[token_embd, h_prev]` then projects back | + +The verifier `scripts/verify-gemma4-assistant-gguf.py` enforces these shape / +dtype invariants and is run automatically by every `run-gemma4-*-mtp-server.sh` +script (skip with `VERIFY_ASSISTANT_GGUF=0`). + +--- + +## 4. Per-step MTP graph (`gemma4-assistant.cpp`) + +`llm_build_gemma4_mtp` builds a single-token, single-sequence graph (`n_tokens += 1`, `n_seqs = 1`, `n_outputs = 1`): + +1. Inputs (registered as `llm_graph_input_mtp`): + - `inp_last_token : I32 [1]` + - `inp_h_prev : F32 [n_embd_backbone, 1]` + - `inp_pos` (standard `build_inp_pos`) + - `inp_attn` (`build_attn_inp_kv_iswa`) +2. Token embedding from the **target's** `tok_embd`, then scaled by + `sqrt(n_embd)` to mirror Gemma 4's input scaling. +3. `concat([tok_e, h_prev], axis=0) → mtp.pre_projection` (collapses the + `2 * n_embd` channel back to `n_embd`). +4. 4 transformer blocks (`mtp.layers[il]`): + - RMSNorm → Q proj → Q-norm → RoPE. + - **Cross-attention** via `build_attn_mtp`: queries from MTP, K/V fetched + from the target's KV cache at `il_kv = last layer in target with the same + attention type` (SWA / full). + - Per HF Gemma 4 quirk (`attention_k_eq_v: true`): even when V was derived + from K, the V slot is **written** with rms-norm-without-scale and + un-rotated, so cross-attn must always read V from cache (not reuse the + post-RoPE K). This is encoded as `use_k_as_v = false`. + - Standard residual + post-norm + GELU FFN + post-FFN norm + per-layer + `out_scale` (if present) + `build_cvec`. +5. Final RMSNorm → `mtp.post_projection` produces the **next-step `h_prev`** + (this is what the host stitches between steps). +6. LM head — dense (tied) **or** centroid-routed for ordered embeddings. +7. Optional `f_final_logit_softcapping`. +8. **In-graph greedy argmax**: `ggml_argmax(cur)` → `I32 [1]`. The final result + exposes three tensors via `llm_graph_result`: + - `t_embd` = `h_post` (next `h_prev`) + - `t_logits` = full-vocab row (kept for diagnostic / `out_logits` API) + - `t_argmax` = greedy token id + +The host reads `t_argmax` (4 bytes) per step instead of pulling the full +F32 `[n_vocab]` row across a backend boundary and running CPU argmax. On +Gemma 4 + Q4_K_XL this alone delivered ~+2-3% throughput +(109.5 → 112.5 tps at n=128; 95.8 → 97.8 tps at n=512), with bit-identical +greedy drafts. The full row is still computed in-graph and is fetched on +demand by passing `out_logits != NULL` to the synchronous +`llama_decode_mtp(...)` API (legacy / diagnostic path), which transparently +falls back to `decode_mtp_sync`. + +--- + +## 5. KV sharing — what is read, what is appended + +The MTP step does **not** allocate or write its own KV. It reads the target's +KV by: + +- `kv_iswa->init_mtp(seq_id, ub)` — produces a memory context whose attention + inputs (`build_attn_inp_kv_iswa`) wire the cross-attn to the target slot for + `seq_id`, with a mask that admits all positions `≤ attn_pos`. +- `attn_pos` is taken from `llama_memory_seq_pos_max(mem, seq_id)` immediately + before submission (post-`seq_rm`). All `n_steps` step positions chosen for + RoPE are strictly **`> attn_pos`**, so the causal/SWA mask uniformly admits + every target cell — that is why a single mask suffices for the whole chained + draft. + +KV-safety contract for asynchronous draft work (see Section 7): + +- `decode_mtp_async` snapshots `h_prev` and `attn_pos` at submit time. +- The target may **append** at positions `> attn_pos` between submit and + `_wait` (this is what the verify decode does), but it must not evict, rewrite + or `seq_rm` cells at positions `≤ attn_pos` until `_wait` returns. +- The current append-only KV cache satisfies this. `common_speculative_cancel` + is invoked at the few server-loop points that *do* mutate KV destructively + (request stop / release; new request `seq_rm`; spec-disabled iterations). + +--- + +## 6. `llama_context` plumbing — `sched_mtp`, worker, and APIs + +The async pipeline lives in `src/llama-context.cpp/.h`. New members on the +context (Phase C of the original plan): + +- `sched_mtp` — a **dedicated** `ggml_backend_sched` for MTP. Created lazily + by `ensure_sched_mtp()` and reserved with a single-token MTP graph (the MTP + graph is invariant in size: `n_tokens = 1`, `n_seqs = 1`, `n_outputs = 1`, + so one reserve covers all subsequent calls). +- `gf_res_prev_mtp` — a **separate** `llm_graph_result` cache. This is the key + reason MTP graph reuse survives target-decode resets, and is the single + largest win of the async refactor. +- A worker thread (`mtp_worker`), `std::mutex` + 2 condition variables + (`mtp_cv_request`, `mtp_cv_response`), and request/response slots + (`std::optional`, `bool mtp_in_flight`, `std::optional`). +- `backend_cfg_mu` — guards shared backend reconfiguration + (`set_threadpool_fn`, `set_n_threads_fns`) so the worker's + `graph_compute_mtp` cannot race the main thread's `graph_compute`. The lock + is held only across cheap setters; the actual `graph_compute_async` calls run + unlocked so target verify and MTP encode can interleave on each scheduler. + +Public C APIs (`include/llama.h`): + +```c +LLAMA_API int32_t llama_decode_mtp_async( + struct llama_context * ctx, + llama_seq_id seq_id, + llama_pos attn_pos, + llama_token last_token, + const float * h_prev, + int32_t n_steps); + +LLAMA_API int32_t llama_decode_mtp_wait( + struct llama_context * ctx, + llama_token * out_drafts, + float * out_h_prev_last); + +// Backward-compatible synchronous facade. If out_logits != NULL falls back to +// decode_mtp_sync (per-step logits captured in-thread). +LLAMA_API int32_t llama_decode_mtp( + struct llama_context * ctx, + llama_seq_id seq_id, llama_pos attn_pos, + llama_token last_token, float * h_prev, + int32_t n_steps, + llama_token * out_drafts, float * out_logits, float * out_h_prev_last); +``` + +Contract: + +- At most **one in-flight request per context**. `_async` while a previous + request has not been `_wait`ed returns `-7`. +- `h_prev` is *copied* into the request → caller may free / reuse immediately. +- Drafts are written into `out_drafts[0..n_steps-1]`; the last `h_prev` is + optionally copied into `out_h_prev_last`. + +Worker loop (`mtp_worker_loop`): waits on `mtp_pending`, runs +`decode_mtp_run` (the per-step chain on `sched_mtp`), publishes +`mtp_completed`, and notifies. `decode_mtp_run` per step: + +1. Build `llama_ubatch` `{token=last_token, embd=h, pos=attn_pos+1+k, output=0}`. +2. `kv_iswa->init_mtp(seq_id, ub)` → memory context. +3. `process_ubatch_mtp` → reuse cached graph if `can_reuse(gparams)` else + rebuild + alloc. +4. `graph_compute_mtp` → `sched_mtp.graph_compute_async` → synchronize. +5. Read `t_argmax` (4 bytes) → `last_token = drafts[k]`; read `t_embd` → + `h` (next `h_prev`). + +On context destruction the worker is signalled via `mtp_worker_stop`, woken, +and joined before tearing down `sched_mtp`. + +--- + +## 7. Speculative driver — pipeline depth-2 + +The host driver lives in `common/speculative.cpp :: +common_speculative_state_mtp`. Its job is to translate the server's +"draft / accept" loop into the right `_async / _wait` calls and to enforce the +KV-safety contract. + +### Depth-2 overlap + +The server normally goes: + +``` + loop: + drafts = common_speculative_draft(...) # produce drafts + target_decode(...) # verify drafts + n_acc, sampled = sample_and_accept_n(...) + common_speculative_accept(spec, n_acc) + seq_rm(...); update slot.sampled / batch.dft index +``` + +Depth-2 inserts a `prepare_next` at the **end** of the iteration, after +`accept` and `seq_rm`: + +``` + common_speculative_prepare_next(spec, slot.sampled) # async submit +``` + +…which calls `llama_decode_mtp_async(...)` for the *next* round using: + +- `attn_pos = seq_pos_max(seq_id)` (post `seq_rm`), +- the real sampled token `slot.sampled` (no optimistic guess), +- `h_prev = embeddings_ith(h_idx)` snapshotted right after sample/accept (see + Section 8 for `h_idx`). + +Then on the next iteration, `common_speculative_draft` checks +`has_pending`. If pending and `pending_n_steps == n_steps`, it goes "lazy": + +``` + llama_decode_mtp_wait(...) # blocks only on whatever is left of MTP +``` + +This **overlaps MTP draft compute with everything that happens between +`accept` and the next `draft`**: token I/O, OAI streaming, slot bookkeeping, +batching, the next prefill if any. The benefit is real because the MTP graph, +while small, is not free — it is `B - 1` sequential single-token decodes +through 4 layers + cross-attn + LM head. + +When `n_steps` changes between iterations (e.g. on the last iteration of a +request), or when the target is about to mutate KV destructively, the driver +**drains** the in-flight request (`mtp_drain_pending_discard`) to keep the +`_async`/`_wait` invariant intact. + +The depth-2 path can be A/B-tested at runtime by exporting +`LLAMA_PIPELINE_DEPTH2=0`, which turns `prepare_next` into a no-op and +restores depth-1 (sync `_async + _wait` inside `draft`). + +### Drain points (server-side) + +`tools/server/server-context.cpp` invokes `common_speculative_cancel` in three +places: + +1. When the iteration **skips** speculative decoding (`n_remaining == 1`, + `n_min` not satisfied, etc.) — otherwise the worker would compute against + KV that is about to change in the upcoming target_decode (we observed Metal + command-buffer status 3 on turbo3 KV before this guard). +2. After `send_final_response` / `slot.release` — the next request will + `seq_rm` and overwrite cells the worker is still reading. +3. In `common_speculative_begin` (new prompt) — the previous generation's + in-flight MTP must not bleed into the next prompt. + +--- + +## 8. The `h_idx` correction + +A subtle correctness issue: `embeddings_ith(-1)` returns the **last batch +output**, which after partial draft acceptance is the hidden state of a +**rejected** draft (computed for the wrong input). Feeding that as `h_prev` +collapses acceptance. + +Fix: after `sample_and_accept_n` the server sets + +```cpp +common_speculative_set_h_idx(slot.spec, slot.i_batch_dft[ids.size() - 1]); +``` + +i.e. it points the next MTP draft at the batch index of the **last accepted +token**. This is honored both in the sync `draft` path and in +`prepare_next` (Section 7). + +--- + +## 9. Adaptive skip-streak + +There are workloads (numbers, code, rare tokens deep in long generations) where +the MTP head is consistently wrong. Drafting still costs ~10 ms with no +accepted tokens. The driver detects this: + +- `prev_n_acc_drafts` snapshot at the end of each `draft`. +- Increment `zero_accept_streak` when `n_acc_drafts` did not move since the + previous call; reset on any non-empty accept. +- After `LLAMA_MTP_SKIP_STREAK_THRESHOLD` consecutive zero-accepts (1..32), + return an empty draft for one batch (server falls back to a single-token + verify), reset the streak, and let the next batch re-arm. +- `skip_streak_last_draft` prevents threshold=1 from oscillating into a + permanent skip. + +Off by default (env unset / `0`). The `MTP_PRESET=throughput` Edge presets do +**not** enable it either — turn it on per-deployment when the workload +warrants. + +--- + +## 10. Diagnostic NDJSON tracer (`LLAMA_MTP_ACC_TRACE`) + +Set `LLAMA_MTP_ACC_TRACE=1` (stderr) or `LLAMA_MTP_ACC_TRACE=/path/to.ndjson` +(append) to enable the `mtp_acc_tracer` in `common/speculative.cpp`. Off by +default at zero overhead (enabled-only L2 reduction over `n_bb` per draft). + +Two events per iteration, paired by `iter`: + +- `mtp_draft` — `iter`, `path` (`sync` / `lazy` / `skip-streak` / `skip-nsteps`), + `seq_id`, `id_last`, `h_idx`, `attn_pos`, `n_steps`, `h_l2` (L2 norm of + `h_prev`), and `drafts[]`. +- `mtp_accept` — `iter`, `n_accepted`, `n_drafted_prev`. + +This is the recommended tool for any acceptance-rate debugging: per-position +acceptance, h_prev stability, `h_idx` selection bias, and depth-2 lazy/sync +distribution all fall out by joining on `iter`. + +--- + +## 11. Operating it — scripts and presets + +### Pre-built assistant GGUFs + +Official Gemma 4 assistant heads, converted with this fork's +`convert_hf_to_gguf.py` (preserves I32 `mtp.token_ordering` for the +centroid-head Edge variants), are published as a Hugging Face collection: + +> [AtomicChat / Gemma 4 Assistant GGUF](https://huggingface.co/collections/AtomicChat/gemma-4-assistant-gguf) +> — F16 / Q8_0 / Q5_K_M / **Q4_K_M** / Q4_K_S quantizations. + +| Target | Assistant repo | Recommended quant | +|---|---|---| +| Gemma 4 E2B | [`AtomicChat/gemma-4-E2B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-E2B-it-assistant-GGUF) | **Q4_K_M** | +| Gemma 4 E4B | [`AtomicChat/gemma-4-E4B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-E4B-it-assistant-GGUF) | **Q4_K_M** | +| Gemma 4 26B-A4B | [`AtomicChat/gemma-4-26B-A4B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-26B-A4B-it-assistant-GGUF) | **Q4_K_M** / Q4_K_S | +| Gemma 4 31B | [`AtomicChat/gemma-4-31B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-31B-it-assistant-GGUF) | **Q4_K_M** / Q4_K_S | + +Q4_K_M is the recommended default: throughput is identical to F16 in the +matrix bench (the head is small enough that bandwidth, not weight precision, +dominates), while VRAM/RAM footprint is ~4× lower. Drop to F16 only if you +are debugging an acceptance regression that you suspect is quant-related; the +verifier `scripts/verify-gemma4-assistant-gguf.py` will refuse to load a +malformed assistant GGUF in either case. + +The repo helpers prefer a quantized assistant under `.scratch/` when one +exists (`gemma-{e2b,e4b,…}-assistant-mtp-Q4_K_M.gguf`) and fall back to F16 +otherwise. Override with `DRAFT_GGUF=…` or pass `--mtp-head` directly. + +### Run scripts + +Helper scripts live under `scripts/`: + +| Script | Target | Notes | +|---|---|---| +| `run-gemma4-mtp-server.sh` | gemma 4 26B | dense LM head; `MTP_PRESET` not used | +| `run-gemma4-31b-mtp-server.sh` | gemma 4 31B | dense LM head | +| `run-gemma4-e2b-mtp-server.sh` | gemma 4 E2B | centroid head; `MTP_PRESET` aware | +| `run-gemma4-e4b-mtp-server.sh` | gemma 4 E4B | centroid head; `MTP_PRESET` aware | +| `run-gemma4-server-turbo.sh` | dense baselines, no MTP | TurboQuant KV demo | +| `quantize-gemma4-edge-assistant-mtp.sh` | quantizer for E2B/E4B assistant | preserves I32 ordering | + +Edge presets (`MTP_PRESET`): + +| Preset | `--draft-block-size` (B) | `--draft-max` | +|---|---:|---:| +| `throughput` | 2 | 6 | +| `lift` | 3 | 8 | +| `balanced` | 3 | 8 | +| `quality` | 4 | 16 | + +Override directly with `DRAFT_BLOCK_SIZE`, `DRAFT_MAX`, +`LLAMA_MTP_SKIP_STREAK_THRESHOLD`. KV typing is taken from `CTK / CTV / CTKD / +CTVD` — both target and assistant inherit the same default (`turbo3`). + +The bench harness `.scratch/bench-matrix.sh` runs the matrix +`{model} × {f16-base, turbo3-base, f16-mtp, turbo3-mtp} × {short=128, long=512} +× 3 runs` against `/v1/chat/completions` with `temperature=0`, +`cache_prompt=false` and `stream=false`, and reports median tps + mean +draft-accept rate. + +--- + +## 12. Latest matrix benchmark (`.scratch/bench-logs/matrix-q4chat.log`) + +Run on 2026-05-07. Q4_K_S assistant heads, draft-block defaults from each +script (`B = 3` for the dense scripts at the time of this run). `accept` is +`draft_n_accepted / draft_n` averaged over 3 runs; `tps` is the median. + +### Bench host + +| Component | Value | +|---|---| +| Machine | MacBook Pro (`Mac16,5`, MX313LL/A) | +| SoC | Apple **M4 Max** — 16 CPU cores (12P + 4E), 40-core GPU | +| Unified memory | 48 GB LPDDR5 | +| OS | macOS 26.3.1 (build 25D2128), Darwin 25.3.0 | +| llama.cpp backend | Metal (full GPU offload: `-ngl 99 -ngld 99`, `-fa on`) | +| Server | local `llama-server` over `127.0.0.1:8080` | +| Client | `python3 urllib` → `/v1/chat/completions`, `temperature=0`, `cache_prompt=false`, `stream=false` | +| Driver | `.scratch/bench-matrix.sh` (3 runs/cell, median tps, mean accept) | + +Single-slot configuration (`--parallel 1 -np 1 --cont-batching`); no other +heavy GPU/CPU workloads were running on the host during the matrix sweep. + +| model | mode | short tps (n=128) | long tps (n=512) | short accept | long accept | +|---|---|---:|---:|---:|---:| +| gemma-26B | f16-base | 81.54 | 83.06 | — | — | +| gemma-26B | turbo3-base | 53.81 | 53.89 | — | — | +| gemma-26B | f16-mtp | **109.49** | **95.75** | 85.9% | 68.9% | +| gemma-26B | turbo3-mtp | 81.91 | 72.17 | 82.3% | 67.9% | +| gemma-31B | f16-base | 14.15 | 15.20 | — | — | +| gemma-31B | turbo3-base | 15.79 | 14.82 | — | — | +| gemma-31B | f16-mtp | **20.24** | **17.30** | 88.0% | 74.6% | +| gemma-31B | turbo3-mtp | 18.67 | 15.68 | 87.0% | 70.8% | + +Key observations: + +- **f16 MTP**: +34 % short / +15 % long over baseline on 26B; +43 % short / + +14 % long on 31B. Acceptance is dominated by short, "essay-y" prompts; long + drafts hit the natural ceiling once content drifts into less predictable + spans. +- **turbo3 MTP**: +52 % short / +34 % long over the turbo3 baseline on 26B + (turbo3 baseline is slower than f16 because the gemma-26B target is + compute-bound at `f16` and bandwidth-helped by turbo3 only when + memory-bound; that asymmetry is not specific to MTP). +- **31B base inversion** (`turbo3-base 15.79 > f16-base 14.15` short): 31B is + bandwidth-bound on this rig, so turbo3 KV beats f16 on the short cell. MTP + still adds value on top of either KV typing. +- **Accept short > accept long** is consistent across the matrix: as decode + drifts away from boilerplate phrasing, the assistant's drafts become less + reliable and `B - 1` chained steps amplify the rejection. + +### How we got here (history within this branch) + +The matrix logs in `.scratch/bench-logs/` show the optimisation journey for the +gemma-26B `f16-mtp` short-prompt cell: + +| Log (mtime, `ls -lt`) | Short tps | Long tps | Short accept | What changed | +|---|---:|---:|---:|---| +| `matrix-run2.log` (01:26) | 70.89 | 76.79 | 55.5% | early async pipeline, sync wrapper | +| `matrix-old.log` (01:41) | 61.88 | 63.98 | 50.0% | depth-1 sync MTP, `h_idx=-1` regression | +| `matrix-q4chat.log` (02:02) | **109.49** | 95.75 | **85.9%** | depth-2 + in-graph argmax + correct `h_idx` | +| `matrix-c-prime.log` (02:50, partial) | 112.30 | 96.69 | 85.9% | identical config, additional run sample | + +The big jump (~62 → ~109 tps short) came from three independent fixes +landing together: + +1. **`h_idx` correction** so MTP feeds the *accepted* hidden state instead of a + rejected draft's output (acceptance jumps from ~50% to ~86%). +2. **Pipeline depth-2 overlap** so MTP work overlaps post-accept bookkeeping + (steady ~+8% throughput at fixed accept). +3. **In-graph argmax** so the host transfers 4 bytes instead of `n_vocab × 4 B` + per step (~+2-3% on top). + +--- + +## 13. Trade-offs and gotchas + +These are the non-obvious failure / regression modes you should keep in mind +when changing or extending this code. + +**Embeddings on the target context.** MTP is meaningless without +`llama_set_embeddings(ctx_tgt, true)`. The server wires this conditionally per +batch (`need_embeddings = need_embd() || mtp_active`). If a future code path +flips embeddings off mid-generation, MTP will silently degrade to drafting +against zero `h_prev`. + +**`h_idx` after partial accept.** Forgetting to call +`common_speculative_set_h_idx` after `sample_and_accept_n` regresses accept +rate to ~50 % on the same workload (matrix-old vs matrix-q4chat). Any new code +path that produces drafts must restore the correct batch index of the *last +accepted* token, not `-1`. + +**KV append-only invariant.** Async MTP correctness depends on `attn_pos` cells +remaining stable until `_wait`. Any new operation that rewrites KV in place +(eviction, sliding-window compaction, retroactive `seq_rm` past `attn_pos`) +must call `common_speculative_cancel` first. The server has three explicit +drain points (Section 7); reuse them rather than inventing a fourth contract. + +**Single in-flight request per context.** This is intentional — multiplexing +MTP across slots requires a sched-per-slot or a request queue with its own +graph-cache. Today a second `_async` returns `-7` and `prepare_next` is a +no-op when one is in flight. With `--parallel > 1` slots run on the same +context: the MTP overlap currently benefits only the slot whose `prepare_next` +won the race; the others fall back to sync. Lifting this is non-trivial +(graph-cache, scheduler, KV snapshot all need per-slot identity). + +**`draft_block_size` vs. `draft_max`.** `draft_block_size` is the **MTP head's +block** (head emits `B - 1` tokens). `draft_max` is the standard llama.cpp +upper bound on draft length the server will accept. For Edge centroid heads +(heavier per-step), small `B` (2-3) usually wins; for the dense 26B/31B, +`B = 3` is the current sweet spot in the matrix bench. + +**Centroid-head `top_k` cost.** Edge MTP runs `top_k` over `n_centroids` and a +routed `get_rows` per draft step. Greedy still materialises the full-vocab row +(masked-fill + scatter) so verify-side argmax stays consistent. +`use_ordered_embeddings` has measurably higher per-step cost than the dense +head; budget `B = 2` (`MTP_PRESET=throughput`) by default on Edge. The Edge +matrix cell is not yet in `matrix-q4chat.log` (the script lists `gemma-E4B` +in `MODELS`, but the row is not present — the GGUF was missing on the bench +host that day). + +**Skip-streak hysteresis.** With `LLAMA_MTP_SKIP_STREAK_THRESHOLD=1` and +without `skip_streak_last_draft`, the driver would skip every other batch +forever as soon as one zero-accept happened. Keep that guard. + +**Backend reconfiguration races.** `set_n_threads` / `set_threadpool` are +process-global on a backend. The `backend_cfg_mu` window in `graph_compute` / +`graph_compute_mtp` is intentionally tiny (only the setters, never the +`graph_compute_async` itself). Lengthening that critical section will block +the worker on every target step and erase the depth-2 win. + +**Vocab compatibility for MTP is laxer than for `--spec-type draft`.** Target +chat templates own stop / EOS tokens; the MTP head only predicts next-token +ids. `common_speculative_are_compatible_mtp` therefore checks `vocab_type`, +size (within `SPEC_VOCAB_MAX_SIZE_DIFFERENCE`) and per-token text equality +from id ≥ 5, but **skips** bos/eos/add_bos/add_eos checks. Don't reuse this +loosened check for non-MTP draft pairings. + +**Optimistic last token (future work).** Submitting `prepare_next` with a +guess of the next sampled token before sample/accept could hide one extra +`llama_decode` on hits. On misses we'd waste the entire MTP block. Not landed +— would need a clear measurement that hit-rate is high enough to justify the +miss cost on this workload. + +--- + +## 14. Quick reference + +```sh +# 26B + Q4 assistant, MTP on TurboQuant3 KV (matches matrix-q4chat row). +scripts/run-gemma4-mtp-server.sh + +# E4B, throughput preset (B=2, max=6), centroid head, optional skip-streak. +LLAMA_MTP_SKIP_STREAK_THRESHOLD=4 \ +MTP_PRESET=throughput \ +scripts/run-gemma4-e4b-mtp-server.sh + +# A/B-test depth-2 overlap vs sync at the same model/config: +LLAMA_PIPELINE_DEPTH2=0 scripts/run-gemma4-mtp-server.sh + +# NDJSON acceptance trace to a file. +LLAMA_MTP_ACC_TRACE=/tmp/mtp.ndjson scripts/run-gemma4-mtp-server.sh + +# Re-run the matrix bench (median over 3 runs per cell). +bash .scratch/bench-matrix.sh | tee .scratch/bench-logs/matrix-$(date +%H%M).log +``` + +Environment knobs: + +| Var | Default | Effect | +|---|---|---| +| `LLAMA_PIPELINE_DEPTH2` | unset (on) | `=0` disables depth-2 overlap; falls back to sync `_async + _wait` inside `draft`. | +| `LLAMA_MTP_SKIP_STREAK_THRESHOLD` | unset / `0` (off) | `1..32` enables zero-accept skip streak. | +| `LLAMA_MTP_ACC_TRACE` | unset (off) | `1` → stderr; any other value → file path (append). | +| `LLAMA_GRAPH_REUSE_DISABLE` | unset (off) | Disables `llm_graph_result::can_reuse`. Useful when changing the MTP graph; disastrous for throughput. | + +Public API entry points: + +```c +llama_model_load_mtp_from_file(model, path, mparams); +llama_model_has_mtp_assistant(model); +llama_model_get_mtp_assistant(model); +llama_model_mtp_n_embd_backbone(model); + +llama_decode_mtp_async(ctx, seq_id, attn_pos, last_token, h_prev, n_steps); +llama_decode_mtp_wait (ctx, out_drafts, out_h_prev_last); +llama_decode_mtp (ctx, ..., out_logits, ...); // sync facade +``` + +Driver entry points (`common/speculative.h`): + +```c +common_speculative_init / _free +common_speculative_set_seq_id // server slot -> target seq id +common_speculative_set_h_idx // last accepted batch idx after accept_n +common_speculative_begin // per-prompt; drains stale MTP +common_speculative_draft // emits drafts (lazy-waits depth-2) +common_speculative_accept // updates stats; emits trace +common_speculative_prepare_next // depth-2: async submit for next round +common_speculative_cancel // drain in-flight MTP +common_speculative_print_stats +``` diff --git a/README.md b/README.md index 4c5396b8016f..0bb8c90a93dc 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,8 @@ MIT, same as upstream llama.cpp. ## Hot topics +- **Gemma 4 MTP speculative decoding: pair a `gemma4` target with the official `gemma4_assistant` head (loaded via `--mtp-head`) for ~+30-50 % short-prompt throughput. See [MTP.md](MTP.md) and the pre-built Q4 assistant GGUFs at the [AtomicChat/Gemma 4 Assistant GGUF collection](https://huggingface.co/collections/AtomicChat/gemma-4-assistant-gguf).** +- **TurboQuant KV cache & weights: WHT-rotated low-bit quantization with backend-native kernels (Metal `TurboFlash`, CUDA, Vulkan, HIP). Use `-ctk turbo3 -ctv turbo3` for ~4.3× KV compression, or quantize weights to `TQ4_1S`/`TQ3_1S`. See [Compression below](#turboquant-kv-cache--weight-compression).** - **Hugging Face cache migration: models downloaded with `-hf` are now stored in the standard Hugging Face cache directory, enabling sharing with other HF tools.** - **[guide : using the new WebUI of llama.cpp](https://github.com/ggml-org/llama.cpp/discussions/16938)** - [guide : running gpt-oss with llama.cpp](https://github.com/ggml-org/llama.cpp/discussions/15396) @@ -230,6 +232,76 @@ range of hardware - locally and in the cloud. The `llama.cpp` project is the main playground for developing new features for the [ggml](https://github.com/ggml-org/ggml) library. +## TurboQuant — KV cache & weight compression + +> **Credits.** TurboQuant in this fork is built on top of the absolutely +> awesome work by **[@TheTom](https://github.com/TheTom)** in +> [TheTom/llama-cpp-turboquant](https://github.com/TheTom/llama-cpp-turboquant). +> Huge thanks for the original WHT-rotated quantization design, the reference +> kernels, and the relentless backend ports — none of this would exist +> without that project. ❤️ + +This fork (`atomic-llama-cpp-turboquant`) packages **TurboQuant** as a family +of WHT-rotated low-bit quantization formats with backend-native kernels. They +target two distinct memory-traffic problems: + +- **KV cache compression** — `TURBO2_0` / `TURBO3_0` / `TURBO4_0` (2/3/4-bit, + WHT + PolarQuant). Selected at runtime via `-ctk` / `-ctv`. +- **Model weight compression** — `TQ3_1S` / `TQ4_1S` (3/4-bit, WHT-rotated + Lloyd-Max with `block_size = 32`). Selected at quantize time as a + `--type` for `llama-quantize`. + +### KV cache types (`-ctk` / `-ctv`) + +| Type | Bits | Compression vs F16 | Notes | +|---|---:|---:|---| +| `turbo2` | 2 | ~6.4× | maximum compression, intended for large-context budgets | +| `turbo3` | 3 | ~4.3× | **recommended default**; Metal `TurboFlash` decode kernel | +| `turbo4` | 4 | ~3.8× | highest accuracy of the family, safest fallback | + +Typical invocation with full GPU offload + Flash-Attention: + +```bash +llama-server -m model.gguf -c 32768 -ngl 99 \ + -ctk turbo3 -ctv turbo3 -fa on +``` + +Pair with `--cache-reuse N` and a long `-c` to see the practical KV-budget +win — TurboQuant typically shifts the OOM ceiling on Apple Silicon / +discrete GPUs by 3-6× at the same context length. + +### Weight quantization types (`llama-quantize`) + +| Type | Bits | Block size | Notes | +|---|---:|---:|---| +| `TQ3_1S` | 3 | 32 | 8-level Lloyd-Max + WHT rotation | +| `TQ4_1S` | 4 | 32 | 16-level Lloyd-Max + WHT rotation; fused Metal/Vulkan MUL_MAT_VEC kernels | + +```bash +# Convert / re-quantize an F16/F32 GGUF to TQ4_1S. +llama-quantize model-f16.gguf model-tq4_1s.gguf TQ4_1S +``` + +`TQ4_1S` typically delivers ~25-35 % size reduction vs Q8_0 with single-digit-% +PPL deltas; on bandwidth-bound models / GPUs it can also be faster than Q8_0 +because of the lighter memory traffic. + +### Backend support + +| Backend | KV `turbo2` / `turbo3` / `turbo4` | Weights `TQ3_1S` / `TQ4_1S` | +|---|---|---| +| Metal (Apple Silicon) | yes; `TurboFlash` flash-attn decode kernel for `turbo3` (off-by-default on Apple10 — see PR #91) | yes (V2.1 fused kernels) | +| CUDA (NVIDIA) | `turbo3` / `turbo4` (full); `turbo2` via reference path | `TQ4_1S` MUL_MAT_VEC | +| Vulkan | `turbo3` KV (FA + coopmat), `SET_ROWS` for `turbo2/4` | `TQ4_1S` (specialised MUL_MAT_VEC, SET_ROWS, CPY) | +| HIP / ROCm | `turbo3` KV; F16-K + TURBO-V mixed dispatch | reference | +| CPU | reference (correctness, not throughput) | reference | + +For combining TurboQuant KV with **Gemma 4 MTP speculative decoding**, see +[MTP.md §11-12](MTP.md). The matrix bench shows that the combo +(`turbo3` KV + MTP) is the right pick when the target model is bandwidth-bound +(e.g. Gemma 4 31B), and that f16-KV + MTP wins when the target is +compute-bound (e.g. Gemma 4 26B-A4B on M4 Max). +
Models @@ -569,6 +641,87 @@ To learn more about model quantization, [read this documentation](tools/quantize
+-
+ Enable TurboQuant KV cache compression (this fork) + + Use a TurboQuant KV-cache type for both K and V — typically with + Flash-Attention enabled — to cut KV memory traffic and footprint at + long contexts. Recommended default is **`turbo3`** (3-bit, ~4.3× vs F16, + accelerated by `TurboFlash` on Metal and dedicated kernels on + CUDA / Vulkan / HIP). + + ```bash + # ~4.3x KV compression vs F16, full GPU offload, Flash-Attn on. + llama-server -m model.gguf -c 32768 \ + -ngl 99 -ctk turbo3 -ctv turbo3 -fa on + ``` + + Pick a stronger compression preset by stepping the bit-width: + + ```bash + -ctk turbo2 -ctv turbo2 # 2-bit KV, ~6.4x vs F16 (highest compression) + -ctk turbo3 -ctv turbo3 # 3-bit KV, ~4.3x (default sweet spot) + -ctk turbo4 -ctv turbo4 # 4-bit KV, ~3.8x (highest accuracy / fallback) + ``` + + See the longer write-up [below](#turboquant-kv-cache--weight-compression) + for weight quantization (`TQ4_1S` / `TQ3_1S`) and the per-backend support + matrix. + +
+ +-
+ Enable Gemma 4 MTP speculative decoding (this fork) + + Pair a `gemma4` target with the official `gemma4_assistant` MTP head. The + head is loaded **into** the target context (no second `llama_context`, + no second KV cache) and runs on a dedicated scheduler so MTP draft compute + overlaps target verification. + + Pre-built assistant GGUFs (recommended **`Q4_K_M`** / `Q4_K_S` for best + speed/quality) are published in the [AtomicChat / Gemma 4 Assistant GGUF + collection](https://huggingface.co/collections/AtomicChat/gemma-4-assistant-gguf): + + | Target model | Assistant (MTP head) GGUF | + |---|---| + | Gemma 4 E2B | [`AtomicChat/gemma-4-E2B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-E2B-it-assistant-GGUF) | + | Gemma 4 E4B | [`AtomicChat/gemma-4-E4B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-E4B-it-assistant-GGUF) | + | Gemma 4 26B-A4B | [`AtomicChat/gemma-4-26B-A4B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-26B-A4B-it-assistant-GGUF) | + | Gemma 4 31B | [`AtomicChat/gemma-4-31B-it-assistant-GGUF`](https://huggingface.co/AtomicChat/gemma-4-31B-it-assistant-GGUF) | + + ```bash + # Manual invocation — works for any of the four targets above. + llama-server \ + -m /path/to/gemma-4-target.gguf \ + --mtp-head /path/to/gemma-4-assistant-Q4_K_M.gguf \ + --spec-type mtp \ + --draft-block-size 3 \ + -c 16384 \ + -ngl 99 -ngld 99 \ + -fa on \ + --host 127.0.0.1 --port 8080 + ``` + + Repo helper scripts pick the right defaults per target (and prefer + a quantized assistant when present under `.scratch/`): + + ```bash + # Dense targets (block size 3 by default). + scripts/run-gemma4-mtp-server.sh # 26B-A4B + scripts/run-gemma4-31b-mtp-server.sh # 31B + + # Edge / centroid-head targets (MTP_PRESET aware: throughput|lift|balanced|quality). + MTP_PRESET=throughput scripts/run-gemma4-e4b-mtp-server.sh + MTP_PRESET=throughput scripts/run-gemma4-e2b-mtp-server.sh + ``` + + Full architecture, async pipeline, KV-safety contract, tuning knobs and + the latest matrix benchmark live in [MTP.md](MTP.md). User-facing CLI + flags (`--spec-type`, `--draft-*`) are documented in + [docs/speculative.md](docs/speculative.md). + +
+ -
Serve an embedding model diff --git a/common/arg.cpp b/common/arg.cpp index 2cbe89447825..0fc6002336c9 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3633,11 +3633,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_N_GPU_LAYERS_DRAFT")); add_opt(common_arg( {"--spec-draft-model", "-md", "--model-draft"}, "FNAME", - "draft model for speculative decoding (default: unused)", + "draft model for speculative decoding, or Gemma 4 MTP assistant GGUF when using --spec-type mtp (default: unused)", [](common_params & params, const std::string & value) { params.speculative.draft.mparams.path = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL")); + add_opt(common_arg( + {"--mtp-head"}, "FNAME", + "alias for Gemma 4 MTP: path to gemma4_assistant GGUF (loaded into the target; use with --spec-type mtp)", + [](common_params & params, const std::string & value) { + params.speculative.draft.mparams.path = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_MTP_HEAD")); add_opt(common_arg( {"--spec-type"}, common_speculative_all_types_str(), string_format("comma-separated list of types of speculative decoding to use (default: %s)\n", diff --git a/common/sampling.h b/common/sampling.h index 49506a00cd8b..11d2179e128c 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -66,8 +66,10 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co // generalized version of common_sampler_sample // -// will cross-reference the sampled tokens with a batch of draft tokens and accept those that match -// if the sampler disagrees at some point, we stop and return the accepted tokens up to now +// will cross-reference greedy-argmax target tokens with a batch of draft tokens and accept those that match +// (draft positions use raw-logit argmax so stochastic params like temperature do not break MTP / greedy drafts) +// then samples the next continuation token with the full chain when all drafts matched +// if the target disagrees at some point, we stop and return the accepted tokens up to now // // common_sampler_sample_n(gsmpl, ctx, { idx }, {}); // diff --git a/common/speculative.cpp b/common/speculative.cpp index 1962cbe828c1..d6b9f9ccd464 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1,9 +1,9 @@ #include "speculative.h" +#include "../src/llama-ext.h" // staging API: llama_set_embeddings_pre_norm / llama_get_embeddings_pre_norm_ith (used by MTP) #include "common.h" #include "ggml.h" #include "llama.h" -#include "../src/llama-ext.h" // staging API: llama_set_embeddings_pre_norm / llama_get_embeddings_pre_norm_ith (used by MTP) #include "log.h" #include "ngram-cache.h" #include "ngram-map.h" @@ -12,24 +12,25 @@ #include #include +#include #include #include #include -#include #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 const std::map common_speculative_type_from_name_map = { - {"none", COMMON_SPECULATIVE_TYPE_NONE}, - {"draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE}, - {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3}, - {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP}, - {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, - {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K}, - {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V}, - {"ngram-mod", COMMON_SPECULATIVE_TYPE_NGRAM_MOD}, - {"ngram-cache", COMMON_SPECULATIVE_TYPE_NGRAM_CACHE} + { "none", COMMON_SPECULATIVE_TYPE_NONE }, + { "draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE }, + { "draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 }, + { "mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP }, + { "draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP }, + { "ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE }, + { "ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K }, + { "ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V }, + { "ngram-mod", COMMON_SPECULATIVE_TYPE_NGRAM_MOD }, + { "ngram-cache", COMMON_SPECULATIVE_TYPE_NGRAM_CACHE } }; static std::string common_speculative_get_devices_str(const std::vector & devices) { @@ -42,23 +43,25 @@ static std::string common_speculative_get_devices_str(const std::vector n_vocab_dft - ? n_vocab_tgt - n_vocab_dft - : n_vocab_dft - n_vocab_tgt; + const int vocab_diff = n_vocab_tgt > n_vocab_dft ? n_vocab_tgt - n_vocab_dft : n_vocab_dft - n_vocab_tgt; if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__); @@ -113,8 +114,7 @@ static bool common_speculative_are_compatible( if (std::strcmp(token_text_tgt, token_text_dft) != 0) { LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__); LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i, - common_token_to_piece(vocab_tgt, i).c_str(), - common_token_to_piece(vocab_dft, i).c_str()); + common_token_to_piece(vocab_tgt, i).c_str(), common_token_to_piece(vocab_dft, i).c_str()); return false; } } @@ -134,21 +134,21 @@ struct common_speculative_impl { uint32_t n_seq; - size_t n_call_begin = 0; // number of times this implementation was called for refresh. - size_t n_call_draft = 0; // number of times this implementation was called for generation. - size_t n_call_accept = 0; // number of times this implementation was called for accumulation. + size_t n_call_begin = 0; // number of times this implementation was called for refresh. + size_t n_call_draft = 0; // number of times this implementation was called for generation. + size_t n_call_accept = 0; // number of times this implementation was called for accumulation. - size_t n_gen_drafts = 0; // number of times a draft or part was generated by this implementation. - size_t n_acc_drafts = 0; // number of times a draft or part was accepted by the target model. - size_t n_gen_tokens = 0; // number of tokens generated by this implementation. - size_t n_acc_tokens = 0; // number of tokens accepted by the target model. + size_t n_gen_drafts = 0; // number of times a draft or part was generated by this implementation. + size_t n_acc_drafts = 0; // number of times a draft or part was accepted by the target model. + size_t n_gen_tokens = 0; // number of tokens generated by this implementation. + size_t n_acc_tokens = 0; // number of tokens accepted by the target model. // TODO: track performance of most recent calls - const bool gen_perf = true; // whether to generate performance stats. + const bool gen_perf = true; // whether to generate performance stats. - int64_t t_begin_us = 0; // total time spent in refresh of this implementation in microseconds. - int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. - int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. + int64_t t_begin_us = 0; // total time spent in refresh of this implementation in microseconds. + int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. + int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. common_speculative_impl(common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {} @@ -176,21 +176,18 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { std::vector smpls; - common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq) - , params(params.draft) - { + common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq), + params(params.draft) { auto * ctx_dft = this->params.ctx_dft; auto * ctx_tgt = this->params.ctx_tgt; LOG_INF("%s: adding speculative implementation 'draft-simple'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); + LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, this->params.n_max, this->params.n_min, + this->params.p_min); LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__, - this->params.n_gpu_layers, - ggml_type_name(this->params.cache_type_k), - ggml_type_name(this->params.cache_type_v), - ctx_tgt ? "yes" : "no", - ctx_dft ? "yes" : "no", + this->params.n_gpu_layers, ggml_type_name(this->params.cache_type_k), + ggml_type_name(this->params.cache_type_v), ctx_tgt ? "yes" : "no", ctx_dft ? "yes" : "no", common_speculative_get_devices_str(this->params.devices).c_str()); batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); @@ -215,8 +212,8 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { smpls.resize(n_seq); for (auto & smpl : smpls) { common_params_sampling params; - params.no_perf = false; - params.top_k = 10; + params.no_perf = false; + params.top_k = 10; params.samplers = { COMMON_SAMPLER_TYPE_TOP_K, }; @@ -240,9 +237,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { } } - ~common_speculative_impl_draft_simple() override { - llama_batch_free(batch); - } + ~common_speculative_impl_draft_simple() override { llama_batch_free(batch); } void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override { // noop @@ -268,7 +263,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { common_batch_clear(batch); // keep track of which sequences are still drafting - int n_drafting = 0; + int n_drafting = 0; std::vector drafting(n_seq); for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { @@ -311,8 +306,8 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { const auto * cur_p = common_sampler_get_candidates(smpl, true); for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", seq_id, k, i, + cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } @@ -329,13 +324,12 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { common_sampler_accept(smpl, id, true); - auto & dp = dparams.at(seq_id); + auto & dp = dparams.at(seq_id); auto & result = *dp.result; result.push_back(id); - if ((params.n_max <= (int) result.size()) || - (dp.n_max > 0 && dp.n_max <= (int) result.size())) { + if ((params.n_max <= (int) result.size()) || (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; continue; @@ -373,19 +367,17 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { // noop } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } }; struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { //common_params_speculative_eagle3 params; - common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) - { + common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) { LOG_INF("%s: adding speculative implementation 'draft-eagle3'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, params.draft.n_max, params.draft.n_min, params.draft.p_min); + LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, params.draft.n_max, params.draft.n_min, + params.draft.p_min); } void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override { @@ -405,13 +397,11 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { // noop } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } }; struct common_speculative_impl_draft_mtp : public common_speculative_impl { - common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) + common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) llama_batch batch; @@ -425,7 +415,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // Per-sequence cross-batch carryover: pair (h_p, x_{p+1}) at MTP pos p+1. // The last h-row of one process() call needs the first token of the NEXT // call to pair with, so it's stashed here until that next call fires. - std::vector> pending_h; // [n_seq][n_embd] + std::vector> pending_h; // [n_seq][n_embd] std::vector i_batch_beg; std::vector i_batch_end; @@ -433,17 +423,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // Hidden rows from the most recent target verification batch, grouped by seq. // Row 0 corresponds to the sampled token, row N to the Nth accepted draft token. std::vector> verify_h; - std::vector verify_h_rows; + std::vector verify_h_rows; // Per-seq draft length from the last draft() call, used in accept() to // roll back ctx_dft's recurrent state past the AR draft's redundant // pre-advancement before process() mirrored the verify batch. std::vector last_n_drafted; - common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) - , params(params.draft) - { + common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq), + params(params.draft) { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set"); @@ -451,20 +440,18 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { n_embd = llama_model_n_embd(llama_get_model(ctx_dft)); LOG_INF("%s: adding speculative implementation 'draft-mtp'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); + LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, + this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__, - this->params.n_gpu_layers, - ggml_type_name(this->params.cache_type_k), - ggml_type_name(this->params.cache_type_v), - ctx_tgt ? "yes" : "no", - ctx_dft ? "yes" : "no", + this->params.n_gpu_layers, ggml_type_name(this->params.cache_type_k), + ggml_type_name(this->params.cache_type_v), ctx_tgt ? "yes" : "no", ctx_dft ? "yes" : "no", common_speculative_get_devices_str(this->params.devices).c_str()); const int32_t n_b = (int32_t) llama_n_batch(ctx_dft); - batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd, /*n_seq_max=*/ 1); + batch = llama_batch_init(/*n_tokens=*/n_b, /*embd=*/n_embd, /*n_seq_max=*/1); // llama_batch_init allocates only one of token/embd; MTP needs both. // TODO: fix, how to call without malloc - batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b); + batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b); smpls.resize(n_seq); for (auto & s : smpls) { @@ -530,14 +517,15 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { if (N <= 0) { return; } - auto * ctx_dft = this->params.ctx_dft; + auto * ctx_dft = this->params.ctx_dft; const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id); if (pos_max < N - 1) { - LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - " - "process() hook may not have run on every prefill ubatch " - "(need_embd / logits=1 on every prompt position?). " - "Drafts may degrade.\n", - __func__, (int) pos_max, N - 1); + LOG_WRN( + "%s: ctx_dft pos_max=%d < N-1=%d - " + "process() hook may not have run on every prefill ubatch " + "(need_embd / logits=1 on every prompt position?). " + "Drafts may degrade.\n", + __func__, (int) pos_max, N - 1); } } @@ -588,7 +576,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // TODO:this is generally true, but would be nice to assert it { const float * h_tgt = llama_get_embeddings_pre_norm(ctx_tgt); - std::memcpy(batch.embd + (size_t) 1 * n_embd, h_tgt, row_bytes * (n_tokens-1)); + std::memcpy(batch.embd + (size_t) 1 * n_embd, h_tgt, row_bytes * (n_tokens - 1)); //{ // // string with seq_ids in the batch @@ -624,7 +612,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { continue; } - const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1; + const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1; verify_h_rows[seq_id] = n_rows; verify_h[seq_id].resize((size_t) n_rows * n_embd); @@ -633,8 +621,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::memcpy(verify_h[seq_id].data() + (size_t) i * n_embd, h, row_bytes); } - std::memcpy(pending_h[seq_id].data(), - verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes); + std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes); } return true; @@ -646,11 +633,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_batch_clear(batch); // keep track of which sequences are still drafting - int n_drafting = 0; + int n_drafting = 0; std::vector drafting(n_seq); - const float * h_row = nullptr; - const size_t row_bytes = (size_t) n_embd * sizeof(float); + const float * h_row = nullptr; + const size_t row_bytes = (size_t) n_embd * sizeof(float); for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { auto & dp = dparams[seq_id]; @@ -666,7 +653,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); h_row = pending_h[seq_id].data(); - std::memcpy(batch.embd + n_embd*(batch.n_tokens - 1), h_row, row_bytes); + std::memcpy(batch.embd + n_embd * (batch.n_tokens - 1), h_row, row_bytes); } int ret = llama_decode(ctx_dft, batch); @@ -696,8 +683,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const auto * cur_p = common_sampler_get_candidates(smpl, true); for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", seq_id, k, i, + cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } @@ -714,7 +701,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_sampler_accept(smpl, id, true); - auto & dp = dparams.at(seq_id); + auto & dp = dparams.at(seq_id); auto & result = *dp.result; result.push_back(id); @@ -726,7 +713,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true); - std::memcpy(batch.embd + n_embd*(batch.n_tokens - 1), h_row, row_bytes); + std::memcpy(batch.embd + n_embd * (batch.n_tokens - 1), h_row, row_bytes); } if (batch.n_tokens == 0) { @@ -767,18 +754,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { return; } - const int32_t i_h = std::min(n_accepted, n_rows - 1); - const size_t row_bytes = (size_t) n_embd * sizeof(float); + const int32_t i_h = std::min(n_accepted, n_rows - 1); + const size_t row_bytes = (size_t) n_embd * sizeof(float); std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } - bool need_embd_pre_norm() const override { - return true; - } + bool need_embd_pre_norm() const override { return true; } }; // state of self-speculation (simple implementation, not ngram-map) @@ -788,16 +771,15 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { // shared across all sequences common_ngram_simple_config config; - common_speculative_impl_ngram_simple( - const common_params_speculative & params, uint32_t n_seq, - common_ngram_simple_config config) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq) - , params(params.ngram_simple) - , config(config) - { + common_speculative_impl_ngram_simple(const common_params_speculative & params, + uint32_t n_seq, + common_ngram_simple_config config) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq), + params(params.ngram_simple), + config(config) { LOG_INF("%s: adding speculative implementation 'ngram-simple'\n", __func__); - LOG_INF("%s: - size_n=%d, size_m=%d, min_hits=%d\n", __func__, - this->params.size_n, this->params.size_m, this->params.min_hits); + LOG_INF("%s: - size_n=%d, size_m=%d, min_hits=%d\n", __func__, this->params.size_n, this->params.size_m, + this->params.min_hits); } void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override { @@ -826,27 +808,23 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { // noop } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } }; struct common_speculative_impl_ngram_map_k : public common_speculative_impl { // n_seq configs std::vector config; - common_speculative_impl_ngram_map_k( - const common_ngram_map & config, - uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, n_seq) - { + common_speculative_impl_ngram_map_k(const common_ngram_map & config, uint32_t n_seq) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, n_seq) { for (uint32_t i = 0; i < n_seq; i++) { this->config.push_back(config); } - LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(this->type).c_str()); - LOG_INF("%s: - size_key=%d, size_value=%d, key_only=%d, min_hits=%d\n", __func__, - config.size_key, config.size_value, config.key_only, config.min_hits); + LOG_INF("%s: adding speculative implementation '%s'\n", __func__, + common_speculative_type_to_str(this->type).c_str()); + LOG_INF("%s: - size_key=%d, size_value=%d, key_only=%d, min_hits=%d\n", __func__, config.size_key, + config.size_value, config.key_only, config.min_hits); } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { @@ -883,9 +861,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { common_ngram_map_accept(config[seq_id], n_accepted); } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } }; struct common_speculative_impl_ngram_mod : public common_speculative_impl { @@ -910,24 +886,23 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { std::vector sinfos; - common_speculative_impl_ngram_mod( - const common_params_speculative & params, - uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq) - , params(params.ngram_mod) - , mod(params.ngram_mod.n_match, 4*1024*1024) - , verbose(std::getenv("LLAMA_TRACE") != nullptr) { + common_speculative_impl_ngram_mod(const common_params_speculative & params, uint32_t n_seq) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq), + params(params.ngram_mod), + mod(params.ngram_mod.n_match, 4 * 1024 * 1024), + verbose(std::getenv("LLAMA_TRACE") != nullptr) { static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t)); LOG_INF("%s: adding speculative implementation 'ngram-mod'\n", __func__); - LOG_INF("%s: - n_match=%d, n_max=%d, n_min=%d\n", __func__, - this->params.n_match, this->params.n_max, this->params.n_min); - LOG_INF("%s: - mod size=%zu (%.3f MB)\n", __func__, - mod.size(), (float)(mod.size_bytes())/1024/1024); + LOG_INF("%s: - n_match=%d, n_max=%d, n_min=%d\n", __func__, this->params.n_match, this->params.n_max, + this->params.n_min); + LOG_INF("%s: - mod size=%zu (%.3f MB)\n", __func__, mod.size(), (float) (mod.size_bytes()) / 1024 / 1024); if (this->params.n_match < 16) { - LOG_WRN("%s: ngram_mod n_match=%d is too small - poor quality is possible, " - "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, this->params.n_match); + LOG_WRN( + "%s: ngram_mod n_match=%d is too small - poor quality is possible, " + "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", + __func__, this->params.n_match); } sinfos.resize(n_seq); @@ -936,7 +911,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { auto & sinfo = sinfos[seq_id]; - sinfo.i_last = 0; + sinfo.i_last = 0; sinfo.n_draft_last = 0; const size_t n = mod.get_n(); @@ -950,7 +925,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { sinfo.i_last = prompt.size() - n; - const double f = (double)mod.get_used() / (double)mod.size(); + const double f = (double) mod.get_used() / (double) mod.size(); LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f); constexpr double f_thold = 0.25; @@ -961,10 +936,8 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { } } - void draft_one( - llama_seq_id seq_id, - common_speculative_draft_params & dparams) { - auto & sinfo = sinfos[seq_id]; + void draft_one(llama_seq_id seq_id, common_speculative_draft_params & dparams) { + auto & sinfo = sinfos[seq_id]; auto & result = *dparams.result; const auto & prompt = *dparams.prompt; @@ -1044,7 +1017,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { // 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; + const double f_acc = (double) n_accepted / (double) sinfo.n_draft_last; if (f_acc < 0.25) { sinfo.n_low++; if (sinfo.n_low >= 5) { @@ -1053,7 +1026,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { } mod.reset(); - sinfo.n_low = 0; + sinfo.n_low = 0; sinfo.i_last = 0; } } else { @@ -1062,9 +1035,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { } } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } }; struct common_speculative_impl_ngram_cache : public common_speculative_impl { @@ -1076,7 +1047,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { bool save_static; struct seq_info { - size_t cache_size = 0; // number of tokens in n-gram cache + size_t cache_size = 0; // number of tokens in n-gram cache common_ngram_cache ngram_cache_context; common_ngram_cache ngram_cache_dynamic; @@ -1085,23 +1056,20 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { std::vector sinfos; - common_speculative_impl_ngram_cache( - const common_params_speculative & params, - uint32_t n_seq, - uint16_t n_draft, - const std::string & path_static, - const std::string & path_dynamic, - bool save_dynamic, - bool save_static) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq) - , params(params.ngram_cache) - , n_draft(n_draft) - , save_dynamic(save_dynamic) - , save_static(save_static) - { + common_speculative_impl_ngram_cache(const common_params_speculative & params, + uint32_t n_seq, + uint16_t n_draft, + const std::string & path_static, + const std::string & path_dynamic, + bool save_dynamic, + bool save_static) : + common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq), + params(params.ngram_cache), + n_draft(n_draft), + save_dynamic(save_dynamic), + save_static(save_static) { LOG_INF("%s: adding speculative implementation 'ngram-cache'\n", __func__); - LOG_INF("%s: - n_draft=%d, cache_static=%s, cache_dynamic=%s\n", __func__, - n_draft, + LOG_INF("%s: - n_draft=%d, cache_static=%s, cache_dynamic=%s\n", __func__, n_draft, path_static.empty() ? "none" : path_static.c_str(), path_dynamic.empty() ? "none" : path_dynamic.c_str()); @@ -1138,10 +1106,8 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { // noop } - void draft_one( - llama_seq_id seq_id, - common_speculative_draft_params & dparams) { - auto & sinfo = sinfos[seq_id]; + void draft_one(llama_seq_id seq_id, common_speculative_draft_params & dparams) { + auto & sinfo = sinfos[seq_id]; auto & result = *dparams.result; const auto & prompt = *dparams.prompt; @@ -1152,13 +1118,11 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { for (size_t j = sinfo.cache_size; j < prompt.size(); ++j) { tokens_new.push_back(prompt[j]); } - tokens_new.push_back(dparams.id_last); // add the last token + tokens_new.push_back(dparams.id_last); // add the last token // Update context ngram cache with new dparams.prompt: - common_ngram_cache_update( - sinfo.ngram_cache_context, - LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, - tokens_new, tokens_new.size(), false); + common_ngram_cache_update(sinfo.ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, tokens_new, + tokens_new.size(), false); sinfo.cache_size = prompt.size() + 1; } @@ -1171,11 +1135,8 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { result.push_back(dparams.id_last); - common_ngram_cache_draft( - inp, result, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, - sinfo.ngram_cache_context, - sinfo.ngram_cache_dynamic, - sinfo.ngram_cache_static); + common_ngram_cache_draft(inp, result, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, sinfo.ngram_cache_context, + sinfo.ngram_cache_dynamic, sinfo.ngram_cache_static); if (result.size() > 0) { // delete first token in result (which is the id_last token) @@ -1205,9 +1166,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { // noop } - bool need_embd() const override { - return false; - } + bool need_embd() const override { return false; } }; struct common_speculative { @@ -1220,9 +1179,8 @@ struct common_speculative { std::vector impl_last; }; -static common_ngram_map get_common_ngram_map( - common_speculative_type type, - const common_params_speculative_ngram_map & config) { +static common_ngram_map get_common_ngram_map(common_speculative_type type, + const common_params_speculative_ngram_map & config) { uint16_t size_key = config.size_n; uint16_t size_value = config.size_m; bool key_only = type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K; @@ -1231,18 +1189,18 @@ static common_ngram_map get_common_ngram_map( return common_ngram_map(size_key, size_value, key_only, min_hits); } -static common_speculative_impl_ngram_cache create_state_ngram_cache( - const common_speculative_config & config, - uint32_t n_seq, - const std::string & path_static, - const std::string & path_dynamic) { - uint16_t n_draft = 8; // TODO get from config? +static common_speculative_impl_ngram_cache create_state_ngram_cache(const common_speculative_config & config, + uint32_t n_seq, + const std::string & path_static, + const std::string & path_dynamic) { + uint16_t n_draft = 8; // TODO get from config? // TODO bool param in common/common.h to set save_static/save_dynamic? - bool save_static = false; + bool save_static = false; bool save_dynamic = false; - common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic); + common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, + save_dynamic); return state; } @@ -1273,16 +1231,26 @@ const char * common_speculative_all_types_str() { std::string common_speculative_type_to_str(common_speculative_type type) { switch (type) { - case COMMON_SPECULATIVE_TYPE_NONE: return "none"; - case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: return "draft-simple"; - case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3"; - case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp"; - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: return "ngram-mod"; - case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: return "ngram-cache"; - default: return "unknown"; + case COMMON_SPECULATIVE_TYPE_NONE: + return "none"; + case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: + return "draft-simple"; + case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: + return "draft-eagle3"; + case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: + return "draft-mtp"; + case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: + return "ngram-simple"; + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: + return "ngram-map-k"; + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: + return "ngram-map-k4v"; + case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: + return "ngram-mod"; + case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: + return "ngram-cache"; + default: + return "unknown"; } } @@ -1294,7 +1262,7 @@ std::vector common_speculative_types_from_names(const s auto type = common_speculative_type_from_name_map.find(name); if (type != common_speculative_type_from_name_map.end()) { if (type->second == COMMON_SPECULATIVE_TYPE_NONE) { - return std::vector { COMMON_SPECULATIVE_TYPE_NONE }; + return std::vector{ COMMON_SPECULATIVE_TYPE_NONE }; } types.push_back(type->second); continue; @@ -1325,14 +1293,14 @@ static uint32_t common_get_enabled_speculative_configs(const std::vector configs = {}; // list of speculative configs to try + std::vector configs = {}; // list of speculative configs to try { uint32_t enabled_configs = common_get_enabled_speculative_configs(params.types); bool has_draft_model_path = !params.draft.mparams.path.empty(); bool has_draft_simple = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE)); - bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3 + bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3 bool has_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr; bool has_ngram_cache = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_CACHE)); @@ -1369,7 +1337,9 @@ common_speculative * common_speculative_init(common_params_speculative & params, has_draft_simple = false; } } else if (has_draft_model_path && !has_mtp && !has_draft_eagle3) { - LOG_WRN("%s: draft model is specified but 'draft' speculative type is not explicitly enabled - enabling it\n", __func__); + LOG_WRN( + "%s: draft model is specified but 'draft' speculative type is not explicitly enabled - enabling it\n", + __func__); has_draft_simple = true; } @@ -1390,61 +1360,61 @@ common_speculative * common_speculative_init(common_params_speculative & params, switch (config.type) { case COMMON_SPECULATIVE_TYPE_NONE: break; - case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: { - impls.push_back(std::make_unique(config.params, n_seq)); - break; - } - case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: { - impls.push_back(std::make_unique(config.params, n_seq)); - break; - } - case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: { - impls.push_back(std::make_unique(config.params, n_seq)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { - common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); - - uint16_t ngram_size_key = ngram_map.size_key; - uint16_t mgram_size_value = ngram_map.size_value; - - auto config_simple = common_ngram_simple_config { - /* .size_ngram = */ ngram_size_key, - /* .size_mgram = */ mgram_size_value - }; - auto state = std::make_unique( - /* .params = */ config.params, - /* .n_seq = */ n_seq, - /* .state = */ config_simple - ); - impls.push_back(std::move(state)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: { - impls.push_back( - std::make_unique( - get_common_ngram_map(config.type, config.params.ngram_map_k), n_seq)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: { - impls.push_back( - std::make_unique( - get_common_ngram_map(config.type, config.params.ngram_map_k4v), n_seq)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: { - impls.push_back( - std::make_unique(config.params, n_seq)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: { - auto state = create_state_ngram_cache( - config, n_seq, - params.ngram_cache.lookup_cache_static, - params.ngram_cache.lookup_cache_dynamic); - impls.push_back(std::make_unique(state)); - break; - } + case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: + { + impls.push_back(std::make_unique(config.params, n_seq)); + break; + } + case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: + { + impls.push_back(std::make_unique(config.params, n_seq)); + break; + } + case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: + { + impls.push_back(std::make_unique(config.params, n_seq)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: + { + common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); + + uint16_t ngram_size_key = ngram_map.size_key; + uint16_t mgram_size_value = ngram_map.size_value; + + auto config_simple = common_ngram_simple_config{ /* .size_ngram = */ ngram_size_key, + /* .size_mgram = */ mgram_size_value }; + auto state = std::make_unique( + /* .params = */ config.params, + /* .n_seq = */ n_seq, + /* .state = */ config_simple); + impls.push_back(std::move(state)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: + { + impls.push_back(std::make_unique( + get_common_ngram_map(config.type, config.params.ngram_map_k), n_seq)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: + { + impls.push_back(std::make_unique( + get_common_ngram_map(config.type, config.params.ngram_map_k4v), n_seq)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: + { + impls.push_back(std::make_unique(config.params, n_seq)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: + { + auto state = create_state_ngram_cache(config, n_seq, params.ngram_cache.lookup_cache_static, + params.ngram_cache.lookup_cache_dynamic); + impls.push_back(std::make_unique(state)); + break; + } default: break; } @@ -1455,11 +1425,9 @@ common_speculative * common_speculative_init(common_params_speculative & params, return nullptr; } - auto * result = new common_speculative { - /* .dparams = */ common_speculative_draft_params_vec(n_seq), - /* .impls = */ std::move(impls), - /* .impl_last = */ std::vector(n_seq, nullptr) - }; + auto * result = new common_speculative{ /* .dparams = */ common_speculative_draft_params_vec(n_seq), + /* .impls = */ std::move(impls), + /* .impl_last = */ std::vector(n_seq, nullptr) }; return result; } @@ -1472,9 +1440,7 @@ void common_speculative_free(common_speculative * spec) { delete spec; } -common_speculative_draft_params & common_speculative_get_draft_params( - common_speculative * spec, - llama_seq_id seq_id) { +common_speculative_draft_params & common_speculative_get_draft_params(common_speculative * spec, llama_seq_id seq_id) { GGML_ASSERT(spec); GGML_ASSERT(seq_id < (llama_seq_id) spec->dparams.size()); @@ -1640,6 +1606,15 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, u } } +void common_speculative_prepare_next(common_speculative * spec, llama_token id_last) { + GGML_UNUSED(spec); + GGML_UNUSED(id_last); +} + +void common_speculative_cancel(common_speculative * spec) { + GGML_UNUSED(spec); +} + void common_speculative_print_stats(const common_speculative * spec) { if (spec == nullptr) { return; @@ -1657,13 +1632,11 @@ void common_speculative_print_stats(const common_speculative * spec) { str_perf = ""; } - LOG_INF("statistics %16s: #calls(b,g,a) = %4zu %6zu %6zu, #gen drafts = %6zu, #acc drafts = %5zu, #gen tokens = %6zu, #acc tokens = %5zu%s\n", - common_speculative_type_to_str(impl->type).c_str(), - impl->n_call_begin, impl->n_call_draft, impl->n_call_accept, - impl->n_gen_drafts, - impl->n_acc_drafts, - impl->n_gen_tokens, - impl->n_acc_tokens, - str_perf.c_str()); + LOG_INF( + "statistics %16s: #calls(b,g,a) = %4zu %6zu %6zu, #gen drafts = %6zu, #acc drafts = %5zu, #gen tokens = " + "%6zu, #acc tokens = %5zu%s\n", + common_speculative_type_to_str(impl->type).c_str(), impl->n_call_begin, impl->n_call_draft, + impl->n_call_accept, impl->n_gen_drafts, impl->n_acc_drafts, impl->n_gen_tokens, impl->n_acc_tokens, + str_perf.c_str()); } } diff --git a/common/speculative.h b/common/speculative.h index f24bac79edb7..290cddcdbf85 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -1,7 +1,7 @@ #pragma once -#include "llama.h" #include "common.h" +#include "llama.h" struct common_speculative; @@ -65,6 +65,10 @@ void common_speculative_draft(common_speculative * spec); // informs the speculative context that n_accepted tokens were accepted by the target model void common_speculative_accept(common_speculative * spec, llama_seq_id, uint16_t n_accepted); +// Compatibility hooks used by server speculative pipeline integration. +void common_speculative_prepare_next(common_speculative * spec, llama_token id_last); +void common_speculative_cancel(common_speculative * spec); + // print statistics about the speculative decoding void common_speculative_print_stats(const common_speculative * spec); diff --git a/docs/development/gemma4-assistant-tensor-inventory.md b/docs/development/gemma4-assistant-tensor-inventory.md new file mode 100644 index 000000000000..cda975e3bdfb --- /dev/null +++ b/docs/development/gemma4-assistant-tensor-inventory.md @@ -0,0 +1,60 @@ +# Gemma 4 Assistant (MTP drafter) — tensor inventory (Stage 0) + +This document records the expected tensor layout for `Gemma4AssistantForCausalLM` / `model_type: gemma4_assistant`, cross-referenced with the MLX-VLM reference implementation (PR #1112). + +## HF config highlights (`google/gemma-4-26B-A4B-it-assistant`) + +- `architectures`: `["Gemma4AssistantForCausalLM"]` +- `model_type`: `gemma4_assistant` +- `backbone_hidden_size`: 2816 (must match paired target Gemma 4 backbone) +- `use_ordered_embeddings`: typically `false` for 26B-A4B / 31B (dense tied LM head); `true` for E2B/E4B (centroid `MaskedEmbedder`) +- `num_centroids`: 2048 (when ordered embeddings) +- `centroid_intermediate_top_k`: 32 +- Nested `text_config`: 4 layers, `layer_types` e.g. three `sliding_attention` + one `full_attention`, `attention_k_eq_v: true`, `vocab_size: 262144`, `hidden_size: 1024`, etc. + +## Expected safetensors names (HF / MLX) + +| Logical component | Typical HF / checkpoint name | GGUF name (this fork) | +|-------------------|------------------------------|------------------------| +| Pre-projection | `pre_projection.weight` | `mtp.pre_projection.weight` | +| Post-projection | `post_projection.weight` | `mtp.post_projection.weight` | +| Token embeddings (inner) | `model.embed_tokens.weight` | `token_embd.weight` | +| Final norm | `model.norm.weight` | `output_norm.weight` | +| Per-layer | `model.layers.{i}.*` | `blk.{i}.*` (via `tensor_mapping.py`) | +| Centroid head (E2B/E4B) | `masked_embedding.centroids.weight` | `mtp.centroids.weight` | +| Token ordering | `masked_embedding.token_ordering` | `mtp.token_ordering` (I32) | +| LM head (if untied) | `lm_head.weight` | `output.weight` | + +## MLX reference files + +- `mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py` — forward, `draft_block`, `sanitize` +- `mlx_vlm/speculative/drafters/gemma4_assistant/masked_embedder.py` — sparse centroid LM head +- `mlx_vlm/speculative/drafters/gemma4_assistant/masks.py` — SWA / full masks + +## Notes for llama.cpp port + +1. MTP consumes **target K/V** for the last sliding-attention and last full-attention layers; the assistant has **no independent KV** in the reference design. The initial C++ integration uses the assistant as a loadable arch with standard KV for bring-up; full KV sharing is wired through `--spec-type mtp` and remains dependent on target/draft pairing and cache layout. +2. `attention_k_eq_v: true` maps to missing `blk.*.attn_v.weight` (V taken from K), matching Gemma 4 handling in `gemma4-iswa.cpp`. +3. Greedy parity with the target requires byte-identical drafting; validate with `tests/test-speculative-mtp` when model paths are available. + +## Centroid / ordered embeddings LM head (E2B / E4B) — implemented in MTP + +When `use_ordered_embeddings=true`, the MTP step uses the masked LM head (same idea as HF `Gemma4AssistantMaskedEmbedder` / MLX `MaskedEmbedder`), implemented in `gemma4_mtp_build_one_step()` in `src/models/gemma4-assistant.cpp`: + +1. `centroid_logits = mul_mat(mtp.centroids, h)` → `[n_centroids, n_tokens]`. +2. `top_k(centroid_logits, centroid_intermediate_top_k)` → centroid indices (I32). +3. `token_ordering` is viewed as `[vsc, n_centroids]` with `vsc = n_vocab / n_centroids` so each centroid column lists `vsc` token ids (matches HF flat layout); `get_rows` gathers candidate token ids for the top centroids. +4. `get_rows(token_embd, ids)` then `mul_mat` yields sparse logits over those ids. +5. Full `[n_vocab, n_tokens]` logits are built for host sampling / tracing: base tensor filled with **`-1e30f`**, then **SET_ROWS** writes the sparse logits at the selected ids (MLX-style `min(selected)-1` is optional if parity tests require it). + +### GGUF / converter layout + +- **`mtp.centroids.weight`**: HF checkpoint stores `[n_centroids, n_embd]` rows; numpy is written as-is so that after GGUF’s dimension packing the loader sees **`[n_embd, n_centroids]`**, matching `create_tensor(..., {n_embd, n_c})` and `mul_mat(mtp_centroids, h)` like the dense tied head. +- **`mtp.token_ordering.weight`**: stored as **I32**, length `n_vocab`. + +Verification: `scripts/verify-gemma4-assistant-gguf.py` checks embedding-length parity, and when `use_ordered_embeddings` is true asserts centroid shape, I32 ordering, and `n_vocab % n_centroids == 0`. + +### Tests / scripts + +- Edge smoke (optional): `LLAMA_MTP_TEST_TARGET_EDGE` + `LLAMA_MTP_TEST_HEAD_EDGE` in `tests/test-speculative-mtp.cpp` (paths must exist locally; CI skips when unset). +- Server helpers: `scripts/run-gemma4-e4b-mtp-server.sh`, `scripts/run-gemma4-e2b-mtp-server.sh`. diff --git a/docs/development/pipeline-depth-2-pure-overlap.md b/docs/development/pipeline-depth-2-pure-overlap.md new file mode 100644 index 000000000000..a2052d4e1707 --- /dev/null +++ b/docs/development/pipeline-depth-2-pure-overlap.md @@ -0,0 +1,35 @@ +# Pipeline depth-2: pure MTP overlap (implementation notes) + +## Summary + +`llama-server` overlaps MTP draft compute with per-iteration post-accept work by +splitting `llama_decode_mtp_async` / `llama_decode_mtp_wait` across server loop +iterations: + +1. End of iteration (after `common_speculative_accept`, prompt rollback, `slot.sampled` update, `llama_memory_seq_rm`): `common_speculative_prepare_next(slot.spec, slot.sampled)`. +2. Start of next iteration: `common_speculative_draft` calls `llama_decode_mtp_wait` first when a request is pending. + +No optimistic token is used; drafts always match the last accepted target token. + +## KV safety (Phase 0) + +- `decode_mtp_async` snapshots `h_prev` and uses `attn_pos = seq_pos_max` after `seq_rm`. +- MTP memory context uses `llama_kv_cache_iswa::init_mtp` / `mtp_slot_info` (assistant read path). +- Target `llama_decode` only appends at positions strictly after the rolled-back window; cells at positions `≤ attn_pos` are not rewritten before `_wait`. + +## API + +- `void common_speculative_prepare_next(common_speculative * spec, llama_token id_last);` — no-op for non-MTP implementations. +- Virtual `prepare_next` on `common_speculative_state` (default empty). + +## Bench decision gate + +Plan target: **≥ +15%** median tps vs baseline on f16-mtp short prompt (3 runs), via `.scratch/bench-mtp.sh`. + +**Status**: not executed in this workspace (no GGUF fixtures under `.scratch/`). Run locally after `scripts/run-gemma4-mtp-server.sh` (or your paths), then: + +```bash +bash .scratch/bench-mtp.sh 127.0.0.1:8080 256 +``` + +Compare accept rate (`draft_n` / `draft_n_accepted`) to ensure parity with pre-change runs. diff --git a/docs/speculative.md b/docs/speculative.md index 041ff58038de..894043f5f13c 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -13,6 +13,133 @@ The `llama-server` application supports several implementations of speculative d A much smaller model (called the _draft model_) generates drafts. A draft model is the most used approach in speculative decoding. +### Gemma 4 MTP assistant (`mtp`) + +For **Gemma 4** targets with the **gemma4_assistant** (MTP head) GGUF, use `--spec-type mtp`. The assistant is **not** a second `llama_context`: weights are loaded into the target model via `llama_model_load_mtp_from_file` (done automatically when using the server/CLI init path). Cross-attention in the MTP graph reads **K/V from the target KV cache** (shared full/sliding layers). + +- Prefer **`--mtp-head /path/to/assistant.gguf`** for clarity; **`--model-draft` (`-md`)** is accepted as a backward-compatible alias (same path field). +- The draft block size \(B\) is **`--draft-block-size`** (the head proposes `B - 1` tokens per round; default 4). +- **`--gpu-layers-draft` / `-ngld`** and **`-ctkd` / `-ctvd`** still apply to how the **assistant tensors** are placed and typed when the assistant GGUF is loaded; the target uses `-ngl` and `-ctk`/`-ctv`. + +Example (paths illustrative). **TurboQuant** KV on the target: `-ctk`/`-ctv`. Assistant-side cache types follow the draft flags if you use them for offload/quant selection. + +```sh +llama-server \ + -m /path/to/gemma-4-target.gguf \ + --mtp-head /path/to/gemma-4-assistant.gguf \ + --spec-type mtp \ + --draft-block-size 4 \ + -c 16384 \ + -ngl 99 -ngld 99 \ + -ctk turbo3 -ctv turbo3 \ + -ctkd turbo3 -ctvd turbo3 \ + -fa on \ + --host 127.0.0.1 --port 8080 +``` + +Repo helper (defaults under `.scratch/`): `scripts/run-gemma4-mtp-server.sh`. + +#### Async MTP draft pipeline + +The MTP draft head runs on a dedicated `ggml_backend_sched` (`sched_mtp`) and a +worker thread, isolated from the target's scheduler. This is exposed via two C +APIs: + +```c +LLAMA_API int32_t llama_decode_mtp_async( + struct llama_context * ctx, + llama_seq_id seq_id, + llama_pos attn_pos, + llama_token last_token, + const float * h_prev, + int32_t n_steps); + +LLAMA_API int32_t llama_decode_mtp_wait( + struct llama_context * ctx, + llama_token * out_drafts, + float * out_h_prev_last); +``` + +Contract: + +- At most **one in-flight request per context**. Calling `_async` while a + previous request has not been waited on returns `-7`. +- `h_prev` is copied into the request, so the caller may free or reuse it + immediately after `_async` returns. +- Target KV positions ≤ `attn_pos` must remain stable until `_wait` returns. + The current append-only KV cache satisfies this as long as no eviction or + overlapping `seq_rm` happens between submit and wait. +- `llama_decode_mtp(...)` is preserved as a backward-compatible synchronous + facade (= `_async` immediately followed by `_wait`). + +Why use the async pair? + +- **Graph-cache isolation**: MTP graph reuse is no longer invalidated by target + decode resets. On Gemma 4 + Q4_K_XL, this alone delivered **~+8% throughput** + in single-slot benchmarks (95.3 → 102.8 tps at `--draft-block-size 3`), with + identical accept rate. +- **Pipeline depth-2 (pure overlap, `llama-server`)**: after target + sample/accept and `llama_memory_seq_rm`, the server calls + `common_speculative_prepare_next(spec, last_accepted_token)`, which submits + `llama_decode_mtp_async` for the *next* round. The blocking + `llama_decode_mtp_wait` is deferred to the start of the next + `common_speculative_draft`, so MTP work can overlap with post-accept + bookkeeping and token I/O. This uses the real sampled token (no optimistic + guess). **KV contract**: submit uses `attn_pos` after `seq_rm`; the next + `llama_decode` appends only positions `> attn_pos`, so backbone cells read by + MTP remain stable until `_wait` (append-only cache). Stale in-flight requests + are drained in `common_speculative_begin` and on skip / param-mismatch paths. +- **In-graph argmax**: the MTP graph publishes a `ggml_argmax` of the final + logits (I32 [1]) via `llm_graph_result::get_argmax()`. Per draft step the host + reads back 4 bytes (one token id) instead of the full F32 [n_vocab] logits + row, and the serial CPU argmax over `n_vocab` is gone. The full logits are + still computed in-graph and can be fetched on demand by passing a non-null + `out_logits` to the synchronous `llama_decode_mtp(...)` API (diagnostic / + legacy path). On Gemma 4 + Q4_K_XL this delivered an additional **~+2-3% + throughput** (`109.5 → 112.5 tps` at `n=128`, `95.8 → 97.8 tps` at `n=512`), + with bit-identical greedy drafts. +- **Future work**: optimistic last-token prediction could hide an additional + `llama_decode` latency on hits but risks wrong drafts on misses. + +#### Diagnostic: per-draft acceptance trace + +Set `LLAMA_MTP_ACC_TRACE` to enable a per-iteration NDJSON tracer in +`common/speculative.cpp`. Each `mtp_draft` event records `iter`, `path` +(`sync` / `lazy` / `skip-streak` / `skip-nsteps`), `seq_id`, `id_last`, +`h_idx`, `attn_pos`, `n_steps`, `h_l2` (L2 norm of the input hidden state), +and the drafted token ids; each `mtp_accept` event records the iteration +counter, `n_accepted`, and `n_drafted_prev`. The host can pair them by `iter` +to reconstruct per-position acceptance, MTP h_prev stability, and +selection-bias breakdowns by `h_idx`. + +```sh +# write trace to stderr +LLAMA_MTP_ACC_TRACE=1 ./llama-server [...] +# write trace to a file +LLAMA_MTP_ACC_TRACE=/tmp/mtp.ndjson ./llama-server [...] +``` + +#### Throughput tuning (Edge assistants / heavy backends) + +Ordered-embeddings MTP (`use_ordered_embeddings`) runs a centroid LM head (`top_k`, routed `get_rows`) per draft step; the greedy graph still materializes a full-vocab logits row (masked fill + scatter) so argmax and optional full-row export stay consistent with verify. + +- **`LLAMA_MTP_SKIP_STREAK_THRESHOLD`**: **off by default** (`0` / unset). If set to `1`–`32`, after that many consecutive batches with **zero** accepted **MTP** drafts, MTP drafting is skipped for one verify-only batch, then retried. +- Helper scripts `scripts/run-gemma4-e4b-mtp-server.sh` and `scripts/run-gemma4-e2b-mtp-server.sh` support **`MTP_PRESET`**: `throughput` (block 2 / max 6), `lift` (block 3 / max 8), `balanced`, or `quality`. Override with `DRAFT_BLOCK_SIZE`, `DRAFT_MAX`, and optionally `LLAMA_MTP_SKIP_STREAK_THRESHOLD`. + +Disabled at zero overhead (no env var or value `0` / empty); when enabled the +extra cost is one `n_bb`-wide L2 reduction per draft and a small NDJSON write. + +#### Reconverting `gemma4_assistant` from Hugging Face + +Example (paths relative to repo root): + +```sh +python convert_hf_to_gguf.py .scratch/gemma-4-26B-A4B-it-assistant \ + --outfile .scratch/gemma-assistant-mtp.gguf --outtype f16 +``` + +Use the resulting GGUF as `--mtp-head` (or `-md`) with `--spec-type mtp`. Older assistant GGUFs with `token_embd.weight` first axis 2816 (backbone width) instead of 1024 will fail load; run `scripts/verify-gemma4-assistant-gguf.py` on the file to check. + ### n-gram Cache (`ngram-cache`) An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences. diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d4e7ac799af2..6ca5a9531249 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -153,6 +153,12 @@ class LLM: DENSE_FEAT_IN_SIZE = "{arch}.{dense}_feat_in" DENSE_FEAT_OUT_SIZE = "{arch}.{dense}_feat_out" + N_CENTROIDS = "{arch}.n_centroids" + CENTROID_TOP_K = "{arch}.centroid_top_k" + N_EMBD_BACKBONE = "{arch}.n_embd_backbone" + USE_ORDERED_EMBEDDINGS = "{arch}.use_ordered_embeddings" + REQUIRES_TARGET_ARCH = "{arch}.requires_target_arch" + class Attention: HEAD_COUNT = "{arch}.attention.head_count" HEAD_COUNT_KV = "{arch}.attention.head_count_kv" @@ -182,6 +188,7 @@ class Attention: KEY_LENGTH_SWA = "{arch}.attention.key_length_swa" VALUE_LENGTH_SWA = "{arch}.attention.value_length_swa" SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" + K_EQ_V = "{arch}.attention.k_eq_v" SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" @@ -430,6 +437,7 @@ class MODEL_ARCH(IntEnum): GEMMA3 = auto() GEMMA3N = auto() GEMMA4 = auto() + GEMMA4_ASSISTANT = auto() GEMMA_EMBEDDING = auto() STARCODER2 = auto() RWKV6 = auto() @@ -862,6 +870,11 @@ class MODEL_TENSOR(IntEnum): NEXTN_HNORM = auto() NEXTN_SHARED_HEAD_HEAD = auto() NEXTN_SHARED_HEAD_NORM = auto() + # Gemma 4 MTP assistant (GGUF names: mtp.*) + MTP_PRE_PROJECTION = auto() + MTP_POST_PROJECTION = auto() + MTP_CENTROIDS = auto() + MTP_TOKEN_ORDERING = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -945,6 +958,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GEMMA3: "gemma3", MODEL_ARCH.GEMMA3N: "gemma3n", MODEL_ARCH.GEMMA4: "gemma4", + MODEL_ARCH.GEMMA4_ASSISTANT: "gemma4_assistant", MODEL_ARCH.GEMMA_EMBEDDING: "gemma-embedding", MODEL_ARCH.STARCODER2: "starcoder2", MODEL_ARCH.RWKV6: "rwkv6", @@ -1407,6 +1421,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm", MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", + MODEL_TENSOR.MTP_PRE_PROJECTION: "mtp.pre_projection", + MODEL_TENSOR.MTP_POST_PROJECTION: "mtp.post_projection", + MODEL_TENSOR.MTP_CENTROIDS: "mtp.centroids", + MODEL_TENSOR.MTP_TOKEN_ORDERING: "mtp.token_ordering", } MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { @@ -2481,6 +2499,26 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PER_LAYER_PROJ_NORM, MODEL_TENSOR.PER_LAYER_POST_NORM, ], + MODEL_ARCH.GEMMA4_ASSISTANT: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.MTP_PRE_PROJECTION, + MODEL_TENSOR.MTP_POST_PROJECTION, + MODEL_TENSOR.MTP_CENTROIDS, + MODEL_TENSOR.MTP_TOKEN_ORDERING, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.LAYER_OUT_SCALE, + ], MODEL_ARCH.GEMMA_EMBEDDING: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index a101382719d0..38b62d8bc4fa 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -799,6 +799,24 @@ def add_clamp_kqv(self, value: float) -> None: def add_shared_kv_layers(self, value: int) -> None: self.add_uint32(Keys.Attention.SHARED_KV_LAYERS.format(arch=self.arch), value) + def add_n_centroids(self, value: int) -> None: + self.add_uint32(Keys.LLM.N_CENTROIDS.format(arch=self.arch), value) + + def add_centroid_top_k(self, value: int) -> None: + self.add_uint32(Keys.LLM.CENTROID_TOP_K.format(arch=self.arch), value) + + def add_n_embd_backbone(self, value: int) -> None: + self.add_uint32(Keys.LLM.N_EMBD_BACKBONE.format(arch=self.arch), value) + + def add_use_ordered_embeddings(self, value: bool) -> None: + self.add_bool(Keys.LLM.USE_ORDERED_EMBEDDINGS.format(arch=self.arch), value) + + def add_attention_k_eq_v(self, value: bool) -> None: + self.add_bool(Keys.Attention.K_EQ_V.format(arch=self.arch), value) + + def add_requires_target_arch(self, value: str) -> None: + self.add_string(Keys.LLM.REQUIRES_TARGET_ARCH.format(arch=self.arch), value) + # if input is array, true means SWA and false means full_attention for each layer def add_sliding_window_pattern(self, value: int | Sequence[bool]) -> None: key = Keys.Attention.SLIDING_WINDOW_PATTERN.format(arch=self.arch) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index c2235cb3b614..0ae11396a7d6 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2268,6 +2268,19 @@ class TensorNameMap: MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: ( "model.layers.{bid}.shared_head.norm", ), + + MODEL_TENSOR.MTP_PRE_PROJECTION: ( + "pre_projection", + ), + MODEL_TENSOR.MTP_POST_PROJECTION: ( + "post_projection", + ), + MODEL_TENSOR.MTP_CENTROIDS: ( + "masked_embedding.centroids", + ), + MODEL_TENSOR.MTP_TOKEN_ORDERING: ( + "masked_embedding.token_ordering", + ), } # architecture-specific block mappings diff --git a/include/llama.h b/include/llama.h index 52eb212466b5..4445dc3778e8 100644 --- a/include/llama.h +++ b/include/llama.h @@ -499,6 +499,20 @@ extern "C" { size_t n_paths, struct llama_model_params params); + // Gemma 4 MTP: load gemma4_assistant GGUF into a gemma4 target (call after llama_model_load_from_file, before llama_init_from_model). + // Returns 0 on success. + LLAMA_API int llama_model_load_mtp_from_file( + struct llama_model * model, + const char * path_mtp, + struct llama_model_params params); + + LLAMA_API const struct llama_model * llama_model_get_mtp_assistant(const struct llama_model * model); + + LLAMA_API bool llama_model_has_mtp_assistant(const struct llama_model * model); + + // Backbone hidden size for MTP input (0 if no MTP assistant is loaded). + LLAMA_API uint32_t llama_model_mtp_n_embd_backbone(const struct llama_model * model); + LLAMA_API void llama_model_save_to_file( const struct llama_model * model, const char * path_model); @@ -574,6 +588,16 @@ extern "C" { // Returns label of classifier output by index (