[yanghaojin]

On GB10, kernels that a CUDA graph runs concurrently slow each other down more than they overlap. Serializing the graph roughly halves the 2048-token prefill of a 27B 4-bit model (6.70 -> 3.05 s) and makes MoE decode about 10% faster. On by default for compute capability 12.1 only, MLX_CUDA_GRAPH_SERIAL=0/1 overrides.

CommandEncoder only adds graph edges from input buffers, so kernels with no data dependency between them run concurrently inside a graph. On a GB10 that costs more than it saves. The same 4-bit quantized matmul (2048 x 17408 x 5120) takes 21.4 ms inside a model's prefill graph and 7.6 ms with graphs off.

This change adds an edge from each node to the one added before it. It only adds dependencies, so outputs are bit-identical.

Turning graphs off gives the same prefill numbers but loses 14-18% on MoE decode (see below), so this keeps the graphs and just serializes them. The existing knobs don't get there either. MLX_MAX_OPS_PER_BUFFER=1 matches the prefill, but MoE decode drops to 36.7 tok/s on the 35B. Larger graphs (MLX_MAX_OPS_PER_BUFFER=100, MLX_MAX_MB_PER_BUFFER=1000) don't change anything.

DGX Spark, 4-bit models (tok/s unless noted):

workload	default	serial	graphs off
27B dense, 2048-token prefill	6.70 s	3.05 s	3.11 s
35B-A3B MoE, 2048-token prefill	1.91 s	1.59 s	1.61 s
27B dense, decode	13.0	13.0	12.5
35B-A3B MoE, decode	62.2	68.3	58.8
80B-class MoE, decode	22.3	24.3	19.9

Reproducing with plain mlx_lm

2048-token prefill of a Qwen3.5-architecture 27B 4-bit model with mlx_lm 0.30.8, three runs each. mlx_lm's own gated-delta loop takes a big chunk of the time here, so the gain is smaller than in the table above. The first run includes warm-up.

MLX_CUDA_GRAPH_SERIAL=0 python bench.py <model>
MLX_CUDA_GRAPH_SERIAL=1 python bench.py <model>

import os, sys, time
import mlx.core as mx
from mlx_lm import load
from mlx_lm.models.cache import make_prompt_cache

model, tok = load(sys.argv[1])
ids = mx.array(tok.encode("The quick brown fox jumps over the lazy dog. " * 300)[:2048])[None]
for i in range(3):
    c = make_prompt_cache(model)
    t = time.perf_counter(); mx.eval(model(ids, cache=c)); dt = time.perf_counter() - t
    print(f"serial={os.environ.get('MLX_CUDA_GRAPH_SERIAL', 'default')} run {i}: {dt:.2f} s")

MLX_CUDA_GRAPH_SERIAL	run 0	run 1	run 2
0	28.97	11.40	11.37
1	10.25	8.86	9.19

I couldn't reproduce the slowdown with a synthetic stack of independent quantized matmuls. It seems to need a real model's graph. I also haven't measured any other GPU, which is why the default is limited to sm_121.

[zcbenz]

CUDA graph being slower because of parallel works seems like a bug in GPU scheduler that that should fixed in driver level, if you could make a standalone repro I would be happy to pass to our contact at NVIDIA.

I don't think we should accept this change because it is not guaranteed to be faster for all kinds of workloads, and a future update of CUDA toolkit could turn it into a performance regression.

In the meanwhile a good way to work around this on client side is to only disable CUDA graphs for prefill, and #4647 will provide an API to allow that.

Also ping @dhiltgen since this looks like a thing that should have special care in inference engines.

[yanghaojin]

The "Reproducing with plain mlx_lm" section in the PR description can be run directly on stock MLX without gbx. The model is publicly available at [GreenBitAI/Qwen3.8-27B-4bit](https://huggingface.co/GreenBitAI/Qwen3.8-27B-4bit), and we're using mlx_lm==0.30.8, as in the PR.

One important detail: on stock MLX, use MLX_USE_CUDA_GRAPHS=1/0 to toggle graphs. The MLX_CUDA_GRAPH_SERIAL=0/1 variable in the PR commands is specific to our patch and has no effect on stock MLX (other than the script's output labels).

On stock MLX main (388f227), GB10, driver 580.95.05, CUDA 13.0, the 2048-token prefill takes 11.1–11.2s with graphs on vs 10.45s off. In gbx, the gap is much larger: 6.70s vs 3.05s.

From nsys --cuda-graph-trace=node, the same qmm_sm80 kernel (gate/up, 2048×17408×5120, grid 4352) takes 12.8ms median with graphs vs 7.4ms without. Gate and up run concurrently on two streams. Small kernels in the GatedDelta loop also slow down significantly (copy_gg_nd: 15.9µs vs 1.8µs on average).

We tried reproducing this in standalone CUDA with two large kernels or a large + small kernel running concurrently, but couldn't reproduce the slowdown.

Once #4647 is merged, we plan to disable CUDA graphs only for prefill at the caller level. Until then, we'll keep the serialization patch in our own build.