iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[KV Connector][Offloading] Certify attention-only hybrids in the canonical portability gate (#51689)

Browse Source 
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Itay Etelis
2026-08-31 15:53:19 +03:00
committed by GitHub
co-authored by Itay Etelis Or Ozeri
parent 4c58a0c398
commit c5d840ff6a
2 changed files with 91 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+56 -66
View File
@@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
HiddenStateCacheSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpec,
KVCacheTensor,
MambaSpec,
MLAAttentionSpec,
@@ -247,8 +248,16 @@ def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig:
)
def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
config = _make_vllm_config()
def _parallelism_agnostic(
kv_cache_groups: list[KVCacheGroupSpec],
*,
canonical: bool = False,
v2: bool = False,
) -> bool:
config = _make_vllm_config(
extra_config={"canonical_layout": True} if canonical else None
)
config.use_v2_model_runner = v2
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
@@ -641,85 +650,66 @@ def test_parallelism_agnostic_for_single_full_attention_group():
assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
_SWA_SPEC = SlidingWindowSpec(
block_size=16,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
sliding_window=128,
)
_SWA_MLA_SPEC = SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
)
def _groups(*specs: KVCacheSpec) -> list[KVCacheGroupSpec]:
return [KVCacheGroupSpec([f"l{i}"], spec) for i, spec in enumerate(specs)]
@pytest.mark.parametrize(
"kv_cache_groups",
[
[KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))],
[
KVCacheGroupSpec(
["l0"],
SlidingWindowSpec(
block_size=16,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
sliding_window=128,
),
)
],
[
KVCacheGroupSpec(["l0"], _full_attention_spec()),
KVCacheGroupSpec(["l1"], _full_attention_spec()),
],
_groups(_mla_spec(head_size=576)),
_groups(_SWA_SPEC),
_groups(_full_attention_spec(), _full_attention_spec()),
],
)
def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
assert not _parallelism_agnostic(kv_cache_groups)
def test_canonical_layout_widens_parallelism_agnostic_to_mla():
"""The canonical layout dedups the TP-replicated MLA latent into one
portable copy, so the gate admits MLA — but only when canonical_layout
is requested."""
mla_groups = [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))]
assert not _parallelism_agnostic(mla_groups)
config = _make_vllm_config(extra_config={"canonical_layout": True})
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=mla_groups,
)
offloading_config = build_offloading_config(config, kv_cache_config)
assert offloading_config.parallel.is_parallelism_agnostic
assert offloading_config.canonical_layout
# hybrid groupings stay out: their non-full-attention layers can only
# derive opaque (exact-topology) mappings
hybrid_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(["l0"], _full_attention_spec()),
KVCacheGroupSpec(["l1"], _full_attention_spec()),
],
)
assert not build_offloading_config(
config, hybrid_config
).parallel.is_parallelism_agnostic
@pytest.mark.parametrize(
("kv_cache_groups", "certified"),
[
pytest.param(_groups(_mla_spec(head_size=576)), True, id="mla-latent"),
pytest.param(
_groups(_full_attention_spec(), _SWA_SPEC), True, id="attention-hybrid"
),
pytest.param(_groups(_SWA_MLA_SPEC), False, id="swa-mla"),
pytest.param(
list(_make_mamba_hybrid_kv_cache_config().kv_cache_groups),
False,
id="mamba-hybrid",
),
],
)
def test_canonical_layout_gate(kv_cache_groups, certified):
"""The canonical layout certifies portability group by group; none of
these shapes are portable in the direct layout."""
assert not _parallelism_agnostic(kv_cache_groups)
assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified
def test_canonical_layout_certifies_v2_model_runner():
"""Canonical bytes are certified per layer against live tensor strides at
registration, so the static gate must not depend on the model-runner
version — the v2 runner is the case the canonical layout exists for."""
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
)
config = _make_vllm_config()
config.use_v2_model_runner = True
assert not build_offloading_config(
config, kv_cache_config
).parallel.is_parallelism_agnostic
config = _make_vllm_config(extra_config={"canonical_layout": True})
config.use_v2_model_runner = True
assert build_offloading_config(
config, kv_cache_config
).parallel.is_parallelism_agnostic
groups = _groups(_full_attention_spec())
assert _parallelism_agnostic(groups, canonical=True, v2=True)
def test_parallelism_agnostic_disabled_on_v2_model_runner():
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+35 -25
View File
@@ -8,7 +8,10 @@ from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
from vllm.v1.kv_cache_interface import (
AttentionSpec,
FullAttentionSpec,
KVCacheSpec,
MLAAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
)
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
@@ -139,35 +142,42 @@ def build_offloading_config(
and parallel_config.decode_context_parallel_size == 1
and parallel_config.prefill_context_parallel_size == 1
)
# Canonical pages are topology-free by construction, so the canonical
# layout widens the gate to every config whose mappings derive portable:
# exactly-sharded or replicated GQA heads (writer rotation) and the
# TP-replicated MLA latent (one canonical copy for all replicas). The
# model-runner version is irrelevant here — certification happens per
# layer against live tensor strides at registration, and create_worker
# fails closed against this flag if any layer cannot be certified.
# Canonical pages are topology-free, so the gate widens to every config
# whose mappings derive portable, group by group; certification happens
# per layer at registration and create_worker fails closed on this flag.
if canonical_layout and not is_parallelism_agnostic:
tp_size = parallel_config.tensor_parallel_size
if isinstance(single_group_spec, FullAttentionSpec):
if type(single_group_spec) is MLAAttentionSpec:
spec_certifiable = (
single_group_spec.tokens_per_state == 1
and single_group_spec.real_page_size_bytes
% single_group_spec.block_size
== 0
total_kv_heads = vllm_config.model_config.get_total_num_kv_heads()
def spec_certifiable(spec: KVCacheSpec) -> bool:
"""Conservative static mirror of _layer_mapping's per-layer checks."""
if not isinstance(spec, AttentionSpec):
return False
if spec.kv_quant_mode.is_per_token_head:
return False
if type(spec) is MLAAttentionSpec:
return (
spec.tokens_per_state == 1
and spec.real_page_size_bytes % spec.block_size == 0
)
else:
total_kv_heads = vllm_config.model_config.get_total_num_kv_heads()
spec_certifiable = (
total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0
)
is_parallelism_agnostic = (
spec_certifiable
and not single_group_spec.kv_quant_mode.is_per_token_head
and parallel_config.decode_context_parallel_size == 1
and parallel_config.prefill_context_parallel_size == 1
and parallel_config.world_size == tp_size
if isinstance(spec, (SlidingWindowMLASpec, MLAAttentionSpec)):
return False
if not isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)):
return False
return (
total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0
) and spec.num_kv_heads == max(1, total_kv_heads // tp_size)
is_parallelism_agnostic = (
len(kv_cache_config.kv_cache_groups) > 0
and all(
spec_certifiable(group.kv_cache_spec)
for group in kv_cache_config.kv_cache_groups
)
and parallel_config.decode_context_parallel_size == 1
and parallel_config.prefill_context_parallel_size == 1
and parallel_config.world_size == tp_size
)
kv_events_config = vllm_config.kv_events_config
cache_dtype = (