Skip to content

test: add a reusable per-architecture token-exact validation harness #496

Description

@inureyes

Part of #493

Context

Each architecture is currently validated ad hoc: dequant a checkpoint offline, run an HF fp32 oracle, then run xla_oracle_check by hand. As more families are added, this manual flow is the bottleneck.

Scope

  • Provide a small reusable tool or example that, given a quantized (MLX 4bit/8bit) or bf16 checkpoint, produces the HF fp32 oracle and runs both gates: the XLA token-exact single-sequence check and the serve reference-exact check (xla_batch_bench).
  • One command should take a checkpoint path and report pass/fail for both gates.
  • Files: a new example/binary under the workspace plus any shared helper it needs.

Acceptance criteria

  • One command validates a checkpoint end to end: offline dequant to f32 oracle, XLA token-exact single-seq, serve reference-exact.
  • Runs against an existing supported model (Llama, Qwen2, or Gemma2) and reports a clean pass.
  • Documented usage so adding a family is turnkey.

Implementation

The harness is split into two reusable tiers so the cheap one gates every change and the expensive one stays opt-in:

  • Structural (byte-exact, pure Rust, no GPU) in src/lib/mlxcel-xla/src/validation.rs: a per-architecture ArchFixture registry + check_arch that parses a checkpoint's config.json, emits each graph, and diffs it against the committed golden assets/<arch>/*.mlir (first-line-diff localized on drift). It honors MLXCEL_XLA_PRECISION (rejects a byte-exact run under f16/bf16, whose emit differs from the f32 goldens). emit_graphs is the golden-less freeze primitive for a brand-new family. This gate reports a clean pass on llama-3.2-1b under cargo test -p mlxcel-xla --lib validation, and the emitter's own byte-exact test now delegates to it (single source of truth).
  • Execution (token-exact + reference-exact): scripts/xla/validate_arch.sh is the one command. It produces the HF fp32 oracle with spike/openxla/oracle_continuation.py (loads the checkpoint in fp32, dequantizing an MLX 4bit/8bit checkpoint offline first with the same affine formula as weights.rs), then runs xla_oracle_check (single-seq greedy == HF oracle) and xla_batch_bench (each batched request == its single-seq reference), reporting pass/fail for both. The two run gates need a real IREE build (xla-iree) and a checkpoint, so a live clean pass is the opt-in GPU/CPU run; the offline-dequant math is self-tested against the Rust loader (oracle_continuation.py --selftest).

Adding a family is documented as turnkey (emit correctly, prove with the one command, optionally freeze goldens, register the fixture) in the validation.rs module docs and src/lib/mlxcel-xla/README.md.

Dependencies

None. Used by Windows B and C to validate new families.

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions