llama : switch the split mode of a live context

[Piggidragon]

Places the model weights again for a different split mode while the context stays alive, so a
request can process its prompt under one split and generate under the other.

Supersedes #58, which did the same thing entirely inside llama-completion by
tearing down and rebuilding everything. That version answered whether the idea is worth anything;
this one is the library version its "for later" section sketched.

Why

Tensor parallelism splits every layer across the devices and pays for a collective per layer. The
cost of that collective scales with the batch, the benefit does not: at a batch of one there is
almost nothing to exchange and both devices work on the same token, while at a prompt-sized batch
the exchange dominates. A layer split is the mirror image - no collective at all, but the devices
run one after the other, so a single token gets no parallelism.

Prefill wants the layer split, generation wants the tensor split, and nothing in the design forces
one choice for a whole request.

Measured on an RTX 4070 + RTX 3060, Qwen3.8-27B-UD-IQ2_M, a 17055-token prompt plus 128 generated
tokens at -c 20480 with an f16 cache. Wall clock is the whole request, model load included:

config	prefill	generation	wall
-sm layer	1000.44 t/s	24.77 t/s	25 s
-sm tensor	576.12 t/s	29.14 t/s	37 s
-psm layer -sm tensor	991.04 t/s	28.86 t/s	27 s
-psm tensor -sm layer (control, wrong way round)	574.64 t/s	24.81 t/s	41 s

The switch keeps 99.1% of layer's prefill and 99.0% of tensor's generation: 27% off the wall clock
against -sm tensor at the same generation rate, or +16.5% generation against -sm layer for two
seconds. The reversed control is slower than either pure mode, which is the check that the gain is
the direction and not the switch.

The switch itself carried 1216 MiB of state and placed the weights of the 9.6 GB model in 2.21 s
from a warm page cache.

What this adds

LLAMA_API bool llama_context_set_split_mode(
        struct llama_context * ctx,
         enum llama_split_mode split_mode,
                  const float * tensor_split);

LLAMA_API bool llama_model_set_split_mode(
          struct llama_model * model,
         enum llama_split_mode split_mode,
                  const float * tensor_split);

LLAMA_API enum llama_split_mode llama_model_get_split_mode(const struct llama_model * model);
LLAMA_API const char * llama_split_mode_name(enum llama_split_mode split_mode);

llama_context_set_split_mode keeps the context object and everything the caller holds on to. The
memory, the logits of the last decode and the output ids are carried across, so the caller can
sample straight after the switch. It goes through four steps:

1. Take the state out. llama_state_get_data writes a layout-independent form of the memory,
   which is what makes the carry work at all: the same bytes restore into a cache that is split by
   layer or split by head. The host copy of the output buffer comes out the same way.
2. Place the weights again. llama_model::set_split_mode frees the current placement first and
   then loads the weights under the new one, so the devices never hold both at once. The model
   object survives, only the tensors inside it are replaced.
3. Build the context again. The backends, the compute scheduler, the output buffer and the
   memory module all follow the devices of the model, so all four are rebuilt through the same code
   the constructor uses - it was factored into init_backends, init_memory and init_sched for
   that.
4. Put the state back, and re-resolve the automatic choices that depend on where an op can run.

A mode that does not fit must not lose the request. If the new placement or the rebuild fails,
the previous split mode is placed again, the context is rebuilt on it and the state goes back in.
The call returns false and the caller still has a working context with its cache intact.

The tool

llama-completion gets --prefill-split-mode / -psm, which takes the same values as
--split-mode. Against #58, the tool side is now nine lines: the model, the context, the sampler
and the chat templates all survive, and the last prompt token no longer has to be held back to
produce logits, because the logits are carried.

-psm is rejected in interactive mode, where there is more than one prompt and the switch happens
only once.

Several sequences at once

The switch takes out and puts back the state of all sequences, not only the one that happens to
be generating, so a context that serves several requests keeps every slot's cache across it. Tested
with four sequences, with a unified cache and with one stream per sequence, in both directions:
each sequence continues exactly like the same sequence in a context that was in the target mode all
along and was handed the same state.

What the switch cannot do is run two split modes at once - the split mode belongs to the model - or
run while a llama_decode is in flight. Those are the real limits for a server, not the state carry.

-psm itself gives nothing to a server, and not because of a missing loop: it fires once, at the
boundary between prompt and generation of one linear request, and that boundary does not exist when
one slot prefills while another decodes. The library call is the part a server needs; the policy that
would decide when to use it is not here.

The reason to want it does survive parallelism, though. llama-batched-bench, same model, -c 32768,
f16 cache, -npp 2048 -ntg 128:

npl	pp layer	pp tensor	layer prefill edge	tg layer	tg tensor	tensor decode edge
1	989.4	598.5	+65.3%	26.47	30.42	+14.9%
2	1046.4	601.5	+74.0%	43.55	48.94	+12.4%
4	1074.1	599.9	+79.0%	59.58	68.63	+15.2%
8	1085.3	600.1	+80.9%	72.57	91.51	+26.1%

Both edges hold at every level, so the two phases still want different modes however many sequences
are in flight. Note that this does not match #58's measurement, where the decode edge shrank from
+33.5% to +23.4% as the batch grew - different quantisation and cache type, so the shape of the
tradeoff is setup-specific and worth measuring before relying on it.

On total throughput at this prompt-to-generation ratio the layer split wins everywhere (315 -> 596
t/s against 285 -> 452 t/s), because prefill dominates the mix. A server that wanted both halves
would have to batch its prefills and its decodes into separate phases first, which is a scheduler
change, not this one.

Device memory

The old placement is freed before the new one is asked for, so the devices never hold both. Peak per
device, sampled at 10 Hz over the whole run, Qwen3.8-27B-UD-IQ2_M at -c 20480:

config	GPU0 peak	GPU1 peak
-sm layer	5806 MiB	6188 MiB
-sm tensor	6418 MiB	6164 MiB
-psm layer -sm tensor	6362 MiB	6180 MiB

The switch costs the per-device maximum of the two modes, not their sum: 6362 against a layer
peak of 5806 and a tensor peak of 6418 on GPU0. The largest context that survives a switch is
therefore smaller than what either pure mode could hold.

When the new mode does not fit anyway, the allocation failure is caught rather than fatal. Checked by
holding 5976 MiB on GPU0 with a helper process, which leaves room for the layer split but not for the
tensor split:

run	result
-sm layer	generates
-sm tensor	cudaMalloc failed: out of memory at load, request lost
-psm layer -sm tensor	same out of memory during the switch, warns, generates under the layer split

The third row is the point: the 17055-token prompt was already in the cache, the tensor placement
failed on a 505 MiB allocation, the layer placement was put back, the context was built on it again
and the cache was restored - and the request finished. That path needs the ggml change below to be
reachable at all.

What is deliberately not here

• Moving tensors between devices instead of reading them again. The weights come from the model
  file, which is in the page cache after the first load. A true device-to-device migration would
  need both placements to exist at the same time, which raises the peak device memory - the opposite
  of what a switch under memory pressure needs.
• Re-splitting the cache in place. The serialized form is layout-independent and already
  exercised; it costs one host round trip of the cache and no new mechanism.
• An automatic policy. completion : allow the prompt to be processed under its own split mode #58's sketch says a cparam that switches on batch size should come after
  the explicit form has been used, and that a server with several slots has no single right answer.
  The explicit call is here; the policy is not.

Tests

tests/test-split-mode-switch.cpp runs against the dummy models that test-llama-archs generates.
Three of them are registered with ctest, following the pattern test-recurrent-state-rollback
already uses; --models DIR sweeps the whole directory, which is how every architecture gets
covered by hand. It skips itself when there is no GPU, so it costs nothing in CI.

The check that matters compares a context that switched into a mode against a context that was in
that mode all along and was handed the same memory state: both then decode the same token and
generate greedily, and the tokens must be identical. On top of that it checks that the serialized
state and the logits of the last decode are byte-identical across the switch, that a switch out and
back leaves the context exactly where it was, that a rejected split mode leaves a working context,
and that the model-level call on its own produces a model that generates like one loaded under that
split mode from the start.

The sweep over the generated models is 883 passed, 0 failed, 207 skipped over 109 models. The skips
are the architectures where a tensor split is not implemented, where the switch is refused and the
context has to keep working - that path is checked too.

One model is left out of the sweep: qwen3-dense.gguf cannot be run at all, because
a plain llama-bench -m qwen3-dense.gguf -sm tensor -n 4 aborts on origin/llama/dev with
GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) while generating. That is a pre-existing
tensor-split bug on a dummy model shape, unrelated to this change, and the only model of the 110
that is affected.

Also here

One ggml change, taken from #58: the meta buffer type asserted on an allocation failure instead of
returning nullptr, so an out of memory under -sm tensor aborted the process and no fallback was
reachable. It now frees what it holds and returns nullptr, the same contract the CUDA buffer type
follows. Without it the fallback above cannot run.

Note on flash attention

llama_init_from_model used to rewrite --flash-attn auto to enabled when the model was loaded
under a tensor split. The context now applies that rule itself, so the requested type survives and a
switch back to a layer split can resolve it again instead of running with flash attention forced on.
-sm tensor -fa auto behaves the same as before, checked against the same model.

Limitations

• The model must not be shared with a second context, and must have no LoRA adapter or control
  vector loaded. Both keep tensors of their own on the devices of the old placement, and the library
  cannot rebuild them.
• The model has to have been loaded from a path. A model loaded from a FILE * or from metadata has
  nowhere to read the weights from.
• Every llama_memory_t taken from the context before the switch is invalid after it.
• A -ts means a share of the layers under a layer split and a share of every tensor under a tensor
  split. The library call takes a tensor_split for that reason; -psm keeps the one the model has.
• common_fit_params still has no implementation for LLAMA_SPLIT_MODE_TENSOR, so a tensor phase
  is not fitted, only caught by the fallback if it does not fit.

---

AI disclosure: Claude Opus 5 wrote the implementation, the test and this description from a design I
own. I have read and understand every line and can defend it without it.

[Piggidragon]

Closing this. The measurements did not support the premise.

Switching the split mode only pays inside a narrow band - the prompt has to be long enough to beat staying in the tensor split, and the generation long enough to beat staying in the layer split. On a 4070 + 3060 with Qwen3.8-27B that band starts at roughly 3800 prompt and 500 generated tokens. Outside it, picking one mode and staying in it wins.

For an agentic workload, where the context is largely prefix-cached and the work is decode-dominated, the tensor split wins on total throughput and there is no phase boundary to exploit.

The one general fix that came out of this - the meta buffer type aborting on an allocation failure instead of reporting it - moved to #57.