Add W8A8 int8 MatMul with optional rotations

CMakeLists.txt

@@ -357,6 +357,16 @@ target_link_libraries(benchmarks libgemma hwy hwy_contrib nlohmann_json::nlohman
 add_executable(debug_prompt evals/debug_prompt.cc)
 target_link_libraries(debug_prompt libgemma hwy hwy_contrib nlohmann_json::nlohmann_json)
 
+add_library(model_comparison evals/model_comparison.cc)
+target_include_directories(model_comparison PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
+
+add_executable(model_comparison_test evals/model_comparison_test.cc)
+target_link_libraries(model_comparison_test model_comparison)
+
+add_executable(gemma_mmlu evals/run_mmlu.cc)
+target_link_libraries(gemma_mmlu libgemma model_comparison hwy hwy_contrib
+                      nlohmann_json::nlohmann_json)
+
 ## Tests
 set(GEMMA_ENABLE_TESTS OFF CACHE BOOL "Enable Gemma tests")
 if (GEMMA_ENABLE_TESTS)
@@ -445,6 +455,27 @@ endif()  # GEMMA_ENABLE_TESTS
 
 ## Tools
 
+# Standalone W8A8 vs BF16 MatMul benchmark (no gtest). The _biased variant
+# forces the encoding x86 uses, to show what the bias correction costs.
+add_executable(bench_matmul_i8 ops/bench_matmul_i8.cc)
+target_link_libraries(bench_matmul_i8 libgemma hwy hwy_contrib)
+
+add_executable(bench_matmul_i8_biased ops/bench_matmul_i8.cc)
+target_compile_definitions(bench_matmul_i8_biased PRIVATE
+                           GEMMA_MM_I8_FORCE_BIASED_B=1)
+target_link_libraries(bench_matmul_i8_biased libgemma hwy hwy_contrib)
+
+# W8A8 correctness, built once per A encoding so that the x86 biased-u8 path
+# is covered on non-x86 hosts as well.
+add_executable(matmul_i8_test ops/matmul_i8_test.cc)
+target_compile_definitions(matmul_i8_test PRIVATE GEMMA_MM_I8_FORCE_BIASED_B=0)
+target_link_libraries(matmul_i8_test libgemma hwy hwy_contrib)
+
+add_executable(matmul_i8_biased_test ops/matmul_i8_test.cc)
+target_compile_definitions(matmul_i8_biased_test PRIVATE
+                           GEMMA_MM_I8_FORCE_BIASED_B=1)
+target_link_libraries(matmul_i8_biased_test libgemma hwy hwy_contrib)
+
 add_executable(migrate_weights io/migrate_weights.cc)
 target_link_libraries(migrate_weights libgemma hwy hwy_contrib)
 

evals/compare_mmlu.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""Compare baseline and compressed gemma_mmlu outputs.
+
+Each input is the stdout captured from gemma_mmlu and may contain unrelated
+lines. Only lines beginning with ``MMLU_RESULT `` are parsed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+
+RESULT_PREFIX = "MMLU_RESULT "
+
+
+def load_results(path: Path) -> dict[int, dict[str, Any]]:
+    results: dict[int, dict[str, Any]] = {}
+    with path.open(encoding="utf-8") as source:
+        for line_number, line in enumerate(source, start=1):
+            if not line.startswith(RESULT_PREFIX):
+                continue
+            try:
+                result = json.loads(line[len(RESULT_PREFIX) :])
+                question_id = int(result["id"])
+                result["correct"] = bool(result["correct"])
+                result["expected"] = str(result["expected"])
+                result["predicted"] = str(result["predicted"])
+            except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error:
+                raise ValueError(
+                    f"{path}:{line_number}: invalid MMLU_RESULT: {error}"
+                ) from error
+            if question_id in results:
+                raise ValueError(f"{path}:{line_number}: duplicate id {question_id}")
+            results[question_id] = result
+
+    if not results:
+        raise ValueError(f"{path}: no {RESULT_PREFIX.strip()} lines found")
+    return results
+
+
+def compare_results(
+    baseline: dict[int, dict[str, Any]], variant: dict[int, dict[str, Any]]
+) -> dict[str, int | float]:
+    baseline_ids = set(baseline)
+    variant_ids = set(variant)
+    if baseline_ids != variant_ids:
+        missing = sorted(baseline_ids - variant_ids)
+        extra = sorted(variant_ids - baseline_ids)
+        raise ValueError(
+            "result IDs differ: "
+            f"missing from variant={missing[:10]}, extra in variant={extra[:10]}"
+        )
+
+    correct_to_incorrect = 0
+    incorrect_to_correct = 0
+    wrong_to_wrong_changes = 0
+    answer_changes = 0
+    baseline_correct = 0
+    variant_correct = 0
+
+    for question_id in sorted(baseline_ids):
+        base = baseline[question_id]
+        changed = variant[question_id]
+        if base["expected"] != changed["expected"]:
+            raise ValueError(
+                f"id {question_id}: expected answers differ: "
+                f"{base['expected']!r} != {changed['expected']!r}"
+            )
+
+        base_correct = base["correct"]
+        changed_correct = changed["correct"]
+        baseline_correct += int(base_correct)
+        variant_correct += int(changed_correct)
+        answer_changed = base["predicted"] != changed["predicted"]
+        answer_changes += int(answer_changed)
+
+        if base_correct and not changed_correct:
+            correct_to_incorrect += 1
+        elif not base_correct and changed_correct:
+            incorrect_to_correct += 1
+        elif not base_correct and not changed_correct and answer_changed:
+            wrong_to_wrong_changes += 1
+
+    samples = len(baseline_ids)
+    flips = correct_to_incorrect + incorrect_to_correct
+    return {
+        "samples": samples,
+        "baseline_correct": baseline_correct,
+        "variant_correct": variant_correct,
+        "baseline_accuracy": baseline_correct / samples,
+        "variant_accuracy": variant_correct / samples,
+        "accuracy_delta": (variant_correct - baseline_correct) / samples,
+        "correct_to_incorrect": correct_to_incorrect,
+        "incorrect_to_correct": incorrect_to_correct,
+        "flips": flips,
+        "flips_fraction": flips / samples,
+        "flips_percent": 100.0 * flips / samples,
+        "wrong_to_wrong_changes": wrong_to_wrong_changes,
+        "answer_changes": answer_changes,
+        "answer_changes_fraction": answer_changes / samples,
+        "answer_changes_percent": 100.0 * answer_changes / samples,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compare baseline and variant gemma_mmlu output streams."
+    )
+    parser.add_argument("baseline", type=Path, help="baseline gemma_mmlu stdout")
+    parser.add_argument("variant", type=Path, help="variant gemma_mmlu stdout")
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    try:
+        metrics = compare_results(
+            load_results(args.baseline), load_results(args.variant)
+        )
+    except (OSError, ValueError) as error:
+        print(f"error: {error}", file=sys.stderr)
+        return 1
+
+    print(f"MMLU_FLIPS {json.dumps(metrics, sort_keys=True)}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evals/compare_mmlu_test.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+from compare_mmlu import compare_results, load_results
+
+
+def result(question_id: int, expected: str, predicted: str) -> dict[str, object]:
+    return {
+        "id": question_id,
+        "expected": expected,
+        "predicted": predicted,
+        "correct": expected == predicted,
+    }
+
+
+class CompareMmluTest(unittest.TestCase):
+    def test_flip_counts(self) -> None:
+        baseline = {
+            1: result(1, "A", "A"),
+            2: result(2, "A", "B"),
+            3: result(3, "A", "C"),
+            4: result(4, "D", "D"),
+        }
+        variant = {
+            1: result(1, "A", "B"),
+            2: result(2, "A", "A"),
+            3: result(3, "A", "D"),
+            4: result(4, "D", "D"),
+        }
+
+        metrics = compare_results(baseline, variant)
+
+        self.assertEqual(metrics["correct_to_incorrect"], 1)
+        self.assertEqual(metrics["incorrect_to_correct"], 1)
+        self.assertEqual(metrics["flips"], 2)
+        self.assertEqual(metrics["flips_percent"], 50.0)
+        self.assertEqual(metrics["wrong_to_wrong_changes"], 1)
+        self.assertEqual(metrics["answer_changes"], 3)
+        self.assertEqual(metrics["answer_changes_percent"], 75.0)
+        self.assertEqual(metrics["accuracy_delta"], 0.0)
+
+    def test_loads_prefixed_results_and_ignores_other_lines(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "run.log"
+            rows = [result(7, "B", "B"), result(8, "C", "A")]
+            path.write_text(
+                "startup noise\n"
+                + "\n".join(f"MMLU_RESULT {json.dumps(row)}" for row in rows)
+                + "\nMMLU_SUMMARY {}\n",
+                encoding="utf-8",
+            )
+
+            loaded = load_results(path)
+
+        self.assertEqual(set(loaded), {7, 8})
+        self.assertTrue(loaded[7]["correct"])
+        self.assertFalse(loaded[8]["correct"])
+
+    def test_requires_matching_question_ids(self) -> None:
+        with self.assertRaisesRegex(ValueError, "result IDs differ"):
+            compare_results(
+                {1: result(1, "A", "A")}, {2: result(2, "A", "A")}
+            )
+
+    def test_requires_matching_expected_answers(self) -> None:
+        with self.assertRaisesRegex(ValueError, "expected answers differ"):
+            compare_results(
+                {1: result(1, "A", "A")}, {1: result(1, "B", "B")}
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()