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, 02:00 UTC
[MRV2] Fix for DS v3.2 (#38030)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 24, 2026 at 14:03 UTC
4b53740d7f3c2216be7a737d585bff37f4975d2e
1 file changed
+5
-2
vllm/v1/worker/gpu/attn_utils.py
+5
-2
@@ -115,9 +115,12 @@ def _reshape_kv_cache(
115
) -> dict[str, torch.Tensor]:
116
kv_caches: dict[str, torch.Tensor] = {}
117
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
118
- kv_cache_spec = kv_cache_group_spec.kv_cache_spec
119
- assert isinstance(kv_cache_spec, AttentionSpec)
118
for layer_name in kv_cache_group_spec.layer_names:
119
+ kv_cache_spec = kv_cache_group_spec.kv_cache_spec
120
+ if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
121
+ kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
122
+ assert isinstance(kv_cache_spec, AttentionSpec)
123
+
124
raw_tensor = kv_cache_raw_tensors[layer_name]
125
assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
126
num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes