fix(mllama): ragged cross-attention test fails on main, and no CI job runs cargo test

[inureyes]

Problem / Background

Two distinct problems, found incidentally while reviewing #930. Neither is caused by that PR.

1. ragged_real_tile_rows_match_reference_masked_full_rows fails deterministically on main

Reproduced on main at 7fe1412d8, in a clean worktree, on an Apple M1 Ultra running macOS 26.5.2:

DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer \
  cargo test --release --lib --features metal,accelerate models::mllama

test models::mllama::text::tests::ragged_real_tile_rows_match_reference_masked_full_rows ... FAILED

thread '...' panicked at src/models/mllama/text.rs:664:9:
assertion `left == right` failed: ragged per-image selection must be byte-identical to the reference-masked full attention
  left: 3.7252903e-9
 right: 0.0

test result: FAILED. 7 passed; 1 failed; 0 ignored; 0 measured; 4719 filtered out

The observed value 3.7252903e-9 is exactly 2^-28. The assertion at src/models/mllama/text.rs:664 demands exact equality:

assert_eq!(
    max_abs_diff(&masked_out, &sliced_out),
    0.0,
    "ragged per-image selection must be byte-identical to the \
     reference-masked full attention"
);

Still reproduces after main advanced to a3c823611, with the identical value, on the same machine:

running 8 tests
test models::mllama::config::tests::text_config_defaults_match_reference ... ok
test models::mllama::config::tests::vision_config_reads_inherited_quantization ... ok
test models::mllama::config::tests::parses_nested_mllama_config ... ok
test models::mllama::text::tests::cached_cross_kv_equals_fresh_recompute ... ok
test models::mllama::text::tests::cross_states_change_invalidates_cache ... ok
test models::mllama::text::tests::full_tile_mask_is_the_unmasked_computation ... ok
test models::mllama::text::tests::ragged_real_tile_rows_match_reference_masked_full_rows ... FAILED
test models::mllama::text::tests::real_tile_rows_match_reference_masked_full_rows ... ok

thread '...' panicked at src/models/mllama/text.rs:664:9:
assertion `left == right` failed: ragged per-image selection must be byte-identical to the reference-masked full attention
  left: 3.7252903e-9
 right: 0.0

test result: FAILED. 7 passed; 1 failed; 0 ignored; 0 measured; 4798 filtered out; finished in 0.08s

That run also confirms the premise behind the "leave the siblings alone" criterion below: real_tile_rows_match_reference_masked_full_rows and full_tile_mask_is_the_unmasked_computation both pass on their exact assert_eq!(..., 0.0).

Nothing between 0ad300c3f7 and a3c823611 touches this test or anything it depends on either: git log 0ad300c3f..HEAD -- src/models/mllama/ is empty, attention_from_ptr in src/lib/mlxcel-core/src/layers.rs last changed in 9c4eeb771 (#149), which predates the test, and MLX_EXPECTED_COMMIT last moved in ac0d86c0c (#270) on 2026-06-14, also before the test landed. The assertion is still verbatim at text.rs:664-669.

2. No workflow runs cargo test on this code, which is why this went unnoticed

.github/workflows/ci.yml defines exactly four jobs, none of which compiles or tests the crate:

Job	Runner	What it runs
changes (Detect changes)	ubuntu-latest	dorny/paths-filter@v4
deny (cargo-deny)	ubuntu-latest	EmbarkStudios/cargo-deny-action@v2
fmt (cargo-fmt)	ubuntu-latest	cargo fmt --all -- --check
cross-repo-refs	ubuntu-latest	python3 scripts/ci/check_cross_repo_refs.py (advisory, does not fail the build)

There is exactly one narrow exception elsewhere, and it does not cover this code. .github/workflows/pipeline-parallel-ci.yml runs cargo clippy -p mlxcel --lib --tests -- -D warnings directly (lines 187-193) and reaches cargo test indirectly through scripts/ci/run-pp-heterogeneous-memory.sh and scripts/ci/run-pp-two-host-logical.sh, whose selectors are distributed::pipeline::, distributed::cluster_init::, distributed::tcp_transport:: and distributed::rdma_transport:: plus two named cases in tests/pipeline_ci_multi_stage_real_models.rs. That job is ubuntu-latest, debug profile, and path-filtered to src/distributed/pipeline/** and its siblings, so it never runs on a model-port PR and never selects models::mllama or main_tests. release.yml runs cargo build --release plus a 10-token generate smoke test, but never the test suite. No workflow anywhere selects the model test modules.

That exception also makes the ci.yml header comment inaccurate in its own right: it states "Clippy and cargo-test do NOT run in any workflow", which was true when it was written and is no longer.

Investigation

The failure is a floating-point reassociation artifact, not a logic error

Evidence, gathered by instrumenting the test in a throwaway worktree:

Probe	kv rows kept	masked columns	max abs diff
A (the failing test)	6, at {0,1,4,5,6,7}	interior {2,3}	3.7252903e-9
B	6, at {0..5}	trailing {6,7}	0.0
C	4, at {0..3}	trailing {4..7}	0.0
D	5, at {0..4}	trailing {5,6,7}	0.0
E	concatenate-built vs directly allocated 6-row input, both unmasked	none	0.0
G	interior-masked 8 rows vs directly allocated contiguous 6 rows	interior {2,3}	3.7252903e-9

Probe E rules out the concatenate construction: building the same 6 rows by concatenating two slices is bit-identical to allocating them contiguously. Probe G then reproduces the failure with no concatenate involved, so the trigger is the interior mask position alone.

Probes B, C and D are the cases where the surviving rows keep their lane positions, because the masked columns are trailing. All three are exactly 0.0, which confirms the premise from PR #622 that exp(logit - 1e9) underflows to exactly 0.0 and contributes exact zeros. Probes A and G are the only cases where the surviving rows move: they sit at lane positions {0,1,4,5,6,7} in the 8-wide reduction but at {0..5} in the 6-wide reduction. The softmax denominator and the value accumulation are therefore summed in a different association order, and f32 addition is not associative, so the two results differ in the last bit.

Magnitude check: the output's max absolute value is 1.3224149e-1, whose ULP is 1.4901161e-8. The difference is 0.25 ULP of the largest output element, and exactly 1 ULP for an element in [2^-5, 2^-4). This is the smallest non-zero difference f32 can represent at that scale.

Conclusion: a tolerance is the correct fix, and the production code is right

The exact-zero expectation over-claims. The reasoning in PR #622 proves equality in exact arithmetic, which is what licenses select_real_tile_states in src/vision/mllama_vl.rs to drop the padding-tile rows. Exact arithmetic does not imply bitwise equality once the surviving lanes change position inside a parallel reduction.

Two further points support this. First, the padding rows carry content of the same magnitude as the real rows, so a genuine masking or row-selection error would show up at order 1e-1, roughly eight orders of magnitude above the observed artifact. Second, only the sliced path ever runs in production: the port threads no text-side mask at all, so the masked branch exists solely inside this test.

What would settle it beyond the above: running the same comparison on a different backend or reduction shape. If the difference is reassociation, it should move or vanish as the kernel tiling changes, rather than staying fixed. That is a confirmation step, not a prerequisite, since probes B through E already isolate the cause.

Attribution: PR #622 is correct, and the test was red on arrival

git blame places the entire test and its assertion on 0ad300c3f7 ("perf(vision/mllama): build cross-attention states from real tiles only (#622)", merged 2026-07-02), which introduced both the test and the code under test. Building and running at that exact commit reproduces the identical failure with the identical value:

test models::mllama::text::tests::ragged_real_tile_rows_match_reference_masked_full_rows ... FAILED
  left: 3.7252903e-9
 right: 0.0

So this is not a later regression, and the attribution to #622 stands. Supporting evidence: nothing has touched src/models/mllama/ since that commit, attention_from_ptr in src/lib/mlxcel-core/src/layers.rs is unchanged since it, and MLX_EXPECTED_COMMIT last moved on 2026-06-14, before the test landed.

The recent model ports are ruled out

The test builds only MllamaTextCrossAttention. build_layer constructs that layer directly from synthetic weights, and its forward uses only UnifiedLinear, RMSNorm and the shared attention_from_ptr kernel. It never constructs llama3::Attention and never applies RoPE, since cross-attention carries no positional encoding on the key side. The recent GPT-family and Helium ports cannot be involved.

The absent gate is the more consequential half, and there are now two instances

Skipping clippy and test in CI is a deliberate posture, documented in the ci.yml header comment and decided in 0fd5fa371 ("chore(ci): drop clippy+test from release pipeline; keep gate local-only (#23)"): both took roughly 30 minutes per run on the shared self-hosted Apple Silicon runner, for failures that make verify catches locally in a fraction of the time. That trade-off is defensible on its own terms, and the repository's agent guidance says the same thing: local and real-model testing is the gate.

The problem is the outcome, not the policy. A deterministic unit-test failure has sat on main since 2026-07-02, and every contributor who has run make verify since then has seen a red suite that is not theirs. That is exactly how a suite stops being a signal.

This is no longer a single anecdote. A second deterministically-failing test was found on main during the same series of model ports: family_order_is_exhaustive in src/main_tests.rs, which asserts that every ModelType::family() string appears in FAMILY_ORDER in src/main.rs. It was failing with missing: ["MiniMax VLM", "Step VLM"], because #800 (MiniMax-M3-VL) and #781 (Step-3) each added a family without adding it to that array. Like the mllama case, it is deterministic, it needs no model weights, it reproduces on a plain cargo test --release --lib, and it sat on main unnoticed. It was repaired in #946, alongside the "Bailing" entry that port genuinely needed, and the fix is verifiable by reverting the three added entries and watching family_order_is_exhaustive fail with missing: ["MiniMax VLM", "Bailing", "Step VLM"].

Two independent deterministic failures reaching main through the same gap is the substantive argument for this half of the issue: the failure mode is systemic rather than a one-off.

CONTRIBUTING.md also makes this worse by describing a gate that does not exist. Lines 45, 46 and 49 still say:

cargo clippy --all-targets --features metal,accelerate -- -D warnings # enforced by CI on self-hosted macOS runner

cargo test --release --features metal,accelerate # enforced by CI on self-hosted macOS runner

CI enforces clippy (with -D warnings) and cargo test on the self-hosted-macos-26-arm64 runner on every PR that touches Rust files.

All three came from da318614b (#14), which genuinely did add that gate. 63d21de0e (#21) then moved it to release.yml, and 0fd5fa371 (#23) removed it entirely, but neither commit updated CONTRIBUTING.md. The claim has been false since 2026-05-18.

Proposed Solution

Half 1: the assertion

Replace the exact assert_eq!(..., 0.0) in ragged_real_tile_rows_match_reference_masked_full_rows with a tolerance assertion, and add a comment explaining why this case tolerates a last-bit difference while its three siblings do not. A bound of 1e-6 sits far below anything a real masking or row-selection error could produce.

Leave real_tile_rows_match_reference_masked_full_rows and full_tile_mask_is_the_unmasked_computation on exact equality. They pass today, they pin the genuinely bitwise-exact property, and weakening them would lose coverage.

Half 2: the gate

Choose explicitly between two options and record the decision:

1. Restore a PR-time cargo test gate on the self-hosted Apple Silicon runner, accepting the runtime cost, optionally narrowed to --lib or to a path-filtered subset so the common case stays cheap.
2. Keep the gate local-only, and instead make a red main impossible to miss: a scheduled run (nightly, for example) on the self-hosted runner that reports failures, so a broken test surfaces within a day rather than on the next contributor's make verify.

Either way, CONTRIBUTING.md must stop claiming a PR-time CI gate that does not exist, and the ci.yml header comment must stop claiming that no workflow runs clippy or cargo test when pipeline-parallel-ci.yml does both on a path-filtered subset.

Acceptance Criteria

• ragged_real_tile_rows_match_reference_masked_full_rows passes on main via DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer cargo test --release --lib --features metal,accelerate models::mllama
• The changed assertion carries a comment recording that the surviving rows change lane position inside the reduction, so the equality holds in exact arithmetic but not bitwise
• real_tile_rows_match_reference_masked_full_rows and full_tile_mask_is_the_unmasked_computation keep their exact assert_eq!(..., 0.0) form and still pass
• No production change to select_real_tile_states or MllamaTextCrossAttention, unless new evidence overturns the tolerance conclusion
• cargo test --release --features metal,accelerate is green on main on Apple Silicon, or each remaining failure is filed as its own issue
• A decision between the PR-time gate and the scheduled run is recorded, with its rationale
• The chosen mechanism is implemented and demonstrated to catch a deliberately failing test before it can reach main unnoticed. Both known instances are usable as the demonstration case: the mllama assertion above, and family_order_is_exhaustive with its FAMILY_ORDER entries reverted.
• CONTRIBUTING.md lines 45, 46 and 49 describe the gate that actually exists
• The ci.yml header comment, CONTRIBUTING.md, and the actual workflow set agree with each other, including the path-filtered clippy and cargo test that pipeline-parallel-ci.yml does run

Technical Considerations

Reproducing locally requires DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer, otherwise cmake fails with xcrun: error: unable to find utility "metal". Use the narrow models::mllama selector, since a full run is long.

The artifact is expected to depend on hardware and kernel choice, because it originates in the reduction shape the Metal SDPA kernel selects. It was verified on an Apple M1 Ultra under macOS 26.5.2. Whether it reproduces on other Apple Silicon parts or under CUDA is unknown, which is a further argument for a tolerance over chasing bitwise equality.

max_abs_diff reduces with max_all(abs(a - b)), so the reported number is the single worst element, not an aggregate.

Implementation context (machine facts, validation method, and a recurring defect class)

Most relevant to this issue: the build and test facts for the current development machine, and the CI note, which is literally the second half of this issue. The DEVELOPER_DIR requirement, the release-profile-only warmth, and the narrow-selector rule are the same ones this issue's reproduction command depends on.

Collected while porting five text-model families in #924, #926, #928, #930, and #946. This
is background for whoever picks this up, not additional scope.

Build and test on the current development machine (Apple M1 Ultra, macOS)

• Every cargo invocation needs DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer, or cmake fails with xcrun: error: unable to find utility "metal". xcode-select still points at CommandLineTools.
• Only the RELEASE profile is warm. A debug cargo check triggers a COLD MLX C++ build that takes many minutes. Use DEVELOPER_DIR=... cargo check --release --lib --tests --features metal,accelerate.
• Narrow test selectors only: DEVELOPER_DIR=... cargo test --release --lib --features metal,accelerate models::<module>. A bare cargo test, cargo test --lib, or cargo clippy --workspace --all-targets runs long enough to stall an agent with a stream-idle watchdog, and a cold cargo build --release can exceed it on the MLX C++ link step alone.
• cargo fmt --all -- --check, never a bare-path fmt check, which produces spurious edition-2024 diffs.
• CI runs no cargo test at all (only changes, cargo-deny, cargo-fmt, and the cross-repo-ref guard). Local runs plus a real checkpoint are the only gate. Two deterministically failing tests were found sitting on main during this series for exactly that reason: see fix(mllama): ragged cross-attention test fails on main, and no CI job runs cargo test #939, and the FAMILY_ORDER entries fixed in feat(models): add Ant Group Ling / Bailing MoE (bailing_moe) text model support #946.
• tracing::warn! is a NO-OP in the mlxcel CLI binary; only src/server/startup.rs installs a subscriber. Use eprintln! for any CLI-facing diagnostic.
• A test that trips an MLX C++ throw aborts the whole test binary with SIGABRT rather than failing cleanly. That is expected, and the abort is stronger evidence than a clean assertion failure. One SIGABRT on the first run after a fresh link that does not reproduce on reruns is a cold Metal-initialization race, observed independently by two reviewers.

Validating model behavior: a token-exact reference oracle

Shape tests do not catch the failures that matter here. A wrong prefill/decode offset, a flat instead of interleaved QKV split, the wrong RoPE convention, or a routing bias applied to the wrong copy all produce correctly-shaped tensors and fluent, plausible output. Only the token id sequence separates them.

What worked: create a scratch venv, pip install mlx-lm, and drive mlx_lm.generate.stream_generate with make_sampler(temp=0.0), printing both text and token ids. Then compare against mlxcel generate. Pass --no-chat-template to mlxcel even when the checkpoint ships a chat template, so both sides see the same raw prompt, and run the templated path separately as a usability check.

Caveat: the oracle is blind wherever mlxcel and mlx-lm make the same choice. GELU was one such case, closed by monkey-patching the reference to the other variant and regenerating. Note mlxcel_core::utils::gelu_approx is erf-based despite its name, while MLX Python's nn.gelu_approx really is the tanh form.

A recurring defect class worth checking for

A value from config.json passes every Rust-side check, violates an undocumented precondition of an MLX C++ entry point, and kills the process at the FIRST FORWARD PASS rather than failing at load. Most cxx bridge functions are declared returning UniquePtr<MlxArray> rather than Result, so the C++ throw is an uncatchable std::terminate, not something catch_unwind contains. The model loads cleanly and the server dies on its first request. This class produced one CRITICAL and five HIGH findings across the five ports.

Specifics established by reading the pinned MLX checkout rather than assuming:

• Gathers do NOT range-check positive indices. take wraps negative indices but an out-of-range positive index silently returns values belonging to no row, and that reaches the logits with no fault. Bound every gather by the real tensor shape, never by a config field.
• slice CLAMPS an out-of-range stop instead of throwing, so a too-wide split silently loses trailing channels.
• fast_rope requires dims even, positive, and no larger than the last axis.
• rms_norm and layer_norm never inspect eps. A NaN or negative eps yields NaN hidden states with no error at all, which is harder to diagnose than a crash.
• matmul throws on an inner-dimension mismatch, which is why an unchecked INPUT axis is fatal even where the row axis bounds only an argmax.
• quantized_matmul divides by bits, so "bits": 0 is a division by zero. On AArch64 the divide returns 0 and std::invalid_argument fires; on x86-64 the hardware raises SIGFPE and kills the process before any exception exists.
• Reconstruct a quantized input width the way MLX does, as scales.shape(-1) * group_size, and check .biases shape equals .scales shape. A self-consistency check cannot catch a checkpoint honestly packed for the wrong hidden_size, because such a checkpoint IS self-consistent.
• Put zero checks BEFORE divisibility checks, because 0.is_multiple_of(0) is true.
• Avoid unbounded probe loops over 0..n_layer in the load path: a huge n_layer hangs with a flat allocation footprint, so no OOM kill rescues it.

src/models/gpt_neox.rs and src/models/helium.rs carry the current ModelArgs::validate / validate_weights shape to copy from.

Conventions that bite

• The // Used by: comments above shared helpers are the designated discovery mechanism for "what breaks if I change this", and they can themselves be stale. One was found listing six callers where there were eight. Verify by grep rather than trusting the comment, and update it when you touch the helper.
• mlxcel arch is the architecture registry. mlxcel list lists downloaded models and will not tell you whether a family is supported.