Fence in command encoder

[zcbenz]

There is a mlx::core::Fence, which is used for synchronizing between streams.

There is also a mlx::core::metal::Fence, which is for synchronizing between command encoders, because the buffers are allocated with unmanaged mode ResourceHazardTrackingModeUntracked. And the CommandEncoder class provides methods to register inputs and outputs to manually track dependencies and setup barriers.

So for a backend that does not have the concept of command buffer encoder, we can just ignore the second type of Fence and there is no need to register inputs/outputs, like the cpu::CommandEncoder, am I understanding it right?

[awni]

That's exactly right. We may want to use multiple CUDA streams (I can't recall how much parallelism it will do if you don't).. but I would probably do the simple thing first (e.g. just launch kernels the default way). Maybe checkout the docs on CUDA streams and just to make sure we design in a way that's extensible to that if needed.

[zcbenz]

CUDA streams run kernels sequentially, like the MLX's cpu stream:
https://docs.nvidia.com/cuda/cuda-c-programming-guide/#streams

I think I will map a MLX gpu stream to a CUDA stream first.

I'm also confused what a stream actually means in MLX.

For cpu backend a stream is a thread that runs commands in sequence, for metal backend a stream is a command queue that executes multiple commands in parallel. I assume for cuda backend we should follow metal backend to run multiple cuda streams in a mlx stream for concurrent kernel execution, but what would a mlx stream mean for cuda backend then, since there is no concept of command queue in cuda?

[awni]

I think I will map a MLX gpu stream to a CUDA stream first.

Makes sense!

For cpu backend a stream is a thread that runs commands in sequence, for metal backend a stream is a command queue that executes multiple commands in parallel.

We will probably also make the CPU stream run tasks in parallel soon.

I'm also confused what a stream actually means in MLX.

The notion of a stream is becoming less relevant and it's crossed my mind several times to try and get rid of it. It used to be more clear what it meant (everything in sequence) but that is changing. Now it is much closer to stream == device. However, there is still at least some cases where it's useful which is why we haven't removed it entirely.

For now a stream maps to a device and more than one stream can map to the same device which gives the user some explicit control over how to paralellize.

[zcbenz]

I looked into how PyTorch handles cuda multi-streams and I think it just uses one cuda stream per device by default. It does provide stream APIs to explicitly run commands in cuda streams, but the discussions I found were mostly about how kernels were not running concurrently when using it.

Quoting pytorch/pytorch#59692 (comment):

If operations on a single stream are big enough to saturate the gpu (as they are in your case), using multiple streams is not expected to help. The situations where multi-stream execution actually helps are pretty limited - you want kernels that don't utilize all the gpu resources (e.g. don't launch many threadblocks), and also you want kernels that run long enough so that overhead of launching and synchronizing streams does not increase total time.

This probably explains why I couldn't find much useful resources about scheduling cuda streams.

[awni]

The situations where multi-stream execution actually helps are pretty limited - you want kernels that don't utilize all the gpu resources (e.g. don't launch many threadblocks), and also you want kernels that run long enough so that overhead of launching and synchronizing streams does not increase total time.

This case is actually comes up more important for inference. It actually makes a difference there which is why we try to do as much in parallel as possible.