forked from Karylab-cklius/vllm
[kv_offload+HMA][13/N]: Enable HMA support (#41445)
This is the final PR in a series to enables HMA support for the offloading connector. The connector advertises `SupportsHMA` and is validated with unit tests and e2e tests. Signed-off-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Or Ozeri
@@ -4,6 +4,7 @@ from collections.abc import Iterable
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.kv_connector.unit.offloading_connector.utils import (
|
||||
generate_store_output,
|
||||
@@ -15,6 +16,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
|
||||
OffloadingConnectorScheduler,
|
||||
)
|
||||
from vllm.v1.core.kv_cache_utils import BlockHash
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheGroupSpec,
|
||||
SlidingWindowSpec,
|
||||
)
|
||||
from vllm.v1.kv_offload.base import (
|
||||
OffloadingEvent,
|
||||
OffloadingManager,
|
||||
@@ -26,32 +32,32 @@ from vllm.v1.request import RequestStatus
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||
def test_offloading_connector(request_runner, async_scheduling: bool):
|
||||
offloaded_block_size = 12
|
||||
gpu_block_size = 4
|
||||
block_size = 4
|
||||
block_size_factor = 3
|
||||
offloaded_block_size = block_size * block_size_factor
|
||||
num_gpu_blocks = 100
|
||||
block_size_factor = offloaded_block_size // gpu_block_size
|
||||
|
||||
runner = request_runner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=gpu_block_size,
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
async_scheduling=async_scheduling,
|
||||
block_size_factor=block_size_factor,
|
||||
)
|
||||
|
||||
# 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.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(list(keys)[1:2])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(list(keys)[1:2])
|
||||
)
|
||||
runner.run(decoded_tokens=[0])
|
||||
|
||||
# add block missing 1 token -> no offload
|
||||
runner.run(
|
||||
decoded_tokens=[0] * (offloaded_block_size - 1),
|
||||
expected_stored_gpu_block_indexes=(3, 4, 5),
|
||||
expected_stored=(3, 4, 5),
|
||||
)
|
||||
runner.manager.prepare_store.assert_not_called()
|
||||
runner.manager.touch.assert_not_called()
|
||||
|
||||
# +1 token -> single block, fail prepare_store
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: None
|
||||
@@ -60,19 +66,19 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
|
||||
|
||||
# 1 more block (+ token for async scheduling)
|
||||
# now set block_hashes_to_store = []
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
|
||||
|
||||
# 1 more block (+ token for kicking off offloading)
|
||||
# now check touch was called with all 6 blocks
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0] * (offloaded_block_size + 1),
|
||||
expected_stored_gpu_block_indexes=(15, 16, 17),
|
||||
expected_stored=(15, 16, 17),
|
||||
)
|
||||
runner.manager.touch.assert_called()
|
||||
block_hashes1 = list(runner.manager.touch.call_args.args[0])
|
||||
@@ -95,23 +101,23 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
|
||||
# terminate request
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
|
||||
expected_stored=tuple(range(6 * block_size_factor)),
|
||||
)
|
||||
|
||||
# full_block_tokens - num_computed_tokens < offloaded_block_size
|
||||
runner.new_request(
|
||||
token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
|
||||
token_ids=[0] * block_size + [1] * (offloaded_block_size - block_size)
|
||||
)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID])
|
||||
runner.manager.lookup.assert_not_called()
|
||||
|
||||
# single block lookup with no hits
|
||||
runner.new_request(token_ids=[1] * offloaded_block_size)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID])
|
||||
runner.manager.lookup.assert_called_once()
|
||||
@@ -119,25 +125,21 @@ 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.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
|
||||
)
|
||||
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.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
|
||||
|
||||
# test take_events
|
||||
def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
|
||||
@@ -167,15 +169,16 @@ 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):
|
||||
offloaded_block_size = 12
|
||||
gpu_block_size = 4
|
||||
block_size = 4
|
||||
block_size_factor = 3
|
||||
offloaded_block_size = block_size * block_size_factor
|
||||
num_gpu_blocks = 100
|
||||
|
||||
runner = request_runner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=gpu_block_size,
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
async_scheduling=async_scheduling,
|
||||
block_size_factor=block_size_factor,
|
||||
)
|
||||
|
||||
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
|
||||
@@ -184,8 +187,8 @@ 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.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0],
|
||||
@@ -193,11 +196,11 @@ def test_request_preemption(request_runner, async_scheduling: bool):
|
||||
)
|
||||
|
||||
# decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
|
||||
decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
|
||||
complete_transfers=False,
|
||||
)
|
||||
|
||||
@@ -208,8 +211,8 @@ def test_request_preemption(request_runner, async_scheduling: bool):
|
||||
runner.run(
|
||||
decoded_tokens=[],
|
||||
complete_transfers=False,
|
||||
expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
|
||||
expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
|
||||
expected_flushed=(0, 1, 2, 3, 4, 5, 6, 7, 8),
|
||||
expected_stored=(0, 1, 2, 3, 4, 5, 6, 7, 8),
|
||||
)
|
||||
|
||||
# restore KV cache space and reset GPU prefix cache
|
||||
@@ -219,17 +222,17 @@ def test_request_preemption(request_runner, async_scheduling: bool):
|
||||
# request should now return from preemption
|
||||
# re-load [0, ..., 8] from the CPU and store [9, 10, 11]
|
||||
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0] * gpu_block_size,
|
||||
expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
|
||||
decoded_tokens=[0] * block_size,
|
||||
expected_loaded=(0, 1, 2, 3, 4, 5, 6, 7, 8),
|
||||
)
|
||||
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=(9, 10, 11),
|
||||
expected_stored=(9, 10, 11),
|
||||
)
|
||||
|
||||
# All stores completed before request_finished -> fence index empty.
|
||||
@@ -238,25 +241,26 @@ def test_request_preemption(request_runner, async_scheduling: bool):
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
|
||||
offloaded_block_size = 12
|
||||
gpu_block_size = 4
|
||||
block_size = 4
|
||||
block_size_factor = 3
|
||||
offloaded_block_size = block_size * block_size_factor
|
||||
num_gpu_blocks = 100
|
||||
|
||||
runner = request_runner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=gpu_block_size,
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
async_scheduling=async_scheduling,
|
||||
block_size_factor=block_size_factor,
|
||||
)
|
||||
|
||||
# store 1 blocks
|
||||
runner.new_request(token_ids=[0] * offloaded_block_size)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=(0, 1, 2),
|
||||
expected_stored=(0, 1, 2),
|
||||
)
|
||||
|
||||
# start a request to load the first block, but don't complete
|
||||
@@ -284,12 +288,12 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
|
||||
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
|
||||
|
||||
# complete transfers
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_loaded_gpu_block_indexes=(0, 1, 2),
|
||||
expected_loaded=(0, 1, 2),
|
||||
)
|
||||
|
||||
# second request will use the GPU prefix cache
|
||||
@@ -301,25 +305,26 @@ 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):
|
||||
offloaded_block_size = 12
|
||||
gpu_block_size = 4
|
||||
block_size = 4
|
||||
block_size_factor = 3
|
||||
offloaded_block_size = block_size * block_size_factor
|
||||
num_gpu_blocks = 100
|
||||
|
||||
runner = request_runner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=gpu_block_size,
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
async_scheduling=async_scheduling,
|
||||
block_size_factor=block_size_factor,
|
||||
)
|
||||
|
||||
# store 1 blocks
|
||||
runner.new_request(token_ids=[0] * offloaded_block_size)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=(0, 1, 2),
|
||||
expected_stored=(0, 1, 2),
|
||||
)
|
||||
|
||||
# start a request to load the first block, but don't complete
|
||||
@@ -345,13 +350,264 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
|
||||
# complete loading request
|
||||
runner.run(
|
||||
decoded_tokens=[],
|
||||
expected_loaded_gpu_block_indexes=(0, 1, 2),
|
||||
expected_loaded=(0, 1, 2),
|
||||
)
|
||||
|
||||
# assert request is deleted
|
||||
assert req_id not in runner.scheduler.requests
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||
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
|
||||
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(
|
||||
["layer0"],
|
||||
FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
["layer1"],
|
||||
SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=sliding_window,
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
runner = request_runner(
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
async_scheduling=async_scheduling,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
|
||||
# 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
|
||||
|
||||
# Blocks [0, 1, 2] miss
|
||||
runner.new_request(token_ids=[0] * block_size * 3)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(decoded_tokens=[0])
|
||||
# _touch called from get_num_new_matched_tokens (2 groups) and
|
||||
# _get_reqs_to_store (2 groups) → 4 touch calls total.
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 4
|
||||
assert len(touch_calls[0].args[0]) == 3
|
||||
assert len(touch_calls[1].args[0]) == 3
|
||||
assert len(touch_calls[2].args[0]) == 3
|
||||
assert len(touch_calls[3].args[0]) == 3
|
||||
|
||||
# store 3 more block
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0] * (block_size * 3 + 2),
|
||||
expected_stored=(0, 1, 2, 3, 4, 5),
|
||||
)
|
||||
|
||||
# touch called from _get_reqs_to_store * 3 blocks, once for each group
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 6
|
||||
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID])
|
||||
|
||||
runner.scheduler.reset_prefix_cache()
|
||||
|
||||
# full 3 blocks hit [0, 1, 2]
|
||||
runner.new_request(token_ids=[0] * (block_size * 3 + 1))
|
||||
runner.manager.lookup.return_value = True
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
# Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2
|
||||
# Group 1 (sliding window, window=2): only the last 2 blocks
|
||||
# are within the window → loads blocks 1,2
|
||||
expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
|
||||
)
|
||||
|
||||
# one touch in get_num_new_matched_tokens x 2 groups
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 2
|
||||
# full attention group touched all 3 blocks
|
||||
assert len(touch_calls[0].args[0]) == 3
|
||||
# sliding window group touched just the last 2 blocks
|
||||
assert len(touch_calls[1].args[0]) == 2
|
||||
|
||||
# 3 blocks are hit on GPU [0, 1, 2]
|
||||
# 1 block loaded [3,]
|
||||
runner.new_request(token_ids=[0] * (block_size * 4 + 1))
|
||||
runner.manager.lookup.return_value = True
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
# Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2
|
||||
# Group 1 (sliding window, window=2): only the last 2 blocks
|
||||
# are within the window → loads blocks 1,2
|
||||
expected_loaded=((0, 3), (1, 3)),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||
def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool):
|
||||
hash_block_size = 4
|
||||
num_gpu_blocks = 100
|
||||
|
||||
# Group 0: block_size=12 (offloaded_block_size=12)
|
||||
# Group 1: block_size=16 (offloaded_block_size=16)
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(
|
||||
["layer0"],
|
||||
FullAttentionSpec(
|
||||
block_size=hash_block_size * 3,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
["layer1"],
|
||||
FullAttentionSpec(
|
||||
block_size=hash_block_size * 4,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
runner = request_runner(
|
||||
block_size=hash_block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
async_scheduling=async_scheduling,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
|
||||
# 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
|
||||
|
||||
# Prompt: 25 tokens, unaligned to both block sizes.
|
||||
# Group 0 blocks: [0, 1], ending_token_offset = 24
|
||||
# Group 1 blocks: [0,], ending_token_offset = 16
|
||||
runner.new_request(token_ids=[0] * 25)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(decoded_tokens=[0])
|
||||
# _touch called from get_num_new_matched_tokens (2 groups) and
|
||||
# _get_reqs_to_store (2 groups) → 4 touch calls total.
|
||||
# Group 0 has 2 offload keys, group 1 has 1.
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 4
|
||||
assert len(touch_calls[0].args[0]) == 2
|
||||
assert len(touch_calls[1].args[0]) == 1
|
||||
assert len(touch_calls[2].args[0]) == 2
|
||||
assert len(touch_calls[3].args[0]) == 1
|
||||
|
||||
# Get to 31 tokens
|
||||
# No further blocks offloaded
|
||||
runner.run(decoded_tokens=[0] * 6, expected_stored=((0, 0), (0, 1), (1, 0)))
|
||||
|
||||
# Get to 32 tokens
|
||||
# Group 0 blocks: [0, 1], ending_token_offset = 24
|
||||
# Group 1 blocks: [0, 1], ending_token_offset = 32
|
||||
runner.run(decoded_tokens=[0])
|
||||
# _get_reqs_to_store touch: only group 1 has a new block to store
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 2
|
||||
assert len(touch_calls[0].args[0]) == 2
|
||||
assert len(touch_calls[1].args[0]) == 2
|
||||
|
||||
# Get to 35 tokens
|
||||
# No further blocks offloaded
|
||||
runner.run(decoded_tokens=[0] * 3, expected_stored=((1, 1),))
|
||||
|
||||
# Get to 36 tokens
|
||||
# Group 0 blocks: [0, 1, 2], ending_token_offset = 36
|
||||
# Group 1 blocks: [0, 1], ending_token_offset = 32
|
||||
runner.run(decoded_tokens=[0])
|
||||
# _get_reqs_to_store touch: only group 0 has a new block to store
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 2
|
||||
assert len(touch_calls[0].args[0]) == 3
|
||||
assert len(touch_calls[1].args[0]) == 2
|
||||
|
||||
# Get to 47 tokens
|
||||
# No further blocks offloaded
|
||||
runner.run(decoded_tokens=[0] * 11, expected_stored=((0, 2),))
|
||||
|
||||
# Get to 48 tokens
|
||||
# Group 0 blocks: [0, 1, 2, 3], ending_token_offset = 4
|
||||
# Group 1 blocks: [0, 1, 2], ending_token_offset = 48
|
||||
runner.run(decoded_tokens=[0])
|
||||
# _get_reqs_to_store touch: both groups have a new block, each with 1 key
|
||||
touch_calls = runner.manager.touch.call_args_list
|
||||
assert len(touch_calls) == 2
|
||||
assert len(touch_calls[0].args[0]) == 4
|
||||
assert len(touch_calls[1].args[0]) == 3
|
||||
|
||||
runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2)))
|
||||
|
||||
# Get to 96 tokens
|
||||
runner.run(
|
||||
decoded_tokens=[0] * 47 + [EOS_TOKEN_ID],
|
||||
expected_stored=((0, 4), (0, 5), (0, 6), (0, 7), (1, 3), (1, 4), (1, 5)),
|
||||
)
|
||||
|
||||
runner.scheduler.reset_prefix_cache()
|
||||
|
||||
# Request with 48 matching tokens
|
||||
# will match 48 tokens (4 block) from the first group
|
||||
# 48 tokens (3 block) from the second group
|
||||
# Total 48 tokens can be loaded
|
||||
runner.new_request(token_ids=[0] * 48)
|
||||
runner.manager.lookup.return_value = True
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0],
|
||||
expected_loaded=((0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2)),
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID])
|
||||
|
||||
# Request with 48+37 matching tokens
|
||||
# 48 tokens will be hit on GPU
|
||||
# extra 32 tokens will be loaded
|
||||
# extra tokens [0, 36] (blocks [4, 5, 6]) from the first group
|
||||
# extra tokens [0, 32] (block [3, 4]) from the second group
|
||||
runner.new_request(token_ids=[0] * (48 + 37))
|
||||
runner.manager.lookup.return_value = True
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[0],
|
||||
expected_loaded=((0, 4), (0, 5), (0, 6), (1, 3), (1, 4)),
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for _maximal_prefix_lookup / _sliding_window_lookup
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -493,26 +749,24 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
|
||||
|
||||
This supports P/D disaggregation where the prefill instance offloads the
|
||||
complete KV cache so a remote decode node can consume it."""
|
||||
offloaded_block_size = 12
|
||||
gpu_block_size = 4
|
||||
block_size_factor = 3
|
||||
offloaded_block_size = gpu_block_size * block_size_factor
|
||||
num_gpu_blocks = 100
|
||||
|
||||
runner = request_runner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=gpu_block_size,
|
||||
block_size_factor=block_size_factor,
|
||||
block_size=gpu_block_size,
|
||||
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)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=(0, 1, 2),
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
|
||||
|
||||
# Reset GPU prefix cache so the next request must load from CPU.
|
||||
runner.scheduler.reset_prefix_cache()
|
||||
@@ -524,22 +778,17 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
|
||||
kv_transfer_params={"do_remote_decode": True},
|
||||
)
|
||||
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)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
|
||||
# Load the first offloaded block from CPU.
|
||||
runner.run(
|
||||
decoded_tokens=[0],
|
||||
expected_loaded_gpu_block_indexes=(0, 1, 2),
|
||||
)
|
||||
runner.run(decoded_tokens=[0], expected_loaded=(0, 1, 2))
|
||||
|
||||
# Store must include ALL 6 GPU blocks (both the loaded prefix and
|
||||
# the newly computed block), not just the 3 new ones.
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2, 3, 4, 5))
|
||||
|
||||
# All stores completed before request_finished -> fence index empty.
|
||||
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
|
||||
|
||||
@@ -553,8 +802,8 @@ 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(
|
||||
offloaded_block_size=12,
|
||||
gpu_block_size=4,
|
||||
block_size_factor=3,
|
||||
block_size=4,
|
||||
num_gpu_blocks=100,
|
||||
async_scheduling=False,
|
||||
)
|
||||
@@ -572,15 +821,15 @@ def test_fence_at_update_state_after_alloc(request_runner):
|
||||
req1 just freed.
|
||||
"""
|
||||
runner = request_runner(
|
||||
offloaded_block_size=4,
|
||||
gpu_block_size=4,
|
||||
block_size_factor=1,
|
||||
block_size=4,
|
||||
num_gpu_blocks=2,
|
||||
async_scheduling=False,
|
||||
)
|
||||
|
||||
runner.new_request(token_ids=[0] * 4)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
|
||||
assert runner.connector_scheduler._block_id_to_pending_jobs
|
||||
@@ -588,14 +837,14 @@ def test_fence_at_update_state_after_alloc(request_runner):
|
||||
runner.scheduler.reset_prefix_cache()
|
||||
runner.new_request(token_ids=[0] * 4)
|
||||
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[],
|
||||
complete_transfers=False,
|
||||
expected_stored_gpu_block_indexes=(0,),
|
||||
expected_flushed_gpu_block_indexes=(0,),
|
||||
expected_stored=(0,),
|
||||
expected_flushed=(0,),
|
||||
)
|
||||
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
|
||||
|
||||
@@ -605,15 +854,15 @@ 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(
|
||||
offloaded_block_size=4,
|
||||
gpu_block_size=4,
|
||||
block_size_factor=1,
|
||||
block_size=4,
|
||||
num_gpu_blocks=2,
|
||||
async_scheduling=False,
|
||||
)
|
||||
|
||||
runner.new_request(token_ids=[0] * 4)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
|
||||
assert runner.connector_scheduler._block_id_to_pending_jobs
|
||||
@@ -621,13 +870,13 @@ def test_fence_at_build_store_jobs(request_runner):
|
||||
runner.scheduler.reset_prefix_cache()
|
||||
runner.new_request(token_ids=[1] * 4)
|
||||
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output([])
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output([])
|
||||
)
|
||||
runner.run(
|
||||
decoded_tokens=[EOS_TOKEN_ID],
|
||||
expected_stored_gpu_block_indexes=(0,),
|
||||
expected_flushed_gpu_block_indexes=(0,),
|
||||
expected_stored=(0,),
|
||||
expected_flushed=(0,),
|
||||
)
|
||||
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
|
||||
|
||||
@@ -636,20 +885,22 @@ 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."""
|
||||
offloaded_block_size = 12
|
||||
gpu_block_size = 4
|
||||
block_size_factor = 3
|
||||
offloaded_block_size = gpu_block_size * block_size_factor
|
||||
runner = request_runner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=4,
|
||||
block_size_factor=block_size_factor,
|
||||
block_size=gpu_block_size,
|
||||
num_gpu_blocks=100,
|
||||
async_scheduling=async_scheduling,
|
||||
)
|
||||
runner.new_request(token_ids=[0] * offloaded_block_size)
|
||||
runner.manager.prepare_store.side_effect = (
|
||||
lambda keys, req_context: generate_store_output(keys)
|
||||
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
|
||||
generate_store_output(keys)
|
||||
)
|
||||
|
||||
# First store: fires when block 0 is fully populated.
|
||||
runner.run(decoded_tokens=[0, 0], expected_stored_gpu_block_indexes=(0, 1, 2))
|
||||
runner.run(decoded_tokens=[0, 0], expected_stored=(0, 1, 2))
|
||||
assert runner.manager.complete_store.call_count == 1
|
||||
first_call_keys = set(runner.manager.complete_store.call_args.args[0])
|
||||
assert len(first_call_keys) == 1
|
||||
@@ -658,7 +909,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),
|
||||
expected_stored_gpu_block_indexes=(3, 4, 5),
|
||||
expected_stored=(3, 4, 5),
|
||||
)
|
||||
assert runner.manager.complete_store.call_count == 1
|
||||
second_call_keys = set(runner.manager.complete_store.call_args.args[0])
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from collections.abc import Iterable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -150,46 +150,61 @@ class MockOffloadingSpec(OffloadingSpec):
|
||||
return specs
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GPUBlock:
|
||||
group_idx: int
|
||||
request_block_offset: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransferSummary:
|
||||
gpu_block_indices: list[int]
|
||||
gpu_blocks: list[GPUBlock]
|
||||
offload_addresses: list[Any]
|
||||
|
||||
|
||||
class RequestRunner:
|
||||
def __init__(
|
||||
self,
|
||||
offloaded_block_size: int,
|
||||
gpu_block_size: int,
|
||||
block_size: int,
|
||||
num_gpu_blocks: int,
|
||||
block_size_factor: int = 1,
|
||||
async_scheduling: bool = True,
|
||||
kv_cache_groups: list[KVCacheGroupSpec] | None = None,
|
||||
):
|
||||
self.offloaded_block_size: int = offloaded_block_size
|
||||
self.gpu_block_size: int = gpu_block_size
|
||||
assert block_size_factor == 1 or kv_cache_groups is None, (
|
||||
"block_size_factor > 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.block_size: int = block_size
|
||||
self.num_gpu_blocks: int = num_gpu_blocks
|
||||
self.async_scheduling: bool = async_scheduling
|
||||
|
||||
self.req_id: int = -1
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
block_size=gpu_block_size, max_num_batched_tokens=1000
|
||||
block_size=block_size,
|
||||
max_num_batched_tokens=1000,
|
||||
disable_hybrid_kv_cache_manager=False,
|
||||
)
|
||||
vllm_config.scheduler_config.async_scheduling = async_scheduling
|
||||
|
||||
extra_config: dict[str, Any] = {
|
||||
"spec_name": "MockOffloadingSpec",
|
||||
"spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
|
||||
}
|
||||
if block_size_factor > 1:
|
||||
extra_config["block_size"] = block_size * block_size_factor
|
||||
|
||||
vllm_config.kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="OffloadingConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={
|
||||
"spec_name": "MockOffloadingSpec",
|
||||
"spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
|
||||
"block_size": offloaded_block_size,
|
||||
},
|
||||
kv_connector_extra_config=extra_config,
|
||||
)
|
||||
|
||||
block_size = vllm_config.cache_config.block_size
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_gpu_blocks,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
if kv_cache_groups is None:
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(
|
||||
["layer"],
|
||||
FullAttentionSpec(
|
||||
@@ -199,7 +214,12 @@ class RequestRunner:
|
||||
dtype=torch.float32,
|
||||
),
|
||||
)
|
||||
],
|
||||
]
|
||||
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_gpu_blocks,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
|
||||
self.num_kv_groups = len(kv_cache_config.kv_cache_groups)
|
||||
@@ -219,11 +239,38 @@ class RequestRunner:
|
||||
|
||||
# register worker kv_caches to enable OffloadingWorker creations
|
||||
# set_current_vllm_config is needed for get_kv_cache_layout() to work
|
||||
with set_current_vllm_config(vllm_config):
|
||||
self.worker_connector.register_cross_layers_kv_cache(
|
||||
kv_cache=torch.empty(0),
|
||||
attn_backend=FlashAttentionBackend,
|
||||
)
|
||||
# Mock get_layers_from_vllm_config so that mock layer names
|
||||
# resolve to layers whose get_attn_backend() returns
|
||||
# FlashAttentionBackend.
|
||||
def _mock_get_layers(_vllm_config, _layer_type, layer_names):
|
||||
mock_layer = MagicMock()
|
||||
mock_layer.get_attn_backend.return_value = FlashAttentionBackend
|
||||
return {name: mock_layer for name in layer_names}
|
||||
|
||||
kv_caches: dict[str, torch.Tensor] = {}
|
||||
for group in kv_cache_groups:
|
||||
spec = group.kv_cache_spec
|
||||
for layer_name in group.layer_names:
|
||||
# Shape follows FlashAttention layout:
|
||||
# (2, num_blocks, block_size, num_kv_heads, head_size)
|
||||
kv_caches[layer_name] = torch.empty(
|
||||
2,
|
||||
num_gpu_blocks,
|
||||
spec.block_size,
|
||||
spec.num_kv_heads,
|
||||
spec.head_size,
|
||||
dtype=spec.dtype,
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1"
|
||||
".offloading.worker.get_layers_from_vllm_config",
|
||||
side_effect=_mock_get_layers,
|
||||
),
|
||||
):
|
||||
self.worker_connector.register_kv_caches(kv_caches)
|
||||
|
||||
# extract connector of scheduler
|
||||
scheduler_connector = self.scheduler.connector
|
||||
@@ -238,10 +285,17 @@ class RequestRunner:
|
||||
assert isinstance(manager, MagicMock)
|
||||
self.manager: MagicMock = manager
|
||||
|
||||
assert len(self.connector_scheduler.config.kv_group_configs) == 1
|
||||
kv_group_config = self.connector_scheduler.config.kv_group_configs[0]
|
||||
assert kv_group_config.gpu_block_size == gpu_block_size
|
||||
assert kv_group_config.offloaded_block_size == offloaded_block_size
|
||||
num_kv_groups = len(kv_cache_config.kv_cache_groups)
|
||||
assert len(self.connector_scheduler.config.kv_group_configs) == num_kv_groups
|
||||
for group_config, kv_cache_group in zip(
|
||||
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
|
||||
)
|
||||
|
||||
# extract OffloadingSpec of worker_connector
|
||||
connector_worker = self.worker_connector.connector_worker
|
||||
@@ -250,18 +304,18 @@ class RequestRunner:
|
||||
assert isinstance(offloading_spec, MockOffloadingSpec)
|
||||
self.offloading_spec: MockOffloadingSpec = offloading_spec
|
||||
|
||||
# mapping (offloading address) -> gpu_block_index
|
||||
self.offloaded: dict[Any, int] = {}
|
||||
# mapping (offloading address) -> GPUBlock
|
||||
self.offloaded: dict[Any, GPUBlock] = {}
|
||||
|
||||
self.completed_loads: list[TransferSummary] = []
|
||||
self.completed_stores: list[TransferSummary] = []
|
||||
self.flushed_gpu_block_indexes: set[int] = set()
|
||||
self.flushed_gpu_blocks: set[GPUBlock] = set()
|
||||
|
||||
# maps {block_id: block_offset}
|
||||
self.gpu_block_index: dict[int, int] = {}
|
||||
# block_id -> GPUBlock
|
||||
self.gpu_blocks: dict[int, GPUBlock] = {}
|
||||
|
||||
init_none_hash(sha256)
|
||||
self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
|
||||
self._block_hasher = get_request_block_hasher(block_size, sha256)
|
||||
|
||||
self._dummy_ctx: ForwardContext = ForwardContext(
|
||||
no_compile_layers={},
|
||||
@@ -297,11 +351,9 @@ class RequestRunner:
|
||||
assert isinstance(src_spec, GPULoadStoreSpec)
|
||||
|
||||
for block_id in src_spec.block_ids:
|
||||
self.flushed_gpu_block_indexes.add(
|
||||
self.gpu_block_index[block_id.item()]
|
||||
)
|
||||
self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
|
||||
|
||||
block_size_factor = self.offloaded_block_size // self.gpu_block_size
|
||||
block_size_factor = self.block_size_factor
|
||||
|
||||
for transfer_spec in self.offloading_spec.get_completed_transfers():
|
||||
src_spec, dst_spec = transfer_spec
|
||||
@@ -317,10 +369,11 @@ class RequestRunner:
|
||||
|
||||
assert isinstance(offload_spec, MockLoadStoreSpec)
|
||||
assert isinstance(gpu_spec, GPULoadStoreSpec)
|
||||
assert len(gpu_spec.group_sizes) == self.num_kv_groups
|
||||
|
||||
gpu_block_indices: list[int] = []
|
||||
gpu_blocks: list[GPUBlock] = []
|
||||
for block_id in gpu_spec.block_ids:
|
||||
gpu_block_indices.append(self.gpu_block_index[block_id.item()])
|
||||
gpu_blocks.append(self.gpu_blocks[block_id.item()])
|
||||
|
||||
# list of (offload_key, sub_block_offset)
|
||||
offload_addresses: list[Any] = []
|
||||
@@ -328,30 +381,49 @@ class RequestRunner:
|
||||
for sub_block_idx in range(block_size_factor):
|
||||
offload_addresses.append((offload_key, sub_block_idx))
|
||||
|
||||
assert gpu_spec.block_indices is not None
|
||||
assert len(gpu_spec.block_indices) == self.num_kv_groups
|
||||
|
||||
gpu_block_offset = 0
|
||||
offload_address_offset = 0
|
||||
for group_size, logical_offset in zip(
|
||||
gpu_spec.group_sizes, gpu_spec.block_indices
|
||||
):
|
||||
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_end_offset = (
|
||||
offload_address_offset + offload_addresses_to_skip + group_size
|
||||
)
|
||||
assert offload_addresses_end_offset <= len(offload_addresses)
|
||||
|
||||
offload_addresses = (
|
||||
offload_addresses[:offload_address_offset]
|
||||
+ offload_addresses[
|
||||
offload_address_offset + offload_addresses_to_skip :
|
||||
]
|
||||
)
|
||||
|
||||
gpu_block_offset += group_size
|
||||
offload_address_offset += group_size
|
||||
|
||||
assert gpu_block_offset == len(gpu_blocks)
|
||||
assert offload_address_offset == len(offload_addresses)
|
||||
|
||||
transfer_summary = TransferSummary(gpu_blocks, offload_addresses)
|
||||
if store:
|
||||
assert len(gpu_block_indices) == len(offload_addresses)
|
||||
|
||||
self.completed_stores.append(
|
||||
TransferSummary(gpu_block_indices, offload_addresses)
|
||||
)
|
||||
self.completed_stores.append(transfer_summary)
|
||||
else:
|
||||
remainder_sub_block_count = len(offload_addresses) - len(
|
||||
gpu_block_indices
|
||||
)
|
||||
assert remainder_sub_block_count >= 0
|
||||
assert remainder_sub_block_count < block_size_factor
|
||||
offload_addresses = offload_addresses[remainder_sub_block_count:]
|
||||
self.completed_loads.append(transfer_summary)
|
||||
|
||||
self.completed_loads.append(
|
||||
TransferSummary(gpu_block_indices, offload_addresses)
|
||||
)
|
||||
|
||||
def _update_gpu_block_idx(self):
|
||||
for blocks in self.scheduler.kv_cache_manager.coordinator.single_type_managers[
|
||||
0
|
||||
].req_to_blocks.values():
|
||||
for block_idx, block in enumerate(blocks):
|
||||
self.gpu_block_index[block.block_id] = block_idx
|
||||
def _update_gpu_blocks(self):
|
||||
for group_idx, manager in enumerate(
|
||||
self.scheduler.kv_cache_manager.coordinator.single_type_managers
|
||||
):
|
||||
for blocks in manager.req_to_blocks.values():
|
||||
for block_idx, block in enumerate(blocks):
|
||||
self.gpu_blocks[block.block_id] = GPUBlock(group_idx, block_idx)
|
||||
|
||||
def _run(self, decoded_tokens: list[int], complete_transfers: bool):
|
||||
"""
|
||||
@@ -376,7 +448,7 @@ class RequestRunner:
|
||||
break
|
||||
|
||||
scheduler_output = self.scheduler.schedule()
|
||||
self._update_gpu_block_idx()
|
||||
self._update_gpu_blocks()
|
||||
|
||||
kv_connector_metadata = scheduler_output.kv_connector_metadata
|
||||
assert kv_connector_metadata is not None
|
||||
@@ -447,13 +519,30 @@ class RequestRunner:
|
||||
if EOS_TOKEN_ID in decoded_tokens:
|
||||
assert not self.scheduler.running
|
||||
|
||||
def _to_gpu_blocks(
|
||||
self, blocks: tuple[int | tuple[int, int], ...]
|
||||
) -> list[GPUBlock]:
|
||||
gpu_blocks: list[GPUBlock] = []
|
||||
for block in blocks:
|
||||
if isinstance(block, int):
|
||||
for group_idx in range(self.num_kv_groups):
|
||||
gpu_blocks.append(
|
||||
GPUBlock(group_idx=group_idx, request_block_offset=block)
|
||||
)
|
||||
else:
|
||||
group_idx, offset = block
|
||||
gpu_blocks.append(
|
||||
GPUBlock(group_idx=group_idx, request_block_offset=offset)
|
||||
)
|
||||
return gpu_blocks
|
||||
|
||||
def run(
|
||||
self,
|
||||
decoded_tokens: list[int],
|
||||
complete_transfers: bool = True,
|
||||
expected_stored_gpu_block_indexes: tuple[int, ...] = (),
|
||||
expected_loaded_gpu_block_indexes: tuple[int, ...] = (),
|
||||
expected_flushed_gpu_block_indexes: tuple[int, ...] = (),
|
||||
expected_stored: tuple[int | tuple[int, int], ...] = (),
|
||||
expected_loaded: tuple[int | tuple[int, int], ...] = (),
|
||||
expected_flushed: tuple[int | tuple[int, int], ...] = (),
|
||||
):
|
||||
"""
|
||||
Runs multiple engine (scheduler + worker) steps.
|
||||
@@ -462,41 +551,49 @@ class RequestRunner:
|
||||
Args:
|
||||
decoded_tokens: the tokens to yield at each step.
|
||||
complete_transfers: complete transfers immediately
|
||||
expected_stored_gpu_block_indexes: GPU block indexes
|
||||
expected_stored: GPU blocks
|
||||
that are expected to be written during the run.
|
||||
expected_loaded_gpu_block_indexes: GPU block indexes
|
||||
expected_loaded: GPU blocks
|
||||
that are expected to be loaded during the run.
|
||||
expected_flushed_gpu_block_indexes: GPU block indexes
|
||||
expected_flushed: GPU blocks
|
||||
that are expected to be flushed during the run.
|
||||
|
||||
A GPU block is either a (group_idx: int, request_block_offset: int)
|
||||
or just request_block_offset: int.
|
||||
The latter case is a convenience for representing all groups.
|
||||
"""
|
||||
|
||||
expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored)
|
||||
expected_loaded_gpu_blocks = self._to_gpu_blocks(expected_loaded)
|
||||
expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed)
|
||||
|
||||
self.manager.reset_mock()
|
||||
self._run(decoded_tokens, complete_transfers)
|
||||
|
||||
loaded_gpu_block_indexes: set[int] = set()
|
||||
loaded_gpu_blocks: set[GPUBlock] = set()
|
||||
for transfer in self.completed_loads:
|
||||
for gpu_block_idx, offloaded_address in zip(
|
||||
transfer.gpu_block_indices, transfer.offload_addresses
|
||||
for gpu_block, offloaded_address in zip(
|
||||
transfer.gpu_blocks, transfer.offload_addresses
|
||||
):
|
||||
loaded_gpu_block_indexes.add(gpu_block_idx)
|
||||
assert gpu_block_idx == self.offloaded[offloaded_address]
|
||||
loaded_gpu_blocks.add(gpu_block)
|
||||
assert gpu_block == self.offloaded[offloaded_address]
|
||||
|
||||
assert set(expected_loaded_gpu_block_indexes) == loaded_gpu_block_indexes
|
||||
assert set(expected_loaded_gpu_blocks) == loaded_gpu_blocks
|
||||
self.completed_loads.clear()
|
||||
|
||||
stored_gpu_block_indexes: set[int] = set()
|
||||
stored_gpu_blocks: set[GPUBlock] = set()
|
||||
for transfer in self.completed_stores:
|
||||
for gpu_block_idx, offloaded_address in zip(
|
||||
transfer.gpu_block_indices, transfer.offload_addresses
|
||||
for gpu_block, offloaded_address in zip(
|
||||
transfer.gpu_blocks, transfer.offload_addresses
|
||||
):
|
||||
stored_gpu_block_indexes.add(gpu_block_idx)
|
||||
self.offloaded[offloaded_address] = gpu_block_idx
|
||||
stored_gpu_blocks.add(gpu_block)
|
||||
self.offloaded[offloaded_address] = gpu_block
|
||||
|
||||
assert set(expected_stored_gpu_block_indexes) == stored_gpu_block_indexes
|
||||
assert set(expected_stored_gpu_blocks) == stored_gpu_blocks
|
||||
self.completed_stores.clear()
|
||||
|
||||
assert set(expected_flushed_gpu_block_indexes) == self.flushed_gpu_block_indexes
|
||||
self.flushed_gpu_block_indexes.clear()
|
||||
assert set(expected_flushed_gpu_blocks) == self.flushed_gpu_blocks
|
||||
self.flushed_gpu_blocks.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -504,13 +601,18 @@ def request_runner():
|
||||
runners = []
|
||||
|
||||
def runner_factory(
|
||||
offloaded_block_size, gpu_block_size, num_gpu_blocks, async_scheduling
|
||||
block_size,
|
||||
num_gpu_blocks,
|
||||
async_scheduling,
|
||||
block_size_factor=1,
|
||||
kv_cache_groups=None,
|
||||
):
|
||||
runner = RequestRunner(
|
||||
offloaded_block_size=offloaded_block_size,
|
||||
gpu_block_size=gpu_block_size,
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=num_gpu_blocks,
|
||||
block_size_factor=block_size_factor,
|
||||
async_scheduling=async_scheduling,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
runners.append(runner)
|
||||
return runner
|
||||
|
||||
@@ -14,13 +14,36 @@ from vllm.config import KVEventsConfig, KVTransferConfig
|
||||
from vllm.distributed.kv_events import BlockStored, KVEventBatch
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
CPU_BLOCK_SIZES = [48]
|
||||
ATTN_BACKENDS = []
|
||||
|
||||
_ATTN_BACKENDS: list[str] = []
|
||||
if current_platform.is_cuda():
|
||||
ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
|
||||
_ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
|
||||
elif current_platform.is_rocm():
|
||||
ATTN_BACKENDS = ["TRITON_ATTN"]
|
||||
_ATTN_BACKENDS = ["TRITON_ATTN"]
|
||||
|
||||
# (model, attn_backend | None, block_size | None, uses_hma)
|
||||
#
|
||||
# - Llama: tested with each attention backend and a custom block_size.
|
||||
# - Gemma-3: HMA (sliding window + full attention), default backend.
|
||||
# - Mamba-130m: HMA (attention-free, acts like sliding_window=1),
|
||||
# default backend. Prefix caching must be force-enabled.
|
||||
# - Falcon-H1-0.5B-Instruct: HMA (parallel SSM/attention in every layer).
|
||||
# After page-size unification the mamba and attention groups have
|
||||
# different block sizes.
|
||||
MODEL_PARAMS: list[tuple[str, str | None, int | None, bool]] = [
|
||||
("meta-llama/Llama-3.2-1B-Instruct", backend, 48, False)
|
||||
for backend in _ATTN_BACKENDS
|
||||
]
|
||||
# HMA / Mamba models are only tested on CUDA (not ROCm).
|
||||
if current_platform.is_cuda():
|
||||
MODEL_PARAMS += [
|
||||
("google/gemma-3-1b-it", None, 48, True),
|
||||
("state-spaces/mamba-130m-hf", None, 48, True),
|
||||
# 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).
|
||||
("tiiuae/Falcon-H1-0.5B-Instruct", None, None, True),
|
||||
]
|
||||
|
||||
# Maximum time (seconds) to wait for the async CPU offload transfer
|
||||
# to complete before giving up.
|
||||
@@ -110,7 +133,7 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _latency_test(llm: LLM, subscriber: MockSubscriber):
|
||||
def _latency_test(llm: LLM, subscriber: MockSubscriber | None):
|
||||
sampling_params = SamplingParams(max_tokens=1)
|
||||
|
||||
num_times_cpu_better_than_cold = 0
|
||||
@@ -118,7 +141,10 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
|
||||
total_cold_time = 0.0
|
||||
total_gpu_hit_time = 0.0
|
||||
total_cpu_hit_time = 0.0
|
||||
prompt_token_ids = [0] * 10001
|
||||
max_model_len = llm.llm_engine.vllm_config.model_config.max_model_len
|
||||
# Use a long prompt that fits within the model's context window.
|
||||
prompt_len = min(10001, max_model_len - 1)
|
||||
prompt_token_ids = [0] * prompt_len
|
||||
for i in tqdm(range(num_tests), desc="Running tests"):
|
||||
prompt_token_ids[0] = i
|
||||
prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
|
||||
@@ -141,10 +167,11 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
|
||||
|
||||
# Verify CPU stored events arrived (offload is done before we
|
||||
# attempt to load from CPU).
|
||||
assert subscriber.get_new_cpu_stored_events(), (
|
||||
f"No CPU stored events received on iteration {i}; "
|
||||
"async offload may not have completed in time"
|
||||
)
|
||||
if subscriber is not None:
|
||||
assert subscriber.get_new_cpu_stored_events(), (
|
||||
f"No CPU stored events received on iteration {i}; "
|
||||
"async offload may not have completed in time"
|
||||
)
|
||||
|
||||
# run generation again - this should trigger loading from CPU
|
||||
start_time = time.time()
|
||||
@@ -163,84 +190,105 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
|
||||
assert num_times_cpu_better_than_cold >= 0.8 * num_tests
|
||||
|
||||
|
||||
def _accuracy_test(llm: LLM, subscriber: MockSubscriber):
|
||||
def _accuracy_test(llm: LLM, subscriber: MockSubscriber | None):
|
||||
sampling_params = SamplingParams(max_tokens=1)
|
||||
cpu_block_size = (
|
||||
llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config[
|
||||
"block_size"
|
||||
]
|
||||
extra_config = (
|
||||
llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config
|
||||
)
|
||||
cpu_block_size = extra_config.get("block_size")
|
||||
if cpu_block_size is None:
|
||||
# No custom offloaded block_size: offloaded blocks match GPU blocks.
|
||||
# Use the hash block_size (cache_config.block_size) for alignment.
|
||||
cpu_block_size = llm.llm_engine.vllm_config.cache_config.block_size
|
||||
|
||||
subscriber.get_new_cpu_stored_events()
|
||||
if subscriber is not None:
|
||||
subscriber.get_new_cpu_stored_events()
|
||||
|
||||
# prepend prompt to be cpu block aligned
|
||||
# Pad prompt so its token count is a multiple of cpu_block_size.
|
||||
# Use the tokenizer directly to avoid expensive llm.generate() calls.
|
||||
tokenizer = llm.get_tokenizer()
|
||||
prompt = "Let's count to 10. One, two, three, four,"
|
||||
while (
|
||||
len(llm.generate(prompt, use_tqdm=False)[0].prompt_token_ids) % cpu_block_size
|
||||
!= 0
|
||||
):
|
||||
while len(tokenizer.encode(prompt)) % cpu_block_size != 0:
|
||||
prompt = ". " + prompt
|
||||
|
||||
assert subscriber.get_new_cpu_stored_events()
|
||||
# Seed the CPU cache with the prompt.
|
||||
llm.generate(prompt, sampling_params, use_tqdm=False)
|
||||
|
||||
test_count = 100
|
||||
success_count = 0
|
||||
for i in range(test_count):
|
||||
if (
|
||||
llm.generate(prompt, sampling_params, use_tqdm=False)[0].outputs[0].text
|
||||
== " five"
|
||||
):
|
||||
success_count += 1
|
||||
if subscriber is not None:
|
||||
assert subscriber.get_new_cpu_stored_events()
|
||||
|
||||
test_count = 20
|
||||
results = llm.generate([prompt] * test_count, sampling_params, use_tqdm=False)
|
||||
success_count = sum(1 for r in results if r.outputs[0].text == " five")
|
||||
assert success_count >= 0.5 * test_count
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
|
||||
def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
|
||||
@pytest.mark.parametrize("model, attn_backend, cpu_block_size, uses_hma", MODEL_PARAMS)
|
||||
def test_cpu_offloading(
|
||||
model: str,
|
||||
attn_backend: str | None,
|
||||
cpu_block_size: int | None,
|
||||
uses_hma: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Tests OffloadingConnector with CPUOffloadingSpec.
|
||||
"""
|
||||
|
||||
# configure OffloadingConnector (spec_name=CPUOffloadingSpec by default)
|
||||
extra_config: dict = {"cpu_bytes_to_use": 500 << 20}
|
||||
if cpu_block_size is not None:
|
||||
extra_config["block_size"] = cpu_block_size
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="OffloadingConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={
|
||||
"cpu_bytes_to_use": 500 << 20,
|
||||
"block_size": cpu_block_size,
|
||||
},
|
||||
kv_connector_extra_config=extra_config,
|
||||
)
|
||||
|
||||
port: int
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("0.0.0.0", 0))
|
||||
port = s.getsockname()[1]
|
||||
# KV events are incompatible with HMA (setting kv_events_config
|
||||
# would force HMA off), so only enable them for non-HMA models.
|
||||
subscriber: MockSubscriber | None = None
|
||||
kv_events_config: KVEventsConfig | None = None
|
||||
if not uses_hma:
|
||||
port: int
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("0.0.0.0", 0))
|
||||
port = s.getsockname()[1]
|
||||
|
||||
events_endpoint = f"tcp://*:{port}"
|
||||
kv_events_config = KVEventsConfig(
|
||||
enable_kv_cache_events=True,
|
||||
publisher="zmq",
|
||||
endpoint=events_endpoint,
|
||||
topic="test",
|
||||
)
|
||||
events_endpoint = f"tcp://*:{port}"
|
||||
kv_events_config = KVEventsConfig(
|
||||
enable_kv_cache_events=True,
|
||||
publisher="zmq",
|
||||
endpoint=events_endpoint,
|
||||
topic="test",
|
||||
)
|
||||
|
||||
# Attention-free / hybrid models disable prefix caching by default
|
||||
# (ModelConfig.is_prefix_caching_supported returns False). Without it,
|
||||
# mamba_block_size falls back to max_model_len, making GPU blocks too
|
||||
# large for any reasonable offloaded block_size. Force-enable it.
|
||||
force_prefix_caching = uses_hma
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Llama-3.2-1B-Instruct",
|
||||
model=model,
|
||||
max_model_len=4096,
|
||||
gpu_memory_utilization=0.5,
|
||||
kv_events_config=kv_events_config,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
attention_config={"backend": attn_backend},
|
||||
**({"attention_config": {"backend": attn_backend}} if attn_backend else {}),
|
||||
# HMA models need explicit opt-in when kv_transfer_config is set
|
||||
**({"disable_hybrid_kv_cache_manager": False} if uses_hma else {}),
|
||||
**({"enable_prefix_caching": True} if force_prefix_caching else {}),
|
||||
# ROCm: batch size 1 to reduce variability
|
||||
**({"max_num_seqs": 1} if current_platform.is_rocm() else {}),
|
||||
)
|
||||
|
||||
events_endpoint = events_endpoint.replace("*", "127.0.0.1")
|
||||
subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
|
||||
if kv_events_config is not None:
|
||||
events_endpoint = events_endpoint.replace("*", "127.0.0.1")
|
||||
subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
|
||||
|
||||
try:
|
||||
_latency_test(llm, subscriber)
|
||||
_accuracy_test(llm, subscriber)
|
||||
finally:
|
||||
subscriber.close()
|
||||
if subscriber is not None:
|
||||
subscriber.close()
|
||||
del llm
|
||||
|
||||
@@ -103,6 +103,7 @@ def create_vllm_config(
|
||||
kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
|
||||
kv_connector: str = "NixlConnector",
|
||||
kv_role: str = "kv_both",
|
||||
disable_hybrid_kv_cache_manager: bool | None = None,
|
||||
) -> VllmConfig:
|
||||
"""Initialize VllmConfig For Testing."""
|
||||
model_config = ModelConfig(
|
||||
@@ -118,6 +119,7 @@ def create_vllm_config(
|
||||
max_model_len=max_model_len,
|
||||
enable_chunked_prefill=enable_chunked_prefill,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
|
||||
)
|
||||
# Cache config, optionally force APC
|
||||
cache_config = CacheConfig(
|
||||
|
||||
@@ -829,7 +829,6 @@ class OffloadingConnectorScheduler:
|
||||
def request_finished(
|
||||
self,
|
||||
request: Request,
|
||||
block_ids: list[int],
|
||||
) -> tuple[bool, dict[str, Any] | None]:
|
||||
"""
|
||||
Called when a request has finished, before its blocks are freed.
|
||||
|
||||
@@ -10,6 +10,7 @@ from vllm.distributed.kv_events import KVCacheEvent
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1 import (
|
||||
KVConnectorBase_V1,
|
||||
KVConnectorRole,
|
||||
SupportsHMA,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
|
||||
@@ -42,7 +43,7 @@ from vllm.v1.outputs import KVConnectorOutput
|
||||
from vllm.v1.request import Request
|
||||
|
||||
|
||||
class OffloadingConnector(KVConnectorBase_V1):
|
||||
class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
|
||||
@property
|
||||
def prefer_cross_layer_blocks(self) -> bool:
|
||||
return True
|
||||
@@ -149,7 +150,15 @@ class OffloadingConnector(KVConnectorBase_V1):
|
||||
block_ids: list[int],
|
||||
) -> tuple[bool, dict[str, Any] | None]:
|
||||
assert self.connector_scheduler is not None
|
||||
return self.connector_scheduler.request_finished(request, block_ids)
|
||||
return self.connector_scheduler.request_finished(request)
|
||||
|
||||
def request_finished_all_groups(
|
||||
self,
|
||||
request: "Request",
|
||||
block_ids: tuple[list[int], ...],
|
||||
) -> tuple[bool, dict[str, Any] | None]:
|
||||
assert self.connector_scheduler is not None
|
||||
return self.connector_scheduler.request_finished(request)
|
||||
|
||||
def take_events(self) -> Iterable[KVCacheEvent]:
|
||||
assert self.connector_scheduler is not None
|
||||
|
||||
Reference in New Issue
Block a user