chore: post-epic cleanup of compat paged kernels, the model-owned paged flag, and the CUDA SSM bf16 parity failures #2230

Part of #2166.

Three cleanups left after the epic, all in mlxcel-core and the model-owned KV families. Implement all three in one PR, run the local gate once at the end (make verify-test-cuda on GB10 covers every item), and keep pushes to a minimum. Each item has its own acceptance criteria.

(a) Remove the reference-only dense-compat paged kernels

Current Behavior

Since #2210 and #2217 no model calls the dense-pointer paged routes; the doc comment at src/lib/mlxcel-core/src/layers.rs:5766-5769 says they "remain as the reference pair their FFI tests pin". What is left:

• C++: paged_decode_attention_dense_compat and paged_decode_attention_rotating_compat (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4112, :4223; declarations src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h:1130, :1142) and their cxx bridge entries (src/lib/mlxcel-core/src/lib.rs:1396, :1411).
• Rust, src/lib/mlxcel-core/src/layers.rs: paged_decode_attention_dense_fallback (:5421), paged_decode_attention_dense_compat (:6071), paged_decode_attention_rotating_fallback (:6138), paged_decode_attention_rotating_compat (:6285), and their validators validate_paged_decode_inputs (:5352) and validate_rotating_paged_decode_inputs (:6093).
• src/lib/mlxcel-core/src/cache.rs: PagedDecodeMetadata (:6343, with from_attention_metadata and from_visible_lengths at :6415-6423), RotatingPagedDecodeMetadata (:6357), and the test paged_decode_metadata_builds_logical_block_tables (:10009). Nothing else uses either type.
• src/lib/mlxcel-core/src/ffi_tests.rs: test_paged_decode_attention_dense_compat_matches_fallback (:347), test_paged_decode_attention_rotating_compat_matches_fallback_after_wrap (:401), and the dense and rotating halves of paged_decode_fallbacks_reject_an_empty_batch (:4264).

Four pooled-kernel tests use paged_decode_attention_dense_fallback as their oracle and must keep their coverage: test_pooled_paged_decode_matches_dense_over_200_steps (ffi_tests.rs:515), test_pooled_paged_decode_sliding_window_via_logical_start (:1071), test_pooled_paged_decode_batch_of_two (:1175), test_pooled_prefill_then_decode_matches_dense (:1282).

Proposed Solution

Delete everything in the list above. Replace the oracle in the four pooled tests with a test-only helper in ffi_tests.rs, dense_reference_decode_attention(q, keys, values, scale) -> UniquePtr<MlxArray>, computing softmax(q k^T * scale) v in float32 over the contiguous [B, H, T, D] reference K/V the tests already build (GQA by repeating KV heads, as the fallback did), with the tests' existing tolerances. Keep the pooled half of the empty-batch test, renamed paged_decode_pooled_fallback_rejects_an_empty_batch. Reword the doc comments that name the removed symbols (layers.rs:5766-5769, src/lib/mlxcel-core/src/cache/attend.rs:55) to describe the current state without the symbol names. ADR bodies (0001, 0008, 0009) and the past-tense mentions in docs/CONTINUOUS_BATCHING.md stay unchanged.

Acceptance Criteria

• rg -n 'paged_decode_attention_(dense|rotating)_(compat|fallback)|PagedDecodeMetadata|validate_(rotating_)?paged_decode_inputs' src/ prints nothing.
• The four pooled tests pass with the new oracle, and fail if the pooled kernel's output is perturbed (check once locally by scaling its output by 1.01 and seeing them fail; do not commit that).
• cargo clippy --workspace --all-targets --features cuda -- -D warnings is clean (no dead-code warnings from the removal), and make verify-kernel-dtype-keys verify-kernel-port-dispatch still pass.

(b) Decide the paged flag of Gemma 3 and Llama 4 by measurement

Current Behavior

Gemma3Wrapper::supports_paged_decode_backend (src/models/gemma3.rs:1620) and Llama4Wrapper::supports_paged_decode_backend (src/models/llama4.rs:1802) return true, so effective_decode_storage_backend (src/server/batch/scheduler/mod.rs:207) resolves auto to paged whenever --parallel is above 1. Both families keep model-owned KV and, since #2217, no attention path reads the pool for them, but they still pay the pool allocation and the per-step sync_sequence_storage mirror (Engine::step, src/lib/mlxcel-core/src/engine/mod.rs:246). The PR #2226 measurement (docs/benchmark_results/unified-engine-final-gb10-2026-10-08.md, "Gemma 3 storage under concurrency"; Gemma-3-1B 4-bit, --parallel 8, 512-token prompts, 128 generated, 3 rounds) had dense 1 to 2 percent higher at concurrency 1, 4 and 8 with overlapping round ranges (for example 342..360 vs 341..359 at concurrency 4), so it is unresolved. Llama 4 was not measured.

Proposed Solution

Measure each family with mlxcel serve --parallel 8 --decode-storage-backend paged|dense and scripts/bench_serving_concurrency.py --concurrency 1,4,8 --max-tokens 128 --metrics, at --prompt-tokens 512 and --prompt-tokens 4096, 5 rounds. Each round runs paged, dense, and paged again as the null arm, in rotating order, each on a fresh server. Models: models/mlx/gemma-3-1b-it-4bit and models/mlx/llama-4-scout-17b-16e-instruct-4bit. Decision rule per family (maintainer: the faster option wins): if dense's median paired delta against paged clears that cell's null spread upward in at least one cell and paged is not faster beyond the null spread in any cell, return false from that family's supports_paged_decode_backend; otherwise keep true and add a doc comment on the method citing the measurement.

Returning false also makes an explicit --decode-storage-backend paged resolve to dense for that family (mod.rs:213-219), and the engine-parity paged row reports n/a for it; record both in the doc comment and in docs/CONTINUOUS_BATCHING.md. Update any test that pins the paged backend for a real Gemma 3 checkpoint (src/server/batch/scheduler_real_batch_parity_tests.rs runs gemma-3-1b-it-4bit "on both backends").

Acceptance Criteria

• Results committed under docs/benchmark_results/ (dated document plus data/ with the raw output, driver script, host, commit and MLX pin), with the per-cell median paired delta and null spread for both models and both prompt lengths.
• Each family's supports_paged_decode_backend matches the decision rule applied to its own numbers, with a doc comment naming the results document.
• If a flag changes: a unit test on effective_decode_storage_backend inputs for that family (supports_paged_decode_backend == false, max_batch_size == 8) asserts Auto and Paged resolve to Dense, and mlxcel serve --parallel 8 -m <that model> logs dense decode storage at startup.

(c) Fix the CUDA SSM bf16 parity failures

Problem / Background

On CUDA (GB10), ssm_update_parity_tests::ssm_update_kernel_matches_graph_step_bf16 and ssm_update_kernel_keys_on_a_log_dtype (src/lib/mlxcel-core/src/ssm_update_parity_tests.rs:340, :355) fail with the new state at normalized RMS about 1.8e-3 to 2.5e-3 against the f32 budget of 1e-5 (tolerance_for, :58). The failing check in both is the granite-4.0-h-tiny shape with activations and A_log in bf16 (seeds 2069 and 2071); the Nemotron-H bf16 check (A_log f32, seed 2070) and every f32 check pass, and the second check of keys_on_a_log_dtype (seed 2072, A_log f32) never runs because the first panics. They fail identically on the pre-epic build 4c44e31 and on main 30ceea2, so make verify-test-cuda is not green (12156 passed, 2 failed). The test and the kernel table were last changed by #2099 and #2179.

Current Behavior

The test's reference (ssm_update_parity_tests.rs:265-302) upcasts every input to float32, A_log included, before forming -exp(A_log), which is what the module doc (:20-25) says the kernel does. The CUDA kernel source (SSM_CUDA_SOURCE, src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp, :676-716) computes -expf((float)A_log[h_idx]) and accumulates the state in float32. dt reaches the kernel from compute_dt_compiled (mlx_cxx_kernels.cpp:858-876, an mlx::core::compiled astype(dt, f32) + dt_bias, log1p(exp(.)), clip), while the reference uses softplus on f32(dt) + f32(dt_bias). Granite's dt_limits of (0.001, 100.0) do not clip, while Nemotron-H's (1e-3, 0.1) clip almost every head, which may be why only granite shows the deviation.

Proposed Solution

Find which side is wrong before changing anything. Run the two tests alone and inside the full mlxcel-core lib suite to separate an in-process ordering effect (JIT module reuse across dtypes, the class ssm_update_kernel_keys_on_a_log_dtype exists for) from a deterministic numeric one. Then compare intermediates for the failing case on CUDA against the reference: dt' (the compute_dt_compiled output against the reference's clipped softplus), exp(-exp(A_log) * dt'), and the state. Fix the stage that deviates: a kernel or compute_dt defect is fixed in mlx_cxx_kernels.cpp (and in the Metal and HIP sources if they share it); a reference that does not compute what the kernel is specified to compute is fixed in the test, and the module doc is corrected to match. A tolerance change is acceptable only with a written derivation in the module doc of why bf16 inputs bound the f32 state error at that level on CUDA but not on Metal or ROCm.

Acceptance Criteria

• Both tests pass on CUDA, alone and in the full suite: cargo test --profile test-fast --features cuda -p mlxcel-core --lib ssm_update_parity_tests -- --test-threads=1.
• The PR states which stage deviated and why, with the measured intermediate deviations before and after.
• A real-checkpoint check: greedy mlxcel generate -m models/mlx/granite-4.0-h-tiny-4bit output is unchanged by the fix or, if the fix changes the kernel numerics, matches the graph path (MLXCEL_SSM_CUDA_KERNEL=0) for the first 64 tokens under MLXCEL_SDPA_DETERMINISTIC=1.
• make verify-test-cuda reports 0 failures. If any test from the fix: six lib tests fail under --features cuda on main #2087 set (closed) fails again, the PR names it with its tracking issue.

Verification

cargo fmt --all -- --check
cargo clippy --workspace --all-targets --features cuda -- -D warnings
make verify-kernel-dtype-keys verify-kernel-port-dispatch
gpu-lock run --tag test -- cargo test --profile test-fast --features cuda -p mlxcel-core --lib -- --test-threads=1 ffi_tests ssm_update_parity_tests
gpu-lock run --tag verify -- make verify-test-cuda

# (c) real checkpoint
cargo build --release --features cuda --bin mlxcel
gpu-lock run --tag granite -- env MLXCEL_SDPA_DETERMINISTIC=1 ./target/release/mlxcel generate -m models/mlx/granite-4.0-h-tiny-4bit -p "Explain the Mamba state update." -n 64
gpu-lock run --tag granite -- env MLXCEL_SDPA_DETERMINISTIC=1 MLXCEL_SSM_CUDA_KERNEL=0 ./target/release/mlxcel generate -m models/mlx/granite-4.0-h-tiny-4bit -p "Explain the Mamba state update." -n 64

# (b) measurement, once, at the end, on a quiet host: one server per arm, for each
# model, prompt length (512, 4096) and round (1..5), arms in rotating order
mlxcel serve -m models/mlx/gemma-3-1b-it-4bit --port 18080 --metrics --parallel 8 --decode-storage-backend paged
python3 scripts/bench_serving_concurrency.py --port 18080 --concurrency 1,4,8 --prompt-tokens 512 --max-tokens 128 --metrics

A pass is a clean gate (0 failures, clippy clean), both SSM tests green on CUDA, and the committed (b) measurement with each family's flag set by the rule above.