lmcache--lmcache/docs/design/integration/vllm/kv_cache_group_edits.md

KV Cache Group Edits

Summary

lmcache/integration/vllm/kv_cache_group_edits.py is the single place where vLLM KV cache groups are re-presented ("edited") before LMCache registration. The connector calls apply_kv_cache_group_edits(kv_cache_config, kv_caches) once in register_kv_caches, and the edited dict feeds both engine-group-info creation (kv_cache_groups.py) and transfer registration.

LMCache derives each group's transfer metadata (block size, page layout, dtype) from the registered tensors, and interprets store/retrieve block IDs in vLLM's scheduler-side block-id space (kv_cache_spec.block_size units). The edits exist to restore one invariant the raw tensors can violate:

The registered tensor's paging granularity must equal the block-id granularity.

Structure

Each case is one KVCacheGroupEdit rule in the module's _EDITS registry: matches(spec, kv_cache) decides structurally — from the vLLM spec kind (get_kv_cache_spec_kind, which also unwraps UniformTypeKVCacheSpecs) and the registered tensor, never from model name or architecture — and apply(spec, kv_cache) produces a view over the same storage. First matching rule wins; unmatched layers pass through. Covering a new group kind means adding one rule.

Model name/arch is deliberately not an input: the same architecture yields different group structures depending on runtime decisions (mamba_cache_mode, attention backend's kernel block size, TP), all of which the config + tensors already resolve.

Edits

Both rules apply only to Mamba-hybrid models (the registry is only consulted when kv_cache_config.has_mamba_layers).

1. Mamba state pages

A Mamba / linear-attention layer (e.g. Qwen3.5 GDN) registers [conv_state, ssm_state] — two tensors with different shapes and dtypes, laid out contiguously in one padded page (conv | ssm | pad). The raw pair trips format discovery (the SSM view starts mid-page). The edit reinterprets each page as one bf16 tensor shaped (num_blocks, 2, block_size, 1, head_size) over the same storage, where head_size is derived so the bytes fill the page exactly.

2. Sub-paged full attention

vLLM unifies page sizes across hybrid groups by inflating the attention logical block size (vllm/platforms/interface.py:_align_hybrid_block_size; Qwen3.5-0.8B: 544), while the attention backend re-pages the physical tensor at its own kernel block size (vllm/v1/worker/utils.py: prepare_kernel_block_sizes; FlashAttention on hybrids: 32). Logical block n then occupies the k = logical/kernel contiguous kernel pages n*k .. n*k+k-1 (vLLM expands the worker-side block table the same way, BlockTable.map_to_kernel_blocks); the scheduler-side block IDs LMCache receives stay logical.

Registering the raw kernel-paged tensor makes LMCache discover block_size == kernel < logical, and _derive_compression_metadata (lmcache/v1/kv_layer_groups.py) misclassifies the group as compressed: only 1/k of each chunk's KV is transferred, addressed against the kernel page space. The edit re-views the tensor as (num_kernel_pages / k, 2, logical_block_size, 1, head_size) — a pure view(), valid because k kernel pages tile each logical page's bytes exactly (enforced; see Invariants).

Startup validation

validate_kv_cache_groups (called at connector init and again at registration) rejects group specs the transfer path cannot serve correctly, with one aggregated error: CrossAttentionSpec, and Mamba with mamba_cache_mode != "align" (no reusable snapshots). Declared slot compression (compress_ratio > 1 / tq_slot_size > 0, e.g. DeepSeek-V4) is not rejected — those groups are served by the compression path in lmcache/v1/kv_layer_groups.py and only skipped by the edits here. Note the compression path still derives per-group ratios from the unified vLLM block size; switching it to per-group block sizes is pending in a separate PR.

Reference: vLLM PR #42828 (Mooncake store HMA support) uses the same validate-and-reject-up-front pattern, and is the reference design for the deferred follow-ups (per-group store/load masks; manager-mirroring hit computation). Caveat for the latter: LMCache's lookup doubles as prefetch, so vLLM's lookup-first-then-trim flow does not map directly. See the module docstring for the full check-when list.

Non-edit: declared compression

Groups whose spec declares slot compression — MLAAttentionSpec. compress_ratio > 1 (DeepSeek-V4 slot packing, storage_block_size < block_size) or TQFullAttentionSpec.tq_slot_size > 0 — genuinely store fewer physical slots than logical tokens. They must reach the compression path in lmcache/v1/kv_layer_groups.py unedited. (DeepSeek-V3.2's fp8_ds_mla cache packs bytes per slot, not slots per block: its specs keep block_size == scheduler block size and compress_ratio == 1, so it never needs an edit either.)

The sub-paged rule's matches excludes declared-compression specs by their own fields (compress_ratio / tq_slot_size), and its apply additionally verifies by byte accounting that k kernel pages tile the logical page's bytes exactly — any undeclared packed layout fails with a loud ValueError rather than being transferred wrongly.

The opaque-page contract

An edited view's dims are addressing metadata only (block id → byte range). The named dims are not semantic: a Mamba view's "K plane" is conv/ssm bytes, and a sub-paged attention view's "K plane" interleaves true K and V at kernel-page granularity (true K is not contiguous across kernel pages, so no logical-block view can have a pure-K plane). The synthetic head shape (1, page_bytes / (2 * block_size * elem)) signals this deliberately.

Byte transport round-trips correctly because store and retrieve share the same bijective block-id → bytes mapping. Consequences:

• Valid: store/retrieve through the MP transfer path on the same engine configuration.
• Not valid for edited groups: content-aware processing (serde compression, blending, head resharding, layout conversion), and sharing cache entries across engines whose attention backends choose different kernel block sizes (the byte order inside a logical page is backend-dependent).

Invariants

• Edits are pure tensor views over the registered storage — never copies.
• A sub-paged view is only produced when kernel_page_bytes * k == spec.page_size_bytes; any mismatch raises ValueError (fail loudly rather than silently transfer a compressed layout).
• After edits, every registered tensor's block dim equals its group's kv_cache_spec.block_size, so the server derives compress_ratio == 1 for these groups.

Code map

Area	File
Edits (this doc)	lmcache/integration/vllm/kv_cache_group_edits.py
Caller	lmcache/integration/vllm/lmcache_mp_connector.py (register_kv_caches)
Compression-ratio derivation (downstream consumer)	lmcache/v1/kv_layer_groups.py
vLLM block-size inflation	vllm/platforms/interface.py (_align_hybrid_block_size)
vLLM kernel-page split + block-table expansion	vllm/v1/worker/utils.py, vllm/v1/worker/block_table.py
End-to-end test	.buildkite/k3_tests/multiprocess/scripts/run-single-test.sh (hma_lm_eval_qwen3_5)

Testing is end-to-end only (the hma_lm_eval_qwen3_5 store-vs-retrieve gsm8k check): the edit internals are expected to change as more group kinds are covered, so tests pin the observable contract — faithful retrieve — rather than view shapes.