Pass the target as the MTP draft's ctx_other so Gemma 4 sidecars load (#91)

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),

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)

c_src/llama_cpp_ex/llama_nif.cpp

@@ -11,6 +11,7 @@
 #include <cmath>
 #include <cstring>
 #include <cerrno>
+#include <optional>
 
 // 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<fine::ResourcePtr<LlamaContext>> ctx_other)
 {
     auto params = llama_context_default_params();
     params.n_ctx           = static_cast<uint32_t>(n_ctx);
@@ -891,6 +893,12 @@ context_create(
     // Speculative decoding / MTP
     params.ctx_type = static_cast<enum llama_context_type>(ctx_type);
     params.n_rs_seq = static_cast<uint32_t>(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<LlamaContext>(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);

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<LlamaContext> ctx_other;
+
     // Shape of the last *successful* llama_decode on this context.
     //
     // `sampler_sample_at/3` takes its index straight from Elixir, and

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

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

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 <target> -hfd <draft> --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 <target> -hfd <draft> --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

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)
 

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