Skip to content

[Hybrid Allocator] Support KV cache groups with different block_size #24949

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 20 commits into
base: main
Choose a base branch
from
Open

[Hybrid Allocator] Support KV cache groups with different block_size #24949

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
23530a4
support different block size
heheda12345 Sep 16, 2025
4b3acea
fix bug
heheda12345 Sep 16, 2025
f03bb61
fix tests
heheda12345 Sep 16, 2025
f061d83
revert
heheda12345 Sep 16, 2025
48a70ae
rename BlockHashListWithBlockSize
heheda12345 Sep 16, 2025
286aeac
fix comments
heheda12345 Sep 16, 2025
5be0625
fix comments
heheda12345 Sep 16, 2025
4ad6559
fix comments
heheda12345 Sep 16, 2025
10df50e
remove unnecessary logic
heheda12345 Sep 16, 2025
aaf8bc9
support block_size alignment
heheda12345 Sep 16, 2025
bab0be5
add test
heheda12345 Sep 16, 2025
20a96f5
Merge commit '17edd8a' into two_dtype_kv_cache
heheda12345 Oct 6, 2025
ff3a21b
ruff
heheda12345 Oct 6, 2025
f988a71
Merge commit 'd6953be' into two_dtype_kv_cache
heheda12345 Oct 6, 2025
2ea1cca
Merge branch 'main' of github.com:vllm-project/vllm into two_dtype_kv…
heheda12345 Oct 6, 2025
b8d368f
fix test
heheda12345 Oct 6, 2025
0d4fad3
remove yapf tag
heheda12345 Oct 6, 2025
856cc74
support block_size alignment
heheda12345 Oct 6, 2025
0c20bc2
reduce diff
heheda12345 Oct 6, 2025
1fe5821
tmp fix on dcp
heheda12345 Oct 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 44 additions & 4 deletions tests/v1/core/test_kv_cache_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1179,7 +1179,9 @@ def test_allocate_with_lookahead():
)

# Test case 1: Requires additional lookahead tokens
kv_cache_manager = KVCacheManager(kv_cache_config=config, max_model_len=100)
kv_cache_manager = KVCacheManager(
kv_cache_config=config, max_model_len=100, hash_block_size=block_size
)
blocks = kv_cache_manager.allocate_slots(
request,
num_new_tokens=3,
Expand All @@ -1188,7 +1190,9 @@ def test_allocate_with_lookahead():
assert len(blocks.get_block_ids()[0]) == 2 # ceil(5/4)=2 blocks

# Test case 2: With precomputed blocks
kv_cache_manager = KVCacheManager(kv_cache_config=config, max_model_len=100)
kv_cache_manager = KVCacheManager(
kv_cache_config=config, max_model_len=100, hash_block_size=block_size
)
# required_blocks = ceil((3 + 2) /4) = 2
blocks = kv_cache_manager.allocate_slots(
request,
Expand All @@ -1199,7 +1203,9 @@ def test_allocate_with_lookahead():

# Test case 3: With precomputed blocks
# required_blocks = ceil((3 + 4) / 4) = 2
kv_cache_manager = KVCacheManager(kv_cache_config=config, max_model_len=100)
kv_cache_manager = KVCacheManager(
kv_cache_config=config, max_model_len=100, hash_block_size=block_size
)
blocks = kv_cache_manager.allocate_slots(
request,
num_new_tokens=3,
Expand Down Expand Up @@ -1367,7 +1373,7 @@ def test_get_kv_cache_config_one_worker():
],
)

# different hidden size
# different hidden size but same type, use UniformTypeKVCacheSpecs
kv_cache_specs_hybrid = {
"layer_1": new_kv_cache_spec(head_size=128),
"layer_2": new_kv_cache_spec(head_size=64),
Expand All @@ -1391,6 +1397,40 @@ def test_get_kv_cache_config_one_worker():
],
)

# Different hidden size and different type, align by different block size
kv_cache_specs_hybrid = {
"layer_1": new_kv_cache_spec(head_size=64),
"layer_2": new_sliding_window_spec(head_size=32),
}
kv_cache_config_hybrid = get_kv_cache_configs(
vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32]
)[0]
assert kv_cache_config_hybrid == KVCacheConfig(
num_blocks=32,
kv_cache_tensors=[
KVCacheTensor(
size=mem_per_block_per_layer * 32, shared_by=["layer_1", "layer_2"]
),
],
kv_cache_groups=[
KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)),
KVCacheGroupSpec(
["layer_2"], new_sliding_window_spec(head_size=32, block_size=32)
),
],
)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you mind adding a test for the mixed dtype case?

    # Different dtype, align by using different block size
    kv_cache_specs_hybrid = {
        'layer_1': new_kv_cache_spec(dtype=torch.float8_e4m3fn),
        'layer_2': new_sliding_window_spec(dtype=torch.bfloat16),
    }
    kv_cache_config_hybrid = get_kv_cache_configs(
        vllm_config, [kv_cache_specs_hybrid],
        [mem_per_block_per_layer * 32])[0]
    assert kv_cache_config_hybrid == KVCacheConfig(
        num_blocks=32 * 2, # 2x blocks because baseline is BF16 (not FP32)
        kv_cache_tensors=[
            KVCacheTensor(size=mem_per_block_per_layer * 32,
                          shared_by=["layer_1", "layer_2"]),
        ],
        kv_cache_groups=[
            KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(dtype=torch.float8_e4m3fn, block_size=32)),
            KVCacheGroupSpec(["layer_2"],
                             new_sliding_window_spec(dtype=torch.bfloat16,
                                                     block_size=16)),
        ],
    )

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

similarly could use new_kv_cache_spec, as it is nothing speficif to new_sliding_window_spec I'd say.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you mind adding a test for the mixed dtype case?

I think there is no difference on mixed dtype & mixed head size from the view of this PR. Feel free to add tests when you are working on mixed dtype support.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

similarly could use new_kv_cache_spec, as it is nothing speficif to new_sliding_window_spec I'd say.

For models only with full attention, we can have a much simpler path because we don't need to ensure all layers have the same page_size_bytes. I'm working on it in another PR.

# different hidden size that cannot be aligned by using different block size
kv_cache_specs_hybrid = {
"layer_1": new_kv_cache_spec(head_size=64),
"layer_2": new_sliding_window_spec(head_size=96),
}

with pytest.raises(NotImplementedError):
get_kv_cache_configs(
vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
)[0]

# Test num_gpu_blocks_override
vllm_config.cache_config.num_gpu_blocks_override = 16
kv_cache_config_override_blocks = get_kv_cache_configs(
Expand Down
91 changes: 89 additions & 2 deletions tests/v1/core/test_prefix_caching.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ def test_prefill(hash_fn):
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

# Complete 3 blocks (48 tokens)
Expand Down Expand Up @@ -254,6 +255,7 @@ def test_prefill_hybrid_model():
make_kv_cache_config_hybrid_model(block_size, 21),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

hash_fn = sha256
Expand Down Expand Up @@ -414,6 +416,7 @@ def test_prefill_plp():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)
# the default hash function is sha256
hash_fn = sha256
Expand Down Expand Up @@ -521,6 +524,7 @@ def test_decode():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

# Complete 3 blocks (48 tokens)
Expand Down Expand Up @@ -583,6 +587,7 @@ def test_evict():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

last_token_id = 5 * 16 + 7
Expand Down Expand Up @@ -641,6 +646,7 @@ def test_hash_block_correct_reuse():
make_kv_cache_config(16, 2),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

# Allocate 1 block and cache it.
Expand Down Expand Up @@ -681,6 +687,7 @@ def test_computed_blocks_not_evicted():
make_kv_cache_config(block_size, 3),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

# Allocate a block and cache it.
Expand Down Expand Up @@ -739,6 +746,7 @@ def test_basic_prefix_caching_disabled():
make_kv_cache_config(block_size, 5),
max_model_len=8192,
enable_caching=False,
hash_block_size=block_size,
)

req1 = make_request(
Expand Down Expand Up @@ -788,6 +796,7 @@ def test_cache_blocks(hash_fn):
block_pool = BlockPool(
num_gpu_blocks=5,
enable_caching=True,
hash_block_size=block_size,
)
# Req:
# Block 0: [0, 1, 2, 3]
Expand Down Expand Up @@ -831,7 +840,9 @@ def test_cache_blocks_multi_group():
This tests that blocks are cached correctly for different kv cache groups.
"""
block_size = 4
block_pool = BlockPool(num_gpu_blocks=10, enable_caching=True)
block_pool = BlockPool(
num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size
)

# Req:
# Block 0/4: [0, 1, 2, 3]
Expand Down Expand Up @@ -919,6 +930,7 @@ def test_mm_prefix_caching():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

# Common prompt tokens (T is text tokens and P is image placeholder tokens)
Expand Down Expand Up @@ -1018,6 +1030,7 @@ def test_cache_key_salting():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

# 3 complete blocks and an incomplete block with 11 tokens.
Expand Down Expand Up @@ -1099,6 +1112,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)
# Complete 3 blocks (48 tokens)
# | Common-0 | Common-1 | Common-2 | ... |
Expand Down Expand Up @@ -1171,6 +1185,7 @@ def test_reset_prefix_cache():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

full_block_token_ids = [i for i in range(3) for _ in range(16)]
Expand Down Expand Up @@ -1211,6 +1226,7 @@ def test_prefix_cache_stats_disabled():
make_kv_cache_config(block_size, 11),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
log_stats=False, # Disable logging stats
)
assert manager.prefix_cache_stats is None
Expand All @@ -1230,7 +1246,7 @@ def test_prefix_cache_stats_disabled():


def test_maybe_evict_cached_block():
pool = BlockPool(num_gpu_blocks=4, enable_caching=True)
pool = BlockPool(num_gpu_blocks=4, enable_caching=True, hash_block_size=16)
block_hash0 = make_block_hash_with_group_id(BlockHash(b"10"), 1000)
block_hash1 = make_block_hash_with_group_id(BlockHash(b"20"), 2000)
block_hash2 = make_block_hash_with_group_id(BlockHash(b"30"), 3000)
Expand Down Expand Up @@ -1291,6 +1307,7 @@ def test_kv_cache_events(blocks_to_cache: int):
max_model_len=8192,
enable_caching=True,
enable_kv_cache_events=True,
hash_block_size=block_size,
)

num_tokens = block_size * blocks_to_cache
Expand Down Expand Up @@ -1346,6 +1363,7 @@ def test_eagle_enabled_removes_last_block():
max_model_len=8192,
enable_caching=True,
use_eagle=True,
hash_block_size=block_size,
)

# Request with 3 full blocks (48 tokens)
Expand Down Expand Up @@ -1378,6 +1396,7 @@ def test_eagle_with_partial_blocks():
max_model_len=8192,
enable_caching=True,
use_eagle=True,
hash_block_size=block_size,
)
# 2 full blocks + 5 tokens (non-divisible length)
token_ids = [0] * (2 * block_size + 5)
Expand Down Expand Up @@ -1417,6 +1436,7 @@ def test_eagle_with_sliding_window():
max_model_len=8192,
enable_caching=True,
use_eagle=True,
hash_block_size=block_size,
)

# 2 full blocks + 5 tokens (non-divisible length)
Expand Down Expand Up @@ -1463,6 +1483,73 @@ def test_eagle_with_sliding_window():
assert num_tokens == 0


def test_different_block_size():
block_size = 16
kv_cache_config = KVCacheConfig(
num_blocks=100,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["layer1"],
FullAttentionSpec(block_size * 2, 1, 1, torch.float32),
),
KVCacheGroupSpec(
["layer2"],
SlidingWindowSpec(
block_size,
1,
1,
torch.float32,
sliding_window=2 * block_size,
),
),
],
)
manager = KVCacheManager(
kv_cache_config=kv_cache_config,
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
)

common_token_ids = [i for i in range(10) for _ in range(block_size)]

req0 = make_request("0", common_token_ids, block_size, sha256)
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
assert not computed_blocks.blocks[0]
assert not computed_blocks.blocks[1]
assert num_computed_tokens == 0
blocks = manager.allocate_slots(
req0, 7 * block_size, len(computed_blocks.blocks[0]) * 16, computed_blocks
)
assert blocks.get_block_ids() == ([1, 2, 3, 4], [5, 6, 7, 8, 9, 10, 11])
req1 = make_request("1", common_token_ids[: 7 * block_size + 1], block_size, sha256)
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
assert len(computed_blocks.blocks[0]) == 3
assert len(computed_blocks.blocks[1]) == 6
assert num_computed_tokens == 6 * 16

req2 = make_request("2", common_token_ids[: 6 * block_size + 1], block_size, sha256)
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
assert len(computed_blocks.blocks[0]) == 3
assert len(computed_blocks.blocks[1]) == 6
assert num_computed_tokens == 6 * 16

# Evict some blocks to make sliding window cache hit length 5*16
# But should return 4 * 16 because full attention cache hit length must be
# a multiple of 32
manager.block_pool.cached_block_hash_to_block.pop(
make_block_hash_with_group_id(req1.block_hashes[6], 1), 11
)
manager.block_pool.cached_block_hash_to_block.pop(
make_block_hash_with_group_id(req1.block_hashes[5], 1), 10
)
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
assert len(computed_blocks.blocks[0]) == 2
assert len(computed_blocks.blocks[1]) == 4
assert num_computed_tokens == 4 * 16


def test_block_lookup_cache_single_block_per_key():
cache = BlockHashToBlockMap()
key0 = BlockHashWithGroupId(b"hash0")
Expand Down
20 changes: 14 additions & 6 deletions tests/v1/core/test_single_type_kv_cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,9 @@ def test_chunked_local_attention_possible_cached_prefix():
attention_chunk_size=4,
)

block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
block_pool = BlockPool(
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
)
manager = get_chunked_local_attention_manager(
chunked_local_attention_spec, block_pool
)
Expand Down Expand Up @@ -111,7 +113,9 @@ def test_sliding_window_possible_cached_prefix():
sliding_window=4,
)

block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
block_pool = BlockPool(
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
)
manager = get_sliding_window_manager(sliding_window_spec, block_pool)

def run_one_case(block_is_cached, expect_length):
Expand Down Expand Up @@ -178,7 +182,7 @@ def test_chunked_local_attention_remove_skipped_blocks():
attention_chunk_size=4,
)

block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=2)

manager = get_chunked_local_attention_manager(attention_spec, block_pool)

Expand Down Expand Up @@ -239,7 +243,7 @@ def test_sliding_window_remove_skipped_blocks():
sliding_window=4,
)

block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=2)

manager = get_sliding_window_manager(sliding_window_spec, block_pool)

Expand Down Expand Up @@ -316,7 +320,9 @@ def test_get_num_blocks_to_allocate():
sliding_window=4, # Placeholder value, not related to test result
)

block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
block_pool = BlockPool(
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
)
manager = get_sliding_window_manager(sliding_window_spec, block_pool)
cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
cached_blocks_2 = [block_pool.null_block for _ in range(5)] + [
Expand All @@ -341,7 +347,9 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
attention_chunk_size=4, # Placeholder value, not related to test result
)

block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
block_pool = BlockPool(
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
)
manager = get_chunked_local_attention_manager(attention_spec, block_pool)
cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
cached_blocks_2 = [block_pool.null_block for _ in range(5)] + [
Expand Down
4 changes: 4 additions & 0 deletions vllm/attention/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,10 @@ def __init__(
f"num_heads ({num_heads}) is not divisible by num_kv_heads ({num_kv_heads})"
)

# TODO in this PR: only for testing now. remove this hardcode later
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

self reminder: remove this

if sliding_window is None:
print("set kv_cache_dtype to fp8_e4m3 for layer", prefix)
kv_cache_dtype = "fp8_e4m3"
# The default k/v_scale is set to 1.0. This is ignored
# when kv-cache is not fp8, and should be used with
# kv-cache in fp8_e5m2. For kv-cache in fp8_e4m3, we
Expand Down
Loading