Support for Batch Generation

[rudrankriyam]

I found the PR by Awni cool in mlx-lm about batch generation and was experimenting with it over the weekend. I was able to implement it with almost same benchmark numbers on my M5 MacBook Pro with Llama 3.2 3B 4-bit:

Batch Size	MLX LM (t/s)	MLX Swift (t/s)
1	61	62
2	122	118
32	349	344

There are subtle improvements that I have not been able to find but I think a review would help me out.

Creating this issue in-case if somebody else is already working on it. If not, I can clean up my branch and send a PR!

[ronaldmannak]

I’d love to see the tps for non-batched inference so you can see the overhead for batch-size = 1. Based on my limited experience with your branch, I’d estimate it’s around 120 TPS, which suggests that the overhead for batch size = 1 is quite significant. I experimented with your code today but wasn’t able to meaningfully improve performance. Given the results from MLX-LM, I suspect this is close to the practical limit for now.

FYI — I’ve forked your code and am working on adding continuous batching (similar to how vLLM does it), which allows new requests to join an already running batch. This should make a big difference for server-side and multi-agent scenarios.

The goal is that when a single request comes in, it runs as a non-batched inference by default (no overhead); but if a second request arrives before the first finishes, the scheduler will upgrade the running inference to batched mode once the prefill phase of the new request completes.

[rudrankriyam]

This one can go in the new mlx-swift-lm package 😄

[ronaldmannak]

Yes! I'm adding a few more optimizations and then we have to figure out how we're going to merge with the new lm package. @rudrankriyam I think I'll finalize my PR with your repo first and then you can create a new PR for the combined work in mlx-swift-lm I guess