kv_cache: reject non-row-addressable layouts and mixed-slice managers#363
Open
yinlin09 wants to merge 3 commits into
Open
kv_cache: reject non-row-addressable layouts and mixed-slice managers#363yinlin09 wants to merge 3 commits into
yinlin09 wants to merge 3 commits into
Conversation
UnpackTorchTensor resolves a tensor's base storage buffer with GetBaseBuffer and then blocks on AwaitBuffer to obtain the PjRtBuffer. AwaitBuffer waits on the buffer's materialization future but does not trigger materialization: for a tensor still carrying a deferred op -- e.g. a KV cache handed to the manager straight after allocation, before any forward pass or sync has run the graph -- nothing enqueues that graph, so the await never completes and construction hangs. Materialize the base node in place before awaiting. It is a no-op when the buffer is already materialized, and updates the node in place so raw transfer still DMAs the live storage buffer rather than a separate copy. This covers every unpack caller (KV cache manager and weight synchronizer), not just the offending vLLM registration path. Add a regression test that constructs a manager over deferred KV caches under DEFER_AND_FUSE and fails via a faulthandler watchdog if the constructor does not return (the constructor holds the GIL, so a hang freezes the process and no Python-thread watchdog could fire).
The nanobind std::string caster only accepts str, so the binding rejected the serialized StartTransferRequest that the KVCacheManager API wrapper itself passes as bytes (pybind11's caster accepted both, which is how earlier builds worked). Take nb::bytes and parse with ParseFromArray, matching the write_block_bytes binding. Also add a KV-cache-manager construction test over host-transferred (cpu_tensor.to(device)) caches, the allocation pattern vLLM uses.
GetMajorSliceByteSize assumes every buffer decomposes into contiguous
per-major-row slices. Two layout families break that silently:
1. Tiled buffers whose major dim is not the outermost physical dim
(e.g. bf16[512,3,2048] with layout {2,0,1}: XLA avoids padding the
3-tap dim by placing the block dim inside the tiled minor pair).
Blocks then interleave at tile granularity - with bf16's (2,1)
sub-tile, adjacent blocks share 32-bit words - and the formula
returns up to the whole buffer as one 'slice'. A KVCacheManager
registering such a layer allocates num_host_blocks x whole-buffer
host memory (observed: +221 GiB for one 20-tensor manager, host
OOM) and no per-block transfer is possible at all.
2. Tiled rank-2 buffers: the rank<3 fallback returns the dense row
stride without consulting tiles, so registration and transfers
proceed with plausible sizes while rows interleave inside tiles -
d2h/h2d then moves scrambled bytes (verified on tpu7x: bf16
[512,6144] round-trips corrupt while untiled shapes round-trip
byte-exact).
Add ValidateMajorSliceAddressable and reject both at manager
construction, and also reject layers whose slice size differs from
layer 0's: host allocation and host block offsets are derived from
bytes_per_block() == layer 0's slice for every layer, so a mixed-slice
manager under- or over-addresses host memory for all other layers.
Callers can register each slice class in its own manager.
A follow-up could support heterogeneous layers properly by sizing and
addressing host buffers per layer (LayerDeviceInfo already stores
physical_size per layer); until then a loud error at registration
replaces host OOM or silent data corruption at transfer time.
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
2 participants
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.
Title: kv_cache: reject non-row-addressable layouts and mixed-slice managers
Summary
GetMajorSliceByteSizeassumes every registered buffer decomposes intocontiguous per-major-row byte slices. Two layout families break that
silently today:
Tiled buffers whose major dim is not the outermost physical dim.
Example: XLA lays out bf16
[512,3,2048]as{2,0,1:T(8,128)(2,1)}(zero padding beats padding the 3-dim), which puts the block dim
inside the tiled minor pair. Blocks interleave at tile granularity —
with the bf16
(2,1)sub-tile, adjacent blocks share 32-bit words —and the formula returns up to the whole buffer as one "slice". A
KVCacheManager registering such a layer allocates
num_host_blocks x whole_bufferhost memory (observed: +221 GiB forone 20-tensor manager, host OOM), and no per-block transfer is
possible for the layout at all.
Tiled rank-2 buffers: the rank<3 fallback returns the dense row
stride without consulting tiles, so registration and transfers
proceed with plausible sizes while rows interleave inside tiles.
Verified on tpu7x: bf16
[512,6144](layout{1,0:T(8,128)(2,1)})and s8
[512,12288]({1,0:T(32,128)(4,1)}) round-trip corruptedthrough d2h/h2d, while untiled shapes round-trip byte-exact.
This adds
ValidateMajorSliceAddressableand rejects both at managerconstruction, and also rejects layers whose slice size differs from
layer 0's: host allocation and host block offsets are all derived from
bytes_per_block()(layer 0's slice), so a mixed-slice manager under-or over-addresses host memory for every other layer. Callers can
register each slice class in its own manager.
Follow-up proposal
Heterogeneous layers could be supported properly by sizing and
addressing host buffers per layer (
LayerDeviceInfoalready storesper-layer
physical_size); until then a loud error at registrationreplaces a host OOM or silent data corruption at transfer time.
Verified on tpu7x with a registration + permuted per-block d2h/h2d
round-trip probe over the shapes above: the two bad families are
rejected with layout-naming errors;
[512,{4,8,16},2048]register withexact slices and round-trip byte-exact.