"""Native vLLM execution of the ALTAY-72M-SKV logical-depth overlay. The 64 physical Qwen3.5 layers remain the only owners of the large projection weights. ALTAY inserts two parameter-tied 3+1 replay superblocks after physical layers 31 and 63: * six Gated-DeltaNet replays share the source layer weights but own six independent recurrent cache slots; * two full-attention replays compute fresh Q/O while reading the already materialized K/V of their physical anchor layers; * only the 28 admitted ALTAY gate/delta tensors are stored additionally. This module is deliberately independent of the parent weight quantization. The source projections keep their native vLLM quantized kernels while the small admitted ALTAY delta tensors stay in BF16/FP32 in every release variant. """ from __future__ import annotations import copy import os from itertools import islice from typing import Any import torch from torch import nn from torch.nn import functional as F from vllm.distributed import get_pp_group from vllm.model_executor.layers.attention import Attention from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.sequence import IntermediateTensors ALTAY_ARCHITECTURE_ID = "LOMONOSOV_ZENIT_ALTAY_72M_SKV_V1" def _logger(): from vllm.logger import init_logger return init_logger("vllm.lomonosov_zenit_altay.overlay") def _alpha_scale(variable: str) -> float: """Диагностический множитель на вклад ветви накладки. По умолчанию 1.0. Зачем. arXiv:2606.11052 («Attention Amnesia in Hybrid LLMs») замерил, что дообучение на цепочках рассуждений систематически ломает дальний поиск у гибридных линейно-внимательных моделей: градиенты смещаются к ближним связям и портят проекции со стороны ЗАПРОСА, тогда как со стороны значений дообучение безвредно. У нас `AltaySharedKVAttentionReplay` правит ровно запрос и выход, а ключи и значения оставляет как есть — та же геометрия риска. Обнуление вклада проверяет это без внешнего чекпоинта: ZENIT_ALTAY_ATTENTION_ALPHA_SCALE=0 снять вклад двух суперблоков внимания ZENIT_ALTAY_RECURRENT_ALPHA_SCALE=0 снять вклад шести рекуррентных повторов Множитель применяется **в месте использования**, а не к самому параметру. Замерено 26.07.2026, почему так: первая редакция умножала `altay_alpha` в конструкторе, но веса накладки грузятся ПОЗЖЕ, через `admitted_named_parameters`, и затирали множитель. Прогон выглядел отрицательным (карта игл совпала знак в знак), а был недействительным. То же самое, что случилось в тот день с хуком распаковки эмбеддингов: правка до того, как веса легли на место, ничего не значит. Это диагностика, а не режим работы: значение, отличное от 1.0, означает, что модель работает не так, как отгружена. """ raw = os.environ.get(variable, "").strip() if not raw: return 1.0 try: scale = float(raw) except ValueError: _logger().warning("%s: игнорирую нечисловое %s=%r", ALTAY_ARCHITECTURE_ID, variable, raw) return 1.0 if scale != 1.0: _logger().warning( "%s: вклад накладки умножен на %.3f по %s — ДИАГНОСТИКА, модель " "работает не так, как отгружена", ALTAY_ARCHITECTURE_ID, scale, variable) return scale PHYSICAL_BASE_LAYERS = 64 LOGICAL_LAYERS = 72 OVERLAY_RANK = 64 OVERLAY_TENSOR_COUNT = 28 EXTRA_RECURRENT_SLOTS = (64, 65, 66, 67, 68, 69) class VllmAltayError(RuntimeError): """The admitted ALTAY runtime contract cannot be executed.""" def _require_parent_topology(text_model: Qwen3_5Model) -> None: layer_types = list(getattr(text_model.config, "layer_types", ())) expected = [ "full_attention" if index % 4 == 3 else "linear_attention" for index in range(PHYSICAL_BASE_LAYERS) ] if len(text_model.layers) != PHYSICAL_BASE_LAYERS or layer_types != expected: raise VllmAltayError( "ALTAY requires the verified 64-layer 3-GDN+1-attention parent" ) def _layer_root(prefix: str) -> str: marker = ".layers." if marker not in prefix: raise VllmAltayError(f"unexpected vLLM layer prefix: {prefix!r}") return prefix.split(marker, 1)[0] def _plain_linear( in_features: int, out_features: int, *, reference: torch.Tensor, ) -> nn.Linear: """Create one unquantized admitted sidecar matrix on the model device.""" dtype = ( reference.dtype if reference.dtype.is_floating_point else torch.bfloat16 ) layer = nn.Linear( in_features, out_features, bias=False, device=reference.device, dtype=dtype, ) # A missing sidecar must be inert until load_weights rejects it. nn.init.zeros_(layer.weight) return layer def _scale_branch(source_layer: Any, value: torch.Tensor, name: str) -> torch.Tensor: if not bool(getattr(source_layer, "layer_scale", False)): return value scale = getattr(source_layer, name).to(value.dtype) if value.ndim == 2: scale = scale[0] return value * (scale + 1) def _combine_layer_output( source_layer: Any, attention_output: torch.Tensor, residual: torch.Tensor, ) -> torch.Tensor: attention_output = _scale_branch( source_layer, attention_output, "attn_layer_scale", ) mlp_input, residual = source_layer.post_attention_layernorm( attention_output, residual, ) mlp_output = source_layer.mlp(mlp_input) mlp_output = _scale_branch(source_layer, mlp_output, "ffn_layer_scale") return residual + mlp_output def _clone_gdn_with_private_state( source: Any, *, cache_index: int, vllm_config: Any, ) -> Any: """Share all GDN parameters while registering a distinct cache identity.""" clone = copy.copy(source) root = _layer_root(source.prefix) prefix = f"{root}.layers.{cache_index}.linear_attn" clone.prefix = prefix clone.layer_idx = cache_index # copy.copy retains the source-bound method object; bind the selected # platform implementation to the clone so it reads clone.prefix/cache. method_name = getattr(source._forward_method, "__name__", "") if not method_name or not hasattr(clone, method_name): raise VllmAltayError("cannot rebind the source GDN forward method") clone._forward_method = getattr(clone, method_name) context = vllm_config.compilation_config.static_forward_context if prefix in context: raise VllmAltayError(f"duplicate ALTAY recurrent cache prefix: {prefix}") context[prefix] = clone return clone class AltayRecurrentReplay(nn.Module): """One parameter-tied GDN+MLP replay with a private recurrent state.""" def __init__( self, source_layer: Any, *, cache_index: int, vllm_config: Any, rank: int = OVERLAY_RANK, ) -> None: super().__init__() object.__setattr__(self, "_source_layer", source_layer) replay = _clone_gdn_with_private_state( source_layer.linear_attn, cache_index=cache_index, vllm_config=vllm_config, ) object.__setattr__(self, "_replay_attention", replay) reference = source_layer.input_layernorm.weight hidden = int(source_layer.linear_attn.hidden_size) self.altay_alpha = nn.Parameter( torch.zeros((), device=reference.device, dtype=torch.float32) ) self.delta_down = _plain_linear(hidden, rank, reference=reference) self.delta_up = _plain_linear(rank, hidden, reference=reference) self._alpha_scale = _alpha_scale("ZENIT_ALTAY_RECURRENT_ALPHA_SCALE") @property def source_layer(self) -> Any: return object.__getattribute__(self, "_source_layer") @property def replay_attention(self) -> Any: return object.__getattribute__(self, "_replay_attention") def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: source = self.source_layer residual = hidden_states normalized = source.input_layernorm(hidden_states) attention_output = torch.empty_like(normalized) self.replay_attention( hidden_states=normalized, output=attention_output, ) transformed = _combine_layer_output(source, attention_output, residual) delta = self.delta_up(F.silu(self.delta_down(hidden_states))) branch = transformed - hidden_states + delta alpha = self.altay_alpha.to(hidden_states.dtype) * self._alpha_scale return hidden_states + alpha * branch class AltaySharedKVAttentionReplay(nn.Module): """Fresh Q/O replay against an anchor layer's already-written K/V.""" def __init__( self, source_layer: Any, *, logical_index: int, vllm_config: Any, rank: int = OVERLAY_RANK, ) -> None: super().__init__() object.__setattr__(self, "_source_layer", source_layer) attention = source_layer.self_attn target = attention.attn.layer_name root = _layer_root(target) prefix = f"{root}.layers.{logical_index}.self_attn.attn" self.shared_attention = Attention( num_heads=attention.num_heads, head_size=attention.head_dim, scale=attention.scaling, num_kv_heads=attention.num_kv_heads, cache_config=vllm_config.cache_config, quant_config=vllm_config.quant_config, prefix=prefix, attn_type=attention.attn.attn_type, kv_sharing_target_layer_name=target, **( { "layer_idx": logical_index, "dual_chunk_attention_config": ( attention.dual_chunk_attention_config ), } if attention.dual_chunk_attention_config else {} ), ) reference = source_layer.input_layernorm.weight hidden = int(attention.hidden_size) qg_width = int(attention.q_size * (2 if attention.attn_output_gate else 1)) o_input = int(attention.num_heads * attention.head_dim) self.altay_alpha = nn.Parameter( torch.zeros((), device=reference.device, dtype=torch.float32) ) self.q_delta_down = _plain_linear(hidden, rank, reference=reference) self.q_delta_up = _plain_linear(rank, qg_width, reference=reference) self.o_delta_down = _plain_linear(o_input, rank, reference=reference) self.o_delta_up = _plain_linear(rank, hidden, reference=reference) self._alpha_scale = _alpha_scale("ZENIT_ALTAY_ATTENTION_ALPHA_SCALE") @property def source_layer(self) -> Any: return object.__getattribute__(self, "_source_layer") def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: source = self.source_layer attention = source.self_attn residual = hidden_states normalized = source.input_layernorm(hidden_states) qkv, _ = attention.qkv_proj(normalized) qg_width = attention.q_size * (2 if attention.attn_output_gate else 1) qg, key, value = qkv.split( [qg_width, attention.kv_size, attention.kv_size], dim=-1, ) qg = qg + self.q_delta_up(F.silu(self.q_delta_down(normalized))) adjusted_qkv = torch.cat((qg, key, value), dim=-1) query, key, value, gate = attention._project_qkv_gate( adjusted_qkv, positions, ) # FlashInfer requires shaped K/V arguments even for a cross-layer # shared cache. The Attention implementation sees # kv_sharing_target_layer_name and therefore skips their cache write; # the actual read is from the physical anchor's existing cache. raw_output = self.shared_attention(query, key, value) gated_output = ( raw_output * torch.sigmoid(gate) if gate is not None else raw_output ) mixed, _ = attention.o_proj(gated_output) mixed = mixed + self.o_delta_up( F.silu(self.o_delta_down(raw_output)) ) transformed = _combine_layer_output(source, mixed, residual) branch = transformed - hidden_states alpha = self.altay_alpha.to(hidden_states.dtype) * self._alpha_scale return hidden_states + alpha * branch class AltayVllmOverlay(nn.Module): """The 28-tensor ALTAY overlay shared by every release quantization.""" def __init__( self, text_model: Qwen3_5Model, *, vllm_config: Any, rank: int = OVERLAY_RANK, ) -> None: super().__init__() self.rank = rank self.sb_a_recurrent = nn.ModuleList( [ AltayRecurrentReplay( text_model.layers[source], cache_index=cache_index, vllm_config=vllm_config, rank=rank, ) for source, cache_index in zip( (28, 29, 30), EXTRA_RECURRENT_SLOTS[:3], strict=True, ) ] ) self.sb_a_attention = AltaySharedKVAttentionReplay( text_model.layers[31], logical_index=70, vllm_config=vllm_config, rank=rank, ) self.sb_b_recurrent = nn.ModuleList( [ AltayRecurrentReplay( text_model.layers[source], cache_index=cache_index, vllm_config=vllm_config, rank=rank, ) for source, cache_index in zip( (60, 61, 62), EXTRA_RECURRENT_SLOTS[3:], strict=True, ) ] ) self.sb_b_attention = AltaySharedKVAttentionReplay( text_model.layers[63], logical_index=71, vllm_config=vllm_config, rank=rank, ) def run_superblock( self, name: str, hidden_states: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: recurrent = getattr(self, f"{name}_recurrent") attention = getattr(self, f"{name}_attention") for replay in recurrent: hidden_states = replay(hidden_states) return attention(hidden_states, positions) def admitted_named_parameters(self): """Yield only the 28 serialized r3 gate/delta tensors. vLLM attention backends may register small implementation-specific scale parameters on ``shared_attention``. Those are runtime state, not ALTAY weights, and are intentionally initialized by vLLM. """ for name, parameter in self.named_parameters(): if ".shared_attention." in f".{name}.": continue yield name, parameter def checkpoint_parameter_names(self) -> set[str]: return { f"model.language_model.altay_overlay.{name}" for name, _ in self.admitted_named_parameters() } def internal_parameter_names(self) -> set[str]: return { f"language_model.model.altay_overlay.{name}" for name, _ in self.admitted_named_parameters() } class LomonosovZenitAltayTextModel(Qwen3_5Model): """Qwen3.5 physical layers plus the admitted logical-depth insertions.""" def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if not get_pp_group().is_first_rank or not get_pp_group().is_last_rank: raise VllmAltayError("ALTAY consumer runtime currently requires PP=1") if intermediate_tensors is not None: raise VllmAltayError("ALTAY consumer runtime does not accept PP tensors") hidden_states = ( inputs_embeds if inputs_embeds is not None else self.embed_input_ids(input_ids) ) residual = None aux_hidden_states = self._maybe_add_hidden_state( [], 0, hidden_states, residual, ) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): hidden_states, residual = layer( positions=positions, hidden_states=hidden_states, residual=residual, ) if layer_idx == 31: combined = hidden_states + residual combined = self.altay_overlay.run_superblock( "sb_a", combined, positions, ) residual = combined hidden_states = torch.zeros_like(combined) elif layer_idx == 63: combined = hidden_states + residual combined = self.altay_overlay.run_superblock( "sb_b", combined, positions, ) residual = combined hidden_states = torch.zeros_like(combined) self._maybe_add_hidden_state( aux_hidden_states, layer_idx + 1, hidden_states, residual, ) hidden_states, _ = self.norm(hidden_states, residual) if aux_hidden_states: return hidden_states, aux_hidden_states return hidden_states def install_vllm_altay( text_model: Qwen3_5Model, *, vllm_config: Any, ) -> AltayVllmOverlay: """Attach the native overlay and six private cache identities in-place.""" _require_parent_topology(text_model) parallel = vllm_config.parallel_config if ( int(parallel.tensor_parallel_size) != 1 or int(parallel.pipeline_parallel_size) != 1 ): raise VllmAltayError( "the certified consumer ALTAY runtime currently requires TP=1, PP=1" ) if hasattr(text_model, "altay_overlay"): raise VllmAltayError("ALTAY overlay is already installed") overlay = AltayVllmOverlay(text_model, vllm_config=vllm_config) text_model.add_module("altay_overlay", overlay) text_model.__class__ = LomonosovZenitAltayTextModel text_model.config.altay_architecture_id = ALTAY_ARCHITECTURE_ID text_model.config.altay_logical_layers = LOGICAL_LAYERS text_model.config.altay_physical_kv_units = PHYSICAL_BASE_LAYERS text_model.config.altay_extra_recurrent_slots = len(EXTRA_RECURRENT_SLOTS) names = overlay.checkpoint_parameter_names() if len(names) != OVERLAY_TENSOR_COUNT: raise VllmAltayError( f"ALTAY overlay tensor count drift: {len(names)}" ) return overlay __all__ = [ "ALTAY_ARCHITECTURE_ID", "AltayVllmOverlay", "LomonosovZenitAltayTextModel", "LOGICAL_LAYERS", "OVERLAY_TENSOR_COUNT", "VllmAltayError", "install_vllm_altay", ]