Skip to content

feat(sdk): add MTP speculative decoding to the llama_cpp backend#1185

Closed
Mengsheng Wu (mengshengwu) wants to merge 2 commits into
mainfrom
feat/llama-cpp-mtp-spec
Closed

feat(sdk): add MTP speculative decoding to the llama_cpp backend#1185
Mengsheng Wu (mengshengwu) wants to merge 2 commits into
mainfrom
feat/llama-cpp-mtp-spec

Conversation

@mengshengwu

Copy link
Copy Markdown
Contributor

Summary

Adds Multi-Token Prediction (MTP) speculative decoding to the llama_cpp backend, wiring llama.cpp's common_speculative (draft-mtp) into the plugin.

  • Load an optional draft/MTP GGUF into an MTP context that shares the target KV cache (ctx_type = LLAMA_CONTEXT_TYPE_MTP, ctx_other = target), pinned to the target's device.
  • Replace the per-token decode loop with a draft → batch-verify → accept loop when a draft model is set; plain decoding is untouched otherwise.
  • Report draft accept counts through the profile.

Public API (FFI change, mirrored across Go / Python / Android per CONTRIBUTING §4):

  • geniex_ModelConfig: spec_draft_model, spec_n_draft
  • geniex_ProfileData: draft_n_total, draft_n_accepted

CLI: --draft-model / --draft-tokens; the acceptance rate is shown in the run summary. Speculative decoding is opt-in (only when a draft model is given) and llama_cpp-only (other runtimes warn and ignore it). Setup failure falls back to plain decoding.

Test plan

  • Builds on Snapdragon X Elite (arm64, qairt off): SDK bridge + geniex_llama_cpp plugin + geniex-bench link clean against the pinned llama.cpp submodule.
  • End-to-end MTP run on X Elite GPU (Adreno/OpenCL) with a Gemma-4 12B main + mtp- drafter: draft-mtp initializes, drafts/verifies/accepts, produces coherent output, and reports the acceptance rate. Baseline (no --draft-model) and MTP paths both exercised from one binary.
  • Draft model on a non-llama_cpp runtime / missing MTP head → warns and runs plain decoding, no crash (graceful fallback verified when the drafter failed to load).

Note: measured acceptance on the community gemmable-4-12b drafter is very low (~5%), so it does not speed up on this pair — expected for a poorly-aligned drafter; the plugin faithfully reports what the drafter produces. Speedup depends on drafter quality, which is out of scope here.

Wire llama.cpp's common_speculative (draft-mtp) into the llama_cpp plugin:
load an optional draft/MTP GGUF into an MTP context that shares the target
KV cache (pinned to the target's device), and replace the per-token decode
loop with a draft-verify loop when a draft model is configured. Draft accept
counts are reported in the profile.

New public config on geniex_ModelConfig (spec_draft_model, spec_n_draft) and
geniex_ProfileData (draft_n_total, draft_n_accepted); mirrored across the Go,
Python, and Android FFI surfaces. CLI gains --draft-model / --draft-tokens and
prints the acceptance rate.

Signed-off-by: Mengsheng Wu <mengshengwu@outlook.com>
Signed-off-by: Mengsheng Wu <mengshengwu@outlook.com>
@mengshengwu

Copy link
Copy Markdown
Contributor Author

Test results

Summary: MTP speculative decoding runs end-to-end on X Elite GPU (Adreno/OpenCL); the draft-verify loop drafts, verifies, accepts, and reports the acceptance rate. Verified on hardware, which the Linux CI build can't exercise.

Model: Mia-AiLab/Gemmable-4-12B-MTP-GGUF — main gemmable-4-12b-Q4_K_M.gguf (7.4 GB) + mtp- drafter (317 MB). Built with -DGENIEX_PLUGIN_QAIRT=OFF, run via geniex-bench --device gpu.

Run decode draft acceptance
baseline (no --draft-model) 6.1 tok/s
MTP (--draft-model ... --draft-tokens 3) 1.6 tok/s 5.4% (14/258)

[spec] draft acceptance = 0.05426 (14 accepted / 258 generated) — the plugin faithfully reports what the drafter produces.

Notes:

  • The MTP draft context must be pinned to the target's device; without it the drafter grabbed HTP0 and the shared-KV MTP graph crashed (ggml-hex: HTP0 dspcall : bad response). Fixed by threading the device selection into setup_speculative — GPU run is clean after that.
  • Acceptance on this community drafter is very low, so it does not speed up on this pair. That's a drafter-quality property, not an implementation issue — a well-aligned Gemma-4 MTP head is where the speedup comes from. The plumbing, accounting, and reporting are what this PR validates.
  • Graceful fallback confirmed: a drafter that fails to load (or a non-llama_cpp runtime) warns and continues with plain decoding.

@mengshengwu

Copy link
Copy Markdown
Contributor Author

CI: all build + lint + test jobs green across linux/windows/android. The single test-python / linux-arm64 failure is an infra flake — apt mirror timeout (ports.ubuntu.com unreachable, exit 100) during dep setup, before any test ran; test-python / windows-arm64 passed on the same FFI. Not related to this diff.

@mengshengwu

Copy link
Copy Markdown
Contributor Author

Superseded by #1195 (generalizes to all speculative types).

@mengshengwu
Mengsheng Wu (mengshengwu) deleted the feat/llama-cpp-mtp-spec branch July 17, 2026 10:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant