[Bugfix] Detect mixed precision in packed KV cache specs (#49623)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Michael Goin committed
Jul 23, 2026 at 23:07 UTC
275556c35cb4ea12ee362a38e222f66f711a8060
2 files changed
+40
-5
tests/v1/core/test_kv_cache_utils.py
+27
@@ -2069,6 +2069,33 @@ def test_generate_scheduler_kv_cache_config():
2069
)
2070
2071
2072
+def test_mixed_precision_kv_cache_with_uniform_type_specs():
2073
+ fp8_spec = new_kv_cache_spec(dtype=torch.float8_e4m3fn)
2074
+ bf16_spec = new_kv_cache_spec(dtype=torch.bfloat16)
2075
+ worker_config = KVCacheConfig(
2076
+ num_blocks=10,
2077
+ kv_cache_tensors=[],
2078
+ kv_cache_groups=[
2079
+ KVCacheGroupSpec(
2080
+ ["fp8_layer"],
2081
+ UniformTypeKVCacheSpecs(
2082
+ block_size=16, kv_cache_specs={"fp8_layer": fp8_spec}
2083
+ ),
2084
+ ),
2085
+ KVCacheGroupSpec(
2086
+ ["bf16_layer"],
2087
+ UniformTypeKVCacheSpecs(
2088
+ block_size=16, kv_cache_specs={"bf16_layer": bf16_spec}
2089
+ ),
2090
+ ),
2091
+ ],
2092
+ )
2093
+ scheduler_config = generate_scheduler_kv_cache_config([worker_config])
2094
+
2095
+ assert worker_config.needs_kv_cache_zeroing
2096
+ assert scheduler_config.needs_kv_cache_zeroing
2097
+
2098
+
2099
def new_mla_spec(cache_dtype_str=None, block_size=16):
2100
# head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576
2101
return MLAAttentionSpec(
vllm/v1/kv_cache_interface.py
+13
-5
@@ -975,11 +975,19 @@ class KVCacheConfig:
975
@property
976
def has_mixed_precision_kv_cache(self) -> bool:
977
"""Whether attention groups store their KV cache at more than one precision."""
978
- kv_cache_precisions = {
979
- (group.kv_cache_spec.dtype, group.kv_cache_spec.kv_quant_mode)
980
- for group in self.kv_cache_groups
981
- if isinstance(group.kv_cache_spec, AttentionSpec)
982
- }
978
+ kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set()
979
+ for group in self.kv_cache_groups:
980
+ group_spec = group.kv_cache_spec
981
+ group_specs = (
982
+ list(group_spec.kv_cache_specs.values())
983
+ if isinstance(group_spec, UniformTypeKVCacheSpecs)
984
+ else [group_spec]
985
+ )
986
+ kv_cache_precisions.update(
987
+ (spec.dtype, spec.kv_quant_mode)
988
+ for spec in group_specs
989
+ if isinstance(spec, AttentionSpec)
990
+ )
991
return len(kv_cache_precisions) > 1
992
993
@property