Performance discrepancy vs torch on cuda

[arogozhnikov]

Hi, I'm exploring mlx on cuda. I've ran a simple test and I see 2x (or more) slower inference with mlx.

Here is an example output:

torch-fwd compiled=False: 0.143911 seconds
torch-fwd compiled=False: 0.042127 seconds
torch-fwd compiled=False: 0.042115 seconds
torch-fwd compiled=False: 0.042109 seconds
torch-fwd compiled=False: 0.042113 seconds
torch-fwd compiled=False: 0.042116 seconds
torch-fwd compiled=False: 0.042115 seconds
torch-fwd compiled=False: 0.042126 seconds
torch-fwd compiled=False: 0.042122 seconds
torch-fwd compiled=False: 0.042262 seconds

torch-fwd compiled=True: 0.416149 seconds
torch-fwd compiled=True: 0.042157 seconds
torch-fwd compiled=True: 0.042157 seconds
torch-fwd compiled=True: 0.042162 seconds
torch-fwd compiled=True: 0.042154 seconds
torch-fwd compiled=True: 0.042320 seconds
torch-fwd compiled=True: 0.042342 seconds
torch-fwd compiled=True: 0.042221 seconds
torch-fwd compiled=True: 0.042153 seconds
torch-fwd compiled=True: 0.042163 seconds

kwargs={'dtype': mlx.core.bfloat16, 'stream': Stream(Device(gpu, 0), 2)}
mlx-fwd: 1.070995 seconds
mlx-fwd: 0.092245 seconds
mlx-fwd: 0.092189 seconds
mlx-fwd: 0.092177 seconds
mlx-fwd: 0.092305 seconds
mlx-fwd: 0.092356 seconds
mlx-fwd: 0.092400 seconds
mlx-fwd: 0.092474 seconds
mlx-fwd: 0.092326 seconds
mlx-fwd: 0.092167 seconds

And here is the minimal example:

import time
from contextlib import contextmanager


@contextmanager
def timeit(label="Elapsed"):
    start = time.perf_counter()
    yield
    end = time.perf_counter()
    print(f"{label}: {end - start:.6f} seconds")


dim_main = 256
dim_hid = 768


def bench_mlx():
    import mlx.core as mx

    def mlx_simple_layer(
        inp,
        proj_ab_W,
        proj_ab_b,
        lin_out_W,
        lin_out_b,
    ):
        hid = mx.addmm(proj_ab_b, inp, proj_ab_W.T)
        return mx.addmm(lin_out_b, hid, lin_out_W.T)

    mlx_simple_layer = mx.compile(mlx_simple_layer)

    kwargs = dict(
        dtype=mx.bfloat16,
        stream=mx.new_stream(mx.Device(type=mx.DeviceType.gpu)),
    )
    print(f"{kwargs=}")

    res = mx.random.normal([1, 1024, 1024, dim_main], **kwargs)

    for _ in range(10):
        with timeit("mlx-fwd"):
            for _ in range(10):
                res = mlx_simple_layer(
                    inp=res,
                    proj_ab_W=mx.ones([dim_hid, dim_main], **kwargs),
                    proj_ab_b=mx.ones(dim_hid, **kwargs),
                    #
                    lin_out_W=mx.ones([dim_main, dim_hid], **kwargs),
                    lin_out_b=mx.ones([dim_main], **kwargs),
                )
            mx.eval(res)  # force sync


def bench_pytorch(compiled: bool):
    import torch
    import torch.nn.functional as F

    def torch_simple_layer(
        inp,
        proj_ab_W,
        proj_ab_b,
        lin_out_W,
        lin_out_b,
    ) -> torch.Tensor:
        hid = F.linear(inp, proj_ab_W, bias=proj_ab_b)
        return F.linear(hid, lin_out_W, bias=lin_out_b)

    if compiled:
        torch_simple_layer = torch.compile(torch_simple_layer)

    kwargs: dict = dict(
        dtype=torch.bfloat16,
        device=torch.device("cuda:0"),
    )
    res = torch.randn((1, 1024, 1024, dim_main), **kwargs)
    for _ in range(10):
        with timeit(f"torch-fwd {compiled=}"):
            for _ in range(10):
                res = torch_simple_layer(
                    inp=res,
                    proj_ab_W=torch.ones((dim_hid, dim_main), **kwargs),
                    proj_ab_b=torch.ones(dim_hid, **kwargs),
                    #
                    lin_out_W=torch.ones((dim_main, dim_hid), **kwargs),
                    lin_out_b=torch.ones(dim_main, **kwargs) * 10000,
                )
            res[0, 0, 0, 0].item()  # force sync


bench_pytorch(compiled=False)
bench_pytorch(compiled=True)
bench_mlx()

Testing setup:

• I use A100
• environment can be installed with uv pip install mlx[cuda]==0.26.3 torch==2.8.0 (it is not simple to find versions that would work on both. I've checked with mlx=0.29 and performance did not change).

Some observations:

• on fp32 speeds of mlx and pytorch are very close
• I've looked at profiler and I think they even use the same funcs from cudnn internally for bf16 matmuls.

[zcbenz]

For simple computations like this the performance is overhead-bound, and mlx builds cuda graph by default which does increases overhead a lot, if you turn off cuda graph by setting env MLX_USE_CUDA_GRAPHS=0 the overhead would be lower and your example should run faster, but for inferencing LLMs the cuda graph is essential for faster inference (#2545).

Also if you increase the dims to large enough that the code becomes memory- or computation-bound, the performance difference should be very small.

[arogozhnikov]

hi @zcbenz

Tested, MLX_USE_CUDA_GRAPHS=0 makes it only slower.

python '/workspaces/models/.personal/test_mlx.py'
kwargs={'dtype': mlx.core.bfloat16, 'stream': Stream(Device(gpu, 0), 2)}
mlx-fwd: 0.824924 seconds
mlx-fwd: 0.091111 seconds
mlx-fwd: 0.091357 seconds
mlx-fwd: 0.091116 seconds
mlx-fwd: 0.091181 seconds


oot@0906c38af594:/workspaces/models# MLX_USE_CUDA_GRAPHS=0 python '/workspaces/models/.personal/test_mlx.py'
kwargs={'dtype': mlx.core.bfloat16, 'stream': Stream(Device(gpu, 0), 2)}
mlx-fwd: 1.091233 seconds
mlx-fwd: 0.098393 seconds
mlx-fwd: 0.097776 seconds
mlx-fwd: 0.098693 seconds
mlx-fwd: 0.099149 seconds

I don't expect this example to be overhead-bounded. Sizes like 1_000_000x256 @ 256x768 shouldn't be overhead-bounded (and you see in my previous example that torch shows 2x performance with/without compilation).

I can confirm that leaving only matmuls still shows >2x difference in performance.

Update: can also confirm difference persists for fp16 (not fp32).

[zcbenz]

Thanks for running the tests, I'll do a profiling to see what is going on.

[zcbenz]

This is the profiling of a single step with PyTorch after warming up, which took 42ms in total and the first gemm took 2.4ms:

---

This is the result of MLX, which took 88ms in total with first gemm took 5.2ms:

The kernel and the setup were exactly the same, what is happening?! 🤯

[awni]

That is very strange.. crazy stuff like this makes me suspect managed memory..

[zcbenz]

So far one difference from PyTorch's cublasLt code I noticed is that, for AddMM PyTorch seems to pass the bias as epilogue to GEMM (i.e. the bias in below graph), while we are passing the bias as C which probably maps to an expensive op.

This is how the parameters are used in cublasLt:

[zcbenz]

The addmm will be on par with PyTorch after #2569, but I think there is still one thing we can improve to make the benchmarking run faster: the overhead for ones is too high.

---

Also the script creates a new stream instead of using the default stream which seems to slow down things a little, I'm interested to find out why.

    kwargs = dict(
        dtype=mx.bfloat16,
        stream=mx.new_stream(mx.Device(type=mx.DeviceType.gpu)),
    )

[zcbenz]

To answer my last comment, the script uses a new stream for inputs:

                res = mlx_simple_layer(
                    inp=res,
                    proj_ab_W=mx.ones([dim_hid, dim_main], **kwargs),
                    proj_ab_b=mx.ones(dim_hid, **kwargs),
                    #
                    lin_out_W=mx.ones([dim_main, dim_hid], **kwargs),
                    lin_out_b=mx.ones([dim_main], **kwargs),
                )

but does not specify a stream for ops:

        hid = mx.addmm(proj_ab_b, inp, proj_ab_W.T)
        return mx.addmm(lin_out_b, hid, lin_out_W.T)

which makes ones and addmm run on different streams because the addmm (and also .T) would run on the default stream if there is no one specified, and the ops can no longer run on the same cuda graph.

So if we change the script to use the default stream, or do something like below to ensure all ops run in the same stream:

with mx.stream(mx.new_stream(mx.Device(type=mx.DeviceType.gpu))):
   ...

We would then get very similar performance with PyTorch:

torch-fwd compiled=False: 0.264826 seconds
torch-fwd compiled=False: 0.044184 seconds
torch-fwd compiled=False: 0.044092 seconds
torch-fwd compiled=False: 0.044132 seconds
torch-fwd compiled=False: 0.044120 seconds
torch-fwd compiled=False: 0.044083 seconds
torch-fwd compiled=False: 0.044075 seconds
torch-fwd compiled=False: 0.044076 seconds
torch-fwd compiled=False: 0.044133 seconds
torch-fwd compiled=False: 0.044097 seconds
torch-fwd compiled=True: 1.110209 seconds
torch-fwd compiled=True: 0.044203 seconds
torch-fwd compiled=True: 0.044224 seconds
torch-fwd compiled=True: 0.044247 seconds
torch-fwd compiled=True: 0.044233 seconds
torch-fwd compiled=True: 0.044297 seconds
torch-fwd compiled=True: 0.044198 seconds
torch-fwd compiled=True: 0.044251 seconds
torch-fwd compiled=True: 0.044287 seconds
torch-fwd compiled=True: 0.044197 seconds
kwargs={'dtype': mlx.core.bfloat16}
mlx-fwd: 1.117719 seconds
mlx-fwd: 0.253744 seconds
mlx-fwd: 0.044851 seconds
mlx-fwd: 0.044823 seconds
mlx-fwd: 0.044756 seconds
mlx-fwd: 0.044705 seconds
mlx-fwd: 0.044739 seconds
mlx-fwd: 0.044785 seconds
mlx-fwd: 0.044723 seconds
mlx-fwd: 0.044731 seconds

[arogozhnikov]

nice find @zcbenz :)