How to Run Performance Benchmarks with Prefill and Generation Speed Tests in oMLX

oMLX provides a built-in benchmarking utility that measures time-to-first-token (TTFT), generation throughput, and continuous batching performance via the BenchmarkRequest API in omlx/admin/benchmark.py.

oMLX ships with a comprehensive benchmarking system designed to evaluate large language model performance across single-request and concurrent scenarios. The benchmark utility is implemented in the admin module and measures both prefill processing speed and token generation throughput while streaming real-time results to clients.

Understanding the oMLX Benchmark Architecture

The benchmark system in [omlx/admin/benchmark.py](https://github.com/jundot/omlx/blob/main/omlx/admin/benchmark.py) executes two distinct test types:

• Single-request tests measure TTFT, generation tokens per second (TPS), prompt processing TPS, and peak memory usage for individual requests at varying prefill lengths.
• Continuous-batching tests measure prompt throughput (pp TPS) and total throughput (tg TPS) when handling concurrent requests at different batch sizes.

A benchmark run follows a strict five-phase execution flow defined in the run_benchmark function:

1. Unload existing models from memory to ensure clean state.
2. Load the target model specified in the request.
3. Warm-up with a minimal prompt to trigger JIT/METAL compilation and prevent timing pollution.
4. Execute single-request tests for each specified prompt_lengths value.
5. Execute batch tests (if the engine supports a scheduler core) using the specified batch_sizes.

Configuring Benchmark Parameters

Valid configurations are constrained by constants defined in omlx/admin/benchmark.py. The VALID_PROMPT_LENGTHS array (lines 32-34) accepts standard values like 1024, 4096, and 8192 tokens. For continuous batching, supported batch_sizes include 2, 4, and 8 (lines 35-36).

The BenchmarkRequest model (lines 38-46) accepts these key fields:

• model_id: The model identifier to benchmark.
• prompt_lengths: List of prefill sizes in tokens for single-request evaluation.
• generation_length: Number of tokens to generate per test (default 128).
• batch_sizes: List of concurrent request counts for throughput testing.

Running Benchmarks via the Admin UI

The web interface provides the simplest way to execute performance tests:

1. Navigate to http://localhost:8080/admin and click the Benchmark tab.
2. Select your target model from the dropdown.
3. Choose one or more Prompt Lengths (e.g., 1024, 8192).
4. Optionally specify Batch Sizes (e.g., 2, 4, 8) for concurrent testing.
5. Click Start Benchmark to initiate the five-phase execution.

The UI displays a live progress log showing model loading, warm-up completion, and per-test metrics including ttft_ms, gen_tps, processing_tps, e2e_latency_s, and peak_memory_bytes. Results export as JSON via the Export JSON button.

Running Benchmarks via the HTTP API

For programmatic access, send a POST request to the benchmark endpoint registered in [omlx/admin/routes.py](https://github.com/jundot/omlx/blob/main/omlx/admin/routes.py) (around line 4742):

curl -X POST http://localhost:8080/admin/api/benchmark \
  -H "Authorization: Bearer $OMLX_API_KEY" \
  -H "Content-Type: application/json" \
  -d '{
    "model_id": "Qwen3-30B-A3B-4bit",
    "prompt_lengths": [1024, 8192],
    "generation_length": 128,
    "batch_sizes": [2, 4, 8]
  }'

The server returns a benchmark ID immediately. Monitor live progress via Server-Sent Events (SSE):

curl -N http://localhost:8080/admin/api/benchmark/events?bench_id=bench-abcdef123456

The SSE stream emits JSON events with type: "progress" updates and type: "result" payloads containing the final metrics.

Core Implementation Details

Prompt Generation Logic

The _generate_prompt function constructs deterministic test inputs of exact token lengths:

# Simplified from omlx/admin/benchmark.py

def _generate_prompt(tokenizer, target_tokens):
    unique_prefix = f"BENCH-{uuid.uuid4().hex} "
    filler = ("The quick brown fox jumps over the lazy dog. "
              "In the realm of artificial intelligence, large language models "
              "have demonstrated remarkable capabilities across diverse tasks. ")
    text = unique_prefix + filler * (target_tokens // 10 + 1)
    tokens = tokenizer.encode(text)[:target_tokens]
    return tokenizer.decode(tokens)

Single-Request Test Timing

The _run_single_test function (lines 77-84) captures precise latency metrics:

async def _run_single_test(engine, prompt, max_tokens, pp_len):
    mx.reset_peak_memory()          # Reset MLX memory tracker

    start = time.perf_counter()
    first_token_time = None
    
    async for out in engine.stream_generate(
        prompt=prompt,
        max_tokens=max_tokens,
        temperature=0.0,
        top_p=1.0
    ):
        if first_token_time is None and out.completion_tokens > 0:
            first_token_time = time.perf_counter()
        last = out
    
    end = time.perf_counter()
    return _compute_single_metrics(
        prompt_tokens=last.prompt_tokens,
        completion_tokens=last.completion_tokens,
        start_time=start,
        first_token_time=first_token_time or end,
        end_time=end,
        peak_memory=mx.get_peak_memory(),
        cached_tokens=last.cached_tokens
    )

Continuous Batching Test

The _run_batch_test function (lines 41-71) submits concurrent requests and aggregates throughput statistics:

async def _run_batch_test(engine, prompts, prompt_tokens, max_tokens, batch_size):
    # Submit batch_size requests concurrently

    results = await asyncio.gather(*[
        _single_request(prompts[i]) for i in range(batch_size)
    ])
    # Compute pp TPS (prompt processing) and tg TPS (token generation)

    # Aggregated per-request TTFT and memory metrics

Important Considerations

Warm-up is automatic. The benchmark includes an internal warm-up phase (lines 56-63) that triggers JIT compilation, eliminating the need for manual dummy requests before timing.

Batch test requirements. Continuous-batching tests execute only if the engine exposes a scheduler core. DFlash-only engines skip this phase with a log notification.

Experimental feature restrictions. If experimental features like DFlash, SpecPrefill, or TurboQuant KV are active, results will not upload to the public community benchmark server (see _upload_to_omlx_ai, lines 1-4).

Summary

• Use omlx/admin/benchmark.py to access the core BenchmarkRequest model and run_benchmark orchestration logic.
• Configure prompt_lengths using values from VALID_PROMPT_LENGTHS (e.g., 1024, 4096, 8192) and batch_sizes of 2, 4, or 8.
• Execute via UI at /admin or programmatically via POST to /admin/api/benchmark with SSE streaming for real-time results.
• Automatic warm-up ensures accurate timing by triggering compilation before measurement begins.
• Memory tracking uses mx.reset_peak_memory() and mx.get_peak_memory() for precise peak consumption reporting.

Frequently Asked Questions

What metrics does the oMLX benchmark return?

The benchmark returns time-to-first-token (ttft_ms), generation throughput (gen_tps), prompt processing throughput (processing_tps), end-to-end latency (e2e_latency_s), and peak memory consumption (peak_memory_bytes). Batch tests additionally report prompt throughput (pp TPS) and total generation throughput (tg TPS).

Why does the benchmark include a warm-up phase?

The warm-up phase (implemented in run_benchmark at lines 56-63) executes a minimal prompt to trigger MLX JIT or Metal shader compilation. This prevents the first real test from being polluted by compilation overhead, ensuring accurate measurements of actual inference performance.

Can I run batch tests on any oMLX engine?

No. Continuous-batching tests require the engine to expose a scheduler core. If you are using a DFlash-only engine or an implementation without scheduling support, the benchmark skips batch testing (logging this decision at lines 13-14) and completes only single-request evaluations.

How do I interpret the continuous-batching results?

Prompt throughput (pp TPS) measures how many prompt tokens the system processes per second across all concurrent requests, while total throughput (tg TPS) measures generated tokens per second. Higher pp TPS indicates efficient parallel prefill processing, while higher tg TPS indicates faster concurrent generation—critical metrics for production serving scenarios.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}