Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

hermes_cli/local_runtime/estimator.py

@@ -0,0 +1,179 @@
+"""Per-layer context-memory estimator + physics check.
+
+The whole-model dense formula misprices 1M-context hybrids by ~100x; the
+per-layer walk fixes that, and every column is measured on real GGUFs:
+
+- full-attention layer: linear in T          (B1: 144.0 KiB/tok on Qwen3-4B
+                                              f16 — formula-exact)
+- SWA layer: capped at the sliding window
+- recurrent layer (n_head_kv == 0): constant (state is ~context-free)
+- q8_0 KV = exactly 34/64 of f16 (holds on CUDA and CPU)
+- weights: exact from the tensor table (within 0.01% of the loader)
+
+The estimator is ADVISORY: fit's allocation is authoritative at launch and
+the touch generation is ground truth after it. Unknown shapes round UP
+(never underestimate memory).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from hermes_cli.local_runtime.gguf import GGUFHeader
+
+# q8_0: 34-byte blocks of 32 f16-equivalent elements (exact).
+_Q8_BYTES_PER_ELEM = 34 / 32
+_F16_BYTES_PER_ELEM = 2.0
+
+# Architectures with a known SWA layer pattern: arch -> fraction of layers
+# that are sliding-window. Unknown SWA archs conservatively treat every
+# layer as full attention (overestimate; safe direction).
+_SWA_LAYER_FRACTION = {"gemma3": 5 / 6, "gemma2": 1 / 2}
+
+# Per-recurrent-layer state allowance (bytes/seq). Deliberately generous —
+# Measured: an entire hybrid slot state is ~99 MB including 8K tokens of
+# full-attn KV, so tens of MiB total is the right order; unknown SSM shapes
+# must never underestimate.
+_RECURRENT_STATE_PER_LAYER = 4 << 20
+
+
+class LayerKind(Enum):
+    FULL = "full"
+    SWA = "swa"
+    RECURRENT = "recurrent"
+
+
+@dataclass
+class ModelProfile:
+    """Everything the policy needs, decoupled from GGUF parsing so the
+    decision-table tests can construct profiles directly (design's
+    verification plan)."""
+
+    name: str
+    weights_bytes: int
+    embd_table_bytes: int
+    n_ctx_train: int
+    layers: list[tuple[LayerKind, int]]   # (kind, kv_bytes_per_token_f16);
+                                          # SWA/recurrent reuse the same
+                                          # per-token figure, capped/ignored
+    swa_window: int = 0
+    moe: bool = False
+    architecture: str = ""
+    n_vocab: int = 0            # prices logits buffers (ubatch x vocab)
+    # Context-cost multiplier. MTP spec decode keeps a small draft
+    # context beside the main one. Calibrated against four measured
+    # server-RSS points on Qwen3.8 Q4 (128K/221K/256K, both postures):
+    # the draft adds ~17% to per-token KV; 1.2 rounds up so the error
+    # stays on the safe side (+250 MiB at 256K, never negative).
+    kv_scale: float = 1.0
+
+    @property
+    def per_token_kv_f16(self) -> int:
+        """Uncapped per-token KV cost (full + SWA share)."""
+        return sum(b for kind, b in self.layers if kind != LayerKind.RECURRENT)
+
+    @property
+    def recurrent_layer_count(self) -> int:
+        return sum(1 for kind, _ in self.layers if kind == LayerKind.RECURRENT)
+
+
+@dataclass
+class HardwareBudget:
+    """Memory the physics check may budget against.
+
+    Budget-source rule: discrete cards may trust the device query
+    (measured honest); unified-memory devices must budget from OS free
+    physical memory minus headroom — their device queries have been
+    observed off by 3x. Callers construct
+    this accordingly; the estimator just consumes it.
+    """
+
+    usable_vram_bytes: int      # live free (discrete) / derived (UMA)
+    total_device_bytes: int
+    ram_available_bytes: int
+    uma: bool = False
+
+
+def profile_from_gguf(header: GGUFHeader) -> ModelProfile:
+    kv_heads = header.head_counts_kv()
+    dk, dv = header.head_dim_k, header.head_dim_v
+    swa_fraction = _SWA_LAYER_FRACTION.get(header.architecture, 0.0)
+    has_swa = header.sliding_window > 0 and swa_fraction > 0
+
+    layers: list[tuple[LayerKind, int]] = []
+    n_attn_seen = 0
+    n_attn_total = sum(1 for h in kv_heads if h > 0)
+    n_swa = round(n_attn_total * swa_fraction) if has_swa else 0
+    for heads in kv_heads:
+        if heads == 0:
+            layers.append((LayerKind.RECURRENT, 0))
+            continue
+        per_token = round(heads * (dk + dv) * _F16_BYTES_PER_ELEM)
+        # Distribute the SWA share across the first n_swa attention layers;
+        # only the full/SWA SPLIT matters to the totals, not which indexes.
+        kind = LayerKind.SWA if n_attn_seen < n_swa else LayerKind.FULL
+        layers.append((kind, per_token))
+        n_attn_seen += 1
+
+    return ModelProfile(
+        name=header.path,
+        weights_bytes=header.tensor_bytes,
+        embd_table_bytes=header.embd_table_bytes,
+        n_ctx_train=header.n_ctx_train,
+        layers=layers,
+        swa_window=header.sliding_window,
+        moe=header.expert_count > 0,
+        architecture=header.architecture,
+        n_vocab=header.n_vocab,
+    )
+
+
+def kv_dtype_factor(flash_attention: bool) -> float:
+    """q8_0 with FA (every backend we ship); f16 on exotic non-FA fallbacks
+    — the 64K guarantee stands either way, the physics check just prices
+    the doubled KV (design: KV dtype is behavior, not config)."""
+    return (_Q8_BYTES_PER_ELEM / _F16_BYTES_PER_ELEM) if flash_attention else 1.0
+
+
+def ctx_bytes(profile: ModelProfile, window: int, *,
+              flash_attention: bool = True) -> int:
+    """Context memory for one window: full layers linear in T, SWA layers
+    capped at the sliding window, recurrent layers constant. Scaled by
+    profile.kv_scale (MTP draft context)."""
+    factor = kv_dtype_factor(flash_attention)
+    total = 0.0
+    for kind, per_token_f16 in profile.layers:
+        if kind == LayerKind.RECURRENT:
+            total += _RECURRENT_STATE_PER_LAYER
+        elif kind == LayerKind.SWA:
+            total += per_token_f16 * factor * min(window, profile.swa_window)
+        else:
+            total += per_token_f16 * factor * window
+    return int(total * profile.kv_scale)
+
+
+@dataclass
+class PhysicsRefusal:
+    """The only true refusal: weights + floor-KV + state exceed VRAM + RAM.
+    The remedy is a smaller quant, never a smaller window."""
+
+    needed_bytes: int
+    available_bytes: int
+    message: str
+
+
+def physics_check(profile: ModelProfile, budget: HardwareBudget,
+                  floor: int, *, flash_attention: bool = True) -> PhysicsRefusal | None:
+    needed = (profile.weights_bytes
+              + ctx_bytes(profile, min(floor, profile.n_ctx_train or floor),
+                          flash_attention=flash_attention))
+    available = budget.usable_vram_bytes + budget.ram_available_bytes
+    if needed > available:
+        gib = 1 << 30
+        return PhysicsRefusal(
+            needed_bytes=needed, available_bytes=available,
+            message=(f"{profile.name}: needs ~{needed / gib:.1f} GiB at the "
+                     f"{floor // 1024}K floor but only ~{available / gib:.1f} GiB "
+                     "of VRAM+RAM exist — try a smaller quant (UD-Q3/Q2)"))
+    return None