fix: Paged block budget test underflows when the paged-v2 workspace reserve exceeds the byte budget #1091

Problem / Background

The test execution::memory_estimate::tests::resolve_block_budget_explicit_bytes_floors_to_block_count fails deterministically on a clean main (0f727572), in 0.00s, on GB10 / CUDA sm_121 / Linux aarch64:

thread '...' panicked at src/execution/memory_estimate.rs:1741:9:
assertion failed: shortfall < 100

This is not one of the TF32-sensitive numeric parity failures on this box. It fails identically with MLX_ENABLE_TF32=0 and reproduces in isolation with --exact.

The failing arithmetic is in the test itself, at src/execution/memory_estimate.rs:1740-1741:

let shortfall = usize::try_from((per_block * 100 - workspace) / per_block).unwrap();
assert!(shortfall < 100);

An earlier assertion in the same test (assert!(workspace > 0, "the minimal config has a derivable geometry")) passes, so workspace is positive. With unsigned arithmetic, (per_block * 100 - workspace) / per_block is at most 99 for any 0 < workspace <= per_block * 100. The assertion can only fail if workspace > per_block * 100, in which case the u64 subtraction wraps and shortfall becomes an enormous number that still converts cleanly through usize::try_from on a 64-bit host.

Root cause: the reserve is larger than the budget on every non-Metal host

The numbers are fully determined, not flaky.

per_block comes from paged_block_bytes(tmp.path(), 32, 32, false). With write_minimal_config (hidden_size 4096, num_attention_heads 32, num_key_value_heads 8, so head_dim 128) that is 2 (K+V) * 32 tokens * 8 kv heads * 128 head_dim * 2 bytes = 131072 per layer, matching the test's own comment. So the budget under test is per_block * 100 = 13107200 bytes (12.5 MiB).

workspace comes from paged_v2_workspace_reserve_bytes(tmp.path(), 32, 1) at src/execution/memory_estimate.rs:921:

target     = mlxcel_core::paged_v2::device_target_ctas()
chunks     = (target * 2).div_ceil(num_kv_heads) + batch
hq         = num_kv_heads * PAGED_V2_MAX_N_REP            (PAGED_V2_MAX_N_REP = 16, line 895)
per_launch = chunks * hq * (head_dim + 1) * size_of::<f32>()
result     = per_launch * PAGED_V2_CONCURRENT_LAUNCHES    (PAGED_V2_CONCURRENT_LAUNCHES = 2, line 885)

device_target_ctas() (src/lib/mlxcel-core/src/paged_v2/plan.rs:435) returns DEFAULT_TARGET_CTAS = 512 on any host where metal_is_available() is false, which is every Linux and CUDA host. Substituting:

chunks     = (1024).div_ceil(8) + 1 = 129
hq         = 8 * 16 = 128
per_launch = 129 * 128 * 129 * 4 = 8520192
workspace  = 8520192 * 2 = 17040384        (16.25 MiB)

17040384 > 13107200 by 3933184 bytes, so the subtraction wraps and shortfall becomes 2^47 - 31 = 140737488355297.

On Apple hosts the target is gpu_core_count * 8 (floored at 64), so the test passes on small parts and fails on large ones. The crossover is at a reported core count of 50: a 48-core part yields a 12813312 byte reserve (passes), a 50-core part yields 13341696 (fails). That, combined with CI not running these tests, is why the failure was never surfaced by the macOS runners.

What is not broken

The non-test implementation of resolve_paged_block_budget (src/execution/memory_estimate.rs:846-878) already guards this correctly at line 876:

let workspace = paged_v2_workspace_reserve_bytes(model_dir, num_layers, batch);
let budget_bytes = budget_bytes.saturating_sub(workspace);
Some(usize::try_from(budget_bytes / per_block).unwrap_or(usize::MAX))

A user-supplied byte budget smaller than the workspace reserve therefore resolves to Some(0), which is the documented "budget rounds below one block" convention, not a wrapped block count. The defect is confined to the test's expectation, which hardcodes an assumption (the reserve is smaller than 100 blocks of KV) that only holds on devices with a small CTA target. Recording this explicitly so the fix does not go chasing a production underflow that is not there.

Secondary consideration worth a look while fixing

The reserve is 16.25 MiB for this geometry on a 512-CTA target, and PAGED_V2_MAX_N_REP = 16 deliberately over-reserves (Llama 3 is 4, so this is roughly a 4x over-reservation for the common case). Combined with the saturating_sub, an explicit --paged-kv-budget in the low-megabyte range now silently resolves to zero blocks on a non-Metal host. That is within the documented contract (the caller decides what to do with Some(0) and must not install a wedging zero budget), but it is a behavior worth confirming is actually reachable and handled by callers, rather than assumed.

Proposed Solution

Make the test's expectation device-independent, and assert the real contract instead of the hardcoded 100-block assumption. Options, in rough order of preference:

• Compute the expected shortfall with saturating_sub (mirroring the implementation) and drop the shortfall < 100 assertion, or replace it with an assertion on the actual contract: the resolved count for a budget that ignored the workspace equals budget.saturating_sub(workspace) / per_block.
• Alternatively, scale the test budget off the measured workspace (for example workspace + per_block * 100) so the "did not account for the workspace" case is expressed relative to the reserve rather than relative to a fixed 12.5 MiB.
• Pin MLXCEL_PAGED_DECODE_V2_TARGET_CTAS for the duration of the test only if the above is not workable. This is the weakest option because device_target_ctas() memoizes through a OnceLock, so an env override set inside a test is order-dependent and will not reliably take effect under a shared test binary.

Add an explicit regression case for the budget-smaller-than-reserve path so the saturating_sub at line 876 is covered directly rather than incidentally.

Acceptance Criteria

• The fix states explicitly where the underflow lives (test expectation, implementation, or both). The evidence above says test-only; confirm or refute it rather than assuming.
• Every place a device-derived workspace reserve is subtracted from a requested budget uses saturating or checked arithmetic, including test helpers and expectation computations, so a budget smaller than the reserve yields 0 blocks rather than a wrapped value.
• resolve_block_budget_explicit_bytes_floors_to_block_count no longer depends on device_target_ctas() being small. It asserts the contract that a budget which did not account for the workspace comes back short by exactly the reserve.
• The test passes on GB10 (device_target_ctas() == 512) and on a device with a small CTA target (an Apple part below 50 reported cores, target 80 to 384).
• A regression test pins the budget-smaller-than-reserve case explicitly and asserts Some(0).
• cargo test --features cuda -- --test-threads=1 execution::memory_estimate is green on GB10.

Technical Considerations

Found while running make verify-test-cuda as the merge gate during unrelated work (the #847/#848/#849 VLM port chain), then control-running main to confirm it was pre-existing. It is one of 5 failures on main at 0f727572. The other 4 are TF32 artifacts and are filed separately. This one is a genuine defect.

Relevant files and symbols:

• src/execution/memory_estimate.rs:846 resolve_paged_block_budget (the saturating_sub at line 876 is correct)
• src/execution/memory_estimate.rs:921 paged_v2_workspace_reserve_bytes
• src/execution/memory_estimate.rs:885 PAGED_V2_CONCURRENT_LAUNCHES = 2, :895 PAGED_V2_MAX_N_REP = 16
• src/execution/memory_estimate.rs:1689-1753 the failing test
• src/lib/mlxcel-core/src/paged_v2/plan.rs:435 device_target_ctas, :76 DEFAULT_TARGET_CTAS = 512, :67 CTAS_PER_UNIT = 8, :80 MIN_TARGET_CTAS = 64

The workspace reserve was introduced by #899 (charging the fused decode v2 workspace to the KV budget before converting the remainder into blocks). The test was written against a geometry where the reserve happened to be small.