Fix the CUDA exit abort and the completion-worker busy loop, and add three CUDA throughput changes

[davidtai]

Summary

This PR fixes the CUDA exit-time abort, stops the CUDA completion-worker busy loop, and adds three CUDA throughput changes from mlx-serve. It is item 1.3 of the mlx-shim program (DBLM-727). Base: main = cb77239. Every commit is signed. Upstream commits keep their upstream author.

#	Commit	What it does	Source
1	20999217 tests: add the CUDA C++ tests from upstream	Adds tests/cuda_tests.cpp (upstream file at 64ea011, unchanged). CUDA builds compile it and link CUDA::cudart, as upstream does.	ml-explore/mlx 64ea011 (file from e0a92d9 ml-explore#4457, 229f5b4 ml-explore#4490)
2	a14f55a8 Leak global CommandEncoder on exit	The global encoder map is a leaked new, so its destructors do not run at exit.	cherry-pick 24c699e (ml-explore#4480)
3	5e03c71e Do not throw from CUDA destructors	~CommandEncoder catches a failed synchronize(). ~CudaHandle, unified_free, free_async and the AtomicEvent deleter do not throw. CUDA hunks only (allocator.cpp, cuda_utils.h, device.cpp, event.cu). The default-stream changes in compile.cpp, stream.{h,cpp}, transforms.cpp are not taken.	9eb3e3a (ml-explore#4514), partial
4	67142db0 Leak thread local streams on exit in main thread	On the main thread the encoder map is a leaked new. CUDA, CPU and Metal hunks all applied without conflicts.	cherry-pick 64ea011 (ml-explore#4576)
5	46cd2e1d Fix completion worker busy loop	The worker kept its batch counter across loops. Before, the wait predicate stayed true after the first batch and the thread spun.	cherry-pick ce916db (ml-explore#4452)
6	52a82aed LRU cache counts consecutive misses	A cache hit resets the miss count. A long-running server no longer fails every new shape. This is the upstream form of the mlx-serve patch mlx-cuda-lru-consecutive-misses. The test case of that patch is not taken: it calls env::get_var, which libmlx.so does not export, so it does not link with BUILD_SHARED_LIBS=ON.	cherry-pick 9ab7999 (ml-explore#4661)
7	f4c5a936 Fused SDPA at head dims 256 and 512	Vector SDPA takes D=256 and 512. cuDNN takes D=256 on sm_120 and later. cuDNN does not take 2 to 15 query rows. Source hunks unchanged. The test hunk was applied with git apply -C1, because in mlx-serve it follows the LRU test that is not taken.	mlx-serve mlx-cuda-sdpa-head-dim-256-512 (ddalcu/mlx-serve b6aaf69). No upstream PR.
8	44847947 Quantized matmul reads the weights once	With qmm_sm80, only 1 row runs qmv; 2 to 7 rows run qmm_sm80. [B, M, K] x against a 2D weight is folded to [B*M, K]. Unchanged.	mlx-serve mlx-cuda-qmm-read-weights-once (b6aaf69). No upstream PR.
9	7a42db35 clear_cache trims the CUDA memory pools	clear_cache synchronizes each free stream and calls cudaMemPoolTrimTo(pool, 0). Not in the original item list. The upstream test "test clear cache trims CUDA pool" (commit 1) failed on spark-2 without it. Only the clear_cache hunk is taken; the Windows WDDM budget is not.	e0a92d9 (ml-explore#4457), partial
10	b1d0ec45 fork.yaml	One section per change, with the upstream ids and mlx-serve patch names.
11	7b80b8b6 style: clang-format	clang-format 21.1.8 (the pre-commit pin) on quantized.cpp and cuda_tests.cpp. Layout only. The output equals the CI "Check Lint" diff (blobs dcc9665, da1e33a). Commits 7 and 8 stay byte-equal to the mlx-serve patches.	review finding 1
12	4d3bdafe QuantizedMatmul returns early on an empty output	After the fold, x [0, M, K] with a 2D weight gave M = 0 and a divide by zero. Return after the output allocation when out.size() == 0, as the CUDA indexing primitives do.	review finding 8
13	a1b60f40 Export InTracing and RetainGraph from a shared libmlx	struct MLX_API InTracing, struct MLX_API RetainGraph in mlx/transforms_impl.h. The Layr-Labs mlx-c calls detail::in_tracing() and detail::retain_graph(), which read InTracing::trace_stack() and RetainGraph::tracing_counter. A hidden-visibility shared libmlx did not export them, so a shared libmlxc had undefined symbols. Only BUILD_SHARED_LIBS=ON shows this; the static SwiftPM Cmlx build never did. Upstream main 243bdf0 does not export them either, so there is no upstream commit to take. Code is unchanged from the stage patch, so the stage guard (struct MLX_API RetainGraph) skips it.	mlx-shim adapters 4d4807c patches/mlx/0001-export-intracing-retaingraph.patch (item 1.1)
14	194a14eb fork.yaml	Review wording (qmm numerics and fp modes, SDPA cut at all head dims, LRU cache list) and a section for commit 13.	review findings 3, 5, 16

fork.yaml sections added: "Export InTracing and RetainGraph from a shared libmlx", "CUDA C++ tests from upstream", "CUDA exit-time abort (upstream backports)", "CUDA completion-worker busy loop (upstream backport)", "CUDA clear_cache trims the memory pools (upstream backport)", "CUDA LRU cache thrashing check counts consecutive misses", "CUDA fused SDPA at head dims 256 and 512 (mlx-serve)", "CUDA quantized matmul reads the weights once (mlx-serve)". The v0.32.2 sync section now says that some of its files carry later fork changes.

python3 scripts/check_forkdiff.py --upstream-ref upstream/main at 194a14eb: rc 0, 190 changed files, 190 described.

Commits 11 to 14 were not built or run on the box. The box receipts below are for b1d0ec45.

Verification on spark-2 (GB10, sm_121, CUDA 13.0)

Folder: /home/eigenlabs/mlx-cbv2-cuda/core-1.3/. Both builds use the same flags: CMake Release, BUILD_SHARED_LIBS=ON, MLX_BUILD_CUDA=ON, MLX_CUDA_ARCHITECTURES=121a-real (the Phase 1 value), MLX_BUILD_GGUF=OFF, MLX_BUILD_TESTS=ON, OpenBLAS 0.3.33 and cuDNN 9.27 wheel from the Phase 1 folder (read-only), nice -n 10, -j 8. base = cb77239, branch = b1d0ec45. All GPU runs held /tmp/mtplx-gpu-exclusive.lock.

Exit abort

Program progs/exit_variants.cpp: one add on the GPU, eval, then exit.

Variant	base rc	branch rc
default stream, return from main	0	0
default stream, std::exit(3) from main	3	3
default stream, op and std::exit(3) on a second thread	3	3
op on main, std::exit(3) from a second thread	3	3
new_thread_unsafe_stream(gpu), return from main	134	0
new_thread_unsafe_stream(gpu) on a second thread, std::exit(3) there	134	3

The base abort text is the Phase 1 text: terminate called after throwing an instance of 'std::runtime_error' what(): cudaStreamSynchronize(stream_) failed: driver shutting down. A thread-unsafe stream lives in the global static encoder map, which the Phase 1 stack destroys in __run_exit_handlers. With the default (thread-local) stream, the abort did not occur on base. Receipts: results/base/exitv-*.std{out,err}, results/{base,branch}/rc.txt.

C++ tests (build/<v>/tests/tests, run with the installed libmlx.so)

	base	branch
test cases	280: 274 pass, 6 fail	283: 277 pass, 6 fail
CUDA cases (test clear cache trims CUDA pool, test fused sdpa at head dims 256 and 512, test quantized matmul over a few rows)	not present	3 of 3 pass (39 assertions)

The 6 failures are the same on both: 4 GGUF cases (GGUF is off by the requested flag), test matrix cholesky and test matrix eigh (CPU linalg). Receipts: results/{base,branch}/tests.stdout, results/branch/tests-cuda.stdout.

Two earlier branch runs failed, and the fixes are in this PR. (a) test clear cache trims CUDA pool failed without commit 9. (b) test quantized matmul over a few rows failed when the binary loaded libmlx.so from the build tree: NVRTC could not find cute/numeric/numeric_types.hpp, because the JIT include path is dirname(libmlx.so)/../include. With the installed library it passes. Receipts: results-run1/branch/.

Quantized matmul, 4-bit, x [rows, 4096] bf16, W [12288, 4096], group 64

GPU kernel time from nsys --cuda-graph-trace=node, 300 serial calls (eval and synchronize after each), median:

Case	base kernel	base us	branch kernel	branch us	Change
6 rows	qmv_kernel	468.2	qmm_sm80_kernel	167.1	-64 %
1 row (run 1 / run 2)	qmv_kernel	84.7 / 83.5	qmv_kernel	82.6 / 85.4	none

Wall time per op, 50 independent calls per eval, median of 21 repeats, two runs (us):

x shape	base run 1	base run 2	branch run 1	branch run 2
1x4096	87.3	88.4	95.1	105.2
2x4096	165.8	154.7	124.8	140.5
4x4096	329.4	293.6	136.5	145.3
6x4096	484.6	431.7	135.2	136.9
7x4096	565.6	517.3	132.3	148.0
8x4096	130.3	137.0	134.6	123.8
16x4096	136.5	146.1	137.6	136.5
6x1x4096 (batched)	515.8	471.6	136.8	141.8

The plan acceptance is "6-row case under 150 us". The kernel time of one serial call is 167.1 us, so that acceptance is not met by the kernel time. The 135 us wall time per op is lower than the kernel time because the 50 calls in one graph have no dependencies and their kernels overlap (nsys: 108 to 124 us GPU span per op, 296 us median kernel duration while overlapped). The mlx-serve figure (485 to 94 us) is from an sm_120 part. A full read of the weights, scales and biases is about 28 MB, which is about 104 us at the GB10 nominal 273 GB/s. Receipts: results/run2b/kern-serial-*.txt, results/run4/*-kern.txt, results/{base,branch}/qmm.csv, results-run1/{base,branch}/qmm.csv.

Completion worker

Program progs/worker_cpu.cpp: 20 s loop of one 1-row 4-bit matmul, eval, synchronize, then a sleep. CPU time per thread from /proc/self/task/*/stat. Runs alternate base and branch (ABAB).

Loop	Metric	base run 1	base run 2	branch run 1	branch run 2
5 ms sleep	busiest thread CPU (s per 20 s)	19.98	19.99	0.78	0.65
5 ms sleep	process CPU (s)	21.13	21.12	1.44	1.27
5 ms sleep	step time, eval + synchronize, median (us)	870.0	893.4	1294.2	1207.8
no sleep	busiest thread CPU (s)	19.97	19.94	3.37	3.18
no sleep	process CPU (s)	28.52	28.51	7.97	7.67
no sleep	step time median (us)	313.8	336.9	510.3	492.7
no sleep	steps in 20 s	48979	47775	34616	35939

The fix removes the spinning thread: process CPU falls by 72 % to 94 %. It also makes each synchronize() slower on this box, by about 175 us with no sleep and about 370 us after a 5 ms sleep. The reason: CommandEncoder::synchronize() waits for a completion handler that the worker thread runs. Before, the spinning thread picked it up at once. Now the thread must wake from its condition variable. Upstream accepted the same trade. A decode loop that synchronizes once per token pays this cost once per token. Receipts: results/run2b/worker-*.txt, results/run3/worker-nosleep-*.txt.

GPU lock hold times (UTC, 2026-10-13)

Run	Start	End	Hold
run 1, base (superseded)	22:27:21.088	22:29:29.787	128.7 s
run 1, branch at 8f80cfa (superseded)	22:29:29.795	22:30:00.231	30.4 s
run 2, base	22:32:20.301	22:32:58.906	38.6 s
run 2, branch	22:32:58.913	22:33:48.563	49.7 s
run 2b, worker ABAB and nsys 6 rows	22:35:34.051	22:37:02.434	88.4 s
run 3, worker ABAB, no sleep	22:38:09.494	22:39:31.934	82.4 s
run 4, nsys 1 row ABAB	22:41:13.940	22:41:25.441	11.5 s

Review findings (open)

From the independent review (verdict MERGE WITH FIXES). Findings 1, 3, 5, 8 and 16 are fixed by commits 11, 12 and 14. These stay open:

#	Finding	State
4	supports_qmm_sm80 (mlx/backend/cuda/quantized/qmm/qmm.cu:64-100) has no mode check. On sm_120/121, mxfp4, nvfp4 and mxfp8 at 2 to 7 rows move from fp_qmv (which qmm.cu:138-139 says is faster on sm_120) to qmm_sm80.	Not tested, not timed. See "Not measured".
6	The 1-row wall time per op rose 8 % to 19 % (87.3 / 88.4 to 95.1 / 105.2 us). The 1-row kernel time did not change. Cause unknown.	A/B with only 46cd2e1d reverted is pending, before the engine pin moves.
7	The exit fix is shown only on C++ programs.	BenchCBv2 to DONE with rc 0 on spark-2 is pending the item 1.1 stage, before item 1.3 is closed.
9	7a42db35 changes clear_cache: it now blocks until the free stream drains and trims the pool to 0. The engine calls clearCache() every 256 tokens in the Gemma 4 MTP loop and after weight load. Cost not measured.	David keeps or drops it. If dropped, also drop the test at tests/cuda_tests.cpp:72-111.
10	The extra synchronize() latency after 46cd2e1d comes from the upstream design (the worker sleeps on a condition variable and synchronize() waits for a handler that it runs). It is not a setting.	None in this PR. Removing the hop needs new code (fork or upstream PR).

Not measured

• Finding 4: fp modes (mxfp4, nvfp4, mxfp8) at 2 to 7 rows on sm_120/121 now run qmm_sm80 instead of fp_qmv. No test and no timing. To keep the change to affine mode only, one condition at mlx/backend/cuda/quantized/quantized.cpp:126 is enough: if (can_use_qmv && (M * B == 1 || (mode_ != QuantizationMode::Affine && M * B < 8))). This restores the old fp_qmv routing for fp modes and keeps affine on qmm_sm80. Not applied.

• SDPA at head dims 256 and 512: only correctness (the C++ test passes). No timing.

• LRU change: no test in this PR (see commit 6).

• The exit fix was not checked through mlx-swift or BenchCBv2. Only the C++ programs above.

• The CPU and Metal hunks of 64ea011 compile on CUDA Linux (CPU part); the Metal part was not built here. Fork macOS CI covers it.

• No Python build.

🤖 Generated with Claude Code

[davidtai]

Wording note for commit f4c5a936 ([CUDA] Fused SDPA at head dims 256 and 512). The commit is kept byte-equal to the mlx-serve patch, so its message is not rewritten.

The cuDNN cut for 2 to 15 query rows (mlx/backend/cuda/scaled_dot_product_attention.cpp:340) applies at all head dims, not only 256 and 512. At head dims 64, 96 and 128, 2 to 3 query rows now take the vector kernel and 4 to 15 query rows take the unfused fallback. Before, cuDNN served them. This is not timed. fork.yaml (commit 194a14eb) now says so.

🤖 Generated with Claude Code

[davidtai]

Box results on spark-2 for the two open review findings (record: work/mlx-shim-item-1.3-benchcbv2.md in the layr-labs work folder):

Check	Result
BenchCBv2 correctness and perf built against a CMake stage of 194a14eb, exit code	0 (134 on cb77239)
BenchCBv2 prefill / decode, Qwen3.5-9B-4bit, B=1	1,048 / 33.0 tok/s (before this branch: 1,075 / 33.2)
qmm 1-row wall time per op, with 46cd2e1d / with it reverted (us, median of 3 runs)	108.8 / 96.7
qmm 6-row wall time per op, same A/B (us)	141.5 / 139.2

Finding 7 (exit abort end to end) is closed. Finding 6 is confirmed: the 1-row cost is the worker-thread wake-up that 46cd2e1d adds to each synchronize(); the engine's decode rate did not change because the cost is paid once per step.