eigencore 0.4.0__py3-none-any.whl

eigencore/engine/sparsity_experiment.py

@@ -0,0 +1,109 @@
+"""
+Activation-sparsity experiment — Stage E slice 3 (PowerInfer/TurboSparse/TEAL-
+style), REPORTED-ONLY.
+
+Idea: if a fraction of FFN neurons are near-zero for a given token, skip their
+contribution (sparse GEMM) to save compute. On bandwidth-bound CPU decode the
+win is small and quant/dequant dominates — hence the kill criterion.
+
+This module measures per-layer activation sparsity on real activations,
+builds a sparse mask plan, and enforces the kill criterion: if the measured
+quality delta exceeds the allowed threshold, the experiment is aborted and no
+sparsity is applied. Nothing here is wired into inference; the backend always
+runs dense today. No quality/performance claim without the owner's runs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+
+def activation_sparsity(activations: np.ndarray, threshold: float = 0.01) -> float:
+    """Fraction of |x| <= threshold. `activations`: (tokens, hidden)."""
+    a = np.abs(np.asarray(activations, dtype=np.float64))
+    if a.size == 0:
+        return 0.0
+    return float((a <= threshold).mean())
+
+
+def layer_sparsity_profile(
+    hidden_states: np.ndarray,  # (layers, tokens, hidden) FFN inputs per layer
+    threshold: float = 0.01,
+) -> np.ndarray:
+    """Per-layer sparsity, (layers,)."""
+    hs = np.asarray(hidden_states, dtype=np.float64)
+    if hs.ndim != 3:
+        raise ValueError("expected (layers, tokens, hidden)")
+    return np.array([activation_sparsity(hs[i], threshold) for i in range(hs.shape[0])])
+
+
+@dataclass
+class SparseMask:
+    layer: int
+    keep: np.ndarray  # bool mask over hidden dim; True = compute
+
+    @property
+    def sparsity(self) -> float:
+        return 1.0 - float(self.keep.mean())
+
+
+def build_sparse_mask(
+    activations: np.ndarray,
+    threshold: float = 0.01,
+) -> SparseMask:
+    """Keep neurons whose |activation| exceeds threshold, averaged over tokens."""
+    a = np.abs(np.asarray(activations, dtype=np.float64))
+    mean_abs = a.mean(axis=0) if a.ndim == 2 else np.abs(a)
+    return SparseMask(layer=-1, keep=mean_abs > threshold)
+
+
+@dataclass
+class SparsityPlan:
+    masks: list[SparseMask]
+    overall_sparsity: float
+    max_quality_delta_allowed: float
+    killed: bool = False
+    kill_reason: str = ""
+    notes: list[str] = field(default_factory=list)
+
+
+def build_sparsity_plan(
+    hidden_states: np.ndarray,  # (layers, tokens, hidden)
+    threshold: float = 0.01,
+    min_sparsity_to_apply: float = 0.5,
+    max_quality_delta_allowed: float = 0.02,
+    measured_quality_delta: float | None = None,
+) -> SparsityPlan:
+    """Build per-layer sparse masks; kill if quality delta exceeds allowed.
+
+    Layers whose measured sparsity is below `min_sparsity_to_apply` get a full
+    keep mask (applying sparsity there saves nothing meaningful).
+    """
+    profile = layer_sparsity_profile(hidden_states, threshold)
+    if measured_quality_delta is not None and measured_quality_delta > max_quality_delta_allowed:
+        return SparsityPlan(
+            masks=[],
+            overall_sparsity=float(profile.mean()),
+            max_quality_delta_allowed=max_quality_delta_allowed,
+            killed=True,
+            kill_reason=(
+                f"quality delta {measured_quality_delta:.4f} > allowed "
+                f"{max_quality_delta_allowed:.4f} — kill criterion hit, no sparsity applied"
+            ),
+        )
+    masks: list[SparseMask] = []
+    for i in range(hidden_states.shape[0]):
+        if profile[i] >= min_sparsity_to_apply:
+            masks.append(build_sparse_mask(hidden_states[i], threshold))
+            masks[-1].layer = i
+        else:
+            masks.append(SparseMask(layer=i, keep=np.ones(hidden_states.shape[2], dtype=bool)))
+    overall = float(sum(m.sparsity for m in masks) / len(masks)) if masks else 0.0
+    return SparsityPlan(
+        masks=masks,
+        overall_sparsity=overall,
+        max_quality_delta_allowed=max_quality_delta_allowed,
+        notes=[f"mean measured sparsity {profile.mean():.1%} at threshold {threshold}"],
+    )

eigencore/engine/speculative.py

@@ -0,0 +1,238 @@
+"""
+Speculative decoding rethought for CPU.
+
+Standard speculative decoding uses a small "draft" model to predict N tokens,
+then the large "verifier" model checks them in a single batch. On GPU this is
+fast because the batch verify step parallelizes well.
+
+On CPU the advantage is different: we exploit the fact that CPU branch prediction
+and sequential processing makes it cheaper to run a tiny draft model for N tokens
+than to run the large model N times. The verify step is still sequential on CPU,
+but we only pay the large-model cost once per accepted draft sequence.
+
+Key CPU-specific adaptations:
+1. Draft model shares the same memory-mapped GGUF file as verifier (different quant)
+2. Acceptance threshold adapts based on observed accept rate
+3. Draft length scales with CPU cache size (longer drafts if L3 cache can hold draft model)
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+
+import numpy as np
+
+
+@dataclass
+class DraftToken:
+    token_id: int
+    logprob: float
+    position: int
+
+
+@dataclass
+class SpeculativeResult:
+    accepted_tokens: list[int]
+    drafted_tokens: int
+    accepted_count: int
+    rejection_position: int | None
+    correction_token: int | None
+    draft_time_ms: float
+    verify_time_ms: float
+    total_time_ms: float
+
+    @property
+    def accept_rate(self) -> float:
+        return self.accepted_count / self.drafted_tokens if self.drafted_tokens > 0 else 0.0
+
+    @property
+    def speedup(self) -> float:
+        if self.total_time_ms <= 0:
+            return 1.0
+        naive_time = self.total_time_ms * (self.drafted_tokens + 1) / max(self.drafted_tokens, 1)
+        return naive_time / self.total_time_ms if self.total_time_ms > 0 else 1.0
+
+
+@dataclass
+class SpeculativeStats:
+    total_drafted: int = 0
+    total_accepted: int = 0
+    total_rounds: int = 0
+    draft_times_ms: list[float] = field(default_factory=list)
+    verify_times_ms: list[float] = field(default_factory=list)
+
+    @property
+    def overall_accept_rate(self) -> float:
+        return self.total_accepted / self.total_drafted if self.total_drafted > 0 else 0.0
+
+    @property
+    def avg_draft_time_ms(self) -> float:
+        return np.mean(self.draft_times_ms) if self.draft_times_ms else 0.0
+
+    @property
+    def avg_verify_time_ms(self) -> float:
+        return np.mean(self.verify_times_ms) if self.verify_times_ms else 0.0
+
+    def summary(self) -> str:
+        return (
+            f"Speculative decoding: {self.total_rounds} rounds | "
+            f"Accept rate: {self.overall_accept_rate:.0%} | "
+            f"Avg draft: {self.avg_draft_time_ms:.1f}ms | "
+            f"Avg verify: {self.avg_verify_time_ms:.1f}ms"
+        )
+
+
+class SpeculativeDecoder:
+    """
+    CPU-optimized speculative decoding engine.
+
+    Takes a draft function (fast, small model) and a verify function (large model)
+    as callables, so it can wrap any underlying inference backend.
+    """
+
+    DEFAULT_DRAFT_LENGTH = 5
+    MIN_DRAFT_LENGTH = 2
+    MAX_DRAFT_LENGTH = 12
+
+    def __init__(
+        self,
+        draft_fn=None,
+        verify_fn=None,
+        initial_draft_length: int = 5,
+        acceptance_threshold: float = 0.3,
+        adaptive: bool = True,
+    ):
+        self.draft_fn = draft_fn
+        self.verify_fn = verify_fn
+        self.draft_length = initial_draft_length
+        self.acceptance_threshold = acceptance_threshold
+        self.adaptive = adaptive
+        self.stats = SpeculativeStats()
+
+    def optimal_draft_length(self, l3_cache_mb: float, draft_model_mb: float) -> int:
+        """
+        Scale draft length based on whether the draft model fits in L3 cache.
+        If it fits, longer drafts are nearly free (no main memory round-trips).
+        """
+        if draft_model_mb <= 0:
+            return self.DEFAULT_DRAFT_LENGTH
+
+        cache_ratio = l3_cache_mb / draft_model_mb
+
+        if cache_ratio >= 1.0:
+            return min(self.MAX_DRAFT_LENGTH, int(self.DEFAULT_DRAFT_LENGTH * 1.5))
+        elif cache_ratio >= 0.5:
+            return self.DEFAULT_DRAFT_LENGTH
+        else:
+            return max(self.MIN_DRAFT_LENGTH, int(self.DEFAULT_DRAFT_LENGTH * 0.6))
+
+    def verify_draft(
+        self,
+        draft_logprobs: np.ndarray,
+        verifier_logprobs: np.ndarray,
+        draft_token_ids: list[int],
+        rng: np.random.Generator | None = None,
+    ) -> SpeculativeResult:
+        """
+        Leviathan/Chen-style speculative rejection sampling.
+
+        `verifier_logprobs` is (n_positions, vocab) log-probs from the target
+        model; `draft_logprobs[i]` is the draft model's log-prob of the drafted
+        token at position i. Token t is accepted with probability
+        min(1, q(t)/p(t)); on rejection the correction token is sampled from
+        the residual distribution norm(max(q - p, 0)), which preserves the
+        target distribution exactly. (The draft distribution is approximated
+        as a point mass at the drafted token, since only its log-prob is
+        passed in — documented limitation.)
+        """
+        start = time.perf_counter()
+        n_draft = len(draft_token_ids)
+        rng = rng or np.random.default_rng()
+
+        if n_draft == 0 or len(draft_logprobs) == 0 or len(verifier_logprobs) == 0:
+            return SpeculativeResult(
+                accepted_tokens=[],
+                drafted_tokens=0,
+                accepted_count=0,
+                rejection_position=None,
+                correction_token=None,
+                draft_time_ms=0,
+                verify_time_ms=0,
+                total_time_ms=0,
+            )
+
+        accepted = []
+        rejection_pos = None
+        correction = None
+
+        for i in range(min(n_draft, len(verifier_logprobs))):
+            token_id = draft_token_ids[i]
+            if i < len(draft_logprobs):
+                draft_lp = draft_logprobs[i]
+            else:
+                break
+            q = np.exp(verifier_logprobs[i] - verifier_logprobs[i].max())
+            q /= q.sum()
+            p_t = float(np.exp(draft_lp))
+            q_t = float(q[token_id]) if token_id < len(q) else 0.0
+            if rng.random() <= min(1.0, q_t / p_t) if p_t > 0 else True:
+                accepted.append(token_id)
+            else:
+                rejection_pos = i
+                # residual distribution: norm(max(q - p, 0)), p = point mass at token
+                residual = np.maximum(q - p_t * np.eye(len(q))[token_id], 0.0)
+                total = residual.sum()
+                if total > 0:
+                    correction = int(rng.choice(len(q), p=residual / total))
+                else:
+                    correction = int(np.argmax(q))
+                break
+
+        elapsed = (time.perf_counter() - start) * 1000
+
+        result = SpeculativeResult(
+            accepted_tokens=accepted,
+            drafted_tokens=n_draft,
+            accepted_count=len(accepted),
+            rejection_position=rejection_pos,
+            correction_token=correction,
+            draft_time_ms=0,
+            verify_time_ms=elapsed,
+            total_time_ms=elapsed,
+        )
+
+        self._update_stats(result)
+
+        if self.adaptive:
+            self._adapt_draft_length(result.accept_rate)
+
+        return result
+
+    def _adapt_draft_length(self, recent_accept_rate: float) -> None:
+        """
+        Adjust draft length based on acceptance rate.
+        High accept rate → try longer drafts (more tokens per round).
+        Low accept rate → shorter drafts (less wasted compute).
+        """
+        if recent_accept_rate > 0.8:
+            self.draft_length = min(self.draft_length + 1, self.MAX_DRAFT_LENGTH)
+        elif recent_accept_rate < 0.3:
+            self.draft_length = max(self.draft_length - 1, self.MIN_DRAFT_LENGTH)
+
+    def _update_stats(self, result: SpeculativeResult) -> None:
+        self.stats.total_drafted += result.drafted_tokens
+        self.stats.total_accepted += result.accepted_count
+        self.stats.total_rounds += 1
+        self.stats.draft_times_ms.append(result.draft_time_ms)
+        self.stats.verify_times_ms.append(result.verify_time_ms)
+
+    def measured_speedup(self, draft_time_ms: float, verify_time_ms: float) -> float:
+        """Speedup computed only from measured times; no theoretical claims."""
+        if verify_time_ms <= 0:
+            return 1.0
+        return (
+            (draft_time_ms + verify_time_ms) / max(verify_time_ms, 1e-9)
+            if draft_time_ms > 0
+            else 1.0
+        )

eigencore/fit/__init__.py

@@ -0,0 +1,34 @@
+"""Phase 2 model-fit estimation."""
+
+from eigencore.fit.estimator import (
+    BYTES_PER_PARAM,
+    DEFAULT_RUNTIME_RESERVE_GB,
+    DEFAULT_SAFETY_MARGIN_GB,
+    FitEstimate,
+    FitVerdict,
+    ModelGeometry,
+    decode_ceiling_tok_s,
+    estimate_fit,
+    estimate_kv_cache_gib,
+    estimate_weights_gib,
+    gpu_layers_for_vram,
+    max_context_for_budget,
+)
+from eigencore.fit.gguf_meta import GGUFModelInfo, read_gguf_model_info
+
+__all__ = [
+    "BYTES_PER_PARAM",
+    "DEFAULT_RUNTIME_RESERVE_GB",
+    "DEFAULT_SAFETY_MARGIN_GB",
+    "GGUFModelInfo",
+    "FitEstimate",
+    "FitVerdict",
+    "ModelGeometry",
+    "estimate_fit",
+    "estimate_kv_cache_gib",
+    "estimate_weights_gib",
+    "gpu_layers_for_vram",
+    "decode_ceiling_tok_s",
+    "max_context_for_budget",
+    "read_gguf_model_info",
+]

eigencore/fit/estimator.py

@@ -0,0 +1,310 @@
+"""
+Model-fit estimation — Phase 2 (honest hardware fit).
+
+Estimates the memory needed to *load and run* a model, separating:
+  - weights        (actual GGUF file size when known, else params × bytes-per-param)
+  - runtime reserve (allocator/backend/OS overhead, configurable)
+  - KV cache       (context-length dependent; needs per-model geometry)
+
+All sizes are **GiB** (1 GiB = 1024**3 bytes) end to end. Without model
+geometry the KV cache is unknowable, so the verdict is `UNKNOWN` rather than
+a made-up heuristic — a Core i7 label alone cannot predict whether a 35B Q4
+model is usable.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+# Approximate bytes per parameter for common GGUF quant types.
+# These are planning estimates, not measurements.
+BYTES_PER_PARAM: dict[str, float] = {
+    "Q4_K_M": 0.55,
+    "Q4_0": 0.50,
+    "Q5_K_M": 0.68,
+    "Q6_K": 0.84,
+    "Q8_0": 1.05,
+    "F16": 2.0,
+    "F32": 4.0,
+}
+
+# KV cache per token, per layer, per KV head, in bytes at fp16.
+# kv_bytes_per_token = layers × kv_heads × head_dim × 2 (K+V) × 2 bytes
+# head_dim = hidden_size / n_heads; for grouped-query attention kv_heads < n_heads.
+FP16_BYTES_PER_ELEMENT = 2
+KV_BYTES_PER_ELEMENT_Q8 = 1  # q8_0 KV cache
+
+DEFAULT_RUNTIME_RESERVE_GB = 2.0
+DEFAULT_SAFETY_MARGIN_GB = 1.0
+
+
+class FitVerdict(str, Enum):
+    FITS_WITH_MARGIN = "fits_with_margin"
+    TIGHT = "tight"
+    DOES_NOT_FIT = "does_not_fit"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class ModelGeometry:
+    """Architecture metadata needed for a KV-cache estimate.
+
+    All fields optional: when absent the estimator falls back to a
+    conservative generic estimate and marks the result as UNKNOWN-precision.
+    """
+
+    num_layers: int | None = None
+    hidden_size: int | None = None
+    num_attention_heads: int | None = None
+    num_kv_heads: int | None = None  # GQA/MQA: defaults to num_attention_heads
+    expert_count: int | None = None  # MoE: total experts per MoE layer
+    expert_used_count: int | None = None  # MoE: experts activated per token
+
+    @property
+    def kv_heads(self) -> int | None:
+        return self.num_kv_heads or self.num_attention_heads
+
+    @property
+    def head_dim(self) -> int | None:
+        if self.hidden_size and self.num_attention_heads:
+            return self.hidden_size // self.num_attention_heads
+        return None
+
+    @property
+    def is_complete(self) -> bool:
+        return all(
+            v is not None for v in (self.num_layers, self.hidden_size, self.num_attention_heads)
+        )
+
+    @property
+    def is_moe(self) -> bool:
+        return bool(self.expert_count and self.expert_used_count)
+
+
+@dataclass
+class FitEstimate:
+    weights_gb: float
+    active_weights_gb: float | None  # MoE only: weights actually read per decode token
+    runtime_reserve_gb: float
+    kv_cache_gb: float
+    total_gb: float
+    available_ram_gb: float
+    headroom_gb: float
+    verdict: FitVerdict
+    context_length: int
+    quant: str
+    resident: bool  # file fully in RAM (False → mmap/SSD paging)
+    decode_ceiling_tok_s: float | None  # bandwidth / active bytes-per-token ceiling
+    gpu_layers: int | None = None  # layers offloaded to GPU, if a VRAM budget was given
+    gpu_weights_gb: float = 0.0  # share of weights resident in VRAM
+    notes: list[str] = field(default_factory=list)
+
+    @property
+    def fits(self) -> bool:
+        return self.verdict in (FitVerdict.FITS_WITH_MARGIN, FitVerdict.TIGHT)
+
+
+def gpu_layers_for_vram(
+    num_layers: int,
+    layer_bytes: int,
+    available_vram_gb: float,
+    kv_cache_gb: float = 0.0,
+    vram_reserve_gb: float = 1.0,
+) -> int:
+    """How many layers fit in VRAM: (vram - reserve - kv) / layer_bytes, capped at num_layers.
+
+    layer_bytes is the per-layer weight size (uniform approximation; real
+    per-layer sizes vary slightly). Returns 0 when nothing fits.
+    """
+    if num_layers <= 0 or layer_bytes <= 0 or available_vram_gb <= 0:
+        return 0
+    budget_bytes = max((available_vram_gb - vram_reserve_gb - kv_cache_gb) * 1024**3, 0.0)
+    return min(num_layers, int(budget_bytes // layer_bytes))
+
+
+def estimate_weights_gib(
+    params_b: float,
+    quant: str,
+    file_size_bytes: int | None = None,
+) -> tuple[float, str | None]:
+    """Weight size in GiB.
+
+    Uses the actual GGUF file size when available (ground truth, includes
+    tokenizer/mmap overhead); otherwise params × bytes-per-param.
+    """
+    if file_size_bytes is not None:
+        return file_size_bytes / 1024**3, None
+    bpp = BYTES_PER_PARAM.get(quant.upper())
+    if bpp is None:
+        return params_b * 0.61, f"unknown quant {quant!r}; assumed ~0.61 bytes/param (Q4_K_M class)"
+    return params_b * bpp, None
+
+
+def estimate_kv_cache_gib(
+    context_length: int,
+    geometry: ModelGeometry | None,
+    params_b: float,
+    kv_quant: str = "fp16",
+) -> tuple[float | None, str | None]:
+    """KV cache size in GiB for a context length.
+
+    layers × kv_heads × head_dim × 2 (K+V) × bytes-per-element × ctx.
+    Returns (None, note) when geometry is incomplete — no heuristic.
+    """
+    if geometry is None or not geometry.is_complete:
+        return None, "no GGUF geometry available; KV cache unknown"
+    bytes_per_elem = (
+        FP16_BYTES_PER_ELEMENT if kv_quant.lower() in ("fp16", "f16") else KV_BYTES_PER_ELEMENT_Q8
+    )
+    per_token = geometry.num_layers * geometry.kv_heads * geometry.head_dim * 2 * bytes_per_elem
+    return per_token * context_length / 1024**3, None
+
+
+def decode_ceiling_tok_s(bandwidth_gbps: float, active_weights_gib: float) -> float:
+    """Optimistic decode ceiling: bandwidth / bytes streamed per token.
+
+    For dense models active == resident weights. For MoE, only the activated
+    experts' weights are read per token (plus shared/non-expert weights, which
+    we approximate as included in the active figure). This is a CEILING —
+    real decode is lower due to cache misses, dequant, sampling overhead.
+    """
+    if active_weights_gib <= 0:
+        return 0.0
+    return bandwidth_gbps / active_weights_gib
+
+
+def estimate_fit(
+    params_b: float,
+    quant: str,
+    context_length: int,
+    available_ram_gb: float,
+    geometry: ModelGeometry | None = None,
+    runtime_reserve_gb: float = DEFAULT_RUNTIME_RESERVE_GB,
+    safety_margin_gb: float = DEFAULT_SAFETY_MARGIN_GB,
+    kv_quant: str = "fp16",
+    file_size_bytes: int | None = None,
+    bandwidth_gbps: float | None = None,
+    active_params_b: float | None = None,
+    available_vram_gb: float = 0.0,
+) -> FitEstimate:
+    """Full memory fit estimate (GiB) for loading + running a model at a context length.
+
+    MoE: pass `active_params_b` (params activated per token) — resident memory
+    still uses the full file size, but the decode ceiling uses active weights.
+    """
+    notes: list[str] = []
+    weights_gb, note = estimate_weights_gib(params_b, quant, file_size_bytes)
+    if note:
+        notes.append(note)
+    if geometry is not None and geometry.is_moe and active_params_b is None:
+        # derive active params from expert ratio when geometry is present
+        assert geometry.expert_count and geometry.expert_used_count
+        active_params_b = params_b * geometry.expert_used_count / geometry.expert_count
+        notes.append(
+            f"MoE: {geometry.expert_count} experts, {geometry.expert_used_count} active/token "
+            f"→ active ≈ {active_params_b:.1f}B params"
+        )
+    active_weights_gb: float | None = None
+    if active_params_b is not None and active_params_b < params_b:
+        active_weights_gb, _ = estimate_weights_gib(active_params_b, quant)
+    kv_gb_raw, note = estimate_kv_cache_gib(context_length, geometry, params_b, kv_quant)
+    if note:
+        notes.append(note)
+    kv_gb = kv_gb_raw or 0.0
+
+    total = weights_gb + runtime_reserve_gb + kv_gb
+    headroom = available_ram_gb - total
+    resident = weights_gb <= available_ram_gb
+
+    if kv_gb_raw is None:
+        verdict = FitVerdict.UNKNOWN
+        notes.append("verdict UNKNOWN: KV cache excluded (no geometry); real usage is higher")
+    elif headroom >= safety_margin_gb:
+        verdict = FitVerdict.FITS_WITH_MARGIN
+    elif headroom >= 0:
+        verdict = FitVerdict.TIGHT
+        notes.append(f"only {headroom:.1f} GB headroom; expect swapping/slowdown under load")
+    else:
+        verdict = FitVerdict.DOES_NOT_FIT
+        notes.append(f"short by {-headroom:.1f} GB at context {context_length}")
+
+    if not resident:
+        notes.append(
+            f"model file ({weights_gb:.1f} GiB) exceeds available RAM ({available_ram_gb:.1f} GiB): "
+            "would run via mmap/SSD paging — expect severe decode slowdown (SSD read rate, not RAM bandwidth)"
+        )
+
+    ceiling = None
+    if bandwidth_gbps is not None:
+        active_gib = active_weights_gb if active_weights_gb is not None else weights_gb
+        ceiling = decode_ceiling_tok_s(bandwidth_gbps, active_gib)
+
+    gpu_layers = None
+    gpu_weights_gb = 0.0
+    if available_vram_gb > 0 and geometry is not None and geometry.num_layers:
+        layer_bytes = int(weights_gb * 1024**3 / geometry.num_layers) if weights_gb > 0 else 0
+        gpu_layers = gpu_layers_for_vram(geometry.num_layers, layer_bytes, available_vram_gb, kv_gb)
+        gpu_weights_gb = gpu_layers * layer_bytes / 1024**3
+        if gpu_layers > 0:
+            notes.append(
+                f"GPU offload: ~{gpu_layers}/{geometry.num_layers} layers "
+                f"(≈{gpu_weights_gb:.1f} GiB) fit in {available_vram_gb:.1f} GiB VRAM; "
+                "rest stays in RAM"
+            )
+        else:
+            notes.append("VRAM present but too small for even one layer after reserve/KV")
+
+    return FitEstimate(
+        weights_gb=weights_gb,
+        active_weights_gb=active_weights_gb,
+        runtime_reserve_gb=runtime_reserve_gb,
+        kv_cache_gb=kv_gb,
+        total_gb=total,
+        available_ram_gb=available_ram_gb,
+        headroom_gb=headroom,
+        verdict=verdict,
+        context_length=context_length,
+        quant=quant,
+        resident=resident,
+        decode_ceiling_tok_s=ceiling,
+        gpu_layers=gpu_layers,
+        gpu_weights_gb=gpu_weights_gb,
+        notes=notes,
+    )
+
+
+def max_context_for_budget(
+    params_b: float,
+    quant: str,
+    available_ram_gb: float,
+    geometry: ModelGeometry | None = None,
+    runtime_reserve_gb: float = DEFAULT_RUNTIME_RESERVE_GB,
+    safety_margin_gb: float = DEFAULT_SAFETY_MARGIN_GB,
+    model_context_limit: int = 128 * 1024,
+    kv_quant: str = "fp16",
+    file_size_bytes: int | None = None,
+) -> int:
+    """Largest context length that fits the RAM budget (power-of-two capped).
+
+    Returns 0 when geometry is unknown — without per-token KV bytes a context
+    recommendation would be a guess.
+    """
+    if geometry is None or not geometry.is_complete:
+        return 0
+    weights_gb, _ = estimate_weights_gib(params_b, quant, file_size_bytes)
+    budget_gb = available_ram_gb - runtime_reserve_gb - safety_margin_gb - weights_gb
+    if budget_gb <= 0:
+        return 0
+    bytes_per_elem = (
+        FP16_BYTES_PER_ELEMENT if kv_quant.lower() in ("fp16", "f16") else KV_BYTES_PER_ELEMENT_Q8
+    )
+    per_token = geometry.num_layers * geometry.kv_heads * geometry.head_dim * 2 * bytes_per_elem
+    if per_token <= 0:
+        return 0
+    max_ctx = int(budget_gb * 1024**3 / per_token)
+    # snap down to a power of two, clamp to model limit and a sane floor
+    ctx = 1
+    while ctx * 2 <= max_ctx and ctx * 2 <= model_context_limit:
+        ctx *= 2
+    return max(ctx, 0)