diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index beb639724e8..6387cf6395e 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -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)) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 6c55b91d8da..017be3bbd5e 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -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) ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 81c00266cfe..25bb664ff40 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -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 diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 73ea5e2be1d..b878e294a6e 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -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, diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 16420b164b2..a323ae7ce52 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -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), ] diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 12dbc57fe97..2f1ce67e9c4 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -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) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index b051f44ecfd..76403a267bc 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -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 ) diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 6f6e0d66196..6c11f2d465f 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -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: _r//_g/.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 - The final segment is .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 diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index 7eb81243524..d3bacc8f326 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -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) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 52e2c836323..01fd295302a 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -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) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 4ac734d957f..2310627e9c8 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -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) diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 37687adce0c..877e9e82ec4 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -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, + ), ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py new file mode 100644 index 00000000000..bf86d02ec46 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -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, + ), + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 410f84c50dd..9c567008de3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -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, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index baa168e1708..58023c39356 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -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(), diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index f896c9cc492..3e0d09a10fd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 260746be051..e82e1e45c00 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 197beca9aec..2fe4bf6a5a7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -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: diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 5a2e3c184d3..17dcd41fbe9 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -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: diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py new file mode 100644 index 00000000000..cd7b3ee2075 --- /dev/null +++ b/vllm/v1/kv_offload/config.py @@ -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 diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index c8b9915a1e5..baf9a66719a 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -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, ) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 4ad6974857c..f6d1c29aff9 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -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, ) diff --git a/vllm/v1/kv_offload/factory.py b/vllm/v1/kv_offload/factory.py index abbc9c0ede7..931fda8308f 100644 --- a/vllm/v1/kv_offload/factory.py +++ b/vllm/v1/kv_offload/factory.py @@ -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. diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index d8fadb09988..b85d4d06979 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -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: diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index d8d17002856..9c51e16432e 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -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, ) diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 0cb37dff0e0..95605fd1735 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -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( diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 5f9e8cdc237..ebf7d878d09 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -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, )