A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 14:30 UTC
[BugFix] Correct max memory usage for multiple KV-cache groups (#36030)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Harry Huang committed
Mar 17, 2026 at 08:38 UTC
45f526d65237d9073a5f3be166b306580687f210
2 files changed
+45
-2
tests/v1/core/test_kv_cache_utils.py
+41
@@ -43,6 +43,7 @@ from vllm.v1.kv_cache_interface import (
43
KVCacheGroupSpec,
44
KVCacheSpec,
45
KVCacheTensor,
46
+ MambaSpec,
47
MLAAttentionSpec,
48
SlidingWindowSpec,
49
UniformTypeKVCacheSpecs,
@@ -157,6 +158,24 @@ def new_chunked_local_attention_spec(
158
)
159
160
161
+def new_mamba_spec(
162
+ block_size=16,
163
+ shapes=((2, 512), (3, 32, 32)),
164
+ dtypes=(torch.float32, torch.float32),
165
+ num_speculative_blocks=2,
166
+ mamba_cache_mode="none",
167
+ page_size_padded=None,
168
+):
169
+ return MambaSpec(
170
+ block_size=block_size,
171
+ shapes=shapes,
172
+ dtypes=dtypes,
173
+ page_size_padded=page_size_padded,
174
+ mamba_cache_mode=mamba_cache_mode,
175
+ num_speculative_blocks=num_speculative_blocks,
176
+ )
177
+
178
+
179
@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
180
def test_none_hash(monkeypatch, hash_fn):
181
import vllm.v1.core.kv_cache_utils
@@ -2010,6 +2029,28 @@ def test_auto_fit_max_model_len():
2029
assert vllm_config.model_config.max_model_len > 0
2030
2031
2032
+def test_auto_fit_max_model_len_with_hybrid():
2033
+ """Test that auto-fit works with hybrid KV cache specs."""
2034
+ # Create config with original_max_model_len=-1 to trigger auto-fit
2035
+ model_config = ModelConfig(max_model_len=8192)
2036
+ # Simulate the user passing -1 by setting original_max_model_len
2037
+ model_config.original_max_model_len = -1
2038
+ vllm_config = VllmConfig(model_config=model_config)
2039
+
2040
+ mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer
2041
+ gamma = 2
2042
+ kv_cache_specs = {
2043
+ "layer_1": new_mamba_spec(num_speculative_blocks=gamma),
2044
+ "layer_2": new_kv_cache_spec(),
2045
+ }
2046
+
2047
+ available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma)
2048
+ _kv_cache_configs = get_kv_cache_configs(
2049
+ vllm_config, [kv_cache_specs], [available_memory]
2050
+ )
2051
+ assert vllm_config.model_config.max_model_len == 1024
2052
+
2053
+
2054
def test_auto_fit_max_model_len_not_triggered():
2055
"""Test that auto-fit is not triggered when original_max_model_len is not -1."""
2056
model_config = ModelConfig(max_model_len=16)
vllm/v1/core/kv_cache_utils.py
+4
-2
@@ -1356,8 +1356,10 @@ def _max_memory_usage_bytes_from_groups(
1356
page_size = get_uniform_page_size(
1357
[group.kv_cache_spec for group in kv_cache_groups]
1358
)
1359
- any_spec = kv_cache_groups[0].kv_cache_spec
1360
- blocks_needed = cdiv(any_spec.max_memory_usage_bytes(vllm_config), page_size)
1359
+ blocks_needed = sum(
1360
+ cdiv(group.kv_cache_spec.max_memory_usage_bytes(vllm_config), page_size)
1361
+ for group in kv_cache_groups
1362
+ )
1363
1364
return group_size * page_size * blocks_needed
1365