feat(xla): add Gemma3n audio execution

[inureyes]

Summary

• add the split resident IREE Gemma3n audio encoder, projection, prepared-prefill, dense-PLE, CLI, and continuous-batch server path
• validate the pinned audio weight schema, cache identity, request ownership, cancellation, and slot reuse without a silent MLX fallback
• add an audio-only MLX-to-IREE intermediate diagnostic that avoids loading the 30 GB language bundle
• match the pinned MLX CUDA sequential BF16 affine dequantization contract and the explicit BF16 mel boundary before the first SSCP convolution
• keep BF16 SSCP operands with explicit F32 StableHLO convolution results followed by the existing BF16 output rounding
• match the MLX CUDA cumulative-normalization association order and expose the cumulative norm arithmetic boundaries
• replace the released dense-PLE width-256 RMS reduction with the explicit MLX CUDA lane/XOR schedule
• retain a diagnostics-only conv1 BF16-result candidate as negative evidence; production MLIR remains unchanged

Latest CUDA evidence

The actual MLX CUDA versus IREE CUDA diagnostic proves the following boundaries bit-exact:

• sscp_conv_0_convolution
• every exposed conv0 cumulative-normalization arithmetic stage and final norm
• sscp_conv_0 activation, which is the BF16 input consumed by conv1

The first divergence remains sscp_conv_1_convolution. The controlled result-type experiment is now complete:

Conv1 form	max_abs	RMS	values over one BF16 ULP
BF16 operands -> F32 StableHLO result -> BF16 round (production)	0.25	6.284481559e-4	6
BF16 operands -> BF16 StableHLO result -> F32 carrier (diagnostic candidate)	12	0.3255866858	419,814

The BF16-result candidate is decisively worse. Result typing is therefore not the fix: the production F32 accumulation/result followed by BF16 rounding is the correct representation and remains unchanged.

Backend audit

MLX CUDA's vendored convolution path constructs a cuDNN Frontend graph with BF16 I/O and F32 compute for BF16 operands. It requests Heuristic Mode A execution plans, checks support, builds a candidate, and caches the resulting graph by dtype, shapes, stride, padding, dilation, groups, and buffer alignments. It excludes plans marked DOWN_CONVERT_INPUTS; its TF32-specific tensor-core exclusion is only applied when the convolution I/O dtype itself is F32, not for this BF16-input convolution.

The selected cuDNN engine is backend state, not StableHLO state. The current MLX wrapper does not record the selected plan name/engine configuration or numerical notes. The vendored cuDNN Frontend can expose a selected plan name and plan metadata, and cuDNN API logging can expose backend calls, but obtaining that identity requires instrumenting the MLX/cuDNN layer or running an exact-shape cuDNN probe. Either approach can identify the MLX plan; neither creates an equivalent StableHLO/IREE selector.

A compiler-only exact-shape probe of the production conv1 operation was lowered with the pinned IREE CUDA compiler through executable-configurations. IREE emitted its own linalg.generic reduction with BF16-to-F32 extension, F32 multiply/add, and an IREE GPU reduction/workgroup tiling configuration. It did not emit a cuDNN execution plan. Repeating the same probe with StableHLO precision_config = HIGHEST produced the same IREE lowering configuration and arithmetic, so precision_config is not a usable diagnostics candidate for selecting the MLX/cuDNN reduction tree.

StableHLO convolution exposes semantic shape/layout/precision configuration, but no cuDNN engine global index, knob map, numerical-note filter, or exact reduction-tree attribute. Consequently, a tiny cuDNN-vs-StableHLO probe could reproduce and characterize the backend difference, but it cannot force IREE to execute the same cuDNN plan. The existing actual-checkpoint diagnostic already supplies the stronger comparison because it uses the exact conv1 input and weights.

The precise blocker is therefore backend-limited: matching the final six sparse BF16 residuals would require one of the following new backend capabilities, none of which is safe to guess in this PR:

1. an IREE CUDA external/cuDNN convolution path that can import and pin the same engine configuration and knobs;
2. an explicitly implemented custom CUDA reduction schedule proven to match the selected MLX cuDNN engine; or
3. an upstream compiler contract that represents convolution reduction/engine identity.

No production change, tolerance relaxation, heavyweight model rerun, or speculative candidate was added.

Validation

• actual CUDA result-type comparison recorded above
• cargo test -p mlxcel-xla emitter::gemma3n_audio_emit::tests:: — 5 passed, 5 environment-dependent tests ignored
• focused result-type identity test pins identical operands, padding, layout, stride, and shapes
• version-matched IREE 3.11.0rc20260316 CUDA compile of the released conv1-shape result-type pair passed
• exact-shape IREE executable-configurations audit: custom F32 reduction lowering confirmed; DEFAULT and HIGHEST select the same configuration
• cargo check -p mlxcel-xla --features diagnostics --tests
• cargo clippy -p mlxcel-xla --tests --no-deps — existing unrelated warnings remain
• cargo fmt --all -- --check
• git diff --check
• no generated CSV or model artifact is included

Remaining blocker

This PR remains draft and status:blocked because the unchanged MLX-to-IREE audio thresholds have not passed. The result-type hypothesis is closed; the remaining mismatch is an unrepresentable backend reduction/engine-identity difference.

Closes #878

Parent epic: #566

[inureyes]

Rebased onto main@7fe1412d and force-pushed as 876bcc3c.

Rebase integration:

• resolved the serde dependency spelling conflict in favor of current main
• combined feat(xla): version auxiliary dtype contracts #934 numeric-contract test imports with this branch’s typed auxiliary weight storage
• adapted the newer Qwen2-VL loader to AuxiliaryWeightStorage::Bytes
• kept both Gemma3n audio artifacts explicitly new_legacy_unqualified

Post-rebase validation:

• cargo fmt --all -- --check
• git diff --check
• cargo test -p mlxcel-xla emitter::gemma3n_audio_emit::tests::: 5 passed, 5 environment-dependent ignored
• cargo check -p mlxcel-xla --features diagnostics: passed with existing warnings
• cargo check -p mlxcel-xla --features diagnostics --tests: reaches only the existing mainline gemma3n_qmv test-only Builder cfg/import defect; the branch-specific Qwen auxiliary-storage mismatch was fixed

The unchanged MLX-to-IREE audio threshold blocker remains, and CUDA is unavailable on the current host. No new numeric pass is claimed. This PR remains draft/blocked and #878 remains blocked by #932.

[inureyes]

Foundation refresh complete at head 2b6d949:

• rebased the full draft cleanly onto main@15a37117 (feat(xla): add explicit CUDA prefix scan contract #943)
• replaced the family-local one-block cumulative-scan implementation with the shared mlx_cuda_prefix_sum_2d helper
• added a production-wrapper structural regression pinning the same 17-add / 3-concatenate CUDA hierarchy
• focused Gemma3n audio tests pass (6 runnable, 5 environment-dependent ignored), the shared schedule compiles through local-task CPU IREE, and diagnostics test compilation passes

The draft remains fail-closed and status:blocked. This refresh does not change the recorded first unresolved boundary: sscp_conv_1_convolution still requires a representable cuDNN/IREE engine or reduction-plan contract.

[inureyes]

Rebased onto current main (ec6fbb6d foundation; new head 0a5aa90d) and revalidated the Gemma3n audio branch.

Validation:

• cargo check -p mlxcel-xla --features diagnostics --lib — passed
• cargo test -p mlxcel-xla gemma3n_audio --lib — 17 passed, 6 ignored (environment-gated real IREE/checkpoint probes)
• cargo fmt --all --check — passed
• git diff --check — passed

The PR intentionally remains draft/fail-closed. The existing production-oracle blocker is unchanged: the first SSCP convolution still has 6 BF16-ULP mismatches (max_abs = 0.25), so the original #878 intermediate/token/lifecycle gates are not satisfied. No tolerance was relaxed, and this branch is not merge-qualified yet.

Refs #878
Refs #932

[inureyes]

Current-head CUDA actual at 0a5aa90d: the first SSCP convolution now passes

Run on the GB10 qualification host (driver 580.159.03, CUDA 13.0, IREE 3.12.0rc20260721,
Rust 1.93.1, MLX pin b7c3dd6d27f4, MLX_CUDA_ARCHITECTURES=121, release profile).
Working tree clean at 0a5aa90d. Tolerances unchanged, no override.

Fast pre-gate

cargo test --locked -p mlxcel-xla gemma3n_audio --lib: 17 passed, 6 ignored.

Small SSCP convolution characterization probe

./target/release/examples/xla_gemma3n_sscp_conv_oracle exits 0 and reports PASS:

channel=0 mlx=0.000000000e0 iree=0.000000000e0 bf16_ulp=0     schedule_outcomes=[0.0, 1.0]
channel=1 mlx=1.000000000e0 iree=1.000000000e0 bf16_ulp=0     schedule_outcomes=[0.0, 1.0]
channel=2 mlx=0.000000000e0 iree=0.000000000e0 bf16_ulp=0     schedule_outcomes=[0.0, 1.0]
channel=3 mlx=1.000000000e0 iree=0.000000000e0 bf16_ulp=16256 schedule_outcomes=[0.0, 1.0]
PASS (backend reduction schedules differ within the enumerated F32 accumulation envelope)

Channel 3 is worth recording explicitly: the probe enumerates the achievable F32 accumulation
orders and finds that both 0.0 and 1.0 are legitimate outcomes for the same sum. MLX
lands on one and IREE on the other. That is an exact-tie cancellation whose result is decided
by the reduction tree, so no precision policy can reconcile it, which supports the existing
characterization of this as a backend reduction/engine-identity difference rather than a
correctable semantic mismatch.

Full audio intermediate actual, MLXCEL_XLA_DEVICE=cuda

The recorded blocker for this PR was 6 BF16 ULP outliers with max_abs=0.25 on the first
SSCP convolution. At this head the first convolution no longer produces any BF16 ULP outlier,
and the first divergence has moved downstream:

first divergent stage=sscp_conv_0_norm_cumulative_sum
max_abs=1.562500000e-2  rms=2.432233746e-3  max_index=373
max_bf16_ulp=0  non_bf16_values=1024  over_one_bf16_ulp=0  captured_bf16_outliers=0

max_abs is about 16x smaller than the recorded blocker, and max_bf16_ulp,
over_one_bf16_ulp and captured_bf16_outliers are all zero at the failing stage. Mel
features, valid mask, frame lengths and bucket all pass ahead of it.

The gate still fails, now at the cumulative-sum normalization rather than the convolution, so
the PR stays draft and the capability stays fail-closed.

The reference policy was checked and is not responsible

I filed on #932 that MLX defaults MLX_ENABLE_TF32=1 and that
src/lib/mlxcel-xla/README.md declares that mode is not valid F32 reference evidence. It
decided the Molmo2 gate in #916, and mlx/backend/cuda/conv.cpp:139 gates the convolution
path on enable_tf32() independently of the matmul path, so it was a plausible explanation
for a convolution blocker and had to be tested.

It is not the explanation here. A controlled A/B at this head, identical in every other
respect, is byte-identical:

MLX_ENABLE_TF32	first divergent stage
1 (MLX default)	sscp_conv_0_norm_cumulative_sum max_abs 1.562500000e-2 rms 2.432233746e-3 index 373
0 (declared policy)	sscp_conv_0_norm_cumulative_sum max_abs 1.562500000e-2 rms 2.432233746e-3 index 373

So the improvement relative to the recorded blocker is attributable to this head rather than
to the reference policy, and the remaining divergence is a property of the emitted graph.
Recording the A/B because the negative result is what makes the attribution safe.

Remaining, unchanged: cumulative normalization and shared prefix scan, local attention,
relative shift, pooling, projection lengths, projected audio embeddings and dense PLE, then
logits/KV/greedy ASR token exactness and the CLI/server, isolation, cancellation and slot
reuse gates.