lablup / lablup/mlxcel

feat(xla): define operator-level MLX/IREE numeric contracts

Open
#932 11 comments 0 reactions 0 assignees View on GitHub
area:architecture area:core area:inference priority:high status:blocked type:enhancement
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Parent and affected work

- Parent epic: #566
- Current default-branch baseline: `ec6fbb6d` (shared foundation through #945)
- Deterministic MLX CUDA Q4 baseline: #923
- Affected family issues and draft PRs: #866 / #920, #868 / #921, #869 / #917, #870 / #918, #871 / #916, #872 / #919, and #878 / #922

## Problem

The remaining OpenXLA multimodal family ports now fail at a shared boundary rather than at missing model topology: MLX CUDA and StableHLO/IREE can execute mathematically equivalent operators with different operand materialization, accumulator precision, reassociation, reduction trees, fused-kernel ordering, or backend-selected convolution plans. The current family implementations encode some of these choices locally, so each full-model oracle run is rediscovering the same backend contract one architecture at a time.

This is not a request to relax tolerances. The existing family gates must remain unchanged, and final greedy output remains token-exact. If an MLX kernel has no deterministic, representable execution contract, the affected XLA capability must remain fail-closed until a reproducible contract is implemented or the reference path itself is made deterministic.

The reference baseline also changed after the current seven branches were cut: #923 fixed a shared-memory race in MLX CUDA `qmm_sm80` and changed the kernel artifact identity. Any Q4 evidence collected before `f7a877e7` must therefore be treated as historical until it is reproduced on the deterministic kernel.

## Current evidence

| Family | Draft PR | First unresolved boundary |
| --- | --- | --- |
| Molmo2 | #916 at `684c6ee1` | Current-head actual local-task gate reaches `projector.output_all`; first failure is flat index 358403 with `max_abs=0.1640625`, `rms=0.0026185848`. The immediately preceding `projector.product` passes (`max_abs=0.037597656`, `rms=0.0002196383`), isolating final dense-projection accumulation/association. |
| Gemma3 VLM | #917 at `9ab13992` | Current-head actual run was bounded before model start by a one-time pinned-MLX cache rebuild. The last valid actual evidence remains `siglip.hidden.block1.output` at `002bb77e`: flat index 658517, 8 failures, `max_abs=0.6672516`; current head passes the exact example check and focused Gemma3 tests but makes no new parity claim. |
| Molmo | #918 at `a4047b22` | Exact assets, hashes, compiler, and target were verified, but the release standalone runner exceeded the bounded five-minute GNU final-link window and produced no executable. No current-head checkpoint comparison ran; intermediate/KV-logit/token/lifecycle gates remain unqualified. |
| Youtu-VL | #919 at `20c2bffd` | Pinned HF and local-task IREE actuals agree through preprocessing, patch projection, and row-wise normalization. First unresolved stage is `layer.7.full`, flat index 310: actual/reference `0.3125`/`0.28515625`, `max_abs=0.203125`, 2,999 unchanged-threshold mismatches. |
| Qwen2.5-VL | #920 at `02d41a18` | Rebased on the shared foundation; focused Qwen2.5 tests pass. This host exposes no NVIDIA device, so post-#923 production MLX CUDA Q4 qmv/qmm evidence has not been reproduced and the family remains fail-closed. |
| Qwen3-VL | #921 at `a4b641e6` | Rebased on the shared foundation; Qwen3 and DeepStack focused tests pass. This host exposes no NVIDIA device, so the stale pre-#923 Q4 comparison is not promoted and current production Q4 parity remains unqualified. |
| Gemma3n audio | #922 at `0a5aa90d` | Shared CUDA-prefix contract is adopted and current diagnostics check plus 17 audio unit tests pass. The production first SSCP convolution still has 6 BF16-ULP residual failures with `max_abs=0.25`; no reproducible cuDNN-plan-equivalent IREE schedule is represented, so the capability remains disabled. |

## Goal

Define and enforce one versioned, operator-level numeric contract for MLX-to-StableHLO/IREE qualification, then use bounded micro-oracles to resolve shared drift before repeating full-model gates.

## Required implementation

- Add a versioned numeric-contract descriptor to the XLA auxiliary artifact identity. It must cover source checkpoint dtype, input/output materialization dtype, contraction accumulator/result dtype, explicit rounding boundaries, reduction or scan association policy where specified, quantized dequantization order, and any backend kernel or algorithm identity required for reproducibility.
- Represent the descriptor in shared XLA configuration and validation code rather than as family-specific strings. Artifact loading must reject a missing or mismatched numeric-contract fingerprint before native invocation.
- Add explicit emitter helpers for the contract-sensitive operations already used by the affected families: affine Q4 dequantization, dense/Q4 matrix multiplication, LayerNorm/RMSNorm reduction and restoration, residual addition, SiLU/GELU materialization, softmax/attention scaling, prefix reductions/scans, and convolution result typing.
- Keep the helpers semantically narrow. A helper may specify operand/result dtype and an explicit association schedule, but it must not silently apply graph-wide F16/BF16 casts or rewrite unrelated operations.
- Build deterministic micro-oracles that run identical operands through the MLX reference operation and the emitted IREE operation, record the first divergent output, and include exact backend/build/kernel identity. At minimum cover Q4 qmv/qmm on the post-#923 kernel, dense matmul, normalization, activation plus projection, residual add, attention/softmax, cumulative normalization, and the Gemma3n SSCP convolution shape.
- Separate three reference claims in reports: exact production-MLX equivalence, exact canonical decomposed-operation equivalence, and mathematical closeness. Only the first claim can satisfy a family production gate unless that family issue explicitly defines the canonical decomposed operation as its production reference.
- For unspecified backend algorithms such as cuDNN convolution plan selection, either import and pin a reproducible engine configuration, implement a proven matching schedule, or keep the affected capability disabled. Do not substitute a tolerance increase or a mathematical proxy for a missing execution contract.
- Rebase the seven affected family branches on the deterministic `f7a877e7` baseline before using their evidence. Re-run only the bounded operator probes first; repeat a heavyweight full-model oracle only after the first divergent operator boundary passes.
- Replace ad hoc family-local precision workarounds with the shared helpers when doing so preserves the family graph contract. Keep family-specific schedules only when the checkpoint architecture genuinely specifies a distinct operation.
- Update diagnostics so every actual report includes the numeric-contract version, MLX revision and patch identity, IREE/compiler version, target/device, selected kernel or algorithm identity when observable, and the exact first divergent tensor coordinate.
- Keep CLI/server capability predicates fail-closed while a required family contract or oracle is unresolved. Merging shared infrastructure must not advertise an unqualified family.

## Merge strategy

Implement the shared descriptor, fingerprint validation, explicit helpers, and micro-oracle harness in a narrow foundation PR first. After that PR lands, rebase the seven draft family PRs sequentially, resolve their shared-file overlap against the foundation, and run the affected bounded probes. A family PR becomes ready only after its original intermediate and token-exact gates pass; diagnostic or fail-closed infrastructure may be split into a separate PR when it is independently complete and does not expose the unfinished family.

## Non-goals

- Raising or removing the existing family-specific absolute, RMS, ULP, or token-exact acceptance thresholds.
- Treating generic StableHLO mathematical equivalence as proof of production MLX equivalence.
- Encoding one GPU model's undocumented kernel schedule as a repository-wide default.
- Re-running all heavyweight model gates before a bounded operator probe shows that the first divergent boundary is fixed.
- Advertising Molmo, Molmo2, Youtu-VL, Qwen2.5-VL, Qwen3-VL, Gemma3 VLM, or Gemma3n audio before its existing issue acceptance criteria pass.

## Validation

- Add unit tests for numeric-contract serialization, fingerprint stability, mismatch rejection, and every explicit materialization boundary.
- Add mutation-sensitive emitter tests that fail when a specified cast, accumulator/result dtype, operation order, or association schedule is removed or reassociated.
- Run deterministic MLX/IREE micro-oracles repeatedly on the production-relevant CUDA target and at least one local CPU target where the operation is supported.
- Reproduce the Qwen2.5-VL and Qwen3-VL Q4 probes after #923 and distinguish a corrected stale reference from a remaining representable contract mismatch.
- Re-run the first-divergence probes for Gemma3 VLM, Molmo2, and Gemma3n audio before any full-model rerun.
- Run the existing `mlxcel-xla` structural/golden tests, native IREE compile/load/invoke checks, text-only XLA architecture oracle, and continuous-batch regression suites.
- Record commands, target/device, compile time, transfers, memory, and first-divergence output without committing generated model artifacts or incidental CSV files.

## Acceptance criteria

- [ ] A versioned operator-level numeric contract is part of every affected auxiliary artifact fingerprint and is rejected on mismatch before invocation.
- [ ] Contract-sensitive emitter operations use shared explicit helpers with mutation-sensitive tests for dtype, rounding, operation order, and required association schedules.
- [ ] Deterministic MLX/IREE micro-oracles cover the listed operation classes and report reproducible backend identity plus the first divergent coordinate.
- [ ] Q4 evidence for #920 and #921 is refreshed on or after #923, with stale pre-fix evidence clearly separated.
- [ ] Each affected family either passes its unchanged original intermediate/token-exact gates or remains fail-closed with a precise unrepresentable backend contract recorded.
- [ ] No tolerance is relaxed and no mathematical proxy is reported as production-backend parity.
- [ ] Existing text-only XLA and continuous-batch behavior remains unchanged and qualified.

Contributor guide

Open the contributing guide

Research direction

Start with the shared XLA configuration and validation code and the auxiliary artifact identity, then inspect the existing mlxcel-xla structural/golden tests and native IREE compile/load/invoke checks. Done means the versioned contract is fingerprinted and rejected on mismatch, bounded micro-oracles cover the listed operations, diagnostics identify the first divergence, and unresolved family capabilities remain fail-closed.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
compilers, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.