chore(ci): compile the OpenXLA feature combinations in CI

[inureyes]

Problem / Background

CI for this repository builds only the default feature set, so anything that lives behind xla-backend, xla-iree, or xla-diagnostics is never compiled by a PR check and breakage lands on main unnoticed. One such break currently blocks every xla-iree build. It was found while rebasing PR #916 onto main at commit 2766abd0, not by any automated gate.

Current Behavior

The OpenXLA serve worker does not handle ModelRequest::PromptCacheWarmup. PR #1154 added the PromptCacheWarmup variant in src/server/model_provider.rs:87 and updated the three workers that match on ModelRequest: src/server/batch/scheduler.rs:3569, src/server/diffusion_worker.rs:195, and src/server/florence2_worker.rs:254. It did not update src/server/batch/xla_worker_admission.rs, whose XlaServeWorker::handle matches on request at roughly line 482. That match is therefore non-exhaustive and any --features xla-iree build fails with error[E0004]: non-exhaustive patterns: ModelRequest::PromptCacheWarmup { .. } not covered. The whole module is gated (src/server/batch/mod.rs:75-78 puts xla_preprocess and xla_worker behind #[cfg(feature = "xla-iree")]), which is exactly why nothing in CI compiled it.

An earlier revision of this issue listed a second break, an unconditional re-export of load_weights_from_dir_with_filter in src/models/mod.rs. That was misattributed and has been removed. On main that re-export does not exist and src/models/sanitize.rs:1256 declares the function as a private fn; both the re-export and the pub(crate) visibility bump come from PR #916's own diff, so it is a review item for that PR rather than breakage on main.

Proposed Solution

Add a ModelRequest::PromptCacheWarmup { .. } arm to XlaServeWorker::handle. The XLA worker owns no prompt cache (only BatchScheduler holds snapshot state), so the correct handling is to drop the job the way src/server/diffusion_worker.rs:195 already does, with a comment saying why, rather than forwarding it or erroring.

The code fix is already prepared and is being submitted as a small standalone PR. The durable problem this issue tracks is the coverage gap that let this land: nothing in .github/workflows/ compiles any XLA feature combination. ci.yml runs deny, fmt, crate-version, kernel-dtype-key, MLX-pin, and cross-repo-ref jobs but no workspace cargo check; pipeline-parallel-ci.yml:189 runs cargo clippy -p mlxcel --lib --tests on default features; nightly-verify.yml:318 runs clippy with metal,accelerate. None of these touch xla-backend, xla-iree, or xla-diagnostics.

Proposal: add at least a compile-only CI job that runs cargo check over the XLA feature combinations, so this class of breakage cannot land silently again.

Scope

In scope: the fix in src/server/batch/xla_worker_admission.rs (already prepared, separate PR), and the design and addition of a compile-only XLA CI job.

Out of scope: running XLA tests in CI, GPU-backed XLA execution in CI, and any change to the OpenXLA backend's behavior.

Implementation Notes

• The cheap half is free, the useful half is not. xla-backend = ["dep:mlxcel-xla"] pulls in the crate without the IREE shim (mlxcel-xla's build script skips it without the iree feature), so cargo check --features xla-backend --all-targets needs no external toolchain and could run on any GitHub-hosted runner today. It would not have caught this break, which lives behind xla-iree.
• The tradeoff is the substance of this issue. xla-iree = ["xla-backend", "mlxcel-xla/iree"] compiles a C shim against a prebuilt IREE runtime and needs IREE_DIST at build time. scripts/iree/setup-cuda.sh provisions it by installing the iree-base-compiler wheel into a private venv and source-building the IREE runtime (runtime only, no LLVM) pinned to the wheel's exact revision. That is far too slow to run per PR from scratch, so a compile-only xla-iree job needs a cached or prebuilt IREE runtime, and deciding where that artifact comes from (an actions cache keyed on the pinned IREE_SHA, a container image, or a self-hosted runner that already has one) is the real decision to make here.
• xla-diagnostics implies cuda. xla-diagnostics = ["cuda", "xla-iree", "mlxcel-xla/diagnostics"], so a job that checks it needs the CUDA toolchain as well, which points at the self-hosted GB10 runner rather than a hosted one.
• Reuse: the existing self-hosted-runner and caching patterns in nightly-verify.yml, rather than a new provisioning scheme.
• Edge case: cargo check without --all-targets misses test-only code, and a dead re-export is only rejected under -D warnings, so the job must carry the same warning policy the other clippy jobs do or it will pass on a tree that the real gates reject.

Acceptance Criteria

• cargo check --features cuda,xla-iree --all-targets succeeds on main with no E0004 and no dead-code warning.
• A CI job exists that compiles at least one XLA feature combination and fails the check run when it does not build.
• The job's feature matrix and its IREE provisioning strategy are written down in the workflow file, including which combinations are deliberately not covered and why.
• Reverting the fix makes the new job fail, demonstrating it actually covers this class of breakage.

Verification

eval "$(bash scripts/iree/setup-cuda.sh --env)"
cargo check --features cuda,xla-iree --all-targets
cargo check --features xla-backend --all-targets

A pass is a clean exit from both, with no warning emitted (the second command needs no IREE toolchain and can be run first as a fast smoke).

Technical Considerations

Related: PR #1154 (introduced PromptCacheWarmup), PR #916 (where this surfaced during rebase).

[inureyes]

Reopened and retitled. PR #1273 landed the compile fix that unblocked --features xla-iree, and the Closes #1270 keyword closed this along with it, but that was only the first of the four acceptance criteria. The remaining three are the substance: a CI job that actually compiles an XLA feature combination, its feature matrix and IREE provisioning strategy written down, and a demonstration that reverting a break makes the job fail.

The old title ("main does not compile under the OpenXLA features") is no longer true, so it has been changed to name the remaining work. The type: label moved from bug to chore for the same reason.

Evidence that the coverage gap is real rather than theoretical: it has now let two separate defects reach main. The PromptCacheWarmup break fixed in #1273, and the link failure filed as #1274, which cargo check cannot catch because it never links.