fix(sampling): temperature_one_support_unchanged pins fused_sample_probs bit patterns that differ by one ulp on Metal and CUDA #1418

Problem / Background

sampling::tests::temperature_one_support_unchanged (added by PR #1391, commit 6ecf756e, closing #1379) calls ffi::fused_sample_probs(&logits, 1.0, top_k, top_p, 0.0) and asserts the raw f32 bit patterns of the returned row against hard-coded golden u32 vectors for three (top_k, top_p) cases. It fails deterministically on the top_k=40, top_p=0.9 case with a one-ulp drift on a handful of non-zero components, on both GPU backends the project gates on, so the macOS nightly and the GB10 CUDA merge gate are both red on main for a test that is not detecting a real regression.

The panic is assertion left == right failed: T=1.0 fused_sample_probs bits moved for top_k=40 top_p=0.9 at src/lib/mlxcel-core/src/sampling.rs:2920:13. Every reported difference is the kernel returning one ulp above the golden. Mapped to the golden vector at src/lib/mlxcel-core/src/sampling.rs:2860-2872: index 26 golden 0x3d3f3109 (1027551497) vs returned 0x3d3f310a (1027551498), index 27 golden 0x3db9baff (1035582207) vs 0x3db9bb00 (1035582208), index 29 golden 0x3cc3616a (1019437418) vs 0x3cc3616b (1019437419), index 47 golden 0x3cc2352f (1019360559) vs 0x3cc23530 (1019360560), index 52 golden 0x3e334131 (1043546417) vs 0x3e334132 (1043546418). The zero pattern (the support) does not move; only the rounding of probabilities inside the support does.

Evidence that this is pre-existing and platform independent:

• Nightly verify (macOS Apple Silicon, metal,accelerate, test-fast profile, step cargo test (test-fast profile, metal,accelerate) which runs make verify-test) failed on exactly this test in run 32760920343 (main at fa4bd5c6, 2026-08-24) and run 32882245154 (main at 681d3656, 2026-08-25).
• On the GB10 CUDA runner the same test fails 3 of 3 isolated runs at main ffca5659 on 2026-08-26 with cargo test --profile test-fast --features cuda -p mlxcel-core --lib sampling::tests::temperature_one_support_unchanged -- --test-threads=1 (exit 101), and it turns make verify-test-cuda red as the only mlxcel-core failure in an otherwise 1563-pass run.
• src/lib/mlxcel-core/src/sampling.rs has not changed since 6ecf756e (git log -1 -- src/lib/mlxcel-core/src/sampling.rs).
• Discovered during the post-merge integration gate of epic feat(embeddings): native embedding and reranking subsystem (epic) #1348, which did not touch sampling.rs and is not the cause.

Current Behavior

The test lives at src/lib/mlxcel-core/src/sampling.rs:2846-2925. It builds a 64-token row with lcg_logits(0x1379_5EED, 64) (sampling.rs:2733), then iterates three cases (top_k, top_p, routed, expected_bits) at sampling.rs:2851-2910: (40, 0.9, true), (0, 0.9, true), (40, 1.0, false). Cases with routed == true are skipped when ffi::sampling_rejection_available() is false (sampling.rs:2912-2914). For each case it converts the fused_sample_probs output to Vec<u32> via ffi::array_to_raw_bytes and does assert_eq!(got, expected, ...) at sampling.rs:2920-2923. Because (40, 0.9) is the first case and assert_eq! panics, the outcomes of (0, 0.9) and (40, 1.0) on the failing platforms are unknown from the CI logs; the fix must cover all three cases, not just the one that currently reports.

The doc comment at sampling.rs:2825-2844 states the goldens were captured from commit e989e45da on the Metal backend and must stay bit-identical because T == 1.0 applies no temperature division on either side of #1379. That bit-identity claim is what fails: the row is produced by fused_sample_probs_impl in src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5978-6025, which runs the filter chain fused_sample_filter_logits (mlx_cxx_bridge.cpp:5260, itself using mlx::core::softmax for the top-p and min-p masks, e.g. mlx_cxx_bridge.cpp:5342) and then returns mlx::core::softmax(astype(x, float32), -1, /*precise=*/true) at mlx_cxx_bridge.cpp:6023-6024. Those are MLX reductions whose reduction order and exp/normalise strategy depend on the backend (Metal vs CUDA), on MLX kernel selection for the row length, and on the MLX version pinned at the time of capture. None of that is under this repository's control, so the low bit of each probability is not a stable contract.

Proposed Solution

Hypothesis, to be confirmed by the implementer rather than treated as a finding: the goldens were captured on one machine at one MLX revision, and the fused kernel's normalisation (sum then divide, or exp then rescale inside softmax(..., precise=true)) is not bit-reproducible across GPU backends or across MLX kernel selections. A one-ulp difference in a probability is within expected floating-point tolerance, so asserting exact bits over-constrains the test. The property #1391 wanted to protect is that T = 1.0 leaves the support and the filtered distribution unchanged, meaning the same zero pattern and the same non-zero entries up to rounding. That property is assertable without stored bits.

Rewrite temperature_one_support_unchanged to assert the property in two exact-plus-tolerance parts, using the same in-process reference pattern that fused_sample_probs_equals_softmax_of_filtered_over_t (sampling.rs:2747-2800) already uses:

1. Support, asserted exactly. For each case compute the reference support from the Rust reference filters on the untempered row, in chain order: top_k_filter (sampling.rs:1310) when top_k > 0, then top_p_filter (sampling.rs:1356) when 0 < top_p < 1, and take kept[i] = probs_row(&filtered)[i].is_finite() (probs_row at sampling.rs:2680). Assert got[i] == 0.0 for every i with !kept[i] and got[i] > 0.0 for every i with kept[i], so the zero pattern is still pinned bit-exactly (0.0f32 has one bit pattern). For the two routed cases keep the existing rationale at sampling.rs:2852-2855: on this row the 0.9 nucleus is 15 tokens and a strict subset of the top-40 set, so the rejection kernel's intersection semantics reduce to the plain nucleus and the chain-order reference support is the correct expectation. State that in the comment and add a debug_assert/assert that the top-p support is a subset of the top-k support for the (40, 0.9) case so the assumption is checked rather than trusted.
2. Values, asserted within tolerance. Compute the f64 reference softmax(row[kept]) with no temperature division (the T = 1.0 contract), exactly as sampling.rs:2775-2789 does with t = 1.0, and assert for every kept i that (f64::from(got[i]) - want[i]).abs() <= 1e-6 * want[i].max(1e-3). This is a relative tolerance of about 1e-6 (roughly 8 f32 ulps at the largest entries, far tighter than any temperature-related regression, which shifts entries at the percent level) with an absolute floor for tiny entries. If the implementer measures a larger platform spread on either backend, widen the tolerance and record the measured spread in the comment; do not reintroduce exact-bit assertions.

Remove the three golden u32 vectors at sampling.rs:2856-2909 and replace the doc comment at sampling.rs:2825-2844 with one that says: the goldens were captured on Metal at e989e45da, they drifted by one ulp on both Metal and CUDA by 2026-08 (link this issue), exact bits are not asserted because the MLX softmax reduction is not bit-reproducible across backends or MLX versions, and the test now asserts support equality exactly plus values within tolerance against an in-process f64 reference. Keep the paragraph explaining why min-p configurations are absent and why (40, 1.0) is unguarded, since both still apply. Keep the routed guard and the sampling_rejection_available() skip unchanged.

Alternative considered and rejected: regenerating the goldens on one backend and keeping assert_eq! on bits. That reproduces the same failure the next time MLX changes a kernel or the other backend rounds differently; it was already captured on Metal once and Metal itself now disagrees with it.

Scope

In scope: src/lib/mlxcel-core/src/sampling.rs, only the temperature_one_support_unchanged test and its doc comment. No change to fused_sample_probs, fused_sample_filter_logits, or any production code.

Out of scope: making mlx::core::softmax bit-reproducible across backends; any change to the rejection kernel intersection semantics; the other sampling tests, which already use tolerance comparisons.

Implementation Notes

• Reuse: the reference-support and f64-softmax construction in fused_sample_probs_equals_softmax_of_filtered_over_t (sampling.rs:2762-2789); the helpers lcg_logits, probs_row, top_k_filter, top_p_filter already in scope for the test module. Consider lifting the support-plus-f64-softmax reference into a small private helper fn reference_filtered_softmax(row: &[f32], t: f32, top_k: i32, top_p: f32, min_p: f32) -> (Vec<bool>, Vec<f64>) in the test module and using it from both tests, so the two do not drift.
• Constraints: the test must remain meaningful on a CPU-only build (the (40, 1.0) case runs everywhere and must keep doing so) and must not depend on which backend is compiled in. Do not add #[ignore] and do not gate the whole test on sampling_rejection_available(); the doc comment at sampling.rs:2839-2844 explains why that was already rejected once.
• Edge cases: a kept entry that is exactly 0.0 after the fused softmax (underflow) would be misclassified as outside the support; the fixed 64-token row spans roughly [-6, 6] so the minimum kept probability is well above f32 subnormal range, but assert got[i] > 0.0 on kept entries explicitly so a future row change fails loudly rather than silently passing. The reference support for (0, 0.9) and (40, 0.9) must be identical on this row (the existing comment says so); assert that too, since it is the premise for reusing the chain-order reference on a routed case.
• Error handling: on failure the assertion message must name the case (top_k, top_p), the token index, the returned value as both f32 and hex bits, and the reference value, so the next platform drift is diagnosable from the CI log without a local reproduction.

Acceptance Criteria

• sampling::tests::temperature_one_support_unchanged passes on the macOS Nightly verify workflow (make verify-test) and on the GB10 CUDA gate (make verify-test-cuda) at the merge commit, with no #[ignore] and no per-backend cfg on the test.
• The support (zero pattern) is still asserted exactly for all three cases: every non-kept token is 0.0 and every kept token is > 0.0.
• Non-zero components are asserted against an in-process f64 reference softmax(row[kept]) with the stated relative tolerance (or a wider, measured one recorded in the comment); no stored golden bit patterns remain in the test.
• The doc comment states which backend produced the original goldens, that they drifted on both Metal and CUDA, and why exact bits are no longer asserted, and references this issue.
• fused_sample_probs_equals_softmax_of_filtered_over_t still passes unchanged in behaviour (shared helper refactor, if done, must not alter its tolerance or cases).
• make verify-test-cuda reports zero failures in mlxcel-core on the merge commit, and the next Nightly verify run after merge is green on the cargo test step.

Verification

# GB10 CUDA runner: the currently failing command, 3 consecutive runs must pass
for i in 1 2 3; do cargo test --profile test-fast --features cuda -p mlxcel-core --lib sampling::tests::temperature_one_support_unchanged -- --test-threads=1 || exit 1; done

# CUDA: the sibling reference test and the whole sampling module
cargo test --profile test-fast --features cuda -p mlxcel-core --lib sampling::tests:: -- --test-threads=1

# macOS Apple Silicon: same via the Makefile wrapper (adds metal,accelerate)
make test-fast FILTER="-p mlxcel-core --lib sampling::tests::temperature_one_support_unchanged"

# Merge gates
make verify-test-cuda   # Linux/NVIDIA
make verify-test        # macOS, what Nightly verify runs

# Lint
cargo clippy --profile test-fast -p mlxcel-core --all-targets -- -D warnings

A pass is: the isolated test exits 0 on both platforms, make verify-test-cuda shows no failing test in mlxcel-core, and the assertion still fails (with the case, index, hex bits and reference in the message) if the tolerance is temporarily set to 0.0 on CUDA, which demonstrates the test is exercising the values and not just the support.

Technical Considerations

Related: #1379 (the temperature reorder), #1391 (the PR that added this test). Discovered during the post-merge gate of epic #1348, which is unrelated to the cause. The failure also blocks any PR that relies on make verify-test-cuda being green, so until this lands, reviewers must read the mlxcel-core failure list rather than the exit code.