"""Model-owned serving defaults for LOMONOSOV ZENIT ALTAY. A consumer who runs ``vllm serve `` must get a working engine, not a tuning exercise. The checkpoint declares ``max_position_embeddings`` of 1,010,000, so vLLM's own default would try to allocate a full million-token KV cache and die on a 32 GiB card. The serving parameters that make the million work are properties of this architecture, not of the user's shell history. This module lets the model carry them. ``config.json`` declares ``lomonosov_serving_profiles``; the plugin installs a wrapper around ``EngineArgs.create_engine_config`` - the single funnel used by both ``vllm serve`` and the offline ``LLM`` API - and fills in only the fields the caller left at their vLLM default. Anything passed explicitly on the command line always wins, and every substitution is logged. Profile selection order: 1. ``ZENIT_SERVING_PROFILE`` environment variable, when set; 2. otherwise, if the caller asked for an explicit ``max_model_len``, the smallest profile that can hold it; 3. otherwise the widest profile whose declared ``requires_free_bytes`` fits the memory this card actually has; 4. otherwise the profile named by ``default_profile``. Step 3 exists because the default context is 1,010,000 tokens, and that window needs more memory than a 24 GiB card owns. Without it the first command a 4090 owner types would end in an out-of-memory error. The device is probed through NVML, which reads the driver without creating a CUDA context in the launcher process; when NVML is unavailable the step is skipped rather than guessed at. """ from __future__ import annotations import json import os from pathlib import Path from typing import Any REPAIR_ID = "ZENIT_MODEL_OWNED_SERVING_PROFILES_V1" CONFIG_KEY = "lomonosov_serving_profiles" ARCHITECTURE_ID = "LomonosovZenitAltayForConditionalGeneration" ENV_PROFILE = "ZENIT_SERVING_PROFILE" # EngineArgs fields that can be filled only when the caller left the vLLM # default in place. "sentinel" is the value that means "not set by the user". _FILLABLE: tuple[tuple[str, Any], ...] = ( ("max_model_len", None), ("kv_cache_dtype", "auto"), ("max_num_batched_tokens", None), ("max_num_seqs", None), ("kv_cache_memory_bytes", None), ("enable_prefix_caching", None), # vLLM decides this one for itself when it is left at None. Its choice # suits many short sequences; a single request being prefilled across a # million tokens is the opposite case, and the profile says so. ("async_scheduling", None), ) # Applied only while the caller has left it at vLLM's default, like the group # below - but it is a boolean whose default is False, so it needs the # default-equal treatment rather than a None sentinel. # # Why the model owns this. vLLM profiles the vision tower at startup by pushing # a dummy batch of images through it, and on this checkpoint that forward pass # dies in humming_gemm with CUDA_ERROR_INVALID_VALUE - the kernel asks for more # shared memory than the card allows. Every measurement all day passed # skip_mm_profiling explicitly and so never met it, while `vllm serve ` # with no arguments met it every single time. Four launch-path bugs were found # this way in one hour; this was the fourth. # Fields whose vLLM default is indistinguishable from an explicit value. They # are applied only when they still equal that default, and always logged. # # ``gpu_memory_utilization`` belongs here for a concrete reason: the million-token # window needs 29.75 GiB, and vLLM's default of 0.9 leaves only 28.7 GiB on a # 32 GiB card. Left alone, the headline context would never be selected on the # hardware it was built for. _DEFAULT_EQUAL_ONLY: tuple[str, ...] = ( "enforce_eager", "swap_space", "gpu_memory_utilization", "skip_mm_profiling", ) # vLLM's own default, used to tell "the caller chose this" from "the caller said # nothing". They are indistinguishable in the dataclass, so the profile treats an # untouched value as unset, the same rule as every other field here. # # Read from vLLM rather than remembered, because remembering it was a real bug: # the constant said 0.9 while vLLM 0.25.1 ships 0.92, so every untouched launch # looked like a deliberate choice, the profile's utilisation was never applied, # and a flagless start on a desktop card died with "free memory is less than # desired GPU memory utilization". The same file already read the live default # in one place and the frozen one in another, which is how the two disagreed. VLLM_DEFAULT_UTILIZATION = 0.9 # fallback only, when vLLM cannot be imported def vllm_default_utilization() -> float: try: from vllm.engine.arg_utils import EngineArgs field = EngineArgs.__dataclass_fields__.get("gpu_memory_utilization") if field is not None and field.default is not None: return float(field.default) except Exception: pass return VLLM_DEFAULT_UTILIZATION # Key a profile may declare to state how much device memory it needs, in bytes, # measured end to end (weights, KV cache, activation peak, CUDA context). BUDGET_KEY = "requires_free_bytes" def _visible_device_index() -> int: """First device vLLM will use, honouring ``CUDA_VISIBLE_DEVICES``.""" visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() if not visible: return 0 first = visible.split(",")[0].strip() try: return int(first) except ValueError: # UUID form - NVML index cannot be derived cheaply return 0 def device_memory() -> tuple[int, int] | None: """Return ``(free_bytes, total_bytes)`` for the first visible GPU. NVML is used on purpose: it answers from the driver without creating a CUDA context, so probing here costs the launcher nothing and does not disturb the worker process that will own the device. ``None`` means "unknown" - every caller must treat that as a reason to skip, never to assume. """ try: import pynvml except ImportError: return None try: pynvml.nvmlInit() except Exception: # noqa: BLE001 - no driver, container without NVML, ... return None try: handle = pynvml.nvmlDeviceGetHandleByIndex(_visible_device_index()) info = pynvml.nvmlDeviceGetMemoryInfo(handle) return int(info.free), int(info.total) except Exception: # noqa: BLE001 return None finally: try: pynvml.nvmlShutdown() except Exception: # noqa: BLE001 pass # How much of the free memory a clamped utilisation may claim. The desktop # compositor grows while the engine loads: measured twice in one evening, a run # that saw 29.06 GiB free at selection had 28.45 GiB by the time vLLM checked, # and refused to start. Two per cent of a 32 GiB card is about 0.6 GiB, which # covers that drift. _UTILIZATION_SAFETY = 0.98 def memory_budget( free_bytes: int, total_bytes: int, gpu_memory_utilization: float | None ) -> int: """Memory a profile may actually claim on this card. vLLM sizes itself against a fraction of *total* memory, but it cannot use what another process already holds - a desktop session, another model. The budget is therefore the smaller of the two. The free side carries the same safety margin the utilisation clamp uses, and it has to, or the two disagree. Measured: with 30.10 GiB free the million profile passed this test by 0.13 GiB, the clamp then lowered utilisation to 0.926 to fit that same free memory, and the engine died inside the kernel launcher with no room left for working buffers. Selection said yes to a profile that clamping had already made impossible. One margin, both places. """ utilization = gpu_memory_utilization if gpu_memory_utilization else 0.9 usable_free = float(free_bytes) * _UTILIZATION_SAFETY return int(min(float(total_bytes) * float(utilization), usable_free)) def _logger(): from vllm.logger import init_logger # vLLM attaches its handler to the "vllm" logger and sets # propagate=False, so a logger named after this package would have # its records dropped: every INFO line below would be invisible to # the person running the model. Naming it under "vllm." puts it # where the configured handler can see it. return init_logger("vllm.lomonosov_zenit_altay.profiles") def _read_model_config(engine_args: Any) -> dict[str, Any] | None: """Read the checkpoint config without forcing a full ModelConfig build.""" model = getattr(engine_args, "model", None) if not model: return None local = Path(str(model)) / "config.json" if local.is_file(): try: return json.loads(local.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None try: # remote or cached repository from vllm.transformers_utils.config import get_config config = get_config( str(model), trust_remote_code=bool(getattr(engine_args, "trust_remote_code", False)), revision=getattr(engine_args, "revision", None), ) except Exception: # noqa: BLE001 - never break startup over a probe return None to_dict = getattr(config, "to_dict", None) return to_dict() if callable(to_dict) else None def _is_our_architecture(config: dict[str, Any]) -> bool: return ARCHITECTURE_ID in (config.get("architectures") or []) def profile_budget( body: dict[str, Any], device: tuple[int, int], caller_utilization: float | None, ) -> int: """Memory this particular profile would have if it were selected. Each profile may raise ``gpu_memory_utilization`` above vLLM's default, so the budget has to be evaluated per profile rather than once: judging the million-token window against a 0.9 budget would reject it on the very card it was built for. A utilization the caller set explicitly always wins. """ free_bytes, total_bytes = device utilization = caller_utilization if utilization is None: utilization = float( body.get("gpu_memory_utilization", VLLM_DEFAULT_UTILIZATION) ) return memory_budget(free_bytes, total_bytes, utilization) def select_profile( spec: dict[str, Any], requested_max_model_len: int | None, device: tuple[int, int] | None = None, caller_utilization: float | None = None, ) -> tuple[str, dict[str, Any]] | None: """Resolve which declared profile applies to this launch. ``device`` is ``(free_bytes, total_bytes)``; ``None`` means the card could not be probed, in which case memory plays no part in the decision. ``caller_utilization`` is the value the caller passed, or ``None`` if they left vLLM's default in place. """ profiles = spec.get("profiles") or {} if not profiles: return None forced = os.environ.get(ENV_PROFILE) if forced: if forced not in profiles: raise ValueError( f"{ENV_PROFILE}={forced!r} is not declared by the model; " f"available: {sorted(profiles)}" ) return forced, profiles[forced] if requested_max_model_len is not None: fits = [ (int(body.get("max_model_len", 0)), name, body) for name, body in profiles.items() if int(body.get("max_model_len", 0)) >= requested_max_model_len ] if fits: fits.sort(key=lambda row: row[0]) _, name, body = fits[0] return name, body widest = max(profiles.items(), key=lambda kv: int(kv[1].get("max_model_len", 0))) return widest[0], widest[1] if device is not None: rated = [ (int(body.get("max_model_len", 0)), name, body) for name, body in profiles.items() if body.get(BUDGET_KEY) is not None ] if rated: affordable = [ row for row in rated if int(row[2][BUDGET_KEY]) <= profile_budget(row[2], device, caller_utilization) ] if affordable: affordable.sort(key=lambda row: row[0], reverse=True) _, name, body = affordable[0] return name, body # Nothing fits. Offer the cheapest profile so the launch fails, if # it must, with a real out-of-memory message from vLLM rather than # with a window this card was never going to hold. rated.sort(key=lambda row: int(row[2][BUDGET_KEY])) _, name, body = rated[0] return name, body default_name = spec.get("default_profile") if default_name in profiles: return default_name, profiles[default_name] return None def clamp_utilization(body: Mapping[str, Any], device: tuple[int, int] | None) -> float | None: """Lower a profile's utilisation to what this card can actually give. A profile's ``gpu_memory_utilization`` is an upper bound measured on a headless card. vLLM refuses to start when ``free < utilization * total``, so on a machine with a desktop the figure that was right for the measurement is wrong for the user: selection correctly steps down from the million to the 393,216 window and then fails anyway, because that profile also carries 0.95. Out of the box, on the most ordinary configuration there is, the model would not load. That is what this prevents. Only ever downward. If the card has room for the profile's figure, the profile's figure stands. """ wanted = body.get("gpu_memory_utilization") if wanted is None or device is None: return None free_bytes, total_bytes = device if not total_bytes: return None affordable = (float(free_bytes) * _UTILIZATION_SAFETY) / float(total_bytes) if affordable >= float(wanted): return None return max(0.5, round(affordable, 3)) def _warn_if_batch_too_small(engine_args: Any, body: dict[str, Any]) -> None: """Say why a small batch will fail, before vLLM says it cryptically. This architecture aligns KV pages across the sixteen full-attention layers and the forty-eight recurrent ones, so vLLM derives a large block size from the recurrent page - 2064 tokens on this checkpoint - and then asserts In Mamba cache align mode, block_size (2064) must be <= max_num_batched_tokens (1024) deep inside engine start-up, naming neither the model nor the way out. The profiles all declare 4096 for exactly this reason, and the caller's explicit value always wins here by design, so the profile cannot fix it silently - but it can name both numbers first. Warns rather than raises: the block size is computed by vLLM later and is not knowable at this point, so a hard limit here would be a guess. The profile's own value is the one number we know to be safe. """ declared = body.get("max_num_batched_tokens") current = getattr(engine_args, "max_num_batched_tokens", None) if declared is None or current is None: return if int(current) >= int(declared): return _logger().warning( "%s: max_num_batched_tokens=%d was passed explicitly, and this profile " "declares %d. KV pages on this architecture are aligned between the " "full-attention and recurrent layers, so vLLM derives a block size of " "about two thousand tokens and refuses to start when the batch is " "smaller than one page. If start-up fails with \"In Mamba cache align " "mode, block_size (...) must be <= max_num_batched_tokens\", this is " "why: raise it to %d or leave it unset and the profile will.", REPAIR_ID, int(current), int(declared), int(declared), ) def apply_profile(engine_args: Any, name: str, body: dict[str, Any]) -> list[str]: """Fill unset EngineArgs fields from the profile. Returns applied fields.""" from vllm.engine.arg_utils import EngineArgs applied: list[str] = [] # Окно нужно проверке позиционного плана: она иначе сверяется с длиной # таблицы углов, а та при mrope вчетверо шире окна и даёт ложную тревогу. window = body.get("max_model_len") or getattr(engine_args, "max_model_len", None) if window: os.environ.setdefault("ZENIT_EFFECTIVE_MAX_MODEL_LEN", str(int(window))) for field, sentinel in _FILLABLE: if field not in body: continue current = getattr(engine_args, field, sentinel) if current != sentinel: continue # caller was explicit - never override value = body[field] if value is None: continue if field == "async_scheduling" and value is False: # Turning it off is right for one sequence walking through a very # long prompt, which is what these profiles are for. It is the # wrong call for someone serving several requests at once, so the # profile yields to that case rather than slowing it down. seqs = getattr(engine_args, "max_num_seqs", None) if seqs is None: seqs = body.get("max_num_seqs") if seqs is not None and int(seqs) > 1: continue setattr(engine_args, field, value) applied.append(f"{field}={value}") defaults = {f.name: f.default for f in EngineArgs.__dataclass_fields__.values()} for field in _DEFAULT_EQUAL_ONLY: if field not in body: continue if getattr(engine_args, field, None) != defaults.get(field): continue value = body[field] setattr(engine_args, field, value) applied.append(f"{field}={value}") _warn_if_batch_too_small(engine_args, body) noosphere_mode = body.get("noosphere_mode") if noosphere_mode and "ZENIT_NOOSPHERE_MODE" not in os.environ: os.environ["ZENIT_NOOSPHERE_MODE"] = str(noosphere_mode) applied.append(f"noosphere_mode={noosphere_mode}") applied.extend(_apply_alloc_conf(body.get("cuda_alloc_conf"))) applied.extend(_apply_position_plan(body.get("position_plan"))) return applied # PyTorch's caching allocator keeps freed blocks in fixed-size segments, and at # the million-token window that costs more than the window has to spare: on a # card with roughly half a gigabyte held by something else, 377 MiB sat reserved # and unallocated while the engine died needing 80 MiB. Expandable segments give # those bytes back - measured, the same launch went from 80 MiB short to 40, and # with the arena at 13,000,000,000 it starts and answers. # # Neither half works alone. Trimming the arena on its own fails harder, not # softer: without expandable segments the allocator cannot find a large enough # contiguous span and the shortfall grows to 272 MiB. _ALLOC_CONF_ENV = "PYTORCH_CUDA_ALLOC_CONF" def _apply_alloc_conf(value: str | None) -> list[str]: """Publish the profile's allocator setting, without overriding the caller. Set before the engine's first CUDA allocation, which is why it lives here rather than in the documentation: a setting the user has to know about is a setting most users will not have, and the million-token window is exactly the case where they would not find out until the engine died. """ if not value: return [] if os.environ.get(_ALLOC_CONF_ENV): return [] # caller was explicit - never replace their choice os.environ[_ALLOC_CONF_ENV] = str(value) return [f"{_ALLOC_CONF_ENV}={value}"] # The checkpoint is trained at 262,144 positions. A window wider than that needs # its far positions aliased onto the trained grid, and which pairs of RoPE # frequencies are aliased is a measured property of the model, so it belongs to # the profile rather than to the person starting the server. _POSITION_PLAN_ENV = { "group": "ZENIT_POSITION_GROUP", "first_pair": "ZENIT_POSITION_GROUP_FIRST_PAIR", "last_pair": "ZENIT_POSITION_GROUP_LAST_PAIR", "faithful_prefix": "ZENIT_POSITION_FAITHFUL_PREFIX", "clamp": "ZENIT_POSITION_CLAMP", "table_positions": "ZENIT_ROPE_TABLE_POSITIONS", } ENV_POSITION_STAGES = "ZENIT_POSITION_STAGES" def _apply_position_plan(plan: dict[str, Any] | None) -> list[str]: """Publish the profile's position plan, without overriding the caller. Anything the caller already set in the environment wins outright: a person testing a different plan must not have it silently replaced by the profile's. A plan may carry ``stages`` instead of the flat fields. The flat form cannot say what measurement asked for: the middle frequency pairs want compressing earlier and more gently than the rest, which took retrieval at 250,334 tokens from 10/14 to 13/14 and cost nothing at a million. When ``stages`` is present the flat fields are not published at all — sending both would leave the receipt describing a map that did not run. """ if not plan: return [] applied = [] stages = plan.get("stages") if stages: if not os.environ.get(ENV_POSITION_STAGES): os.environ[ENV_POSITION_STAGES] = json.dumps( stages, separators=(",", ":")) applied.append(ENV_POSITION_STAGES) return applied for key, name in _POSITION_PLAN_ENV.items(): if key not in plan: continue if os.environ.get(name): continue os.environ[name] = str(plan[key]) applied.append(f"{name}={plan[key]}") return applied def _gib(value: int | float) -> str: return f"{float(value) / (1024 ** 3):.2f} GiB" def _report_memory_choice( name: str, body: dict[str, Any], spec: dict[str, Any], budget: int ) -> None: """Say plainly which window this card gets, and why.""" logger = _logger() need = body.get(BUDGET_KEY) default_name = spec.get("default_profile") window = body.get("max_model_len") if need is not None and int(need) > budget: logger.warning( "%s: this GPU offers %s, and the smallest declared profile %r needs " "%s. Startup is likely to run out of memory. A card with more " "memory, or a smaller build of this model, is required.", REPAIR_ID, _gib(budget), name, _gib(int(need)), ) return if name != default_name: logger.warning( "%s: the default profile %r does not fit this GPU (%s available), " "so %r was selected instead: context %s tokens. Pass an explicit " "--max-model-len, or set %s, to override.", REPAIR_ID, default_name, _gib(budget), name, f"{int(window):,}" if window else "unknown", ENV_PROFILE, ) else: logger.info( "%s: profile %r fits this GPU (%s available, %s required): " "context %s tokens.", REPAIR_ID, name, _gib(budget), _gib(int(need)) if need is not None else "unstated", f"{int(window):,}" if window else "unknown", ) def install_model_owned_serving_defaults() -> bool: """Wrap ``EngineArgs.create_engine_config`` once, for this architecture.""" from vllm.engine.arg_utils import EngineArgs if getattr(EngineArgs, "_zenit_serving_profiles_installed", False): return False original = EngineArgs.create_engine_config def patched(self, *args: Any, **kwargs: Any): try: config = _read_model_config(self) if config and _is_our_architecture(config): spec = config.get(CONFIG_KEY) or {} requested = getattr(self, "max_model_len", None) caller_chose_window = ( requested is not None or bool(os.environ.get(ENV_PROFILE)) ) # The card is probed either way. Two different questions are being # answered with it, and conflating them was a real defect: asking # for a window with --max-model-len, or forcing a profile, used to # skip the probe entirely, so the utilisation clamp below never # fired and the profile's 0.95 went through unchanged. On any card # with a desktop session that is more than is free, and vLLM # refuses to start - measured on a 5090 holding 1.8 GiB of desktop: # "Free memory (29.57/31.39 GiB) is less than desired GPU memory # utilization (0.95, 29.82 GiB)". Choosing a window is the caller's # business; how much of the card is free is not. device = device_memory() current = getattr(self, "gpu_memory_utilization", None) caller_utilization = ( float(current) if current is not None and current != vllm_default_utilization() else None ) # Memory decides the window only when the caller left it open. selection_device = None if caller_chose_window else device selected = select_profile( spec, requested, selection_device, caller_utilization ) if selected is not None: name, body = selected if selection_device is not None: _report_memory_choice( name, body, spec, profile_budget(body, selection_device, caller_utilization), ) if caller_utilization is None: lowered = clamp_utilization(body, device) if lowered is not None: asked = float(body.get("gpu_memory_utilization")) body = dict(body) body["gpu_memory_utilization"] = lowered _logger().warning( "%s: profile %r asks for gpu_memory_utilization " "%.2f, but only %.2f GiB of %.2f is free on this " "card, so %.3f is used instead. Close other GPU " "programs, or run headless, to get the full " "figure back.", REPAIR_ID, name, asked, device[0] / (1024 ** 3), device[1] / (1024 ** 3), lowered, ) applied = apply_profile(self, name, body) logger = _logger() if applied: logger.info( "%s: profile %r supplied %s " "(explicit command-line values were preserved)", REPAIR_ID, name, ", ".join(applied), ) else: logger.info( "%s: profile %r selected; every value was already set " "by the caller", REPAIR_ID, name, ) except ValueError: raise except Exception: # noqa: BLE001 - defaults must never break startup _logger().exception("%s failed; using vLLM defaults", REPAIR_ID) return original(self, *args, **kwargs) EngineArgs.create_engine_config = patched EngineArgs._zenit_serving_profiles_installed = True return True __all__ = [ "ARCHITECTURE_ID", "BUDGET_KEY", "CONFIG_KEY", "ENV_PROFILE", "REPAIR_ID", "apply_profile", "device_memory", "install_model_owned_serving_defaults", "memory_budget", "select_profile", ]