Add fixed MatMul scheduling controls (step 2 of #1002)

BUILD.bazel

@@ -802,6 +802,7 @@ cc_library(
         ":basics",
         ":configs",
         ":mat",
+        ":matmul_env",
         ":threading_context",
         "//compression:types",
         "//io",
@@ -1206,13 +1207,30 @@ cc_binary(
     ],
 )
 
+cc_library(
+    name = "model_comparison",
+    srcs = ["evals/model_comparison.cc"],
+    hdrs = ["evals/model_comparison.h"],
+    deps = ["@nlohmann_json//:json"],
+)
+
+cc_test(
+    name = "model_comparison_test",
+    srcs = ["evals/model_comparison_test.cc"],
+    deps = [
+        ":model_comparison",
+        "@googletest//:gtest_main",
+    ],
+)
+
 cc_binary(
     name = "gemma_mmlu",
     srcs = ["evals/run_mmlu.cc"],
     deps = [
         ":args",
         ":benchmark_helper",
         ":gemma_lib",
+        ":model_comparison",
         "//io",
         "@highway//:hwy",
         "@highway//:profiler",

CMakeLists.txt

@@ -351,6 +351,14 @@ 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})
+target_link_libraries(model_comparison PRIVATE nlohmann_json::nlohmann_json)
+
+add_executable(gemma_mmlu evals/run_mmlu.cc)
+target_link_libraries(gemma_mmlu PRIVATE 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)
@@ -359,6 +367,23 @@ enable_testing()
 find_package(GTest REQUIRED)
 include(GoogleTest)
 
+# Gemma's SIMD tests need Highway's test helpers even when its own test suite
+# is disabled (enabling both suites introduces duplicate target names).
+if(NOT TARGET hwy_test)
+  add_library(hwy_test ${highway_SOURCE_DIR}/hwy/tests/test_util.cc)
+  target_link_libraries(hwy_test PUBLIC hwy)
+endif()
+
+add_executable(model_comparison_test evals/model_comparison_test.cc)
+target_link_libraries(model_comparison_test PRIVATE model_comparison GTest::Main)
+gtest_discover_tests(model_comparison_test)
+find_package(Python3 3.10 COMPONENTS Interpreter QUIET)
+if(Python3_Interpreter_FOUND)
+  add_test(NAME model_comparison_python_test
+    COMMAND ${Python3_EXECUTABLE} -B -m unittest discover
+            -s ${CMAKE_CURRENT_SOURCE_DIR}/evals -p compare*_test.py)
+endif()
+
 # Tests use HWY_IS_TEST=1 with `hwy/foreach_target.h`, which expands the test
 # TU for every attainable SIMD target (including EMU128). libgemma's
 # foreach_target sources (flash_attention.cc, attention.cc, gemma.cc, vit.cc)

deepseek/run_dsv4.cc

@@ -161,7 +161,7 @@ int Main(int argc, char** argv) {
   consumed.AbortIfUnconsumed();
 
   ThreadingContext ctx(args.threading);
-  MatMulEnv env(ctx);
+  MatMulEnv env(ctx, args.inference.MatMulSchedule());
   Gemma gemma(args, ctx);
   KVCache kv_cache(gemma.Config(), args.inference, ctx.allocator);
 

evals/README.md

@@ -0,0 +1,152 @@
+# Model comparison and MMLU
+
+`gemma_mmlu` scores multiple-choice questions and emits machine-readable results.
+`compare_models.py` runs a root model and one or more targets **serially**, then
+writes `comparison.json` and `comparison.md`. It accepts arbitrary weight files,
+CLI arguments, and environment overrides; no quantization method is required.
+Models must use the same tokenizer and vocabulary for full-vocabulary KL.
+
+## Build and test
+
+From the repository root, with the usual CMake dependencies and GoogleTest:
+
+```sh
+cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \
+  -DGEMMA_ENABLE_TESTS=ON -DHWY_ENABLE_TESTS=OFF
+cmake --build build --target gemma_mmlu model_comparison_test -j2
+ctest --test-dir build -R 'ModelComparison|model_comparison_python' --output-on-failure
+python3 -B -m unittest discover -s evals -p 'compare*_test.py'
+```
+
+Python 3.10+ is required; the scripts use only the standard library. CMake
+registers Python tests when an interpreter is available. Optional cross-entropy
+runs also require `cmake --build build --target single_benchmark -j2`.
+Bazel targets `//:gemma_mmlu` and `//:model_comparison_test` are also provided.
+
+## Compare models
+
+Save a configuration such as `comparison.json`:
+
+```json
+{
+  "mmlu": "gemma/evals/mmlu.json",
+  "max_questions": 0,
+  "root": {
+    "name": "root",
+    "weights": "root.sbs",
+    "args": ["--num_threads", "6", "--pin", "1"],
+    "env": {}
+  },
+  "targets": [
+    {
+      "name": "target",
+      "weights": "target.sbs",
+      "args": ["--num_threads", "6", "--pin", "1"],
+      "env": {}
+    }
+  ]
+}
+```
+
+Paths for weights, MMLU, and optional `cross_entropy` text are relative to the
+configuration file. Extra CLI arguments retain their ordinary interpretation
+relative to the runner's working directory. Names must be unique and contain
+only letters, digits, underscores, dots, or hyphens, starting with a letter,
+digit, or underscore. The runner owns input, output, weights, question-limit,
+cross-entropy, and verbosity flags; these cannot be overridden in `args`.
+Environment overrides extend the calling process's environment. Record or clear
+any inherited settings relevant to the experiment, and use equivalent runtime
+settings across models when comparing performance.
+
+```sh
+python3 -B evals/compare_models.py --config comparison.json \
+  --build_dir build --output_dir eval-results/run-1
+```
+
+The output directory must be new or empty. Each model retains its stdout and
+stderr. The root additionally writes `<name>.root-kl.jsonl`, reused by all
+targets. A failed process or invalid result stops the series without producing a
+successful comparison report; existing logs remain available for diagnosis.
+`max_questions: 0` runs all samples; a positive value selects the first N.
+To measure run-to-run variation, use the same weights and settings for root
+and target. For reproducible MatMul schedules, add `--matmul_schedule fixed` to
+both models' `args` arrays. The available modes are:
+
+- `auto` (default): measure candidates and choose the fastest.
+- `fixed`: choose the first legal candidate without comparing timings.
+- `fixed_min_k`: choose the legal candidate with the fewest K partitions,
+  retaining candidate order to break ties. This reduces intermediate output
+  rounding but may be slower.
+
+Fixed modes also fix activation-conversion scheduling and, when enabled, the
+oneDNN BRGeMM candidate. A single candidate is measured once, then reused.
+The policy is set when `MatMulEnv` is constructed and cannot change within that
+environment. C++ callers can pass `MMSchedule::kFixed` or
+`MMSchedule::kFixedMinK` as its second constructor argument.
+
+For a self-comparison, keep the binary, SIMD target, thread/topology settings,
+batch sizes, weights, and prompts identical. Fixed scheduling removes variation
+from timing-based MatMul candidate selection; it does not promise identical
+results across hardware or builds. `--deterministic` separately controls sampling.
+The comparison report retains each model's scheduling option in its arguments.
+
+## Scoring and metrics
+
+Each question uses the dataset prompt, the runner's answer-only instruction,
+and the model's chat template. The runner scores the **first generated token**.
+It accepts single-token encodings of `A`–`D`, with or without a leading space,
+and uses the highest logit among each label's available spellings. A label with
+no single-token encoding causes an error. Equal label scores select A before B,
+then C, then D. This replaces the previous runner's free-form answer decoding;
+results from the two protocols should not be mixed.
+
+- `MMLU_RESULT`: question ID, expected/predicted label, correctness, four label
+  logits, their softmax probabilities, and the winning margin. These
+  probabilities normalize the four label scores; they are not full-vocabulary
+  token probabilities. Target runs also report `full_vocab_kl`.
+- `MMLU_SUMMARY`: number answered, number correct, and accuracy as a fraction.
+- `MMLU_KL_SUMMARY`: mean, median, p95, and maximum of
+  `KL(root || target)` in nats over the **whole vocabulary at the first answer
+  token**, after the model's logit soft cap. Percentiles use linear interpolation.
+  This is not KL averaged over every token of all possible answer strings.
+- **Flips**: correct-to-incorrect plus incorrect-to-correct transitions.
+  Wrong-to-wrong changes are separate; `answer_changes` counts all changed
+  labels. Aggregate accuracy can hide changes in individual answers.
+- `MMLU_TIMING`: time in `Generate`, time inside the evaluation sampling callback,
+  and their difference (`inference_seconds`). The latter includes MatMul
+  candidate measurements; it excludes prompt construction, reference
+  serialization/parsing, and result printing.
+- **Wall time** includes process startup, model loading, evaluation, and reference
+  I/O. It is not directly comparable between writing and reading references.
+  Peak RSS is sampled from Linux `/proc` every 20 ms and can miss short-lived
+  peaks; unavailable measurements are shown as null/—.
+- Optional **cross entropy** comes from `single_benchmark`: total bits divided
+  by input tokens, including its uniform first-token cost. The report uses the
+  final cumulative entropy and throughput, and compares bits/token to root.
+
+The bundled 83-question fixture is a small regression dataset, not the complete
+MMLU benchmark. A single run is not a statistically reliable speed measurement.
+
+## Reference files
+
+References are versioned JSON Lines, with one metadata header followed by one
+record per question. The header contains the exact dataset text and hex-encoded
+serialized tokenizer bytes. Records store IDs, expected labels, tokenized
+prompts, and float32 logits that round-trip without precision loss. Exact
+comparisons avoid hash collisions and reject mismatched datasets, tokenizers,
+question limits, or chat-template tokenization. Truncated, extra, malformed, or
+non-finite records fail explicitly.
+
+Text is inspectable and portable, but larger than binary; expect hundreds of
+MiB for an 83-question, 262k-vocabulary run. Targets stream one question at a
+time. Reference I/O is kept outside the reported `Generate` time.
+
+For a standalone root/target pair:
+
+```sh
+build/gemma_mmlu --weights root.sbs --input gemma/evals/mmlu.json \
+  --reference_out root.jsonl > root.out
+build/gemma_mmlu --weights target.sbs --input gemma/evals/mmlu.json \
+  --reference_in root.jsonl > target.out
+python3 -B evals/compare_mmlu.py root.out target.out
+```

evals/benchmark_helper.cc

@@ -37,7 +37,9 @@
 namespace gcpp {
 
 GemmaEnv::GemmaEnv(const GemmaArgs& args)
-    : ctx_(args.threading), env_(ctx_), gemma_(args, ctx_) {
+    : ctx_(args.threading),
+      env_(ctx_, args.inference.MatMulSchedule()),
+      gemma_(args, ctx_) {
   const ModelConfig& config = gemma_.Config();
 
   if (args.inference.verbosity >= 2) {

evals/compare_mmlu.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+# SPDX-License-Identifier: Apache-2.0
+"""Compare baseline and target 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) :])
+                if not isinstance(result, dict):
+                    raise ValueError("result must be an object")
+                question_id = result["id"]
+                if type(question_id) is not int:
+                    raise ValueError("id must be an integer")
+                if result["expected"] not in ("A", "B", "C", "D") or result[
+                    "predicted"
+                ] not in ("A", "B", "C", "D"):
+                    raise ValueError("expected/predicted must be A, B, C, or D")
+                if type(result["correct"]) is not bool or result["correct"] != (
+                    result["expected"] == result["predicted"]
+                ):
+                    raise ValueError("correct must match expected == 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]:
+    if not baseline or not variant:
+        raise ValueError("cannot compare empty results")
+    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())