I’d like to implement sparse attention KV cache policies (such as H2O) in vLLM. One issue I am running into is that a page (or block) needs to extend over all heads of a model.
In H2O or other sparse attention methods, KV vectors can be in the cache for some heads, but not for others. This does not need any special scaled dot product attention (SDPA), because we can always use torch.scatter and torch.gather to distribute or extract dense content.
I could implement H2O on the page level, but this would mean if KV vectors are evicted from the cache, this must happen for all heads.
Any pointer to how sparse attention can be properly implemented in vLLM, would probably help me. Or is this library simply restricted to serve non-sparse attention only?
vLLM currently assumes that KV cache blocks (pages) are allocated and managed at the sequence level, and all heads in a layer share the same block layout. This means that, as of now, vLLM does not natively support per-head or per-layer sparse KV cache eviction (such as H2O) where different heads could have different tokens retained in cache. Instead, eviction and compaction are performed at the block level, affecting all heads together. This limitation is discussed in several RFCs and issues, which note that supporting more advanced, head-wise or layer-wise sparse KV cache management would require significant changes to the memory layout and block manager abstractions in vLLM. Current workarounds (like using block_size=1) are inefficient and not scalable for production use. The system is not restricted to non-sparse attention, but true per-head sparse KV cache is not yet supported in the mainline codebase. See the discussion in RFC: Support KV Cache Compaction, RFC: Sparse KV cache management framework, and RFC: Support sparse KV cache framework.
If you want to implement H2O or similar methods, you would currently need to evict KV cache at the block/page level, which applies to all heads. Per-head selective eviction is not supported, and implementing it would require a redesign of the block manager and memory layout to allow non-uniform cache across heads and layers. Would you like a more detailed breakdown of the technical blockers and possible directions for extending vLLM to support per-head sparse KV cache?
Thanks for these comments!
I implemented H2O on my own in a fairly simple way. I am using a fixed-size buffer of shape (batch_size, n_head, cache_length, head_size), which can store information (last axis) for n_headheads and cache_length tokens.
And I use torch.scatter and torch.gather to write into or read from this buffer. This means that I can easily cache KV info for different tokens, depending on batch and head.
Wouldn’t that be something that vLLM could do on the block level?