mission-control-v2/backend/services/fit.py

"""
Hardware-Fit-Mathe (VRAM/RAM, tps-Schätzung) für APUs mit Unified Memory
(Bosgame M5 / Strix Halo). Portiert aus Mission Control v1 (hw_math.py).
"""

import re

# Bytes pro Parameter je GGUF-Quant (Annahme).
QUANT_BYTES_PER_PARAM = {
    "Q2_K": 0.35, "Q3_K_S": 0.38, "Q3_K_M": 0.42, "Q3_K_L": 0.45,
    "Q4_0": 0.50, "Q4_1": 0.55, "Q4_K_S": 0.50, "Q4_K_M": 0.55,
    "Q5_0": 0.62, "Q5_1": 0.68, "Q5_K_S": 0.62, "Q5_K_M": 0.65,
    "Q6_K": 0.75, "Q8_0": 1.00, "F16": 2.00, "BF16": 2.00,
    "MXFP4": 0.55, "FP8": 1.05, "AWQ": 0.55,
}


def estimate_memory_gb(params_b: float, quant: str, ctx: int) -> float:
    """Geschätzter Speicherbedarf in GB (Gewichte + Kontext-KV).
    KV-Cache skaliert NICHT linear mit den Gesamt-Parametern (er hängt an
    Layern × KV-Heads, gedämpft durch GQA) → sqrt-Skalierung, kalibriert am
    gemessenen Punkt Hermes-4-14B @ 128K ≈ 19 GB KV."""
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.65)
    weights = params_b * bpp
    context_vram = (ctx / 8192) * (max(params_b, 7) / 7) ** 0.5 * 0.84
    return weights + context_vram


def extract_active_params_b(name: str) -> float | None:
    """Aktive Parameter bei MoE ('30B-A3B' → 3.0). None bei Dense."""
    m = re.search(r"(?<![a-zA-Z])a(\d+(?:\.\d+)?)b\b", name.lower())
    return float(m.group(1)) if m else None


def estimate_speed(req_gb: float, sys_ram_gb: float, moe_active_ratio: float = 1.0) -> float:
    """Geschätzte t/s anhand der ~273 GB/s Bandbreite der APU.
    moe_active_ratio = aktive/gesamt Params; < 1 bei MoE."""
    bw = 273 if sys_ram_gb > 8 else 70
    if req_gb <= 0:
        return 0.0
    raw_tps = (bw / req_gb) * 0.55
    if moe_active_ratio < 0.8:
        raw_tps *= (1.0 / moe_active_ratio) ** 0.5
    return raw_tps


def evaluate_fit(params_b: float, quant: str, ctx: int, sys_ram_gb: float, name: str = "") -> dict:
    """Fit für ein Shared-Memory-System (APU). name → MoE-Erkennung (optional)."""
    req_gb = estimate_memory_gb(params_b, quant, ctx)
    active_b = extract_active_params_b(name) if name else None
    moe_ratio = (active_b / params_b) if (active_b and params_b > 0) else 1.0
    tps = estimate_speed(req_gb, sys_ram_gb, moe_ratio)
    usable_ram = max(sys_ram_gb - 4.0, 0)
    if req_gb > usable_ram:
        fit_level, text = "too_tight", "Zu groß (OOM)"
    elif req_gb > usable_ram * 0.8:
        fit_level, text = "marginal", "Könnte knapp werden"
    else:
        fit_level, text = "perfect", "Passt perfekt"
    return {"level": fit_level, "text": text, "req_gb": round(req_gb, 1), "tps": round(tps, 0)}


def extract_params_b(name: str) -> float:
    """Parametergröße (Mrd.) aus Repo-/Dateiname. 8x7B (MoE) → 56."""
    moe = re.search(r"(\d+)x(\d+(?:\.\d+)?)[bB]", name)
    if moe:
        return float(moe.group(1)) * float(moe.group(2))
    m = re.search(r"(\d+(?:\.\d+)?)[bB](?![a-zA-Z])", name)
    return float(m.group(1)) if m else 7.0


_NICE_CTX = [2048, 4096, 8192, 16384, 32768, 49152, 65536, 98304, 131072]


def max_ctx_in_budget(params_b: float, quant: str, budget_gb: float) -> int:
    """Größter 'schöner' Kontext, dessen Gewichte + KV in budget_gb passen.
    Budget-basierter Kern → wird von der setup-bewussten ctx-Vergabe
    (services.budget) mit dem ECHTEN freien Budget gefüttert."""
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.65)
    weights = params_b * bpp
    ctx_budget = budget_gb - weights
    if ctx_budget <= 0:
        return 2048
    # KV pro 8k — EXAKTE Inverse von estimate_memory_gb (sqrt, kalibriert an
    # Hermes-14B@128K≈19GB). Vorher linear → für große Modelle viel zu konservativ.
    per_8k = (max(params_b, 7) / 7) ** 0.5 * 0.84
    raw_ctx = (ctx_budget / per_8k) * 8192
    best = _NICE_CTX[0]
    for c in _NICE_CTX:
        if c <= raw_ctx:
            best = c
    return best


def max_ctx_for(params_b: float, quant: str, sys_ram_gb: float) -> int:
    """Roh-Obergrenze: größter Kontext für dieses Modell ALLEIN gegen den
    Gesamt-RAM (80 % nutzbar). Ignoriert bewusst das übrige Setup —
    setup-bewusst rechnet services.budget.setup_aware_ctx."""
    return max_ctx_in_budget(params_b, quant, max(sys_ram_gb - 4.0, 0) * 0.8)