"""Model-owned KV grouping for the 54 GDN + 16 attention ALTAY stack. vLLM 0.25.1 chooses 16 as the native hybrid group size. That removes padding from the full-attention family, but pads 54 GDN states to 64. The first ALTAY repair instead used a group size of 18 to minimize the *aggregate* slot count (72 instead of 80). That objective was wrong for a million-token request: GDN groups retain one recurrent state per request, while every padded full-attention slot is paid at every token. Group size 18 therefore padded the 16-layer attention family to 18 and wasted roughly 1.5 GiB at 1,010,000 tokens with the release TurboQuant profile. A group size of 8 keeps all 16 full-attention layers exact and moves the only two padding slots to the fixed-size GDN family (54 -> 56). All 70 real cache owners and their contents are unchanged. Only cache-pool grouping changes. The module is installed from the model's vLLM plugin before cache planning, so the optimization travels with every supported quantization. """ from __future__ import annotations import logging import os from collections import defaultdict from dataclasses import dataclass from typing import Any, Callable logger = logging.getLogger(__name__) ARCHITECTURE_ID = "LomonosovZenitAltayForConditionalGeneration" ENV_NAME = "ZENIT_ALTAY_KV_LAYOUT" TRACE_ENV_NAME = "ZENIT_ALTAY_KV_LAYOUT_TRACE_PATH" REPAIR_ID = "ALTAY72_KV_ATTENTION_EXACT_G8_V2" EXPECTED_FAMILY_COUNTS = (16, 54) COMPACT_GROUP_SIZE = 8 @dataclass(frozen=True) class LayoutPlan: family_counts: tuple[int, ...] group_size: int actual_layers: int group_count: int allocated_slots: int padding_slots: int def make_layout_plan( family_counts: tuple[int, ...], group_size: int, ) -> LayoutPlan: """Return exact slot accounting for a proposed uniform group size.""" if not family_counts or any(count <= 0 for count in family_counts): raise ValueError("family counts must be positive") if group_size <= 0: raise ValueError("group size must be positive") group_count = sum((count + group_size - 1) // group_size for count in family_counts) actual_layers = sum(family_counts) allocated_slots = group_count * group_size return LayoutPlan( family_counts=tuple(family_counts), group_size=group_size, actual_layers=actual_layers, group_count=group_count, allocated_slots=allocated_slots, padding_slots=allocated_slots - actual_layers, ) LEGACY_PLAN = make_layout_plan(EXPECTED_FAMILY_COUNTS, 16) COMPACT_PLAN = make_layout_plan(EXPECTED_FAMILY_COUNTS, COMPACT_GROUP_SIZE) # Состав семей может измениться намеренно — например когда RazorAttention даёт # части слоёв полного внимания короткое окно, и шестнадцать слоёв становятся # семью плюс девятью. Проверка ниже нарочно строгая: она отвергает ЛЮБОЙ # незнакомый состав, потому что молча раскладывать кэш не так, как # сертифицировано, хуже отказа. Поэтому новый состав надо объявить явно. _expected_counts: tuple[int, ...] = EXPECTED_FAMILY_COUNTS _expected_plan: LayoutPlan = COMPACT_PLAN def expect_family_counts(counts: tuple[int, ...]) -> LayoutPlan: """Объявить другой состав семей и пересчитать под него план. Возвращает новый план. Вызывать до подъёма движка; вызывающий обязан знать, почему состав изменился, — самодеятельности здесь быть не должно. """ global _expected_counts, _expected_plan _expected_counts = tuple(sorted(counts)) _expected_plan = make_layout_plan(_expected_counts, COMPACT_GROUP_SIZE) logger.info( "%s: ожидаемый состав семей изменён на %s, план: %d групп, %d слотов, " "%d из них добивка", REPAIR_ID, _expected_counts, _expected_plan.group_count, _expected_plan.allocated_slots, _expected_plan.padding_slots, ) return _expected_plan def _mode() -> str: return os.environ.get(ENV_NAME, "attention_exact").strip().lower() def _trace(event: str, **payload: Any) -> None: """Optionally persist process-local integration evidence. The trace is disabled by default. Certification probes opt in with an absolute path so an EngineCore ``spawn`` cannot silently lose the patch. """ destination = os.environ.get(TRACE_ENV_NAME) if not destination: return try: import json row = { "event": event, "pid": os.getpid(), "repair_id": REPAIR_ID, **payload, } with open(destination, "a", encoding="utf-8") as stream: stream.write( json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" ) stream.flush() except Exception: # A diagnostic trace must never alter cache-planner semantics. logger.exception("%s could not write its integration trace", REPAIR_ID) def _architecture_names(vllm_config: Any) -> set[str]: model_config = getattr(vllm_config, "model_config", None) names: set[str] = set() resolved = getattr(model_config, "architecture", None) if isinstance(resolved, str): names.add(resolved) declared = getattr(model_config, "architectures", ()) if isinstance(declared, (list, tuple)): names.update(name for name in declared if isinstance(name, str)) hf_config = getattr(model_config, "hf_config", None) hf_declared = getattr(hf_config, "architectures", ()) if isinstance(hf_declared, (list, tuple)): names.update(name for name in hf_declared if isinstance(name, str)) return names def _compact_groups(kv_cache_spec: dict[str, Any]) -> list[Any] | None: """Build ALTAY groups with 18 slots, or return None on a foreign layout.""" from vllm.v1.core import kv_cache_utils as core hidden_type = getattr(core, "HiddenStateCacheSpec", ()) if hidden_type and any( isinstance(spec, hidden_type) for spec in kv_cache_spec.values() ): # Keep this release optimization narrowly scoped. A future vLLM that # exposes separate HiddenStateCacheSpec objects must be re-certified. return None unified = core.unify_kv_cache_spec_page_size(dict(kv_cache_spec)) same_type_layers: dict[Any, list[str]] = defaultdict(list) for layer_name, layer_spec in unified.items(): same_type_layers[layer_spec].append(layer_name) counts = tuple(sorted(len(layers) for layers in same_type_layers.values())) if counts != _expected_counts: return None grouped_layer_names: list[list[str]] = [] for layers in same_type_layers.values(): number_of_groups = (len(layers) + COMPACT_GROUP_SIZE - 1) // COMPACT_GROUP_SIZE for index in range(number_of_groups): names = layers[index::number_of_groups] if names: grouped_layer_names.append(names) groups = core.create_kv_cache_group_specs(unified, grouped_layer_names) plan = make_layout_plan(counts, max(len(group.layer_names) for group in groups)) if plan != _expected_plan: raise RuntimeError(f"unexpected ALTAY compact layout: {plan!r}") return groups def install_altay72_compact_kv_layout(*, force: bool = False) -> bool: """Install the architecture-scoped vLLM cache planner patch once.""" from vllm.v1.core import kv_cache_utils as core if getattr(core, "_zenit_altay72_compact_layout", False): if force: core._zenit_altay72_compact_layout_force = True _trace( "install_reused", force=force, forced_by_model=bool( getattr(core, "_zenit_altay72_compact_layout_force", False) ), ) return False original: Callable[..., list[Any]] = core.get_kv_cache_groups original_uniform: Callable[..., list[Any]] = ( core._get_kv_cache_groups_uniform_page_size ) def patched(vllm_config: Any, kv_cache_spec: dict[str, Any]) -> list[Any]: forced_by_model = bool( getattr(core, "_zenit_altay72_compact_layout_force", False) ) if ( not forced_by_model and ARCHITECTURE_ID not in _architecture_names(vllm_config) ): return original(vllm_config, kv_cache_spec) mode = _mode() if mode in {"legacy", "off", "0", "false"}: return original(vllm_config, kv_cache_spec) if getattr( getattr(vllm_config, "scheduler_config", None), "disable_hybrid_kv_cache_manager", False, ): return original(vllm_config, kv_cache_spec) try: groups = _compact_groups(kv_cache_spec) except Exception: if mode in {"required", "strict"}: raise logger.exception("%s failed; retaining vLLM native layout", REPAIR_ID) return original(vllm_config, kv_cache_spec) if groups is None: if mode in {"required", "strict"}: raise RuntimeError( f"{REPAIR_ID} expected cache-family counts " f"{_expected_counts}" ) logger.warning( "%s did not recognize the cache layout; retaining vLLM native layout", REPAIR_ID, ) return original(vllm_config, kv_cache_spec) _trace( "top_level_active", mode=mode, group_lengths=[len(group.layer_names) for group in groups], allocated_slots=_expected_plan.allocated_slots, padding_slots=_expected_plan.padding_slots, ) if mode in {"required", "strict"}: logger.warning( "%s integration proof: top-level cache grouping is active", REPAIR_ID, ) logger.info( "%s active: actual=%d, slots=%d->%d, padding=%d->%d", REPAIR_ID, _expected_plan.actual_layers, LEGACY_PLAN.allocated_slots, _expected_plan.allocated_slots, LEGACY_PLAN.padding_slots, _expected_plan.padding_slots, ) return groups def patched_uniform(kv_cache_spec: dict[str, Any]) -> list[Any]: """Patch the exact private symbol used by vLLM's native planner. vLLM imports ``get_kv_cache_configs`` into EngineCore before general plugins are registered. Its function globals still resolve through ``kv_cache_utils``, but patching this private grouping primitive as well makes the integration independent of that import timing. """ mode = _mode() forced_by_model = bool( getattr(core, "_zenit_altay72_compact_layout_force", False) ) if mode in {"legacy", "off", "0", "false"}: return original_uniform(kv_cache_spec) if not forced_by_model and mode not in {"required", "strict"}: return original_uniform(kv_cache_spec) try: groups = _compact_groups(kv_cache_spec) except Exception: if mode in {"required", "strict"}: raise logger.exception( "%s private planner failed; retaining vLLM native layout", REPAIR_ID, ) return original_uniform(kv_cache_spec) if groups is None: if mode in {"required", "strict"}: raise RuntimeError( f"{REPAIR_ID} expected cache-family counts " f"{_expected_counts} in private planner" ) return original_uniform(kv_cache_spec) _trace( "private_uniform_active", mode=mode, group_lengths=[len(group.layer_names) for group in groups], allocated_slots=_expected_plan.allocated_slots, padding_slots=_expected_plan.padding_slots, ) if mode in {"required", "strict"}: logger.warning( "%s integration proof: private uniform-page planner is active", REPAIR_ID, ) return groups core.get_kv_cache_groups = patched core._get_kv_cache_groups_uniform_page_size = patched_uniform core._zenit_altay72_compact_layout = True core._zenit_altay72_compact_layout_force = bool(force) core._zenit_altay72_compact_layout_original = original core._zenit_altay72_compact_layout_uniform_original = original_uniform _trace("install_new", force=force, mode=_mode()) if _mode() in {"required", "strict"}: logger.warning( "%s integration proof: installer active (force=%s)", REPAIR_ID, force, ) return True __all__ = [ "COMPACT_PLAN", "LEGACY_PLAN", "REPAIR_ID", "TRACE_ENV_NAME", "LayoutPlan", "install_altay72_compact_kv_layout", "make_layout_plan", ]