Architecture: shared self-MTP draft context ownership for parallel slots

[joelfarthing]

Hi @ikawrakow, I've been investigating the TAG_SERVER_SPEC_REWORK TODO in
examples/server/server-context.cpp:3549 on current origin/main:

rework to have a single draft llama_context shared across all slots

#1119 introduced the server TODO. I saw the current -np guardrail in examples/server/server-context.cpp and the related #1804/#1901 history; this is an architecture question about whether the right next step is to replace that guardrail with server-owned shared self-MTP context plus per-slot speculative state, not a request to bypass the guardrail as-is.

I have this in a public experimental branch, but I'm intentionally not treating that branch as a PR:

https://github.com/joelfarthing/ik_llama.cpp/tree/filament/mtp-shared-context-experiment-20260530

I also split out a smaller local lifecycle branch that I can push if this is the right direction:

https://github.com/joelfarthing/ik_llama.cpp/tree/filament/mtp-shared-context-lifecycle

The lifecycle split is meant to be the architectural core only:

• server_context owns the shared self-MTP companion context.
• Slots keep the per-slot state: seq id, draft tokens, validation indices, recurrent checkpoint state, hidden-state validity, sampler state, and counters.
• Target-only state changes explicitly clear/trim/shift companion state across release, cancellation, prompt reuse, prompt-cache restore, slot restore, context shift, KV clear, and rewind paths.
• The current upstream -np > 1 MTP guardrail is preserved for the general case; the shared self-MTP path is an explicit env-gated exception rather than a broad speculative-decoding rewrite.

The range-checkpoint result is why I think this is worth discussing. In the shared parallel self-MTP shape, the old fallback path made recurrent checkpoint saves expensive. Deferring checkpoint saves to isolated validation ranges and
using gpu-fallback instead of CPU gave the cleanest signal on MoE Qwen3.6:

• MoE Qwen3.6 median max request time went from 2.20s to 1.65s.
• MoE checkpoint save went from 630.872ms to 28.425ms.
• Dense Qwen3.6 also reduced checkpoint-save cost, but I'm not claiming a stable dense wall-time win because restore/redecode still dominates there.
• Normal single-slot MTP already selects per-step and should stay on that path.

I'm intentionally holding back the batched companion-draft prototype for now. It does build locally as a separate branch, but I don't think it should be the first review surface. It touches batching, sequence ids, hidden-state row shape, and shared companion state all at once. I'd rather figure out the ownership and lifecycle shape first.

The three questions I would most like your feedback on:

1. Is "server owns the shared self-MTP context, slots own only per-slot speculative state" the right ownership model for ik_llama?

2. For recurrent shared parallel MTP, is it reasonable to force gpu-fallback until per-step checkpoint buffers are truly per-slot, and should the range-checkpoint change be a stacked follow-up or part of the first shared context PR?

3. Would you prefer the first reviewable PR to be only the env-gated lifecyce branch, with batched companion draft kept out, or would you split this differently?

I can reshape the branch around whatever direction and review path would be easiest. My initial goal here is to avoid dumping a broad experimental branch into review before the ownership model is settled.

Thanks!

[SamuelOliveirads]

I'm not Ik, but I can contribute a little to the discussion.

1. I think the way forward for speculative could involve shared-context, so if you only want to work with MTP, it would be helpful to start with that step as long as the seq-route is guaranteed. I think mainline already has this part set up, I didn't check closely but it could be a reference.
2. To keep the PR simple, it might be possible or if you prefer to work only with MTP via GLM or Gemma 4, you can isolate the flow generally and then work on specific features to expand the support.
3. I would argue that it should deliver some immediate value instead of isolating it in a variable. As I mentioned in topic 2 this could be a PR that implements the shared context and then expands it to more models or methods for speculative.

[joelfarthing]

Thanks Sam, this is really helpful. I’ll reshape around shared context as the first review surface with concrete MTP value, rather than treating it as only an env-gated experiment.

To make sure I’m interpreting the sequence-route requirement correctly: would you consider it satisfied if the shared companion path uses slot.id as the draft sequence id throughout, and mirrors target/companion cleanup on release, cancel, prompt reuse, context shift, prompt-cache restore, and rollback?

For immediate value, would you prefer the first PR target one narrow MTP path, such as GLM/Gemma 4, with batched companion draft and range-checkpoint performance work left as follow-ups?

Also, if there’s a specific mainline path you had in mind as the reference, a pointer would be much appreciated.

[SamuelOliveirads]

Yes, I think the approach you described for slot.id would be a good direction to take. You can start with GLM if you prefer, it’s the simplest in terms of specifics, then expand to the other architectures, and only then, if that’s still your goal, expand to all the specs.

For reference, you can use this PR: ggml-org/llama.cpp#22838