epic: CUDA (MLX backend) prefill and decode acceleration program #623

Summary

Program of work to raise prefill and decode throughput of the MLX backend on CUDA (primary validation target: NVIDIA GB10, sm_121, 122 GB unified LPDDR5x at ~273 GB/s). Based on the 2026-06-17 GB10 benchmark sweep (benchmarks/cuda_gb10_2026-06-17.csv, 147 models) compared against the same-week Metal M1 Ultra sweep, plus a code audit of the generation hot path, the CUDA binding layer, and the serving layer.

Diagnosis (baseline)

Dense single-stream decode is already close to the memory-bandwidth roofline on GB10:

Consequences:

• For 7B+ dense models, kernel micro-tuning has at most ~15% headroom. The big single-stream levers are fewer bytes (KV quantization, lower-bit weights) and fewer steps (speculative decoding).
• Models under ~2B are launch/dispatch-overhead bound (67-75% attained). Pipelining and host-overhead work pays off there.
• Several model families are far below roofline on CUDA and represent outright defects rather than tuning targets:

Code-audit findings that shape the plan:

• The pinned MLX commit a6ec7123 (2026-06-10) predates upstream CUDA work that targets exactly these gaps: qmv_wide small-batch quantized matvec (#3764), RoPE without copy (#3704), JIT-compiled qmm/gather_gemm plus kernel cache (#3706, #3576, #3587), fused SDPA for asymmetric Q/V head dims (#3637), gather_qmm NAX kernel-name fix (#3632), and disabling CUDA managed memory on Tegra (#3701). GB10 is a Tegra-family unified-memory SoC, so #3701 alone may shift all numbers.
• The server BatchScheduler decode loop is synchronous per step (eval + item_i32, src/server/batch/scheduler.rs:4778), while the CLI loop already overlaps graph build with execution via async_eval_pair (src/lib/mlxcel-core/src/generate.rs:1157-1260).
• The fused decode-MoE CUDA kernel is proven (+55% on qwen3-30b-a3b) but its activation threshold default (MLXCEL_FUSED_MOE_MAX_DFF=4096) is tuned for Metal; the measured CUDA crossover is ~13-14k (docs/benchmark_results/fused-moe-decode-kernel-design.md).
• The native paged-attention kernel is Metal-only (src/lib/mlx-cpp/turbo/paged_attention.cpp); CUDA always takes the gather-then-SDPA fallback.
• Serving throughput machinery (continuous batching, chunked prefill, prompt cache, paged KV) exists but ships disabled: effective batch size defaults to 1, --max-batch-prefill defaults to 1, prompt cache is off unless configured.
• Host-side per-token costs: full-history re-detokenization every token (src/server/model_worker.rs:2014), 3 mutex locks + JSON serialization per streamed token, prompt tokenization on the single generation thread.

Phases and sub-issues

Implementation order. Each phase gate is a benchmark comparison against the phase's entry baseline.

Phase 0: measurement foundation

• perf(bench): long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry #624 Long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry

Phase 1: upstream sync and re-baseline (gate for Phase 3 scoping)

• chore(mlx): upgrade pinned MLX commit, rebase CUDA overlays, GB10 re-baseline and outlier triage #625 Upgrade pinned MLX commit, rebase CUDA overlays, full GB10 re-baseline and outlier triage

Phase 2: low-risk defaults and knobs (parallelizable after Phase 1)

• perf(cuda/moe): backend-specific default for the fused decode-MoE Dff threshold #626 CUDA-specific fused decode-MoE Dff threshold default
• perf(cuda): allocator policy: wire set_cache_limit and re-evaluate periodic clear_memory_cache #627 CUDA allocator policy: wire set_cache_limit, revisit periodic clear_memory_cache
• perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #628 Serving-throughput defaults: parallel decode, batched prefill, prompt cache

Phase 3: targeted CUDA gap fixes (scope re-confirmed by Phase 1 triage; close any item upstream already fixed)

• perf(cuda/moe): fix MoE prefill collapse on CUDA (gather_qmm / grouped GEMM path) #629 MoE prefill collapse on CUDA (gather_qmm / grouped GEMM path)
• fix(cuda/quant): repair nvfp4 execution path on CUDA #630 nvfp4 execution path repair on CUDA
• perf(cuda/ssm): close the hybrid-SSM decode gap on CUDA #631 Hybrid-SSM decode gap on CUDA

Phase 4: dispatch and host overhead

• perf(server): port lookahead async_eval pipelining into BatchScheduler decode #632 Port lookahead async_eval pipelining into BatchScheduler decode
• perf(server): incremental detokenization and per-token streaming overhead reduction #633 Incremental detokenization and per-token streaming overhead reduction

Phase 5: long-context and bandwidth reduction

• perf(cuda/attn): port the native paged-attention decode kernel to CUDA #634 Port native paged-attention kernel to CUDA
• perf(cuda/kv): validate and enable quantized KV cache modes on CUDA #635 Validate and enable quantized KV cache modes on CUDA
• perf(cuda): single-dtype decode graph: eliminate per-token AsType conversions #636 Single-dtype decode graph on CUDA (AsType elimination)

Phase 6: larger bets

• perf(cuda/quant): sm_120/sm_121-tuned quantized GEMM for Blackwell prefill #637 sm_120/sm_121-tuned quantized GEMM for Blackwell prefill
• perf(speculative): CUDA drafter pairing matrix and continuous-batching integration #638 Speculative decoding: CUDA pairing matrix and continuous-batching integration

Working agreements for implementers

• Hardware: CUDA items require a CUDA machine; the reference target is GB10 (sm_121). Checkpoints are expected under ./models/ as used by scripts/bench_all_models.sh and the CSVs in benchmarks/.
• Build: cargo build --release --features cuda. The first build compiles MLX from source via CMake (src/lib/mlx-cpp/CMakeLists.txt); expect a long cold build. CI does not run the test suite, so cargo test --features cuda locally is the merge gate.
• Benchmarks: scripts/bench_decode.sh (same-process decode harness, binary mlxcel-bench-decode), scripts/bench_all_models.sh (sweep, appends CSV under benchmarks/), scripts/bench_mlxlm.py (cross-check vs Python mlx-lm).
• Every sub-issue must land integrated into the real code path it targets (CLI and/or server), not as a standalone module. Benchmarks or metrics proving the effect are part of done.
• Do not regress Metal. Anything backend-conditional must key off runtime backend detection (existing pattern: !mlx::core::metal::is_available() in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp) or build features, and keep the Metal defaults unchanged.

Non-goals

• XLA backend work (separate track, see spike/openxla/).
• Multi-GPU tensor parallelism (scaffolded in src/lib/mlxcel-core/src/streams.rs for issues feat: single-node multi-GPU tensor parallelism to load and serve models larger than one GPU (CUDA/DGX) #486/feat: multi-GPU tensor-parallel runtime with per-rank device placement and cross-GPU collectives #488; not part of this epic).
• Apple Silicon (Metal) performance work, except for guarding against regressions.