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

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).