# SPDX-License-Identifier: Apache-2.0 """Position remapping for contexts past the trained window, and a smaller table. Two independent things live here because both are the same one-line hook: the rotary cos/sin table. **Group compression.** The checkpoint was trained at 262,144 positions and the declared 1,010,000 window is a training-free extension. Measured on real prose with fourteen needles: 92.9% at 33K, 78.6% at 250K on the model's own positions, 0% at a million. Following Jet-Long (arXiv:2607.07740), remote positions are aliased onto the trained grid by ``f(p) = p // G`` with ``G = ceil(L / 262144)``, so every angle the model sees is one it was trained on. Jet-Long computes the alias per relative distance and keeps a RoPE-faithful local window, which needs two attention passes and an inclusion-exclusion merge; that is not reachable without rewriting the TurboQuant prefill kernel. What is reachable is the same aliasing on the *frequency* axis instead of the distance axis. DPE (arXiv:2504.18857) measures that RoPE's frequency subspaces do not share one effective length: the fastest pairs stay usable well past the trained window, while the middle ones lose retrieval the moment they leave it. So the fast pairs keep true positions - preserving the local resolution that Jet-Long's local window exists to protect, inside a single attention pass - and only the slower pairs are aliased. Which pairs, and by how much, is set by the environment and must be measured, not assumed: Jet-Long's own ablation shows that removing the local protection entirely (their w0=0) collapses RULER from 76.65 to 4.10 on a pure-softmax model, and whether 48 Gated DeltaNet layers out of 64 substitute for it here is an open question. **Table size.** MRotaryEmbedding multiplies the table by four, because in Qwen2.5-VL the position index tracks video duration. For a 1,010,000-token text window that is 4,040,000 x 64 in bf16 = 517 MB of device memory for positions no text prompt can reach. Capping it returns three quarters of that. Both are off unless asked for, so the default path is byte-identical to stock. """ from __future__ import annotations import logging import math import os logger = logging.getLogger("vllm.lomonosov_zenit_altay.positional") REPAIR_ID = "ZENIT_POSITIONAL_EXTENSION_V1" ENV_GROUP = "ZENIT_POSITION_GROUP" ENV_FIRST_PAIR = "ZENIT_POSITION_GROUP_FIRST_PAIR" ENV_LAST_PAIR = "ZENIT_POSITION_GROUP_LAST_PAIR" ENV_FAITHFUL_PREFIX = "ZENIT_POSITION_FAITHFUL_PREFIX" ENV_CLAMP = "ZENIT_POSITION_CLAMP" ENV_LOGIT_SCALE = "ZENIT_ROPE_LOGIT_SCALE" ENV_TABLE_POSITIONS = "ZENIT_ROPE_TABLE_POSITIONS" # A multi-stage plan, as JSON: # # [{"from": 196608, "group": 2, "first_pair": 15, "last_pair": 21}, # {"from": 262144, "group": 4, "first_pair": 0, "last_pair": 31}] # # The single-stage variables above cannot express what measurement asked for. The # middle frequency pairs want their positions compressed *earlier and more gently* # than the rest — measured at 250,334 tokens, that took retrieval from 10/14 to # 13/14 — while every pair still needs the coarse division beyond the trained # window so a million-token prompt lands on angles the model has seen. # # Breakpoints are read against the ORIGINAL position, not against the running # mapped value. Composing the maps instead would make the second stage never fire: # by the time a position has been halved it no longer reaches the 262,144 mark. ENV_STAGES = "ZENIT_POSITION_STAGES" CHUNK_POSITIONS = 262_144 # Обученный диапазон углов: `original_max_position_embeddings` из конфигурации. # Совпадает с размером куска численно, но это разные величины, и путать их # нельзя: один — сколько строк считать за раз, другой — до какой позиции модель # видела углы при обучении. Проверка `check_trained_span` сверяется со вторым. TRAINED_POSITIONS = 262_144 ENV_TRAINED = "ZENIT_TRAINED_POSITIONS" # Фактическое окно обслуживания. Выставляет `serving_profiles.apply_profile`; # нужно только проверке `check_trained_span`, на построение таблицы не влияет. ENV_WINDOW_HINT = "ZENIT_EFFECTIVE_MAX_MODEL_LEN" _installed = False _original_compute = None def _read_int(name: str, default: int = 0) -> int: raw = os.environ.get(name, "").strip() if not raw: return default try: return int(raw) except ValueError: logger.warning("%s: ignoring non-integer %s=%r", REPAIR_ID, name, raw) return default def _read_float(name: str, default: float = 0.0) -> float: raw = os.environ.get(name, "").strip() if not raw: return default try: return float(raw) except ValueError: logger.warning("%s: ignoring non-numeric %s=%r", REPAIR_ID, name, raw) return default def stage_plan(pairs: int) -> list[dict] | None: """Parse ``ZENIT_POSITION_STAGES`` into normalised, ordered stages. Returns None when unset or unusable, so the caller falls back to the single-stage variables. A malformed plan is refused loudly rather than half-applied: a position map that is *almost* what was measured is worth nothing, because every number we have was measured against a specific one. """ import json raw = os.environ.get(ENV_STAGES, "").strip() if not raw: return None try: parsed = json.loads(raw) except ValueError as exc: logger.warning("%s: ignoring unparseable %s (%s)", REPAIR_ID, ENV_STAGES, exc) return None if not isinstance(parsed, list) or not parsed: logger.warning("%s: %s must be a non-empty list", REPAIR_ID, ENV_STAGES) return None stages: list[dict] = [] for entry in parsed: if not isinstance(entry, dict) or "from" not in entry or "group" not in entry: logger.warning("%s: stage %r needs 'from' and 'group'", REPAIR_ID, entry) return None try: start = int(entry["from"]) group = int(entry["group"]) first = max(0, int(entry.get("first_pair", 0))) last = int(entry.get("last_pair", pairs - 1)) except (TypeError, ValueError): logger.warning("%s: stage %r has non-integer fields", REPAIR_ID, entry) return None if last < 0: last = pairs - 1 last = min(last, pairs - 1) if start < 0 or group < 1 or first > last: logger.warning("%s: stage %r is out of range", REPAIR_ID, entry) return None stages.append({"from": start, "group": group, "first_pair": first, "last_pair": last}) stages.sort(key=lambda s: s["from"]) return stages def _map_scalar(position: int, stages: list[dict]) -> int: """То же отображение, что у `_piecewise`, но для одного числа. Нужно затем, чтобы посчитать `f(L)` на старте и не строить ради этого таблицу на миллион строк. Логика повторена буква в букву: порог сверяется с ИСХОДНОЙ позицией, деление целочисленное, смещения накапливаются. """ prev_from, out_at_prev, group = 0, 0, 1 for stage in stages: start = int(stage["from"]) if position < start: return out_at_prev + (position - prev_from) // group out_at_prev += (start - prev_from) // group prev_from = start group = int(stage["group"]) return out_at_prev + (position - prev_from) // group def check_trained_span(by_set: dict[tuple, list[int]], stages: list[dict], positions: int) -> list[str]: """Влезает ли последняя позиция окна в обученный диапазон. Замерено 26.07.2026, зачем это нужно. Релизный план держал верным префикс ровно в 262 144 позиции — то есть тратил всё обученное окно — и лишь потом начинал сжимать. Что бы ни шло дальше, оно ложилось сверху, и последний токен неизбежно оказывался за краем: на 400 000 в 296 893 (+13%), на миллионе в 447 858 (+71%). Удержание на миллионе при этом стояло на 14.3%. Ошибка целиком арифметическая и видна до запуска, но её ничто не проверяло. Теперь проверяет. Порог не выдумывается: он равен `original_max_position_embeddings` из конфигурации. Возвращает список строк с жалобами; пустой список — всё внутри. """ trained = _read_int(ENV_TRAINED, TRAINED_POSITIONS) if trained <= 0 or positions <= 0: return [] # Сверяться надо с окном, а не с длиной таблицы. Замерено 26.07.2026: при # mrope таблица строится на 4 040 004 позиции, то есть вчетверо шире окна # 1 010 001, и первая редакция этой проверки жаловалась на ЛЮБОЙ план, # включая верный (f(4 040 003) «за обученными на 1441%»). Ни один промпт до # таких позиций не доходит: их занимают три секции mrope. # # Окно приходит из профиля через `ZENIT_EFFECTIVE_MAX_MODEL_LEN`. Если его # нет — сверяемся по длине таблицы и говорим об этом в жалобе, чтобы # завышенный процент нельзя было принять за настоящий. window = _read_int(ENV_WINDOW_HINT) by_table = window <= 0 if by_table: window = positions last = window - 1 complaints = [] for key, cols in sorted(by_set.items(), key=lambda kv: min(kv[1])): if not key: mapped = last # пары без ступеней идут на истинных позициях else: mapped = _map_scalar(last, [stages[i] for i in key]) if mapped <= trained: continue over = mapped - trained line = (f"pairs {min(cols)}..{max(cols)}: f({last:,}) = {mapped:,}, " f"beyond the trained {trained:,} by {over:,} " f"({100.0 * over / trained:.0f}%)") if by_table: line += ("; measured against the ROPE TABLE length, not the serving " "window - set ZENIT_EFFECTIVE_MAX_MODEL_LEN for an exact check") # Подсказка, а не догадка: при группе g последней ступени верный # префикс P должен удовлетворять P + (L-P)//g <= trained. if key: g = int(stages[key[-1]]["group"]) if g > 1: limit = int((trained - last / g) / (1 - 1 / g)) line += (f"; with group {g} the faithful prefix must not exceed " f"{max(limit, 0):,}") complaints.append(line) return complaints def _piecewise(rows, stages: list[dict]): """Map original positions through ordered stages, breakpoints on the input. ``f(p) = out_at_prev + (p - prev_from) // group``, where the walk advances through every stage whose ``from`` the position has passed. Integer division matches the single-stage path exactly, so a one-element plan reproduces it. """ import torch out = torch.zeros_like(rows) prev_from = 0.0 out_at_prev = 0.0 group = 1 covered = torch.zeros_like(rows, dtype=torch.bool) for stage in stages: start = float(stage["from"]) # Positions before this boundary are finished by the running segment. here = (~covered) & (rows < start) if bool(here.any()): out[here] = out_at_prev + torch.div( rows[here] - prev_from, group, rounding_mode="floor") covered |= here out_at_prev = out_at_prev + (start - prev_from) // group prev_from = start group = int(stage["group"]) rest = ~covered if bool(rest.any()): out[rest] = out_at_prev + torch.div( rows[rest] - prev_from, group, rounding_mode="floor") return out def group_plan(pairs: int) -> tuple[int, int, int] | None: """(group, first pair, last pair inclusive), or None when disabled. ``pairs`` is the number of frequency pairs, i.e. rotary_dim // 2. Bounds are clamped rather than rejected so a plan written for one head size does not silently do nothing on another. """ group = _read_int(ENV_GROUP) if (group <= 1 and _read_int(ENV_CLAMP) <= 0 and _read_float(ENV_LOGIT_SCALE, 1.0) == 1.0): return None group = max(group, 1) first = max(0, _read_int(ENV_FIRST_PAIR, 0)) last = _read_int(ENV_LAST_PAIR, pairs - 1) if last < 0: last = pairs - 1 last = min(last, pairs - 1) if first > last: logger.warning( "%s: empty pair range %d..%d, group compression disabled", REPAIR_ID, first, last, ) return None return group, first, last def _build_cache_staged(module, inv_freq, pairs: int, positions: int, stages: list[dict]): """Build the cos/sin table from a multi-stage plan. Each frequency pair gets its own piecewise map, built from the stages that name it. Pairs no stage names keep true positions. Pairs sharing a stage set are computed once and shared, so the usual two-stage plan costs two passes rather than thirty-two. """ import torch by_set: dict[tuple, list[int]] = {} for pair in range(pairs): applicable = tuple( i for i, s in enumerate(stages) if s["first_pair"] <= pair <= s["last_pair"] ) by_set.setdefault(applicable, []).append(pair) described = ", ".join( f"pairs {min(cols)}..{max(cols)} -> " + (" then ".join(f"//{stages[i]['group']} from {stages[i]['from']:,}" for i in key) if key else "true positions") for key, cols in sorted(by_set.items(), key=lambda kv: min(kv[1])) ) logger.info("%s: staged position plan: %s", REPAIR_ID, described) # Жаловаться громко, но не падать: план может быть заведомо выходящим за # обученное окно — так мерили и так замерена цена промаха. Молчать нельзя, # потому что именно молчание стоило миллиону 14.3%. for complaint in check_trained_span(by_set, stages, positions): logger.warning("%s: position map leaves the trained range - %s", REPAIR_ID, complaint) slices = [] for start in range(0, positions, CHUNK_POSITIONS): stop = min(start + CHUNK_POSITIONS, positions) rows = torch.arange(start, stop, dtype=torch.float) mapped = rows[:, None].repeat(1, pairs) for key, cols in by_set.items(): if not key: continue aliased = _piecewise(rows, [stages[i] for i in key]) index = torch.tensor(cols, dtype=torch.long) mapped[:, index] = aliased[:, None].expand(-1, len(cols)) freqs = mapped * inv_freq[None, :] slices.append(torch.cat((freqs.cos(), freqs.sin()), dim=-1)) cache = torch.cat(slices, dim=0) # Everything below mirrors the single-stage path exactly. It has to: the two # branches must differ only in how positions are mapped, or a comparison # between them measures the plumbing instead of the map. tau = _read_float(ENV_LOGIT_SCALE, 1.0) if tau != 1.0: cache = cache * math.sqrt(tau) logger.info("%s: rotary logit temperature tau=%.4f (table x%.4f)", REPAIR_ID, tau, math.sqrt(tau)) mscale = float(getattr(module, "mscale", 1.0) or 1.0) if mscale != 1.0: cache = cache * mscale logger.warning( "%s: staged compression on top of YaRN mscale=%.4f - unmeasured " "combination", REPAIR_ID, mscale, ) return cache def _build_cache(module): """Rebuild the cos/sin table with per-pair position aliasing.""" import torch inv_freq = module._compute_inv_freq(module.base) pairs = int(inv_freq.numel()) positions = int(module.max_position_embeddings) stages = stage_plan(pairs) if stages is not None: return _build_cache_staged(module, inv_freq, pairs, positions, stages) plan = group_plan(pairs) if plan is None: return None group, first, last = plan # An optional faithful prefix. Jet-Long protects a RoPE-faithful window # around each query, which is a *relative* construction needing two # attention passes. A table can only protect a fixed range of absolute # positions - and the right range is the beginning, because every prompt # starts at zero. With ``faithful_prefix`` set to the trained window, a # prompt that fits inside it runs on exactly the positions the model was # trained on, and only what lies beyond is aliased: # # f(p) = p for p < w # f(p) = w + (p - w) // G for p >= w # # Without this, a profile built for a million-token window would hand # coarsened positions to a ten-thousand-token prompt, which is a loss for # no gain - the aliasing exists only to keep *distant* angles trained. faithful = max(0, _read_int(ENV_FAITHFUL_PREFIX)) # Clamping instead of dividing: f(p) = min(p, clamp - 1). Beyond the trained # window every relative distance collapses onto one angle, so positional # resolution out there is zero - but no angle is ever out of distribution, # which is the opposite corner of the trade-off from dividing. Worth trying # because 48 of the 64 layers carry no positions at all and manage: hybrids # that reach a million (Solar Open 2) do it with no positional encoding in # their softmax layers. Here it costs one line and one measurement. clamp_at = max(0, _read_int(ENV_CLAMP)) # Built in row chunks: the full [4,040,000 x 32] float32 intermediates would # cost half a gigabyte each, three at a time, for a table that ends up in # bf16 anyway. slices = [] for start in range(0, positions, CHUNK_POSITIONS): stop = min(start + CHUNK_POSITIONS, positions) rows = torch.arange(start, stop, dtype=torch.float) if clamp_at: aliased = rows.clamp(max=float(clamp_at - 1)) elif faithful: beyond = torch.div( (rows - faithful).clamp(min=0), group, rounding_mode="floor" ) aliased = torch.where(rows < faithful, rows, faithful + beyond) else: aliased = torch.div(rows, group, rounding_mode="floor") mapped = rows[:, None].repeat(1, pairs) mapped[:, first:last + 1] = aliased[:, None].expand(-1, last + 1 - first) freqs = mapped * inv_freq[None, :] slices.append(torch.cat((freqs.cos(), freqs.sin()), dim=-1)) cache = torch.cat(slices, dim=0) # Length-adaptive attention temperature. Scaling the cos/sin table scales # both the rotated query and the rotated key, so a table factor of sqrt(tau) # multiplies the rotary part of the logit by tau. Derivation: with one target # exceeding N-1 distractors by a margin Delta, the target's softmax mass is # preserved as N grows past the trained window W when # tau*(N) = 1 + log((N-1)/(W-1)) / Delta, tau*(N<=W) = 1. # At N = 1,010,000 with the margin implied by our own needle curve # (Delta/sigma = 4.666) this gives 1.2890 - and YaRN ships 1.1349^2 = 1.2880. # So YaRN's number is about right for exactly a million and wrong everywhere # else, which is what our 250K ladder measured when removing it helped. # # Honest limitation: only 64 of 256 head dimensions rotate, so this scales # the rotary contribution to the logit and leaves the other three quarters # alone. It is a partial temperature, not a clean one. YaRN has the same flaw. tau = _read_float(ENV_LOGIT_SCALE, 1.0) if tau != 1.0: cache = cache * math.sqrt(tau) logger.info("%s: rotary logit temperature tau=%.4f (table x%.4f)", REPAIR_ID, tau, math.sqrt(tau)) mscale = float(getattr(module, "mscale", 1.0) or 1.0) if mscale != 1.0: # Stock YaRN folds its logit gain into the table. Kept for faithfulness, # and reported: aliasing on top of a rescaled table is not a # configuration any measurement here has blessed. cache = cache * mscale logger.warning( "%s: group compression on top of YaRN mscale=%.4f - unmeasured " "combination", REPAIR_ID, mscale, ) logger.info( "%s: pairs %d..%d of %d over %d positions, %s, mscale=%.4f", REPAIR_ID, first, last, pairs, positions, f"clamped at {clamp_at}" if clamp_at else f"aliased by //{group}, faithful prefix {faithful}", mscale, ) return cache def install() -> bool: """Patch the multimodal rotary embedding. Idempotent.""" global _installed, _original_compute if _installed: return True try: from vllm.model_executor.layers.rotary_embedding.mrope import ( MRotaryEmbedding, ) except Exception as error: # pragma: no cover - stock vLLM always has it logger.warning("%s: not installed (%s)", REPAIR_ID, error) return False _original_compute = MRotaryEmbedding._compute_cos_sin_cache def _compute_cos_sin_cache(self): # The video-duration headroom is a property of the table, not of the # window, so it is capped before anything else touches the size. wanted = _read_int(ENV_TABLE_POSITIONS) if wanted > 0 and wanted < int(self.max_position_embeddings): logger.info( "%s: rotary table capped at %d positions (was %d)", REPAIR_ID, wanted, self.max_position_embeddings, ) self.max_position_embeddings = wanted try: cache = _build_cache(self) except Exception as error: logger.warning( "%s: falling back to the stock table (%s)", REPAIR_ID, error ) cache = None if cache is not None: return cache return _original_compute(self) MRotaryEmbedding._compute_cos_sin_cache = _compute_cos_sin_cache _installed = True return True def describe() -> str: """Name for the plugin's startup log; the plugin appends the state.""" group = _read_int(ENV_GROUP) table = _read_int(ENV_TABLE_POSITIONS) clamp_at = _read_int(ENV_CLAMP) tau = _read_float(ENV_LOGIT_SCALE, 1.0) if group <= 1 and table <= 0 and clamp_at <= 0 and tau == 1.0: return f"{REPAIR_ID}(idle)" parts = [] if clamp_at > 0: parts.append(f"clamp={clamp_at}") if tau != 1.0: parts.append(f"tau={tau:g}") if group > 1: first = max(0, _read_int(ENV_FIRST_PAIR, 0)) last = _read_int(ENV_LAST_PAIR, -1) span = f"{first}..{last}" if last >= 0 else f"{first}..end" parts.append(f"group=//{group} pairs={span}") faithful = _read_int(ENV_FAITHFUL_PREFIX) if faithful > 0: parts.append(f"faithful={faithful}") if table > 0: parts.append(f"table={table}") return f"{REPAIR_ID}({','.join(parts)})"