spec : cost-aware MTP draft length, MTP draft vocabulary subset, flat get_rows and shape-aware graph key - #58
Open
SimonTeixidor wants to merge 5 commits into
Open
SimonTeixidor wants to merge 5 commits into
SimonTeixidor wants to merge 5 commits into
Conversation
…che key get_rows: for float rows with ne00 <= 32, launch over (row, element) instead of one workgroup per row. f32 ne00 = 1 over 248320 rows measured 986 us -> 20 us on gfx1151. graph key: nodes[0] alone gives the same key to graphs that share one meta arena (different topologies, or one topology at a varying verify width), so those graphs never replay. Mix the node count, the last node and three node shapes into the key. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC
With --spec-draft-mtp-vocab N, the MTP context copies the LM head rows of token ids < N plus control and user-defined tokens into a smaller head when it is created. The qwen35 and qwen35moe MTP graphs compute the draft logits over that head and scatter them into a full-vocabulary row of -inf (ggml_fill + ggml_set_rows). Only the draft changes: the draft cannot propose other tokens, verification is unchanged. The subset is built only when token ids follow BPE merge order, so that low ids are the frequent tokens; otherwise the option is ignored with an error. Probe contexts with unallocated weights skip the build. New field llama_context_params::mtp_draft_vocab (0 = off). Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC
With draft-mtp and --spec-draft-n-min > 0, the draft length is chosen token by token between n_min and n_max. A position is kept, and the next one drafted, when its expected value in output tokens times the current time per output token covers its verify column cost and draft step cost. Acceptance per drafter confidence, draft step time, time per output token and the verify cost are all measured online, so no model-specific tuning is needed. The verify cost V(w) is timed per round from the end of the draft to accept() and kept per verify width with decay, on top of a level term that follows context depth. V is fit non-decreasing in the width with a weak linear prior, and the column cost of position j is the smoothed difference V(j+1) - V(j), clamped to a small positive minimum. A verify width next to the stopping point that has had almost no recent samples is drafted now and then, rate-limited. Only single-head MTP with its own KV cache (qwen35, qwen35moe). --spec-draft-n-max alone keeps the fixed draft length. A summary line per request is logged with the speculative statistics and the learned verify and column costs. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC
dzannotti
requested changes
Sep 15, 2026
Collaborator
There was a problem hiding this comment.
Correctness blocker: mtp_draft_vocab_build() stores the first subset in mutable llama_model state. The Qwen MTP graph builders select model.mtp_draft_head from that shared pointer and output shape, without checking the requesting context's mtp_draft_vocab.
Reproduction by inspection: create MTP context A with mtp_draft_vocab=N, then context B over the same model with the default 0 (or a different N). B silently uses A's subset, so its behavior depends on initialization order.
Please make the subset context/config-owned, or key an immutable cache by requested N and backend/buffer ownership, and gate graph use on that context's setting. Add regressions for A=N then B=0 and A=N1 then B=N2, including simultaneous contexts.
This also bundles draft-vocabulary API, cost controller, get_rows kernel, and CUDA graph cache-key work. Please split it; most is generic llama.cpp/CUDA work and belongs in upstream staging. Diff whitespace check passed. I have not independently reproduced the performance claim; resolve the correctness blocker first.
comment generated by my clanker Codex
The subset was stored in mutable llama_model state by the first MTP context that asked for one, and the qwen35 and qwen35moe MTP graphs used it whenever it existed. A second context over the same model with mtp_draft_vocab = 0, or a different N, silently drafted over the first context's subset. The subset is now an immutable llama_mtp_draft_vocab obtained with llama_model::mtp_draft_vocab_get(n_keep). The model keeps weak references keyed by (n_keep, buffer type) under a mutex, so contexts that ask for the same N at the same time share one copy, and it is freed with the last context holding it. Each context holds its own reference and passes it to its graphs through llm_graph_params; the MTP graphs use only that one, and graph reuse compares it. llama_cparams::mtp_draft_vocab records the N actually in use (0 when no subset was built: probe contexts with unallocated weights, a tokenizer that fails the merge-order guard, or a model whose MTP block has its own LM head, which now logs a warning instead of building an unused subset). Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC
Contexts over one model with different mtp_draft_vocab settings (N then 0, 0 then N, N1 then N2, the same N twice, two contexts created from two threads at once, a subset requested again after its last context was freed) are decoded interleaved with the same inputs. A context with N > 0 must have -inf exactly outside its own draft vocabulary and match a full-vocabulary context inside it; a context with N = 0 must have finite logits everywhere. Needs a qwen35 or qwen35moe model with an MTP head (-m), so it is built but not registered with ctest. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC
Author
|
Fixed the reported bug and added a test. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Branch
pr/mtp-costaware-draft(2eeae5f84), 3 commits onhalo/master69946438a.Overview
8a1f66cd3ggml-cuda : flat get_rows for small rows, shape-aware graph cache key. get_rows on float rows withne00 <= 32launches over (row, element), not one workgroup per row. The graph key adds node count, last node and three node shapes: graphs sharing a meta arena got one key fromnodes[0]and never replayed. Both unconditional.94c706a82spec : MTP draft vocabulary subset (--spec-draft-mtp-vocab N). The MTP head only scores token ids below N, which makes each draft step much cheaper while losing little acceptance. Qwen's BPE token ids follow merge order, so lower ids are roughly the more frequent tokens. On Qwen3.6-35B-A3B with N = 65536, the draft step went from 3.13 to 1.42 ms and acceptance from 74.4 % to 72.5 % (57.07 -> 63.98 t/s). Verification still uses the full vocabulary. For models whose ids do not follow merge order the flag is ignored with an error.2eeae5f84spec : cost-aware MTP draft length (--spec-draft-n-min > 0). Improves on the existing--spec-draft-adaptive, which sizes drafts from measured acceptance, by adding a cost model and per-token acceptance statistics. The draft step time, the verify cost per extra draft position and the time per output token are all measured online, so nothing needs tuning per model or context depth. For every draft token it estimates whether drafting one more is worth it, or whether the extra draft step and verify column are likely to cost more than the expected payoff. Enabled by--spec-draft-n-min, up to--spec-draft-n-max.Measurements
Stock
654803517vs this branch, same image.llama-serverwith production sampling (temperature 0.6, thinking on),llama-benchywith a 2048-token prompt and 2000 generated tokens, n = 4 per cell in mirrored blocks. Stock--spec-draft-n-min 1 --spec-draft-n-max 10 --spec-draft-adaptivevs this branch--spec-draft-n-min 1 --spec-draft-n-max 10, with and without--spec-draft-mtp-vocab 65536. Each cell: tg t/s (sd), acceptance, accepted tokens per draft step.The branch is ahead in every cell and gets more accepted tokens per draft step. Only cost-aware + subset at depth 32000 is statistically resolved. With sampled text, acceptance varies from 33 % to 82 % between runs of the same arm.
In an informal test with pi in non-interactive mode on three coding tasks, cost-aware drafting was 11 % faster than stock adaptive: 30.66 t/s against 27.54 t/s.
Correctness:
--spec-type draft-mtp --spec-draft-n-max 3, greedy: tokens, log-probs and draft counters identical on both.n_max 3diverge at step 97 on the 13480-token prompt.test-backend-ops -b ROCm0: GET_ROWS 215/215, SET_ROWS 159/159, expansion 4/4, scatter/gather at real shapes 6/6. Scatter vs old gather expansion: identical draft counters and output hashes.Raw results
llama-benchy --base-url http://127.0.0.1:18091/v1 --model q --pp 2048 --tg 2000 --exact-tg --depth D --runs 2 --no-warmup --skip-coherence --no-adapt-prompt --latency-mode api --tokenizer <local Qwen3.8-27B tokenizer>(0.4.0), depths 0 and 32000, prompts seeded per depth; one warm-up chat request per server. Server:-fa on -c 131072 -b 1024 -ub 1024 --parallel 1 -ngl 999 --load-mode none -ctkd q8_0 -ctvd q8_0,--temp 0.6 --top-p 0.95 --top-k 20 --min-p 0,qwen-sharp.jinja, thinking on, one server per arm. Drafted/accepted are/metricsdeltas. Arms ran in two blocks, A and B, in mirrored order.Per run:
Prompt processing:
llama-bench -m Qwen3.8-27B-UD-IQ4_XS.gguf -p 2048 -d 0,12000,32000 -b 1024 -ub 1024 -n 0 -r 2 -ngl 99 -fa on -ctk f16 -ctv f16 --load-mode none -o jsonl, blocks (stock, branch) and (branch, stock).Token generation: same with
-p 0 -n 128 -d 0,32000.Correctness commands:
llama-debug --save-logits -b 16384 -ub 512(fourtests/corpusfiles and a 13480-token concatenation);llama-perplexity -c 512 -b/-ub 512 --chunks 4KLD-base files (PPL prose 5.4246, code 1.9202, structured 2.4814, numeric 1.9349 on both builds); greedyllama-server/completion, 256 tokens, top-10 log-probs. Speculative checks at-b/-ub 512. Scatter/gather: Qwen3.8-27B,--spec-draft-n-max 3 --spec-draft-mtp-vocab 65536, temp 0, 2 x 512 tokens.Speculative draft counters with identical
--spec-type draft-mtp --spec-draft-n-max 3(identical on both builds): 27B 151/308, 189/198, 80/140; 35B-A3B 178/230, 81/136, 137/350 (code, long, prose). New flags on the 13480-token prompt match stockn_max 3over all 128 tokens. Scatter vs gather draft counters: 326/555, 328/548.Additional information
llama_model, used only whenn_outputs == 1.Requirements
654803517. The branch was then rebased onto69946438a, whose 6 new commits touch onlyggml/rocmfpx, one Vulkan shader andtests/CMakeLists.txt; not re-measured.llama-benchyand subset-head output not compared.llama-benchPP at depth 64000 not run;llama-benchyat depths 0 and 32000 only, n = 4 per cell.test-backend-opsnot run (GET_ROWS, SET_ROWS only).--parallel > 1(log line sums all slots) and on the CLI.https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC