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
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.
Part of #493
Context
Each architecture is currently validated ad hoc: dequant a checkpoint offline, run an HF fp32 oracle, then run
xla_oracle_checkby hand. As more families are added, this manual flow is the bottleneck.Scope
xla_batch_bench).Acceptance criteria
Implementation
The harness is split into two reusable tiers so the cheap one gates every change and the expensive one stays opt-in:
src/lib/mlxcel-xla/src/validation.rs: a per-architectureArchFixtureregistry +check_archthat parses a checkpoint'sconfig.json, emits each graph, and diffs it against the committed goldenassets/<arch>/*.mlir(first-line-diff localized on drift). It honorsMLXCEL_XLA_PRECISION(rejects a byte-exact run underf16/bf16, whose emit differs from the f32 goldens).emit_graphsis the golden-less freeze primitive for a brand-new family. This gate reports a clean pass onllama-3.2-1bundercargo test -p mlxcel-xla --lib validation, and the emitter's own byte-exact test now delegates to it (single source of truth).scripts/xla/validate_arch.shis the one command. It produces the HF fp32 oracle withspike/openxla/oracle_continuation.py(loads the checkpoint in fp32, dequantizing an MLX 4bit/8bit checkpoint offline first with the same affine formula asweights.rs), then runsxla_oracle_check(single-seq greedy == HF oracle) andxla_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.rsmodule docs andsrc/lib/mlxcel-xla/README.md.Dependencies
None. Used by Windows B and C to validate new families.