[KV Offload] Reshape the transfer data model: per group specs and offloaded side alignment offset

[hickeyma]

Purpose

This PR reshapes how a single KV cache transfer is described between the scheduler (which decides what to move) and the worker (which performs the copy). It does NOT change how data is copied, stored, or sent. It instead changes the shape of the descriptor that flows between them.

What is changed:

1. Removes packing: A transfer used to pack all KV cache groups into one flat block_ids array which was confusing with a parallel group_sizes array. We now emit one spec per KV cache group.
2. Relocates the alignment offset: "where does this group's first GPU block sit" offset (block_indices) moves off the GPU spec and onto the offloaded (CPU) spec via set_gpu_block_offset(...). Offloaded is the side that the offset actually describes.
3. Introduces an explicit, direction tagged transfer unit: (GroupTransfer +TransferSpec) replacing the anonymous (src, dst) tuple.

The result is a descriptor whose structure matches the physical reality of the transfer, removing two kinds of conditions that must always hold and were previously enforced by runtime asserts and recomputed in frequently executed code.

Partial #33689

Tasks from #33689:

• Remove GPULoadStoreSpec.group_sizes, and instead have a LoadStoreSpec per group.
• Replace GPULoadStore.block_indices with an offloaded-side LoadStoreSpec.set_gpu_block_offset(...)
• Introduce submit_store() / submit_load() methods on OffloadingWorker, replacing the need for is_store flag ([KV Offload] Replace OffloadingHandler with OffloadingWorker #45053)

How KV offloading transfers work

vLLM v1 can offload KV cache blocks from GPU to CPU RAM & storage. The control flow:

┌────────────────────────────┐         ┌──────────────────────────────┐
│  Scheduler process          │         │  Worker process               │
│  (OffloadingConnector-       │         │  (OffloadingConnector-        │
│   Scheduler)                 │         │   Worker)                     │
│                              │         │                               │
│  OffloadingManager           │         │  OffloadingWorker             │
│   .prepare_store / load  ────┼─ spec ─▶│   .transfer_async             │
│   decides WHICH offloaded    │         │    → SingleDirectionOffloading│
│   blocks + allocates them    │         │      Handler does the copy    │
└────────────────────────────┘         │      (swap_blocks kernel)      │
                                         └──────────────────────────────┘

The transfer spec is the contract crossing that boundary. The scheduler builds it and the worker's handler consumes it to issue the actual block copies.

The issue that makes the spec non-trivial "block size mismatch"

Offloaded blocks can be larger than GPU blocks. Defined by block_size_factor = offloaded_block_size / gpu_block_size (an integer ≥ 1). One offloaded CPU block therefore holds block_size_factor GPU sized sub-blocks.

When a request's first offloadable GPU block does not land on an offloaded block boundary, the worker must skip the leading sub-blocks of the first offloaded block:

gpu_block_offset = 5, block_size_factor = 4
blocks_to_skip   = 5 % 4 = 1

offloaded block (holds 4 GPU sub-blocks):
   ┌────┬────┬────┬────┐
   │ s0 │ s1 │ s2 │ s3 │
   └────┴────┴────┴────┘
     ^skip   ^─── copy these GPU sub-blocks ───^

With hybrid / HMA models there are multiple KV cache groups (e.g. a full-attention group and a sliding-window group), and each group has its own first block offset. So the skip is inherently per group. It is because of this that this redesign is crucial.

Problem: the old data model packed unrelated things together

Before

# vllm/v1/kv_offload/base.py  (before)
TransferSpec = tuple[LoadStoreSpec, LoadStoreSpec]   # (src, dst) — direction implicit

class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
    def __init__(self, block_ids, group_sizes, block_indices):
        super().__init__(block_ids)
        assert sum(group_sizes) == len(block_ids)      # packing invariant
        assert len(block_indices) == len(group_sizes)  # parallel-array invariant
        self.group_sizes   = group_sizes   # how many GPU blocks per group
        self.block_indices = block_indices # first-block GPU offset per group

A single GPULoadStoreSpec flattened every group into one block_ids array, then used two parallel arrays to recover the structure:

• group_sizes[i] — how many blocks belong to group i (the packing).
• block_indices[i] — the logical GPU-block offset of group i's first block (the alignment skip source).

        ┌──────────────── one GPULoadStoreSpec ────────────────┐
block_ids:   [ b0  b1  b2 | b3  b4 | b5  b6  b7  b8 ]
group_sizes: [     3      |   2    |       4        ]
block_indx:  [     0      |   5    |       0        ]
              └ group 0 ──┘└ grp1 ─┘└─── group 2 ───┘
              (must be reconstructed by the consumer)

Why this is a problem

Issue	Detail
Recomputed in frequently executed code	The worker has to walk group_sizes to re-slice the flat array back into per group views on every transfer (src_offset/dst_offset). It then needs to assert src_offset == num_src_blocks to prove it walked it correctly. The structure is discarded at construction and rebuilt at consumption.
Constraints enforced at runtime	sum(group_sizes) == len(block_ids) and len(block_indices) == len(group_sizes) are structural truths that an assert can only catch. They are not preventable.
block_indices live on the wrong object	The skip only ever applies to the offloaded side (the GPU side always has block_size_factor == 1, so its skip is always 0). Yet the offset was stored on the GPU spec, forcing the worker to reach into gpu_spec.block_indices and compute the offloaded side skip from it. The data described one side but lived on the other.
Direction is implicit	(src, dst) carried no label. The worker inferred direction from a separate gpu_to_cpu flag on itself and from which side happened to be a GPULoadStoreSpec. Reading the spec alone could not tell you store vs load.

---

Design: a per group, direction tagged descriptor

After

 # vllm/v1/kv_offload/base.py  (after)

class LoadStoreSpec(ABC):
    def __init__(self) -> None:
        self._gpu_block_offset: int = 0   # 0 == block-aligned (common case)

    def set_gpu_block_offset(self, offset: int) -> "LoadStoreSpec":
        """Record this group's first-block GPU offset. Returns self for chaining."""
        self._gpu_block_offset = offset
        return self

    @property
    def gpu_block_offset(self) -> int:
        return self._gpu_block_offset


class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
    def __init__(self, block_ids: list[int]):
        super().__init__()
        self.block_ids = np.array(block_ids, dtype=np.int64)


class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
    """One KV cache group's GPU blocks. No group_sizes, no block_indices."""
    def __init__(self, block_ids: list[int]):
        super().__init__(block_ids)


@dataclass
class GroupTransfer:
    """One KV cache group's transfer descriptor. Direction-neutral."""
    gpu_spec: GPULoadStoreSpec  # always block_size_factor == 1
    offload_spec: LoadStoreSpec # carries gpu_block_offset (the alignment)

Direction is now expressed by which method is called on OffloadingWorker introduced in #45053:

# OffloadingWorker class (vllm/v1/kv_offload/base.py)
def submit_store(self, job_id: int, groups: Sequence[GroupTransfer]) -> bool: ...
def submit_load(self, job_id: int, groups: Sequence[GroupTransfer]) -> bool: ...

The same example from above is restructured (the parallel arrays are gone) and each group is a self contained pair. The offset rides on the offloaded spec it describes:

submit_store(job_id, groups=[
   GroupTransfer(gpu_spec=GPU[b0,b1,b2],    offload_spec=CPU[c0]@offset=0),
   GroupTransfer(gpu_spec=GPU[b3,b4],       offload_spec=CPU[c1]@offset=5),
   GroupTransfer(gpu_spec=GPU[b5,b6,b7,b8], offload_spec=CPU[c2,c3]@offset=0),
])
                                                           └ gpu_block_offset

Mapping old to new

Old	New
group_sizes[i]	len(groups[i].gpu_spec.block_ids)
block_indices[i]	groups[i].offload_spec.gpu_block_offset
TransferSpec = (src, dst) tuple	groups: Sequence[GroupTransfer] passed to submit_store / submit_load
gpu_to_cpu inferred from spec types + flag	direction encoded in the method called (submit_store vs submit_load)
packed flat block_ids	per-group block_ids on each spec
gpu_block_offset on BlockIDsLoadStoreSpec	gpu_block_offset on base LoadStoreSpec. Any medium can carry it

What is the value?

1. Structure matches reality: A transfer is a set of per group block movements, each with its own alignment. The new types say exactly that. The old types said "one big array plus two additional arrays you must zip back together correctly."
2. Removes potential bugs: The packing condition (sum(group_sizes) == len(block_ids)) and the parallel array condition (len(block_indices) == len(group_sizes)) no longer exist and hence remove potential issues. The runtime asserts that policed them are deleted.
3. The frequently executed code path gets shorter: The handler no longer re-slices a flat array or maintains offset cursors per transfer. Fewer operations, less state, on every store and load.
4. Correct ownership: gpu_block_offset now lives on the spec it describes (offloaded side). A reader no longer has to know "the skip is on the GPU spec but only applies to the CPU side."
5. Explicit direction: removes a layer of inference (spec type sniffing + a separate gpu_to_cpu flag) and makes logs/repr/debugging legible.
6. Medium agnostic alignment offset: gpu_block_offset lives on the base LoadStoreSpec so that file or object based backends can carry it without inheriting from BlockIDsLoadStoreSpec.

Test Plan

Test Result

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[orozery]

Thanks @hickeyma !

Why are we coupling the offloaded medium to use BlockIDsLoadStoreSpec?
We want to support arbitrary LoadStoreSpec, not necessarily blockIDs based ones.
One example is a GDS backend (instead of CPU backend) which will map blocks to files instead of block IDs.
So we want to keep it a general LoadStoreSpec.

BTW I think there's a chance that it this PR may be simplified if we first tackle another road map item:

Remove kv_offload/worker/worker.py. Replace OffloadingHandler with OffloadingWorker in abstract.py with abstract functions: submit_store, submit_load (replacing transfer_async).

[hickeyma]

Thanks @orozery for the feedback.

Why are we coupling the offloaded medium to use BlockIDsLoadStoreSpec?
We want to support arbitrary LoadStoreSpec, not necessarily blockIDs based ones.
One example is a GDS backend (instead of CPU backend) which will map blocks to files instead of block IDs.
So we want to keep it a general LoadStoreSpec.

You are right. gpu_block_offset is alignment metadata and not a block ID concept. I'll push a fix for that.

BTW I think there's a chance that it this PR may be simplified if we first tackle another road map item

Ok, I'll push a PR for Task 7 and we can iterate on that.

[hickeyma]

@orozery Pused PR #45053 for Task 7.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hickeyma.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[hickeyma]

Why are we coupling the offloaded medium to use BlockIDsLoadStoreSpec?
We want to support arbitrary LoadStoreSpec, not necessarily blockIDs based ones.
One example is a GDS backend (instead of CPU backend) which will map blocks to files instead of block IDs.
So we want to keep it a general LoadStoreSpec.

@orozery Fixed

[orozery]

This forces the OffloadingManager to be group-aware.
I suggest a different approach, which also aligns well with the P2P tier work by @liranschour and @ronensc:

Goal

Keep OffloadingManager completely group-unaware. The scheduler — which already iterates per group — handles the splitting.

API changes to OffloadingManager

Method	Signature	Purpose
prepare_store	(keys, req_context) -> Collection[OffloadKey] | None	Allocate space, return accepted keys (or None if full). No spec, no evicted_keys.
prepare_load	(keys, req_context) -> None	Pin blocks for reading (protect from eviction). No spec returned.
get_spec (new)	(keys) -> LoadStoreSpec	Return medium-specific addressing for the given keys. Called per group by the scheduler.

[hickeyma]

Thanks @orozery, this seems a better approach.

You're right that the per group return value pulls group awareness into the manager where it doesn't belong. prepare_write([key], ...) ends up returning a list of empty specs with a single non-empty one and the caller has to index back in with get_offload_group_idx(key) to find the one that matters. The manager is handling something that is really the scheduler's job as it already walks per group. Pulling get_spec out as its own thing reads much cleaner.

It also helps that this lines up with the P2P tier work, so I'd rather we converge on one manager API than have it drift. A couple of things I'd like to nail down with you (and @liranschour / @ronensc) before I refactor so we are all on the same path:

• evicted_keys: dropping it from prepare_store sounds right since the CPU manager already emits eviction events via take_events(). I want to double check that nothing downstream still relies on the returned list before removing it?
• get_spec granularity: called once per group with that group's keys returning a single LoadStoreSpec. Does that match what you and the P2P side need, or do you want it to take a batch?

If that all sounds good I'll update this PR rather than leaving the manager half migrated. Happy to break it out into a follow-up instead if you'd prefer to keep this one focused on the descriptor reshape. I supoose whatever works best for the P2P timeline. Let me know what you think.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hickeyma.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork