Repository navigation
fix(paged): paged_v2 chunked launch fails only in the full suite, pointing at cross-test state #1054
Description
Problem / Background
paged_v2::launch::launch_tests::chunk_size_does_not_change_the_answer fails in the full mlxcel-core lib suite, but passes when run on its own.
paged_v2::launch::launch_tests::chunk_size_does_not_change_the_answer
assertion failed: pages_per_chunk 1 relative error 55896.2
at src/lib/mlxcel-core/src/paged_v2/launch_tests.rs:324
The tolerance at that assertion site is 2e-3. An error of roughly 55896x is far outside any precision explanation, and is consistent with reading stale or wrongly shaped state rather than with an arithmetic inaccuracy. Raising the tolerance is not the fix and must not be proposed as one.
This is a pre-existing failure, not a regression from any recent change.
Environment
- GB10 (DGX Spark), sm_121
mainat6276f556, release profile,--features cuda- MLX pin
2c46b953db88965c4270cc7306eda6887a3247f2
Why this was invisible until now
It was surfaced by:
cargo test --release --features cuda -p mlxcel-core --lib -- --test-threads=1
which reports 1406 passed and 4 failed in 88 seconds. Without --test-threads=1 the suite aborts with SIGABRT partway through, because libtest drives MLX concurrently from many host threads. That is #1048. Because the suite never reached a verdict on this platform, this failure reached main unnoticed.
The other failures in the same run: #1049 covers the two fused RoPE parity failures, and #1053 covers paged_v2::sparse::sparse_tests::an_f16_cache_matches_the_reference_within_its_own_precision.
Key evidence
It passes in isolation, and that is the whole point of this issue.
- Run entirely alone, it passes in 1.73 seconds.
- It also passes inside a scoped 92-test
paged_v2run. - It fails only in the full
--libsuite.
So the chunking path itself is very likely correct, and something earlier in the suite leaves state behind that breaks it. The suspect is shared mutable process state, not the chunk-size logic.
This is the material difference from #1053, which reproduces deterministically in isolation and is therefore an ordinary wrong-result bug. The two need materially different investigations and must not be merged into one.
Note that pages_per_chunk 1 is the first iteration of the loop over [1, 2, 3, 4, 7, 16, 64], so the assertion aborts before the larger chunk sizes are evaluated. Whether they would also fail is unknown and worth capturing during the investigation.
Proposed Solution
Identify which preceding test induces the failure, then fix the leaked state, either by making the state properly scoped or by making the affected code independent of it.
Investigation approach:
- Bisect by module. Run the full
--libsuite with--test-threads=1and progressively restrict the test filter to narrow which preceding module or test must be present for the failure to appear. Serialized execution makes the ordering reproducible, so this bisection is deterministic. - Candidate process-global state this repo already has:
- The autotune memo,
src/lib/mlxcel-core/src/autotune/mod.rs. A memoised tactic selected under a different shape or dtype by an earlier test would be reused here and could produce a wrongly shaped launch. - Cached kernel or graph state, including anything captured or compiled once and keyed loosely enough to collide across tests.
- Environment variables set by a preceding test.
src/lib/mlxcel-core/src/test_support/env_lock.rsexists precisely because of this hazard, so a test that mutates the environment without taking the lock, or that fails to restore it, is a live candidate.
- The autotune memo,
- Once the inducing test is known, confirm the mechanism by reproducing the failure from exactly two tests run in sequence. That converts a suite-ordering anecdote into a minimal reproducer.
Acceptance Criteria
- The inducing predecessor is identified: a minimal ordered pair or short sequence of tests that reproduces the failure.
- The leaked state is named specifically, for example the memo entry, cache key, or environment variable involved.
- Root cause fixed at the source of the leak, so that the state is correctly scoped, correctly keyed, or correctly restored. A test-only workaround such as reordering or ignoring the test is not acceptable.
- Assessed and recorded whether the leak affects production behavior as well as the test suite. If it does, the production path is fixed too, not only the test path.
- The assertion tolerance at
src/lib/mlxcel-core/src/paged_v2/launch_tests.rs:324is unchanged at2e-3, or is tightened, never loosened. -
cargo test --release --features cuda -p mlxcel-core --lib -- --test-threads=1passes this test in the full suite. - Regression coverage prevents the same class of leak from recurring silently, for example by asserting the state is reset or by adding the missing lock usage.
Technical Considerations
If the leaked state turns out to be the autotune memo or a kernel or graph cache, the same leak can plausibly affect other tests that happen to run after the same predecessor and simply have tolerances loose enough to hide it. Worth scanning for other suite-only discrepancies once the mechanism is understood.
Related: #1048 (the reason this is only now visible), #1049, #1053.
Activity
fix(test): make the CUDA merge gate runnable and serialized (#1052)
fix(paged): key the paged decode v2 JIT cache on input dtypes (#1058)
inureyes commentedon Aug 7, 2026
The leaked state, named
It is MLX's CUDA JIT module cache, and the key is too weak.
mlx/backend/cuda/custom_kernel.cppnames a kernel"custom_kernel_" + name + template_arguments_hash(template_args), whilebuild_kernelin the same file generates every buffer parameter's type from the runtime dtype of the corresponding input (dtype_to_cuda_type(arr.dtype())).cu::get_jit_module(mlx/backend/cuda/jit_module.cpp) then memoises the compiled module under exactly that name in a process-globalstd::unordered_mapand invokes the source builder only on a cache miss.So a launch whose
template_argsare all ints hashes to one name for every input dtype. The first dtype to compile wins for the life of the process, and every later call at a different dtype reads its buffers through the wrong pointer type. Nothing throws.Metal is unaffected:
mlx/backend/common/metal_kernel.cppappends oneget_type_string(arr.dtype())per input to the kernel name, with the comment "The generated source depends on the dtypes of the inputs and outputs ... Include them in the kernel name so that a given name always maps to the same source." That asymmetry is why this is CUDA-only.This matches the candidate this issue nominated third ("cached kernel or graph state, including anything captured or compiled once and keyed loosely enough to collide across tests"), and it is the same mechanism as #1053. The two symptoms differ only in which predecessor supplies the colliding dtype, which is why they were correctly filed separately.
Does it affect production, or only tests
Production, and the clearest case is the sampler.
gumbel_max_sample_accepts(turbo/sampling.cpp) explicitly admitsfloat32,float16andbfloat16, while its template args were{TgSize, NumSplits}andNumSplitsdepends only on(batch, vocab). Two models with the same vocabulary but different logits dtypes therefore shared one compiled kernel on CUDA, and the second one sampled from a buffer read through the wrong pointer type.rejection_sample_accepts(turbo/sampling_rejection.cpp) is wider still: it places no dtype restriction at all.What landed
#1053 (PR #1058) fixed the two paged decode v2 kernels. This issue covers the rest of the class:
turbo/paged_attention.cpp(the v1 kernel): addsQType,KVType,VType.turbo/sampling.cpp: addsLogitsType.turbo/sampling_rejection.cpp: addsProbsType,ParamsType.Audited and already correct, needing no change:
turbo/fused_norm.cpp({"T", x.dtype()},{"TW", weight.dtype()}),turbo/fused_rope_append.cpp({"T", qkv.dtype()}), and the fused decode-MoE kernels incpp/mlx_cxx_kernels.cpp({"T", T}). Out of scope because they are Metal-only launchers with nocuda_kernel(call, and Metal's key already carries the dtypes:turbo/sparse_v_sdpa.cppand the three launches inturbo/turbo4_delegated_sdpa.cpp.Preventing recurrence
scripts/ci/check_kernel_dtype_keys.py, wired intomake verifyand a newkernel dtype keysCI job. Everytemplate_argsinitialiser in a file that also contains acuda_kernel(call must name at least one input dtype, inline or through a local bound from.dtype().There is deliberately no allowlist. Metal-only launchers are excluded by the absence of
cuda_kernel(in the file rather than by a hand-maintained list, so adding a CUDA port to one brings it under the check automatically. That is the inverted-rule shapescripts/ci/check_crate_versions.pyalready uses, chosen for the same reason: a list would go stale exactly when it mattered.The check was verified against a negative control (removing
{"LogitsType", logits.dtype()}makes it fail with that file and line), because a guard that cannot fail is not a guard.What is NOT established here, and needs a CUDA session
This issue's first acceptance criterion asks for a minimal ordered pair of tests that reproduces the failure. That is not delivered, and I will not guess at it.
nvccon this host and no reachable CUDA node (re-verified 2026-08-07:rexy.office.lablupandindominus.office.lablupboth fail to resolve their bastion).mlxcel-core, the v2 partial kernel is reached only frompaged_v2tests, and their template-arg tuples do not overlap across dtypes:launch_testsusesPageSize=32atDim64 and 128,sparse_testsreshapes the pool soPageSize=1, andcascade_launch_testsusesDim=16, PageSize=8. The merge kernel looked promising, sincepages_per_chunk 1is the first loop value and the only one that runs a merge at all, but the one caller outsidepaged_v2is MLAsplit_kv, whose test usesDIM=5and whose production path casts to f32 atsplit_kv.rs:278before the merge. So that hypothesis is refuted, not confirmed.mlxcel-corelib suite serialized reports 1415 passed, 0 failed in 92.23s,chunk_size_does_not_change_the_answerincluded.A GB10 session should run
cargo test --release --features cuda -p mlxcel-core --lib -- --test-threads=1on currentmainand report whether the failure is gone. #1053's fix plausibly removed it, since it re-keyed both the partial and the merge kernel, but "plausibly" is the honest word: without a run, whether that was this issue's specific predecessor is unknown.docs(reports): add technical reports for the #1053/#1054/#1045/#1040 …
fix(kernels): key every CUDA JIT launch on its input dtypes (lablup#1059
fix(sampling): filter on the untempered distribution and scale by tem…