completion : allow the prompt to be processed under its own split mode - #58
Closed
Piggidragon wants to merge 3 commits into
Closed
completion : allow the prompt to be processed under its own split mode#58Piggidragon wants to merge 3 commits into
Piggidragon wants to merge 3 commits into
Conversation
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
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
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:
Plain The allocation failure fix from this branch (2c72f2d) is worth keeping on its own and moved to #57. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 65536with a q8_0 device-resident cache:-smtensor -ts 50,50layerLayer 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 isprocessed under it, then the model is reloaded under
--split-modeand the cache is carriedacross, 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_datawrites a layout-independent form: this was verified separately bywriting a
--prompt-cacheunder-sm layerand restoring it under-sm tensor, whichproduced byte-identical output.
The switch:
llama_state_seq_get_data,llama_state_seq_set_data,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 are12282 MiB and 12288 MiB:
-sm layer-sm tensor -ts 50,50-psm layer -sm tensor -ts 50,50The 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
-psmis therefore smaller than under either pure mode. At-c 98304the same pattern holds (layer 10815/11833, tensor 11537/11467,
-psm11497/11825), andneither mode dominates - at
-c 114688 -ts 50,50layer runs out on device 1 while tensorstill fits.
Three things follow, and this PR handles them:
The fit has to run again for the second mode.
common_fit_paramswrites the placement itfinds back into
params.tensor_splitand the buft overrides. On the second load it then seesthem set and aborts with
tensor_split already set by user, so the generation mode inheritedthe 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_ctxto the first phase so thestate 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
-tstuned for theprefill mode killed the request after the prompt was already processed:
-c 65536 -ts 57,43fits 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 tensoraborted the process and no fallback wasreachable. 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:-ngl/-ts,-psm layer -sm tensor-sm tensor-ngl 99 -ts 50,50 -psm layer -sm tensor-ngl 99 -ts 57,43 -psm layer -sm tensor-ngl 99 -sm layerResult
Same setup, wall clock for the whole request - the 58.3k-token prompt plus 128 generated
tokens, model load included:
-sm tensor-sm layer-psm layer -sm tensor-psm tensor -sm layer(control, wrong way round)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 tensorthat is 40% off the wall clock for thesame generation rate.
Against
-sm layerit costs the 5 s reload and buys +33% generation, so it pays foritself after about 300 generated tokens - below that, plain
-sm layeris better. Thereversed 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 bufferwhile 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-psmrequires--parallel 1and is rejected otherwise. This is not a missing loop, it iswhat 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: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=2and by +26% atnpl=8. The batch of one that this PR optimises is the largest gapthere is.
Limitations
llama-completion, and only-np 1. A server would want this per request, which needsthe library version below.
common_fit_paramshas no implementation forLLAMA_SPLIT_MODE_TENSORand returnsearly (
common/fit.cpp:183). The placement is therefore fitted only for a layer phase; atensor phase loads with the default distribution, and what keeps that from failing hard is
the fallback above, not a calculation.
-tsapplies to both phases and means different things in each - a share of thelayers under a layer split, a share of every tensor under a tensor split. There is no
--prefill-tensor-split; a-tsthat suits one phase can make the other one not fit, andthe fallback is what catches that.
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 alayer split for the prompt inherits that, so
-psm layeris only as trustworthy as-sm layeris on a given model.It rebuilds the sampler, the chat templates, the memory handle and the vocab. Combining
-psmwith--prompt-cacheis untested.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_LAYERa device holdswhole layers; under
..._TENSORit holds a slice of every layer. Both placements are alreadycomputed at load time - the first by
llama_model::load_tensorsthroughget_layer_buft_list,the second by
llama_meta_device_get_split_state. A migration walks the tensors, computes thedestination 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-offloadthere is nothing to do: the cachesits 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_tensoralready performs for a host-resident cachecopy - the same segment and granularity rules apply, so this is a re-use, not a new mechanism.
3. Rebuilding the scheduler.
llama_contextownssched, the compute buffers and thegraph-reuse state, all derived from the backend set. The sequence is:
synchronize(), dropthe 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 anexplicit 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 1above say what it is worth: atnpl=8the generation edge is still +23%, butthe 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.