examples.test_batch_sync

Test generate_batch_sync() performance.

View the complete module source at #L1-L105.

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.test_batch_sync

Test generate_batch_sync() performance.

examples.test_batch_sync.main

main()

Source code in examples/test_batch_sync.py

def main():
    from mlx_lm import load
    from vllm_mlx import EngineCore, EngineConfig, SamplingParams, SchedulerConfig

    MODEL = "mlx-community/Qwen3-0.6B-8bit"
    print(f"Loading {MODEL}...")
    model, tokenizer = load(MODEL)

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

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

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

    print("\n" + "=" * 70)
    print("BATCH SIZE SCALING TEST: generate_batch_sync()")
    print("=" * 70)
    print(f"{'Batch':>6} | {'Time':>8} | {'Tokens':>7} | {'Tok/s':>8} | {'% README':>8}")
    print("-" * 70)

    for multiplier in [1, 2, 4, 8, 16]:
        # Create fresh engine for each test to avoid cache state issues
        config = EngineConfig(
            scheduler_config=SchedulerConfig(
                max_num_seqs=256,
                prefill_batch_size=8,
                completion_batch_size=32,
            )
        )
        engine = EngineCore(model, tokenizer, config)

        prompts = base_prompts * multiplier
        formatted = [format_prompt(p) for p in prompts]

        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)
        throughput = total_tokens / elapsed
        pct = throughput / 1003.7 * 100

        print(f"{len(prompts):>6} | {elapsed:>7.2f}s | {total_tokens:>7} | {throughput:>7.1f} | {pct:>7.1f}%")

    print("-" * 70)
    print(f"README benchmark: 1003.7 tok/s (5 prompts, 50 max_tokens)")

    # Async comparison
    print("\n" + "=" * 70)
    print("ASYNC generate() COMPARISON (5 prompts)")
    print("=" * 70)

    async def run_async():
        config = EngineConfig(
            scheduler_config=SchedulerConfig(
                max_num_seqs=256,
                prefill_batch_size=8,
                completion_batch_size=32,
            )
        )
        engine = EngineCore(model, tokenizer, config)
        await engine.start()

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

            start = time.perf_counter()
            tasks = [engine.generate(p, params) for p in formatted]
            outputs = await asyncio.gather(*tasks)
            elapsed = time.perf_counter() - start

            total_tokens = sum(r.completion_tokens for r in outputs)
            return total_tokens, elapsed
        finally:
            await engine.stop()

    tokens, elapsed = asyncio.run(run_async())
    throughput = tokens / elapsed
    pct = throughput / 1003.7 * 100

    print(f"Tokens: {tokens}")
    print(f"Time: {elapsed:.2f}s")
    print(f"Throughput: {throughput:.1f} tok/s ({pct:.1f}% of README)")

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.test_batch_sync.main · function

examples.test_batch_sync.main() -> not annotated

Function main calls print, load, SamplingParams, EngineConfig.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Function main calls print, load, SamplingParams, EngineConfig. No direct raise statement appears in this definition.

View source #L8-L101.

examples.test_batch_sync.main.format_prompt · nested function

examples.test_batch_sync.main.format_prompt(p) -> not annotated

Nested Function main.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 main.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 #L24-L29.

examples.test_batch_sync.main.run_async · nested function

async examples.test_batch_sync.main.run_async() -> not annotated

Nested Function main.run_async calls EngineConfig, SchedulerConfig, EngineCore, engine.start; awaits asynchronous work; returns (total_tokens, elapsed).

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: (total_tokens, elapsed)

Exceptions and behavior

Nested Function main.run_async calls EngineConfig, SchedulerConfig, EngineCore, engine.start; awaits asynchronous work; returns (total_tokens, elapsed). No direct raise statement appears in this definition.

View source #L71-L93.

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
main	function	main() -> not annotated	Function main calls print, load, SamplingParams, EngineConfig.	#L8-L101
main.format_prompt	nested function	main.format_prompt(p) -> not annotated	Nested Function main.format_prompt calls tokenizer.apply_chat_template; returns tokenizer.apply_chat_template([{'role': 'user', 'content': p}], tokenize=False, add_generation_prompt=True).	#L24-L29
main.run_async	nested function	async main.run_async() -> not annotated	Nested Function main.run_async calls EngineConfig, SchedulerConfig, EngineCore, engine.start; awaits asynchronous work; returns (total_tokens, elapsed).	#L71-L93