MLX Batch Generation Benchmark

Benchmarking and utilities for parallel LLM inference on Apple Silicon using MLX.

Overview

This project explores the batch generation capabilities of MLX for running multiple LLM inference requests simultaneously. We test how many concurrent "threads" can efficiently run on a single model and measure the throughput gains.

Model tested: openai/gpt-oss-20b (21B parameters, 3.6B active - MoE architecture)

Key Findings

Short Generation (100 tokens per prompt)

Batch Size	Total TPS	Per-Prompt TPS	Speedup	Memory
1	97 t/s	97 t/s	1.00x	13.76 GB
2	136 t/s	68 t/s	1.40x	13.76 GB
4	194 t/s	48 t/s	1.99x	13.76 GB
8	273 t/s	34 t/s	2.80x	13.76 GB
16	281 t/s	18 t/s	2.88x	13.76 GB
24	295 t/s	12 t/s	3.03x	13.76 GB
32	293 t/s	9 t/s	3.01x	13.76 GB

Optimal batch size: 24 - Peak throughput of ~295 tokens/second (3x speedup over single)

Long Generation (1500 tokens per prompt - ~1000 words)

Batch Size	Total TPS	Per-Prompt TPS	Speedup	Time per Batch
1	98 t/s	98 t/s	1.00x	15s
2	123 t/s	61 t/s	1.25x	24s
4	158 t/s	39 t/s	1.60x	38s
8	188 t/s	23 t/s	1.91x	64s
16	144 t/s	9 t/s	1.46x	167s

Optimal batch size: 8 - Peak throughput of ~188 tokens/second for long content

Why the Difference?

• Short prompts: GPU compute is the bottleneck → larger batches help
• Long prompts: Memory bandwidth becomes the bottleneck → KV cache grows with sequence length
• Batch 16+ with long prompts: Performance degrades due to memory pressure

Installation

pip install mlx mlx-lm

Usage

Benchmark Script

# Basic benchmark with default settings
python mlx_batch_benchmark.py

# Custom batch sizes
python mlx_batch_benchmark.py --batch-sizes 1,2,4,8,16,24,32

# Long story generation (unique prompts, avoids KV cache benefits)
python mlx_batch_benchmark.py --long-stories --max-tokens 1500 --batch-sizes 1,2,4,8,16

# More runs for stable averages
python mlx_batch_benchmark.py --runs 5

# Different model
python mlx_batch_benchmark.py --model mlx-community/Llama-3.2-3B-Instruct-4bit

Options:

Flag	Default	Description
--model	openai/gpt-oss-20b	Model path or HuggingFace repo
--batch-sizes	1,2,4,8,16	Comma-separated batch sizes to test
--max-tokens	100	Maximum tokens per prompt
--runs	3	Number of runs per batch size
--long-stories	false	Use unique 1000-word story prompts

Async Generation Patterns

MLX uses Metal which doesn't support true multi-threaded inference. However, several patterns are available:

# Streaming demo (token-by-token, interruptible)
python mlx_async_generate.py stream

# Queue-based server (on-demand submission, sequential processing)
python mlx_async_generate.py queue

# Interruptible generation demo
python mlx_async_generate.py interrupt

Queue Server Pattern

from mlx_async_generate import GenerationServer, load_model

load_model("openai/gpt-oss-20b")

server = GenerationServer()
server.start()

# Submit prompts anytime - non-blocking
def on_complete(result):
    print(f"Done in {result['time']:.2f}s: {result['response'][:100]}")

server.submit("Write a haiku about coding", max_tokens=100, callback=on_complete)
server.submit("Explain quantum computing", max_tokens=200, callback=on_complete)

# Results arrive via callback as each completes
# Or poll: result = server.get_result()

server.stop()

Streaming with Early Stop

from mlx_async_generate import StreamingServer, load_model

load_model()
streaming = StreamingServer()

# Can be stopped anytime
result = streaming.generate(
    "Write a long story...",
    max_tokens=1000,
    on_token=lambda t: print(t, end="", flush=True)
)

# From another thread: streaming.stop()

Architecture Insights

How MLX Batch Generation Works

1. Batched KV Cache: All prompts share the same forward pass, with separate KV caches
2. Padded Sequences: Shorter sequences are padded to match the longest
3. Single Metal Command Buffer: All operations execute in one GPU pipeline

Trade-offs

Approach	Total Throughput	Latency to First Result	Use Case
Batch (large)	High (295 t/s)	High (wait for slowest)	Background processing
Batch (small)	Medium (190 t/s)	Medium	Balanced
Sequential	Low (98 t/s)	Low (immediate)	Interactive
Streaming	Low (98 t/s)	Lowest (token-by-token)	Chat/UI

Memory Behavior

• Model weights: ~13.76 GB (constant)
• KV cache scales with: batch_size × sequence_length × num_layers × hidden_dim
• Memory stays flat during generation (MLX pre-allocates)

Benchmark Results Interpretation

Metrics Explained

• Gen TPS: Total tokens generated per second across all prompts
• Per-Prompt TPS: Effective speed each individual prompt experiences
• Prompt TPS: Speed of processing input prompts (prefill phase)
• Speedup: Improvement over single-prompt baseline

Example Output

======================================================================
MLX Batch Generation Benchmark
======================================================================
Model: openai/gpt-oss-20b
Max tokens per prompt: 100
Runs per batch size: 2
Batch sizes to test: [1, 2, 4, 8]
Long story mode: False
======================================================================

Loading model...
Model loaded in 2.50s
Active memory: 13.76 GB
Peak memory: 14.61 GB

======================================================================
Running benchmarks...
======================================================================

Testing batch size: 1
  Generation TPS: 97.45 tokens/sec (total)
  Per-prompt TPS: 97.45 tokens/sec
  Prompt processing: 600.27 tokens/sec
  Avg generation time: 1.03s
  Avg tokens generated: 100
  Current memory: 13.76 GB

Testing batch size: 8
  Generation TPS: 273.02 tokens/sec (total)
  Per-prompt TPS: 34.13 tokens/sec
  Prompt processing: 1386.10 tokens/sec
  Avg generation time: 2.93s
  Avg tokens generated: 800
  Current memory: 13.76 GB

======================================================================
RESULTS SUMMARY
======================================================================
Batch    Gen TPS      Per-Prompt   Prompt TPS   Speedup    Time(s)
----------------------------------------------------------------------
1        97.45        97.45        600.27       1.00x      1.03
8        273.02       34.13        1386.10      2.80x      2.93
======================================================================

Optimal batch size: 8
Peak throughput: 273.02 tokens/sec
Speedup over single: 2.80x

Recommendations

For Maximum Throughput

• Use batch_generate with batch size 16-24 for short generations
• Use batch size 8 for long generations (1000+ words)

For Lowest Latency

• Use stream_generate for real-time token streaming
• Use sequential processing via GenerationServer

For Production Services

• Implement request queuing with the GenerationServer pattern
• Consider dynamic batching based on incoming request rate
• Monitor memory with mx.get_active_memory()

Files

File	Description
mlx_batch_benchmark.py	Throughput benchmarking script
mlx_async_generate.py	Async patterns (queue, streaming, interrupt)
README.md	This documentation

Environment

• Hardware: Apple Silicon (M1/M2/M3/M4)
• Framework: MLX 0.30.1, mlx-lm 0.30.0
• Model: GPT-OSS 20B (MoE, 3.6B active parameters)

References

• MLX-LM GitHub
• MLX Documentation
• GPT-OSS 20B on HuggingFace

License

MIT