"""Native vLLM ALTAY model class with in-model NOOSPHERE preprocessing.""" from __future__ import annotations import json import logging import os from pathlib import Path from typing import Any import torch from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.models.qwen3_5 import ( Qwen3_5ForConditionalGeneration, Qwen3_5ProcessingInfo, ) from vllm.model_executor.models.qwen3_vl import ( Qwen3VLDummyInputsBuilder, Qwen3VLMultiModalProcessor, ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.renderers import TokenizeParams from .generate import ( memory_config, protected_token_ids, resolve_mode, ) from .noosphere.additive import build_additive_evidence_plan from .tq_embedding_quant import install_quantized_embedding from .vllm_altay import install_vllm_altay from .vllm_kv_layout import install_altay72_compact_kv_layout from .vllm_turboquant import install_turboquant_hybrid_shape_repair # vLLM attaches its handler to the "vllm" logger and sets propagate=False, so a # logger named after this package has every record dropped. That is why the # NOOSPHERE receipts below were invisible in the engine log even while they were # being emitted, and why a run that engaged the mechanism looked identical to one # that never called it. logger = logging.getLogger("vllm.lomonosov_zenit_altay.model") LOGICAL_CONTEXT_TOKENS = 1_010_000 def _emit_receipt(receipt: dict[str, Any]) -> None: """Log every receipt and optionally persist one for certification.""" payload = json.dumps(receipt, ensure_ascii=False, sort_keys=True) logger.info("ZENIT_NOOSPHERE_RECEIPT %s", payload) destination = os.environ.get("ZENIT_NOOSPHERE_RECEIPT_PATH") if not destination: return # The input processor runs in the front-end process and the model in the # engine core. With one shared path each overwrote the other, so only the # last writer survived and neither could be relied on. Suffix by stage and # pid so every receipt is kept. stage = str(receipt.get("status", "receipt")).lower()[:40] path = Path(destination) path = path.with_name(f"{path.stem}.{stage}.{os.getpid()}{path.suffix}") path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(payload + "\n", encoding="utf-8") os.replace(temporary, path) class LomonosovZenitAltayMultiModalProcessor(Qwen3VLMultiModalProcessor): """Preserve raw input and defer additive evidence to the model class.""" def _cached_apply_hf_processor(self, inputs: Any, timing_ctx: Any): prompt_ids, mm_info, is_update_applied = super()._cached_apply_hf_processor( inputs, timing_ctx, ) config_proxy = type("_ModelProxy", (), {"config": self.info.get_hf_config()})() mode = resolve_mode(config_proxy, os.environ.get("ZENIT_NOOSPHERE_MODE")) config = memory_config(config_proxy) if mode != "auto" or len(prompt_ids) <= config.hot_tokens: receipt = { "status": "RAW_OR_SHORT_CONTEXT", "mode": mode, "logical_raw_tokens": len(prompt_ids), "hot_tokens": config.hot_tokens, "architecture_owned": True, } self.last_noosphere_receipt = receipt # Emit this branch too. Previously it only set the attribute, so a # run in which NOOSPHERE declined to engage was indistinguishable # from a run in which this processor was never called at all - and # vLLM's offline `LLM` API does skip the multimodal processor for # text-only token prompts. Silence is not evidence of engagement. _emit_receipt(receipt) return prompt_ids, mm_info, is_update_applied receipt = { "status": "ADDITIVE_RAW_CONTEXT_DEFERRED_TO_MODEL", "mode": "additive", "logical_raw_tokens": len(prompt_ids), "model_visible_raw_tokens": len(prompt_ids), "raw_tokens_removed": 0, "raw_tokens_reordered": 0, "raw_tokens_replaced": 0, "architecture_owned": True, "quantization_independent": True, "vllm_native_model": True, } self.last_noosphere_receipt = receipt _emit_receipt(receipt) return prompt_ids, mm_info, is_update_applied class LomonosovZenitAltayProcessingInfo(Qwen3_5ProcessingInfo): """Expose the logical context limit to the renderer, not the KV allocator.""" def get_default_tok_params(self) -> TokenizeParams: return TokenizeParams( max_total_tokens=LOGICAL_CONTEXT_TOKENS, do_lower_case=False, add_special_tokens=True, ) @MULTIMODAL_REGISTRY.register_processor( LomonosovZenitAltayMultiModalProcessor, info=LomonosovZenitAltayProcessingInfo, dummy_inputs=Qwen3VLDummyInputsBuilder, ) class LomonosovZenitAltayForConditionalGeneration( Qwen3_5ForConditionalGeneration ): """Qwen3.5 kernels plus native ALTAY depth and additive model memory.""" noosphere_model_owned = True def __init__(self, *, vllm_config: Any, prefix: str = "model"): # vLLM starts EngineCore with ``spawn`` after parent-side plugin # discovery. Install the cache planner inside that process as well, # before cache groups are created after model loading. install_altay72_compact_kv_layout(force=True) # Must run inside the EngineCore process before KV-cache tensors are # reshaped. This makes raw TurboQuant operation download-and-run, # without the historical sitecustomize/PYTHONPATH launcher repair. install_turboquant_hybrid_shape_repair(vllm_config) super().__init__(vllm_config=vllm_config, prefix=prefix) # The upstream Qwen3.5 implementation builds the embedding without a # quant_config, so a checkpoint shipping a quantised embedding table # cannot load. Rebuild it here, before weights arrive; a no-op when the # checkpoint does not declare one. install_quantized_embedding(self, vllm_config) install_vllm_altay( self.language_model.model, vllm_config=vllm_config, ) config_proxy = type("_ModelProxy", (), {"config": self.config})() self._noosphere_mode = resolve_mode( config_proxy, os.environ.get("ZENIT_NOOSPHERE_MODE"), ) self._noosphere_prompt_buffer: torch.Tensor | None = None self._noosphere_prompt_tokens = 0 self._noosphere_bias_indices: torch.Tensor | None = None self._noosphere_bias_values: torch.Tensor | None = None self._noosphere_additive_receipt: dict[str, Any] | None = None self._noosphere_prompt_completed = False self._noosphere_last_decode_tokens = 0 self._noosphere_last_prefill_tokens = 0 self._noosphere_declined_logged = False self._noosphere_counts_logged = False self._noosphere_skip_logged = False self._noosphere_pending_ids: torch.Tensor | None = None @staticmethod def _noosphere_forward_token_counts() -> tuple[int, int, int | None]: """Return (decode, prefill, requests) from the active vLLM context.""" if not is_forward_context_available(): return 0, 0, None context = get_forward_context() metadata: Any = context.attn_metadata if isinstance(metadata, list): metadata = metadata[0] if metadata else {} if isinstance(metadata, dict): metadata = next(iter(metadata.values()), None) decode_tokens = int(getattr(metadata, "num_decode_tokens", 0) or 0) # ``num_prefill_tokens`` is a FlashAttention-metadata field. The # TurboQuant backend this checkpoint ships with does not define it, so # reading it with a default of zero returned zero on every forward pass: # the raw-prompt capture below never ran, _noosphere_prompt_tokens stayed # at zero, and _build_noosphere_additive_plan therefore exited on its # first guard. The mechanism this model is named for could not engage at # all with its own attention backend, silently, because getattr does not # complain about a missing field. # # Derive it from what TurboQuantMetadata does publish instead: # num_actual_tokens counts every token in the batch excluding padding, # and decode tokens come first, so the remainder is the prefill. prefill_tokens = int(getattr(metadata, "num_prefill_tokens", 0) or 0) if not prefill_tokens: actual_tokens = int(getattr(metadata, "num_actual_tokens", 0) or 0) prefill_tokens = max(actual_tokens - decode_tokens, 0) descriptor = context.batch_descriptor requests = ( int(descriptor.num_reqs) if descriptor is not None and descriptor.num_reqs is not None else None ) return decode_tokens, prefill_tokens, requests def _reset_noosphere_additive(self) -> None: self._noosphere_prompt_tokens = 0 self._noosphere_bias_indices = None self._noosphere_bias_values = None self._noosphere_additive_receipt = None self._noosphere_prompt_completed = False self._noosphere_declined_logged = False def _ensure_noosphere_prompt_buffer(self, device: torch.device) -> torch.Tensor: buffer = self._noosphere_prompt_buffer if buffer is None or buffer.device != device: buffer = torch.empty( LOGICAL_CONTEXT_TOKENS, dtype=torch.int32, device=device, ) self._noosphere_prompt_buffer = buffer return buffer def _build_noosphere_additive_plan(self) -> None: # Every one of these used to be a bare return, and each of them has in # fact fired in production while the run looked healthy: the mechanism # this model is named for declined to engage and said nothing. Name the # reason instead - once per prompt, so a million-token prefill does not # turn the log into a transcript. hot_tokens = memory_config(self).hot_tokens reason = None if self._noosphere_mode != "auto": reason = f"mode={self._noosphere_mode}" elif self._noosphere_bias_indices is not None: reason = "plan already built for this prompt" elif self._noosphere_prompt_tokens <= hot_tokens: reason = ( f"captured {self._noosphere_prompt_tokens} raw tokens, which is " f"within the literal hot window of {hot_tokens}" ) if reason is not None: if not self._noosphere_declined_logged: self._noosphere_declined_logged = True logger.info("NOOSPHERE additive plan not built: %s", reason) return buffer = self._noosphere_prompt_buffer if buffer is None: return logical_tokens = self._noosphere_prompt_tokens raw_ids = [ int(value) for value in buffer[:logical_tokens].detach().cpu().tolist() ] plan = build_additive_evidence_plan( raw_ids, config=memory_config(self), protected_token_ids=protected_token_ids(self), ) self._noosphere_bias_indices = torch.tensor( plan.token_ids, dtype=torch.long, device=buffer.device, ) self._noosphere_bias_values = torch.tensor( plan.token_biases, dtype=torch.float32, device=buffer.device, ) receipt = dict(plan.receipt) receipt.update( { "architecture_owned": True, "quantization_independent": True, "vllm_native_model": True, "gpu_raw_token_buffer_bytes": int( buffer.element_size() * logical_tokens ), } ) self._noosphere_additive_receipt = receipt _emit_receipt(receipt) def _capture_noosphere_raw_prompt( self, input_ids: torch.Tensor | None, positions: torch.Tensor, ) -> None: if self._noosphere_mode != "auto" or input_ids is None: # Both of these are silent refusals on the hot path, and one of them # is easy to hit without noticing: a multimodal model class can be # called with precomputed embeddings and no input_ids at all, and # then the raw prompt is never captured however healthy the run # looks. Say which one fired, once. if not self._noosphere_skip_logged: self._noosphere_skip_logged = True logger.info( "NOOSPHERE raw capture skipped: mode=%s input_ids=%s", self._noosphere_mode, "absent" if input_ids is None else "present", ) return decode_tokens, prefill_tokens, requests = ( self._noosphere_forward_token_counts() ) self._noosphere_last_decode_tokens = decode_tokens self._noosphere_last_prefill_tokens = prefill_tokens if not self._noosphere_counts_logged: # One line, on the first forward that reaches here: without it the # difference between "the counts are wrong" and "the guards rejected # the plan" is invisible from the outside. self._noosphere_counts_logged = True logger.info( "NOOSPHERE first forward: mode=%s decode=%d prefill=%d requests=%s", self._noosphere_mode, decode_tokens, prefill_tokens, requests, ) # A 1.01M request has physical concurrency one on RTX 5090. Refuse # ambiguous interleaved batches instead of mixing separate prompts. if requests is not None and requests != 1: return if prefill_tokens: # A prefill chunk arriving after a decode step is the first chunk of # the next prompt, so that is where the accumulator is cleared. The # previous marker was set in ``compute_logits``, which runs after # *every* chunk rather than at the end of the prompt: the counter was # therefore reset before each chunk and held only the last one. On a # 100,029-token prompt the runtime reported "captured 957 raw # tokens" - exactly the final chunk - and declined to build the plan # because 957 is inside the literal hot window. Using the decode step # as the boundary keeps this free of a GPU->CPU sync on positions. if self._noosphere_prompt_completed: self._reset_noosphere_additive() if input_ids.numel() < decode_tokens + prefill_tokens: # With prompt embeddings enabled the runner embeds only a subset # of the batch's ids, and capturing that subset would record a # prompt the model never saw. Refuse, out loud. if not self._noosphere_declined_logged: self._noosphere_declined_logged = True logger.info( "NOOSPHERE raw capture skipped: got %d ids for a batch of " "%d (%d decode + %d prefill)", input_ids.numel(), decode_tokens + prefill_tokens, decode_tokens, prefill_tokens, ) return prefill_ids = input_ids[ decode_tokens : decode_tokens + prefill_tokens ] start = self._noosphere_prompt_tokens stop = start + int(prefill_ids.numel()) if stop > LOGICAL_CONTEXT_TOKENS: raise RuntimeError( "NOOSPHERE raw capture exceeds the 1,010,000-token contract" ) buffer = self._ensure_noosphere_prompt_buffer(input_ids.device) buffer[start:stop].copy_(prefill_ids.to(dtype=torch.int32)) self._noosphere_prompt_tokens = stop if stop == LOGICAL_CONTEXT_TOKENS: self._build_noosphere_additive_plan() elif decode_tokens: # For shorter long prompts, the first decode step is the first # unambiguous prompt-boundary signal exposed to the model class. self._build_noosphere_additive_plan() self._noosphere_prompt_completed = True def embed_input_ids(self, input_ids: torch.Tensor, *args: Any, **kwargs: Any): """Hold on to the token ids; the capture itself happens in forward. Two facts about this architecture put the ids and the counts in different places, and each of them silently disabled the capture in turn. First, the model is multimodal, so vLLM's runner embeds the tokens before calling ``forward`` and passes ``input_ids=None`` - the runtime reported "raw capture skipped: mode=auto input_ids=absent" on every prompt. Second, this method runs *outside* the forward context, so the attention metadata that says how many tokens are prefill and how many are decode is not there yet: capturing here reported "decode=0 prefill=0 requests=None" and did nothing. So the ids are parked here and consumed in ``forward``, which does run inside the forward context. A reference is enough - forward follows immediately, before the runner refills its input buffer. """ self._noosphere_pending_ids = input_ids return super().embed_input_ids(input_ids, *args, **kwargs) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, intermediate_tensors: Any | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs: object, ): # The ids come either straight from the runner (text-only path) or from # embed_input_ids, which saw them a moment earlier on the multimodal # path. Either way the capture runs here, inside the forward context, # where the prefill/decode split is knowable. pending = self._noosphere_pending_ids self._noosphere_pending_ids = None self._capture_noosphere_raw_prompt( input_ids if input_ids is not None else pending, positions ) return super().forward( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **kwargs, ) def compute_logits( self, hidden_states: torch.Tensor, ) -> torch.Tensor | None: logits = super().compute_logits(hidden_states) indices = self._noosphere_bias_indices values = self._noosphere_bias_values if ( logits is not None and logits.numel() and indices is not None and values is not None and indices.numel() ): logits[:, indices] = ( logits[:, indices] + values.to(dtype=logits.dtype) ) return logits def load_weights(self, weights): overlay = self.language_model.model.altay_overlay expected_checkpoint = overlay.checkpoint_parameter_names() seen_checkpoint: set[str] = set() def tapped_weights(): for name, tensor in weights: if name in expected_checkpoint: seen_checkpoint.add(name) yield name, tensor loaded = super().load_weights(tapped_weights()) missing_checkpoint = expected_checkpoint - seen_checkpoint if missing_checkpoint: raise RuntimeError( "ALTAY sidecar is incomplete or absent: " + ", ".join(sorted(missing_checkpoint)[:4]) ) expected_internal = overlay.internal_parameter_names() missing_internal = expected_internal - set(loaded) if missing_internal: raise RuntimeError( "vLLM did not load the complete ALTAY sidecar: " + ", ".join(sorted(missing_internal)[:4]) ) self.altay_overlay_loaded = True return loaded __all__ = [ "LomonosovZenitAltayForConditionalGeneration", "LomonosovZenitAltayMultiModalProcessor", "LomonosovZenitAltayProcessingInfo", ]