Skip to content

completion : allow the prompt to be processed under its own split mode - #58

Closed
Piggidragon wants to merge 3 commits into
GenerelSchwerz:llama/devfrom
Piggidragon:mgpu/split-mode-switch
Closed

completion : allow the prompt to be processed under its own split mode#58
Piggidragon wants to merge 3 commits into
GenerelSchwerz:llama/devfrom
Piggidragon:mgpu/split-mode-switch

Conversation

@Piggidragon

@Piggidragon Piggidragon commented Aug 30, 2026

Copy link
Copy Markdown

Processes the prompt under one split mode and generates under another, because the two phases
do not want the same thing from a multi-GPU split.

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.

Measured on an RTX 4070 + RTX 3060, Qwen3.8-27B-UD-Q5_K_M, a 58.3k-token prompt at
-c 65536 with a q8_0 device-resident cache:

-sm prefill generation
tensor -ts 50,50 507.92 t/s 19.76 t/s
layer 807.57 t/s 14.80 t/s

Layer prefills 59% faster, tensor generates 34% faster.

Prefill wants the layer split, generation wants the tensor split, and the gap is large in
both directions. Nothing in the design forces one choice for a whole request.

What this does

--prefill-split-mode (-psm) takes the same values as --split-mode. The prompt is
processed under it, then the model is reloaded under --split-mode and the cache is carried
across, so generation runs under the other mode.

It is tool-level only - no library changes, apart from one allocation fix described under
Memory. The sequence state already survives a change of split mode, because
llama_state_seq_get_data writes a layout-independent form: this was verified separately by
writing a --prompt-cache under -sm layer and restoring it under -sm tensor, which
produced byte-identical output.

The switch:

  1. reads the sequence state with llama_state_seq_get_data,
  2. releases the model and context, so the two never hold device memory at once,
  3. reloads under the generation split mode, fitting the placement again for that mode,
  4. writes the state back with llama_state_seq_set_data,
  5. replays the prompt into the fresh sampler, so the repetition penalties are not empty.

The last prompt token is deliberately held back from the first phase. The state carries
the cache but not the logits, and the sampler needs those; letting the reloaded model decode
that one token produces them. Removing the last cache cell and replaying it instead does not
work on every model - Qwen3.5 rejects it with an M-RoPE position constraint.

Memory

The two modes do not place the same bytes on the same device, so the switch has to be sized
for both. Peak per device from nvidia-smi, same model, -c 65536, q8_0 cache. The cards are
12282 MiB and 12288 MiB:

config CUDA0 CUDA1
-sm layer 10143 11289
-sm tensor -ts 50,50 10931 10863
-psm layer -sm tensor -ts 50,50 10897 11281

The switch costs the per-device maximum of both modes, not the larger of the two totals:
device 0 has to hold what tensor needs while device 1 holds what layer needs. The largest
context that runs under -psm is therefore smaller than under either pure mode. At -c 98304
the same pattern holds (layer 10815/11833, tensor 11537/11467, -psm 11497/11825), and
neither mode dominates - at -c 114688 -ts 50,50 layer runs out on device 1 while tensor
still fits.

Three things follow, and this PR handles them:

The fit has to run again for the second mode. common_fit_params writes the placement it
finds back into params.tensor_split and the buft overrides. On the second load it then sees
them set and aborts with tensor_split already set by user, so the generation mode inherited
the placement of the prefill mode. That was an out of memory in the common case. The switch
now keeps the arguments and lets the fit run again, and pins n_ctx to the first phase so the
state is restored into a context of the same size.

A mode that does not fit must not lose the request. If the generation mode cannot be
loaded, the prompt is still in the state that was just read, so the tool reloads under the
prefill mode and finishes there with a warning. Before, an explicit -ts tuned for the
prefill mode killed the request after the prompt was already processed: -c 65536 -ts 57,43
fits under layer (11577/9839) and does not fit under tensor, and a 58.3k-token prompt was
discarded after 80 seconds of work.

The meta buffer type has to report an allocation failure. It asserted 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 already follows. This is the only change outside the tool.

Verified, -c 65536, q8_0 cache, peak per device, all five producing the same output:

run rc peak CUDA0/1 switched fell back
auto, no -ngl/-ts, -psm layer -sm tensor 0 10897/10829 yes no
auto, -sm tensor 0 10937/10867 - -
-ngl 99 -ts 50,50 -psm layer -sm tensor 0 10897/11281 yes no
-ngl 99 -ts 57,43 -psm layer -sm tensor 0 11871/9899 yes yes
-ngl 99 -sm layer 0 10143/11289 - -

Result

Same setup, wall clock for the whole request - the 58.3k-token prompt plus 128 generated
tokens, model load included:

config wall generation
-sm tensor 153 s 19.76 t/s
-sm layer 86 s 14.80 t/s
-psm layer -sm tensor 91 s 19.62 t/s
-psm tensor -sm layer (control, wrong way round) 156 s 14.82 t/s

It gets both halves: layer's prefill and tensor's generation, the latter within 0.7% of what
a pure tensor run reaches. Against -sm tensor that is 40% off the wall clock for the
same generation rate.

Against -sm layer it costs the 5 s reload and buys +33% generation, so it pays for
itself after about 300 generated tokens - below that, plain -sm layer is better. The
reversed control is slower than either pure mode, which is the check that the gain is the
direction and not the reload.

The switch carried 2085 MiB of state and the reload took 5.2 s, measured between the log line
and the first sampled token.

Cost

The switch is a full model reload. Measured at about 5 s for a 19.8 GB model from a warm page
cache, of which the device upload is the bulk. The state itself is cheap to move and, with
--no-kv-offload, never leaves host memory at all. It is held in host memory as one buffer
while the second model loads - 2085 MiB for the 58.3k-token prompt above, linear in the
context.

So the switch pays off when the prefill saving exceeds the reload, which for a prompt of this
size it does by an order of magnitude, and it does not pay off for short prompts. It is
opt-in and off by default.

Why -np 1

-psm requires --parallel 1 and is rejected otherwise. This is not a missing loop, it is
what a global mode switch can do at all. The trigger is a single linear walk over one prompt,
and with several slots there is no instant where none of them is generating; the split mode
belongs to the model, so a batch that mixes a prefill chunk with decode tokens has no single
right mode; and the state carry is one sequence, so a reload would silently destroy the cache
of every other one.

The gain also shrinks with the batch, which is the argument above running in reverse.
llama-batched-bench, -c 32768, q8_0 cache, -npp 2048 -ntg 128:

npl pp layer pp tensor tg layer tg tensor tensor tg edge
1 950.2 591.2 18.85 25.17 +33.5%
2 1011.0 590.9 33.09 41.78 +26.3%
4 1034.1 591.3 46.91 58.35 +24.4%
8 1041.5 591.3 57.60 71.10 +23.4%

Tensor's generation edge falls from +33.5% to +23.4% while layer's prefill edge grows from
+61% to +76%. On total throughput at this prompt to generation ratio layer is already ahead at
npl=2 and by +26% at npl=8. The batch of one that this PR optimises is the largest gap
there is.

Limitations

  • Only llama-completion, and only -np 1. A server would want this per request, which needs
    the library version below.
  • common_fit_params has no implementation for LLAMA_SPLIT_MODE_TENSOR and returns
    early (common/fit.cpp:183). The placement is therefore fitted only for a layer phase; a
    tensor phase loads with the default distribution, and what keeps that from failing hard is
    the fallback above, not a calculation.
  • An explicit -ts applies to both phases and means different things in each - a share of the
    layers under a layer split, a share of every tensor under a tensor split. There is no
    --prefill-tensor-split; a -ts that suits one phase can make the other one not fit, and
    the fallback is what catches that.
  • A layer split is unreliable on this model, independently of this PR: it intermittently
    produces all-! output, at 2-5 runs in 10. Reported in detail on ggml-meta : split a host-resident KV cache by head #57. Using a
    layer split for the prompt inherits that, so -psm layer is only as trustworthy as
    -sm layer is on a given model.
  • The reload discards anything derived from the old model that this tool does not rebuild.
    It rebuilds the sampler, the chat templates, the memory handle and the vocab. Combining
    -psm with --prompt-cache is untested.
  • The reload is a full one. A migration that moves only the weights that change device would
    be cheaper, and is the main thing the library version buys.

The library version, for later

The tool version answers whether the idea is worth anything. If it is, the same thing belongs
in the library, where it can be per request and cheaper. Sketch, in the order it would have to
be built:

1. Weight migration instead of a reload. Under LLAMA_SPLIT_MODE_LAYER a device holds
whole layers; under ..._TENSOR it holds a slice of every layer. Both placements are already
computed at load time - the first by llama_model::load_tensors through get_layer_buft_list,
the second by llama_meta_device_get_split_state. A migration walks the tensors, computes the
destination placement, allocates the destination buffers, copies through host memory where
there is no P2P, and frees the source. For a 2-GPU 50/50 layer-to-tensor move about half the
model changes device each way, so it moves less than a reload re-uploads, and it does not
touch the file at all. The cheap first cut is to keep the reload and only add the API, since
a reload from a warm page cache is already close to the transfer cost.

2. Keeping the cache in place. With --no-kv-offload there is nothing to do: the cache
sits in host memory in one global layout and the split happens at delivery, which is why the
state round-trip works today. A device-resident cache has to be re-split by head, which is the
operation ggml_backend_meta_buffer_set_tensor already performs for a host-resident cache
copy - the same segment and granularity rules apply, so this is a re-use, not a new mechanism.

3. Rebuilding the scheduler. llama_context owns sched, the compute buffers and the
graph-reuse state, all derived from the backend set. The sequence is: synchronize(), drop
the scheduler and its buffers, rebuild through the same path the constructor uses, invalidate
the reused graph. The reserve step has to run again because the workspace sizes differ between
the modes.

4. The API. llama_context_set_split_mode(ctx, mode) returning bool is enough for an
explicit caller. An automatic policy would be a cparam - switch when a batch crosses a size
threshold in either direction - but that should come after the explicit form has been used.
Whatever the API, it has to answer what the tool answers with a fallback: a mode that does not
fit must leave the caller with a usable context, not a freed one.

5. A policy that does not thrash. Per request the rule is simple, prompt then generation.
A server with several slots decoding while another prefills has no single answer, and would
need either a hysteresis or a rule that only switches when the whole batch agrees. The numbers
under -np 1 above say what it is worth: at npl=8 the generation edge is still +23%, but
the prefill loss is larger, so a policy that always picks tensor for decoding would be slower
overall than never switching.

Risks to be aware of before starting. The meta backend rotates a fixed set of per-buffer
compute containers and carries a FIXME saying so; that rotation assumes a stable buffer set
and is the first thing a mid-session rebuild would disturb. Anything holding a pointer into
the old model - samplers hold the vocab, adapters, mtmd contexts - has to be revalidated, and
the library cannot see those, so the API has to make the invalidation explicit.

A large batch pays for the collective a tensor split needs and a single token
does not, so the two phases of a request can prefer different split modes.
--prefill-split-mode processes the prompt under one mode, then carries the cache
through a reload into the other to generate.

Assisted-by: Claude Opus 5
@Piggidragon
Piggidragon marked this pull request as draft August 30, 2026 20:24
An out of memory in a simple buft aborted the process instead of returning
nullptr, so the caller could not fall back.

Assisted-by: Claude Opus 5
The fit writes the placement it found back into tensor_split and the buft
overrides, so the second load kept the placement of the first split mode and
could not fit the device memory. Keep the arguments and fit again, pin n_ctx
so both phases agree, and generate under the prefill split mode if the
generation one does not fit.

The switch reloads the model and carries one sequence, so require -np 1.

Assisted-by: Claude Opus 5
@Piggidragon

Copy link
Copy Markdown
Author

Closing this together with #59.

Measured on a 4070 + 3060 with Qwen3.8-27B-UD-IQ2_M, a 17055 token prompt and 128 generated tokens - the case this PR targets:

wall
-sm layer 25 s
-psm layer -sm tensor 27 s
-sm tensor 37 s

Plain -sm layer beats the switch. The decode saving over 128 tokens is 128/24.77 - 128/28.86 = 0.73 s, while the switch costs 3.1 s. It only turns positive above roughly 3800 prompt and 500 generated tokens at the same time.

The allocation failure fix from this branch (2c72f2d) is worth keeping on its own and moved to #57.

@Piggidragon
Piggidragon deleted the mgpu/split-mode-switch branch August 31, 2026 21:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant