tinymind-native-8b-remote-handoff/bundle/evaluation/int6_precision_tradeoff.py

"""Evidence-gated INT6 precision-vs-throughput tradeoff report."""

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_PRECISION_REPORT = ROOT / "reports" / "int6_precision_ladder" / "int6_precision_ladder_report.json"
DEFAULT_TFW_REPORT = ROOT / "reports" / "tfw_optimizer_int4_vs_int6_realdata" / "tfw_optimizer_report.json"


def _load(path: str | Path) -> dict:
    p = Path(path)
    return json.loads(p.read_text(encoding="utf-8-sig")) if p.exists() else {}


def _candidate(report: dict, name: str) -> dict:
    for row in report.get("candidates", []):
        if row.get("name") == name:
            return row
    return {}


def _safe_float(value: object, default: float = 0.0) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def _ratio(numerator: float, denominator: float) -> float:
    return numerator / denominator if denominator > 0 else 0.0


def build_int6_precision_tradeoff(
    out_dir: str | Path,
    precision_report: str | Path = DEFAULT_PRECISION_REPORT,
    tfw_report: str | Path = DEFAULT_TFW_REPORT,
    min_mae_reduction_pct: float = 25.0,
    min_tfw_ratio_vs_int4: float = 0.45,
) -> dict:
    precision = _load(precision_report)
    tfw = _load(tfw_report)
    int4_precision = precision.get("formats", {}).get("int4", {})
    int6_precision = precision.get("formats", {}).get("int6", {})
    int4_error = int4_precision.get("error", {})
    int6_error = int6_precision.get("error", {})
    int4_tfw = _candidate(tfw, "int4_sparse_tensor_core_imma_sp")
    int6_bridge = _candidate(tfw, "int6_bridge_imma_fast")

    int4_mae = _safe_float(int4_error.get("mean_abs_error"))
    int6_mae = _safe_float(int6_error.get("mean_abs_error"))
    int4_rmse = _safe_float(int4_error.get("root_mean_square_error"))
    int6_rmse = _safe_float(int6_error.get("root_mean_square_error"))
    int4_max_abs = _safe_float(int4_error.get("max_abs_error"))
    int6_max_abs = _safe_float(int6_error.get("max_abs_error"))
    int4_bytes = _safe_float(int4_precision.get("artifact_bytes_reference"))
    int6_bytes = _safe_float(int6_precision.get("artifact_bytes_reference"))
    int4_tfw_value = _safe_float(int4_tfw.get("avg_effective_tops_per_watt"))
    int6_tfw_value = _safe_float(int6_bridge.get("avg_effective_tops_per_watt"))

    mae_reduction_pct = max(0.0, (1.0 - _ratio(int6_mae, int4_mae)) * 100.0) if int4_mae > 0 else 0.0
    rmse_reduction_pct = max(0.0, (1.0 - _ratio(int6_rmse, int4_rmse)) * 100.0) if int4_rmse > 0 else 0.0
    max_abs_reduction_pct = max(0.0, (1.0 - _ratio(int6_max_abs, int4_max_abs)) * 100.0) if int4_max_abs > 0 else 0.0
    tfw_ratio = _ratio(int6_tfw_value, int4_tfw_value)
    artifact_size_ratio = _ratio(int6_bytes, int4_bytes)

    precision_gate = mae_reduction_pct >= min_mae_reduction_pct and int6_mae < int4_mae
    runtime_gate = (
        int6_bridge.get("passed") is True
        and int6_bridge.get("metric_mode") == "real_data"
        and int6_bridge.get("real_data_movement_measured") is True
        and int6_bridge.get("power_measurement_representative") is True
    )
    tfw_gate = tfw_ratio >= min_tfw_ratio_vs_int4
    int6_tradeoff_wins = precision_gate and runtime_gate and tfw_gate

    # The score intentionally rewards measured precision and penalizes size and
    # TF/W loss. It is only a local selector, not a world-rank metric.
    precision_gain = max(_ratio(int4_mae, int6_mae), 1.0) if int6_mae > 0 else 0.0
    tradeoff_score = precision_gain * max(tfw_ratio, 0.0) / max(artifact_size_ratio, 1.0)

    report = {
        "schema_version": "tinymind-int6-precision-tradeoff-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "inputs": {
            "precision_report": str(precision_report),
            "tfw_report": str(tfw_report),
        },
        "thresholds": {
            "min_mae_reduction_pct": min_mae_reduction_pct,
            "min_tfw_ratio_vs_int4": min_tfw_ratio_vs_int4,
        },
        "precision": {
            "int4_mae": int4_mae,
            "int6_mae": int6_mae,
            "mae_reduction_pct": mae_reduction_pct,
            "int4_rmse": int4_rmse,
            "int6_rmse": int6_rmse,
            "rmse_reduction_pct": rmse_reduction_pct,
            "int4_max_abs_error": int4_max_abs,
            "int6_max_abs_error": int6_max_abs,
            "max_abs_reduction_pct": max_abs_reduction_pct,
        },
        "runtime": {
            "int4_tops_per_watt": int4_tfw_value,
            "int6_bridge_tops_per_watt": int6_tfw_value,
            "int6_vs_int4_tfw_ratio": tfw_ratio,
            "int6_metric_mode": int6_bridge.get("metric_mode"),
            "real_data_movement_measured": int6_bridge.get("real_data_movement_measured") is True,
            "power_measurement_representative": int6_bridge.get("power_measurement_representative") is True,
        },
        "size": {
            "int4_artifact_bytes_reference": int4_bytes,
            "int6_artifact_bytes_reference": int6_bytes,
            "int6_vs_int4_artifact_size_ratio": artifact_size_ratio,
        },
        "decision": {
            "precision_gate_passed": precision_gate,
            "runtime_gate_passed": runtime_gate,
            "tfw_gate_passed": tfw_gate,
            "int6_precision_tradeoff_winner": int6_tradeoff_wins,
            "tradeoff_score": tradeoff_score,
            "recommended_use": (
                "Use INT6 bridge for accuracy-sensitive paths where lower quantization drift matters."
                if int6_tradeoff_wins
                else "Keep INT4 as default until INT6 passes precision, real-data runtime, and TF/W-ratio gates."
            ),
        },
        "claim_gate": {
            "local_precision_tradeoff_claim_allowed": int6_tradeoff_wins,
            "world_best_precision_or_tfw_claim_allowed": False,
            "reason": "This is a local precision/runtime tradeoff report, not an external standardized leaderboard result.",
        },
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    json_path = out / "int6_precision_tradeoff_report.json"
    md_path = out / "int6_precision_tradeoff_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    p = report["precision"]
    r = report["runtime"]
    d = report["decision"]
    return "\n".join(
        [
            "# TinyMind INT6 Precision Tradeoff",
            "",
            f"- INT6 precision-tradeoff winner: {d['int6_precision_tradeoff_winner']}",
            f"- INT4 MAE: {p['int4_mae']:.8f}",
            f"- INT6 MAE: {p['int6_mae']:.8f}",
            f"- MAE reduction: {p['mae_reduction_pct']:.2f}%",
            f"- INT4 TOPS/W: {r['int4_tops_per_watt']:.6f}",
            f"- INT6 bridge TOPS/W: {r['int6_bridge_tops_per_watt']:.6f}",
            f"- INT6/INT4 TF/W ratio: {r['int6_vs_int4_tfw_ratio']:.4f}",
            f"- Tradeoff score: {d['tradeoff_score']:.6f}",
            f"- Recommendation: {d['recommended_use']}",
            "- World-best claim allowed: false",
            "",
        ]
    )