test(cuda): run the deferred CUDA arm for the merged ROCm routing work on a GB10 node

[inureyes]

Part of #1801

Problem / Background

The changes listed below are already merged to main. This issue is a verification pass, not a merge gate: it can find a regression that already shipped, and a failure here means opening a fix issue against the specific change, not reverting the set blindly.

The CUDA arm was deferred at merge time because no CUDA host was reachable from the ROCm development machine or from the Metal host, and the self-hosted CUDA CI runner has been down (its jobs sat queued and expired at the 24 hour timeout). Rather than leaving an implied "someone should check CUDA" note in each PR, the whole deferred set is collected here to be run in one pass when a GB10 node is available.

What needs the CUDA arm

Change	Why CUDA is at risk	What would go wrong
refactor(core): route custom kernels by GPU backend kind (#1803, PR #1869)	Ten call sites that chose a fused kernel with !mlx::core::metal::is_available() now ask a GpuKernelBackend enum, and four availability predicates now call custom_kernels_available(). On a CUDA build the enum resolves to Cuda and the predicates to true, so the selection is equivalent by construction.	If it is not, CUDA silently takes a Metal kernel or a graph fallback. A fallback keeps the output identical and only changes speed, so an output comparison alone cannot see it.
The same PR's Result change	paged_attention_decode, paged_attention_decode_v2_partial and paged_attention_merge_states now return Result through cxx instead of a bare value, and their launchers refuse when the backend has no port.	The refusal branch is unreachable on CUDA, but the Result wrapping is new code on a hot path.
paged_decode_backend() in src/lib/mlxcel-core/src/layers.rs	It now asks custom_kernels_available() before Metal and CUDA.	If that predicate ever answered false on a CUDA host, every fused paged decode would fall back to gather with identical output and lower throughput.
feat(rocm): port the BitNet bitlinear_matmul kernel (#1862, PR #1870)	The kernel selection gained a third arm and a new bitlinear_kernel_available() predicate.	CUDA must still pick the CUDA kernel, and the BitNet tests must run rather than skip.
feat(rocm): implement FFT through hipFFT (#1825, PR #1861)	Only src/lib/mlx-cpp/patches-rocm/ plus a hipfft entry in the link list and docs. A CUDA build copies no file from that directory.	Low risk; included so the run covers the whole merged set.

Procedure

This is written for someone with no context from these PRs, on a GB10 node with a working CUDA toolchain and a clone of this repository.

1. Ask the binary which backend it resolved

The merged work ships its own probe, so the central risk of this whole issue does not have to be inferred from counters. gpu_kernel_backend() in src/lib/mlx-cpp/turbo/gpu_backend.cpp resolves the backend once behind a static, and when MLXCEL_DEBUG_KERNEL_BACKEND is set in the environment it prints exactly one line to stderr:

[mlxcel] custom kernel backend: cuda

When the resolved backend has no custom kernel ports it appends (no custom kernel ports; using MLX graph fallbacks). Run the arm binary once with the variable set and read that line:

MLXCEL_DEBUG_KERNEL_BACKEND=1 ./target/release/mlxcel generate --model models/mlx/Meta-Llama-3.1-8B-Instruct-4bit --prompt hi -n 8 2>&1 | grep 'custom kernel backend'

Any run that touches a kernel-availability predicate triggers it, and a short generate does (fused_norm.cpp and sampling.cpp both call custom_kernels_available()), so no special model or flag is needed. Pass condition: the line reads cuda with no suffix. That single line settles the risk the rest of this issue is built around, which is that a new predicate answers wrongly and every fused path quietly falls back to an MLX graph with identical output and only lower throughput. It costs one run, so do it first: if it fails, nothing below is worth running, and the failure itself is the bug report.

2. Build both arms

The baseline is 56d82fbc (fix: refuse byte-identical arms in the A/B output gate (#1867)), the last commit on main before this ROCm series began landing; 4641dc7a (feat(rocm): implement FFT through hipFFT, PR #1861) is the first of the series on main. The arm is main at or after that point. Name the two binaries target/release/mlxcel.before (baseline) and target/release/mlxcel (arm).

git checkout 56d82fbc
cargo build --release --features cuda --bin mlxcel
/bin/cp target/release/mlxcel target/release/mlxcel.before
git checkout main
cargo build --release --features cuda --bin mlxcel 2>&1 | tee /tmp/build-arm.log

The check that matters is in the second build's log: src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp and the files under src/lib/mlx-cpp/turbo/ (paged_attention.cpp, paged_attention_v2.cpp, paged_attention_v2_merge.cpp) must actually recompile. A cargo build that no-ops after a bridge .cpp edit has been observed on this project, and when it happens the arm binary is secretly identical to the baseline, at which point every check below passes for the wrong reason. Grep the log for those filenames and keep the matching lines as evidence.

3. The test gate

This is the strongest of the automated checks, because the tests that exercise these launchers call them directly rather than through a model, so no token floor or default-off gate can quietly skip them.

git checkout 56d82fbc && make verify-test-cuda 2>&1 | tee /tmp/test-before.log
git checkout main && make verify-test-cuda 2>&1 | tee /tmp/test-after.log
grep '^test result:' /tmp/test-before.log > /tmp/result-before.txt
grep '^test result:' /tmp/test-after.log > /tmp/result-after.txt
diff /tmp/result-before.txt /tmp/result-after.txt

Pass condition: the same tests pass in both runs, and nothing that passed before fails after. The groups whose results matter most are paged_v2::*, mla::split_kv, cache::paged_batch_decode and models::bitnet::*. A BitNet group that goes from passing to skipped is a failure, not a pass.

4. Greedy output equality

./scripts/ab_output_equality.sh \
  --baseline target/release/mlxcel.before \
  --arm target/release/mlxcel \
  --model models/mlx/Qwen3-30B-A3B-4bit \
  --model models/mlx/Meta-Llama-3.1-8B-Instruct-4bit \
  --prompt "Explain how paged attention differs from a contiguous KV cache." \
  -n 64

The script pins --temp 0 and --show-reasoning on both arms, so the comparison covers every generated token and is not a comparison between two samples. It also runs the baseline twice as a self-consistency control and reports INCONCLUSIVE rather than FAIL when the baseline does not reproduce itself. It refuses to start when the two binaries are byte-identical; that refusal is the no-op build from step 2 surfacing, not a fault in the script. Exit status: 0 all pairs equal, 1 at least one pair differs, 2 inconclusive with no differences.

5. Reach the three changed paged bridges

mlxcel generate never reaches any of them: it is single-sequence and dense, and the fused paged path needs a pool-backed cache. Each bridge has its own production route, and they are not the same route, so "run a server" is not a single answer for all three.

Bridge function	How production reaches it
paged_attention_decode (v1, src/lib/mlx-cpp/turbo/paged_attention.cpp)	Only through mlxcel tune's paged-decode-splits op (src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs:252, driven from src/commands/tune.rs:369). Its other caller is PagedBlockPool::paged_decode_fused (src/lib/mlxcel-core/src/cache/paged.rs:2304), whose only non-test caller is paged_decode_attention_pooled (src/lib/mlxcel-core/src/layers.rs:5406), and that function has no caller anywhere in the tree outside its own doc comments and tests: #720 and ADR 0001 retired it to a library-only API, and production models call paged_decode_attention_dense_compat (a different, unchanged bridge function) instead.
paged_attention_decode_v2_partial (paged_attention_v2.cpp)	The server's pool-backed batched decode, which is the #899 production entry point and default on: src/models/llama3.rs:735 and src/models/qwen3.rs:355 call paged_batch_decode_attention, which reaches PagedBlockPool::paged_decode_batched (src/lib/mlxcel-core/src/cache/paged.rs:1957), then launch_v2, run_decode_v2, and V2Context::launch_with_lse (src/lib/mlxcel-core/src/paged_v2/launch.rs:123). It is gated by the paged_v2::dispatch token floors, which MLXCEL_PAGED_ATTENTION_NATIVE=1 bypasses, and it needs paged decode storage, which effective_decode_storage_backend (src/server/batch/scheduler/mod.rs:216) resolves: for a requested Auto or Paged it returns Paged when max_batch_size > 1 && supports_batching && supports_paged_decode_backend, otherwise Dense. All three hold by default here, since --parallel defaults to -1 which resolves to 4 slots (resolve_n_parallel, src/server/cli_input.rs:1210) and --max-batch-size inherits it, and Llama3 and Qwen3 both return true from supports_paged_decode_backend (src/models/llama3.rs:1556, src/models/qwen3.rs:1127). So a default server already selects the paged path; --decode-storage-backend paged only pins it. --kv-unified is a KV budget flag and selects nothing here.
paged_attention_merge_states (paged_attention_v2_merge.cpp)	Two routes. The same v2 path as above, but only when the plan emits more than one chunk per request (plan.needs_merge, src/lib/mlxcel-core/src/paged_v2/launch.rs:152), so a long-context request is needed, not just any request. And DeepSeek MLA: src/lib/mlxcel-core/src/mla/split_kv.rs:305 through absorbed_decode_split_kv, reached from src/models/deepseek_v2.rs:552 behind MLXCEL_MLA_SPLIT_KV, which makes an ordinary generate run reach it.

So run three things rather than one:

# v1 paged_attention_decode: the only production route is the tuner.
MLXCEL_DEBUG_KERNEL_BACKEND=1 ./target/release/mlxcel tune --op paged-decode-splits

# v2 partial (+ merge on long contexts): the batched server path, not --kv-unified.
MLXCEL_PAGED_ATTENTION_NATIVE=1 MLXCEL_DEBUG_KERNEL_BACKEND=1 ./target/release/mlxcel-server --model models/mlx/Qwen3-30B-A3B-4bit --max-batch-size 4 --decode-storage-backend paged --metrics

# merge_states on its own route: DeepSeek MLA split-KV, reachable from plain generate.
MLXCEL_MLA_SPLIT_KV=1 ./target/release/mlxcel generate --model models/mlx/DeepSeek-V2-Lite-4bit --prompt "..." -n 64 --temp 0

For the server arm, send several concurrent chat completions with long prompts so more than one request decodes at once and each clears the dispatch floor, then repeat the identical sequence against the baseline binary. The workspace test gate in step 3 covers all three bridges directly and is the check to trust if a route above cannot be driven on this node.

The paged arm's pass condition is the decode counters, not output equality

Output equality cannot see this failure. The gather fallback is output-identical to the fused path by construction, so an arm that silently stopped launching the fused kernel scores a clean PASS in step 4. There is an in-tree instrument for exactly this question, so use it rather than inferring.

paged_batch_decode_stats() (src/lib/mlxcel-core/src/cache/paged_batch_decode.rs:213) returns the process-wide counters v2_launches, gather_fallbacks, declines and cascade_launches (plus the cascade token / member sums and cascade_failures). The scheduler publishes them every tick: run_loop.rs:415 calls update_paged_decode_gauges(mlxcel_core::cache::paged_batch_decode_stats()), and /metrics renders three of them (src/server/routes/metrics.rs:225-229), which is why the server command above carries --metrics (the route is mounted unconditionally but the handler answers 501 without the flag, src/server/app.rs:475):

curl -s localhost:8080/metrics | grep mlxcel_paged_decode_launches_total
# mlxcel_paged_decode_launches_total{path="fused_v2"} N
# mlxcel_paged_decode_launches_total{path="gather"}   N
# mlxcel_paged_decode_launches_total{path="cascade"}  N

These are values rather than log lines on purpose, and the doc comment on PagedBlockPool::paged_decode_batched (src/lib/mlxcel-core/src/cache/paged.rs:1948-1951) records why: the outcome "is a value rather than a log line because the first cut of this used tracing::debug! and a whole production benchmark sweep then compared gather against gather without anything saying so." That is this issue's failure mode with a different cause, which is what makes the counters the pass condition here rather than a nice-to-have.

Pass condition for the paged server arm, read after the identical request sequence on each arm:

• v2_launches (path="fused_v2") is non-zero on both arms and moves together across the two runs. Zero on either arm means the fused path never ran and everything else measured gather against gather.
• gather_fallbacks (path="gather") does not rise on the arm alone. This is the counter the regression this PR could introduce actually moves: PR refactor(core): route custom kernels by GPU backend kind #1869 added an early return at the top of PagedBlockPool::paged_decode_batched (src/lib/mlxcel-core/src/cache/paged.rs:1981-1988) that returns Ok((None, PagedDecodeOutcome::NotServable("the GPU backend has no fused paged-attention kernel port"))) when custom_kernels_available() is false, and paged_batch_decode_attention folds an Ok((None, _)) into gather_fallbacks (paged_batch_decode.rs:375), not into declines.
• declines (not published to /metrics; read it from a test or a debug build) is the pre-pool reject count, bumped only when batch_is_servable or the single-token shape check rejects the batch before the pool is touched (paged_batch_decode.rs:328). It should not differ between arms either, but the new early return does not touch it, so it is corroboration and not the primary signal.

The server log corroborates the number with a reason. report_once (paged_batch_decode.rs:152) announces the first occurrence of each distinct outcome kind through tracing::info!, deliberately at info rather than debug, so the line paged decode v2: gather: batch not servable (the GPU backend has no fused paged-attention kernel port) appears without anyone enabling a debug filter. Read it as corroboration only: the flags are keyed by outcome kind, and NotServable is one kind (index 8, src/lib/mlxcel-core/src/paged_v2/outcome.rs:130), so an ordinary batched prefill declining with NotServable("not a single-token decode step") earlier in the run can claim the slot and suppress the kernel-port line. The counter cannot be pre-empted that way.

The two instruments answer different questions and both are needed. MLXCEL_DEBUG_KERNEL_BACKEND from step 1 says what the backend predicate resolved to; the counters say whether the fused path actually ran. A build can print cuda and still serve every step from gather (dispatch floor, multi-slab layer, MLXCEL_PAGED_ATTENTION_NATIVE kill switch), and the counters alone cannot say which of those it was. Post both.

Step 3 already exercises the same instrument at test level: tests/cascade_decode_dispatch.rs:185 reads paged_batch_decode_stats() and asserts it agrees with the returned outcome, on the grounds that otherwise the /metrics view of which kernel ran is decorative.

Two corrections to an earlier revision of this issue, recorded so they are not reintroduced. It said the server arm with --kv-unified and MLXCEL_PAGED_ATTENTION_NATIVE=1 exercised paged_attention_decode; it does not, because that bridge sits behind the library-only pooled entry point and the server never calls it. And --kv-unified is not what enables the paged pool: it only collapses the parallel context slots to one for sizing (effective_parallel_context_slots, src/server/startup.rs:788) and the batch scheduler does not read it at all.

Acceptance Criteria

• The [mlxcel] custom kernel backend: line from step 1 is posted and reads cuda with no fallback suffix.
• The two test result: summaries from step 3 are posted and show no new failures and no group that went from passing to skipped.
• The A/B script's exit status and its per-pair verdict lines from step 4 are posted.
• The build log evidence from step 2 is posted, showing mlx_cxx_kernels.cpp and the src/lib/mlx-cpp/turbo/ sources actually recompiling in the arm build.
• Each of the three routes in step 5 is either driven and reported, or explicitly recorded as not runnable on this node with the reason.
• The mlxcel_paged_decode_launches_total fused / gather / cascade split from step 5 is posted for both arms, showing fused_v2 non-zero on both and gather not rising on the arm alone.
• Any difference found is filed as its own issue, naming the arm it appeared on and the exact command that produced it.

Note

A plain "all equal" is not sufficient on its own. It needs the step 2 build evidence and the step 1 backend line beside it, because the two most likely failure modes in this set both produce exactly an "all equal" result: a no-op rebuild that leaves both arms pointing at the same code, and a predicate that answers wrongly so every fused path silently takes an MLX graph fallback with identical output and lower throughput. For the paged server arm it also needs the step 5 counter split, which is the only one of the three that observes the fused path actually running.

[inureyes]

The Metal half of the same criterion is measured and recorded in #1881: docs/benchmark_results/kernel-backend-kind-metal-m1ultra-2026-09-13.md. Metal selects the same kernels. Four arms, with what was not measured named.

Four traps cost a re-run each there, and all four apply unchanged here.

• An EQUAL row whose backend probe stayed silent proves nothing: the two binaries agreed on a path neither changed. Run with MLXCEL_DEBUG_KERNEL_BACKEND=1 and treat a silent arm as failed. Three of five checkpoints were silent on the first pass, because FUSED_ADD_RMSNORM_DEFAULT and FUSED_ROPE_APPEND_DEFAULT are opt-in, falcon-mamba does not use ssm_update_kernel (falcon_h1, granitemoehybrid, nemotron_h, plamo2 do), and split_kv_enabled() short-circuits on absorbed_enabled().
• Without MLXCEL_PAGED_ATTENTION_NATIVE=1 the paged arm declines at the 4096-token dispatch floor and compares gather against gather, which is the failure the paged_decode_batched doc comment records from the first Route production batched paged decode through the fused v2 kernel, retiring the gather-then-SDPA hot path #899 sweep.
• /metrics is off unless --metrics is passed, and curl -f fails silently, so a missing counter file reads as a zero.
• The discriminator is gather_fallbacks, not declines. The gate at the head of paged_decode_batched returns Ok((None, ...)), which the caller folds into gather_fallbacks; declines counts only pre-pool rejections and is not exported to the gauges.

scripts/paged_decode_counter_ab.sh in #1881 encodes all four and is backend-agnostic, so the CUDA arm should be a single invocation against a GB10-built pair. Also worth knowing before reading its output: mlxcel_paged_decode_launches_total counts layer decode steps, not tokens, so it does not scale with context.