diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d8cf13..6b032b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,36 @@ # Changelog +## Unreleased + +### Fixed + +- **Gemma 4 MTP sidecars failed at `MTP.init/2` with `"failed to create + context"`** (#91). llama.cpp's `gemma4-assistant` head builds its draft + context *over* the target's — it aliases the target's KV cells and reads the + target's token embeddings — and the constructor throws + `Gemma4Assistant requires ctx_other to be set` when + `llama_context_params.ctx_other` is null. The NIF built every context from + `llama_context_default_params()` and never exposed the field, so the Qwen 3.8 + sidecar path from v0.8.44 (whose `qwen35` head ignores the pointer) worked + while a Gemma 4 pair could not get past context creation. Reproduced with + Unsloth `gemma-4-E4B-it-Q4_K_M.gguf` + `MTP/mtp-gemma-4-E4B-it-Q8_0.gguf` + on Metal: validation passes (head `n_layer_nextn == 4`, widths match at + 2560), then `llama_init_from_model` fails. `Context.create/2` gains + `:ctx_other` (a `%Context{}`, forwarded as `params.ctx_other`; the NIF holds + the peer's resource so it outlives the context that aliases it), and + `MTP.init/2` passes the target as the draft's `:ctx_other` unconditionally — + the same thing upstream's `common_speculative_init_result` does, so the + in-file and Qwen sidecar paths are unchanged (the pointer is only read for + `gemma4-assistant`, `eagle3` and `dflash`). With the fix the Gemma 4 pair + loads, drafts (53% acceptance at `n_draft: 3`, greedy) and the + `:mtp_sidecar` suite passes **6/6** against it, including greedy + equivalence with plain decoding on the target. One assertion in that suite + pinned Qwen 3.8 rather than the binding — `timing_us.ckpt > 0` — and now + keys off the target's `seq_rm` kind: `:full` targets must be billed for + snapshots, everything else must not. Creating a draft context on a Gemma 4 + sidecar *without* a target still fails; that is llama.cpp's contract, and + `Context.create/2` documents it. + ## v0.8.49 This section also covers v0.8.44 (b10435, PR #86), v0.8.45 (b10582, #87), diff --git a/README.md b/README.md index 1451e87..76258f3 100644 --- a/README.md +++ b/README.md @@ -646,6 +646,7 @@ Upstream llama.cpp implements more speculative types behind the same `common_spe ### Models with MTP heads - [`ggml-org/Qwen3.8-27B-GGUF`](https://huggingface.co/ggml-org/Qwen3.8-27B-GGUF) — **sidecar layout**: the target (`Qwen3.8-27B-Q4_K_M.gguf`, ~18 GB) carries *no* head, and `mtp-Qwen3.8-27B-Q4_0.gguf` (~1.6 GB) carries nothing else. Load both and pass the head as `draft_model:`. +- [`unsloth/gemma-4-E4B-it-GGUF`](https://huggingface.co/unsloth/gemma-4-E4B-it-GGUF) — **sidecar layout**, different head architecture: the target (`gemma-4-E4B-it-Q4_K_M.gguf`, ~5 GB) carries no head and `MTP/mtp-gemma-4-E4B-it-Q8_0.gguf` (~0.1 GB) is a `gemma4-assistant` that shares the target's KV cache. Same `draft_model:` call; `MTP.init/2` wires the target in as the draft's `:ctx_other`, which this architecture requires and Qwen's ignores. 53% greedy acceptance at `n_draft: 3` on M4 Max. Fixed in #91. - [`ggml-org/Qwen3.6-35B-A3B-MTP-GGUF`](https://huggingface.co/ggml-org/Qwen3.6-35B-A3B-MTP-GGUF) (recommended: `Q4_K_M`, ~21 GB) - [`ggml-org/Qwen3.6-27B-MTP-GGUF`](https://huggingface.co/ggml-org/Qwen3.6-27B-MTP-GGUF) - [`unsloth/Qwen3.6-35B-A3B-MTP-GGUF`](https://huggingface.co/unsloth/Qwen3.6-35B-A3B-MTP-GGUF) diff --git a/c_src/llama_cpp_ex/llama_nif.cpp b/c_src/llama_cpp_ex/llama_nif.cpp index e629541..0f3b2df 100644 --- a/c_src/llama_cpp_ex/llama_nif.cpp +++ b/c_src/llama_cpp_ex/llama_nif.cpp @@ -11,6 +11,7 @@ #include #include #include +#include // The ggml RPC backend is opt-in at build time (LLAMA_RPC=1). GGML_USE_RPC is // set by the Makefile, not inherited from cmake: ggml puts it on the `ggml` @@ -845,7 +846,8 @@ context_create( bool kv_unified, // Speculative decoding / MTP int64_t ctx_type, - int64_t n_rs_seq) + int64_t n_rs_seq, + std::optional> ctx_other) { auto params = llama_context_default_params(); params.n_ctx = static_cast(n_ctx); @@ -891,6 +893,12 @@ context_create( // Speculative decoding / MTP params.ctx_type = static_cast(ctx_type); params.n_rs_seq = static_cast(n_rs_seq); + // Peer context for architectures whose draft shares state with the target. + // gemma4-assistant throws at construction without it (llama-context.cpp, + // LLM_ARCH_GEMMA4_ASSISTANT); qwen35's constructor ignores it. Upstream's + // common_speculative_init_result sets it for every draft context, so the + // MTP module always passes the target here. + params.ctx_other = ctx_other ? (*ctx_other)->ctx : nullptr; // For embedding models, n_ubatch must equal n_batch if (embeddings) { @@ -904,6 +912,9 @@ context_create( auto res = fine::make_resource(ctx, model); res->kv_unified = kv_unified; + if (ctx_other) { + res->ctx_other = std::move(*ctx_other); + } return fine::Ok(std::move(res)); } FINE_NIF(context_create, ERL_NIF_DIRTY_JOB_CPU_BOUND); diff --git a/c_src/llama_cpp_ex/llama_nif.h b/c_src/llama_cpp_ex/llama_nif.h index 11481b0..9ae8632 100644 --- a/c_src/llama_cpp_ex/llama_nif.h +++ b/c_src/llama_cpp_ex/llama_nif.h @@ -60,6 +60,13 @@ class LlamaContext { // side effect, so it must never be called on a live context. bool kv_unified = false; + // The peer context passed as llama_context_params.ctx_other, or null. + // Held so the peer outlives this context: architectures that use it + // (gemma4-assistant) build this context's KV cache *over* the peer's — + // llama_kv_cache aliases `other->v_cells_impl` and shares layer tensors — + // so freeing the peer first would leave this context reading freed cells. + fine::ResourcePtr ctx_other; + // Shape of the last *successful* llama_decode on this context. // // `sampler_sample_at/3` takes its index straight from Elixir, and diff --git a/docs/release-guide.md b/docs/release-guide.md index 73557d2..a3016a5 100644 --- a/docs/release-guide.md +++ b/docs/release-guide.md @@ -161,6 +161,14 @@ LLAMA_SMOKE_MTP_MODEL=~/Downloads/Qwen3.8-27B-Q4_K_M.gguf \ LLAMA_SMOKE_MTP_DRAFT_MODEL=~/Downloads/mtp-Qwen3.8-27B-Q4_0.gguf \ mix test --include mtp_sidecar +# Same tag, other head architecture: gemma4-assistant needs the target as the +# draft's ctx_other (#91), which Qwen's head ignores — so a Qwen pair alone +# cannot tell you that wiring still works. +GGML_METAL_NO_RESIDENCY=1 \ +LLAMA_SMOKE_MTP_MODEL=~/Downloads/gemma-4-E4B-it-Q4_K_M.gguf \ +LLAMA_SMOKE_MTP_DRAFT_MODEL=~/Downloads/mtp-gemma-4-E4B-it-Q8_0.gguf \ + mix test --include mtp_sidecar + # :rpc_live needs an RPC build AND a reachable worker, and must run with no # model tag beside it — see test/test_helper.exs for why combining them aborts. # The worker can be local: another BEAM running LlamaCppEx.RPC.Server, or diff --git a/lib/llama_cpp_ex/context.ex b/lib/llama_cpp_ex/context.ex index 97aa7c4..cb3da06 100644 --- a/lib/llama_cpp_ex/context.ex +++ b/lib/llama_cpp_ex/context.ex @@ -59,7 +59,8 @@ defmodule LlamaCppEx.Context do :embeddings, :pooling_type, :ctx_type, - :n_rs_seq + :n_rs_seq, + :ctx_other ] @doc """ @@ -139,6 +140,15 @@ defmodule LlamaCppEx.Context do * `:n_rs_seq` - Number of recurrent-state snapshots per sequence to retain for partial rollback of speculative drafts. `0` (default) disables rollback. For an MTP draft context, set this to your intended max draft length (e.g. `3`). + * `:ctx_other` - Peer `t:t/0` this context is paired with, forwarded as + `llama_context_params.ctx_other`. Some MTP head architectures build the draft + context *over* the target's: `gemma4-assistant` shares the target's KV cache + and token embeddings and refuses to construct without it, so a draft + `Context.create/2` on a Gemma 4 sidecar fails with + `"Gemma4Assistant requires ctx_other to be set"` unless the target is passed + here. `qwen35` ignores it. `LlamaCppEx.MTP.init/2` always passes the target, + matching upstream's `common_speculative_init_result`. The peer is kept alive + for as long as this context exists. Defaults to `nil`. """ @spec create(LlamaCppEx.Model.t(), keyword()) :: {:ok, t()} | {:error, String.t()} @@ -183,6 +193,7 @@ defmodule LlamaCppEx.Context do # Speculative decoding / MTP ctx_type = Keyword.get(opts, :ctx_type, :default) |> ctx_type_to_int() n_rs_seq = Keyword.get(opts, :n_rs_seq, 0) + ctx_other = Keyword.get(opts, :ctx_other) case LlamaCppEx.NIF.context_create( model_ref, @@ -212,7 +223,8 @@ defmodule LlamaCppEx.Context do swa_full, kv_unified, ctx_type, - n_rs_seq + n_rs_seq, + ctx_other_ref(ctx_other) ) do {:ok, ref} -> {:ok, %__MODULE__{ref: ref, model: model}} {:error, _} = error -> error @@ -313,4 +325,10 @@ defmodule LlamaCppEx.Context do defp ctx_type_to_int(:default), do: 0 defp ctx_type_to_int(:mtp), do: 1 defp ctx_type_to_int(n) when is_integer(n), do: n + + # A bare reference is deliberately not accepted: the peer must be a live + # `%Context{}` so the caller cannot hand the NIF a ref whose resource has + # already been collected, and so the FunctionClauseError names the option. + defp ctx_other_ref(nil), do: nil + defp ctx_other_ref(%__MODULE__{ref: ref}), do: ref end diff --git a/lib/llama_cpp_ex/mtp.ex b/lib/llama_cpp_ex/mtp.ex index 1fa7304..5883800 100644 --- a/lib/llama_cpp_ex/mtp.ex +++ b/lib/llama_cpp_ex/mtp.ex @@ -34,10 +34,12 @@ defmodule LlamaCppEx.MTP do **In the target GGUF** (e.g. `ggml-org/Qwen3.6-35B-A3B-MTP-GGUF`) — pass just the model, as above. - **In a sidecar GGUF** — pass it as `:draft_model`. This is how Qwen 3.8 ships: - `Qwen3.8-27B-Q4_K_M.gguf` carries no head at all (`n_layer_nextn == 0`) and - `mtp-Qwen3.8-27B-Q4_0.gguf` carries nothing else. It is the binding's - equivalent of upstream's `-hf -hfd --spec-type draft-mtp`. + **In a sidecar GGUF** — pass it as `:draft_model`. This is how Qwen 3.8 and + Gemma 4 ship: `Qwen3.8-27B-Q4_K_M.gguf` carries no head at all + (`n_layer_nextn == 0`) and `mtp-Qwen3.8-27B-Q4_0.gguf` carries nothing else; + likewise `gemma-4-E4B-it-Q4_K_M.gguf` plus `mtp-gemma-4-E4B-it-Q8_0.gguf`. It + is the binding's equivalent of upstream's + `-hf -hfd --spec-type draft-mtp`. {:ok, target} = LlamaCppEx.load_model("Qwen3.8-27B-Q4_K_M.gguf", n_gpu_layers: 999, load_mtp: true) @@ -46,6 +48,16 @@ defmodule LlamaCppEx.MTP do {:ok, mtp} = LlamaCppEx.MTP.init(target, draft_model: head, n_draft: 1) + The two head architectures differ in what the draft context needs from the + target. Qwen's (`qwen35`) is self-contained. Gemma 4's (`gemma4-assistant`) + shares the target's KV cache and token embeddings, so its draft context is + built *over* the target's — `init/2` passes the target as the draft's + `:ctx_other` (see `LlamaCppEx.Context.create/2`) for both, as upstream does. + A consequence worth knowing: a Gemma 4 sidecar cannot be given a context of + its own — `Context.create(head)` without `:ctx_other` fails with + `"Gemma4Assistant requires ctx_other to be set"`. That is llama.cpp's + contract, not a binding limitation. + > #### Speculation is not always a win on hybrid models {: .warning} > > A model that mixes recurrent (SSM) layers with attention ones — Qwen 3.8 is @@ -212,12 +224,17 @@ defmodule LlamaCppEx.MTP do defp do_init(model, head_model, opts, n_draft) do base_ctx_opts = forwardable_context_opts(opts) main_opts = Keyword.merge(base_ctx_opts, ctx_type: :default) - # Match upstream server: MTP draft context is created with n_rs_seq=0. - # The MTP impl handles state rollback internally via cached hidden - # states (pending_h / verify_h), not via recurrent-state snapshots. - draft_opts = Keyword.merge(base_ctx_opts, ctx_type: :mtp, n_rs_seq: 0) with {:ok, main_ctx} <- Context.create(model, main_opts), + # Match upstream's common_speculative_init_result: the draft context is + # created with n_rs_seq=0 (the MTP impl rolls back via cached hidden + # states, not recurrent-state snapshots) and with the target as + # ctx_other — unconditionally. gemma4-assistant builds its KV cache over + # the target's and throws at construction without the peer; qwen35 + # ignores it. Passing it always is what makes a Gemma 4 sidecar load + # through the same path as a Qwen 3.8 one. + draft_opts = + Keyword.merge(base_ctx_opts, ctx_type: :mtp, n_rs_seq: 0, ctx_other: main_ctx), {:ok, mtp_ctx} <- Context.create(head_model, draft_opts), {:ok, spec_ref} <- LlamaCppEx.NIF.speculative_init(main_ctx.ref, mtp_ctx.ref, n_draft) do diff --git a/lib/llama_cpp_ex/nif.ex b/lib/llama_cpp_ex/nif.ex index 99f6363..b989445 100644 --- a/lib/llama_cpp_ex/nif.ex +++ b/lib/llama_cpp_ex/nif.ex @@ -98,7 +98,8 @@ defmodule LlamaCppEx.NIF do _swa_full, _kv_unified, _ctx_type, - _n_rs_seq + _n_rs_seq, + _ctx_other ), do: :erlang.nif_error(:not_loaded) diff --git a/test/mtp_model_test.exs b/test/mtp_model_test.exs index 8daedd9..9a5d9e3 100644 --- a/test/mtp_model_test.exs +++ b/test/mtp_model_test.exs @@ -295,8 +295,12 @@ defmodule LlamaCppEx.MTPSidecarTest do # A hybrid target (Qwen 3.8: SSM layers beside attention ones) cannot roll back # part of a sequence natively, so the loop snapshots the recurrent state every # iteration. That cost is the difference between speculation paying off and not, - # so it gets its own bucket rather than hiding inside :other. - test "timing_us reports a ckpt bucket", %{session: session} do + # so it gets its own bucket rather than hiding inside :other. An attention-only + # target (Gemma 4) trims any range natively and must never pay for a snapshot — + # the NIF gates the checkpoint on common_context_can_seq_rm == FULL at init. + test "timing_us bills ckpt only when the target cannot trim partially", %{ + session: session + } do assert {:ok, _} = MTP.generate(session, "Count to ten:", max_tokens: 24, temp: 0.0) timing = MTP.stats(session).timing_us @@ -305,8 +309,15 @@ defmodule LlamaCppEx.MTPSidecarTest do assert Map.has_key?(timing, key), "timing_us is missing #{inspect(key)}" end - assert timing.ckpt > 0, - "a hybrid target should have paid for at least one recurrent-state snapshot" + case LlamaCppEx.TestModels.seq_rm_kind(:mtp) do + :full -> + assert timing.ckpt > 0, + "a hybrid target should have paid for at least one recurrent-state snapshot" + + kind -> + assert timing.ckpt == 0, + "a #{inspect(kind)} target trims natively and must not be billed for snapshots" + end # The named buckets are carved out of total, never billed twice on top of it. assert timing.draft + timing.verify + timing.sample + timing.ckpt <= timing.total