test(cuda): run the deferred CUDA arm for the merged ROCm routing work on a GB10 node
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
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:
```bash
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).
```bash
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.
```bash
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
```bash
./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:
```bash
# 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`):
```bash
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 #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.
Contributor guide
Research direction
Start on a GB10 node with the backend probe in src/lib/mlx-cpp/turbo/gpu_backend.cpp, then build the baseline and arm with the commands in the issue and verify the listed bridge files recompile. Run make verify-test-cuda, scripts/ab_output_equality.sh, and the paged routes through tune, the batched server, and DeepSeek MLA. Done means the backend reports cuda, tests remain passing, outputs match, and paged counters show fused launches without unexpected fallbacks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- git, rust
- Domain
- backend, performance, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100