Skip to content

test: precision accuracy gate and per-device default precision for the OpenXLA backend (#449) #515

Description

@inureyes

Summary

Add a precision accuracy gate that compares low-precision OpenXLA output against the f32 / HF temp-0 reference across several prompts (token-exactness, and a divergence / perplexity margin where exactness is too strict), then use it to make f16 the default precision on GPU targets (Metal, CUDA) while keeping f32 where it does not help (CPU).

Scope

  • A gate harness (test or bench) that runs the OpenXLA path at f32 vs f16 over a small prompt set and asserts the accuracy criterion. Wire it into the test suite so precision changes are guarded.
  • Device-aware default precision: f16 on GPU HAL devices, f32 on local-task, overridable. Applied in the device / precision selection path (alongside default_device).

Integration (required for done)

The default precision is actually applied in the generation path per device (not just a helper), and the gate runs as part of cargo test. After this issue, MLXCEL_BACKEND=xla MLXCEL_XLA_DEVICE=metal runs f16 by default and the gate guards it.

Validation

  • Gate passes for f16 on the prompt set; a deliberately bad precision setting fails it.
  • tok/s improvement confirmed as the new default on Metal.

Depends on S1. Part of #513

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)status:doneCompleted

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions