Skip to content
Merged
Changes from 1 commit
Commits
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
34 changes: 33 additions & 1 deletion vllm/v1/kv_cache_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -939,6 +939,38 @@ class KVCacheConfig:
def has_mamba_layers(self) -> bool:
return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)

@property
def has_block_dropping_layers(self) -> bool:
"""Any group uses an attention type that frees KV blocks mid-request
(sliding-window or chunked-local)."""
return any(
isinstance(g.kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec))
for g in self.kv_cache_groups
)

@property
def has_quantized_kv_cache(self) -> bool:
"""Any group stores its KV cache in a quantized dtype (FP8/NVFP4)."""
return any(
isinstance(g.kv_cache_spec, AttentionSpec)
and g.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE
for g in self.kv_cache_groups
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I feel the condition is a little bit off. The fundamental issue is that we have multiple KV cache groups stored at different precisions, right? Any way to capture that more directly?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Agreed. Switched to has_mixed_precision_kv_cache to directly capture the case. Thank you for the comment!


@property
def needs_kv_cache_zeroing(self) -> bool:
return self.has_mamba_layers
"""Whether newly allocated KV cache blocks must be zeroed before use.

Required in two cases:
- Mamba layers, whose state is read before being fully written.
- Hybrid configs mixing a block-dropping attention type with a quantized
KV cache. Groups share the same tensors, so a page freed mid-request by
a block-dropping group can be reallocated and read before every slot is
written this step (partial-block tail, alignment padding). Stale
quantized bytes decode to NaN/Inf and poison attention; zeroing makes
such slots read as 0. Uniform/unquantized configs are unaffected.
"""
needs_quantized_reuse_zeroing = (
self.has_block_dropping_layers and self.has_quantized_kv_cache
)
return self.has_mamba_layers or needs_quantized_reuse_zeroing
Loading