biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V0 Deprecation] Refactor kv cache from list to element (#37487)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-03-23 23:10:11 -04:00
committed by GitHub
parent de99d91ece
commit c59a132f96
27 changed files with 70 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/compile/passes/test_fusion_attn.py
View File

@@ -127,7 +127,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
raw_tensor = raw_tensor.view(kv_cache_shape)
kv_cache = raw_tensor.permute(*inv_order)
self.attn.kv_cache = [kv_cache]
self.attn.kv_cache = kv_cache
# Build attn metadata
self.attn_metadata = self.builder.build(

6
tests/compile/passes/test_rope_kvcache_fusion.py
View File

@@ -148,7 +148,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
raw_tensor = raw_tensor.view(kv_cache_shape)
kv_cache = raw_tensor.permute(*inv_order)
self.attn.kv_cache = [kv_cache]
self.attn.kv_cache = kv_cache
# Build attn metadata
attn_metadata = self.builder.build(
@@ -295,7 +295,7 @@ def test_rope_kvcache_fusion(
}
q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
attn_layer = forward_context.no_compile_layers[model.layer_name]
kv_cache_unfused = attn_layer.kv_cache[0]
kv_cache_unfused = attn_layer.kv_cache
del dummy
torch._dynamo.mark_dynamic(qkv, 0)
@@ -309,7 +309,7 @@ def test_rope_kvcache_fusion(
}
q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
attn_layer = forward_context.no_compile_layers[model.layer_name]
kv_cache_fused = attn_layer.kv_cache[0]
kv_cache_fused = attn_layer.kv_cache
del dummy
assert fusion_pass.matched_count == 1

6
tests/v1/e2e/general/test_mamba_prefix_cache.py
View File

@@ -258,8 +258,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):
mamba_kv_cache_dict[
num_computed_tokens - num_computed_tokens % BLOCK_SIZE
] = (
kv_cache[0][0][block_id].clone(),
kv_cache[0][1][block_id].clone(),
kv_cache[0][block_id].clone(),
kv_cache[1][block_id].clone(),
)
last_num_computed_tokens = num_computed_tokens
@@ -302,7 +302,7 @@ def get_fake_process_mamba_fn(
mamba_layer_name = kv_cache_config.kv_cache_groups[
mamba_group_id
].layer_names[0]
mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[0][-1]
mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[-1]
mamba_block_table = input_batch.block_table.block_tables[
mamba_group_id
].block_table.cpu[0]

34
tests/v1/worker/test_gpu_model_runner.py
View File

@@ -670,8 +670,8 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
runner.initialize_kv_cache(kv_cache_config)
layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
layer_0_kv = vllm_ctx[layer_0].kv_cache
layer_1_kv = vllm_ctx[layer_1].kv_cache
# check layer 1 kv cache does NOT share memory with layer 0
assert id(layer_1_kv) != id(layer_0_kv)
@@ -740,8 +740,8 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
runner.initialize_kv_cache(kv_cache_config)
kv_cache_config_after_init = runner.kv_cache_config
layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
layer_0_kv = vllm_ctx[layer_0].kv_cache
layer_1_kv = vllm_ctx[layer_1].kv_cache
# check layer 1 kv cache shares memory with layer 0
assert id(layer_1_kv) == id(layer_0_kv)
@@ -864,9 +864,9 @@ def test_hybrid_attention_mamba_tensor_shapes():
np.random.shuffle(ind)
blocks0, blocks1 = ind[: (num_blocks // 2)], ind[(num_blocks // 2) :]
attn_shape = vllm_ctx[layer_0].kv_cache[0].shape
conv_shape = vllm_ctx[layer_2].kv_cache[0][0].shape
ssm_shape = vllm_ctx[layer_2].kv_cache[0][1].shape
attn_shape = vllm_ctx[layer_0].kv_cache.shape
conv_shape = vllm_ctx[layer_2].kv_cache[0].shape
ssm_shape = vllm_ctx[layer_2].kv_cache[1].shape
# assert we are using FlashInfer
assert attn_shape[0] % num_blocks == 0
@@ -905,21 +905,21 @@ def test_hybrid_attention_mamba_tensor_shapes():
kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio
for layer in [layer_0, layer_1]:
# attention: kv_cache[0][kernel_block_idx, kv_idx, ...]
# attention: kv_cache[kernel_block_idx, kv_idx, ...]
for i, kernel_block in enumerate(kernel_blocks_for_attention):
vllm_ctx[layer].kv_cache[0][kernel_block, :] = attn_blocks_constant[i]
vllm_ctx[layer].kv_cache[kernel_block, :] = attn_blocks_constant[i]
# fill mamba blocks with constants using kernel block indices
for layer in [layer_2, layer_3, layer_4, layer_5]:
# mamba: kv_cache[0][component][kernel_block_idx, ...]
# mamba: kv_cache[component][kernel_block_idx, ...]
for i, kv_block in enumerate(kv_blocks_for_mamba):
vllm_ctx[layer].kv_cache[0][0][kv_block, :] = conv_blocks_constant[i]
vllm_ctx[layer].kv_cache[0][1][kv_block, :] = ssm_blocks_constant[i]
vllm_ctx[layer].kv_cache[0][kv_block, :] = conv_blocks_constant[i]
vllm_ctx[layer].kv_cache[1][kv_block, :] = ssm_blocks_constant[i]
# verify attention and mamba contents are correct
for layer in [layer_0, layer_1]:
for i, kernel_block in enumerate(kernel_blocks_for_attention):
actual_kv = vllm_ctx[layer].kv_cache[0][kernel_block, :]
actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]
expected = attn_blocks_constant[i]
# Check K and V separately
@@ -928,8 +928,8 @@ def test_hybrid_attention_mamba_tensor_shapes():
for layer in [layer_2, layer_3, layer_4, layer_5]:
for i, kv_block in enumerate(kv_blocks_for_mamba):
actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]
actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]
expected_conv = conv_blocks_constant[i]
expected_ssm = ssm_blocks_constant[i]
@@ -938,8 +938,8 @@ def test_hybrid_attention_mamba_tensor_shapes():
for layer in [layer_2, layer_3, layer_4, layer_5]:
for i, kv_block in enumerate(kv_blocks_for_mamba):
actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]
actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]
expected_conv = conv_blocks_constant[i]
expected_ssm = ssm_blocks_constant[i]
assert torch.equal(actual_conv, expected_conv)

20
tests/v1/worker/test_utils.py
View File

@@ -23,10 +23,10 @@ def test_bind_kv_cache(default_vllm_config):
}
runner_kv_caches: list[torch.Tensor] = []
bind_kv_cache(kv_cache, ctx, runner_kv_caches)
assert ctx["layers.0.self_attn"].kv_cache[0] is kv_cache["layers.0.self_attn"]
assert ctx["layers.1.self_attn"].kv_cache[0] is kv_cache["layers.1.self_attn"]
assert ctx["layers.2.self_attn"].kv_cache[0] is kv_cache["layers.2.self_attn"]
assert ctx["layers.3.self_attn"].kv_cache[0] is kv_cache["layers.3.self_attn"]
assert ctx["layers.0.self_attn"].kv_cache is kv_cache["layers.0.self_attn"]
assert ctx["layers.1.self_attn"].kv_cache is kv_cache["layers.1.self_attn"]
assert ctx["layers.2.self_attn"].kv_cache is kv_cache["layers.2.self_attn"]
assert ctx["layers.3.self_attn"].kv_cache is kv_cache["layers.3.self_attn"]
assert runner_kv_caches[0] is kv_cache["layers.0.self_attn"]
assert runner_kv_caches[1] is kv_cache["layers.1.self_attn"]
@@ -50,8 +50,8 @@ def test_bind_kv_cache_non_attention(default_vllm_config):
runner_kv_caches: list[torch.Tensor] = []
bind_kv_cache(kv_cache, ctx, runner_kv_caches)
assert ctx["model.layers.20.attn"].kv_cache[0] is kv_cache["model.layers.20.attn"]
assert ctx["model.layers.28.attn"].kv_cache[0] is kv_cache["model.layers.28.attn"]
assert ctx["model.layers.20.attn"].kv_cache is kv_cache["model.layers.20.attn"]
assert ctx["model.layers.28.attn"].kv_cache is kv_cache["model.layers.28.attn"]
assert runner_kv_caches[0] is kv_cache["model.layers.20.attn"]
assert runner_kv_caches[1] is kv_cache["model.layers.28.attn"]
@@ -74,14 +74,14 @@ def test_bind_kv_cache_draft_model(default_vllm_config):
runner_kv_caches: list[torch.Tensor] = []
bind_kv_cache(kv_cache, ctx, runner_kv_caches)
assert ctx["model.layers.0.attn"].kv_cache[0] is kv_cache["model.layers.0.attn"]
assert ctx["model.layers.1.attn"].kv_cache[0] is kv_cache["model.layers.1.attn"]
assert ctx["model.layers.0.attn"].kv_cache is kv_cache["model.layers.0.attn"]
assert ctx["model.layers.1.attn"].kv_cache is kv_cache["model.layers.1.attn"]
assert (
ctx["draft_model.layers.0.attn"].kv_cache[0]
ctx["draft_model.layers.0.attn"].kv_cache
is kv_cache["draft_model.layers.0.attn"]
)
assert (
ctx["draft_model.layers.1.attn"].kv_cache[0]
ctx["draft_model.layers.1.attn"].kv_cache
is kv_cache["draft_model.layers.1.attn"]
)