karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[KV Offload] Define clean backend configuration boundary (#48150)

Browse Source 
Signed-off-by: Change72 <changg@nvidia.com>
Signed-off-by: Chang Guo <cguo51@asu.edu>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Or Ozeri <or@ozery.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Chang Guo
2026-07-16 13:27:05 +03:00
committed by GitHub
co-authored by Codex Claude Fable 5 Or Ozeri OpenAI Codex
parent 8c3393f373
commit a9531edfa6
27 changed files with 1287 additions and 855 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+25 -22
View File
@@ -8,6 +8,9 @@ import torch
from tests.v1.kv_connector.unit.utils import create_vllm_config
from vllm.config import KVEventsConfig, KVTransferConfig
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
OffloadingEventGroupSpec,
OffloadingEventsTracker,
@@ -70,15 +73,15 @@ def _group_config(
*,
group_idx: int = 0,
block_size: int = 4,
block_size_factor: int = 1,
sliding_window_size_in_blocks: int | None = None,
blocks_per_chunk: int = 1,
sliding_window_size_in_chunks: int | None = None,
) -> GroupOffloadConfig:
return GroupOffloadConfig(
group_idx=group_idx,
gpu_block_size=block_size,
offloaded_block_size=block_size * block_size_factor,
hash_block_size_factor=block_size_factor,
sliding_window_size_in_blocks=sliding_window_size_in_blocks,
tokens_per_block=block_size,
tokens_per_chunk=block_size * blocks_per_chunk,
hashes_per_chunk=blocks_per_chunk,
sliding_window_size_in_chunks=sliding_window_size_in_chunks,
kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
)
@@ -90,7 +93,7 @@ def _record_chunks(
num_chunks: int,
) -> list[OffloadKey]:
keys: list[OffloadKey] = []
hbf = group_config.hash_block_size_factor
hbf = group_config.hashes_per_chunk
for chunk_idx in range(num_chunks):
tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
assert tail_hash is not None
@@ -149,14 +152,14 @@ def test_take_events_publishes_routable_block_stored():
def test_take_events_factor_gt_1_chunk_store_and_remove():
block_size = 4
block_size_factor = 3
blocks_per_chunk = 3
tracker = _tracker()
group_config = _group_config(
block_size=block_size, block_size_factor=block_size_factor
block_size=block_size, blocks_per_chunk=blocks_per_chunk
)
req = _request(
block_hashes=[_hash(i) for i in range(6)],
token_count=block_size * block_size_factor * 2,
token_count=block_size * blocks_per_chunk * 2,
)
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
@@ -169,17 +172,17 @@ def test_take_events_factor_gt_1_chunk_store_and_remove():
expected_chunk_hashes = [
_wire_hash(_hash(i))
for i in range(
chunk_idx * block_size_factor,
(chunk_idx + 1) * block_size_factor,
chunk_idx * blocks_per_chunk,
(chunk_idx + 1) * blocks_per_chunk,
)
]
assert event.block_hashes == expected_chunk_hashes
assert event.block_size == block_size
assert len(event.token_ids) == block_size * block_size_factor
assert len(event.token_ids) == block_size * blocks_per_chunk
if chunk_idx == 0:
assert event.parent_block_hash is None
else:
assert event.parent_block_hash == _wire_hash(_hash(block_size_factor - 1))
assert event.parent_block_hash == _wire_hash(_hash(blocks_per_chunk - 1))
expected_hashes.extend(expected_chunk_hashes)
assert len(tracker._pending_event_metadata) == 2
@@ -194,12 +197,12 @@ def test_take_events_factor_gt_1_chunk_store_and_remove():
def test_take_events_factor_gt_1_store_is_order_independent():
block_size_factor = 3
blocks_per_chunk = 3
tracker = _tracker()
group_config = _group_config(block_size_factor=block_size_factor)
group_config = _group_config(blocks_per_chunk=blocks_per_chunk)
req = _request(
block_hashes=[_hash(i) for i in range(6)],
token_count=4 * block_size_factor * 2,
token_count=4 * blocks_per_chunk * 2,
)
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
unknown_key = make_offload_key(_hash(12345), 0)
@@ -244,7 +247,7 @@ def test_take_events_opt_out_keeps_placeholders():
def test_record_store_skips_sliding_window_group():
tracker = _tracker()
group_config = _group_config(sliding_window_size_in_blocks=2)
group_config = _group_config(sliding_window_size_in_chunks=2)
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
@@ -258,8 +261,8 @@ def test_record_store_skips_sliding_window_group():
def test_take_events_groups_removed_hashes_by_kv_group():
tracker = _tracker()
group0_config = _group_config(group_idx=0, block_size_factor=2)
group1_config = _group_config(group_idx=1, block_size_factor=2)
group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
group1_config = _group_config(group_idx=1, blocks_per_chunk=2)
req0 = _request(block_hashes=[_hash(0), _hash(1)], token_count=8)
req1 = _request(block_hashes=[_hash(10), _hash(11)], token_count=8)
key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0]
@@ -293,7 +296,7 @@ def test_take_events_supports_restore_after_eviction():
assert not tracker._pending_event_metadata
req.all_token_ids = [5, 6, 7, 8]
tracker.record_store(req, group_config, offload_block_idx=0, offload_key=key)
tracker.record_store(req, group_config, chunk_idx=0, offload_key=key)
second_store = list(tracker.take_events([_stored_event([key])]))
assert len(second_store) == 1
@@ -351,4 +354,4 @@ def test_tiering_rejects_self_describing_kv_events():
)
with pytest.raises(ValueError, match="TieringOffloadingSpec"):
TieringOffloadingSpec(vllm_config, kv_cache_config)
TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+167 -169
View File
@@ -105,20 +105,20 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_offloading_connector(request_runner, async_scheduling: bool):
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# 3 blocks, store just the middle block (skip first and last)
# blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(list(keys)[1:2])
)
@@ -126,7 +126,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
# add block missing 1 token -> no offload
runner.run(
decoded_tokens=[0] * (offloaded_block_size - 1),
decoded_tokens=[0] * (tokens_per_chunk - 1),
expected_stored=(3, 4, 5),
)
runner.manager.touch.assert_not_called()
@@ -141,7 +141,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
)
runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1))
# 1 more block (+ token for kicking off offloading)
# now check touch was called with all 6 blocks
@@ -149,7 +149,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
generate_store_output(keys)
)
runner.run(
decoded_tokens=[0] * (offloaded_block_size + 1),
decoded_tokens=[0] * (tokens_per_chunk + 1),
expected_stored=(15, 16, 17),
)
runner.manager.touch.assert_called()
@@ -160,7 +160,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.run(decoded_tokens=[EOS_TOKEN_ID])
# create a new request differing only on the last token
runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1])
runner.run(decoded_tokens=[0])
runner.manager.touch.assert_called()
block_hashes2 = list(runner.manager.touch.call_args.args[0])
@@ -173,12 +173,12 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
# terminate request
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored=tuple(range(6 * block_size_factor)),
expected_stored=tuple(range(6 * blocks_per_chunk)),
)
# full_block_tokens - num_computed_tokens < offloaded_block_size
# full_block_tokens - num_computed_tokens < tokens_per_chunk
runner.new_request(
token_ids=[0] * block_size + [1] * (offloaded_block_size - block_size)
token_ids=[0] * block_size + [1] * (tokens_per_chunk - block_size)
)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
@@ -187,7 +187,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.manager.lookup.assert_not_called()
# single block lookup with no hits
runner.new_request(token_ids=[1] * offloaded_block_size)
runner.new_request(token_ids=[1] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
)
@@ -196,7 +196,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
# single block lookup with a hit
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
)
@@ -204,9 +204,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
# single block lookup with a hit in a middle block
runner.new_request(
token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
)
runner.new_request(token_ids=[0] * tokens_per_chunk * 2 + [1] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
)
@@ -217,15 +215,15 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_request_preemption(request_runner, async_scheduling: bool):
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
@@ -233,7 +231,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
# 2 blocks, store all, without flushing
# blocks = [0, 1, 2], [3, 4, 5]
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -247,7 +245,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
generate_store_output(keys)
)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
decoded_tokens=[0] * (2 * tokens_per_chunk - block_size),
complete_transfers=False,
)
@@ -297,14 +295,14 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
still arrive afterward for already-submitted transfer jobs.
"""
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
runner = request_runner(
block_size=block_size,
num_gpu_blocks=100,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# Record the order of per-request connector calls on the (mocked) manager.
@@ -324,10 +322,10 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
# Decode a couple of blocks, keeping every transfer in flight
# (complete_transfers=False) so no store completes while the request runs.
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
runner.run(decoded_tokens=[0], complete_transfers=False)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size),
decoded_tokens=[0] * (2 * tokens_per_chunk),
complete_transfers=False,
)
@@ -347,7 +345,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
runner.run(
decoded_tokens=[],
complete_transfers=True,
expected_stored=tuple(range(4 * block_size_factor)),
expected_stored=tuple(range(4 * blocks_per_chunk)),
)
# on_request_finished is issued exactly once.
@@ -364,19 +362,19 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -387,7 +385,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
# start a request to load the first block, but don't complete
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
runner.run(
decoded_tokens=[],
@@ -399,7 +397,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
assert transfer_jobs
# start a new request to load the same first block
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
runner.run(
decoded_tokens=[],
@@ -428,19 +426,19 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_abort_loading_requests(request_runner, async_scheduling: bool):
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -451,7 +449,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
# start a request to load the first block, but don't complete
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
runner.run(
decoded_tokens=[],
@@ -483,7 +481,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bool):
block_size = 4
num_gpu_blocks = 100
# sliding_window=8 -> 2 offloaded blocks (block_size_factor=1)
# sliding_window=8 -> 2 offloaded chunks (blocks_per_chunk=1)
sliding_window = 8
kv_cache_groups = [
@@ -518,8 +516,8 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
# Verify group configs: group 0 = full attention, group 1 = sliding window
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
assert len(kv_group_configs) == 2
assert kv_group_configs[0].sliding_window_size_in_blocks is None
assert kv_group_configs[1].sliding_window_size_in_blocks == 2
assert kv_group_configs[0].sliding_window_size_in_chunks is None
assert kv_group_configs[1].sliding_window_size_in_chunks == 2
# Blocks [0, 1, 2] miss
runner.new_request(token_ids=[0] * block_size * 3)
@@ -587,16 +585,16 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool):
hash_block_size = 4
tokens_per_hash = 4
num_gpu_blocks = 100
# Group 0: block_size=12 (offloaded_block_size=12)
# Group 1: block_size=16 (offloaded_block_size=16)
# Group 0: block_size=12 (tokens_per_chunk=12)
# Group 1: block_size=16 (tokens_per_chunk=16)
kv_cache_groups = [
KVCacheGroupSpec(
["layer0"],
FullAttentionSpec(
block_size=hash_block_size * 3,
block_size=tokens_per_hash * 3,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
@@ -605,7 +603,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
KVCacheGroupSpec(
["layer1"],
FullAttentionSpec(
block_size=hash_block_size * 4,
block_size=tokens_per_hash * 4,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
@@ -614,7 +612,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
]
runner = request_runner(
block_size=hash_block_size,
block_size=tokens_per_hash,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
@@ -623,10 +621,10 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
# Verify group configs
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
assert len(kv_group_configs) == 2
assert kv_group_configs[0].gpu_block_size == 12
assert kv_group_configs[0].offloaded_block_size == 12
assert kv_group_configs[1].gpu_block_size == 16
assert kv_group_configs[1].offloaded_block_size == 16
assert kv_group_configs[0].tokens_per_block == 12
assert kv_group_configs[0].tokens_per_chunk == 12
assert kv_group_configs[1].tokens_per_block == 16
assert kv_group_configs[1].tokens_per_chunk == 16
# Prompt: 25 tokens, unaligned to both block sizes.
# Group 0 blocks: [0, 1], ending_token_offset = 24
@@ -932,20 +930,20 @@ class TestSlidingWindowLookup:
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):
"""With REQUEST_LEVEL policy, all blocks are stored — including prefix hits."""
gpu_block_size = 4
block_size_factor = 3
offloaded_block_size = gpu_block_size * block_size_factor
tokens_per_block = 4
blocks_per_chunk = 3
tokens_per_chunk = tokens_per_block * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size_factor=block_size_factor,
block_size=gpu_block_size,
blocks_per_chunk=blocks_per_chunk,
block_size=tokens_per_block,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# Store 1 offloaded block (3 GPU blocks) via a normal request.
runner.new_request(token_ids=[0] * offloaded_block_size)
# Store 1 offloaded chunk (3 GPU blocks) via a normal request.
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -962,14 +960,14 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling
policy=OffloadPolicy.REQUEST_LEVEL
)
# New request with 2 offloaded blocks; first matches what's in CPU.
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
# New request with 2 offloaded chunks; first matches what's in CPU.
runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
# Load the first offloaded block from CPU.
# Load the first offloaded chunk from CPU.
runner.run(decoded_tokens=[0], expected_loaded=(0, 1, 2))
# Store must include ALL 6 GPU blocks (both the loaded prefix and
@@ -989,7 +987,7 @@ def test_loads_do_not_populate_fence_index(request_runner):
"""Loads don't populate _block_id_to_pending_jobs (protected by
delay_free_blocks while in flight)."""
runner = request_runner(
block_size_factor=3,
blocks_per_chunk=3,
block_size=4,
num_gpu_blocks=100,
async_scheduling=False,
@@ -1008,7 +1006,7 @@ def test_fence_at_update_state_after_alloc(request_runner):
req1 just freed.
"""
runner = request_runner(
block_size_factor=1,
blocks_per_chunk=1,
block_size=4,
num_gpu_blocks=2,
async_scheduling=False,
@@ -1059,7 +1057,7 @@ def test_fence_at_build_store_jobs(request_runner):
reusing a finished request's pending-store block is flushed by
_build_store_jobs's fence."""
runner = request_runner(
block_size_factor=1,
blocks_per_chunk=1,
block_size=4,
num_gpu_blocks=2,
async_scheduling=False,
@@ -1108,16 +1106,16 @@ def test_fence_at_build_store_jobs(request_runner):
def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
"""complete_store fires per-job, not deferred to request finish.
Each call carries only that store's keys."""
gpu_block_size = 4
block_size_factor = 3
offloaded_block_size = gpu_block_size * block_size_factor
tokens_per_block = 4
blocks_per_chunk = 3
tokens_per_chunk = tokens_per_block * blocks_per_chunk
runner = request_runner(
block_size_factor=block_size_factor,
block_size=gpu_block_size,
blocks_per_chunk=blocks_per_chunk,
block_size=tokens_per_block,
num_gpu_blocks=100,
async_scheduling=async_scheduling,
)
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -1131,7 +1129,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
# Second store: fires when block 1 is fully populated, with different keys.
runner.run(
decoded_tokens=[0] * (offloaded_block_size + 1),
decoded_tokens=[0] * (tokens_per_chunk + 1),
expected_stored=(3, 4, 5),
)
assert runner.manager.complete_store.call_count == 1
@@ -1148,25 +1146,25 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
"""Validates max_offload_tokens: type coercion, boundary values, and capping.
Setup: 3 offloaded blocks × 3 GPU blocks each = 9 GPU block offsets (0–8).
Setup: 3 offloaded chunks × 3 GPU blocks each = 9 GPU block offsets (0–8).
"""
gpu_block_size = 4
block_size_factor = 3
offloaded_block_size = gpu_block_size * block_size_factor # 12
tokens_per_block = 4
blocks_per_chunk = 3
tokens_per_chunk = tokens_per_block * blocks_per_chunk # 12
num_gpu_blocks = 100
all_offsets = (0, 1, 2, 3, 4, 5, 6, 7, 8)
def make_runner():
return request_runner(
block_size=gpu_block_size,
block_size=tokens_per_block,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
def setup(r, max_offload_tokens):
r.new_request(
token_ids=[0] * offloaded_block_size * 3,
token_ids=[0] * tokens_per_chunk * 3,
kv_transfer_params={"max_offload_tokens": max_offload_tokens},
)
r.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -1228,9 +1226,9 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
setup(r, 0)
r.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
# positive int cap -> limits offload to first 2 offloaded blocks (offsets 0–5)
# positive int cap -> limits offload to first 2 chunks (offsets 0–5)
r = make_runner()
setup(r, 24) # 24 tokens = 2 offloaded blocks × 12 tokens each
setup(r, 24) # 24 tokens = 2 offloaded chunks × 12 tokens each
r.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored=(0, 1, 2, 3, 4, 5),
@@ -1242,8 +1240,8 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
def test_offload_prompt_only(request_runner, async_scheduling: bool):
"""offload_prompt_only=True offloads prompt blocks but never decode blocks.
Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill
4 more offloaded blocks. The flag clamps the offloadable token count to the
Setup: a 2-chunk prompt followed by enough decode tokens to fill
4 more offloaded chunks. The flag clamps the offloadable token count to the
prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever
eligible for store; the decode blocks (offsets >= 6) are skipped.
@@ -1253,18 +1251,18 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):
subtleties. The decode steps are still enough for the prompt store to
complete and show up in expected_stored.
"""
gpu_block_size = 4
block_size_factor = 3
offloaded_block_size = gpu_block_size * block_size_factor # 12
tokens_per_block = 4
blocks_per_chunk = 3
tokens_per_chunk = tokens_per_block * blocks_per_chunk # 12
num_prompt_blocks = 2
num_decode_blocks = 4
prompt_offsets = (0, 1, 2, 3, 4, 5)
runner = request_runner(
block_size=gpu_block_size,
block_size=tokens_per_block,
num_gpu_blocks=100,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
extra_config_overrides={"offload_prompt_only": True},
)
@@ -1272,9 +1270,9 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):
generate_store_output(keys)
)
runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks)
runner.new_request(token_ids=[0] * tokens_per_chunk * num_prompt_blocks)
runner.run(
decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks),
decoded_tokens=[0] * (tokens_per_chunk * num_decode_blocks),
expected_stored=prompt_offsets,
)
@@ -1291,21 +1289,21 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_reset_cache(request_runner, async_scheduling: bool):
"""reset_cache flushes in-flight loads, calls manager.reset_cache(), resets
next_stored_block_idx for active requests and clears job tracking."""
next_stored_chunk_idx for active requests and clears job tracking."""
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# Store 1 offloaded block (3 GPU blocks) to CPU.
runner.new_request(token_ids=[0] * offloaded_block_size)
# Store 1 offloaded chunk (3 GPU blocks) to CPU.
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -1317,7 +1315,7 @@ def test_reset_cache(request_runner, async_scheduling: bool):
# Reset GPU prefix cache then start a request that loads from CPU.
# Leave the load in-flight so that reset_cache must flush it.
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
@@ -1335,11 +1333,11 @@ def test_reset_cache(request_runner, async_scheduling: bool):
# Record job counter to verify the reset counter is set correctly.
job_counter_before_reset = runner.connector_scheduler._job_counter
# After update_state_after_alloc, next_stored_block_idx is advanced to
# After update_state_after_alloc, next_stored_chunk_idx is advanced to
# skip the loaded prefix; reset_cache must bring it back to 0.
for req_status in runner.connector_scheduler._req_status.values():
for group_state in req_status.group_states:
assert group_state.next_stored_block_idx > 0
assert group_state.next_stored_chunk_idx > 0
# Reset the cache
runner.connector_scheduler.reset_cache()
@@ -1354,18 +1352,18 @@ def test_reset_cache(request_runner, async_scheduling: bool):
# All internal job tracking must be cleared.
assert not runner.connector_scheduler._jobs
assert not runner.connector_scheduler._block_id_to_pending_jobs
if runner.connector_scheduler._blocks_being_loaded is not None:
assert not runner.connector_scheduler._blocks_being_loaded
if runner.connector_scheduler._chunks_being_loaded is not None:
assert not runner.connector_scheduler._chunks_being_loaded
# Job reset counter must equal the job counter so that completions for
# pre-reset jobs arriving from workers are silently discarded.
assert runner.connector_scheduler._stale_job_threshold == job_counter_before_reset
# next_stored_block_idx must be reset to 0 for every active request so
# next_stored_chunk_idx must be reset to 0 for every active request so
# that post-reset stores restart from block 0.
for req_status in runner.connector_scheduler._req_status.values():
for group_state in req_status.group_states:
assert group_state.next_stored_block_idx == 0
assert group_state.next_stored_chunk_idx == 0
@pytest.mark.parametrize("async_scheduling", [True, False])
@@ -1376,14 +1374,14 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
without calling on_request_finished twice.
"""
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
runner = request_runner(
block_size=block_size,
num_gpu_blocks=100,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
finalized: list[str] = []
@@ -1396,10 +1394,10 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
# Decode a couple of blocks and keep every transfer in flight, so the
# request has pending store jobs.
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
runner.run(decoded_tokens=[0], complete_transfers=False)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size),
decoded_tokens=[0] * (2 * tokens_per_chunk),
complete_transfers=False,
)
@@ -1430,7 +1428,7 @@ def test_pending_transfer_defers_prefix_lookup():
With async scheduling, a preempted request's store can be flushed by the
worker before the scheduler consumes its completion. If the request is
re-admitted in that window, the connector should defer it instead of
looking up offloaded blocks and later asserting when a load is queued while
looking up offloaded chunks and later asserting when a load is queued while
the store job is still tracked.
"""
scheduler = object.__new__(OffloadingConnectorScheduler)
@@ -1464,27 +1462,27 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
re-admission path must defer while the scheduler still tracks the store.
"""
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
runner = request_runner(
block_size=block_size,
num_gpu_blocks=100,
async_scheduling=True,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
num_free_blocks_empty = free_block_queue.num_free_blocks
req_id = "0"
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
runner.run(decoded_tokens=[0], complete_transfers=False)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
decoded_tokens=[0] * (2 * tokens_per_chunk - block_size),
complete_transfers=False,
)
@@ -1530,8 +1528,8 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
- Group 0: full attention (MLA-like), block_size=16
- Group 1: SWA, block_size=4, sliding_window=8
alignment_block_count = 16 / 4 = 4 SWA blocks per alignment segment.
sliding_window_size_in_blocks = ceil(8 / 4) = 2.
alignment_chunk_count = 16 / 4 = 4 SWA blocks per alignment segment.
sliding_window_size_in_chunks = ceil(8 / 4) = 2.
Within each segment of 4 SWA blocks, only the trailing 2 are stored.
With 32 tokens (2 full-attn blocks, 8 SWA blocks):
@@ -1574,17 +1572,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
kv_cache_groups=kv_cache_groups,
)
# Verify config: alignment_block_count computed correctly
# Verify config: alignment_chunk_count computed correctly
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
assert len(kv_group_configs) == 2
# Group 0: full attention -> no alignment skip
assert kv_group_configs[0].alignment_block_count is None
assert kv_group_configs[0].sliding_window_size_in_blocks is None
assert kv_group_configs[0].offloaded_block_size == full_attn_block_size
# Group 1: SWA -> alignment_block_count = 16/4 = 4, tail = 2
assert kv_group_configs[1].alignment_block_count == 4
assert kv_group_configs[1].sliding_window_size_in_blocks == 2
assert kv_group_configs[1].offloaded_block_size == swa_block_size
assert kv_group_configs[0].alignment_chunk_count is None
assert kv_group_configs[0].sliding_window_size_in_chunks is None
assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size
# Group 1: SWA -> alignment_chunk_count = 16/4 = 4, tail = 2
assert kv_group_configs[1].alignment_chunk_count == 4
assert kv_group_configs[1].sliding_window_size_in_chunks == 2
assert kv_group_configs[1].tokens_per_chunk == swa_block_size
# Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks
# (block_size=4). Decode 1 token to kick off processing (stores are
@@ -1602,9 +1600,9 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
# Group 0 (full attn, block_size=16): 2 offloaded blocks
# Group 0 (full attn, block_size=16): 2 offloaded chunks
# -> GPU blocks (0, 0) and (0, 1)
# Group 1 (SWA, block_size=4): 8 offloaded blocks, skip first 2
# Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2
# per segment of 4:
# Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3)
# Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7)
@@ -1625,7 +1623,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
# Group 0: full prefix lookup hits 2 offloaded blocks
# Group 0: full prefix lookup hits 2 offloaded chunks
# -> loads GPU blocks (0, 0), (0, 1)
# Group 1: sliding window lookup finds trailing 2 from last segment
# (blocks 6, 7 which were stored)
@@ -1684,7 +1682,7 @@ def test_stale_sliding_window_block_after_prepare_store_failure(
runner.new_request(token_ids=[0] * block_size * 3)
# First step: prepare_store FAILS -> offloading delayed.
# next_stored_block_idx stays at 0, block_ids[0] still holds the
# next_stored_chunk_idx stays at 0, block_ids[0] still holds the
# original block_id for position 0.
runner.manager.prepare_store.side_effect = lambda keys, req_context: None
runner.run(decoded_tokens=[0])
@@ -1719,19 +1717,19 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
"""When skip_reading_prefix_cache=True, the offloading connector must not
load any blocks from CPU even if a matching prefix is cached there."""
block_size = 4
block_size_factor = 3
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 3
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# Populate the CPU offload cache with one block.
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -1747,7 +1745,7 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
# The offloading connector must not load anything from CPU, but must
# still offload the freshly computed blocks (state management intact).
runner.new_request(
token_ids=[0] * offloaded_block_size,
token_ids=[0] * tokens_per_chunk,
skip_reading_prefix_cache=True,
)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -2219,7 +2217,7 @@ class TestEagle:
def test_full_attn_store_excludes_trailing_decode_block(
self, request_runner, async_scheduling: bool
):
"""Eagle full-attention group excludes the trailing block only while
"""Eagle full-attention group excludes the trailing chunk only while
decoding.
Setup: 2 groups — group 0 is normal full-attention, group 1 is
@@ -2229,8 +2227,8 @@ class TestEagle:
draft-layer KV is volatile until the next block starts).
"""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 1
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
kv_cache_groups = [
@@ -2260,7 +2258,7 @@ class TestEagle:
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
@@ -2268,7 +2266,7 @@ class TestEagle:
assert not kv_group_configs[0].is_eagle_group
assert kv_group_configs[1].is_eagle_group
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -2292,7 +2290,7 @@ class TestEagle:
self, request_runner, async_scheduling: bool
):
"""Eagle sliding-window group stores all prompt blocks but excludes
the trailing block while decoding."""
the trailing chunk while decoding."""
block_size = 4
sliding_window = 8
num_gpu_blocks = 100
@@ -2321,7 +2319,7 @@ class TestEagle:
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
assert len(kv_group_configs) == 1
assert kv_group_configs[0].is_eagle_group
assert kv_group_configs[0].sliding_window_size_in_blocks == 2
assert kv_group_configs[0].sliding_window_size_in_chunks == 2
runner.new_request(token_ids=[0] * block_size * 3)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -2340,8 +2338,8 @@ class TestEagle:
"""An eagle group with a single-block prompt stores it at the end of
prefill: prompt blocks are stable, so no tail is held back."""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 1
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
kv_cache_groups = [
@@ -2362,10 +2360,10 @@ class TestEagle:
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -2378,18 +2376,18 @@ class TestEagle:
"""Eagle store must not drop interior blocks across prefill chunks.
Regression: the trailing-block exclusion (num_blocks - 1) was applied
when collecting keys, but next_stored_block_idx advanced by the
non-decremented count, so the trailing block of every chunked-prefill
when collecting keys, but next_stored_chunk_idx advanced by the
non-decremented count, so the trailing chunk of every chunked-prefill
chunk was skipped and never re-considered. With the harness chunk budget
(1000 tokens) and block_size 4, a prompt longer than one chunk lost the
block at the chunk boundary, leaving a permanent gap that caps prefix
reuse at the first hole. Only the trailing decode block may be held
back; all other blocks must be stored exactly once (no duplicates from
next_stored_block_idx regressing at the prefill->decode transition).
next_stored_chunk_idx regressing at the prefill->decode transition).
"""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 1
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 1000
kv_cache_groups = [
@@ -2409,13 +2407,13 @@ class TestEagle:
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group
# Prompt spans more than one prefill chunk (chunk budget 1000 tokens).
num_blocks = 256
runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks)
runner.new_request(token_ids=[0] * tokens_per_chunk * num_blocks)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -2429,7 +2427,7 @@ class TestEagle:
for b in t.gpu_blocks
)
# The stored blocks must be contiguous from 0: no interior block is
# dropped at a chunk boundary. (The bug left a gap at offloaded block
# dropped at a chunk boundary. (The bug left a gap at offloaded chunk
# 249, the tail of the first 1000-token chunk.)
assert offsets == list(range(len(offsets))), (
f"interior hole in stored blocks: {offsets}"
@@ -2439,14 +2437,14 @@ class TestEagle:
def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool):
"""Eagle group constrains load: convergence tightens both groups.
Store 3 offloaded blocks per group (all prompt blocks, so the eagle
Store 3 offloaded chunks per group (all prompt chunks, so the eagle
group stores all 3 as well). Then a new request loads from CPU. The
eagle group pops its trailing hit block on load, tightening the hit
to 2 blocks for both groups.
"""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 1
tokens_per_chunk = block_size * blocks_per_chunk
num_gpu_blocks = 100
kv_cache_groups = [
@@ -2476,10 +2474,10 @@ class TestEagle:
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -2497,7 +2495,7 @@ class TestEagle:
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1])
runner.new_request(token_ids=[0] * tokens_per_chunk * 3 + [1])
runner.manager.lookup.return_value = LookupResult.HIT
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output([])
@@ -2527,8 +2525,8 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
GPU blocks before the store completes.
"""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 1
tokens_per_chunk = block_size * blocks_per_chunk
# Use 2 GPU blocks so the second run reuses the same blocks,
# triggering a fence-based flush of the in-flight job from run 1.
@@ -2536,11 +2534,11 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
block_size=block_size,
num_gpu_blocks=2,
async_scheduling=False,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# 4 prompt tokens → 1 GPU block (block 0)
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -2576,7 +2574,7 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
# Run 2: block reuse triggers fence-based flush → cleanup.
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
@@ -2603,26 +2601,26 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):
- Run 3: block reuse → both jobs flushed via fence
"""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
blocks_per_chunk = 1
tokens_per_chunk = block_size * blocks_per_chunk
# 4 GPU blocks: block 0 is null, blocks 1-3 are usable.
runner = request_runner(
block_size=block_size,
num_gpu_blocks=4,
async_scheduling=False,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
)
# Prompt: 4 tokens → block 1
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
# Run 1: 4 decoded tokens → block 2 full → job_0 created for block 1.
runner.run(
decoded_tokens=[0] * offloaded_block_size,
decoded_tokens=[0] * tokens_per_chunk,
complete_transfers=False,
)
assert len(runner.connector_scheduler._jobs) >= 1
@@ -2630,7 +2628,7 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):
# Run 2: 4 more tokens + EOS → block 3 full → more jobs created.
# Request finishes → all jobs registered in fence.
runner.run(
decoded_tokens=[0] * offloaded_block_size + [EOS_TOKEN_ID],
decoded_tokens=[0] * tokens_per_chunk + [EOS_TOKEN_ID],
complete_transfers=False,
)
num_jobs = len(runner.connector_scheduler._jobs)
@@ -2638,7 +2636,7 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):
# Run 3: block reuse → fence flushes both jobs.
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+4 -3
View File
@@ -96,11 +96,12 @@ def _make_worker(kv_cache_config: KVCacheConfig):
)
spec = MagicMock(spec=OffloadingSpec)
spec.kv_cache_config = kv_cache_config
spec.vllm_config = MagicMock()
spec.get_worker.return_value = MagicMock()
worker = OffloadingConnectorWorker(spec=spec)
worker = OffloadingConnectorWorker(
spec=spec,
kv_cache_config=kv_cache_config,
)
worker.worker = MagicMock()
return worker, spec
tests/v1/kv_connector/unit/offloading_connector/utils.py
+17 -19
View File
@@ -17,7 +17,6 @@ from vllm import SamplingParams
from vllm.config import (
KVEventsConfig,
KVTransferConfig,
VllmConfig,
set_current_vllm_config,
)
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
@@ -57,6 +56,7 @@ from vllm.v1.kv_offload.base import (
TransferResult,
make_offload_key,
)
from vllm.v1.kv_offload.config import OffloadingConfig
from vllm.v1.request import Request
from vllm.v1.structured_output import StructuredOutputManager
@@ -123,8 +123,8 @@ class MockOffloadingWorker(OffloadingWorker):
class MockOffloadingSpec(OffloadingSpec):
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
super().__init__(vllm_config, kv_cache_config)
def __init__(self, config: OffloadingConfig):
super().__init__(config)
self.manager = MagicMock(spec=OffloadingManager)
self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
@@ -175,17 +175,17 @@ class RequestRunner:
self,
block_size: int,
num_gpu_blocks: int,
block_size_factor: int = 1,
blocks_per_chunk: int = 1,
async_scheduling: bool = True,
kv_cache_groups: list[KVCacheGroupSpec] | None = None,
extra_config_overrides: dict[str, Any] | None = None,
):
assert block_size_factor == 1 or kv_cache_groups is None, (
"block_size_factor > 1 requires all groups to have the same "
assert blocks_per_chunk == 1 or kv_cache_groups is None, (
"blocks_per_chunk > 1 requires all groups to have the same "
"block size, so kv_cache_groups must be None (use default group)"
)
self.block_size_factor: int = block_size_factor
self.blocks_per_chunk: int = blocks_per_chunk
self.block_size: int = block_size
self.num_gpu_blocks: int = num_gpu_blocks
self.async_scheduling: bool = async_scheduling
@@ -208,8 +208,8 @@ class RequestRunner:
# opt-out tests override this to cover the legacy placeholders.
"self_describing_kv_events": True,
}
if block_size_factor > 1:
extra_config["block_size"] = block_size * block_size_factor
if blocks_per_chunk > 1:
extra_config["block_size"] = block_size * blocks_per_chunk
if extra_config_overrides:
extra_config.update(extra_config_overrides)
@@ -313,11 +313,9 @@ class RequestRunner:
self.connector_scheduler.config.kv_group_configs,
kv_cache_config.kv_cache_groups,
):
gpu_block_size = kv_cache_group.kv_cache_spec.block_size
assert group_config.gpu_block_size == gpu_block_size
assert (
group_config.offloaded_block_size == gpu_block_size * block_size_factor
)
tokens_per_block = kv_cache_group.kv_cache_spec.block_size
assert group_config.tokens_per_block == tokens_per_block
assert group_config.tokens_per_chunk == tokens_per_block * blocks_per_chunk
# extract OffloadingSpec of worker_connector
connector_worker = self.worker_connector.connector_worker
@@ -389,7 +387,7 @@ class RequestRunner:
for block_id in dst_spec.block_ids:
self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
block_size_factor = self.block_size_factor
blocks_per_chunk = self.blocks_per_chunk
for src_spec, dst_spec in self.offloading_spec.get_completed_transfers():
if isinstance(src_spec, GPULoadStoreSpec):
@@ -412,7 +410,7 @@ class RequestRunner:
# list of (offload_key, sub_block_offset)
offload_addresses: list[Any] = []
for offload_key in offload_spec.offload_keys:
for sub_block_idx in range(block_size_factor):
for sub_block_idx in range(blocks_per_chunk):
offload_addresses.append((offload_key, sub_block_idx))
assert gpu_spec.block_indices is not None
@@ -426,7 +424,7 @@ class RequestRunner:
gpu_block_end_offset = gpu_block_offset + group_size
assert gpu_block_end_offset <= len(gpu_blocks)
offload_addresses_to_skip = logical_offset % block_size_factor
offload_addresses_to_skip = logical_offset % blocks_per_chunk
offload_addresses_end_offset = (
offload_address_offset + offload_addresses_to_skip + group_size
)
@@ -651,14 +649,14 @@ def request_runner():
block_size,
num_gpu_blocks,
async_scheduling,
block_size_factor=1,
blocks_per_chunk=1,
kv_cache_groups=None,
extra_config_overrides=None,
):
runner = RequestRunner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
extra_config_overrides=extra_config_overrides,
tests/v1/kv_connector/unit/test_offloading_connector.py
+1 -1
View File
@@ -45,7 +45,7 @@ if current_platform.is_cuda():
# Falcon-H1: parallel hybrid (every layer has both attention and SSM).
# The mamba and attention groups end up with different GPU block sizes
# after page-size unification, so we leave cpu_block_size=None
# (block_size_factor stays 1).
# (blocks_per_chunk stays 1).
("tiiuae/Falcon-H1-0.5B-Instruct", None, None, True),
]
tests/v1/kv_offload/cpu/test_gpu_worker.py
+20 -20
View File
@@ -23,7 +23,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
NUM_GPU_BLOCKS = [64]
NUM_CPU_BLOCKS = [256]
GPU_PAGE_SIZES = [512, 1024]
BLOCK_SIZE_FACTORS = [1, 3]
BLOCKS_PER_CHUNK_VALUES = [1, 3]
NUM_TENSORS = [4]
SEEDS = [0]
DEVICE_TYPE = current_platform.device_type
@@ -35,7 +35,7 @@ NUM_MAPPINGS_PER_GROUP = [2]
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
@pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES)
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
@@ -48,7 +48,7 @@ def test_transfer(
gpu_to_cpu: bool,
num_mappings: int,
gpu_page_size_bytes: int,
block_size_factor: int,
blocks_per_chunk: int,
num_gpu_blocks: int,
num_cpu_blocks: int,
num_tensors: int,
@@ -92,7 +92,7 @@ def test_transfer(
mmap_region: SharedOffloadRegion | None = None
if use_shared_memory:
cpu_page_size = round_up(
gpu_page_size_bytes * num_tensors * block_size_factor,
gpu_page_size_bytes * num_tensors * blocks_per_chunk,
SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
)
mmap_region = SharedOffloadRegion(
@@ -105,25 +105,25 @@ def test_transfer(
worker = CPUOffloadingWorker(
kv_caches=kv_caches,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
num_cpu_blocks=num_cpu_blocks,
mmap_region=mmap_region,
)
# select block mappings
gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * block_size_factor)
gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * blocks_per_chunk)
cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
# expand cpu blocks to gpu-page granularity for uniform comparison:
# each cpu block maps to block_size_factor consecutive sub-blocks
# each cpu block maps to blocks_per_chunk consecutive sub-blocks
cpu_blocks_expanded = [
cpu_block * block_size_factor + j
cpu_block * blocks_per_chunk + j
for cpu_block in cpu_blocks
for j in range(block_size_factor)
for j in range(blocks_per_chunk)
]
# maybe skip some GPU blocks to test reading/writing from the middle of a CPU block
blocks_to_skip = block_size_factor - 1
blocks_to_skip = blocks_per_chunk - 1
if blocks_to_skip > 0:
gpu_blocks = gpu_blocks[blocks_to_skip:]
cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]
@@ -214,7 +214,7 @@ def test_transfer(
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
@pytest.mark.parametrize("num_mappings_per_group", NUM_MAPPINGS_PER_GROUP)
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
@pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES)
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
@pytest.mark.parametrize("seed", SEEDS)
@@ -225,7 +225,7 @@ def test_transfer_multi_group(
gpu_to_cpu: bool,
num_mappings_per_group: int,
gpu_page_size_bytes: int,
block_size_factor: int,
blocks_per_chunk: int,
num_gpu_blocks: int,
num_cpu_blocks: int,
seed: int,
@@ -234,7 +234,7 @@ def test_transfer_multi_group(
"""Test transfers with three KV cache groups:
- Group 0: aligned transfer with num_mappings_per_group blocks
- Group 1: zero blocks (empty group)
- Group 2: unaligned CPU->GPU transfer (logical_offset=block_size_factor-1,
- Group 2: unaligned CPU->GPU transfer (logical_offset=blocks_per_chunk-1,
causing the implementation to skip source sub-blocks) with
num_mappings_per_group blocks
"""
@@ -275,7 +275,7 @@ def test_transfer_multi_group(
worker = CPUOffloadingWorker(
kv_caches=canonical_kv_caches,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
num_cpu_blocks=num_cpu_blocks,
)
@@ -283,7 +283,7 @@ def test_transfer_multi_group(
group_sizes_in_cpu_blocks = [num_mappings_per_group, 0, num_mappings_per_group]
total_cpu_blocks = sum(group_sizes_in_cpu_blocks)
total_gpu_blocks_needed = total_cpu_blocks * block_size_factor
total_gpu_blocks_needed = total_cpu_blocks * blocks_per_chunk
gpu_blocks_all = random.sample(range(num_gpu_blocks), total_gpu_blocks_needed)
cpu_blocks_all = random.sample(range(num_cpu_blocks), total_cpu_blocks)
@@ -293,7 +293,7 @@ def test_transfer_multi_group(
gpu_offset = 0
cpu_offset = 0
for size in group_sizes_in_cpu_blocks:
gpu_count = size * block_size_factor
gpu_count = size * blocks_per_chunk
gpu_blocks_per_group.append(gpu_blocks_all[gpu_offset : gpu_offset + gpu_count])
cpu_blocks_per_group.append(cpu_blocks_all[cpu_offset : cpu_offset + size])
gpu_offset += gpu_count
@@ -302,15 +302,15 @@ def test_transfer_multi_group(
# expand cpu blocks to gpu-page granularity
cpu_blocks_expanded_per_group = [
[
cpu_block * block_size_factor + j
cpu_block * blocks_per_chunk + j
for cpu_block in cpu_blocks
for j in range(block_size_factor)
for j in range(blocks_per_chunk)
]
for cpu_blocks in cpu_blocks_per_group
]
# skip sub-blocks from group 2 to test unaligned transfers.
sub_blocks_to_skip = block_size_factor - 1 # e.g. 2 when block_size_factor=3
sub_blocks_to_skip = blocks_per_chunk - 1 # e.g. 2 when blocks_per_chunk=3
if sub_blocks_to_skip > 0:
gpu_blocks_per_group[2] = gpu_blocks_per_group[2][
sub_blocks_to_skip:-sub_blocks_to_skip
@@ -347,7 +347,7 @@ def test_transfer_multi_group(
cpu_blocks_expanded_per_group, gpu_blocks_per_group
)
]
num_dst_sub_blocks = num_cpu_blocks * block_size_factor
num_dst_sub_blocks = num_cpu_blocks * blocks_per_chunk
else:
handler = worker._load_handler
src_spec = CPULoadStoreSpec(cpu_blocks)
tests/v1/kv_offload/test_factory.py
+329 -13
View File
@@ -11,17 +11,33 @@ These tests verify:
4. Error paths — unregistered specs, missing config, duplicate registration.
"""
from typing import cast
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.config import KVTransferConfig
from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
)
from vllm.platforms import current_platform
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
SlidingWindowSpec,
)
from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec
from vllm.v1.kv_offload.base import (
CanonicalKVCaches,
OffloadingHistogramMetadata,
OffloadingManager,
OffloadingSpec,
OffloadingWorker,
)
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
from vllm.v1.kv_offload.factory import OffloadingSpecFactory
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
@@ -39,6 +55,17 @@ def restore_registry():
OffloadingSpecFactory._registry = original
def _get_extra_config(config: VllmConfig) -> dict:
assert config.kv_transfer_config is not None
return config.kv_transfer_config.kv_connector_extra_config
def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec:
return OffloadingSpecFactory.create_spec(
build_offloading_config(config, kv_cache_config)
)
def _make_vllm_config(
spec_name: str | None = "CPUOffloadingSpec",
cpu_bytes_to_use: int | None = None,
@@ -95,6 +122,46 @@ def _make_vllm_config(
)
def _make_layout_vllm_config(
spec_name: str = "CPUOffloadingSpec",
cpu_bytes_to_use: int | None = None,
extra_config: dict | None = None,
tensor_parallel_size: int = 1,
pipeline_parallel_size: int = 1,
prefill_context_parallel_size: int = 1,
decode_context_parallel_size: int = 1,
) -> VllmConfig:
config = MagicMock()
config.cache_config.block_size = 16
config.cache_config.enable_prefix_caching = True
config.cache_config.prefix_match_unit = None
config.cache_config.cache_dtype = torch.float16
config.model_config.model = "test-model"
world_size = (
tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
)
with patch.object(current_platform, "device_count", return_value=world_size):
config.parallel_config = ParallelConfig(
tensor_parallel_size=tensor_parallel_size,
pipeline_parallel_size=pipeline_parallel_size,
prefill_context_parallel_size=prefill_context_parallel_size,
decode_context_parallel_size=decode_context_parallel_size,
)
config.kv_events_config = None
config.use_v2_model_runner = False
connector_extra_config = dict(extra_config or {})
connector_extra_config["spec_name"] = spec_name
if cpu_bytes_to_use is not None:
connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
config.kv_transfer_config = KVTransferConfig(
kv_connector="OffloadingConnector",
kv_role="kv_both",
kv_connector_extra_config=connector_extra_config,
)
return cast(VllmConfig, config)
def _make_kv_cache_config():
"""Build a minimal KVCacheConfig with one KV cache tensor."""
num_blocks = 16
@@ -122,6 +189,78 @@ def _make_kv_cache_config():
)
def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
num_blocks = 4
if packed:
kv_cache_tensors = [
KVCacheTensor(
size=64,
shared_by=[layer_name],
block_stride=16,
)
for layer_name in ("layer0", "layer1")
]
else:
kv_cache_tensors = [
KVCacheTensor(size=40, shared_by=["layer0"]),
KVCacheTensor(size=24, shared_by=["layer1"]),
]
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=[
KVCacheGroupSpec(
["layer0", "layer1"],
FullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
)
],
)
def _make_hybrid_kv_cache_config() -> KVCacheConfig:
return KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(size=40, shared_by=["full_layer"]),
KVCacheTensor(size=24, shared_by=["mla_layer"]),
],
kv_cache_groups=[
KVCacheGroupSpec(
["full_layer"],
FullAttentionSpec(
block_size=12,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["mla_layer"],
MLAAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
),
),
],
)
class SingleArgExternalOffloadingSpec(OffloadingSpec):
def get_manager(self) -> OffloadingManager:
raise NotImplementedError
def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
raise NotImplementedError
# ---------------------------------------------------------------------------
# Pre-registration integrity (CI sentinel)
# ---------------------------------------------------------------------------
@@ -154,7 +293,7 @@ def test_tiering_spec_registered():
def test_get_spec_cls_returns_registered_class():
"""Registered spec_name returns correct class."""
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
assert spec_cls is CPUOffloadingSpec
@@ -162,7 +301,7 @@ def test_get_spec_cls_default_to_cpu():
"""Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
config = _make_vllm_config(spec_name=None)
config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
assert spec_cls is CPUOffloadingSpec
@@ -176,16 +315,193 @@ def test_create_cpu_offloading_spec_end_to_end():
Verifies:
- cpu_bytes_to_use validation and num_blocks calculation
- block_size % hash_block_size assertion
- block_size % tokens_per_hash assertion
- spec instance is CPUOffloadingSpec
"""
config = _make_vllm_config(cpu_bytes_to_use=65536)
kv_cache_config = _make_kv_cache_config()
spec = OffloadingSpecFactory.create_spec(config, kv_cache_config)
spec = _create_spec(config, kv_cache_config)
assert isinstance(spec, CPUOffloadingSpec)
assert spec.num_blocks > 0
@pytest.mark.parametrize("packed", [False, True])
def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
cpu_bytes_to_use = 1920
config = _make_layout_vllm_config(
cpu_bytes_to_use=cpu_bytes_to_use,
extra_config={"block_size": 32},
tensor_parallel_size=3,
pipeline_parallel_size=2,
)
spec = _create_spec(config, _make_sizing_kv_cache_config(packed))
assert isinstance(spec, CPUOffloadingSpec)
assert spec.cpu_page_size_per_worker == 32
assert spec.kv_bytes_per_chunk == 192
assert spec.num_blocks == cpu_bytes_to_use // 192
def test_cpu_spec_rejects_partially_packed_tensor_layout():
config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.kv_cache_tensors[0].block_stride = 16
with pytest.raises(AssertionError):
_create_spec(config, kv_cache_config)
def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.num_blocks = 0
kv_cache_config.kv_cache_tensors[0].block_stride = 16
spec = _create_spec(config, kv_cache_config)
assert isinstance(spec, CPUOffloadingSpec)
assert spec.cpu_page_size_per_worker == 0
assert spec.kv_bytes_per_chunk == 0
assert spec.num_blocks == 0
def test_tiering_spec_aligns_row_size():
alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
cpu_bytes_to_use = alignment * 3
config = _make_layout_vllm_config(
spec_name="TieringOffloadingSpec",
cpu_bytes_to_use=cpu_bytes_to_use,
extra_config={"block_size": 32},
tensor_parallel_size=3,
pipeline_parallel_size=2,
)
spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False))
assert isinstance(spec, TieringOffloadingSpec)
assert spec.cpu_page_size_per_worker == 32
assert spec.kv_bytes_per_chunk == alignment
assert spec.num_blocks == cpu_bytes_to_use // alignment
def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():
config = _make_layout_vllm_config(
cpu_bytes_to_use=65536,
extra_config={"block_size": 64},
prefill_context_parallel_size=2,
)
spec = _create_spec(config, _make_kv_cache_config())
assert spec.tokens_per_block == (32,)
assert spec.tokens_per_hash == 32
assert spec.blocks_per_chunk == 2
def test_offloading_config_preserves_data_parallel_index():
config = _make_layout_vllm_config()
config.parallel_config.data_parallel_index = 2
offloading_config = build_offloading_config(config, _make_kv_cache_config())
assert offloading_config.parallel.data_parallel_index == 2
def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes():
config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
config.cache_config.block_size = 4
spec = _create_spec(config, _make_hybrid_kv_cache_config())
assert spec.tokens_per_block == (12, 16)
assert spec.tokens_per_hash == 4
assert spec.blocks_per_chunk == 1
def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec:
return FullAttentionSpec(
block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32
)
def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
config = _make_layout_vllm_config()
kv_cache_config = KVCacheConfig(
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
)
offloading_config = build_offloading_config(config, kv_cache_config)
return offloading_config.parallel.is_parallelism_agnostic
def test_parallelism_agnostic_for_single_full_attention_group():
assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
@pytest.mark.parametrize(
"kv_cache_groups",
[
# MLA latent KV is replicated per rank, never head-sharded.
[
KVCacheGroupSpec(
["l0"],
MLAAttentionSpec(
block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
),
)
],
# Sliding window is not full attention.
[
KVCacheGroupSpec(
["l0"],
SlidingWindowSpec(
block_size=16,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
sliding_window=128,
),
)
],
# Hybrid model: more than one KV cache group.
[
KVCacheGroupSpec(["l0"], _full_attention_spec()),
KVCacheGroupSpec(["l1"], _full_attention_spec()),
],
],
)
def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
assert not _parallelism_agnostic(kv_cache_groups)
def test_parallelism_agnostic_disabled_on_v2_model_runner():
config = _make_layout_vllm_config()
config.use_v2_model_runner = True
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
)
offloading_config = build_offloading_config(config, kv_cache_config)
assert not offloading_config.parallel.is_parallelism_agnostic
def test_create_dynamic_spec_receives_translated_config():
config = _make_layout_vllm_config(
spec_name="SingleArgExternalOffloadingSpec",
extra_config={
"spec_module_path": "tests.v1.kv_offload.test_factory",
},
)
kv_cache_config = _make_kv_cache_config()
offloading_config = build_offloading_config(config, kv_cache_config)
spec = OffloadingSpecFactory.create_spec(offloading_config)
assert isinstance(spec, SingleArgExternalOffloadingSpec)
assert spec.config is offloading_config
# ---------------------------------------------------------------------------
# Dynamic import via spec_module_path
# ---------------------------------------------------------------------------
@@ -205,7 +521,7 @@ def test_dynamic_load_via_spec_module_path():
config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
"vllm.v1.kv_offload.cpu.spec"
)
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
assert spec_cls is CPUOffloadingSpec
@@ -218,12 +534,12 @@ def test_unregistered_spec_without_module_path_raises():
"""spec_name not in registry + no spec_module_path → ValueError."""
config = _make_vllm_config(spec_name="NonexistentSpec")
with pytest.raises(ValueError, match="Unsupported spec type"):
OffloadingSpecFactory.get_spec_cls(config)
OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
# create_spec should also fail (calls get_spec_cls internally)
kv_cache_config = _make_kv_cache_config()
with pytest.raises(ValueError, match="Unsupported spec type"):
OffloadingSpecFactory.create_spec(config, kv_cache_config)
_create_spec(config, kv_cache_config)
def test_cpu_spec_missing_cpu_bytes_to_use_raises():
@@ -232,7 +548,7 @@ def test_cpu_spec_missing_cpu_bytes_to_use_raises():
config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
kv_cache_config = _make_kv_cache_config()
with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
OffloadingSpecFactory.create_spec(config, kv_cache_config)
_create_spec(config, kv_cache_config)
def test_duplicate_registration_raises():
@@ -253,7 +569,7 @@ def test_build_metric_definitions_below_threshold():
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
config = _make_vllm_config(store_threshold=1)
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
metrics = spec_cls.build_metric_definitions(
config.kv_transfer_config.kv_connector_extra_config
)
@@ -266,7 +582,7 @@ def test_build_metric_definitions_allocation_size_histogram():
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
config = _make_vllm_config(store_threshold=0)
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
metrics = spec_cls.build_metric_definitions(
config.kv_transfer_config.kv_connector_extra_config
)
@@ -291,7 +607,7 @@ def test_build_metric_definitions_returns_counter_at_threshold():
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
config = _make_vllm_config(store_threshold=2)
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
metrics = spec_cls.build_metric_definitions(
config.kv_transfer_config.kv_connector_extra_config
)
tests/v1/kv_offload/test_file_mapper.py
+106 -144
View File
@@ -4,80 +4,60 @@
from unittest.mock import MagicMock
import torch
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
MLAAttentionSpec,
SlidingWindowSpec,
)
from vllm.v1.kv_offload.base import (
OffloadingSpec,
make_offload_key,
from vllm.v1.kv_offload.base import OffloadingSpec, make_offload_key
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
OffloadingConfig,
OffloadingGroupConfig,
OffloadingModelConfig,
OffloadingParallelConfig,
)
from vllm.v1.kv_offload.file_mapper import FileMapper
# ---------------------------------------------------------------------------
# Shared mocks (mirrors test_fs_tier.py pattern)
# ---------------------------------------------------------------------------
_MOCK_VLLM_CONFIG = MagicMock()
_MOCK_VLLM_CONFIG.model_config.model = "test-model"
_MOCK_VLLM_CONFIG.cache_config.block_size = 16
_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32"
_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.rank = 0
_MOCK_KV_CACHE_CONFIG = MagicMock()
_MOCK_KV_CACHE_CONFIG.kv_cache_groups = []
_MOCK_OFFLOADING_SPEC = MagicMock(spec=OffloadingSpec)
_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG
_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
_MOCK_OFFLOADING_SPEC.block_size_factor = 1
# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------
def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
"""Helper to create FileMapper with customizable mock config."""
# Create a copy of the mock config to avoid modifying the global one
mock_vllm_config = MagicMock()
mock_vllm_config.model_config.model = kwargs.get("model_name", "test-model")
mock_vllm_config.cache_config.block_size = kwargs.get("hash_block_size", 16)
mock_vllm_config.cache_config.cache_dtype = (
f"torch.{kwargs.get('dtype', 'float16')}"
"""Build a FileMapper from a mocked spec carrying a hand-built config."""
config = OffloadingConfig(
groups=tuple(
OffloadingGroupConfig(
tokens_per_block=tokens_per_block,
layer_names=(layer_name,),
)
for tokens_per_block, layer_name in kwargs.get("groups", ())
),
worker_kv_bytes_per_block=0,
enable_kv_cache_events=False,
extra_config={},
engine_id="test-engine",
model=OffloadingModelConfig(
name=kwargs.get("model_name", "test-model"),
dtype=kwargs.get("dtype", "float16"),
),
cache=OffloadingCacheConfig(
tokens_per_hash=kwargs.get("tokens_per_hash", 16),
blocks_per_chunk=kwargs.get("blocks_per_chunk", 1),
),
parallel=OffloadingParallelConfig(
rank=kwargs.get("rank", 0),
world_size=kwargs.get("world_size", 1),
tp_size=kwargs.get("tp_size", 1),
pp_size=kwargs.get("pp_size", 1),
pcp_size=kwargs.get("pcp_size", 1),
dcp_size=kwargs.get("dcp_size", 1),
data_parallel_index=0,
is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False),
),
)
mock_vllm_config.parallel_config.tensor_parallel_size = kwargs.get("tp_size", 1)
mock_vllm_config.parallel_config.pipeline_parallel_size = kwargs.get("pp_size", 1)
mock_vllm_config.parallel_config.prefill_context_parallel_size = kwargs.get(
"pcp_size", 1
)
mock_vllm_config.parallel_config.decode_context_parallel_size = kwargs.get(
"dcp_size", 1
)
mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0)
mock_vllm_config.use_v2_model_runner = kwargs.get("use_v2_model_runner", False)
mock_kv_cache_config = MagicMock()
mock_kv_cache_config.kv_cache_groups = kwargs.get("kv_cache_groups", [])
mock_offloading_spec = MagicMock(spec=OffloadingSpec)
mock_offloading_spec.vllm_config = mock_vllm_config
mock_offloading_spec.kv_cache_config = mock_kv_cache_config
mock_offloading_spec.block_size_factor = kwargs.get("block_size_factor", 1)
spec = MagicMock(spec=OffloadingSpec)
spec.config = config
return FileMapper.from_offloading_spec(
root_dir=kwargs.get("root_dir", "/tmp/cache"),
offloading_spec=mock_offloading_spec,
gpu_blocks_per_file=mock_offloading_spec.block_size_factor,
offloading_spec=spec,
blocks_per_file=config.cache.blocks_per_chunk,
parallel_agnostic=kwargs.get("parallel_agnostic", False),
)
@@ -92,7 +72,7 @@ def test_get_file_name_full_structure():
Path must match: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin
Concretely:
- The segment immediately after base_path must end with `_r0`
- The segment immediately after base_path must end with `_r3`
- The next segment is the first 3 hex chars of the block hash
- The next segment is <2 hex chars>_g<group_idx>
- The final segment is <full hash hex>.bin
@@ -105,7 +85,7 @@ def test_get_file_name_full_structure():
path = fm.get_file_name(key)
expected_path = (
"/tmp/cache/test-model_588656ebcc66_r3/000/10_g2/0001020304050607.bin"
"/tmp/cache/test-model_42b94bdc9933_r3/000/10_g2/0001020304050607.bin"
)
assert path == expected_path
@@ -114,19 +94,30 @@ def test_get_run_config_fields():
fm = make_mapper_from_offloading_spec(
model_name="my-model",
dtype="bfloat16",
tp_size=2,
tp_size=4,
pp_size=3,
pcp_size=2,
dcp_size=2,
groups=((64, "layer0"),),
tokens_per_hash=64,
blocks_per_chunk=3,
)
cfg = fm.get_run_config()
assert cfg == {
"model_name": "my-model",
"hash_block_size": 16,
"gpu_blocks_per_file": 1,
"tp_size": 2,
"pp_size": 1,
"pcp_size": 1,
"dcp_size": 1,
"tokens_per_hash": 64,
"blocks_per_file": 3,
"tp_size": 4,
"pp_size": 3,
"pcp_size": 2,
"dcp_size": 2,
"dtype": "bfloat16",
"kv_cache_groups": [],
"kv_cache_groups": [
{
"tokens_per_block": 64,
"layer_names": ["layer0"],
}
],
"inference_engine": "vllm",
}
@@ -137,90 +128,61 @@ def test_get_config_file_path():
assert config_path == f"{fm.base_path}/config.json"
# ---------------------------------------------------------------------------
# parallel_agnostic: honored only for a single non-MLA full-attention group
# ---------------------------------------------------------------------------
def _full_attention_group() -> KVCacheGroupSpec:
return KVCacheGroupSpec(
layer_names=["layer0"],
kv_cache_spec=FullAttentionSpec(
block_size=16, num_kv_heads=4, head_size=128, dtype=torch.float32
),
)
def _sliding_window_group() -> KVCacheGroupSpec:
return KVCacheGroupSpec(
layer_names=["layer0"],
kv_cache_spec=SlidingWindowSpec(
block_size=16,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
sliding_window=128,
),
)
def test_parallel_agnostic_enabled_for_single_full_attention():
# tp/rank are collapsed out of the namespace so the cache is shared
# across tensor-parallel sizes.
def test_hybrid_file_identity_uses_resolved_tokens_per_hash():
# For heterogeneous groups the namespace records the resolved hash
# granularity (GCD of the group block sizes), which is the actual
# granularity of the offload block hashes.
fm = make_mapper_from_offloading_spec(
tp_size=2,
groups=((12, "full_layer"), (16, "mla_layer")),
tokens_per_hash=4,
)
assert fm.fields["tokens_per_hash"] == 4
assert fm.fields["kv_cache_groups"] == [
{"tokens_per_block": 12, "layer_names": ["full_layer"]},
{"tokens_per_block": 16, "layer_names": ["mla_layer"]},
]
# ---------------------------------------------------------------------------
# parallel_agnostic: opt-in honored only when the config marks the layout
# parallelism-agnostic (predicate computation is covered in test_factory.py)
# ---------------------------------------------------------------------------
def test_parallel_agnostic_collapses_namespace_when_config_allows():
fm = make_mapper_from_offloading_spec(
tp_size=4,
pp_size=3,
pcp_size=2,
dcp_size=2,
rank=1,
kv_cache_groups=[_full_attention_group()],
is_parallelism_agnostic=True,
parallel_agnostic=True,
)
assert fm.fields["tp_size"] == 1
assert fm.fields["pp_size"] == 1
assert fm.fields["pcp_size"] == 1
assert fm.fields["dcp_size"] == 1
assert fm.rank == 0
def test_parallel_agnostic_disabled_for_multiple_groups():
# More than one KV-cache group (hybrid model) => keep per-layout namespacing.
fm = make_mapper_from_offloading_spec(
tp_size=2,
kv_cache_groups=[_full_attention_group(), _full_attention_group()],
parallel_agnostic=True,
)
assert fm.fields["tp_size"] == 2
def test_parallel_agnostic_disabled_for_non_full_attention():
# Single group but not full attention (sliding window) => keep namespacing.
fm = make_mapper_from_offloading_spec(
tp_size=2,
kv_cache_groups=[_sliding_window_group()],
parallel_agnostic=True,
)
assert fm.fields["tp_size"] == 2
def test_parallel_agnostic_excludes_mla():
# MLA latent KV is replicated per rank, so its offloaded blocks are not
# parallelism-invariant: the opt-in must not collapse tp/rank.
group = KVCacheGroupSpec(
layer_names=["layer0"],
kv_cache_spec=MLAAttentionSpec(
block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
),
)
fm = make_mapper_from_offloading_spec(
tp_size=2, rank=1, kv_cache_groups=[group], parallel_agnostic=True
)
assert fm.fields["tp_size"] == 2
assert fm.rank == 1
def test_parallel_agnostic_disabled_on_v2_model_runner():
# V2's KV layout is not known to be parallelism-invariant: don't collapse.
def test_parallel_agnostic_ignored_when_config_disallows():
fm = make_mapper_from_offloading_spec(
tp_size=2,
rank=1,
kv_cache_groups=[_full_attention_group()],
use_v2_model_runner=True,
is_parallelism_agnostic=False,
parallel_agnostic=True,
)
assert fm.fields["tp_size"] == 2
assert fm.rank == 1
def test_namespace_kept_without_parallel_agnostic_opt_in():
fm = make_mapper_from_offloading_spec(
tp_size=2,
rank=1,
is_parallelism_agnostic=True,
parallel_agnostic=False,
)
assert fm.fields["tp_size"] == 2
assert fm.rank == 1
tests/v1/kv_offload/tiering/p2p/test_data_transport.py
+1 -1
View File
@@ -47,7 +47,7 @@ class TestDataTransportBase:
def test_config_fingerprint_deterministic(self):
"""Same config fields → same fingerprint."""
view = self._make_view()
fields = {"model": "llama", "dtype": "float16", "block_size_factor": 1}
fields = {"model": "llama", "dtype": "float16", "blocks_per_chunk": 1}
with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
t1 = NixlTransport("test:1", view, config_fields=fields)
t2 = NixlTransport("test:2", view, config_fields=fields)
tests/v1/kv_offload/tiering/p2p/test_manager.py
+3 -3
View File
@@ -1419,9 +1419,9 @@ class TestBindHostPortDefaults:
or SimpleNamespace(),
)
spec = SimpleNamespace(
block_size_factor=1,
vllm_config=SimpleNamespace(
parallel_config=SimpleNamespace(data_parallel_index=dp_index)
blocks_per_chunk=1,
config=SimpleNamespace(
parallel=SimpleNamespace(data_parallel_index=dp_index)
),
)
mgr = P2PSecondaryTierManager(spec, memoryview(b""), **kwargs)
tests/v1/kv_offload/tiering/test_fs_tier.py
+34 -22
View File
@@ -22,11 +22,18 @@ from vllm.distributed.kv_events import MEDIUM_FS
from vllm.v1.kv_offload.base import (
LookupResult,
OffloadingEvent,
OffloadingKVEventsConfig,
OffloadKey,
ReqContext,
ScheduleEndContext,
make_offload_key,
)
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
OffloadingConfig,
OffloadingModelConfig,
OffloadingParallelConfig,
)
from vllm.v1.kv_offload.tiering.base import JobMetadata
from vllm.v1.kv_offload.tiering.fs.manager import (
FileSystemTierManager,
@@ -41,35 +48,40 @@ _BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096.
_DTYPE: torch.dtype = torch.float32
_CTX = ReqContext(req_id="test")
_MOCK_VLLM_CONFIG = MagicMock()
_MOCK_VLLM_CONFIG.model_config.model = "test-model"
_MOCK_VLLM_CONFIG.cache_config.block_size = 16
_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32"
_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1
_MOCK_VLLM_CONFIG.parallel_config.rank = 0
_MOCK_KV_CACHE_CONFIG = MagicMock()
_MOCK_KV_CACHE_CONFIG.kv_cache_groups = []
_MOCK_OFFLOADING_SPEC = MagicMock()
_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG
_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
_MOCK_OFFLOADING_SPEC.block_size_factor = 1
def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock:
"""Mock spec with an explicit global KV events flag."""
spec = MagicMock()
spec.vllm_config = _MOCK_VLLM_CONFIG
spec.kv_cache_config = _MOCK_KV_CACHE_CONFIG
spec.block_size_factor = 1
spec.kv_events_config.enable_kv_cache_events = enable_kv_cache_events
spec.config = OffloadingConfig(
groups=(),
worker_kv_bytes_per_block=0,
enable_kv_cache_events=enable_kv_cache_events,
extra_config={},
engine_id="test-engine",
model=OffloadingModelConfig(name="test-model", dtype="float32"),
cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
parallel=OffloadingParallelConfig(
rank=0,
world_size=1,
tp_size=1,
pp_size=1,
pcp_size=1,
dcp_size=1,
data_parallel_index=0,
is_parallelism_agnostic=False,
),
)
spec.blocks_per_chunk = 1
spec.kv_events_config = OffloadingKVEventsConfig(
enable_kv_cache_events=enable_kv_cache_events,
self_describing_kv_events=False,
)
return spec
_MOCK_OFFLOADING_SPEC = _make_offloading_spec(enable_kv_cache_events=False)
def key(n: int) -> OffloadKey:
return make_offload_key(n.to_bytes(8, "big"), 0)
tests/v1/kv_offload/tiering/test_obj_tier.py
+30 -15
View File
@@ -19,11 +19,18 @@ import torch
from vllm.v1.kv_offload.base import (
LookupResult,
OffloadingKVEventsConfig,
OffloadKey,
ReqContext,
ScheduleEndContext,
make_offload_key,
)
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
OffloadingConfig,
OffloadingModelConfig,
OffloadingParallelConfig,
)
from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig
from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager
@@ -33,24 +40,30 @@ from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManag
# ---------------------------------------------------------------------------
def _make_vllm_config():
return SimpleNamespace(
model_config=SimpleNamespace(model="test/model"),
cache_config=SimpleNamespace(block_size=16, cache_dtype="float16"),
parallel_config=SimpleNamespace(
tensor_parallel_size=1,
pipeline_parallel_size=1,
prefill_context_parallel_size=1,
decode_context_parallel_size=1,
def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig:
return OffloadingConfig(
groups=(),
worker_kv_bytes_per_block=0,
enable_kv_cache_events=enable_kv_cache_events,
extra_config={},
engine_id="test-engine",
model=OffloadingModelConfig(name="test/model", dtype="float16"),
cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
parallel=OffloadingParallelConfig(
rank=0,
world_size=1,
tp_size=1,
pp_size=1,
pcp_size=1,
dcp_size=1,
data_parallel_index=0,
is_parallelism_agnostic=False,
),
use_v2_model_runner=False,
)
_OFFLOADING_SPEC = SimpleNamespace(
vllm_config=_make_vllm_config(),
kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
config=_make_offloading_config(enable_kv_cache_events=False),
)
_STORE_CONFIG = {
@@ -182,9 +195,11 @@ class MockNixlAgent:
def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace:
"""Offloading spec stub with an explicit global KV events flag."""
return SimpleNamespace(
vllm_config=_make_vllm_config(),
kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
kv_events_config=SimpleNamespace(enable_kv_cache_events=enable_kv_cache_events),
config=_make_offloading_config(enable_kv_cache_events),
kv_events_config=OffloadingKVEventsConfig(
enable_kv_cache_events=enable_kv_cache_events,
self_describing_kv_events=False,
),
)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+131
View File
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Translate vLLM KV cache metadata for native offloading backends."""
from typing import TYPE_CHECKING
from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
OffloadingConfig,
OffloadingGroupConfig,
OffloadingModelConfig,
OffloadingParallelConfig,
)
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor
def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool:
"""Return whether a KV cache tensor uses a packed block stride."""
return bool(kv_cache_tensor.block_stride)
def build_offloading_config(
vllm_config: "VllmConfig",
kv_cache_config: "KVCacheConfig",
) -> OffloadingConfig:
"""Translate vLLM configuration into the native offloading boundary."""
kv_transfer_config = vllm_config.kv_transfer_config
assert kv_transfer_config is not None
extra_config = kv_transfer_config.kv_connector_extra_config
assert kv_transfer_config.engine_id is not None
engine_id = kv_transfer_config.engine_id
parallel_config = vllm_config.parallel_config
context_parallel_factor = (
parallel_config.decode_context_parallel_size
* parallel_config.prefill_context_parallel_size
)
groups = tuple(
OffloadingGroupConfig(
tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor),
layer_names=tuple(group.layer_names),
)
for group in kv_cache_config.kv_cache_groups
)
_, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config)
for group in groups:
assert group.tokens_per_block % tokens_per_hash == 0, (
f"tokens_per_block={group.tokens_per_block} not divisible by "
f"tokens_per_hash={tokens_per_hash}. "
f"Hybrid models (e.g. Mamba+Attention) need "
f"--enable-prefix-caching to align block sizes."
)
blocks_per_chunk = 1
tokens_per_chunk = extra_config.get("block_size")
if tokens_per_chunk is not None:
tokens_per_chunk_int = int(tokens_per_chunk)
unique_tokens_per_block = {group.tokens_per_block for group in groups}
assert len(unique_tokens_per_block) == 1, (
"If 'block_size' is specified in kv_connector_extra_config, "
"there must be at least one KV cache group, "
"and all groups must have the same block size."
)
tokens_per_block = unique_tokens_per_block.pop()
assert tokens_per_chunk_int % tokens_per_block == 0
blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
worker_kv_bytes_per_block = 0
if kv_cache_config.num_blocks > 0:
packed_tensors = tuple(
is_kv_cache_tensor_packed(tensor)
for tensor in kv_cache_config.kv_cache_tensors
)
is_packed = any(packed_tensors)
assert not is_packed or all(packed_tensors)
total_gpu_kv_bytes = (
kv_cache_config.kv_cache_tensors[0].size
if is_packed
else sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors)
)
worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
# Only a single non-MLA full-attention group is parallelism-invariant:
# MLA latent KV is replicated per rank (never head-sharded), and the V2
# model runner's KV layout is not known to be parallelism-invariant.
single_group = (
kv_cache_config.kv_cache_groups[0].kv_cache_spec
if len(kv_cache_config.kv_cache_groups) == 1
else None
)
is_parallelism_agnostic = (
not vllm_config.use_v2_model_runner
and single_group is not None
and isinstance(single_group, FullAttentionSpec)
and not isinstance(single_group, MLAAttentionSpec)
)
kv_events_config = vllm_config.kv_events_config
return OffloadingConfig(
groups=groups,
worker_kv_bytes_per_block=worker_kv_bytes_per_block,
enable_kv_cache_events=(
kv_events_config is not None and kv_events_config.enable_kv_cache_events
),
extra_config=extra_config,
engine_id=engine_id,
model=OffloadingModelConfig(
name=vllm_config.model_config.model,
dtype=str(vllm_config.cache_config.cache_dtype).replace("torch.", ""),
),
cache=OffloadingCacheConfig(
tokens_per_hash=tokens_per_hash,
blocks_per_chunk=blocks_per_chunk,
),
parallel=OffloadingParallelConfig(
rank=parallel_config.rank,
world_size=parallel_config.world_size,
tp_size=parallel_config.tensor_parallel_size,
pp_size=parallel_config.pipeline_parallel_size,
pcp_size=parallel_config.prefill_context_parallel_size,
dcp_size=parallel_config.decode_context_parallel_size,
data_parallel_index=parallel_config.data_parallel_index,
is_parallelism_agnostic=is_parallelism_agnostic,
),
)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+12 -12
View File
@@ -101,7 +101,7 @@ class OffloadingEventsTracker:
self,
req: Request,
group_config: "GroupOffloadConfig",
offload_block_idx: int,
chunk_idx: int,
offload_key: OffloadKey,
) -> None:
"""Snapshot the KV cache event payload for one offloaded chunk.
@@ -111,9 +111,9 @@ class OffloadingEventsTracker:
"""
if not self.self_describing_enabled:
return
if group_config.sliding_window_size_in_blocks is not None:
if group_config.sliding_window_size_in_chunks is not None:
return
meta = self._build_event_metadata(req, group_config, offload_block_idx)
meta = self._build_event_metadata(req, group_config, chunk_idx)
self._pending_event_metadata[offload_key] = meta
def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
@@ -142,19 +142,19 @@ class OffloadingEventsTracker:
self,
req: Request,
group_config: "GroupOffloadConfig",
offload_block_idx: int,
chunk_idx: int,
) -> _OffloadEventMetadata:
"""Build the payload snapshot for one offloaded chunk: its
constituent per-block hashes, the whole chunk's tokens, and the
per-block ``block_size``."""
hbf = group_config.hash_block_size_factor
hbf = group_config.hashes_per_chunk
assert hbf > 0
assert offload_block_idx >= 0
assert chunk_idx >= 0
# per-block token count (= the GPU/hash block size)
sub_block_size = group_config.offloaded_block_size // hbf
tokens_per_hash = group_config.tokens_per_chunk // hbf
# chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash
# is the chunk's OffloadKey.
first_hash_idx = offload_block_idx * hbf
first_hash_idx = chunk_idx * hbf
last_hash_idx = first_hash_idx + hbf
assert first_hash_idx >= 0
assert last_hash_idx <= len(req.block_hashes)
@@ -164,7 +164,7 @@ class OffloadingEventsTracker:
chunk_hashes.append(block_hash)
assert len(chunk_hashes) == hbf
if group_config.sliding_window_size_in_blocks is not None:
if group_config.sliding_window_size_in_chunks is not None:
# record_store filters these out before calling this helper.
raise AssertionError("self-describing events only support full attention")
@@ -175,8 +175,8 @@ class OffloadingEventsTracker:
parent_block_hash = req.block_hashes[first_hash_idx - 1]
assert parent_block_hash is not None
tok_start = offload_block_idx * group_config.offloaded_block_size
tok_end = tok_start + group_config.offloaded_block_size
tok_start = chunk_idx * group_config.tokens_per_chunk
tok_end = tok_start + group_config.tokens_per_chunk
assert tok_end <= len(req.all_token_ids)
token_ids = tuple(req.all_token_ids[tok_start:tok_end])
@@ -190,7 +190,7 @@ class OffloadingEventsTracker:
block_hashes=tuple(chunk_hashes),
parent_block_hash=parent_block_hash,
token_ids=token_ids,
block_size=sub_block_size,
block_size=tokens_per_hash,
lora_id=lora_id,
lora_name=lora_name,
extra_keys=None,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
+1 -1
View File
@@ -321,10 +321,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
spec_cls = OffloadingSpecFactory.get_spec_cls(vllm_config)
kv_transfer_config = vllm_config.kv_transfer_config
assert kv_transfer_config is not None
extra_config = kv_transfer_config.kv_connector_extra_config
spec_cls = OffloadingSpecFactory.get_spec_cls(extra_config)
self._offloading_metric_metadata: dict[str, OffloadingMetricMetadata] = {
**spec_cls.build_metric_definitions(extra_config),
**get_connector_metric_definitions(),
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+217 -207
View File
@@ -6,6 +6,7 @@ from dataclasses import dataclass, field
from itertools import islice
from typing import Any, NamedTuple
from vllm.config import VllmConfig
from vllm.distributed.kv_events import KVCacheEvent
from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
@@ -31,6 +32,7 @@ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheSpec,
MambaSpec,
SlidingWindowSpec,
@@ -74,32 +76,32 @@ class TransferJobStatus:
class GroupOffloadConfig(NamedTuple):
group_idx: int
gpu_block_size: int
offloaded_block_size: int
hash_block_size_factor: int
tokens_per_block: int
tokens_per_chunk: int
hashes_per_chunk: int
# KV cache spec metadata propagated onto emitted BlockStored events so
# KV-aware consumers can classify and filter the group.
kv_event_group_spec: OffloadingEventGroupSpec
# None below means full attention
sliding_window_size_in_blocks: int | None
# Number of this group's offloaded blocks per full-attention alignment
# segment. Used to skip storing SWA blocks that can never serve a load
sliding_window_size_in_chunks: int | None
# Number of this group's offloaded chunks per full-attention alignment
# segment. Used to skip storing SWA chunks that can never serve a load
# hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes
# than the MLA full-attention group).
# None for full-attention groups or when the optimization doesn't apply.
alignment_block_count: int | None = None
# True for EAGLE/MTP draft-model attention groups. The trailing block
alignment_chunk_count: int | None = None
# True for EAGLE/MTP draft-model attention groups. The trailing chunk
# of these groups is volatile and lacks a stable hash, so it must
# be excluded from store and load scheduling.
is_eagle_group: bool = False
def get_sliding_window_size_in_blocks(
kv_cache_spec: KVCacheSpec, offloaded_block_size: int
def get_sliding_window_size_in_chunks(
kv_cache_spec: KVCacheSpec, tokens_per_chunk: int
) -> int | None:
if isinstance(kv_cache_spec, SlidingWindowSpec):
assert kv_cache_spec.sliding_window > 0
return cdiv(kv_cache_spec.sliding_window, offloaded_block_size)
return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)
if isinstance(kv_cache_spec, MambaSpec):
# Mamba depends on a single state
@@ -109,116 +111,121 @@ def get_sliding_window_size_in_blocks(
return None
def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None:
def resolve_mamba_align_size(
spec: "OffloadingSpec", kv_cache_config: KVCacheConfig
) -> int | None:
"""Scan all KV cache groups in *spec* and return the single mamba alignment
size, or None if no group requires mamba alignment.
For MambaSpec groups in "align" cache mode the hit window must be rounded
down to a multiple of the offloaded block size. Asserts that all such
down to a multiple of the offloaded chunk size. Asserts that all such
groups agree on the same value.
"""
mamba_align_size: int | None = None
for idx, gpu_block_size in enumerate(spec.gpu_block_size):
kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec
for idx, tokens_per_block in enumerate(spec.tokens_per_block):
kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align":
offload_block_size = gpu_block_size * spec.block_size_factor
assert mamba_align_size is None or mamba_align_size == offload_block_size
mamba_align_size = offload_block_size
tokens_per_chunk = tokens_per_block * spec.blocks_per_chunk
assert mamba_align_size is None or mamba_align_size == tokens_per_chunk
mamba_align_size = tokens_per_chunk
return mamba_align_size
class SchedulerOffloadConfig(NamedTuple):
kv_group_configs: tuple[GroupOffloadConfig, ...]
block_size_factor: int
blocks_per_chunk: int
num_workers: int
offload_prompt_only: bool
@classmethod
def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
def from_spec(
cls,
spec: OffloadingSpec,
vllm_config: VllmConfig,
kv_cache_config: KVCacheConfig,
) -> "SchedulerOffloadConfig":
# Determine the alignment token count from the full-attention group(s).
# This is the offloaded_block_size of the full-attention group; load
# This is the tokens_per_chunk of the full-attention group; load
# hits are always aligned to this boundary, so SWA blocks earlier in
# each segment can never serve a load hit. Relevant for hybrid
# architectures like DeepSeek V4 (MLA + SWA groups).
full_attn_offloaded_block_sizes: set[int] = set()
for idx, gpu_block_size in enumerate(spec.gpu_block_size):
kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec
sw = get_sliding_window_size_in_blocks(
kv_spec, gpu_block_size * spec.block_size_factor
full_attn_tokens_per_chunk: set[int] = set()
for idx, tokens_per_block in enumerate(spec.tokens_per_block):
kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
sw = get_sliding_window_size_in_chunks(
kv_spec, tokens_per_block * spec.blocks_per_chunk
)
if sw is None:
full_attn_offloaded_block_sizes.add(
gpu_block_size * spec.block_size_factor
)
full_attn_tokens_per_chunk.add(tokens_per_block * spec.blocks_per_chunk)
# Only apply the optimization if there's a single consistent
# full-attention alignment size.
alignment_tokens: int | None = None
if len(full_attn_offloaded_block_sizes) == 1:
alignment_tokens = full_attn_offloaded_block_sizes.pop()
if len(full_attn_tokens_per_chunk) == 1:
alignment_tokens = full_attn_tokens_per_chunk.pop()
def _alignment_block_count(
offloaded_block_size: int,
sliding_window_size_in_blocks: int | None,
def _alignment_chunk_count(
tokens_per_chunk: int,
sliding_window_size_in_chunks: int | None,
) -> int | None:
if alignment_tokens is None or sliding_window_size_in_blocks is None:
if alignment_tokens is None or sliding_window_size_in_chunks is None:
return None
if alignment_tokens <= offloaded_block_size:
if alignment_tokens <= tokens_per_chunk:
return None
per_segment = alignment_tokens // offloaded_block_size
if sliding_window_size_in_blocks >= per_segment:
per_segment = alignment_tokens // tokens_per_chunk
if sliding_window_size_in_chunks >= per_segment:
return None
return per_segment
eagle_groups = {
idx
for idx, g in enumerate(spec.kv_cache_config.kv_cache_groups)
for idx, g in enumerate(kv_cache_config.kv_cache_groups)
if g.is_eagle_group
}
use_eagle = (
spec.vllm_config.speculative_config is not None
and spec.vllm_config.speculative_config.use_eagle()
vllm_config.speculative_config is not None
and vllm_config.speculative_config.use_eagle()
)
if use_eagle and not eagle_groups:
eagle_groups = set(range(len(spec.kv_cache_config.kv_cache_groups)))
eagle_groups = set(range(len(kv_cache_config.kv_cache_groups)))
if eagle_groups:
logger.info(
"KV offloading: EAGLE/MTP draft attention groups %s "
"detected. The trailing block of these groups will be "
"detected. The trailing chunk of these groups will be "
"excluded from offloading due to volatility.",
sorted(eagle_groups),
)
return cls(
num_workers=spec.vllm_config.parallel_config.world_size,
num_workers=vllm_config.parallel_config.world_size,
kv_group_configs=tuple(
GroupOffloadConfig(
group_idx=idx,
gpu_block_size=gpu_block_size,
offloaded_block_size=gpu_block_size * spec.block_size_factor,
hash_block_size_factor=(
(gpu_block_size * spec.block_size_factor)
// spec.hash_block_size
tokens_per_block=tokens_per_block,
tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk,
hashes_per_chunk=(
(tokens_per_block * spec.blocks_per_chunk)
// spec.tokens_per_hash
),
sliding_window_size_in_blocks=(
sw := get_sliding_window_size_in_blocks(
spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec,
gpu_block_size * spec.block_size_factor,
sliding_window_size_in_chunks=(
sw := get_sliding_window_size_in_chunks(
kv_cache_config.kv_cache_groups[idx].kv_cache_spec,
tokens_per_block * spec.blocks_per_chunk,
)
),
alignment_block_count=_alignment_block_count(
gpu_block_size * spec.block_size_factor, sw
alignment_chunk_count=_alignment_chunk_count(
tokens_per_block * spec.blocks_per_chunk, sw
),
kv_event_group_spec=get_offloading_event_group_spec(
spec.kv_cache_config.kv_cache_groups[idx]
kv_cache_config.kv_cache_groups[idx]
),
is_eagle_group=idx in eagle_groups,
)
for idx, gpu_block_size in enumerate(spec.gpu_block_size)
for idx, tokens_per_block in enumerate(spec.tokens_per_block)
),
block_size_factor=spec.block_size_factor,
blocks_per_chunk=spec.blocks_per_chunk,
offload_prompt_only=spec.offload_prompt_only,
)
@@ -227,11 +234,11 @@ class SchedulerOffloadConfig(NamedTuple):
class RequestGroupState:
offload_keys: list[OffloadKey] = field(default_factory=list)
block_ids: list[int] = field(default_factory=list)
# index of next block (of size offloaded_block_size) to offload
next_stored_block_idx: int = 0
# number of offloaded blocks hit (including GPU prefix cache)
# Index of the next chunk to offload.
next_stored_chunk_idx: int = 0
# Number of offloaded chunks hit (including GPU prefix cache)
# when the request first started
num_hit_blocks: int = 0
num_hit_chunks: int = 0
@dataclass(slots=True)
@@ -278,11 +285,11 @@ class RequestOffloadState:
):
for req_block_hash in islice(
self.req.block_hashes,
group_config.hash_block_size_factor * len(group_state.offload_keys)
+ group_config.hash_block_size_factor
group_config.hashes_per_chunk * len(group_state.offload_keys)
+ group_config.hashes_per_chunk
- 1,
None,
group_config.hash_block_size_factor,
group_config.hashes_per_chunk,
):
group_state.offload_keys.append(
make_offload_key(req_block_hash, group_config.group_idx)
@@ -298,46 +305,46 @@ class RequestOffloadState:
for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
group_state.block_ids.extend(new_blocks)
def storable_blocks(
def storable_chunks(
self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
) -> int:
"""Number of leading offloaded blocks eligible for store.
"""Number of leading offloaded chunks eligible for store.
For eagle/MTP groups the volatile trailing block of the offloadable
For eagle/MTP groups the volatile trailing chunk of the offloadable
range is excluded while decoding: the draft-layer KV of the last
accepted position may be rewritten after spec-token rejection. During
prefill the trailing block is stable (the draft input for a chunk's
prefill the trailing chunk is stable (the draft input for a chunk's
last position is the next prompt token), so it is stored immediately.
The exclusion must be applied consistently everywhere
``next_stored_block_idx`` is derived: otherwise the trailing block of
``next_stored_chunk_idx`` is derived: otherwise the trailing chunk of
each step is skipped on collection but jumped over by
``next_stored_block_idx``, so it is never re-considered and a
``next_stored_chunk_idx``, so it is never re-considered and a
permanent hole breaks prefix-reuse lookup.
"""
num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
num_chunks = num_offloadable_tokens // group_config.tokens_per_chunk
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
if group_config.is_eagle_group and is_decoding:
num_blocks = max(0, num_blocks - 1)
return num_blocks
num_chunks = max(0, num_chunks - 1)
return num_chunks
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
# max(): at the prefill->decode transition of a block-aligned prompt,
# storable_blocks drops by one (the eagle exclusion kicks in), and the
# index must not move backwards past already-stored blocks.
# max(): at the prefill->decode transition of a chunk-aligned prompt,
# storable_chunks drops by one (the eagle exclusion kicks in), and the
# index must not move backwards past already-stored chunks.
for group_config, group_state in zip(
self.config.kv_group_configs, self.group_states
):
group_state.next_stored_block_idx = max(
group_state.next_stored_block_idx,
self.storable_blocks(group_config, num_offloadable_tokens),
group_state.next_stored_chunk_idx = max(
group_state.next_stored_chunk_idx,
self.storable_chunks(group_config, num_offloadable_tokens),
)
def update_num_hit_blocks(self, num_cached_tokens: int) -> None:
def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
for group_config, group_state in zip(
self.config.kv_group_configs, self.group_states
):
group_state.num_hit_blocks = (
num_cached_tokens // group_config.offloaded_block_size
group_state.num_hit_chunks = (
num_cached_tokens // group_config.tokens_per_chunk
)
@@ -354,22 +361,26 @@ class OffloadingConnectorScheduler:
def __init__(
self,
spec: OffloadingSpec,
vllm_config: VllmConfig,
kv_cache_config: KVCacheConfig,
):
self.config = SchedulerOffloadConfig.from_spec(spec)
self.config = SchedulerOffloadConfig.from_spec(
spec, vllm_config, kv_cache_config
)
self.manager: OffloadingManager = spec.get_manager()
self._connector_stats = OffloadingConnectorStats()
full_attention_groups: list[int] = []
sliding_window_groups: list[int] = []
for group_config in self.config.kv_group_configs:
if group_config.sliding_window_size_in_blocks is None:
if group_config.sliding_window_size_in_chunks is None:
full_attention_groups.append(group_config.group_idx)
else:
sliding_window_groups.append(group_config.group_idx)
# sort sliding window groups by window size in decreasing order
def _sliding_window_sort_key(i: int) -> int:
val = self.config.kv_group_configs[i].sliding_window_size_in_blocks
val = self.config.kv_group_configs[i].sliding_window_size_in_chunks
assert val is not None
return val
@@ -378,7 +389,9 @@ class OffloadingConnectorScheduler:
# used by _lookup
self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups)
self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups
self._mamba_align_size: int | None = resolve_mamba_align_size(spec)
self._mamba_align_size: int | None = resolve_mamba_align_size(
spec, kv_cache_config
)
self._req_status: dict[ReqId, RequestOffloadState] = {}
self._current_batch_load_jobs: dict[int, TransferJob] = {}
@@ -386,9 +399,9 @@ class OffloadingConnectorScheduler:
# GPU block IDs allocated in the current engine step
self._current_batch_allocated_block_ids: set[int] = set()
# if GPU prefix caching is enabled,
# track loaded blocks to avoid redundant loads
self._blocks_being_loaded: set[OffloadKey] | None = (
set() if spec.vllm_config.cache_config.enable_prefix_caching else None
# Track loaded chunks to avoid redundant loads.
self._chunks_being_loaded: set[OffloadKey] | None = (
set() if vllm_config.cache_config.enable_prefix_caching else None
)
# Job ID counter shared by loads and stores.
@@ -434,7 +447,7 @@ class OffloadingConnectorScheduler:
def _maximal_prefix_lookup(
self, keys: Iterable[OffloadKey], req_context: ReqContext
) -> int | None:
"""Return the number of consecutive offloaded blocks from the start,
"""Return the number of consecutive offloaded chunks from the start,
or None if the backend deferred a lookup."""
hit_count = 0
defer_lookup = False
@@ -490,18 +503,18 @@ class OffloadingConnectorScheduler:
for group_config, group_state in zip(
self.config.kv_group_configs, req_status.group_states
):
if group_config.sliding_window_size_in_blocks is None:
if group_config.sliding_window_size_in_chunks is None:
self.manager.touch(group_state.offload_keys, req_status.req_context)
else:
# we aim to keep just blocks that are necessary to hit
# the original request (+ decoded blocks)
blocks_to_skip = max(
# Keep only chunks needed to hit the original request, plus
# decoded chunks.
chunks_to_skip = max(
0,
group_state.num_hit_blocks
- group_config.sliding_window_size_in_blocks,
group_state.num_hit_chunks
- group_config.sliding_window_size_in_chunks,
)
self.manager.touch(
group_state.offload_keys[blocks_to_skip:],
group_state.offload_keys[chunks_to_skip:],
req_status.req_context,
)
@@ -531,7 +544,7 @@ class OffloadingConnectorScheduler:
defer_lookup = False
lookup_groups = self._lookup_groups
# Tracks which eagle groups have already popped their volatile trailing block
# Tracks which eagle groups have already popped their volatile trailing chunk
# in the current convergence iteration. Reset when a non-eagle group
# tightens the hit boundary, requiring a fresh pop.
eagle_verified: set[int] = set()
@@ -544,79 +557,76 @@ class OffloadingConnectorScheduler:
group_idx
]
group_state: RequestGroupState = req_status.group_states[group_idx]
offloaded_block_size = group_config.offloaded_block_size
tokens_per_chunk = group_config.tokens_per_chunk
offload_keys = group_state.offload_keys
assert (
len(offload_keys)
>= req_status.req.num_tokens // offloaded_block_size
len(offload_keys) >= req_status.req.num_tokens // tokens_per_chunk
)
is_eagle_unverified = (
group_config.is_eagle_group and group_idx not in eagle_verified
)
# Constrain to block-aligned boundary for this group
# Constrain to a chunk-aligned boundary for this group.
max_hit_size_tokens = min(
max_hit_size_tokens, len(offload_keys) * offloaded_block_size
max_hit_size_tokens, len(offload_keys) * tokens_per_chunk
)
if max_hit_size_tokens - num_computed_tokens < offloaded_block_size:
# we can only load less than a block, better skip
if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk:
# We can only load less than a chunk, so skip.
return 0
sliding_window_size_in_blocks = (
group_config.sliding_window_size_in_blocks
sliding_window_size_in_chunks = (
group_config.sliding_window_size_in_chunks
)
# For eagle groups, query one extra block that will be popped.
# For eagle groups, query one extra chunk that will be popped.
# We only need to increase the query size for sliding window groups.
query_max = max_hit_size_tokens
if is_eagle_unverified and sliding_window_size_in_blocks is not None:
if is_eagle_unverified and sliding_window_size_in_chunks is not None:
query_max = min(
max_hit_size_tokens + offloaded_block_size,
len(offload_keys) * offloaded_block_size,
max_hit_size_tokens + tokens_per_chunk,
len(offload_keys) * tokens_per_chunk,
)
num_blocks = min(
cdiv(query_max, offloaded_block_size), len(offload_keys)
)
start_block_idx = num_computed_tokens // offloaded_block_size
offload_keys = offload_keys[start_block_idx:num_blocks]
num_chunks = min(cdiv(query_max, tokens_per_chunk), len(offload_keys))
start_chunk_idx = num_computed_tokens // tokens_per_chunk
offload_keys = offload_keys[start_chunk_idx:num_chunks]
# end index (in the sliced offload_keys) up to which we
# have backend-confirmed hits
num_hit_blocks: int | None
if sliding_window_size_in_blocks is None:
num_hit_blocks = self._maximal_prefix_lookup(
num_hit_chunks: int | None
if sliding_window_size_in_chunks is None:
num_hit_chunks = self._maximal_prefix_lookup(
offload_keys, req_status.req_context
)
else:
required_window = sliding_window_size_in_blocks
required_window = sliding_window_size_in_chunks
if is_eagle_unverified:
required_window += 1
num_hit_blocks = self._sliding_window_lookup(
num_hit_chunks = self._sliding_window_lookup(
offload_keys,
required_window,
req_status.req_context,
)
if num_hit_blocks == 0:
if num_hit_chunks == 0:
return 0
if num_hit_blocks is None:
if num_hit_chunks is None:
defer_lookup = True
else:
if is_eagle_unverified:
num_hit_blocks -= 1
num_hit_chunks -= 1
eagle_verified.add(group_idx)
max_hit_size_tokens = min(
max_hit_size_tokens,
offloaded_block_size * (start_block_idx + num_hit_blocks),
tokens_per_chunk * (start_chunk_idx + num_hit_chunks),
)
new_num_hit_tokens = max_hit_size_tokens - num_computed_tokens
if new_num_hit_tokens < offloaded_block_size:
# we can only load less than a block, better skip
if new_num_hit_tokens < tokens_per_chunk:
# We can only load less than a chunk, so skip.
return 0
if new_num_hit_tokens < num_hit_tokens:
@@ -632,7 +642,7 @@ class OffloadingConnectorScheduler:
# sliding window works with the new_num_hit_tokens
lookup_groups = self._sliding_window_groups
looked_up_sliding_window |= sliding_window_size_in_blocks is not None
looked_up_sliding_window |= sliding_window_size_in_chunks is not None
num_hit_tokens = new_num_hit_tokens
if defer_lookup:
@@ -642,28 +652,28 @@ class OffloadingConnectorScheduler:
)
return None
# possibly delay request if any of the hit blocks is already being loaded
if self._blocks_being_loaded:
# Possibly delay the request if any hit chunk is already being loaded.
if self._chunks_being_loaded:
for group_config, group_state in zip(
self.config.kv_group_configs, req_status.group_states
):
offloaded_block_size = group_config.offloaded_block_size
sliding_window_size_in_blocks = (
group_config.sliding_window_size_in_blocks
tokens_per_chunk = group_config.tokens_per_chunk
sliding_window_size_in_chunks = (
group_config.sliding_window_size_in_chunks
)
offload_keys = group_state.offload_keys
num_blocks = cdiv(
num_computed_tokens + num_hit_tokens, offloaded_block_size
num_chunks = cdiv(
num_computed_tokens + num_hit_tokens, tokens_per_chunk
)
start_block_idx = num_computed_tokens // offloaded_block_size
offload_keys = offload_keys[start_block_idx:num_blocks]
if sliding_window_size_in_blocks is not None:
offload_keys = offload_keys[-sliding_window_size_in_blocks:]
if any(key in self._blocks_being_loaded for key in offload_keys):
# hit blocks are being loaded, delay request
start_chunk_idx = num_computed_tokens // tokens_per_chunk
offload_keys = offload_keys[start_chunk_idx:num_chunks]
if sliding_window_size_in_chunks is not None:
offload_keys = offload_keys[-sliding_window_size_in_chunks:]
if any(key in self._chunks_being_loaded for key in offload_keys):
# Hit chunks are being loaded, so delay the request.
logger.debug(
"Delaying request %s since some of its"
" blocks are already being loaded",
" chunks are already being loaded",
req_status.req.request_id,
)
return None
@@ -740,7 +750,7 @@ class OffloadingConnectorScheduler:
req_status.deferred_lookup_start_time = lookup_start
else:
self._maybe_observe_lookup_async_delay(req_status)
req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))
req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0))
self._touch(req_status)
@@ -771,10 +781,10 @@ class OffloadingConnectorScheduler:
block.block_id for block in group_blocks if block.block_id != 0
)
gpu_block_size = group_config.gpu_block_size
offloaded_block_size = group_config.offloaded_block_size
tokens_per_block = group_config.tokens_per_block
tokens_per_chunk = group_config.tokens_per_chunk
offload_keys = group_state.offload_keys
num_gpu_blocks = cdiv(num_cached_tokens, gpu_block_size)
num_gpu_blocks = cdiv(num_cached_tokens, tokens_per_block)
assert len(group_blocks) >= num_gpu_blocks
num_locally_computed_gpu_blocks = num_gpu_blocks
@@ -786,24 +796,24 @@ class OffloadingConnectorScheduler:
assert (
num_locally_computed_tokens
<= num_locally_computed_gpu_blocks * gpu_block_size
<= num_locally_computed_gpu_blocks * tokens_per_block
)
num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks
if group_config.sliding_window_size_in_blocks is not None:
if group_config.sliding_window_size_in_chunks is not None:
assert (
num_pending_gpu_blocks
<= group_config.sliding_window_size_in_blocks
* self.config.block_size_factor
<= group_config.sliding_window_size_in_chunks
* self.config.blocks_per_chunk
)
num_blocks = cdiv(num_cached_tokens, offloaded_block_size)
assert len(offload_keys) >= num_blocks
num_chunks = cdiv(num_cached_tokens, tokens_per_chunk)
assert len(offload_keys) >= num_chunks
if num_pending_gpu_blocks:
start_block_idx = (
num_locally_computed_gpu_blocks // self.config.block_size_factor
start_chunk_idx = (
num_locally_computed_gpu_blocks // self.config.blocks_per_chunk
)
keys_to_load.extend(offload_keys[start_block_idx:num_blocks])
keys_to_load.extend(offload_keys[start_chunk_idx:num_chunks])
dst_block_ids.extend(
block.block_id
@@ -814,11 +824,11 @@ class OffloadingConnectorScheduler:
group_sizes.append(num_pending_gpu_blocks)
block_indices.append(num_locally_computed_gpu_blocks)
# Skip prefix-hit blocks for block-level policy; for
# request-level, next_stored_block_idx stays at 0 so all
# blocks (including hits) are offloaded.
# Skip prefix-hit chunks for block-level policy; for
# request-level, next_stored_chunk_idx stays at 0 so all
# chunks (including hits) are offloaded.
if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:
group_state.next_stored_block_idx = num_blocks
group_state.next_stored_chunk_idx = num_chunks
src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
dst_spec = GPULoadStoreSpec(
@@ -841,8 +851,8 @@ class OffloadingConnectorScheduler:
is_store=False,
)
if self._blocks_being_loaded is not None:
self._blocks_being_loaded.update(keys_to_load)
if self._chunks_being_loaded is not None:
self._chunks_being_loaded.update(keys_to_load)
def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:
"""
@@ -877,16 +887,16 @@ class OffloadingConnectorScheduler:
# Zero out stale block_ids in sliding window groups' pending-store
# positions. Only sliding window groups can have stale entries (blocks
# freed by remove_skipped_blocks then reallocated). Only positions in
# [next_stored_block_idx * bsf, end) need checking where end is the
# [next_stored_chunk_idx * bsf, end) need checking where end is the
# pre-extend length: earlier positions were already offloaded, later
# ones are fresh allocations from this step.
if self._sliding_window_groups and self._current_batch_allocated_block_ids:
block_size_factor = self.config.block_size_factor
blocks_per_chunk = self.config.blocks_per_chunk
for req_id, req_status in self._req_status.items():
ends = new_block_ids_end.get(req_id)
for i, grp_idx in enumerate(self._sliding_window_groups):
group_state = req_status.group_states[grp_idx]
start = group_state.next_stored_block_idx * block_size_factor
start = group_state.next_stored_chunk_idx * blocks_per_chunk
end = ends[i] if ends is not None else len(group_state.block_ids)
for j in range(start, end):
if (
@@ -899,7 +909,7 @@ class OffloadingConnectorScheduler:
self,
scheduler_output: SchedulerOutput,
) -> dict[int, TransferJob]:
block_size_factor = self.config.block_size_factor
blocks_per_chunk = self.config.blocks_per_chunk
store_jobs: dict[int, TransferJob] = {}
for req_id in scheduler_output.num_scheduled_tokens:
req_status = self._req_status.get(req_id)
@@ -915,59 +925,59 @@ class OffloadingConnectorScheduler:
if max_offload_tokens is not None:
num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
# Skip decode-phase blocks: clamp to the prompt length so only
# prefill (prompt) blocks become eligible for store. next_stored_idx
# never advances past this boundary, so decode blocks are never
# Skip decode-phase chunks: clamp to the prompt length so only
# prefill chunks become eligible for store. next_stored_chunk_idx
# never advances past this boundary, so decode chunks are never
# queued in this or any later step.
if self.config.offload_prompt_only:
num_offloadable_tokens = min(
num_offloadable_tokens, req.num_prompt_tokens
)
# Filter out blocks skipped due to sliding window attention / SSM
# Filter out chunks skipped due to sliding window attention / SSM
# or unreachable by the load path's alignment constraints.
new_offload_keys: list[OffloadKey] = []
for group_config, group_state in zip(
self.config.kv_group_configs, req_status.group_states
):
num_blocks = req_status.storable_blocks(
num_chunks = req_status.storable_chunks(
group_config, num_offloadable_tokens
)
start_block_idx = group_state.next_stored_block_idx
if num_blocks <= start_block_idx:
start_chunk_idx = group_state.next_stored_chunk_idx
if num_chunks <= start_chunk_idx:
continue
offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
# For each block to offload, take the last corresponding GPU block.
# e.g. if block size factor is 3 and GPU block IDs are
# 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8.
offload_keys = group_state.offload_keys[start_chunk_idx:num_chunks]
# For each chunk, take the last corresponding GPU block. For
# blocks_per_chunk=3 and GPU block IDs 1 5 6 7 2 4 9 3 8,
# this selects GPU blocks 6 4 8.
# A block_id of 0 means either a sliding window / SSM skip
# or a stale entry that was zeroed out — skip it either way.
offload_block_ids = group_state.block_ids[
start_block_idx * block_size_factor
+ block_size_factor
- 1 : num_blocks * block_size_factor : block_size_factor
start_chunk_idx * blocks_per_chunk
+ blocks_per_chunk
- 1 : num_chunks * blocks_per_chunk : blocks_per_chunk
]
assert len(offload_keys) == len(offload_block_ids)
alignment_block_count = group_config.alignment_block_count
tail = group_config.sliding_window_size_in_blocks
alignment_chunk_count = group_config.alignment_chunk_count
tail = group_config.sliding_window_size_in_chunks
for key_idx, (offload_key, block_id) in enumerate(
zip(offload_keys, offload_block_ids)
):
if block_id == 0:
continue
# Skip SWA blocks that can never serve a load hit:
# Skip SWA chunks that can never serve a load hit:
# within each full-attention alignment segment, only the
# trailing `tail` blocks are reachable by
# trailing `tail` chunks are reachable by
# _sliding_window_lookup. For DeepSeek V4 with 100K
# tokens this reduces SWA stores by ~78%.
if alignment_block_count is not None:
if alignment_chunk_count is not None:
assert tail is not None
abs_block_idx = start_block_idx + key_idx
pos_in_segment = abs_block_idx % alignment_block_count
if pos_in_segment < alignment_block_count - tail:
abs_chunk_idx = start_chunk_idx + key_idx
pos_in_segment = abs_chunk_idx % alignment_chunk_count
if pos_in_segment < alignment_chunk_count - tail:
continue
new_offload_keys.append(offload_key)
@@ -982,7 +992,7 @@ class OffloadingConnectorScheduler:
self._connector_stats.increase_counter(
_ConnectorMetricName.ALLOCATION_FAILURE
)
logger.warning("Request %s: cannot store blocks", req_id)
logger.warning("Request %s: cannot store chunks", req_id)
continue
if not store_output.keys_to_store:
@@ -1002,29 +1012,29 @@ class OffloadingConnectorScheduler:
self.config.kv_group_configs, req_status.group_states
):
is_sliding_window = (
group_config.sliding_window_size_in_blocks is not None
group_config.sliding_window_size_in_chunks is not None
)
num_blocks = req_status.storable_blocks(
num_chunks = req_status.storable_chunks(
group_config, num_offloadable_tokens
)
start_block_idx = group_state.next_stored_block_idx
start_chunk_idx = group_state.next_stored_chunk_idx
block_ids = group_state.block_ids
num_group_blocks = 0
start_gpu_block_idx: int | None = None
for idx, offload_key in enumerate(
group_state.offload_keys[start_block_idx:num_blocks]
group_state.offload_keys[start_chunk_idx:num_chunks]
):
if offload_key not in keys_to_store:
continue
offloaded_block_idx = start_block_idx + idx
chunk_idx = start_chunk_idx + idx
self._events_tracker.record_store(
req, group_config, offloaded_block_idx, offload_key
req, group_config, chunk_idx, offload_key
)
gpu_block_idx = offloaded_block_idx * block_size_factor
for i in range(block_size_factor):
gpu_block_idx = chunk_idx * blocks_per_chunk
for i in range(blocks_per_chunk):
block_id = block_ids[gpu_block_idx + i]
if block_id == 0:
continue
@@ -1039,8 +1049,8 @@ class OffloadingConnectorScheduler:
group_sizes.append(num_group_blocks)
block_indices.append(start_gpu_block_idx or 0)
group_state.next_stored_block_idx = max(
group_state.next_stored_block_idx, num_blocks
group_state.next_stored_chunk_idx = max(
group_state.next_stored_chunk_idx, num_chunks
)
src_spec = GPULoadStoreSpec(
@@ -1076,7 +1086,7 @@ class OffloadingConnectorScheduler:
)
logger.debug(
"Request %s offloading %s blocks upto %d tokens (job %d)",
"Request %s offloading %s chunks upto %d tokens (job %d)",
req_id,
len(keys_to_store),
num_offloadable_tokens,
@@ -1198,8 +1208,8 @@ class OffloadingConnectorScheduler:
self.manager.complete_store(job_status.keys, req_status.req_context)
else:
self.manager.complete_load(job_status.keys, req_status.req_context)
if self._blocks_being_loaded:
self._blocks_being_loaded.difference_update(job_status.keys)
if self._chunks_being_loaded:
self._chunks_being_loaded.difference_update(job_status.keys)
if self._block_id_to_pending_jobs:
# Sliding window blocks are tracked from store creation
# and must be cleaned up unconditionally.
@@ -1289,7 +1299,7 @@ class OffloadingConnectorScheduler:
yield from self._events_tracker.take_events(self.manager.take_events())
def reset_cache(self) -> None:
"""Reset the offloading manager cache, evicting all stored blocks."""
"""Reset the offloading manager cache, evicting all stored chunks."""
# reset_cache cannot be called in the middle of a schedule step
assert not self._current_batch_load_jobs
@@ -1306,10 +1316,10 @@ class OffloadingConnectorScheduler:
# Reset offloading manager cache
self.manager.reset_cache()
# Reset store progress so active requests re-offload from block 0
# Reset store progress so active requests re-offload from chunk 0.
for status in self._req_status.values():
for group_state in status.group_states:
group_state.next_stored_block_idx = 0
group_state.next_stored_chunk_idx = 0
status.transfer_jobs.clear()
# Discard jobs and save job_counter to be able to discard worker responses
@@ -1323,8 +1333,8 @@ class OffloadingConnectorScheduler:
# Note: _current_batch_jobs_to_flush is intentionally NOT cleared.
# The load flush IDs collected above must be delivered to workers.
if self._blocks_being_loaded is not None:
self._blocks_being_loaded.clear()
if self._chunks_being_loaded is not None:
self._chunks_being_loaded.clear()
def shutdown(self) -> None:
self.manager.shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+14 -5
View File
@@ -10,10 +10,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingWorkerMetadata,
ReqId,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
is_kv_cache_tensor_packed,
)
from vllm.logger import init_logger
from vllm.v1.attention.backend import AttentionBackend
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheConfig,
MambaSpec,
UniformTypeKVCacheSpecs,
)
@@ -33,8 +37,13 @@ logger = init_logger(__name__)
class OffloadingConnectorWorker:
"""Implementation of Worker side methods"""
def __init__(self, spec: OffloadingSpec):
def __init__(
self,
spec: OffloadingSpec,
kv_cache_config: KVCacheConfig,
):
self.spec = spec
self.kv_cache_config = kv_cache_config
self.worker: OffloadingWorker | None = None
# job_id -> req_id for in-flight loads.
@@ -50,7 +59,7 @@ class OffloadingConnectorWorker:
def register_kv_caches(
self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
):
kv_cache_config = self.spec.kv_cache_config
kv_cache_config = self.kv_cache_config
num_blocks = kv_cache_config.num_blocks
# Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
@@ -58,7 +67,7 @@ class OffloadingConnectorWorker:
# General (non-packed) layouts size the tensor at page_size_bytes per
# manager block, so page_size_bytes is the correct offloading stride.
layer_is_packed: dict[str, bool] = {
ln: bool(kv_tensor.block_stride)
ln: is_kv_cache_tensor_packed(kv_tensor)
for kv_tensor in kv_cache_config.kv_cache_tensors
for ln in kv_tensor.shared_by
}
@@ -142,7 +151,7 @@ class OffloadingConnectorWorker:
(
t
for t in kv_cache_config.kv_cache_tensors
if t.block_stride and t.shared_by
if is_kv_cache_tensor_packed(t) and t.shared_by
),
None,
)
@@ -237,7 +246,7 @@ class OffloadingConnectorWorker:
num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim)
assert num_blocks_physical_dim == 0
kv_cache_groups = self.spec.kv_cache_config.kv_cache_groups
kv_cache_groups = self.kv_cache_config.kv_cache_groups
assert len(kv_cache_groups) == 1
kv_cache_spec = kv_cache_groups[0].kv_cache_spec
num_layers = len(kv_cache_groups[0].layer_names)
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+9 -3
View File
@@ -23,6 +23,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingConnectorMetadata,
OffloadingWorkerMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
OffloadPromMetrics,
@@ -56,14 +59,17 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
):
super().__init__(vllm_config, role, kv_cache_config)
spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config)
offloading_config = build_offloading_config(vllm_config, kv_cache_config)
spec = OffloadingSpecFactory.create_spec(offloading_config)
self.connector_scheduler: OffloadingConnectorScheduler | None = None
self.connector_worker: OffloadingConnectorWorker | None = None
if role == KVConnectorRole.SCHEDULER:
self.connector_scheduler = OffloadingConnectorScheduler(spec)
self.connector_scheduler = OffloadingConnectorScheduler(
spec, vllm_config, kv_cache_config
)
elif role == KVConnectorRole.WORKER:
self.connector_worker = OffloadingConnectorWorker(spec)
self.connector_worker = OffloadingConnectorWorker(spec, kv_cache_config)
def shutdown(self) -> None:
if self.connector_worker is not None:
vllm/v1/kv_offload/base.py
+9 -56
View File
@@ -14,14 +14,13 @@ import numpy as np
import torch
from vllm.logger import init_logger
from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
)
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_offload.config import OffloadingConfig
# `OffloadKey` identifies an offloaded block. It combines a block hash with
# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
@@ -482,22 +481,15 @@ class OffloadingSpec(ABC):
"""Return Prometheus metric definitions emitted by this spec."""
return {}
def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
def __init__(self, config: OffloadingConfig):
logger.warning(
"Initializing OffloadingSpec. This API is experimental and "
"subject to change in the future as we iterate the design."
)
self.vllm_config = vllm_config
self.kv_cache_config = kv_cache_config
kv_transfer_config = vllm_config.kv_transfer_config
assert kv_transfer_config is not None
self.extra_config = kv_transfer_config.kv_connector_extra_config
kv_events_config = vllm_config.kv_events_config
self.config = config
self.extra_config = config.extra_config
self.kv_events_config = OffloadingKVEventsConfig(
enable_kv_cache_events=(
kv_events_config is not None and kv_events_config.enable_kv_cache_events
),
enable_kv_cache_events=config.enable_kv_cache_events,
self_describing_kv_events=bool(
self.extra_config.get("self_describing_kv_events", False)
),
@@ -511,48 +503,9 @@ class OffloadingSpec(ABC):
self.extra_config.get("offload_prompt_only", True)
)
parallel_config = vllm_config.parallel_config
context_parallel_factor = (
parallel_config.decode_context_parallel_size
* parallel_config.prefill_context_parallel_size
)
# gpu block size per group
self.gpu_block_size: tuple[int, ...] = tuple(
kv_cache_group.kv_cache_spec.block_size * context_parallel_factor
for kv_cache_group in kv_cache_config.kv_cache_groups
)
# hash_block_size must match what the scheduler uses for
# Request.block_hashes (resolved via resolve_kv_cache_block_sizes).
_, self.hash_block_size = resolve_kv_cache_block_sizes(
kv_cache_config, vllm_config
)
for block_size in self.gpu_block_size:
assert block_size % self.hash_block_size == 0, (
f"gpu_block_size={block_size} not divisible by "
f"hash_block_size={self.hash_block_size}. "
f"Hybrid models (e.g. Mamba+Attention) need "
f"--enable-prefix-caching to align block sizes."
)
# offloaded_block_size / gpu_block_size
self.block_size_factor: int = 1
offloaded_block_size = self.extra_config.get("block_size")
if offloaded_block_size is not None:
offloaded_block_size_int = int(offloaded_block_size)
gpu_block_sizes = set(self.gpu_block_size)
assert len(gpu_block_sizes) == 1, (
"If 'block_size' is specified in kv_connector_extra_config, "
"there must be at least one KV cache group, "
"and all groups must have the same block size."
)
gpu_block_size = gpu_block_sizes.pop()
assert offloaded_block_size_int % gpu_block_size == 0
self.block_size_factor = offloaded_block_size_int // gpu_block_size
self.tokens_per_block = tuple(group.tokens_per_block for group in config.groups)
self.tokens_per_hash = config.cache.tokens_per_hash
self.blocks_per_chunk = config.cache.blocks_per_chunk
@abstractmethod
def get_manager(self) -> OffloadingManager:
vllm/v1/kv_offload/config.py
+70
View File
@@ -0,0 +1,70 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Normalized configuration consumed by native offloading backends."""
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any
@dataclass(frozen=True)
class OffloadingGroupConfig:
# Total token span covered by one block across all workers
# (accounts for context parallelism).
tokens_per_block: int
# Layer names belonging to this group.
layer_names: tuple[str, ...]
@dataclass(frozen=True)
class OffloadingModelConfig:
# Model identifier (e.g. HuggingFace model path).
name: str
# KV cache data type (e.g. "float16").
dtype: str
@dataclass(frozen=True)
class OffloadingCacheConfig:
# Tokens per block hash.
tokens_per_hash: int
# Blocks coalesced into one offload chunk.
blocks_per_chunk: int
@dataclass(frozen=True)
class OffloadingParallelConfig:
# Worker index in [0, world_size). 0 on the scheduler side.
rank: int
# Total number of workers.
world_size: int
# Tensor parallel size.
tp_size: int
# Pipeline parallel size.
pp_size: int
# Prefill context parallel size.
pcp_size: int
# Decode context parallel size.
dcp_size: int
# Data parallel replica index of this engine.
data_parallel_index: int
# True when concatenating a block's data across all workers yields
# the same result regardless of the parallelism configuration.
is_parallelism_agnostic: bool
@dataclass(frozen=True)
class OffloadingConfig:
groups: tuple[OffloadingGroupConfig, ...]
# KV bytes stored by one worker per block.
worker_kv_bytes_per_block: int
# Whether the scheduler emits KV cache events. When true,
# the offloading backend should emit events as well.
enable_kv_cache_events: bool
# Offloading-specific configuration from kv_connector_extra_config.
extra_config: Mapping[str, Any]
# Unique identifier for this engine, distinct per DP rank.
engine_id: str
model: OffloadingModelConfig
cache: OffloadingCacheConfig
parallel: OffloadingParallelConfig
vllm/v1/kv_offload/cpu/gpu_worker.py
+28 -32
View File
@@ -72,7 +72,7 @@ class Transfer:
def compute_sub_block_ptrs(
block_ids: np.ndarray,
block_size_factor: int,
blocks_per_chunk: int,
output: np.ndarray,
tensor: torch.Tensor,
skip_count: int = 0,
@@ -80,38 +80,38 @@ def compute_sub_block_ptrs(
"""
Compute byte pointers for sub-blocks of the given block IDs.
Each block in block_ids contains block_size_factor sub-blocks.
Each block in block_ids contains blocks_per_chunk sub-blocks.
The pointer for sub-block j of block b is:
base_ptr + b * row_stride + j * sub_block_size
base_ptr + b * row_stride + j * block_page_size
where sub_block_size = tensor.shape[1] // block_size_factor (gpu page size).
where block_page_size = tensor.shape[1] // blocks_per_chunk (gpu page size).
This handles tensors where row_stride != block_size_factor * sub_block_size
This handles tensors where row_stride != blocks_per_chunk * block_page_size
(e.g. non-contiguous CPU tensors).
Args:
block_ids: array of block IDs at the tensor's native granularity.
block_size_factor: number of sub-blocks per block.
blocks_per_chunk: number of sub-blocks per block.
output: pre-allocated pointer array to write pointers into.
tensor: the source or destination tensor.
skip_count: sub-blocks to skip in the first block.
"""
assert skip_count < block_size_factor
assert skip_count < blocks_per_chunk
num_sub_blocks = len(output)
base_ptr = tensor.data_ptr()
row_stride = tensor.stride(0)
if block_size_factor == 1:
if blocks_per_chunk == 1:
# Fast path: 1:1 mapping, no sub-block expansion needed.
output[:] = base_ptr + block_ids.astype(np.uint64)[:num_sub_blocks] * row_stride
return
# Vectorized expansion for block_size_factor > 1.
assert tensor.shape[1] % block_size_factor == 0
sub_block_size = tensor.shape[1] // block_size_factor
sub_offsets = np.arange(block_size_factor, dtype=np.uint64) * sub_block_size
# (num_blocks, 1) + (1, block_size_factor) -> (num_blocks, block_size_factor)
# Vectorized expansion for blocks_per_chunk > 1.
assert tensor.shape[1] % blocks_per_chunk == 0
block_page_size = tensor.shape[1] // blocks_per_chunk
sub_offsets = np.arange(blocks_per_chunk, dtype=np.uint64) * block_page_size
# (num_blocks, 1) + (1, blocks_per_chunk) -> (num_blocks, blocks_per_chunk)
all_ptrs = (
base_ptr + block_ids.astype(np.uint64)[:, np.newaxis] * row_stride
) + sub_offsets[np.newaxis, :]
@@ -175,7 +175,7 @@ class SingleDirectionOffloadingHandler:
self,
gpu_tensors: list[torch.Tensor],
cpu_tensors: list[torch.Tensor],
block_size_factor: int,
blocks_per_chunk: int,
kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
gpu_to_cpu: bool,
mmap_region: SharedOffloadRegion | None = None,
@@ -205,7 +205,7 @@ class SingleDirectionOffloadingHandler:
assert cpu_tensor.device.type == "cpu"
_, gpu_page_size = gpu_tensor.shape
_, cpu_page_size = cpu_tensor.shape
assert cpu_page_size == gpu_page_size * block_size_factor
assert cpu_page_size == gpu_page_size * blocks_per_chunk
self.src_tensors: list[torch.Tensor] = (
gpu_tensors if gpu_to_cpu else cpu_tensors
@@ -220,9 +220,9 @@ class SingleDirectionOffloadingHandler:
)
# GPU blocks may be smaller
# cpu_page_size = gpu_page_size * block_size_factor.
self.src_block_size_factor = 1 if self.gpu_to_cpu else block_size_factor
self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1
# cpu_page_size = gpu_page_size * blocks_per_chunk.
self.src_blocks_per_chunk = 1 if self.gpu_to_cpu else blocks_per_chunk
self.dst_blocks_per_chunk = blocks_per_chunk if self.gpu_to_cpu else 1
# mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
self._mmap_region = mmap_region
@@ -313,20 +313,16 @@ class SingleDirectionOffloadingHandler:
if group_size == 0:
continue
src_logical_blocks_to_skip = block_idx % self.src_block_size_factor
dst_logical_blocks_to_skip = block_idx % self.dst_block_size_factor
src_logical_blocks_to_skip = block_idx % self.src_blocks_per_chunk
dst_logical_blocks_to_skip = block_idx % self.dst_blocks_per_chunk
src_logical_blocks_count = group_size + src_logical_blocks_to_skip
dst_logical_blocks_count = group_size + dst_logical_blocks_to_skip
dst_blocks_count = cdiv(
dst_logical_blocks_count, self.dst_block_size_factor
)
dst_blocks_count = cdiv(dst_logical_blocks_count, self.dst_blocks_per_chunk)
dst_end_offset = dst_offset + dst_blocks_count
assert dst_end_offset <= num_dst_blocks
src_blocks_count = cdiv(
src_logical_blocks_count, self.src_block_size_factor
)
src_blocks_count = cdiv(src_logical_blocks_count, self.src_blocks_per_chunk)
src_end_offset = src_offset + src_blocks_count
assert src_end_offset <= num_src_blocks
@@ -339,14 +335,14 @@ class SingleDirectionOffloadingHandler:
compute_sub_block_ptrs(
group_src,
self.src_block_size_factor,
self.src_blocks_per_chunk,
all_src[op_idx:end_idx],
self.src_tensors[t_idx],
skip_count=src_logical_blocks_to_skip,
)
compute_sub_block_ptrs(
group_dst,
self.dst_block_size_factor,
self.dst_blocks_per_chunk,
all_dst[op_idx:end_idx],
self.dst_tensors[t_idx],
skip_count=dst_logical_blocks_to_skip,
@@ -476,7 +472,7 @@ class CPUOffloadingWorker(OffloadingWorker):
def __init__(
self,
kv_caches: CanonicalKVCaches,
block_size_factor: int,
blocks_per_chunk: int,
num_cpu_blocks: int,
mmap_region: SharedOffloadRegion | None = None,
):
@@ -492,7 +488,7 @@ class CPUOffloadingWorker(OffloadingWorker):
gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view(
(-1, gpu_page_size_bytes)
)
cpu_page_size_bytes = gpu_page_size_bytes * block_size_factor
cpu_page_size_bytes = gpu_page_size_bytes * blocks_per_chunk
if mmap_region is not None:
cpu_tensor = mmap_region.create_next_view(cpu_page_size_bytes)
@@ -518,7 +514,7 @@ class CPUOffloadingWorker(OffloadingWorker):
self._store_handler = SingleDirectionOffloadingHandler(
gpu_tensors=gpu_tensors,
cpu_tensors=cpu_tensors,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
kv_cache_groups_data_refs=kv_caches.group_data_refs,
gpu_to_cpu=True,
mmap_region=mmap_region,
@@ -527,7 +523,7 @@ class CPUOffloadingWorker(OffloadingWorker):
self._load_handler = SingleDirectionOffloadingHandler(
gpu_tensors=gpu_tensors,
cpu_tensors=cpu_tensors,
block_size_factor=block_size_factor,
blocks_per_chunk=blocks_per_chunk,
kv_cache_groups_data_refs=kv_caches.group_data_refs,
gpu_to_cpu=False,
)
vllm/v1/kv_offload/cpu/spec.py
+16 -31
View File
@@ -4,10 +4,8 @@ from typing import Any
from typing_extensions import override
from vllm.config import VllmConfig
from vllm.platforms import current_platform
from vllm.utils.math_utils import round_up
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_offload.base import (
CanonicalKVCaches,
OffloadingCounterMetadata,
@@ -18,6 +16,7 @@ from vllm.v1.kv_offload.base import (
OffloadingSpec,
OffloadingWorker,
)
from vllm.v1.kv_offload.config import OffloadingConfig
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -73,8 +72,8 @@ class CPUOffloadingSpec(OffloadingSpec):
)
return definitions
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
super().__init__(vllm_config, kv_cache_config)
def __init__(self, config: OffloadingConfig):
super().__init__(config)
cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use")
if not cpu_bytes_to_use:
@@ -82,42 +81,28 @@ class CPUOffloadingSpec(OffloadingSpec):
"cpu_bytes_to_use must be specified in kv_connector_extra_config"
)
world_size = vllm_config.parallel_config.world_size
world_size = config.parallel.world_size
self.num_blocks = 0
self.kv_bytes_per_offloaded_block = 0
self.kv_bytes_per_chunk = 0
self.cpu_page_size_per_worker = 0
assert kv_cache_config is not None
if kv_cache_config.num_blocks > 0 and world_size > 0:
is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors)
assert not is_packed or all(
t.block_stride for t in kv_cache_config.kv_cache_tensors
)
total_gpu_kv_bytes = (
kv_cache_config.kv_cache_tensors[0].size
if is_packed
else sum(t.size for t in kv_cache_config.kv_cache_tensors)
)
kv_bytes_per_block = (
total_gpu_kv_bytes // kv_cache_config.num_blocks
) * world_size
kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
if config.worker_kv_bytes_per_block > 0 and world_size > 0:
kv_bytes_per_block = config.worker_kv_bytes_per_block * world_size
kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk
# calculate cpu_page_size_per_worker
self.cpu_page_size_per_worker = kv_bytes_per_offloaded_block // world_size
self.cpu_page_size_per_worker = kv_bytes_per_chunk // world_size
# calculate num_blocks
aligned_kv_bytes_per_offloaded_block = round_up(
kv_bytes_per_offloaded_block, self.BLOCK_SIZE_ALIGNMENT
)
self.num_blocks = (
int(cpu_bytes_to_use) // aligned_kv_bytes_per_offloaded_block
aligned_kv_bytes_per_chunk = round_up(
kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT
)
self.num_blocks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk
# Expose aligned_kv_bytes_per_offloaded_block as
# kv_bytes_per_offloaded_block. Note that this might contain
# Expose aligned_kv_bytes_per_chunk as
# kv_bytes_per_chunk. Note that this might contain
# some padding. i.e. each offloaded block is of the form,
# |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
self.kv_bytes_per_offloaded_block = aligned_kv_bytes_per_offloaded_block
self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk
# scheduler-side
self._manager: OffloadingManager | None = None
@@ -150,7 +135,7 @@ class CPUOffloadingSpec(OffloadingSpec):
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
return CPUOffloadingWorker(
kv_caches=kv_caches,
block_size_factor=self.block_size_factor,
blocks_per_chunk=self.blocks_per_chunk,
num_cpu_blocks=self.num_blocks,
)
vllm/v1/kv_offload/factory.py
+8 -22
View File
@@ -1,15 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib
from collections.abc import Callable
from typing import TYPE_CHECKING
from collections.abc import Callable, Mapping
from typing import Any
from vllm.logger import init_logger
from vllm.v1.kv_offload.base import OffloadingSpec
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_offload.config import OffloadingConfig
logger = init_logger(__name__)
@@ -30,10 +27,7 @@ class OffloadingSpecFactory:
cls._registry[name] = loader
@classmethod
def get_spec_cls(cls, config: "VllmConfig") -> type[OffloadingSpec]:
kv_transfer_config = config.kv_transfer_config
assert kv_transfer_config is not None
extra_config = kv_transfer_config.kv_connector_extra_config
def get_spec_cls(cls, extra_config: Mapping[str, Any]) -> type[OffloadingSpec]:
spec_name = extra_config.get("spec_name", "CPUOffloadingSpec")
if spec_name in cls._registry:
spec_cls = cls._registry[spec_name]()
@@ -47,19 +41,11 @@ class OffloadingSpecFactory:
return spec_cls
@classmethod
def create_spec(
cls,
config: "VllmConfig",
kv_cache_config: "KVCacheConfig",
) -> OffloadingSpec:
kv_transfer_config = config.kv_transfer_config
assert kv_transfer_config is not None
spec_name = kv_transfer_config.kv_connector_extra_config.get(
"spec_name", "CPUOffloadingSpec"
)
spec_cls = cls.get_spec_cls(config)
def create_spec(cls, config: OffloadingConfig) -> OffloadingSpec:
spec_name = config.extra_config.get("spec_name", "CPUOffloadingSpec")
spec_cls = cls.get_spec_cls(config.extra_config)
logger.info("Creating offloading spec with name: %s", spec_name)
return spec_cls(config, kv_cache_config)
return spec_cls(config)
# Register various specs here.
vllm/v1/kv_offload/file_mapper.py
+19 -35
View File
@@ -4,7 +4,6 @@
import hashlib
import json
from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
from vllm.v1.kv_offload.base import (
OffloadingSpec,
OffloadKey,
@@ -25,8 +24,8 @@ class FileMapper:
self,
root_dir: str,
model_name: str,
hash_block_size: int,
gpu_blocks_per_file: int,
tokens_per_hash: int,
blocks_per_file: int,
tp_size: int,
pp_size: int,
pcp_size: int,
@@ -49,8 +48,8 @@ class FileMapper:
self.rank: int = rank
self.fields: dict = {
"model_name": model_name,
"hash_block_size": hash_block_size,
"gpu_blocks_per_file": gpu_blocks_per_file,
"tokens_per_hash": tokens_per_hash,
"blocks_per_file": blocks_per_file,
"tp_size": tp_size,
"pp_size": pp_size,
"pcp_size": pcp_size,
@@ -66,47 +65,32 @@ class FileMapper:
cls,
root_dir: str,
offloading_spec: OffloadingSpec,
gpu_blocks_per_file: int = 1,
blocks_per_file: int = 1,
parallel_agnostic: bool = False,
) -> "FileMapper":
"""Build a FileMapper from an OffloadingSpec."""
vllm_config = offloading_spec.vllm_config
kv_cache_config = offloading_spec.kv_cache_config
parallel_config = vllm_config.parallel_config
dtype = str(vllm_config.cache_config.cache_dtype).replace("torch.", "")
config = offloading_spec.config
kv_cache_groups = [
{
"block_size": group.kv_cache_spec.block_size,
"tokens_per_block": group.tokens_per_block,
"layer_names": list(group.layer_names),
}
for group in kv_cache_config.kv_cache_groups
for group in config.groups
]
# Only a single full-attention group is parallelism-invariant. MLA is
# excluded: its latent KV is replicated per rank, never head-sharded.
# The V2 model runner is excluded: its KV layout is not known to be
# parallelism-invariant.
groups = kv_cache_config.kv_cache_groups
spec = groups[0].kv_cache_spec if len(groups) == 1 else None
parallel_agnostic = (
parallel_agnostic
and not vllm_config.use_v2_model_runner
and isinstance(spec, FullAttentionSpec)
and not isinstance(spec, MLAAttentionSpec)
)
parallel = config.parallel
return cls(
root_dir=root_dir,
model_name=vllm_config.model_config.model,
hash_block_size=vllm_config.cache_config.block_size,
gpu_blocks_per_file=gpu_blocks_per_file,
tp_size=parallel_config.tensor_parallel_size,
pp_size=parallel_config.pipeline_parallel_size,
pcp_size=parallel_config.prefill_context_parallel_size,
dcp_size=parallel_config.decode_context_parallel_size,
rank=parallel_config.rank,
dtype=dtype,
model_name=config.model.name,
tokens_per_hash=config.cache.tokens_per_hash,
blocks_per_file=blocks_per_file,
tp_size=parallel.tp_size,
pp_size=parallel.pp_size,
pcp_size=parallel.pcp_size,
dcp_size=parallel.dcp_size,
rank=parallel.rank,
dtype=config.model.dtype,
kv_cache_groups=kv_cache_groups,
parallel_agnostic=parallel_agnostic,
parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic),
)
def get_file_name(self, key: OffloadKey) -> str:
vllm/v1/kv_offload/tiering/fs/manager.py
+3 -3
View File
@@ -113,8 +113,8 @@ class FileSystemTierManager(SecondaryTierManager):
):
"""
Args:
offloading_spec: contains the vllm_config, kv_cache_config
and block_size_factor.
offloading_spec: Contains normalized offloading configuration and
blocks_per_chunk.
primary_kv_view: Memoryview of the primary tier's CPU KV cache.
tier_type: Tier type identifier, set by SecondaryTierFactory.
root_dir: Root directory for block files.
@@ -150,7 +150,7 @@ class FileSystemTierManager(SecondaryTierManager):
self.file_mapper = FileMapper.from_offloading_spec(
root_dir=root_dir,
offloading_spec=offloading_spec,
gpu_blocks_per_file=offloading_spec.block_size_factor,
blocks_per_file=offloading_spec.blocks_per_chunk,
parallel_agnostic=True,
)
vllm/v1/kv_offload/tiering/p2p/manager.py
+4 -4
View File
@@ -127,8 +127,8 @@ class P2PSecondaryTierManager(SecondaryTierManager):
configuration reference.
Args:
offloading_spec: Owning ``OffloadingSpec`` (provides
``vllm_config`` and the offloaded block layout).
offloading_spec: Owning ``OffloadingSpec`` (provides normalized
model, parallel, and cache layout configuration).
primary_kv_view: Memoryview over the CPU primary tier; the
NIXL agent registers this region for RDMA transfers.
tier_type: Tier identifier (defaults to ``"p2p"``).
@@ -164,7 +164,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
# One control socket per DP replica: offset the base by the global
# data-parallel index so replicas on a host don't collide (mirrors
# NIXL). For DP=1 the index is 0, leaving the base port unchanged.
dp_index = offloading_spec.vllm_config.parallel_config.data_parallel_index
dp_index = offloading_spec.config.parallel.data_parallel_index
port = int(port) + dp_index
# Two decoupled identities:
# _local_id (``host:port``): the ZMQ control identity that peers
@@ -181,7 +181,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
config_fields = FileMapper.from_offloading_spec(
root_dir="",
offloading_spec=offloading_spec,
gpu_blocks_per_file=offloading_spec.block_size_factor,
blocks_per_file=offloading_spec.blocks_per_chunk,
parallel_agnostic=True,
).get_run_config()
self._data: DataTransport = NixlTransport(
vllm/v1/kv_offload/tiering/spec.py
+9 -12
View File
@@ -36,14 +36,13 @@ from typing import Any
import torch
from typing_extensions import override
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_offload.base import (
CanonicalKVCaches,
OffloadingManager,
OffloadingMetricMetadata,
)
from vllm.v1.kv_offload.config import OffloadingConfig
from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
@@ -87,8 +86,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
metrics.update(tier_cls.build_metric_definitions(tier_config))
return metrics
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
super().__init__(vllm_config, kv_cache_config)
def __init__(self, config: OffloadingConfig):
super().__init__(config)
# Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
self._manager: OffloadingManager | None = None
if self.kv_events_config.self_describing_kv_events:
@@ -110,10 +109,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
# engine_id is unique per DP replica (suffixed with _dp{rank} in both
# the Ray and multiprocessing paths), so it names a per-replica offload
# region. Non-None is guaranteed by OffloadingSpec.__init__.
assert vllm_config.kv_transfer_config is not None
assert vllm_config.kv_transfer_config.engine_id is not None
self._engine_id: str = vllm_config.kv_transfer_config.engine_id
# region.
self._engine_id = config.engine_id
@override
def get_manager(self) -> OffloadingManager:
@@ -134,7 +131,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
engine_id=self._engine_id,
num_blocks=self.num_blocks,
rank=None,
kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
kv_bytes_per_block=self.kv_bytes_per_chunk,
cpu_page_size=self.cpu_page_size_per_worker,
)
self._scheduler_mmap = scheduler_mmap
@@ -196,18 +193,18 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
# Fold the global physical device index into the replica-local
# [0, world_size) slot range.
world_size = self.vllm_config.parallel_config.world_size
world_size = self.config.parallel.world_size
rank = torch.accelerator.current_device_index() % world_size
worker_mmap = SharedOffloadRegion(
engine_id=self._engine_id,
num_blocks=self.num_blocks,
rank=rank,
kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
kv_bytes_per_block=self.kv_bytes_per_chunk,
cpu_page_size=self.cpu_page_size_per_worker,
)
return CPUOffloadingWorker(
kv_caches=kv_caches,
block_size_factor=self.block_size_factor,
blocks_per_chunk=self.blocks_per_chunk,
num_cpu_blocks=self.num_blocks,
mmap_region=worker_mmap,
)