diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index afb0a9bda5..99d2ccfb82 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -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(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 70919ce24e..1faa4f5495 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -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 = (