Pass target context as ctx_other so Gemma4 E4B MTP can init

CHANGELOG.md

@@ -95,6 +95,16 @@ The one tag that is not green is `:mtp_cancel`, and it moved: see Changed.
   target/draft hidden-width mismatch — upstream compares those with a
   `GGML_ASSERT`, which is an unconditional `ggml_abort` and would take the VM
   down instead of returning an error.
+- **`Context.create/2` accepts optional `:ctx_other`.** Pass an existing
+  `%Context{}` and the NIF forwards it as `llama_context_params.ctx_other`.
+  `MTP.init/2` always sets it on the draft. Gemma4 E4B's `gemma4-assistant`
+  sidecar needs that link at construction; Qwen's constructor leaves
+  `cparams.ctx_other` as `nullptr`, so the NIF default (omitted → `nullptr`)
+  still works for them.
+- **`:mtp_sidecar` gained a second fixture pair** behind
+  `LLAMA_SMOKE_MTP_E4B_MODEL` / `LLAMA_SMOKE_MTP_E4B_DRAFT_MODEL`. When those
+  vars are unset the E4B module skips, so `--include mtp_sidecar` with only the
+  Qwen pair stays green.
 - **`stats/1` reports `timing_us.ckpt`.** Recurrent-state save/restore, which
   only hybrid models pay, was previously folded into `:other` — a bucket whose
   documented cause is Metal GPU-sync waits. On Qwen 3.8 (48 SSM layers to 16

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:`.
+- [`ggml-org/gemma-4-E4B-it-GGUF`](https://huggingface.co/ggml-org/gemma-4-E4B-it-GGUF) — **sidecar layout**: the target (`gemma-4-E4B-it-Q4_K_M.gguf`) is `gemma4` with *no* nextn head; `mtp-gemma-4-E4B-it-Q8_0.gguf` is a separate `gemma4-assistant` sidecar. Same `draft_model:` API. llama.cpp requires the target as `ctx_other` for `gemma4-assistant`; `MTP.init/2` always passes it (Qwen's constructor ignores it).
 - [`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)
@@ -660,7 +661,7 @@ Acceptance on a 0.8B target is not representative of production throughput —
 drafting is nearly as expensive as decoding at that size — so use these to
 exercise the path, not to measure it.
 
-A regular (non-MTP) quant will fail at `LlamaCppEx.MTP.init/2` — some GGUF in the pair must contain the MTP head's tensors. To check a file before loading it, look for a `*.nextn_predict_layers` key and `blk.N.nextn.*` tensors in its metadata. When the publisher ships the head separately (Qwen 3.8), that sidecar is the file with those tensors and the target legitimately has none; pass it as `draft_model:` rather than looking for a combined build.
+A regular (non-MTP) quant will fail at `LlamaCppEx.MTP.init/2` — some GGUF in the pair must contain the MTP head's tensors. To check a file before loading it, look for a `*.nextn_predict_layers` key and `blk.N.nextn.*` tensors in its metadata. The E4B sidecar also has those, and is additionally a `gemma4-assistant` architecture. When the publisher ships the head separately, that sidecar is the file with the head tensors and the target legitimately has none; pass it as `draft_model:` rather than looking for a combined build.
 
 The model must also be loaded with `load_mtp: true` (see below). Upstream gates those tensors behind a load-time flag that defaults to off, and they cannot be attached afterwards, so `MTP.init/2` refuses a model loaded without it rather than letting the omission surface later as `verify decode failed: code=-1`.
 
@@ -741,6 +742,34 @@ a head whose hidden width does not match the target's — that last one because
 upstream compares the two with a `GGML_ASSERT`, which aborts the VM rather than
 failing the call.
 
+#### Sidecar head: Gemma4 E4B (`gemma4` target + `gemma4-assistant` draft)
+
+Same `draft_model:` API as Qwen 3.8, two files. llama.cpp requires the target
+as `ctx_other` for `gemma4-assistant`. `MTP.init/2` always passes it.
+
+```elixir
+:ok = LlamaCppEx.init()
+
+{:ok, target} =
+  LlamaCppEx.load_model(
+    Path.expand("~/Downloads/gemma-4-E4B-it-Q4_K_M.gguf"),
+    n_gpu_layers: 999,
+    load_mtp: true
+  )
+
+{:ok, head} =
+  LlamaCppEx.load_model(
+    Path.expand("~/Downloads/mtp-gemma-4-E4B-it-Q8_0.gguf"),
+    n_gpu_layers: 999,
+    load_mtp: true
+  )
+
+{:ok, mtp} = LlamaCppEx.MTP.init(target, draft_model: head, n_draft: 3, n_ctx: 8192)
+
+{:ok, text} = LlamaCppEx.MTP.generate(mtp, "Explain MTP in one paragraph.", max_tokens: 200)
+IO.puts(text)
+```
+
 #### Synchronous generate (collect to a string)
 
 ```elixir

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,9 @@ 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);
+    if (ctx_other) {
+        params.ctx_other = (*ctx_other)->ctx;
+    }
 
     // For embedding models, n_ubatch must equal n_batch
     if (embeddings) {
@@ -902,7 +907,12 @@ context_create(
         return fine::Error(std::string("failed to create context"));
     }
 
-    auto res = fine::make_resource<LlamaContext>(ctx, model);
+    fine::ResourcePtr<LlamaContext> res;
+    if (ctx_other) {
+        res = fine::make_resource<LlamaContext>(ctx, model, std::move(*ctx_other));
+    } else {
+        res = fine::make_resource<LlamaContext>(ctx, model);
+    }
     res->kv_unified = kv_unified;
     return fine::Ok(std::move(res));
 }

c_src/llama_cpp_ex/llama_nif.h

@@ -44,6 +44,9 @@ class LlamaContext {
 public:
     llama_context* ctx;
     fine::ResourcePtr<LlamaModel> model;
+    // Fine keep for the raw llama_context* passed as params.ctx_other.
+    // gemma4-assistant stores that pointer in cparams; Qwen does not.
+    fine::ResourcePtr<LlamaContext> ctx_other;
 
     // Reusable explicit batch for the decode-side NIFs (batch_eval,
     // batch_eval_sample, decode_token, prefill), allocated once on first use
@@ -112,6 +115,10 @@ class LlamaContext {
     LlamaContext(llama_context* c, fine::ResourcePtr<LlamaModel> m)
         : ctx(c), model(std::move(m)) {}
 
+    LlamaContext(llama_context* c, fine::ResourcePtr<LlamaModel> m,
+                 fine::ResourcePtr<LlamaContext> other)
+        : ctx(c), model(std::move(m)), ctx_other(std::move(other)) {}
+
     // Returns the reusable batch with capacity for at least n tokens
     // (per-token seq-id capacity 1 — all decode builders use single-seq
     // entries). Contents are stale; the caller fills 0..n-1 and n_tokens.
@@ -127,6 +134,10 @@ class LlamaContext {
     ~LlamaContext() {
         if (batch_capacity > 0) llama_batch_free(batch);
         if (ctx) llama_free(ctx);
+        ctx = nullptr;
+        // llama_free the draft first: gemma4-assistant still holds
+        // cparams.ctx_other until that returns. Then drop the Fine keep.
+        ctx_other = {};
     }
 
     LlamaContext(const LlamaContext&) = delete;

docs/release-guide.md

@@ -152,6 +152,16 @@ 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
 
+# The Qwen pair is required: MTPSidecarTest has no skip and calls path!/1.
+# Add the E4B pair to also run MTPE4BSidecarTest. Unset E4B vars skip that
+# module, so the Qwen-only command above stays green.
+GGML_METAL_NO_RESIDENCY=1 \
+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 \
+LLAMA_SMOKE_MTP_E4B_MODEL=~/Downloads/gemma-4-E4B-it-Q4_K_M.gguf \
+LLAMA_SMOKE_MTP_E4B_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

@@ -25,9 +25,13 @@ defmodule LlamaCppEx.Context do
   """
 
   @enforce_keys [:ref, :model]
-  defstruct [:ref, :model]
+  defstruct [:ref, :model, :ctx_other]
 
-  @type t :: %__MODULE__{ref: reference(), model: LlamaCppEx.Model.t()}
+  @type t :: %__MODULE__{
+          ref: reference(),
+          model: LlamaCppEx.Model.t(),
+          ctx_other: t() | nil
+        }
 
   @tuning_option_keys [
     :n_threads,
@@ -59,7 +63,8 @@ defmodule LlamaCppEx.Context do
     :embeddings,
     :pooling_type,
     :ctx_type,
-    :n_rs_seq
+    :n_rs_seq,
+    :ctx_other
   ]
 
   @doc """
@@ -138,7 +143,13 @@ defmodule LlamaCppEx.Context do
       decoding via `LlamaCppEx.MTP`.
     * `: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`).
+      MTP draft context, use `0` — the MTP implementation handles rollback
+      internally via cached hidden states (`pending_h` / `verify_h`), not
+      recurrent-state snapshots.
+    * `:ctx_other` - An existing `%Context{}` whose raw pointer is passed as
+      `llama_context_params.ctx_other`. Optional; omitted or `nil` is
+      `nullptr`. `gemma4-assistant` requires it at construction. `MTP.init/2`
+      always sets it on the draft.
 
   """
   @spec create(LlamaCppEx.Model.t(), keyword()) :: {:ok, t()} | {:error, String.t()}
@@ -183,39 +194,50 @@ 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)
-
-    case LlamaCppEx.NIF.context_create(
-           model_ref,
-           n_ctx,
-           n_batch,
-           n_ubatch,
-           n_threads,
-           n_threads_batch,
-           embeddings,
-           pooling_type,
-           n_seq_max,
-           type_k,
-           type_v,
-           flash_attn,
-           offload_kqv,
-           op_offload,
-           rope_scaling_type,
-           rope_freq_base,
-           rope_freq_scale,
-           yarn_ext_factor,
-           yarn_attn_factor,
-           yarn_beta_fast,
-           yarn_beta_slow,
-           yarn_orig_ctx,
-           attention_type,
-           no_perf,
-           swa_full,
-           kv_unified,
-           ctx_type,
-           n_rs_seq
-         ) do
-      {:ok, ref} -> {:ok, %__MODULE__{ref: ref, model: model}}
-      {:error, _} = error -> error
+    ctx_other = Keyword.get(opts, :ctx_other)
+
+    case ctx_other_for_nif(ctx_other) do
+      {:error, _} = error ->
+        error
+
+      {:ok, ctx_other_ref, ctx_other_struct} ->
+        case LlamaCppEx.NIF.context_create(
+               model_ref,
+               n_ctx,
+               n_batch,
+               n_ubatch,
+               n_threads,
+               n_threads_batch,
+               embeddings,
+               pooling_type,
+               n_seq_max,
+               type_k,
+               type_v,
+               flash_attn,
+               offload_kqv,
+               op_offload,
+               rope_scaling_type,
+               rope_freq_base,
+               rope_freq_scale,
+               yarn_ext_factor,
+               yarn_attn_factor,
+               yarn_beta_fast,
+               yarn_beta_slow,
+               yarn_orig_ctx,
+               attention_type,
+               no_perf,
+               swa_full,
+               kv_unified,
+               ctx_type,
+               n_rs_seq,
+               ctx_other_ref
+             ) do
+          {:ok, ref} ->
+            {:ok, %__MODULE__{ref: ref, model: model, ctx_other: ctx_other_struct}}
+
+          {:error, _} = error ->
+            error
+        end
     end
   end
 
@@ -231,9 +253,9 @@ defmodule LlamaCppEx.Context do
   Returns the number of recurrent-state snapshots per sequence available for
   partial rollback of speculative drafts.
 
-  `0` means the context does not support partial rollback (e.g. a regular target
-  context with `n_rs_seq: 0`). For an MTP draft context created with
-  `n_rs_seq: N`, this returns at most `N`.
+  `0` means the context does not support partial rollback. MTP drafts are
+  created with `n_rs_seq: 0`; rollback is via cached hidden states, not
+  recurrent-state snapshots.
   """
   @spec n_rs_seq(t()) :: non_neg_integer()
   def n_rs_seq(%__MODULE__{ref: ref}), do: LlamaCppEx.NIF.context_n_rs_seq(ref)
@@ -313,4 +335,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
+
+  defp ctx_other_for_nif(nil), do: {:ok, nil, nil}
+  defp ctx_other_for_nif(%__MODULE__{} = ctx), do: {:ok, ctx.ref, ctx}
+
+  defp ctx_other_for_nif(other),
+    do: {:error, ":ctx_other must be a LlamaCppEx.Context, got: #{inspect(other)}"}
 end

lib/llama_cpp_ex/model.ex

@@ -199,8 +199,9 @@ defmodule LlamaCppEx.Model do
 
   @doc """
   Returns the output-side embedding width — the row width an MTP draft head
-  consumes. Equal to `n_embd/1` for every architecture currently in tree; it is
-  a distinct number because `LlamaCppEx.MTP` matches it across the target and a
+  consumes. Usually equal to `n_embd/1`; for `gemma4-assistant` the checkpoint
+  has a small `n_embd` but `n_embd_out` is the target's hidden width. It is a
+  distinct number because `LlamaCppEx.MTP` matches it across the target and a
   separate drafter GGUF.
   """
   @spec n_embd_out(t()) :: integer()

lib/llama_cpp_ex/mtp.ex

@@ -46,6 +46,12 @@ defmodule LlamaCppEx.MTP do
 
       {:ok, mtp} = LlamaCppEx.MTP.init(target, draft_model: head, n_draft: 1)
 
+  **Gemma4 E4B sidecar** — same `:draft_model` API: a `gemma4` target plus a
+  `gemma4-assistant` head. `init/2` always passes the target as `:ctx_other`
+  on the draft. `gemma4-assistant` keeps that pointer (shared KV and the
+  target's token embeddings); Qwen's constructor leaves `cparams.ctx_other`
+  as `nullptr`.
+
   > #### Speculation is not always a win on hybrid models {: .warning}
   >
   > A model that mixes recurrent (SSM) layers with attention ones — Qwen 3.8 is
@@ -112,17 +118,18 @@ defmodule LlamaCppEx.MTP do
 
     * `:draft_model` - A separate `LlamaCppEx.Model` holding the MTP head, for
       checkpoints that ship it as a sidecar GGUF rather than inside the target
-      file (Qwen 3.8 is the current example: `Qwen3.8-27B-Q4_K_M.gguf` plus
-      `mtp-Qwen3.8-27B-Q4_0.gguf`). It must be loaded with `load_mtp: true`.
+      file (Qwen 3.8: `Qwen3.8-27B-Q4_K_M.gguf` plus `mtp-Qwen3.8-27B-Q4_0.gguf`;
+      Gemma4 E4B: a `gemma4` target plus a `gemma4-assistant` head). It must be
+      loaded with `load_mtp: true`.
       Defaults to `nil`, meaning the head is expected inside the target model
       and the draft context is built against it.
     * `:n_draft` - Max draft tokens generated per iteration. Defaults to `3`.
       Larger values mean fewer model forward passes but lower per-iteration
       acceptance; 2–4 is the sweet spot in practice.
     * `:n_ctx` - Context size for both contexts. Defaults to `2048`.
-    * Any `LlamaCppEx.Context` option (e.g. `:n_threads`, `:flash_attn`,
-      `:type_k`/`:type_v`, `:offload_kqv`). The same options are applied to
-      both the target and draft contexts.
+    * Any `Context` tuning option (e.g. `:n_threads`, `:flash_attn`,
+      `:type_k`/`:type_v`, `:offload_kqv`). Applied to both contexts.
+      `:ctx_other` is not a caller option; the draft always gets the target.
 
   Returns `{:ok, %MTP{}}` or `{:error, reason}`.
   """
@@ -218,7 +225,8 @@ defmodule LlamaCppEx.MTP do
     draft_opts = Keyword.merge(base_ctx_opts, ctx_type: :mtp, n_rs_seq: 0)
 
     with {:ok, main_ctx} <- Context.create(model, main_opts),
-         {:ok, mtp_ctx} <- Context.create(head_model, draft_opts),
+         {:ok, mtp_ctx} <-
+           Context.create(head_model, Keyword.merge(draft_opts, ctx_other: main_ctx)),
          {:ok, spec_ref} <-
            LlamaCppEx.NIF.speculative_init(main_ctx.ref, mtp_ctx.ref, n_draft) do
       {:ok,

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)