examples.test_batching

Example: Test continuous batching with vllm-mlx.

View the complete module source at #L1-L195.

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_batching

Example: Test continuous batching with vllm-mlx.

This script demonstrates the continuous batching capability by sending multiple concurrent requests and measuring throughput.

Usage

python examples/test_batching.py python examples/test_batching.py --model mlx-community/Qwen2.5-3B-Instruct-4bit python examples/test_batching.py --num-requests 10 --max-tokens 50

examples.test_batching.run_single_request async

run_single_request(engine: AsyncEngineCore, request_id: str, prompt: str, sampling_params: SamplingParams) -> dict

Run a single request and collect timing.

Source code in examples/test_batching.py

async def run_single_request(
    engine: AsyncEngineCore,
    request_id: str,
    prompt: str,
    sampling_params: SamplingParams,
) -> dict:
    """Run a single request and collect timing."""
    start = time.perf_counter()

    await engine.add_request(
        prompt=prompt,
        sampling_params=sampling_params,
        request_id=request_id,
    )

    tokens = []
    ttft = None  # Time to first token

    async for output in engine.stream_outputs(request_id):
        if ttft is None:
            ttft = time.perf_counter() - start
        tokens.extend(output.new_token_ids)

        if output.finished:
            break

    total_time = time.perf_counter() - start

    return {
        "request_id": request_id,
        "prompt_length": len(prompt.split()),
        "num_tokens": len(tokens),
        "ttft": ttft,
        "total_time": total_time,
        "tokens_per_second": len(tokens) / total_time if total_time > 0 else 0,
        "output_text": output.output_text if output else "",
    }

examples.test_batching.run_concurrent_requests async

run_concurrent_requests(engine: AsyncEngineCore, prompts: List[str], sampling_params: SamplingParams) -> List[dict]

Run multiple requests concurrently.

Source code in examples/test_batching.py

async def run_concurrent_requests(
    engine: AsyncEngineCore,
    prompts: List[str],
    sampling_params: SamplingParams,
) -> List[dict]:
    """Run multiple requests concurrently."""
    tasks = []
    for i, prompt in enumerate(prompts):
        task = run_single_request(engine, f"req-{i}", prompt, sampling_params)
        tasks.append(task)

    return await asyncio.gather(*tasks)

examples.test_batching.print_results

print_results(results: List[dict], total_time: float)

Print benchmark results.

Source code in examples/test_batching.py

def print_results(results: List[dict], total_time: float):
    """Print benchmark results."""
    print("\n" + "=" * 60)
    print("CONTINUOUS BATCHING BENCHMARK RESULTS")
    print("=" * 60)

    total_tokens = sum(r["num_tokens"] for r in results)
    avg_ttft = sum(r["ttft"] for r in results if r["ttft"]) / len(results)
    avg_time = sum(r["total_time"] for r in results) / len(results)

    print(f"\nRequests: {len(results)}")
    print(f"Total tokens generated: {total_tokens}")
    print(f"Total wall time: {total_time:.2f}s")
    print(f"Throughput: {total_tokens / total_time:.1f} tokens/s")
    print(f"Avg time to first token: {avg_ttft * 1000:.1f}ms")
    print(f"Avg request time: {avg_time:.2f}s")

    print("\n" + "-" * 60)
    print("Per-request details:")
    print("-" * 60)
    for r in results:
        print(
            f"  {r['request_id']}: {r['num_tokens']} tokens, "
            f"{r['total_time']:.2f}s, {r['tokens_per_second']:.1f} tok/s"
        )

    print("\n" + "-" * 60)
    print("Sample outputs:")
    print("-" * 60)
    for r in results[:3]:  # Show first 3
        text = r["output_text"][:100] + "..." if len(r["output_text"]) > 100 else r["output_text"]
        print(f"  {r['request_id']}: {text}")

examples.test_batching.main async

main()

Source code in examples/test_batching.py

async def main():
    parser = argparse.ArgumentParser(description="Test continuous batching")
    parser.add_argument(
        "--model",
        type=str,
        default="mlx-community/Llama-3.2-1B-Instruct-4bit",
        help="Model to use",
    )
    parser.add_argument(
        "--num-requests",
        type=int,
        default=5,
        help="Number of concurrent requests",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=30,
        help="Max tokens per request",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Sampling temperature",
    )
    args = parser.parse_args()

    print(f"Loading model: {args.model}")
    from mlx_lm import load

    model, tokenizer = load(args.model)

    # Configure scheduler for batching
    scheduler_config = SchedulerConfig(
        max_num_seqs=32,
        prefill_batch_size=8,
        completion_batch_size=16,
    )

    engine_config = EngineConfig(
        model_name=args.model,
        scheduler_config=scheduler_config,
    )

    # Sample prompts
    prompts = [
        "What is the capital of France?",
        "Explain quantum computing in simple terms.",
        "Write a haiku about programming.",
        "What are the benefits of exercise?",
        "How do airplanes fly?",
        "What is machine learning?",
        "Describe the water cycle.",
        "What is photosynthesis?",
        "Explain gravity simply.",
        "What causes rain?",
    ][: args.num_requests]

    sampling_params = SamplingParams(
        max_tokens=args.max_tokens,
        temperature=args.temperature,
    )

    print(f"\nRunning {len(prompts)} concurrent requests...")
    print(f"Max tokens: {args.max_tokens}, Temperature: {args.temperature}")

    async with AsyncEngineCore(model, tokenizer, engine_config) as engine:
        # Let engine warm up
        await asyncio.sleep(0.1)

        start = time.perf_counter()
        results = await run_concurrent_requests(engine, prompts, sampling_params)
        total_time = time.perf_counter() - start

    print_results(results, total_time)

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_batching.run_single_request · function

async examples.test_batching.run_single_request(engine: AsyncEngineCore, request_id: str, prompt: str, sampling_params: SamplingParams) -> dict

Run a single request and collect timing.

Parameters

Name	Type	Required	Default	Description
engine	AsyncEngineCore	yes	none	Required positional or keyword input.
request_id	str	yes	none	Required positional or keyword input.
prompt	str	yes	none	Required positional or keyword input.
sampling_params	SamplingParams	yes	none	Required positional or keyword input.

Returns

• Type: dict
• Direct return expressions: {'request_id': request_id, 'prompt_length': len(prompt.split()), 'num_tokens': len(tokens), 'ttft': ttft, 'total_time':…

Exceptions and behavior

Function run_single_request calls time.perf_counter, engine.add_request, engine.stream_outputs, tokens.extend; awaits asynchronous work; returns {'request_id': request_id, 'prompt_length': len(prompt.split()), 'num_tokens': len(tokens), 'ttft': ttft, 'total_time':…. No direct raise statement appears in this definition.

View source #L29-L65.

examples.test_batching.run_concurrent_requests · function

async examples.test_batching.run_concurrent_requests(engine: AsyncEngineCore, prompts: List[str], sampling_params: SamplingParams) -> List[dict]

Run multiple requests concurrently.

Parameters

Name	Type	Required	Default	Description
engine	AsyncEngineCore	yes	none	Required positional or keyword input.
prompts	List[str]	yes	none	Required positional or keyword input.
sampling_params	SamplingParams	yes	none	Required positional or keyword input.

Returns

• Type: List[dict]
• Direct return expressions: await asyncio.gather(*tasks)

Exceptions and behavior

Function run_concurrent_requests calls enumerate, run_single_request, tasks.append, asyncio.gather; awaits asynchronous work; returns await asyncio.gather(*tasks). No direct raise statement appears in this definition.

View source #L68-L79.

examples.test_batching.print_results · function

examples.test_batching.print_results(results: List[dict], total_time: float) -> not annotated

Print benchmark results.

Parameters

Name	Type	Required	Default	Description
results	List[dict]	yes	none	Required positional or keyword input.
total_time	float	yes	none	Required positional or keyword input.

Returns

• Type: not annotated

Exceptions and behavior

Function print_results calls print, sum, len. No direct raise statement appears in this definition.

View source #L82-L113.

examples.test_batching.main · function

async examples.test_batching.main() -> not annotated

Function main calls argparse.ArgumentParser, parser.add_argument, parser.parse_args, print; awaits asynchronous work.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Function main calls argparse.ArgumentParser, parser.add_argument, parser.parse_args, print; awaits asynchronous work. No direct raise statement appears in this definition.

View source #L116-L191.

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
run_single_request	function	async run_single_request(engine: AsyncEngineCore, request_id: str, prompt: str, sampling_params: SamplingParams) -> dict	Run a single request and collect timing.	#L29-L65
run_concurrent_requests	function	async run_concurrent_requests(engine: AsyncEngineCore, prompts: List[str], sampling_params: SamplingParams) -> List[dict]	Run multiple requests concurrently.	#L68-L79
print_results	function	print_results(results: List[dict], total_time: float) -> not annotated	Print benchmark results.	#L82-L113
main	function	async main() -> not annotated	Function main calls argparse.ArgumentParser, parser.add_argument, parser.parse_args, print; awaits asynchronous work.	#L116-L191