feat(mp): layer-major KV staging and first-layer request release - #23
Draft
youngrok-XCENA wants to merge 10 commits into
Draft
feat(mp): layer-major KV staging and first-layer request release#23youngrok-XCENA wants to merge 10 commits into
youngrok-XCENA wants to merge 10 commits into
Conversation
The lmcache-driven retrieve path stages whole chunks and scatters every layer in one burst, so the earliest layer is only usable once the last one has landed. Nothing downstream can overlap a request's prefill with its own transfer. Add a layer-major path beside it. The host object keeps its [kv, layer, token, hidden] layout: a layer slice is kv_size contiguous byte ranges inside the chunk, so slicing needs no change to the store path, no new key, and no per-deployment layout choice. What it does need is a layer selector on the scatter launch, because LaunchVar only carried a block range and the kernel derived the layer axis from the shape descriptor -- one launch always covered all layers. The kernel now takes layer_offset and staged_layers: blockIdx.z walks the staged slice, the engine side is addressed at layer_offset + slot, and the staged side uses the slot with staged_layers as its stride divisor. Both default to "all layers", which reproduces the current addressing exactly, so the chunk-major path is bit-identical and stays the default. _run_object_group_transfer_plan_layer_major issues one native plan per slice and calls back after each, which is where a caller records the per-slice completion event. It is deliberately narrower than the chunk-major planner: one kernel group per object group, no GDS objects. Hybrid models and GDS deployments keep the chunk-major path. The temp staging buffer is still sized for every layer and a slice uses its prefix; sizing it down to the slice (which is what removes the chunk-size-proportional GPU buffer at boot) is left to a follow-up.
Layer-major staging only pays off if the request goes back to the engine early. Today an async-load request leaves WAITING_FOR_REMOTE_KVS when get_finished reports it, and that report waits for the whole retrieve; by the time the forward pass runs, every layer is already resident, wait_for_layer_load has nothing to block on, and prefill runs strictly after the transfer. The per-layer machinery is then inert -- layer-major transfer without early release measures the same as chunk-major. LayerArrivalGate owns that decision, separately from any transport so it can be tested without a GPU or a server: the transport records each slice as it lands, the engine thread asks whether enough layers have arrived to release (one, by default) and which event gates the layer it is about to read. A slice covering N layers costs one wait rather than N, and both threads share one lock. The one-layer default is measured, not guessed: on an in-process connector with this structure the first layer arrives in 2.7-3.1 ms against 89 ms for a full transfer, at every concurrency from 1 to 8, so the release condition is met long before the transfer ends and waiting for a second slice only delays the start. Also declare requires_piecewise_for_cudagraph. Per-layer synchronization cannot be captured in a CUDA graph, so a full-graph replay drops the wait and the model can read KV still in flight. vLLM's default FULL_AND_PIECEWISE runs prefill batches piecewise so the wait does fire, but --cudagraph-mode FULL would lose it silently. The declaration is a classmethod called before any server handshake, so it reads lmcache.mp.layerwise_overlap from extra_config; that flag has to be set on the vLLM side as well as the server side. The ZMQ per-slice completion frames that feed record_slice are not here yet, so nothing enables this path end to end.
The previous two commits added layer-major staging and the release policy but never wired them, so the new planner was unreachable and no deployment could reach it. Add the switch. --retrieve-layers-per-stage N (server config retrieve_layers_per_stage, default 0) selects the path per transfer. 0 keeps chunk-major, so the default is byte-for-byte the behaviour before this series. N > 0 stages the layer axis in slices of N layers. Off by default on purpose: layer-major staging is not a free win. When the per-layer transfer runs well ahead of the per-layer compute, the compute stream stalls at every layer and that wait is shared with the co-scheduled requests in the batch. On an in-process connector with the same structure the sign flipped with prompt length and chunk size -- 50 ms saved at prompt 16k / chunk 256, 38 ms lost at prompt 64k / chunk 256. A deployment should measure before enabling it. Retrieve only. Stores keep chunk-major, where no per-layer consumer exists to overlap with. Retrieves fall back to chunk-major, logging the reason once, when the object group holds several kernel groups (hybrid models), when the batch has GDS-backed objects, or when the native object-group transfer extension is missing -- so enabling the knob on an unsupported deployment degrades rather than fails. Note that turning the knob on does not yet make first-token time drop: the per-slice completion frames that would let the engine start on layer 0 are still missing, so a request is still released only once the whole retrieve is done. What the knob buys today is the A/B control arm for that change -- layer-major transfer without early release.
Layer-major staging moved KV a slice at a time, but a request still went back to the engine only once its whole retrieve had finished, so by the time the forward pass ran every layer was already resident and nothing overlapped. Turning the knob on bought reordered copies and no gain. This closes that. Two facts have to cross the process boundary, and they cannot travel together. The per-layer event says "these bytes are on the GPU". It cannot also say "the server got this far", because an event that has not been recorded reports *complete* -- a worker polling pre-made events would release a request whose KV was never copied. That is why the existing single-event path only touches its event after the server's response arrives, which is exactly the wait being removed here. So a small shared-memory board carries the second fact: one monotonic count per in-flight retrieve, written by the server after it records a slice's events and read by the worker. Each slot sits on its own cache line and has one writer and one reader, so neither side takes a lock. - Server: after enqueuing a slice, record its layers' events, then raise the board count. Order matters -- the count is what makes those events safe to wait on, so it is published last. Failing to set this up costs the overlap, not correctness: the retrieve proceeds and the worker waits for all of it. - Worker: a pool of slots, each owning one board entry and num_layers events built once and reused, so a retrieve costs a slot rather than 32 driver round trips. When the pool is empty a retrieve runs the old way. - get_finished reports a request once its first layer lands, keeping the future in a draining set so the transfer's real completion and any failure are still collected and the slot is freed then. - wait_for_layer_load blocks on each in-flight retrieve's event for the layer about to be read, one wait per slice rather than per layer, with a timeout that logs rather than hanging the model-runner thread. Enabled by lmcache.mp.layerwise_overlap on the vLLM side together with the server's --retrieve-layers-per-stage. The two are still separate settings; a deployment has to set both, and whether the worker should instead learn it from the server at registration is worth deciding in review. Not yet exercised on a GPU: the board is covered by a real cross-process test, and the pool, dispatch and staging arithmetic by unit tests, but no run has moved a byte of KV through this path.
…LM side Turning the feature on took two settings that had to agree: the server's --retrieve-layers-per-stage and the connector's lmcache.mp.layerwise_overlap. Setting one without the other did something silently useless -- the server alone reordered copies with nothing waiting on them, the connector alone had no progress to track and only constrained the CUDA graph mode. Keep one setting, on the vLLM side, and let it carry the width: lmcache.mp.layerwise_overlap is now a layer count (0 = off, true = 1). The worker sends the width with every retrieve, so the server has no setting of its own; drop the flag, the config field, and the context property. The vLLM side has to be the one that survives, not the server: the CUDA graph declaration is a classmethod called before any server handshake, so it can only read deployer-supplied config. Putting the switch there means the graph mode and the transfer width come from the same value and cannot disagree. resolve_layers_per_stage reads what a deployer would plausibly write -- an int, a bool, or a decimal string -- and treats anything else as off with a warning rather than failing startup.
The retrieve handler grew three parameters for layer-major staging -- the
arrival board, the per-layer event handles and the slice width -- but the
protocol still declared five payload classes. The MP server compares handler
signatures against the protocol when it registers handlers, so it raised at
startup and never came up, with or without the feature enabled:
Handler for RequestType.RETRIEVE expects 5 arguments, but got 8
ValueError: Handler signature does not match for request type: RETRIEVE
Widening the declaration alone is not enough. Both ends reject a payload whose
length differs from the declaration, so appending the three fields only when
layer-major is on would have made every chunk-major retrieve fail instead. The
payload is now always eight items long and a chunk-major retrieve fills the
last three with (None, None, 0).
The other callers of RETRIEVE -- the SGLang and TensorRT-LLM adapters and the
server bench helper -- send the same three placeholders.
The worker builds its per-layer arrival slots from the transport's event
backend, reading it as a public attribute:
event_backend = getattr(transfer_ctx, "event_backend", None)
device = getattr(transfer_ctx, "device", None)
LMCacheDrivenTransferContext keeps both under private names and exposes
neither, so getattr always returned None. Layer-major retrieval therefore never
engaged in any deployment: the worker logged that it was enabled, then took the
fallback and ran an ordinary chunk-major retrieve, leaving only
Layer-major overlap is enabled but the transfer context exposes no event
backend or no layers were registered; retrieves run without per-layer
progress
Adding the two properties is the fix that keeps the adapter off another class's
private members, which SLF enforcement in lmcache/v1/multiprocess/ requires.
The worker pools its per-layer events and exports each handle once, because
exporting is a driver round trip that should not be repeated per retrieve. It
then hands those handles to the transport, which treated them as event objects
and exported them again:
AttributeError: 'bytes' object has no attribute 'ipc_handle'
The transport now takes layer_event_handles -- already-exported bytes -- and
puts them on the wire unchanged, which is what the pool's contract was.
run_object_group_transfer_plan gained layers_per_stage and on_layer_batch, and
the stub this test substitutes for it kept the old signature, so both of its
retrieves failed:
TypeError: fake_transfer() got an unexpected keyword argument
'layers_per_stage'
The file is not one this branch otherwise touches, so the failure only shows up
when the whole suite runs.
Two modules, covering the two halves the unit tests could not reach because they stub the device and the server. test_layer_major_retrieve_gpu speaks the message protocol to a real server and compares GPU bytes. Every slice width retrieves into its own destination range, so a width that dropped a layer, wrote it twice or placed it at the wrong offset differs from the stored blocks there. Widths 1, 2, 3, 8 and 32 all match byte for byte across 32 layers, and width 1 matches the chunk-major path layer for layer: the knob changes when bytes arrive, not which. test_layer_major_worker_release_gpu drives the real worker adapter against a real server, which is what exercises the arrival pool, the release decision and the per-layer wait. There is no vLLM; the engine is replaced by the loop a worker runs -- poll get_finished, then walk the layers calling wait_for_layer_load before reading each one. On a 1 GiB retrieve the request comes back after 1.7 ms with one layer per slice against 20.8 ms chunk-major, and all 32 layers still read correct KV afterwards. A second case asserts the release really preceded the copy finishing, by checking that the released request sits in the draining set with its transfer future unresolved -- a state the chunk-major path cannot reach. Each test gets its own server. A retrieve reads what its own lookup unlocked and these harnesses have no session teardown to release that lock, so a shared server carries one test's locks into the next.
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
1 participant
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.
What this PR does / why we need it:
In MP mode a cache-hit request has its KV moved a chunk at a time, and it is
released from the waiting queue only after the whole transfer has landed. Its
prefill therefore starts after the transfer ends, and the first token costs the
whole transfer plus the whole compute.
This PR moves KV a layer slice at a time and releases the request on its
first slice, so the compute of the earlier layers runs while the later ones
are still copying. One setting turns it on:
The value is layers per slice:
0(default) keeps the chunk-major pathuntouched,
truemeans one layer at a time,8means eight. Ints, bools anddecimal strings are accepted; anything else logs a warning and reads as off
rather than failing startup. The LMCache server has no setting of its own --
the worker sends the width with every retrieve, so the two cannot disagree.
Measurements
Llama-3.1-8B-Instruct on one H200, L1 in host DRAM, lveval
factrecall_en16k(15,154 tokens per request), 20 documents, cache-hit rounds. These are
measurements of this branch, not of a sibling implementation.
Enabling the knob also changes vLLM's cuda-graph mode, because the connector
declares
requires_piecewise_for_cudagraph. A two-arm comparison cannotseparate that from the overlap itself, so the baseline below is the knob off
with graphs forced to PIECEWISE. Against vLLM's default
FULL_AND_PIECEWISE, that control arm differs by −2.4% / +6.2% / +3.5% infirst-token time at concurrency 1 / 4 / 8, so the graph mode is not what the
numbers below are showing.
The gain is bounded by whichever of the two is shorter, since the point is to
hide one inside the other. The LMCache chunk size moves that ratio: a larger
chunk leaves a larger unmatched tail on a hit, which is the prefill the
arriving layers have to overlap with. Sweeping it, at concurrency 1:
Last two columns are the change in mean first-token time against the baseline;
negative is faster.
ceilingismin(transfer, compute) / (transfer + compute),the most the overlap could remove. Hit fraction comes from each run's own
cumulative external hit rate (the workload runs one populate pass and two hit
passes, so per-request hit is 1.5x the cumulative figure); compute is the
populate pass's first-token time scaled by the missed fraction; transfer is the
hit tokens' KV size at the retrieve rate measured on this machine, 20.6 ms per
GiB.
The gain peaks where the two are the same length, which is what the
mechanism predicts. All four metrics at that point (chunk 2048, 8 layers per
slice):
End-to-end and throughput move less than first-token time because 50 output
tokens of decode dominate the request. Time per output token does not move at
all, which is what should happen: retrieval is a prefill-time concern.
Where it loses. Two regimes, and both argue for keeping the default off:
When the transfer dwarfs the compute there is nothing to hide and the slicing
only adds cost. And above concurrency 1 we could not find a configuration
that wins: chunk 4096 goes from −14.9% at concurrency 1 to +18.5% at 4.
Splitting one request's transfer into pieces lets concurrent requests' pieces
interleave on the transfer path, delaying each request's completion. Wider
slices reduce the damage but did not remove it. We have not instrumented the
interleaving directly, so that explanation is inference from the shape.
Caveats worth stating: one pass per cell, no repeats, so small differences such
as the −1.7% at chunk 256 are not evidence of a sign. Only one point (1.32) sits
near the balance, so the peak's location is bracketed rather than pinned. And
the ratio was moved with the chunk size, which also changes the storage layout,
the staging buffer size and the prefill batch budget.
Design changes
Two things change: the unit KV is moved in, and when the request is
released.
The unit -- a chunk used to be the smallest thing transferred and scattered, so
there was no way to pick layers. A layer slice can now be that unit, and the
storage layout is untouched: a chunk's memory already has a layer axis, so a
slice is a few contiguous ranges inside it. No new keys, no server-wide layout
choice, no discriminator for mixed deployments.
The release point -- it used to be when the transfer finished. It is now when
the first slice lands. This is where the gain comes from. Slicing the transfer
without moving the release leaves every layer already present when compute
starts, so there is nothing left to overlap; that is the shape upstream LMCache#4460
has, and its author's measurements are a wash for the same reason.
flowchart TB subgraph a["Before — a chunk at a time"] direction LR A1["chunks staged<br/>whole"] --> A2["every layer<br/>scattered at once"] --> A3["request released<br/>when transfer ends"] --> A4["prefill"] end subgraph b["After — a layer slice at a time"] direction LR B1["one slice<br/>staged"] --> B2["that slice's layers<br/>scattered"] --> B3["slice arrival<br/>published"] B3 --> B4["request released<br/>on its first slice"] --> B5["prefill"] B3 -.->|"remaining slices"| B1 B5 -.->|"waits per layer"| B3 end a ~~~ bThe scope is deliberately narrow. Only object groups with a single kernel group
(full attention) and non-GDS deployments take the new path. Hybrid models put
layer groups of different shapes in one object and would each need slicing
separately; GDS deployments are a different path already. Both keep the old
behaviour and log the reason once.
Implementation Details
Kernel (
csrc/cuda/) --LaunchVargainslayer_offsetandstaged_layers. The kernel already walked layers asblockIdx.zand thestaging offset was
[kv, nl, T, D], so only the axes were split:blockIdx.zis the position inside the staged slice, the engine-side address is
layer_offset + slot, and the staging-side address uses the slot withstaged_layersin place ofshape_desc.nl. The grid's z extent becomesstaged_layers. Both fields default to0(= all layers), and in that case theaddress arithmetic is identical to before; pybind exposes them with the same
defaults, so existing call sites are unchanged.
Staging (
lmcache/v1/gpu_connector/gpu_ops.py) --build_layer_staging_copiescopies a layer range rather than a whole object.An object is
[kv, nl, T, D], so a slice iskv_sizecontiguous ranges (onefor MLA) and the destination buffer is filled as
[kv, layer_count, T, D].host_offsetadvances with the range actually read, because the native memcpyuses it to split on host-pinned boundaries.
Planning (
.../modules/lmcache_driven_transfer.py) --_run_object_group_transfer_plan_layer_majorsits beside the existingfunction rather than inside it, issuing one native plan per slice and calling
on_layer_batch(layer_start, layer_count)after each. That callback is wherethe completion event is recorded. About 40 lines of planning logic are
duplicated; that is deliberate, to keep the paths from interleaving. Happy to
merge them behind a parameter if you prefer sharing.
The staging buffer is still sized for all layers and a slice uses its prefix.
Sizing it down to the slice -- which is what would remove the startup GPU
reservation proportional to chunk size -- is left for a follow-up.
Release policy (
.../layer_arrival.py) --LayerArrivalGateholds "howmany layers before releasing" and "which event layer L waits on". The transport
is injected, so the policy is testable without a GPU or a server. The transfer
thread calls
record_slice; the engine thread callsrelease_readyandevent_for_layer; both touch the same state, so it is lock-protected. A slicecovering N layers costs one wait, not N.
Cross-process arrival (
layer_arrival_board.py,layer_arrival_pool.py) --two facts have to cross the process boundary and they cannot share a carrier. A
per-layer event says "these bytes are on the GPU". It cannot also say "the
server got this far", because an event that has not been recorded reports
complete when queried -- so a worker polling pre-made events would release a
request whose KV was never copied. The existing single-event path avoids this by
touching the event only after the server's response arrives, which is exactly
the wait being removed here.
So the second fact rides a small shared-memory board: one monotonic count per
in-flight retrieve, raised by the server after it records the slice's event
and read by the worker. Each slot sits on its own cache line with a single
writer and a single reader, so neither side takes a lock. On the worker side the
slots are pooled -- a slot owns one board entry plus one event per layer, built
once and reused -- so a retrieve costs one slot rather than 32 driver round
trips. When the pool is empty the retrieve runs the old way.
Release and per-layer wait (
vllm_multi_process_adapter.py) --get_finishedreports a request whose first slice has landed and keeps itsfuture in a draining set, so the transfer's real completion and any failure are
still collected, and the slot is returned then.
wait_for_layer_loadwaits onthe event for the layer about to be read, once per slice rather than per
layer, and logs after a timeout instead of hanging, since the model-runner
thread must not hold the whole batch.
CUDA graphs (
lmcache_mp_connector.py) -- per-layer synchronization is notcaptured in a CUDA graph, so a full-graph replay skips the wait entirely and the
model can read KV that has not arrived. That is a correctness issue, not a
performance one, hence the
requires_piecewise_for_cudagraphdeclaration.vLLM's default
FULL_AND_PIECEWISEroutes prefill-bearing batches through thepiecewise graph and the wait fires, but
--cudagraph-mode FULLwould silentlylose the barrier.
This method is a classmethod called before any handshake with the server,
which is why the setting lives on the vLLM side. It cannot read server config,
so a server-side setting would give the graph mode and the transfer width
different sources that can disagree. One value feeds both.
Special notes for your reviewers:
Four seams have to line up, and three of them fail quietly. Worth
knowing if you carry any of this into a different shape:
declare eight. The server compares handler signatures against the protocol
at registration and will not start otherwise, with or without the feature
enabled.
only when the feature is on -- a chunk-major retrieve sends
(None, None, 0).If the transport keeps that private,
getattrreturns None, the adapterlogs that the feature is enabled and then silently falls back to
chunk-major. This one is invisible except in one warning line.
must not export them again.
What is verified. Two GPU test modules, plus the existing suite:
tests/v1/multiprocess/+tests/v1/test_vllm_mp_adapter.pySlice width costs scale with slice count, not layer count: on a 1 GiB
retrieve the whole transfer grows by 0.3 ms at 32 layers per slice, 1.0 ms at
8, and 7.6 ms at 1.
Base is not
dev. Our internaldev(f82f6fd) is 34 commits behindupstream
dev(09bc14c), and four of those touch files this PR changes.Rebasing onto the older base would put this on a different shape of the
code, so the base branch here is a snapshot of upstream
dev09bc14c.Relationship to #4460.
That PR changes the storage layout to layer-interleaved and has cross-process
completion notification, but no early release -- its completion check
looks at the last frame and the last layer's event, so
get_finishedonlyfrees a request after the whole transfer, and
wait_for_layer_loadpassesstraight through. This PR leaves the storage layout alone, carries the
notification on a shared-memory count instead of ZMQ partial frames, and
releases on the first slice. If you would rather have one path, the early
release and the arrival board here are the parts worth taking; the layer
range fields on the launch descriptor are independent of the layout question.
Follow-ups not in this PR -- size the staging buffer to the slice (drops
the startup GPU reservation proportional to chunk size); find whether
serializing transfers on one stream removes the concurrency loss; pick a
default slice width; repeat the sweep on a slower storage tier, where the
transfer is longer and the balance sits at a more common chunk size.
If applicable: