vllm_mlx.model_runner

MLX Model Runner for vLLM.

View the complete module source at #L1-L476.

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.

vllm_mlx.model_runner

MLX Model Runner for vLLM.

This module implements the model runner that bridges vLLM's request handling with mlx-lm's inference capabilities.

Includes low-level optimizations: - mx.compile() for kernel fusion - Memory bandwidth optimization - Prefill chunking for L2 cache efficiency

vllm_mlx.model_runner.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.model_runner.SamplerOutput dataclass

SamplerOutput(token_ids: list[int], logprobs: list[dict] | None = None)

Output from sampling.

vllm_mlx.model_runner.SamplerOutput.token_ids instance-attribute

token_ids: list[int]

vllm_mlx.model_runner.SamplerOutput.logprobs class-attribute instance-attribute

logprobs: list[dict] | None = None

vllm_mlx.model_runner.MLXModelRunnerOutput dataclass

MLXModelRunnerOutput(req_id_to_token_ids: dict[str, list[int]], req_id_to_logprobs: dict[str, list[dict]] | None = None, num_tokens_generated: int = 0, generation_time_s: float = 0.0)

Output from MLX model runner, compatible with vLLM's ModelRunnerOutput.

vllm_mlx.model_runner.MLXModelRunnerOutput.req_id_to_token_ids instance-attribute

req_id_to_token_ids: dict[str, list[int]]

vllm_mlx.model_runner.MLXModelRunnerOutput.req_id_to_logprobs class-attribute instance-attribute

req_id_to_logprobs: dict[str, list[dict]] | None = None

vllm_mlx.model_runner.MLXModelRunnerOutput.num_tokens_generated class-attribute instance-attribute

num_tokens_generated: int = 0

vllm_mlx.model_runner.MLXModelRunnerOutput.generation_time_s class-attribute instance-attribute

generation_time_s: float = 0.0

vllm_mlx.model_runner.MLXModelRunner

MLXModelRunner(vllm_config: VllmConfig, enable_optimizations: bool = True)

Model runner that uses mlx-lm for inference.

This class handles: - Model loading via mlx-lm - Converting vLLM requests to mlx-lm format - Running inference and returning results in vLLM format - KV cache management (delegated to mlx-lm)

Optimizations: - mx.compile() for kernel fusion (fuses multiple ops into single Metal kernel) - Memory optimization for bandwidth efficiency - Prefill chunking for L2 cache utilization

Initialize MLX model runner.

Parameters:

• vllm_config (VllmConfig) –

  vLLM configuration

• enable_optimizations (bool, default: True ) –

  Whether to enable low-level optimizations

Source code in vllm_mlx/model_runner.py

def __init__(self, vllm_config: "VllmConfig", enable_optimizations: bool = True):
    """
    Initialize MLX model runner.

    Args:
        vllm_config: vLLM configuration
        enable_optimizations: Whether to enable low-level optimizations
    """
    self.vllm_config = vllm_config
    self.model_config = vllm_config.model_config
    self.cache_config = vllm_config.cache_config
    self.scheduler_config = vllm_config.scheduler_config

    # mlx-lm model and tokenizer
    self.model = None
    self.tokenizer = None
    self._loaded = False

    # Sampler for generation
    self._sampler = None

    # Cache for prompt processing
    self._prompt_cache = None

    # KV cache blocks
    self._num_cache_blocks = 0

    # Optimization settings
    self._enable_optimizations = enable_optimizations
    self._compiled_forward = None  # Compiled model forward pass
    self._hardware_info = None  # Detected hardware profile

    logger.info(f"MLXModelRunner initialized for model: {self.model_config.model}")
    logger.info(
        f"Low-level optimizations: {'ENABLED' if enable_optimizations else 'disabled'}"
    )

vllm_mlx.model_runner.MLXModelRunner.vllm_config instance-attribute

vllm_config = vllm_config

vllm_mlx.model_runner.MLXModelRunner.model_config instance-attribute

model_config = vllm_config.model_config

vllm_mlx.model_runner.MLXModelRunner.cache_config instance-attribute

cache_config = vllm_config.cache_config

vllm_mlx.model_runner.MLXModelRunner.scheduler_config instance-attribute

scheduler_config = vllm_config.scheduler_config

vllm_mlx.model_runner.MLXModelRunner.model instance-attribute

model = None

vllm_mlx.model_runner.MLXModelRunner.tokenizer instance-attribute

tokenizer = None

vllm_mlx.model_runner.MLXModelRunner._loaded instance-attribute

_loaded = False

vllm_mlx.model_runner.MLXModelRunner._sampler instance-attribute

_sampler = None

vllm_mlx.model_runner.MLXModelRunner._prompt_cache instance-attribute

_prompt_cache = None

vllm_mlx.model_runner.MLXModelRunner._num_cache_blocks instance-attribute

_num_cache_blocks = 0

vllm_mlx.model_runner.MLXModelRunner._enable_optimizations instance-attribute

_enable_optimizations = enable_optimizations

vllm_mlx.model_runner.MLXModelRunner._compiled_forward instance-attribute

_compiled_forward = None

vllm_mlx.model_runner.MLXModelRunner._hardware_info instance-attribute

_hardware_info = None

vllm_mlx.model_runner.MLXModelRunner.load_model

load_model() -> None

Load model using mlx-lm with optimizations.

Source code in vllm_mlx/model_runner.py

def load_model(self) -> None:
    """Load model using mlx-lm with optimizations."""
    if self._loaded:
        return

    try:
        from mlx_lm import load

        model_name = self.model_config.model

        logger.info(f"Loading model with mlx-lm: {model_name}")
        start_time = time.time()

        self.model, self.tokenizer = load(
            model_name,
            tokenizer_config={
                "trust_remote_code": self.model_config.trust_remote_code,
            },
        )

        load_time = time.time() - start_time
        logger.info(f"Model loaded in {load_time:.2f}s")

        self._loaded = True

        # Create default sampler
        self._create_default_sampler()

        # Apply low-level optimizations
        if self._enable_optimizations:
            self._apply_optimizations()

    except ImportError:
        raise ImportError(
            "mlx-lm is required for MLX model runner. "
            "Install with: pip install mlx-lm"
        )
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise

vllm_mlx.model_runner.MLXModelRunner._apply_optimizations

_apply_optimizations() -> None

Apply low-level optimizations for maximum performance.

Source code in vllm_mlx/model_runner.py

def _apply_optimizations(self) -> None:
    """Apply low-level optimizations for maximum performance."""
    try:
        from vllm_mlx.optimizations import (
            configure_memory_optimization,
            detect_hardware,
        )

        # Detect hardware and apply memory optimization
        self._hardware_info = detect_hardware()
        logger.info(f"Hardware detected: {self._hardware_info.chip_name}")
        logger.info(f"Memory: {self._hardware_info.total_memory_gb:.1f} GB")
        logger.info(f"Bandwidth: {self._hardware_info.memory_bandwidth_gbs} GB/s")

        # Configure memory settings
        configure_memory_optimization()

        # Compile the model forward pass for kernel fusion
        self._setup_compiled_forward()

    except Exception as e:
        logger.warning(f"Failed to apply optimizations: {e}")

vllm_mlx.model_runner.MLXModelRunner._setup_compiled_forward

_setup_compiled_forward() -> None

Setup compiled forward pass using mx.compile() for kernel fusion.

This fuses multiple operations into single Metal kernels, reducing kernel launch overhead and improving throughput.

Source code in vllm_mlx/model_runner.py

def _setup_compiled_forward(self) -> None:
    """
    Setup compiled forward pass using mx.compile() for kernel fusion.

    This fuses multiple operations into single Metal kernels,
    reducing kernel launch overhead and improving throughput.
    """
    if self.model is None:
        return

    try:
        # Compile the model's __call__ method
        # This creates fused Metal kernels for the forward pass
        if hasattr(self.model, "__call__"):
            self._compiled_forward = mx.compile(self.model.__call__)
            logger.info("Compiled forward pass enabled (mx.compile kernel fusion)")
        else:
            logger.warning(
                "Model does not have __call__ method, skipping compilation"
            )

    except Exception as e:
        logger.warning(f"Failed to compile forward pass: {e}")
        self._compiled_forward = None

vllm_mlx.model_runner.MLXModelRunner._create_default_sampler

_create_default_sampler() -> None

Create default sampler for generation.

Source code in vllm_mlx/model_runner.py

def _create_default_sampler(self) -> None:
    """Create default sampler for generation."""
    try:
        from mlx_lm.sample_utils import make_sampler

        self._sampler = make_sampler(
            temp=0.7,
            top_p=0.9,
        )
    except ImportError:
        logger.warning("Could not create sampler, using defaults")

vllm_mlx.model_runner.MLXModelRunner.initialize_cache

initialize_cache(num_blocks: int) -> None

Initialize KV cache.

Source code in vllm_mlx/model_runner.py

def initialize_cache(self, num_blocks: int) -> None:
    """Initialize KV cache."""
    self._num_cache_blocks = num_blocks
    logger.info(f"KV cache initialized with {num_blocks} blocks")

vllm_mlx.model_runner.MLXModelRunner.get_kv_cache_spec

get_kv_cache_spec() -> dict

Get KV cache specification.

Source code in vllm_mlx/model_runner.py

def get_kv_cache_spec(self) -> dict:
    """Get KV cache specification."""
    return {
        "num_blocks": self._num_cache_blocks,
        "block_size": self.cache_config.block_size,
    }

vllm_mlx.model_runner.MLXModelRunner.get_cache_block_size_bytes

get_cache_block_size_bytes() -> int

Calculate cache block size in bytes.

Source code in vllm_mlx/model_runner.py

def get_cache_block_size_bytes(self) -> int:
    """Calculate cache block size in bytes."""
    if not self._loaded or self.model is None:
        return 0

    # Get model config
    config = getattr(self.model, "config", None)
    if config is None:
        return 0

    head_size = getattr(config, "head_dim", 64)
    num_kv_heads = getattr(
        config, "num_key_value_heads", getattr(config, "num_attention_heads", 32)
    )
    num_layers = getattr(config, "num_hidden_layers", 32)
    block_size = self.cache_config.block_size

    # 2 for K and V, 2 bytes for float16
    return 2 * block_size * num_layers * num_kv_heads * head_size * 2

vllm_mlx.model_runner.MLXModelRunner.warm_up

warm_up() -> None

Warm up model with a test generation.

Source code in vllm_mlx/model_runner.py

def warm_up(self) -> None:
    """Warm up model with a test generation."""
    if not self._loaded:
        self.load_model()

    logger.info("Warming up model...")

    try:
        from mlx_lm import generate

        # Simple warm-up generation
        _ = generate(
            self.model,
            self.tokenizer,
            prompt="Hello",
            max_tokens=5,
            verbose=False,
        )
        logger.info("Model warm-up complete")

    except Exception as e:
        logger.warning(f"Warm-up failed (non-critical): {e}")

vllm_mlx.model_runner.MLXModelRunner.execute_model

execute_model(scheduler_output: SchedulerOutput) -> MLXModelRunnerOutput

Execute model inference for scheduled requests.

Parameters:

• scheduler_output (SchedulerOutput) –

  Contains requests to process

Returns:

• MLXModelRunnerOutput –

  MLXModelRunnerOutput with generated tokens

Source code in vllm_mlx/model_runner.py

def execute_model(
    self,
    scheduler_output: "SchedulerOutput",
) -> MLXModelRunnerOutput:
    """
    Execute model inference for scheduled requests.

    Args:
        scheduler_output: Contains requests to process

    Returns:
        MLXModelRunnerOutput with generated tokens
    """
    if not self._loaded:
        raise RuntimeError("Model not loaded. Call load_model() first.")

    start_time = time.time()
    req_id_to_token_ids: dict[str, list[int]] = {}
    total_tokens = 0

    # Process new requests
    for req_data in scheduler_output.scheduled_new_reqs:
        req_id = req_data.req_id
        prompt_token_ids = req_data.prompt_token_ids

        # Generate tokens for this request
        generated_ids = self._generate_for_request(
            prompt_token_ids=prompt_token_ids,
            sampling_params=req_data.sampling_params,
            max_tokens=1,  # Generate one token at a time for streaming
        )

        req_id_to_token_ids[req_id] = generated_ids
        total_tokens += len(generated_ids)

    # Process running requests (continue generation)
    for req_id in scheduler_output.scheduled_running_reqs:
        # For running requests, we continue generation
        # This is simplified - in practice we'd use KV cache
        generated_ids = self._continue_generation(req_id)
        if generated_ids:
            req_id_to_token_ids[req_id] = generated_ids
            total_tokens += len(generated_ids)

    generation_time = time.time() - start_time

    return MLXModelRunnerOutput(
        req_id_to_token_ids=req_id_to_token_ids,
        num_tokens_generated=total_tokens,
        generation_time_s=generation_time,
    )

vllm_mlx.model_runner.MLXModelRunner._prefill_with_chunking

_prefill_with_chunking(input_ids: array, cache: Optional[Any] = None) -> tuple[array, Any]

Process prompt with optimal chunking for L2 cache efficiency.

Long prompts are broken into chunks that fit in L2 cache, maximizing memory bandwidth utilization during prefill.

Parameters:

• input_ids (array) –

  Input token IDs [1, seq_len]

• cache (Optional[Any], default: None ) –

  Optional existing KV cache

Returns:

• tuple[array, Any] –

  Tuple of (logits, updated_cache)

Source code in vllm_mlx/model_runner.py

def _prefill_with_chunking(
    self,
    input_ids: mx.array,
    cache: Optional[Any] = None,
) -> tuple[mx.array, Any]:
    """
    Process prompt with optimal chunking for L2 cache efficiency.

    Long prompts are broken into chunks that fit in L2 cache,
    maximizing memory bandwidth utilization during prefill.

    Args:
        input_ids: Input token IDs [1, seq_len]
        cache: Optional existing KV cache

    Returns:
        Tuple of (logits, updated_cache)
    """
    try:
        from vllm_mlx.optimizations import get_optimal_prefill_size
    except ImportError:
        # Fallback if optimizations module not available
        def get_optimal_prefill_size(seq_len):
            return min(512, seq_len)

    seq_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids)
    chunk_size = get_optimal_prefill_size(seq_len)

    # Reshape if needed
    if len(input_ids.shape) == 1:
        input_ids = input_ids.reshape(1, -1)

    # Use compiled forward if available, otherwise use model directly
    forward_fn = self._compiled_forward if self._compiled_forward else self.model

    if seq_len <= chunk_size:
        # Process entire sequence at once
        return forward_fn(input_ids, cache=cache)

    # Process in chunks for large prompts
    for i in range(0, seq_len, chunk_size):
        chunk = input_ids[:, i : i + chunk_size]
        logits, cache = forward_fn(chunk, cache=cache)
        mx.eval(cache)  # Force evaluation to free intermediate memory

    return logits, cache

vllm_mlx.model_runner.MLXModelRunner._generate_for_request

_generate_for_request(prompt_token_ids: list[int], sampling_params: Any, max_tokens: int = 1) -> list[int]

Generate tokens for a single request.

Uses optimizations when enabled: - Compiled forward pass (kernel fusion) - Prefill chunking for long prompts

Parameters:

• prompt_token_ids (list[int]) –

  Input token IDs

• sampling_params (Any) –

  Sampling parameters

• max_tokens (int, default: 1 ) –

  Maximum tokens to generate

Returns:

• list[int] –

  List of generated token IDs

Source code in vllm_mlx/model_runner.py

def _generate_for_request(
    self,
    prompt_token_ids: list[int],
    sampling_params: Any,
    max_tokens: int = 1,
) -> list[int]:
    """
    Generate tokens for a single request.

    Uses optimizations when enabled:
    - Compiled forward pass (kernel fusion)
    - Prefill chunking for long prompts

    Args:
        prompt_token_ids: Input token IDs
        sampling_params: Sampling parameters
        max_tokens: Maximum tokens to generate

    Returns:
        List of generated token IDs
    """
    try:
        from mlx_lm.generate import generate_step
        from mlx_lm.sample_utils import make_sampler

        # Create sampler from sampling params
        temp = getattr(sampling_params, "temperature", 0.7)
        top_p = getattr(sampling_params, "top_p", 0.9)
        sampler = make_sampler(temp=temp, top_p=top_p)

        # Convert token IDs to MLX array
        prompt = mx.array(prompt_token_ids)

        generated_ids = []

        # Generate tokens
        for token_info in generate_step(
            prompt=prompt,
            model=self.model,
            max_tokens=max_tokens,
            sampler=sampler,
        ):
            if hasattr(token_info, "token"):
                generated_ids.append(token_info.token)
            elif isinstance(token_info, tuple) and len(token_info) > 0:
                generated_ids.append(token_info[0])

            if len(generated_ids) >= max_tokens:
                break

        return generated_ids

    except Exception as e:
        logger.error(f"Generation failed: {e}")
        return []

vllm_mlx.model_runner.MLXModelRunner._continue_generation

_continue_generation(req_id: str) -> list[int]

Continue generation for an existing request.

This is a placeholder - in a full implementation, we would use cached KV states to continue generation efficiently.

Source code in vllm_mlx/model_runner.py

def _continue_generation(self, req_id: str) -> list[int]:
    """
    Continue generation for an existing request.

    This is a placeholder - in a full implementation, we would
    use cached KV states to continue generation efficiently.
    """
    # For now, return empty - full implementation would track state
    return []

vllm_mlx.model_runner.MLXModelRunner.decode_tokens

decode_tokens(token_ids: list[int]) -> str

Decode token IDs to text.

Source code in vllm_mlx/model_runner.py

def decode_tokens(self, token_ids: list[int]) -> str:
    """Decode token IDs to text."""
    if self.tokenizer is None:
        return ""
    return self.tokenizer.decode(token_ids)

vllm_mlx.model_runner.MLXModelRunner.get_model_info

get_model_info() -> dict

Get information about the loaded model and optimizations.

Source code in vllm_mlx/model_runner.py

def get_model_info(self) -> dict:
    """Get information about the loaded model and optimizations."""
    info = {
        "loaded": self._loaded,
        "model_name": self.model_config.model,
        "optimizations_enabled": self._enable_optimizations,
    }

    if self._loaded and self.model is not None:
        config = getattr(self.model, "config", None)
        if config:
            info.update(
                {
                    "vocab_size": getattr(config, "vocab_size", None),
                    "hidden_size": getattr(config, "hidden_size", None),
                    "num_layers": getattr(config, "num_hidden_layers", None),
                    "num_heads": getattr(config, "num_attention_heads", None),
                }
            )

        # Add optimization status
        info["optimizations"] = {
            "kernel_fusion": self._compiled_forward is not None,
            "memory_optimized": self._hardware_info is not None,
        }

        if self._hardware_info:
            info["hardware"] = {
                "chip": self._hardware_info.chip_name,
                "memory_gb": self._hardware_info.total_memory_gb,
                "bandwidth_gbs": self._hardware_info.memory_bandwidth_gbs,
                "gpu_cores": self._hardware_info.gpu_cores,
                "prefill_chunk_size": self._hardware_info.optimal_prefill_size,
            }

    return info

vllm_mlx.model_runner.MLXModelRunner.__repr__

__repr__() -> str

Source code in vllm_mlx/model_runner.py

def __repr__(self) -> str:
    status = "loaded" if self._loaded else "not loaded"
    opt_status = "optimized" if self._compiled_forward else "standard"
    return f"<MLXModelRunner model={self.model_config.model} status={status} mode={opt_status}>"

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.

vllm_mlx.model_runner.SamplerOutput · class

vllm_mlx.model_runner.SamplerOutput(token_ids: list[int], logprobs: list[dict] | None = None)

Output from sampling.

Parameters

Name	Type	Required	Default	Description
token_ids	list[int]	yes	none	Required constructor field.
logprobs	list[dict] \| None	no	None	Optional constructor field; defaults to None.

Returns

• Constructs: vllm_mlx.model_runner.SamplerOutput

Exceptions and behavior

Class SamplerOutput declares 0 direct member(s). No direct raise statement appears in this definition.

View source #L29-L33.

vllm_mlx.model_runner.MLXModelRunnerOutput · class

vllm_mlx.model_runner.MLXModelRunnerOutput(req_id_to_token_ids: dict[str, list[int]], req_id_to_logprobs: dict[str, list[dict]] | None = None, num_tokens_generated: int = 0, generation_time_s: float = 0.0)

Output from MLX model runner, compatible with vLLM's ModelRunnerOutput.

Parameters

Name	Type	Required	Default	Description
req_id_to_token_ids	dict[str, list[int]]	yes	none	Required constructor field.
req_id_to_logprobs	dict[str, list[dict]] \| None	no	None	Optional constructor field; defaults to None.
num_tokens_generated	int	no	0	Optional constructor field; defaults to 0.
generation_time_s	float	no	0.0	Optional constructor field; defaults to 0.0.

Returns

• Constructs: vllm_mlx.model_runner.MLXModelRunnerOutput

Exceptions and behavior

Class MLXModelRunnerOutput declares 0 direct member(s). No direct raise statement appears in this definition.

View source #L37-L50.

vllm_mlx.model_runner.MLXModelRunner · class

vllm_mlx.model_runner.MLXModelRunner(vllm_config: 'VllmConfig', enable_optimizations: bool = True)

Model runner that uses mlx-lm for inference.

Parameters

Name	Type	Required	Default	Description
vllm_config	'VllmConfig'	yes	none	vLLM configuration
enable_optimizations	bool	no	True	Whether to enable low-level optimizations

Returns

• Constructs: vllm_mlx.model_runner.MLXModelRunner

Exceptions and behavior

Class MLXModelRunner declares 16 direct member(s). No direct raise statement appears in this definition.

View source #L53-L476.

vllm_mlx.model_runner.MLXModelRunner.__init__ · method

vllm_mlx.model_runner.MLXModelRunner.__init__(vllm_config: 'VllmConfig', enable_optimizations: bool = True) -> not annotated

Initialize MLX model runner.

Parameters

Name	Type	Required	Default	Description
vllm_config	'VllmConfig'	yes	none	vLLM configuration
enable_optimizations	bool	no	True	Whether to enable low-level optimizations

Returns

• Type: not annotated

Exceptions and behavior

Method MLXModelRunner.__init__ updates self.vllm_config, self.model_config, self.cache_config, self.scheduler_config; calls logger.info. No direct raise statement appears in this definition.

View source #L69-L104.

vllm_mlx.model_runner.MLXModelRunner.load_model · method

vllm_mlx.model_runner.MLXModelRunner.load_model() -> None

Load model using mlx-lm with optimizations.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MLXModelRunner.load_model updates self.model, self.tokenizer, self._loaded; calls logger.info, time.time, load, self._create_default_sampler; can raise ImportError; returns None. Directly raised exceptions: ImportError.

View source #L106-L145.

vllm_mlx.model_runner.MLXModelRunner._apply_optimizations · method

vllm_mlx.model_runner.MLXModelRunner._apply_optimizations() -> None

Apply low-level optimizations for maximum performance.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method MLXModelRunner._apply_optimizations updates self._hardware_info; calls detect_hardware, logger.info, configure_memory_optimization, self._setup_compiled_forward. No direct raise statement appears in this definition.

View source #L147-L168.

vllm_mlx.model_runner.MLXModelRunner._setup_compiled_forward · method

vllm_mlx.model_runner.MLXModelRunner._setup_compiled_forward() -> None

Setup compiled forward pass using mx.compile() for kernel fusion.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MLXModelRunner._setup_compiled_forward updates self._compiled_forward; calls hasattr, mx.compile, logger.info, logger.warning; returns None. No direct raise statement appears in this definition.

View source #L170-L193.

vllm_mlx.model_runner.MLXModelRunner._create_default_sampler · method

vllm_mlx.model_runner.MLXModelRunner._create_default_sampler() -> None

Create default sampler for generation.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method MLXModelRunner._create_default_sampler updates self._sampler; calls make_sampler, logger.warning. No direct raise statement appears in this definition.

View source #L195-L205.

vllm_mlx.model_runner.MLXModelRunner.initialize_cache · method

vllm_mlx.model_runner.MLXModelRunner.initialize_cache(num_blocks: int) -> None

Initialize KV cache.

Parameters

Name	Type	Required	Default	Description
num_blocks	int	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method MLXModelRunner.initialize_cache updates self._num_cache_blocks; calls logger.info. No direct raise statement appears in this definition.

View source #L207-L210.

vllm_mlx.model_runner.MLXModelRunner.get_kv_cache_spec · method

vllm_mlx.model_runner.MLXModelRunner.get_kv_cache_spec() -> dict

Get KV cache specification.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict
• Direct return expressions: {'num_blocks': self._num_cache_blocks, 'block_size': self.cache_config.block_size}

Exceptions and behavior

Method MLXModelRunner.get_kv_cache_spec returns {'num_blocks': self._num_cache_blocks, 'block_size': self.cache_config.block_size}. No direct raise statement appears in this definition.

View source #L215-L220.

vllm_mlx.model_runner.MLXModelRunner.get_cache_block_size_bytes · method

vllm_mlx.model_runner.MLXModelRunner.get_cache_block_size_bytes() -> int

Calculate cache block size in bytes.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: 0; 2 * block_size * num_layers * num_kv_heads * head_size * 2

Exceptions and behavior

Method MLXModelRunner.get_cache_block_size_bytes calls getattr; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L222-L240.

vllm_mlx.model_runner.MLXModelRunner.warm_up · method

vllm_mlx.model_runner.MLXModelRunner.warm_up() -> None

Warm up model with a test generation.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method MLXModelRunner.warm_up calls self.load_model, logger.info, generate, logger.warning. No direct raise statement appears in this definition.

View source #L242-L263.

vllm_mlx.model_runner.MLXModelRunner.execute_model · method

vllm_mlx.model_runner.MLXModelRunner.execute_model(scheduler_output: 'SchedulerOutput') -> MLXModelRunnerOutput

Execute model inference for scheduled requests.

Parameters

Name	Type	Required	Default	Description
scheduler_output	'SchedulerOutput'	yes	none	Contains requests to process

Returns

• Type: MLXModelRunnerOutput
• Direct return expressions: MLXModelRunnerOutput(req_id_to_token_ids=req_id_to_token_ids, num_tokens_generated=total_tokens, generation_time_s=gene…

Exceptions and behavior

Method MLXModelRunner.execute_model calls RuntimeError, time.time, self._generate_for_request, len; can raise RuntimeError; returns MLXModelRunnerOutput(req_id_to_token_ids=req_id_to_token_ids, num_tokens_generated=total_tokens, generation_time_s=gene…. Directly raised exceptions: RuntimeError.

View source #L265-L315.

vllm_mlx.model_runner.MLXModelRunner._prefill_with_chunking · method

vllm_mlx.model_runner.MLXModelRunner._prefill_with_chunking(input_ids: mx.array, cache: Optional[Any] = None) -> tuple[mx.array, Any]

Process prompt with optimal chunking for L2 cache efficiency.

Parameters

Name	Type	Required	Default	Description
input_ids	mx.array	yes	none	Input token IDs [1, seq_len]
cache	Optional[Any]	no	None	Optional existing KV cache

Returns

• Type: tuple[mx.array, Any]
• Direct return expressions: forward_fn(input_ids, cache=cache); (logits, cache)

Exceptions and behavior

Method MLXModelRunner._prefill_with_chunking calls len, get_optimal_prefill_size, input_ids.reshape, forward_fn; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L317-L362.

vllm_mlx.model_runner.MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size · nested function

vllm_mlx.model_runner.MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size(seq_len) -> not annotated

Nested Function MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size calls min; returns min(512, seq_len).

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: min(512, seq_len)

Exceptions and behavior

Nested Function MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size calls min; returns min(512, seq_len). No direct raise statement appears in this definition.

View source #L339-L340.

vllm_mlx.model_runner.MLXModelRunner._generate_for_request · method

vllm_mlx.model_runner.MLXModelRunner._generate_for_request(prompt_token_ids: list[int], sampling_params: Any, max_tokens: int = 1) -> list[int]

Generate tokens for a single request.

Parameters

Name	Type	Required	Default	Description
prompt_token_ids	list[int]	yes	none	Input token IDs
sampling_params	Any	yes	none	Sampling parameters
max_tokens	int	no	1	Maximum tokens to generate

Returns

• Type: list[int]
• Direct return expressions: generated_ids; []

Exceptions and behavior

Method MLXModelRunner._generate_for_request calls getattr, make_sampler, mx.array, generate_step; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L364-L418.

vllm_mlx.model_runner.MLXModelRunner._continue_generation · method

vllm_mlx.model_runner.MLXModelRunner._continue_generation(req_id: str) -> list[int]

Continue generation for an existing request.

Parameters

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

Returns

• Type: list[int]
• Direct return expressions: []

Exceptions and behavior

Method MLXModelRunner._continue_generation returns []. No direct raise statement appears in this definition.

View source #L420-L428.

vllm_mlx.model_runner.MLXModelRunner.decode_tokens · method

vllm_mlx.model_runner.MLXModelRunner.decode_tokens(token_ids: list[int]) -> str

Decode token IDs to text.

Parameters

Name	Type	Required	Default	Description
token_ids	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: ''; self.tokenizer.decode(token_ids)

Exceptions and behavior

Method MLXModelRunner.decode_tokens calls self.tokenizer.decode; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L430-L434.

vllm_mlx.model_runner.MLXModelRunner.get_model_info · method

vllm_mlx.model_runner.MLXModelRunner.get_model_info() -> dict

Get information about the loaded model and optimizations.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict
• Direct return expressions: info

Exceptions and behavior

Method MLXModelRunner.get_model_info calls getattr, info.update; returns info. No direct raise statement appears in this definition.

View source #L436-L471.

vllm_mlx.model_runner.MLXModelRunner.__repr__ · method

vllm_mlx.model_runner.MLXModelRunner.__repr__() -> str

Method MLXModelRunner.__repr__ returns f'<MLXModelRunner model={self.model_config.model} status={status} mode={opt_status}>'.

Parameters

This callable has no explicit inputs.

Returns

• Type: str
• Direct return expressions: f'<MLXModelRunner model={self.model_config.model} status={status} mode={opt_status}>'

Exceptions and behavior

Method MLXModelRunner.__repr__ returns f'<MLXModelRunner model={self.model_config.model} status={status} mode={opt_status}>'. No direct raise statement appears in this definition.

View source #L473-L476.

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
SamplerOutput	class	SamplerOutput(token_ids: list[int], logprobs: list[dict] \| None = None)	Output from sampling.	#L29-L33
MLXModelRunnerOutput	class	MLXModelRunnerOutput(req_id_to_token_ids: dict[str, list[int]], req_id_to_logprobs: dict[str, list[dict]] \| None = None, num_tokens_generated: int = 0, generation_time_s: float = 0.0)	Output from MLX model runner, compatible with vLLM's ModelRunnerOutput.	#L37-L50
MLXModelRunner	class	MLXModelRunner(vllm_config: 'VllmConfig', enable_optimizations: bool = True)	Model runner that uses mlx-lm for inference.	#L53-L476
MLXModelRunner.__init__	method	MLXModelRunner.__init__(vllm_config: 'VllmConfig', enable_optimizations: bool = True) -> not annotated	Initialize MLX model runner.	#L69-L104
MLXModelRunner.load_model	method	MLXModelRunner.load_model() -> None	Load model using mlx-lm with optimizations.	#L106-L145
MLXModelRunner._apply_optimizations	method	MLXModelRunner._apply_optimizations() -> None	Apply low-level optimizations for maximum performance.	#L147-L168
MLXModelRunner._setup_compiled_forward	method	MLXModelRunner._setup_compiled_forward() -> None	Setup compiled forward pass using mx.compile() for kernel fusion.	#L170-L193
MLXModelRunner._create_default_sampler	method	MLXModelRunner._create_default_sampler() -> None	Create default sampler for generation.	#L195-L205
MLXModelRunner.initialize_cache	method	MLXModelRunner.initialize_cache(num_blocks: int) -> None	Initialize KV cache.	#L207-L210
MLXModelRunner.get_kv_cache_spec	method	MLXModelRunner.get_kv_cache_spec() -> dict	Get KV cache specification.	#L215-L220
MLXModelRunner.get_cache_block_size_bytes	method	MLXModelRunner.get_cache_block_size_bytes() -> int	Calculate cache block size in bytes.	#L222-L240
MLXModelRunner.warm_up	method	MLXModelRunner.warm_up() -> None	Warm up model with a test generation.	#L242-L263
MLXModelRunner.execute_model	method	MLXModelRunner.execute_model(scheduler_output: 'SchedulerOutput') -> MLXModelRunnerOutput	Execute model inference for scheduled requests.	#L265-L315
MLXModelRunner._prefill_with_chunking	method	MLXModelRunner._prefill_with_chunking(input_ids: mx.array, cache: Optional[Any] = None) -> tuple[mx.array, Any]	Process prompt with optimal chunking for L2 cache efficiency.	#L317-L362
MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size	nested function	MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size(seq_len) -> not annotated	Nested Function MLXModelRunner._prefill_with_chunking.get_optimal_prefill_size calls min; returns min(512, seq_len).	#L339-L340
MLXModelRunner._generate_for_request	method	MLXModelRunner._generate_for_request(prompt_token_ids: list[int], sampling_params: Any, max_tokens: int = 1) -> list[int]	Generate tokens for a single request.	#L364-L418
MLXModelRunner._continue_generation	method	MLXModelRunner._continue_generation(req_id: str) -> list[int]	Continue generation for an existing request.	#L420-L428
MLXModelRunner.decode_tokens	method	MLXModelRunner.decode_tokens(token_ids: list[int]) -> str	Decode token IDs to text.	#L430-L434
MLXModelRunner.get_model_info	method	MLXModelRunner.get_model_info() -> dict	Get information about the loaded model and optimizations.	#L436-L471
MLXModelRunner.__repr__	method	MLXModelRunner.__repr__() -> str	Method MLXModelRunner.__repr__ returns f'<MLXModelRunner model={self.model_config.model} status={status} mode={opt_status}>'.	#L473-L476