examples.benchmark_all_models

Benchmark all text models for README.

View the complete module source at #L1-L148.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

examples.benchmark_all_models

Benchmark all text models for README.

examples.benchmark_all_models.benchmark_model

benchmark_model(model_name: str)

Benchmark a single model and return results.

Source code in examples/benchmark_all_models.py

def benchmark_model(model_name: str):
    """Benchmark a single model and return results."""
    from mlx_lm import load
    from vllm_mlx import EngineCore, EngineConfig, SamplingParams, SchedulerConfig

    base_prompts = [
        "What is 2+2?",
        "Name 3 colors.",
        "What is Python?",
        "Capital of Japan?",
        "Who wrote Hamlet?",
    ]

    params = SamplingParams(max_tokens=50, temperature=0.7)

    print(f"\n{'='*60}")
    print(f"Benchmarking: {model_name}")
    print("="*60)

    print("Loading model...")
    model, tokenizer = load(model_name)

    def format_prompt(p):
        return tokenizer.apply_chat_template(
            [{"role": "user", "content": p}],
            tokenize=False,
            add_generation_prompt=True,
        )

    formatted = [format_prompt(p) for p in base_prompts]

    # Create engine with explicit ID for tracking
    config = EngineConfig(
        scheduler_config=SchedulerConfig(
            max_num_seqs=256,
            prefill_batch_size=8,
            completion_batch_size=32,
        )
    )
    engine = EngineCore(model, tokenizer, config)

    try:
        # Test 1: Single request throughput
        print("\n1. Single request throughput...")
        single_times = []
        single_tokens = []
        for p in formatted[:3]:
            start = time.perf_counter()
            result = engine.generate_batch_sync([p], params)[0]
            elapsed = time.perf_counter() - start
            single_times.append(elapsed)
            single_tokens.append(result.completion_tokens)

        single_tps = sum(single_tokens) / sum(single_times)
        print(f"   Single: {single_tps:.1f} tok/s")

        # Test 2: Batch throughput (5 concurrent)
        print("2. Batch throughput (5 concurrent)...")
        engine.scheduler.reset()

        # Warmup
        _ = engine.generate_batch_sync(formatted[:1], params)

        # Reset for clean measurement
        engine.scheduler.reset()

        start = time.perf_counter()
        results = engine.generate_batch_sync(formatted, params)
        elapsed = time.perf_counter() - start

        total_tokens = sum(r.completion_tokens for r in results)
        batch_tps = total_tokens / elapsed
        print(f"   Batch:  {batch_tps:.1f} tok/s")

        speedup = batch_tps / single_tps

        # Test 3: Speed measurement
        print("3. Generation speed...")
        engine.scheduler.reset()

        start = time.perf_counter()
        result = engine.generate_batch_sync([formatted[0]], SamplingParams(max_tokens=30, temperature=0.0))[0]
        elapsed = time.perf_counter() - start

        ttft_ms = elapsed / result.completion_tokens * 1000 if result.completion_tokens > 0 else 0
        gen_tps = result.completion_tokens / elapsed if elapsed > 0 else 0

        print(f"   TTFT:   ~{ttft_ms:.1f}ms (estimated)")
        print(f"   Speed:  {gen_tps:.1f} tok/s")

        return {
            "model": model_name.split("/")[-1],
            "single_tps": single_tps,
            "batch_tps": batch_tps,
            "speedup": speedup,
            "ttft_ms": ttft_ms,
            "gen_tps": gen_tps,
        }
    finally:
        # Always close the engine to release model ownership
        engine.close()

examples.benchmark_all_models.main

main()

Source code in examples/benchmark_all_models.py

def main():
    models = [
        "mlx-community/Llama-3.2-1B-Instruct-4bit",
        "mlx-community/Llama-3.2-3B-Instruct-4bit",
        "mlx-community/Qwen3-0.6B-8bit",
        "mlx-community/Qwen3-30B-A3B-4bit",
        "mlx-community/Qwen2.5-1.5B-Instruct-4bit",
    ]

    results = []
    for model_name in models:
        try:
            result = benchmark_model(model_name)
            results.append(result)
        except Exception as e:
            print(f"Error benchmarking {model_name}: {e}")
            import traceback
            traceback.print_exc()

    # Print summary
    print("\n" + "="*80)
    print("BENCHMARK RESULTS SUMMARY")
    print("="*80)

    print("\n### Continuous Batching Results\n")
    print("| Model | Single | Batch (5 req) | Speedup |")
    print("|-------|--------|---------------|---------|")
    for r in results:
        print(f"| {r['model']} | {r['single_tps']:.1f} tok/s | {r['batch_tps']:.1f} tok/s | **{r['speedup']:.2f}x** |")

    print("\n### Generation Speed\n")
    print("| Model | TTFT | Speed |")
    print("|-------|------|-------|")
    for r in results:
        print(f"| {r['model']} | ~{r['ttft_ms']:.1f}ms | {r['gen_tps']:.1f} tok/s |")

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

examples.benchmark_all_models.benchmark_model · function

examples.benchmark_all_models.benchmark_model(model_name: str) -> not annotated

Benchmark a single model and return results.

Parameters

Name	Type	Required	Default	Description
model_name	str	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: {'model': model_name.split('/')[-1], 'single_tps': single_tps, 'batch_tps': batch_tps, 'speedup': speedup, 'ttft_ms': t…

Exceptions and behavior

Function benchmark_model calls SamplingParams, print, load, format_prompt; returns {'model': model_name.split('/')[-1], 'single_tps': single_tps, 'batch_tps': batch_tps, 'speedup': speedup, 'ttft_ms': t…. No direct raise statement appears in this definition.

View source #L7-L107.

examples.benchmark_all_models.benchmark_model.format_prompt · nested function

examples.benchmark_all_models.benchmark_model.format_prompt(p) -> not annotated

Nested Function benchmark_model.format_prompt calls tokenizer.apply_chat_template; returns tokenizer.apply_chat_template([{'role': 'user', 'content': p}], tokenize=False, add_generation_prompt=True).

Parameters

Name	Type	Required	Default	Description
p	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: tokenizer.apply_chat_template([{'role': 'user', 'content': p}], tokenize=False, add_generation_prompt=True)

Exceptions and behavior

Nested Function benchmark_model.format_prompt calls tokenizer.apply_chat_template; returns tokenizer.apply_chat_template([{'role': 'user', 'content': p}], tokenize=False, add_generation_prompt=True). No direct raise statement appears in this definition.

View source #L29-L34.

examples.benchmark_all_models.main · function

examples.benchmark_all_models.main() -> not annotated

Function main calls benchmark_model, results.append, print, traceback.print_exc.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Function main calls benchmark_model, results.append, print, traceback.print_exc. No direct raise statement appears in this definition.

View source #L110-L144.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
benchmark_model	function	benchmark_model(model_name: str) -> not annotated	Benchmark a single model and return results.	#L7-L107
benchmark_model.format_prompt	nested function	benchmark_model.format_prompt(p) -> not annotated	Nested Function benchmark_model.format_prompt calls tokenizer.apply_chat_template; returns tokenizer.apply_chat_template([{'role': 'user', 'content': p}], tokenize=False, add_generation_prompt=True).	#L29-L34
main	function	main() -> not annotated	Function main calls benchmark_model, results.append, print, traceback.print_exc.	#L110-L144