vllm_mlx.mllm_batch_generator

MLLM Batch Generator for multimodal continuous batching.

View the complete module source at #L1-L3073.

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

MLLM Batch Generator for multimodal continuous batching.

This module implements continuous batching for Multimodal Language Models (MLLMs) like Qwen3-VL, following the same architecture as LLM continuous batching but adapted for vision models.

Key insight: VLM models have a model.language_model which is a standard LLM. After the initial forward pass with vision encoding, text generation uses only the language model - which CAN be batched using the same BatchKVCache pattern.

Architecture: 1. Vision inputs are processed per-request (not batched) 2. Initial VLM forward pass extracts cross-attention states / encoder outputs 3. Language model generation is batched using BatchKVCache (like LLM batching)

vllm_mlx.mllm_batch_generator.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.mllm_batch_generator.PrefillAbortedError

PrefillAbortedError(request_id: str)

Bases: Exception

Raised when a prefill is aborted due to client disconnect.

Source code in vllm_mlx/mllm_batch_generator.py

def __init__(self, request_id: str):
    self.request_id = request_id
    super().__init__(f"Prefill aborted for request {request_id}")

vllm_mlx.mllm_batch_generator.PrefillAbortedError.request_id instance-attribute

request_id = request_id

vllm_mlx.mllm_batch_generator.MLLMBatchRequest dataclass

MLLMBatchRequest(uid: int, request_id: str, prompt: str, images: Optional[List[str]] = None, videos: Optional[List[str]] = None, audio: Optional[List[str]] = None, max_tokens: int = 256, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 0, min_p: float = 0.0, presence_penalty: float = 0.0, repetition_penalty: float = 1.0, logits_processors: Optional[List[Callable]] = None, input_ids: Optional[array] = None, pixel_values: Optional[array] = None, attention_mask: Optional[array] = None, image_grid_thw: Optional[array] = None, extra_kwargs: Dict[str, Any] = dict(), is_text_only: bool = False, num_tokens: int = 0, output_tokens: List[int] = list(), vision_encoded: bool = False, cross_attention_states: Optional[Any] = None, encoder_outputs: Optional[Any] = None)

Request data for MLLM batch processing.

Contains all information needed to process a multimodal request within the batch generator.

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.uid instance-attribute

uid: int

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.request_id instance-attribute

request_id: str

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.prompt instance-attribute

prompt: str

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.images class-attribute instance-attribute

images: Optional[List[str]] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.videos class-attribute instance-attribute

videos: Optional[List[str]] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.audio class-attribute instance-attribute

audio: Optional[List[str]] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.max_tokens class-attribute instance-attribute

max_tokens: int = 256

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.temperature class-attribute instance-attribute

temperature: float = 0.7

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.top_p class-attribute instance-attribute

top_p: float = 0.9

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.top_k class-attribute instance-attribute

top_k: int = 0

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.min_p class-attribute instance-attribute

min_p: float = 0.0

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.presence_penalty class-attribute instance-attribute

presence_penalty: float = 0.0

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.repetition_penalty class-attribute instance-attribute

repetition_penalty: float = 1.0

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.logits_processors class-attribute instance-attribute

logits_processors: Optional[List[Callable]] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.input_ids class-attribute instance-attribute

input_ids: Optional[array] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.pixel_values class-attribute instance-attribute

pixel_values: Optional[array] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.attention_mask class-attribute instance-attribute

attention_mask: Optional[array] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.image_grid_thw class-attribute instance-attribute

image_grid_thw: Optional[array] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.extra_kwargs class-attribute instance-attribute

extra_kwargs: Dict[str, Any] = field(default_factory=dict)

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.is_text_only class-attribute instance-attribute

is_text_only: bool = False

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.num_tokens class-attribute instance-attribute

num_tokens: int = 0

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.output_tokens class-attribute instance-attribute

output_tokens: List[int] = field(default_factory=list)

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.vision_encoded class-attribute instance-attribute

vision_encoded: bool = False

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.cross_attention_states class-attribute instance-attribute

cross_attention_states: Optional[Any] = None

vllm_mlx.mllm_batch_generator.MLLMBatchRequest.encoder_outputs class-attribute instance-attribute

encoder_outputs: Optional[Any] = None

vllm_mlx.mllm_batch_generator.MLLMBatchResponse dataclass

MLLMBatchResponse(uid: int, request_id: str, token: int, logprobs: array, finish_reason: Optional[str] = None, prompt_cache: Optional[Callable[[], List[Any]]] = None, from_draft: bool = False, mtp_attempted: bool = False, mtp_attempted_count: int = 0)

Response from a batch generation step.

Contains the generated token and metadata for a single request.

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.uid instance-attribute

uid: int

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.request_id instance-attribute

request_id: str

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.token instance-attribute

token: int

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.logprobs instance-attribute

logprobs: array

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.finish_reason class-attribute instance-attribute

finish_reason: Optional[str] = None

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.prompt_cache class-attribute instance-attribute

prompt_cache: Optional[Callable[[], List[Any]]] = None

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.from_draft class-attribute instance-attribute

from_draft: bool = False

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.mtp_attempted class-attribute instance-attribute

mtp_attempted: bool = False

vllm_mlx.mllm_batch_generator.MLLMBatchResponse.mtp_attempted_count class-attribute instance-attribute

mtp_attempted_count: int = 0

vllm_mlx.mllm_batch_generator.MLLMBatch dataclass

MLLMBatch(uids: List[int], request_ids: List[str], y: array, logprobs: List[array], max_tokens: List[int], num_tokens: List[int], cache: List[Any], requests: List[MLLMBatchRequest], logits_processors: Optional[List[Optional[List[Callable]]]] = None, samplers: Optional[List[Optional[Callable]]] = None)

Represents an active batch of MLLM requests.

Manages the batch state including tokens, caches, and metadata for all requests being processed together.

vllm_mlx.mllm_batch_generator.MLLMBatch.uids instance-attribute

uids: List[int]

vllm_mlx.mllm_batch_generator.MLLMBatch.request_ids instance-attribute

request_ids: List[str]

vllm_mlx.mllm_batch_generator.MLLMBatch.y instance-attribute

y: array

vllm_mlx.mllm_batch_generator.MLLMBatch.logprobs instance-attribute

logprobs: List[array]

vllm_mlx.mllm_batch_generator.MLLMBatch.max_tokens instance-attribute

max_tokens: List[int]

vllm_mlx.mllm_batch_generator.MLLMBatch.num_tokens instance-attribute

num_tokens: List[int]

vllm_mlx.mllm_batch_generator.MLLMBatch.cache instance-attribute

cache: List[Any]

vllm_mlx.mllm_batch_generator.MLLMBatch.requests instance-attribute

requests: List[MLLMBatchRequest]

vllm_mlx.mllm_batch_generator.MLLMBatch.logits_processors class-attribute instance-attribute

logits_processors: Optional[List[Optional[List[Callable]]]] = None

vllm_mlx.mllm_batch_generator.MLLMBatch.samplers class-attribute instance-attribute

samplers: Optional[List[Optional[Callable]]] = None

vllm_mlx.mllm_batch_generator.MLLMBatch.__len__

__len__() -> int

Source code in vllm_mlx/mllm_batch_generator.py

def __len__(self) -> int:
    return len(self.uids)

vllm_mlx.mllm_batch_generator.MLLMBatch.filter

filter(keep_idx: List[int]) -> None

Filter batch to keep only requests at specified indices.

Parameters:

• keep_idx (List[int]) –

  Indices of requests to keep

Source code in vllm_mlx/mllm_batch_generator.py

def filter(self, keep_idx: List[int]) -> None:
    """
    Filter batch to keep only requests at specified indices.

    Args:
        keep_idx: Indices of requests to keep
    """
    self.uids = [self.uids[k] for k in keep_idx]
    self.request_ids = [self.request_ids[k] for k in keep_idx]
    self.logprobs = [self.logprobs[k] for k in keep_idx]
    self.max_tokens = [self.max_tokens[k] for k in keep_idx]
    self.num_tokens = [self.num_tokens[k] for k in keep_idx]
    self.requests = [self.requests[k] for k in keep_idx]
    if self.logits_processors is not None:
        self.logits_processors = [self.logits_processors[k] for k in keep_idx]
    if self.samplers is not None:
        self.samplers = [self.samplers[k] for k in keep_idx]

    keep_idx_array = mx.array(keep_idx, mx.int32)
    self.y = self.y[keep_idx_array]

    # Filter cache entries
    for c in self.cache:
        if hasattr(c, "filter"):
            c.filter(keep_idx_array)

vllm_mlx.mllm_batch_generator.MLLMBatch.extend

extend(other: MLLMBatch) -> None

Extend this batch with another batch.

Parameters:

• other (MLLMBatch) –

  Batch to merge into this one

Source code in vllm_mlx/mllm_batch_generator.py

def extend(self, other: "MLLMBatch") -> None:
    """
    Extend this batch with another batch.

    Args:
        other: Batch to merge into this one
    """
    self.uids.extend(other.uids)
    self.request_ids.extend(other.request_ids)
    self.y = mx.concatenate([self.y, other.y])
    self.logprobs.extend(other.logprobs)
    self.num_tokens.extend(other.num_tokens)
    self.max_tokens.extend(other.max_tokens)
    self.requests.extend(other.requests)

    # Extend logits_processors
    if self.logits_processors is not None or other.logits_processors is not None:
        # At this point self.uids already includes other.uids from extend above
        self_len = len(self.uids) - len(other.uids)
        self_lp = self.logits_processors or [None] * self_len
        other_lp = other.logits_processors or [None] * len(other.uids)
        self.logits_processors = list(self_lp) + list(other_lp)

    # Extend samplers
    if self.samplers is not None or other.samplers is not None:
        self_len = len(self.uids) - len(other.uids)
        self_s = self.samplers or [None] * self_len
        other_s = other.samplers or [None] * len(other.uids)
        self.samplers = list(self_s) + list(other_s)

    # Extend cache - handle both BatchKVCache (.keys/.values) and
    # ArraysCache (.cache list) from hybrid models like Qwen3.5. Some
    # cache integrations, such as quantized SDPA caches, expose state only
    # through empty()/extend() and do not publish .keys.
    for c, o in zip(self.cache, other.cache):
        if c is not None and o is not None and hasattr(c, "extend"):
            try:
                has_kv = hasattr(c, "keys") and c.keys is not None
                has_arrays = hasattr(c, "cache")
                has_extendable_state = hasattr(c, "empty") and not c.empty()
                if has_kv or has_arrays or has_extendable_state:
                    c.extend(o)
            except Exception as e:
                logger.warning(f"Failed to extend cache: {e}")

vllm_mlx.mllm_batch_generator.MLLMBatch.extract_cache

extract_cache(idx: int) -> List[Any]

Extract cache for a single request (for prefix caching).

Handles BatchRotatingKVCache negative left_padding bug: during generation with rotation, left_padding becomes negative, causing extract() to use Python negative indexing and truncate the buffer to only generation tokens instead of the full window.

Source code in vllm_mlx/mllm_batch_generator.py

def extract_cache(self, idx: int) -> List[Any]:
    """
    Extract cache for a single request (for prefix caching).

    Handles BatchRotatingKVCache negative left_padding bug:
    during generation with rotation, left_padding becomes negative,
    causing extract() to use Python negative indexing and truncate
    the buffer to only generation tokens instead of the full window.
    """
    from mlx_lm.models.cache import (
        BatchRotatingKVCache,
        RotatingKVCache,
    )

    result = []
    for c in self.cache:
        if not hasattr(c, "extract"):
            result.append(None)
        elif isinstance(c, BatchRotatingKVCache):
            # Custom extraction: clamp left_padding to >= 0
            cache = RotatingKVCache(c.max_size)
            padding = max(0, c.left_padding[idx].item())
            offset = c.offset[idx].item()
            cache.keys = c.keys[idx : idx + 1]
            cache.values = c.values[idx : idx + 1]
            cache._idx = c._idx
            if c.rotated:
                cache.keys = mx.roll(cache.keys, -c._idx, axis=2)
                cache.values = mx.roll(cache.values, -c._idx, axis=2)
                cache._idx = c.max_size
            cache.keys = mx.contiguous(cache.keys[:, :, padding : cache._idx])
            cache.values = mx.contiguous(cache.values[:, :, padding : cache._idx])
            cache.offset = offset
            cache._idx = cache.keys.shape[2]
            cache.step = getattr(c, "step", c.max_size)
            cache.keep = getattr(c, "keep", 0)
            result.append(cache)
        else:
            result.append(c.extract(idx))
    return result

vllm_mlx.mllm_batch_generator.MLLMBatchStats

MLLMBatchStats()

Statistics for MLLM batch generation.

Source code in vllm_mlx/mllm_batch_generator.py

def __init__(self):
    self.prompt_tokens: int = 0
    self.prompt_time: float = 0
    self.generation_tokens: int = 0
    self.generation_time: float = 0
    self.vision_encoding_time: float = 0
    self.num_images_processed: int = 0
    self.peak_memory: float = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.prompt_tokens instance-attribute

prompt_tokens: int = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.prompt_time instance-attribute

prompt_time: float = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.generation_tokens instance-attribute

generation_tokens: int = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.generation_time instance-attribute

generation_time: float = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.vision_encoding_time instance-attribute

vision_encoding_time: float = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.num_images_processed instance-attribute

num_images_processed: int = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.peak_memory instance-attribute

peak_memory: float = 0

vllm_mlx.mllm_batch_generator.MLLMBatchStats.prompt_tps property

prompt_tps: float

Return measured multimodal prompt throughput in tokens per second.

vllm_mlx.mllm_batch_generator.MLLMBatchStats.generation_tps property

generation_tps: float

Return measured decode throughput in tokens per second.

vllm_mlx.mllm_batch_generator.MLLMBatchStats.to_dict

to_dict() -> Dict[str, Any]

Return token, timing, vision, and peak-memory statistics.

Source code in vllm_mlx/mllm_batch_generator.py

def to_dict(self) -> Dict[str, Any]:
    """Return token, timing, vision, and peak-memory statistics."""

    return {
        "prompt_tokens": self.prompt_tokens,
        "prompt_time": self.prompt_time,
        "prompt_tps": self.prompt_tps,
        "generation_tokens": self.generation_tokens,
        "generation_time": self.generation_time,
        "generation_tps": self.generation_tps,
        "vision_encoding_time": self.vision_encoding_time,
        "num_images_processed": self.num_images_processed,
        "peak_memory": self.peak_memory,
    }

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator

MLLMBatchGenerator(model: Module, processor: Any, mm_processor: Optional[MultimodalProcessor] = None, max_tokens: int = 256, stop_tokens: Optional[set] = None, sampler: Optional[Callable[[array], array]] = None, prefill_batch_size: int = 4, completion_batch_size: int = 16, prefill_step_size: int = 1024, enable_vision_cache: bool = True, vision_cache_size: int = 100, prefix_cache_config: Optional[MemoryCacheConfig] = None, max_kv_size: int = 0)

Batch generator for Vision Language Models.

This class manages continuous batching for MLLM requests:

1. Vision Encoding Phase:
2. Process images/videos through vision encoder (per-request)
3. Extract vision features and merge with text embeddings

4. Store cross-attention states for language model

5. Language Generation Phase:

6. Use language model with BatchKVCache for batched generation
7. Generate tokens for all requests simultaneously
8. Same pattern as LLM BatchGenerator

Example

generator = MLLMBatchGenerator(model, processor) uids = generator.insert([request1, request2]) while responses := generator.next(): ... for resp in responses: ... print(f"Request {resp.request_id}: token={resp.token}")

Initialize MLLM batch generator.

Parameters:

• model (Module) –

  The VLM model (must have model.language_model)

• processor (Any) –

  The VLM processor for tokenization and image processing

• mm_processor (Optional[MultimodalProcessor], default: None ) –

  Optional MultimodalProcessor for input preparation

• max_tokens (int, default: 256 ) –

  Default max tokens per request

• stop_tokens (Optional[set], default: None ) –

  Set of stop token IDs

• sampler (Optional[Callable[[array], array]], default: None ) –

  Sampling function (default: argmax)

• prefill_batch_size (int, default: 4 ) –

  Max requests to prefill together

• completion_batch_size (int, default: 16 ) –

  Max requests for completion batching

• prefill_step_size (int, default: 1024 ) –

  Tokens to process per prefill step

• enable_vision_cache (bool, default: True ) –

  Enable vision embedding caching

• vision_cache_size (int, default: 100 ) –

  Max entries in vision cache

• prefix_cache_config (Optional[MemoryCacheConfig], default: None ) –

  Config for KV prefix cache (text-only requests)

• max_kv_size (int, default: 0 ) –

  Maximum KV cache size per sequence (0 = unbounded)

Source code in vllm_mlx/mllm_batch_generator.py

def __init__(
    self,
    model: nn.Module,
    processor: Any,
    mm_processor: Optional[MultimodalProcessor] = None,
    max_tokens: int = 256,
    stop_tokens: Optional[set] = None,
    sampler: Optional[Callable[[mx.array], mx.array]] = None,
    prefill_batch_size: int = 4,  # Smaller for MLLM due to vision overhead
    completion_batch_size: int = 16,  # Can be larger for text generation
    prefill_step_size: int = 1024,
    enable_vision_cache: bool = True,
    vision_cache_size: int = 100,
    prefix_cache_config: Optional[MemoryCacheConfig] = None,
    max_kv_size: int = 0,
):
    """
    Initialize MLLM batch generator.

    Args:
        model: The VLM model (must have model.language_model)
        processor: The VLM processor for tokenization and image processing
        mm_processor: Optional MultimodalProcessor for input preparation
        max_tokens: Default max tokens per request
        stop_tokens: Set of stop token IDs
        sampler: Sampling function (default: argmax)
        prefill_batch_size: Max requests to prefill together
        completion_batch_size: Max requests for completion batching
        prefill_step_size: Tokens to process per prefill step
        enable_vision_cache: Enable vision embedding caching
        vision_cache_size: Max entries in vision cache
        prefix_cache_config: Config for KV prefix cache (text-only requests)
        max_kv_size: Maximum KV cache size per sequence (0 = unbounded)
    """
    self.model = model
    self.processor = processor
    self.mm_processor = mm_processor
    self.max_kv_size = max_kv_size

    # Get language model for text generation
    self.language_model = getattr(model, "language_model", model)

    # Check if this is actually a VLM with separate language model
    self.is_vlm = hasattr(model, "language_model")
    if self.is_vlm:
        logger.info(
            "MLLMBatchGenerator: Using VLM's language_model for batched generation"
        )
    else:
        logger.warning(
            "MLLMBatchGenerator: Model does not have language_model, using model directly"
        )

    # Patch attention for BatchKVCache compatibility
    from .patches.qwen3_5_mllm import patch_qwen35_attention_for_batching
    from .patches.gemma4_mllm import patch_gemma4_attention_for_batching
    from .patches.glm4v_moe_mllm import patch_glm4v_moe_for_batching

    patch_qwen35_attention_for_batching()
    patch_gemma4_attention_for_batching()
    patch_glm4v_moe_for_batching()

    self.max_tokens = max_tokens
    self.stop_tokens = stop_tokens or set()
    self.sampler = sampler or (lambda x: mx.argmax(x, axis=-1))

    self.prefill_batch_size = prefill_batch_size
    self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
    self.prefill_step_size = prefill_step_size

    # Request management
    self.unprocessed_requests: List[MLLMBatchRequest] = []
    self.active_batch: Optional[MLLMBatch] = None
    self.uid_counter = 0

    # Statistics
    self._stats = MLLMBatchStats()

    # Error responses for requests that failed during preprocessing
    self._pending_error_responses: List[MLLMBatchResponse] = []

    # Per-request prefill progress: request_id → (processed_tokens, total_tokens)
    self._prefill_progress: Dict[str, Tuple[int, int]] = {}

    # Aborted request IDs — checked between prefill chunks to allow
    # early termination when a client disconnects during long prefill.
    # Set operations are GIL-protected, safe across event-loop and
    # executor threads.
    self._aborted_request_ids: set = set()

    # Deferred removal queue — UIDs scheduled for removal from another
    # thread (typically the event loop on client disconnect).  The
    # actual removal, which mutates `active_batch` and touches MLX
    # arrays, must happen on the scheduler thread to avoid a race with
    # an in-flight forward pass.  Metal asserts ("encodeSignalEvent
    # with uncommitted encoder") if two threads submit GPU work on the
    # same stream concurrently.  See `schedule_removal` /
    # `process_pending_removals`.
    self._pending_removal_uids: set = set()
    self._pending_removal_lock = threading.Lock()

    # Vision embedding cache for repeated images
    self.vision_cache = VisionEmbeddingCache(
        max_pixel_entries=vision_cache_size,
        max_encoding_entries=vision_cache_size // 2,
        enabled=enable_vision_cache,
    )
    if enable_vision_cache:
        logger.info(
            f"MLLMBatchGenerator: Vision cache enabled (size={vision_cache_size})"
        )

    # KV prefix cache for text-only requests
    self.prefix_cache: Optional[MemoryAwarePrefixCache] = None
    if prefix_cache_config is not None:
        self.prefix_cache = MemoryAwarePrefixCache(
            model=self.language_model,
            config=prefix_cache_config,
        )
        logger.info("MLLMBatchGenerator: KV prefix cache enabled")

    # Normalize chat template for prefix-cache stability.
    # Qwen3.5 chat template retroactively changes formatting of earlier
    # assistant messages based on last_query_index (position of last
    # non-tool user message).  When a user text message is appended,
    # last_query_index jumps forward, removing <think> blocks from
    # earlier assistant turns — shifting tokens mid-sequence and
    # breaking prefix match.  Fix: always use plain format for
    # historical assistant turns (thinking is still added by the
    # generation prompt at the end).
    self._normalize_chat_template_for_prefix_cache()

    # Compute think-suffix length for prefix cache key stripping.
    # Models with enable_thinking=True add <think>\n to the generation
    # prompt.  This breaks prefix cache (stored key ends with <think>
    # but next request has actual response at that position).
    # Stripping the suffix from cache keys enables clean PREFIX match.
    self._think_suffix_len = self._compute_think_suffix_len()

    # Generation stream
    if MLLMBatchGenerator._stream is None:
        MLLMBatchGenerator._stream = mx.new_stream(mx.default_device())

    # Memory management
    self._old_wired_limit = None
    if mx.metal.is_available():
        self._old_wired_limit = mx.set_wired_limit(
            mx.device_info()["max_recommended_working_set_size"]
        )

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._stream class-attribute instance-attribute

_stream = None

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.model instance-attribute

model = model

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.processor instance-attribute

processor = processor

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.mm_processor instance-attribute

mm_processor = mm_processor

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.max_kv_size instance-attribute

max_kv_size = max_kv_size

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.language_model instance-attribute

language_model = getattr(model, 'language_model', model)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.is_vlm instance-attribute

is_vlm = hasattr(model, 'language_model')

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.max_tokens instance-attribute

max_tokens = max_tokens

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.stop_tokens instance-attribute

stop_tokens = stop_tokens or set()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.sampler instance-attribute

sampler = sampler or (lambda x: mx.argmax(x, axis=-1))

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.prefill_batch_size instance-attribute

prefill_batch_size = prefill_batch_size

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.completion_batch_size instance-attribute

completion_batch_size = max(completion_batch_size, prefill_batch_size)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.prefill_step_size instance-attribute

prefill_step_size = prefill_step_size

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.unprocessed_requests instance-attribute

unprocessed_requests: List[MLLMBatchRequest] = []

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.active_batch instance-attribute

active_batch: Optional[MLLMBatch] = None

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.uid_counter instance-attribute

uid_counter = 0

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._stats instance-attribute

_stats = MLLMBatchStats()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._pending_error_responses instance-attribute

_pending_error_responses: List[MLLMBatchResponse] = []

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._prefill_progress instance-attribute

_prefill_progress: Dict[str, Tuple[int, int]] = {}

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._aborted_request_ids instance-attribute

_aborted_request_ids: set = set()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._pending_removal_uids instance-attribute

_pending_removal_uids: set = set()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._pending_removal_lock instance-attribute

_pending_removal_lock = threading.Lock()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.vision_cache instance-attribute

vision_cache = VisionEmbeddingCache(max_pixel_entries=vision_cache_size, max_encoding_entries=vision_cache_size // 2, enabled=enable_vision_cache)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.prefix_cache instance-attribute

prefix_cache: Optional[MemoryAwarePrefixCache] = None

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._think_suffix_len instance-attribute

_think_suffix_len = self._compute_think_suffix_len()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._old_wired_limit instance-attribute

_old_wired_limit = None

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._normalize_chat_template_for_prefix_cache

_normalize_chat_template_for_prefix_cache() -> None

Patch chat template so historical assistant turns are prefix-stable.

Qwen3.5's chat template computes last_query_index — the position of the last non-tool-response user message — and conditionally wraps assistant turns after that index in <think>...\n</think>\n\n. When a new user text message is appended, last_query_index jumps forward, retroactively removing these <think> wrappers from earlier assistant turns. This shifts tokens mid-sequence and breaks prefix cache.

Fix: replace the conditional with the plain (ELSE) branch so ALL historical assistant messages use <|im_start|>assistant\ncontent without any injected <think> block. The generation prompt still adds <think>\n at the very end, so the model generates thinking.

Source code in vllm_mlx/mllm_batch_generator.py

def _normalize_chat_template_for_prefix_cache(self) -> None:
    """Patch chat template so historical assistant turns are prefix-stable.

    Qwen3.5's chat template computes ``last_query_index`` — the position
    of the last non-tool-response user message — and conditionally wraps
    assistant turns after that index in ``<think>...\\n</think>\\n\\n``.
    When a new user text message is appended, ``last_query_index`` jumps
    forward, retroactively removing these ``<think>`` wrappers from
    earlier assistant turns.  This shifts tokens mid-sequence and breaks
    prefix cache.

    Fix: replace the conditional with the plain (ELSE) branch so ALL
    historical assistant messages use ``<|im_start|>assistant\\ncontent``
    without any injected ``<think>`` block.  The generation prompt still
    adds ``<think>\\n`` at the very end, so the model generates thinking.
    """
    if self.prefix_cache is None:
        return  # No prefix cache — no need to normalize

    # Find the chat template.  VLM processors (e.g. Qwen3VLProcessor)
    # keep a SEPARATE copy of chat_template from their tokenizer — both
    # must be patched.  The processor's copy is used by
    # BatchedEngine._apply_chat_template() (text rendering), while the
    # tokenizer's copy is used by _compute_think_suffix_len().
    tokenizer = getattr(self.processor, "tokenizer", self.processor)
    # Prefer the processor's own template (it's the one used for rendering)
    template = getattr(self.processor, "chat_template", None)
    if not template:
        template = getattr(tokenizer, "chat_template", None)
    if not template or "last_query_index" not in template:
        return  # Not affected

    import re

    # The pattern in Qwen3.5 template:
    #   {%- if loop.index0 > ns.last_query_index %}
    #       {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
    #   {%- else %}
    #       {{- '<|im_start|>' + message.role + '\n' + content }}
    #   {%- endif %}
    #
    # Replace with just the ELSE branch (always plain format).
    pattern = (
        r"\{%-\s*if\s+loop\.index0\s*>\s*ns\.last_query_index\s*%\}"
        r".*?"
        r"\{%-\s*else\s*%\}"
        r"\s*(\{\{-.*?content.*?\}\})"
        r"\s*\{%-\s*endif\s*%\}"
    )
    new_template = re.sub(pattern, r"\1", template, flags=re.DOTALL)
    if new_template != template:
        # Patch ALL copies: processor, tokenizer, and any dict variants.
        if hasattr(self.processor, "chat_template"):
            self.processor.chat_template = new_template
        tokenizer.chat_template = new_template
        logger.info(
            "[prefix_cache] Normalized chat template: removed "
            "last_query_index conditional for prefix-stable assistant turns"
        )
    else:
        logger.debug(
            "[prefix_cache] Chat template has last_query_index but "
            "regex did not match — template may use a different pattern"
        )

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._compute_think_suffix_len

_compute_think_suffix_len() -> int

Compute how many extra tokens enable_thinking=True adds at the END.

Compares the generation prompt suffix with and without enable_thinking to find the think-tag suffix length (typically <think>\n = 2 tokens for Qwen3/Qwen3.5).

Returns 0 if the template doesn't support enable_thinking.

Source code in vllm_mlx/mllm_batch_generator.py

def _compute_think_suffix_len(self) -> int:
    """Compute how many extra tokens enable_thinking=True adds at the END.

    Compares the generation prompt suffix with and without
    ``enable_thinking`` to find the think-tag suffix length
    (typically ``<think>\\n`` = 2 tokens for Qwen3/Qwen3.5).

    Returns 0 if the template doesn't support ``enable_thinking``.
    """
    try:
        # Find something with apply_chat_template
        applicator = None
        for candidate in [
            getattr(self.processor, "tokenizer", None),
            self.processor,
        ]:
            if candidate is not None and hasattr(candidate, "apply_chat_template"):
                applicator = candidate
                break

        if applicator is None:
            return 0

        dummy = [{"role": "user", "content": "x"}]

        try:
            text_with = applicator.apply_chat_template(
                dummy,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=True,
            )
            text_without = applicator.apply_chat_template(
                dummy,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False,
            )
        except TypeError:
            return 0

        # Check if enable_thinking adds a known think tag at the end.
        # enable_thinking may also change the system prompt, so we can't
        # simply compare lengths — we look at the ending instead.
        for tag in ["<think>\n", "<think>"]:
            if text_with.endswith(tag) and not text_without.endswith(tag):
                tokenizer = getattr(self.processor, "tokenizer", self.processor)
                suffix_tokens = tokenizer.encode(tag)
                base_tokens = tokenizer.encode("")
                suffix_len = len(suffix_tokens) - len(base_tokens)
                if suffix_len > 0:
                    logger.info(
                        f"[think_suffix] Detected think tag "
                        f"'{tag.strip()}' = {suffix_len} token(s)"
                    )
                return max(0, suffix_len)

        return 0
    except Exception:
        return 0

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.close

close() -> None

Release resources and reset wired limit.

Source code in vllm_mlx/mllm_batch_generator.py

def close(self) -> None:
    """Release resources and reset wired limit."""
    if self._old_wired_limit is not None:
        mx.synchronize(MLLMBatchGenerator._stream)
        mx.set_wired_limit(self._old_wired_limit)
        self._old_wired_limit = None

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.abort_prefill

abort_prefill(request_id: str) -> None

Signal that a request's prefill should be aborted.

Called from the event loop thread when a client disconnects. The prefill loop checks this set between chunks and raises PrefillAbortedError to exit early.

Source code in vllm_mlx/mllm_batch_generator.py

def abort_prefill(self, request_id: str) -> None:
    """Signal that a request's prefill should be aborted.

    Called from the event loop thread when a client disconnects.
    The prefill loop checks this set between chunks and raises
    PrefillAbortedError to exit early.
    """
    self._aborted_request_ids.add(request_id)
    logger.info(f"[abort_prefill] Marked {request_id} for prefill abort")

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.schedule_removal

schedule_removal(uids: List[int]) -> None

Thread-safe deferred removal of UIDs from the batch.

Safe to call from any thread (typically the event loop during client-disconnect cleanup). The actual remove(), which creates mx.array instances and filters the KV cache, runs on the scheduler thread via :meth:process_pending_removals at the next batch boundary. This avoids the Metal encodeSignalEvent: uncommitted encoder crash that occurs when two threads submit GPU work on the same stream concurrently.

Source code in vllm_mlx/mllm_batch_generator.py

def schedule_removal(self, uids: List[int]) -> None:
    """Thread-safe deferred removal of UIDs from the batch.

    Safe to call from any thread (typically the event loop during
    client-disconnect cleanup).  The actual `remove()`, which creates
    ``mx.array`` instances and filters the KV cache, runs on the
    scheduler thread via :meth:`process_pending_removals` at the next
    batch boundary.  This avoids the Metal ``encodeSignalEvent:
    uncommitted encoder`` crash that occurs when two threads submit
    GPU work on the same stream concurrently.
    """
    with self._pending_removal_lock:
        self._pending_removal_uids.update(uids)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.process_pending_removals

process_pending_removals() -> None

Remove any UIDs enqueued via :meth:schedule_removal.

MUST be called from the scheduler thread only, at a safe point (e.g. the start of :meth:MLLMScheduler.step before any forward pass has been issued). Safe to call even when the queue is empty (no-op).

Source code in vllm_mlx/mllm_batch_generator.py

def process_pending_removals(self) -> None:
    """Remove any UIDs enqueued via :meth:`schedule_removal`.

    MUST be called from the scheduler thread only, at a safe point
    (e.g. the start of :meth:`MLLMScheduler.step` before any forward
    pass has been issued).  Safe to call even when the queue is
    empty (no-op).
    """
    # Swap the pending set under a lock so enqueues from other threads
    # cannot be dropped between snapshot and clear.
    with self._pending_removal_lock:
        if not self._pending_removal_uids:
            return
        pending = self._pending_removal_uids
        self._pending_removal_uids = set()

    uids = list(pending)
    self.remove(uids)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.__del__

__del__()

Source code in vllm_mlx/mllm_batch_generator.py

def __del__(self):
    try:
        self.close()
    except Exception:
        pass

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.insert

insert(requests: List[MLLMBatchRequest]) -> List[int]

Insert requests for batch processing.

Parameters:

• requests (List[MLLMBatchRequest]) –

  List of MLLMBatchRequest to process

Returns:

• List[int] –

  List of UIDs assigned to requests

Source code in vllm_mlx/mllm_batch_generator.py

def insert(
    self,
    requests: List[MLLMBatchRequest],
) -> List[int]:
    """
    Insert requests for batch processing.

    Args:
        requests: List of MLLMBatchRequest to process

    Returns:
        List of UIDs assigned to requests
    """
    uids = []
    for req in requests:
        req.uid = self.uid_counter
        self.uid_counter += 1
        self.unprocessed_requests.append(req)
        uids.append(req.uid)

    # Sort by estimated complexity (no images = simpler)
    self.unprocessed_requests = sorted(
        self.unprocessed_requests,
        key=lambda x: (
            0 if not x.images and not x.videos and not x.audio else 1,
            len(x.images or []) + len(x.videos or []) + len(x.audio or []),
        ),
    )

    logger.debug(f"Inserted {len(requests)} requests, UIDs: {uids}")
    return uids

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.remove

remove(uids: List[int]) -> None

Remove requests from processing.

Parameters:

• uids (List[int]) –

  List of UIDs to remove

Source code in vllm_mlx/mllm_batch_generator.py

def remove(self, uids: List[int]) -> None:
    """
    Remove requests from processing.

    Args:
        uids: List of UIDs to remove
    """
    uid_set = set(uids)

    # Remove from active batch
    if self.active_batch is not None:
        keep_idx = [
            i for i, uid in enumerate(self.active_batch.uids) if uid not in uid_set
        ]
        if keep_idx:
            self.active_batch.filter(keep_idx)
        else:
            self.active_batch = None

    # Remove from unprocessed
    self.unprocessed_requests = [
        r for r in self.unprocessed_requests if r.uid not in uid_set
    ]

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._preprocess_request

_preprocess_request(request: MLLMBatchRequest) -> None

Preprocess a single MLLM request (vision encoding).

This prepares the inputs by: 1. Processing images/videos through the processor 2. Tokenizing the prompt with image tokens 3. Running vision encoder to get features

Uses vision cache to skip processing for repeated images. Idempotent: if input_ids is already set, returns immediately.

Parameters:

• request (MLLMBatchRequest) –

  Request to preprocess

Source code in vllm_mlx/mllm_batch_generator.py

def _preprocess_request(self, request: MLLMBatchRequest) -> None:
    """
    Preprocess a single MLLM request (vision encoding).

    This prepares the inputs by:
    1. Processing images/videos through the processor
    2. Tokenizing the prompt with image tokens
    3. Running vision encoder to get features

    Uses vision cache to skip processing for repeated images.
    Idempotent: if input_ids is already set, returns immediately.

    Args:
        request: Request to preprocess
    """
    # Already preprocessed (e.g. by early executor offloading in
    # _process_loop or chunked prefill interleaving).  Only skip for
    # text-only requests; media requests need pixel/audio cache lookup
    # even if input_ids was set.
    if (
        request.input_ids is not None
        and not request.images
        and not request.videos
        and not request.audio
    ):
        return

    from mlx_vlm.utils import prepare_inputs

    tic = time.perf_counter()

    # Collect all images (including video frames) and audio inputs
    all_images = []
    all_audio = []

    if request.images:
        from .models.mllm import process_image_input

        for img in request.images:
            try:
                path = process_image_input(img)
                all_images.append(path)
            except Exception as e:
                logger.warning(f"Failed to process image: {e}")

    if request.videos:
        from .models.mllm import (
            process_video_input,
            extract_video_frames_smart,
            save_frames_to_temp,
            DEFAULT_FPS,
            MAX_FRAMES,
        )

        for video in request.videos:
            try:
                video_path = process_video_input(video)
                frames = extract_video_frames_smart(
                    video_path,
                    fps=DEFAULT_FPS,
                    max_frames=MAX_FRAMES,
                )
                frame_paths = save_frames_to_temp(frames)
                all_images.extend(frame_paths)
            except Exception as e:
                logger.warning(f"Failed to process video: {e}")

    if request.audio:
        from .models.mllm import process_audio_input

        for audio in request.audio:
            try:
                path = process_audio_input(audio)
                all_audio.append(path)
            except Exception as e:
                logger.warning(f"Failed to process audio: {e}")

    # Check pixel cache first
    cached_pixels = None
    if not all_audio:
        cached_pixels = self.vision_cache.get_pixel_cache(
            all_images, request.prompt
        )
    if cached_pixels is not None:
        # Cache hit - use cached pixel values
        request.input_ids = cached_pixels.input_ids
        request.pixel_values = cached_pixels.pixel_values
        request.attention_mask = cached_pixels.attention_mask
        request.image_grid_thw = cached_pixels.image_grid_thw
        request.extra_kwargs = dict(cached_pixels.extra_kwargs)

        logger.debug(
            f"Pixel cache HIT for request {request.request_id}: "
            f"saved {cached_pixels.processing_time:.2f}s"
        )
        return

    # Cache miss - process images
    # Get model config
    model_config = getattr(self.model, "config", None)
    image_token_index = (
        getattr(model_config, "image_token_index", None) if model_config else None
    )

    # Prepare inputs using mlx_vlm
    inputs = prepare_inputs(
        self.processor,
        images=all_images if all_images else None,
        audio=all_audio if all_audio else None,
        prompts=request.prompt,
        image_token_index=image_token_index,
    )

    request.input_ids = inputs.get("input_ids")
    request.pixel_values = inputs.get("pixel_values")
    request.attention_mask = inputs.get("attention_mask")

    # Extract extra kwargs
    request.extra_kwargs = {
        k: v
        for k, v in inputs.items()
        if k not in ["input_ids", "pixel_values", "attention_mask"]
    }
    request.image_grid_thw = request.extra_kwargs.pop("image_grid_thw", None)

    processing_time = time.perf_counter() - tic

    # Store in pixel cache for future reuse
    if all_images and not all_audio and request.pixel_values is not None:
        self.vision_cache.set_pixel_cache(
            images=all_images,
            prompt=request.prompt,
            pixel_values=request.pixel_values,
            input_ids=request.input_ids,
            attention_mask=request.attention_mask,
            image_grid_thw=request.image_grid_thw,
            extra_kwargs=request.extra_kwargs,
            processing_time=processing_time,
        )

    self._stats.num_images_processed += len(all_images)
    self._stats.vision_encoding_time += processing_time

    # Mark text-only requests (eligible for prefix cache)
    request.is_text_only = not bool(all_images or all_audio)

    logger.debug(
        f"Preprocessed request {request.request_id}: "
        f"{len(all_images)} images, {len(all_audio)} audio clips, "
        f"{request.input_ids.size if request.input_ids is not None else 0} tokens "
        f"({processing_time:.2f}s)"
    )

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._copy_prefix_cache staticmethod

_copy_prefix_cache(cache_list)

Create shallow copies of cache objects to prevent mutation of stored prefix cache.

MLX arrays are immutable and safe to share, but cache objects have mutable Python attributes (offset, _idx) that get modified by update_and_fetch(). Without copying, the stored prefix cache entry is corrupted after each use.

Source code in vllm_mlx/mllm_batch_generator.py

@staticmethod
def _copy_prefix_cache(cache_list):
    """Create shallow copies of cache objects to prevent mutation of stored prefix cache.

    MLX arrays are immutable and safe to share, but cache objects have mutable
    Python attributes (offset, _idx) that get modified by update_and_fetch().
    Without copying, the stored prefix cache entry is corrupted after each use.
    """
    from mlx_lm.models.cache import KVCache, RotatingKVCache

    copies = []
    for c in cache_list:
        if isinstance(c, RotatingKVCache):
            new_c = RotatingKVCache(c.max_size, c.keep)
            new_c.step = c.step
            new_c.keys = c.keys
            new_c.values = c.values
            new_c.offset = c.offset
            new_c._idx = c._idx
            copies.append(new_c)
        elif isinstance(c, KVCache):
            new_c = KVCache()
            new_c.step = c.step
            new_c.keys = c.keys
            new_c.values = c.values
            new_c.offset = c.offset
            copies.append(new_c)
        else:
            copies.append(c)
    return copies

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._has_empty_rotating_cache staticmethod

_has_empty_rotating_cache(cache_list)

Check if any RotatingKVCache layer has no data (keys=None).

This happens when prefix cache stores a long response where all sliding-window entries were trimmed (entries_to_keep=0). Using such a cache produces garbage — fall through to full prefill.

Source code in vllm_mlx/mllm_batch_generator.py

@staticmethod
def _has_empty_rotating_cache(cache_list):
    """Check if any RotatingKVCache layer has no data (keys=None).

    This happens when prefix cache stores a long response where all
    sliding-window entries were trimmed (entries_to_keep=0).
    Using such a cache produces garbage — fall through to full prefill.
    """
    from mlx_lm.models.cache import RotatingKVCache

    for c in cache_list:
        if isinstance(c, RotatingKVCache) and c.keys is None:
            return True
    return False

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._trim_rotating_caches staticmethod

_trim_rotating_caches(cache_list)

Trim RotatingKVCache buffers restored from prefix cache.

Prefix cache stores the full KV state (offset may exceed max_size for sliding-window layers). RotatingKVCache._update_in_place computes new_size = min(step, max_size - prev) which goes negative when prev > max_size, crashing with "Negative dimensions not allowed".

Trimming the buffer to max_size and clamping offset/idx prevents this.

Source code in vllm_mlx/mllm_batch_generator.py

@staticmethod
def _trim_rotating_caches(cache_list):
    """Trim RotatingKVCache buffers restored from prefix cache.

    Prefix cache stores the full KV state (offset may exceed max_size for
    sliding-window layers).  RotatingKVCache._update_in_place computes
    ``new_size = min(step, max_size - prev)`` which goes negative when
    ``prev > max_size``, crashing with "Negative dimensions not allowed".

    Trimming the buffer to max_size and clamping offset/idx prevents this.
    """
    from mlx_lm.models.cache import RotatingKVCache

    for layer_cache in cache_list:
        if not isinstance(layer_cache, RotatingKVCache):
            continue
        if layer_cache.keys is None:
            layer_cache.offset = 0
            continue
        buf_len = layer_cache.keys.shape[2]
        if buf_len > layer_cache.max_size:
            trim_size = buf_len - layer_cache.max_size
            layer_cache.keys = layer_cache._trim(trim_size, layer_cache.keys)
            layer_cache.values = layer_cache._trim(trim_size, layer_cache.values)
            layer_cache._idx = layer_cache.max_size
        layer_cache.offset = min(layer_cache.offset, layer_cache.max_size)
        # Defensive: ensure size() <= keys.shape[2] to prevent merge crash.
        # Prefix cache trimming can create offset > keys.shape[2] when
        # a supersequence/LCP trim crosses the max_size boundary.
        buf_len = layer_cache.keys.shape[2]
        if min(layer_cache.offset, layer_cache.max_size) > buf_len:
            logger.warning(
                f"RotatingKVCache offset ({layer_cache.offset}) > "
                f"buffer ({buf_len}), capping to buffer size"
            )
            layer_cache.offset = buf_len

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._run_chunked_text_prefill

_run_chunked_text_prefill(request: MLLMBatchRequest, cache: List[Any]) -> array

Run prefill in chunks for text-only requests, reporting real progress.

Processes input_ids in prefill_step_size chunks through the language model, updating _prefill_progress after each chunk so the status endpoint can report accurate prefill percentage.

Returns:

• array –

  Logits from the last chunk (same contract as _run_vision_encoding).

Source code in vllm_mlx/mllm_batch_generator.py

def _run_chunked_text_prefill(
    self, request: MLLMBatchRequest, cache: List[Any]
) -> mx.array:
    """
    Run prefill in chunks for text-only requests, reporting real progress.

    Processes input_ids in prefill_step_size chunks through the language
    model, updating ``_prefill_progress`` after each chunk so the status
    endpoint can report accurate prefill percentage.

    Returns:
        Logits from the last chunk (same contract as _run_vision_encoding).
    """
    input_ids = request.input_ids
    if input_ids.ndim == 1:
        input_ids = input_ids[None, :]

    total = input_ids.shape[1]
    step = self.prefill_step_size

    # Short prompt — process in one shot (no chunking overhead)
    if total <= step:
        self._prefill_progress[request.request_id] = (total, total)
        output = self.language_model(input_ids, cache=cache)
        request.vision_encoded = True
        # Release preprocessed inputs after encoding (issue #442)
        request.pixel_values = None
        request.attention_mask = None
        request.image_grid_thw = None
        request.extra_kwargs.clear()
        if hasattr(output, "logits"):
            return output.logits
        return output

    logger.info(
        f"[chunked_prefill] Starting {request.request_id[:12]}: "
        f"{total} tokens, step={step}"
    )

    # Process all chunks except the last
    processed = 0
    chunk_count = 0
    while processed + step < total:
        # Check for abort between chunks (client disconnect)
        if request.request_id in self._aborted_request_ids:
            self._aborted_request_ids.discard(request.request_id)
            logger.info(
                f"[chunked_prefill] Aborted {request.request_id} at "
                f"{processed}/{total} tokens"
            )
            raise PrefillAbortedError(request.request_id)

        chunk = input_ids[:, processed : processed + step]
        self.language_model(chunk, cache=cache)
        # Eval ALL cache types to break the lazy graph between chunks.
        # ArraysCache (e.g. GatedDeltaNet) has .state; KVCache (full
        # attention) has .keys/.values. Hybrid models like Qwen3.5 use
        # both. Skipping either type lets the computation graph grow
        # across chunks → OOM on long prompts.
        _eval_prompt_cache(cache)
        processed += step
        chunk_count += 1
        self._prefill_progress[request.request_id] = (processed, total)

        # Log progress every 10 chunks so operators can see prefill
        # is progressing (not hanging) during long prompts.
        if chunk_count % 10 == 0:
            logger.info(
                f"[chunked_prefill] {request.request_id[:12]}: "
                f"chunk {chunk_count}, {processed}/{total} tokens"
            )

        # Release Metal buffer pool periodically.  Full-attention layers
        # produce attention score buffers that grow each chunk (1024 ×
        # growing_context).  Old smaller buffers can't be reused, so the
        # pool accumulates O(N²) memory without clearing.
        if chunk_count % 4 == 0:
            mx.clear_cache()

    # Last chunk — return logits for sampling
    last_chunk = input_ids[:, processed:]
    output = self.language_model(last_chunk, cache=cache)
    request.vision_encoded = True
    # Release preprocessed inputs after encoding (issue #442)
    request.pixel_values = None
    request.attention_mask = None
    request.image_grid_thw = None
    request.extra_kwargs.clear()
    self._prefill_progress[request.request_id] = (total, total)

    if chunk_count > 0:
        logger.info(
            f"[chunked_prefill] Completed {request.request_id[:12]}: "
            f"{total} tokens in {chunk_count + 1} chunks"
        )

    if hasattr(output, "logits"):
        return output.logits
    return output

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._run_vision_encoding

_run_vision_encoding(request: MLLMBatchRequest, cache: Optional[List[Any]] = None) -> array

Run the initial VLM forward pass to encode vision and get first logits.

This runs the full VLM model (vision + language) on the prompt, which encodes the images and fills the provided KV cache.

Parameters:

• request (MLLMBatchRequest) –

  Preprocessed request with input_ids and pixel_values

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

  KV cache list for the language model. If provided, the language model writes its KV state directly into this cache during the forward pass.

Returns:

• array –

  Logits from the forward pass

Source code in vllm_mlx/mllm_batch_generator.py

def _run_vision_encoding(
    self, request: MLLMBatchRequest, cache: Optional[List[Any]] = None
) -> mx.array:
    """
    Run the initial VLM forward pass to encode vision and get first logits.

    This runs the full VLM model (vision + language) on the prompt,
    which encodes the images and fills the provided KV cache.

    Args:
        request: Preprocessed request with input_ids and pixel_values
        cache: KV cache list for the language model. If provided, the
               language model writes its KV state directly into this cache
               during the forward pass.

    Returns:
        Logits from the forward pass
    """
    # Build model call kwargs
    kwargs = dict(request.extra_kwargs)

    if request.pixel_values is not None:
        kwargs["pixel_values"] = request.pixel_values
    if request.attention_mask is not None:
        kwargs["attention_mask"] = request.attention_mask
    if request.image_grid_thw is not None:
        kwargs["image_grid_thw"] = request.image_grid_thw

    # Run full VLM forward pass with cache.
    # The VLM passes cache= through to self.language_model(),
    # so the language model writes KV state directly into our cache.
    input_ids = request.input_ids
    if input_ids.ndim == 1:
        input_ids = input_ids[None, :]

    output = self.model(input_ids, cache=cache, **kwargs)
    request.vision_encoded = True

    # Release preprocessed vision inputs now that they have been encoded
    # into the KV cache.  pixel_values can be hundreds of MB for multi-
    # image requests; holding them pins Metal buffers for the entire
    # generation duration (issue #442).
    request.pixel_values = None
    request.attention_mask = None
    request.image_grid_thw = None
    request.extra_kwargs.clear()

    # Handle LanguageModelOutput or plain tensor
    if hasattr(output, "logits"):
        return output.logits
    return output

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._process_prompts

_process_prompts(requests: List[MLLMBatchRequest]) -> MLLMBatch

Process a batch of requests through vision encoding and initial prefill.

For MLLM, this is more complex than LLM: 1. Preprocess each request (tokenize, process images) 2. Run vision encoding per-request with individual KVCache objects 3. Merge individual caches into a BatchKVCache for generation

Parameters:

• requests (List[MLLMBatchRequest]) –

  Requests to process

Returns:

• MLLMBatch –

  MLLMBatch ready for generation

Source code in vllm_mlx/mllm_batch_generator.py

def _process_prompts(self, requests: List[MLLMBatchRequest]) -> MLLMBatch:
    """
    Process a batch of requests through vision encoding and initial prefill.

    For MLLM, this is more complex than LLM:
    1. Preprocess each request (tokenize, process images)
    2. Run vision encoding per-request with individual KVCache objects
    3. Merge individual caches into a BatchKVCache for generation

    Args:
        requests: Requests to process

    Returns:
        MLLMBatch ready for generation
    """
    from mlx_lm.models.cache import make_prompt_cache
    from mlx_lm.sample_utils import make_logits_processors, make_sampler

    tic = time.perf_counter()

    # Preprocess all requests (per-request error handling)
    failed_requests = []
    for req in requests:
        try:
            self._preprocess_request(req)
        except Exception as e:
            logger.error(
                f"Failed to preprocess request {req.request_id}: "
                f"{type(e).__name__}: {e}"
            )
            failed_requests.append(req)

    # Remove failed requests from batch and create error responses
    if failed_requests:
        for req in failed_requests:
            requests.remove(req)
            self._pending_error_responses.append(
                MLLMBatchResponse(
                    uid=req.uid,
                    request_id=req.request_id,
                    token=0,
                    logprobs=mx.zeros(1),
                    finish_reason="error",
                )
            )

    if not requests:
        # All requests failed
        return None

    logits_processors_by_request: dict[str, Optional[List[Callable]]] = {}
    samplers_by_request: dict[str, Optional[Callable]] = {}
    for req in requests:
        need_rep = req.repetition_penalty and req.repetition_penalty != 1.0
        need_pres = req.presence_penalty and req.presence_penalty != 0.0
        combined: List[Callable] = []
        if need_rep or need_pres:
            lp_kwargs = {}
            if need_rep:
                lp_kwargs["repetition_penalty"] = req.repetition_penalty
            if need_pres:
                lp_kwargs["presence_penalty"] = req.presence_penalty
            combined.extend(make_logits_processors(**lp_kwargs))
            logger.info(
                f"[sampling] request={req.request_id[:12]} "
                f"rep_penalty={req.repetition_penalty} "
                f"pres_penalty={req.presence_penalty}"
            )
        if req.logits_processors:
            combined.extend(req.logits_processors)
            logger.info(
                f"[sampling] request={req.request_id[:12]} "
                f"extra_logits_processors={len(req.logits_processors)}"
            )
        logits_processors_by_request[req.request_id] = combined or None

        samplers_by_request[req.request_id] = make_sampler(
            temp=req.temperature,
            top_p=req.top_p,
            top_k=req.top_k,
            min_p=req.min_p,
        )
        logger.info(
            f"[sampling] request={req.request_id[:12]} "
            f"temp={req.temperature} top_p={req.top_p} "
            f"top_k={req.top_k} min_p={req.min_p}"
        )

    def _sample_first_token(req: MLLMBatchRequest, logits: mx.array):
        sample_logits = logits
        processors = logits_processors_by_request.get(req.request_id)
        if processors:
            empty_tokens = mx.array([], dtype=mx.uint32)
            for processor in processors:
                sample_logits = processor(empty_tokens, sample_logits)

        logprobs = sample_logits - mx.logsumexp(
            sample_logits, axis=-1, keepdims=True
        )
        sampler = samplers_by_request.get(req.request_id) or self.sampler
        sampled = sampler(logprobs)
        mx.eval(sampled, logprobs)
        return sampled, logprobs

    total_prompt_tokens = sum(
        req.input_ids.size if req.input_ids is not None else 1 for req in requests
    )
    self._stats.prompt_tokens += total_prompt_tokens

    # Log large prompts for monitoring (was previously a hard check that
    # caused infinite retry loops when requests exceeded the limit).
    max_batch_tokens = self.prefill_step_size * len(requests)
    if total_prompt_tokens > max_batch_tokens:
        logger.warning(
            f"Large batch prefill: {total_prompt_tokens} tokens "
            f"(step_size={self.prefill_step_size}, requests={len(requests)}). "
            f"Processing may be slow."
        )

    # Run vision encoding for each request with its own KVCache.
    # Vision encoding cannot be batched because each request may have
    # different images/pixel values. We pass a per-request KVCache to
    # the VLM so the language model writes its KV state directly into it.
    #
    # For text-only requests, we check the prefix cache first. If there's
    # a hit, we skip the full VLM forward and run only the language model
    # on the remaining (uncached) tokens.
    first_tokens = []
    all_logprobs = []
    per_request_caches = []

    aborted_requests = []
    for req in requests:
        try:
            # Check abort before starting prefill
            if req.request_id in self._aborted_request_ids:
                self._aborted_request_ids.discard(req.request_id)
                raise PrefillAbortedError(req.request_id)

            # Try prefix cache for all requests (text-only and multimodal).
            # VLM forward writes the same KV state as language model forward
            # for text tokens, so cached KV from a previous VLM run is valid.
            # However, if the remaining (uncached) tokens contain image
            # placeholders, we must fall back to VLM forward instead of
            # running them through the language model alone.
            cached_kv = None
            remaining_ids = None
            if self.prefix_cache is not None and req.input_ids is not None:
                input_ids_list = req.input_ids.reshape(-1).tolist()
                # Strip think suffix from lookup key so stored entries
                # (also stripped) match as clean PREFIX.
                S = self._think_suffix_len
                lookup_ids = input_ids_list[:-S] if S > 0 else input_ids_list
                cached_kv, remaining_ids = self.prefix_cache.fetch(lookup_ids)
                # Append think suffix back to remaining so the model
                # sees the full generation prompt (<think>\n).
                if cached_kv is not None and S > 0:
                    remaining_ids = list(remaining_ids) + input_ids_list[-S:]

                # If remaining tokens contain image placeholders, the
                # language-model-only path cannot handle them — clear the
                # cache hit so we fall through to full VLM forward.
                if cached_kv is not None and remaining_ids:
                    img_tok = getattr(
                        getattr(self.model, "config", None),
                        "image_token_index",
                        None,
                    )
                    if img_tok is not None and img_tok in remaining_ids:
                        cached_kv = None
                        remaining_ids = None

            # Detect empty RotatingKVCache in cached entry — if any sliding-window
            # layer has keys=None (all entries trimmed), the cache is unusable.
            # Fall through to full prefill instead of producing garbage.
            if cached_kv is not None and self._has_empty_rotating_cache(cached_kv):
                logger.warning(
                    f"Prefix cache hit for {req.request_id} has empty "
                    f"RotatingKVCache layers — falling through to full prefill"
                )
                cached_kv = None
                remaining_ids = None

            if cached_kv is not None and remaining_ids:
                # Prefix/LCP match — run language model on remaining tokens.
                # Copy cache to prevent mutation of stored prefix cache entry.
                request_cache = self._copy_prefix_cache(cached_kv)
                self._trim_rotating_caches(request_cache)
                remaining = mx.array(remaining_ids)[None, :]
                cached_count = len(input_ids_list) - len(remaining_ids)
                total_tokens = len(input_ids_list)
                remaining_count = len(remaining_ids)

                with mx.stream(MLLMBatchGenerator._stream):
                    step = self.prefill_step_size
                    if remaining_count <= step:
                        # Short remaining — process in one shot
                        self._prefill_progress[req.request_id] = (
                            total_tokens,
                            total_tokens,
                        )
                        logits = self.language_model(remaining, cache=request_cache)
                    else:
                        # Chunked prefill on remaining tokens
                        self._prefill_progress[req.request_id] = (
                            cached_count,
                            total_tokens,
                        )
                        processed = 0
                        chunk_count = 0
                        while processed + step < remaining_count:
                            # Check for abort between chunks
                            if req.request_id in self._aborted_request_ids:
                                self._aborted_request_ids.discard(req.request_id)
                                logger.info(
                                    f"[chunked_prefill] Aborted {req.request_id} "
                                    f"at {cached_count + processed}/{total_tokens} tokens"
                                )
                                raise PrefillAbortedError(req.request_id)

                            chunk = remaining[:, processed : processed + step]
                            self.language_model(chunk, cache=request_cache)
                            # Eval ALL cache types (see _run_chunked_text_prefill)
                            _eval_prompt_cache(request_cache)
                            processed += step
                            chunk_count += 1
                            self._prefill_progress[req.request_id] = (
                                cached_count + processed,
                                total_tokens,
                            )
                            if chunk_count % 4 == 0:
                                mx.clear_cache()
                        # Last chunk — return logits
                        remaining = remaining[:, processed:]
                        logits = self.language_model(remaining, cache=request_cache)
                        self._prefill_progress[req.request_id] = (
                            total_tokens,
                            total_tokens,
                        )

                    if hasattr(logits, "logits"):
                        logits = logits.logits

                    last_logits = logits[:, -1, :]

                    sampled, logprobs = _sample_first_token(req, last_logits)

                    first_tokens.append(sampled.item())
                    all_logprobs.append(logprobs.squeeze(0))

                per_request_caches.append(request_cache)
                req.vision_encoded = True
                logger.debug(
                    f"Prefix cache hit for {req.request_id}: "
                    f"cached={cached_count}, "
                    f"remaining={remaining_count}"
                )

            elif cached_kv is not None and not remaining_ids:
                # Exact/supersequence match — cache has all prompt tokens,
                # but we still need logits for the last position.
                # Trim by 1 so re-running the last token produces correct
                # logits for the next-token prediction.
                # _trim_cache_offset creates new cache objects (safe for
                # stored entry).
                request_cache = _trim_cache_offset(cached_kv, 1)
                last_token = req.input_ids[:, -1:]
                total_tokens = len(input_ids_list)
                self._prefill_progress[req.request_id] = (
                    total_tokens,
                    total_tokens,
                )

                with mx.stream(MLLMBatchGenerator._stream):
                    logits = self.language_model(last_token, cache=request_cache)
                    if hasattr(logits, "logits"):
                        logits = logits.logits

                    last_logits = logits[:, -1, :]

                    sampled, logprobs = _sample_first_token(req, last_logits)

                    first_tokens.append(sampled.item())
                    all_logprobs.append(logprobs.squeeze(0))

                per_request_caches.append(request_cache)
                req.vision_encoded = True
                logger.debug(
                    f"Prefix cache exact hit for {req.request_id}: "
                    f"all {total_tokens} tokens cached"
                )

            else:
                # Cache miss — full forward pass
                request_cache = make_prompt_cache(
                    self.language_model,
                    max_kv_size=self.max_kv_size or None,
                )

                with mx.stream(MLLMBatchGenerator._stream):
                    # Text-only: chunked prefill with real progress tracking
                    # Multimodal: atomic VLM forward (vision encoder needs full input)
                    if req.is_text_only:
                        logits = self._run_chunked_text_prefill(
                            req, cache=request_cache
                        )
                    else:
                        logits = self._run_vision_encoding(req, cache=request_cache)

                    # Extract last token logits
                    last_logits = logits[:, -1, :]

                    sampled, logprobs = _sample_first_token(req, last_logits)

                    first_tokens.append(sampled.item())
                    all_logprobs.append(logprobs.squeeze(0))

                per_request_caches.append(request_cache)

        except PrefillAbortedError:
            aborted_requests.append(req)
            self._prefill_progress.pop(req.request_id, None)
            self._pending_error_responses.append(
                MLLMBatchResponse(
                    uid=req.uid,
                    request_id=req.request_id,
                    token=0,
                    logprobs=mx.zeros(1),
                    finish_reason="abort",
                )
            )

    # Remove aborted requests — they have no entries in the parallel
    # lists (first_tokens, all_logprobs, per_request_caches)
    if aborted_requests:
        for req in aborted_requests:
            requests.remove(req)
        mx.clear_cache()
        if not requests:
            return None

    # Merge per-request caches into batched caches.
    # Both KVCache.merge() and ArraysCache.merge() produce batch-aware
    # caches that support filter/extend/extract for continuous batching.
    #
    # Fix: RotatingKVCache._update_concat does NOT trim on first call —
    # if prompt length > max_size, the buffer grows beyond max_size.
    # BatchRotatingKVCache.merge() then hits a shape mismatch when
    # copying via _temporal_order (full buffer) into a max_size slice.
    # Trim buffer to max_size before merging.
    from mlx_lm.models.cache import RotatingKVCache

    for rc in per_request_caches:
        self._trim_rotating_caches(rc)
        for layer_cache in rc:
            if isinstance(layer_cache, RotatingKVCache):
                if layer_cache.keys is not None:
                    # Normalize wrapped rotating cache for merge:
                    # after rotation _idx wraps around but merge()
                    # expects _idx == actual buffer size.
                    # Use keys.shape[2] (actual entries) NOT size()
                    # which can be inconsistent after prefix cache trim
                    # (size() = min(offset, max_size) but buffer may
                    # have fewer entries when trimmed).
                    actual_buf = layer_cache.keys.shape[2]
                    if layer_cache._idx != actual_buf and actual_buf > 0:
                        layer_cache.keys = layer_cache._temporal_order(
                            layer_cache.keys
                        )
                        layer_cache.values = layer_cache._temporal_order(
                            layer_cache.values
                        )
                        layer_cache._idx = actual_buf

    try:
        batch_cache = [
            per_request_caches[0][layer_idx].merge(
                [c[layer_idx] for c in per_request_caches]
            )
            for layer_idx in range(len(per_request_caches[0]))
        ]
    except Exception as e:
        sample_type = type(per_request_caches[0][0]).__name__
        logger.error(
            f"Failed to merge per-request caches ({sample_type}): "
            f"{type(e).__name__}: {e}"
        )
        raise

    # Create initial y (first generated tokens)
    y = mx.array(first_tokens)

    batch_logits_processors = [
        logits_processors_by_request.get(req.request_id) for req in requests
    ]
    has_any_lp = any(batch_logits_processors)
    batch_samplers = [samplers_by_request.get(req.request_id) for req in requests]
    has_any_sampler = any(batch_samplers)

    self._stats.prompt_time += time.perf_counter() - tic

    # Release preprocessed vision inputs for all requests now that
    # they have been encoded into the batch KV cache.  pixel_values,
    # input_ids, etc. can be hundreds of MB per request; holding them
    # for the entire generation duration pins Metal buffers (issue #442).
    for req in requests:
        req.pixel_values = None
        req.attention_mask = None
        req.image_grid_thw = None
        req.extra_kwargs.clear()

    return MLLMBatch(
        uids=[req.uid for req in requests],
        request_ids=[req.request_id for req in requests],
        y=y,
        logprobs=all_logprobs,
        max_tokens=[req.max_tokens for req in requests],
        num_tokens=[0] * len(requests),
        cache=batch_cache,
        requests=requests,
        logits_processors=batch_logits_processors if has_any_lp else None,
        samplers=batch_samplers if has_any_sampler else None,
    )

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._step

_step(input_tokens: array, cache: List[Any], logits_processors: Optional[List[Optional[List[Callable]]]] = None, output_tokens: Optional[List[List[int]]] = None, samplers: Optional[List[Optional[Callable]]] = None) -> Tuple[array, List[array]]

Run one generation step through the language model.

Parameters:

• input_tokens (array) –

  Input tokens [batch_size, 1] or [batch_size]

• cache (List[Any]) –

  BatchKVCache for the language model

• logits_processors (Optional[List[Optional[List[Callable]]]], default: None ) –

  Per-request logits processors (e.g. repetition penalty)

• output_tokens (Optional[List[List[int]]], default: None ) –

  Per-request generated tokens so far (needed by processors)

• samplers (Optional[List[Optional[Callable]]], default: None ) –

  Per-request sampler functions (for top_k/min_p)

Returns:

• Tuple[array, List[array]] –

  Tuple of (sampled tokens, logprobs list)

Source code in vllm_mlx/mllm_batch_generator.py

def _step(
    self,
    input_tokens: mx.array,
    cache: List[Any],
    logits_processors: Optional[List[Optional[List[Callable]]]] = None,
    output_tokens: Optional[List[List[int]]] = None,
    samplers: Optional[List[Optional[Callable]]] = None,
) -> Tuple[mx.array, List[mx.array]]:
    """
    Run one generation step through the language model.

    Args:
        input_tokens: Input tokens [batch_size, 1] or [batch_size]
        cache: BatchKVCache for the language model
        logits_processors: Per-request logits processors (e.g. repetition penalty)
        output_tokens: Per-request generated tokens so far (needed by processors)
        samplers: Per-request sampler functions (for top_k/min_p)

    Returns:
        Tuple of (sampled tokens, logprobs list)
    """
    # Ensure correct shape
    if input_tokens.ndim == 1:
        input_tokens = input_tokens[:, None]

    # Run language model only (not full VLM)
    output = self.language_model(input_tokens, cache=cache)

    # Handle LanguageModelOutput or plain tensor
    if hasattr(output, "logits"):
        logits = output.logits
    else:
        logits = output

    logits = logits[:, -1, :]

    # Apply per-request logits processors (repetition penalty etc.)
    if logits_processors and output_tokens and any(logits_processors):
        processed_logits = []
        for e in range(logits.shape[0]):
            sample_logits = logits[e : e + 1]
            if logits_processors[e]:
                # ``output_tokens[e]`` already contains all generated
                # tokens including the current step's input token (built
                # by the caller as ``req.output_tokens + [token]``).
                full_context = output_tokens[e]
                for processor in logits_processors[e]:
                    sample_logits = processor(mx.array(full_context), sample_logits)
            processed_logits.append(sample_logits)
        logits = mx.concatenate(processed_logits, axis=0)

    # Sample — per-request samplers for top_k/min_p support
    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
    if samplers and any(samplers):
        sampled_list = []
        for e in range(logprobs.shape[0]):
            s = samplers[e] if samplers[e] else self.sampler
            sampled_list.append(s(logprobs[e : e + 1]))
        sampled = mx.concatenate(sampled_list, axis=0)
    else:
        sampled = self.sampler(logprobs)

    return sampled, list(logprobs)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._next

_next() -> List[MLLMBatchResponse]

Internal next() implementation.

Returns:

• List[MLLMBatchResponse] –

  List of MLLMBatchResponse for this step

Source code in vllm_mlx/mllm_batch_generator.py

def _next(self) -> List[MLLMBatchResponse]:
    """
    Internal next() implementation.

    Returns:
        List of MLLMBatchResponse for this step
    """
    tic = time.perf_counter()

    prompt_processing = False
    batch = self.active_batch
    num_active = len(batch) if batch else 0

    # Only start a new batch when there is no active batch generating.
    # Per-request KV caches are created during vision encoding and then
    # merged into a single BatchKVCache. Merging into an active batch
    # mid-generation would cause shape mismatches in attention layers,
    # so queued requests wait until the current batch finishes.
    # Exception: text-only requests can be extended into an active batch
    # via the elif branch below (they skip vision encoding entirely).
    if num_active == 0:
        requests = self.unprocessed_requests[: self.completion_batch_size]

        if len(requests) == 0:
            self.active_batch = None
            return []

        try:
            # Save count before _process_prompts which modifies
            # `requests` in-place via .remove() for failed items.
            num_to_consume = len(requests)
            new_batch = self._process_prompts(requests)
            self.unprocessed_requests = self.unprocessed_requests[num_to_consume:]
            self.active_batch = new_batch
            prompt_processing = True
        except Exception as e:
            logger.error(
                f"Failed to process batch of {len(requests)} prompts: "
                f"{type(e).__name__}: {e}",
                exc_info=True,
            )
            # Remove failed requests to avoid infinite retry loop
            self.unprocessed_requests = self.unprocessed_requests[len(requests) :]
            for req in requests:
                self._pending_error_responses.append(
                    MLLMBatchResponse(
                        uid=req.uid,
                        request_id=req.request_id,
                        token=0,
                        logprobs=mx.zeros(1),
                        finish_reason="error",
                    )
                )

    # Mid-batch extend: text-only requests can join an active batch
    # without vision encoding (no shape mismatch risk).
    elif self.unprocessed_requests:
        text_only = [
            r for r in self.unprocessed_requests if not r.images and not r.videos
        ][: self.completion_batch_size]

        if text_only:
            try:
                # Capture UIDs before _process_prompts modifies
                # text_only in-place via .remove() for failed items.
                all_uids = {r.uid for r in text_only}
                new_batch = self._process_prompts(text_only)
                # Remove ALL requested (both successful and failed)
                self.unprocessed_requests = [
                    r for r in self.unprocessed_requests if r.uid not in all_uids
                ]
                if new_batch is not None:
                    batch.extend(new_batch)
                prompt_processing = True
            except Exception as e:
                logger.warning(
                    f"Failed to extend batch with text-only requests: "
                    f"{type(e).__name__}: {e}"
                )
                # Remove failed requests to avoid infinite retry loop
                processed_uids = {r.uid for r in text_only}
                self.unprocessed_requests = [
                    r
                    for r in self.unprocessed_requests
                    if r.uid not in processed_uids
                ]
                for req in text_only:
                    self._pending_error_responses.append(
                        MLLMBatchResponse(
                            uid=req.uid,
                            request_id=req.request_id,
                            token=0,
                            logprobs=mx.zeros(1),
                            finish_reason="error",
                        )
                    )

    # Collect any pending error responses (from failed preprocessing)
    error_responses = []
    if self._pending_error_responses:
        error_responses = list(self._pending_error_responses)
        self._pending_error_responses.clear()

    # Generate next token for active batch
    batch = self.active_batch
    if batch is None:
        return error_responses

    y, logprobs = batch.y, batch.logprobs
    output_tokens = None
    if batch.logits_processors:
        y_list = y.tolist()
        output_tokens = [
            list(req.output_tokens) + [token]
            for req, token in zip(batch.requests, y_list)
        ]
    batch.y, batch.logprobs = self._step(
        y[:, None],
        batch.cache,
        batch.logits_processors,
        output_tokens,
        batch.samplers,
    )
    mx.async_eval(batch.y, batch.logprobs)

    y = y.tolist()
    toc = time.perf_counter()

    if prompt_processing and num_active == 0:
        # Pure prompt processing (new batch, no prior generation)
        self._stats.prompt_time += toc - tic
    else:
        # Generation step — even if a new request was extended into the
        # batch, the dominant cost is generating for all existing requests.
        self._stats.generation_time += toc - tic

    # Build responses and track finished
    keep_idx = []
    end_idx = []
    responses = []

    for i, (token, uid, request_id, num_tok, max_tok, req) in enumerate(
        zip(
            y,
            batch.uids,
            batch.request_ids,
            batch.num_tokens,
            batch.max_tokens,
            batch.requests,
        )
    ):
        num_tok += 1
        batch.num_tokens[i] = num_tok
        req.num_tokens = num_tok
        req.output_tokens.append(token)

        if batch.logits_processors and _processors_can_retire(
            batch.logits_processors[i]
        ):
            remaining_processors, retired_count = _drop_retired_processors(
                batch.logits_processors[i]
            )
            if retired_count > 0:
                # Keep the per-request slot but replace an empty processor
                # stack with None. The next `_mtp_step` uses any([None]) ==
                # False, so a fully retired request becomes MTP-eligible
                # without changing batch alignment.
                batch.logits_processors[i] = remaining_processors
                logger.info(
                    "[MTP-MLLM] request=%s retired %d processor(s); "
                    "mtp_eligible_next_step=%s",
                    request_id[:12],
                    retired_count,
                    remaining_processors is None,
                )

        finish_reason = None
        cache_fn = None

        if token in self.stop_tokens:
            finish_reason = "stop"
            end_idx.append(i)
        elif num_tok >= max_tok:
            finish_reason = "length"
            end_idx.append(i)
        else:
            keep_idx.append(i)

        if finish_reason is not None:
            # Extract cache for this request
            cache_fn = lambda idx=i: batch.extract_cache(idx)
            # Cleanup prefill progress tracking
            self._prefill_progress.pop(request_id, None)

        responses.append(
            MLLMBatchResponse(
                uid=uid,
                request_id=request_id,
                token=token,
                logprobs=logprobs[i],
                finish_reason=finish_reason,
                prompt_cache=cache_fn,
            )
        )

    # Store caches for finished text-only requests BEFORE filtering
    self._maybe_store_prefix_cache(batch, end_idx)

    # Remove finished requests from batch
    if end_idx:
        if keep_idx:
            batch.filter(keep_idx)
        else:
            self.active_batch = None

    self._stats.generation_tokens += len(responses)
    return error_responses + responses

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.next

next() -> List[MLLMBatchResponse]

Generate next token for all requests in the batch.

Returns:

• List[MLLMBatchResponse] –

  List of MLLMBatchResponse, one per active request

Source code in vllm_mlx/mllm_batch_generator.py

def next(self) -> List[MLLMBatchResponse]:
    """
    Generate next token for all requests in the batch.

    Returns:
        List of MLLMBatchResponse, one per active request
    """
    with mx.stream(MLLMBatchGenerator._stream):
        return self._next()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.stats

stats() -> MLLMBatchStats

Get generation statistics.

Returns:

• MLLMBatchStats –

  MLLMBatchStats with timing and token counts

Source code in vllm_mlx/mllm_batch_generator.py

def stats(self) -> MLLMBatchStats:
    """
    Get generation statistics.

    Returns:
        MLLMBatchStats with timing and token counts
    """
    self._stats.peak_memory = mx.get_peak_memory() / 1e9
    return self._stats

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._maybe_store_prefix_cache

_maybe_store_prefix_cache(batch: MLLMBatch, end_indices: List[int]) -> None

Store KV caches for finished text-only requests into prefix cache.

Must be called BEFORE batch.filter() so that indices are still valid.

Source code in vllm_mlx/mllm_batch_generator.py

def _maybe_store_prefix_cache(
    self, batch: MLLMBatch, end_indices: List[int]
) -> None:
    """Store KV caches for finished text-only requests into prefix cache.

    Must be called BEFORE batch.filter() so that indices are still valid.
    """
    if self.prefix_cache is None or not end_indices:
        return
    for i in end_indices:
        req = batch.requests[i]
        if req.input_ids is not None:
            try:
                extracted = batch.extract_cache(i)
                input_ids_list = req.input_ids.reshape(-1).tolist()
                # Store prompt-only KV: trim generated tokens (+ think
                # suffix) so the stored offset equals key length exactly.
                # The exact-match path trims by 1 at fetch time to
                # re-derive logits for the last prompt token.
                output_count = batch.num_tokens[i]
                S = self._think_suffix_len
                total_trim = output_count + S
                prompt_cache = _trim_cache_offset(extracted, total_trim)
                cache_key = input_ids_list[:-S] if S > 0 else input_ids_list
                self.prefix_cache.store(cache_key, prompt_cache)
            except Exception as e:
                logger.warning(
                    f"Failed to store prefix cache for {req.request_id}: {type(e).__name__}: {e}"
                )

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_prefill_progress

get_prefill_progress(request_id: str) -> Optional[Tuple[int, int]]

Return (processed_tokens, total_tokens) or None.

Source code in vllm_mlx/mllm_batch_generator.py

def get_prefill_progress(self, request_id: str) -> Optional[Tuple[int, int]]:
    """Return (processed_tokens, total_tokens) or None."""
    return self._prefill_progress.get(request_id)

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_vision_cache_stats

get_vision_cache_stats() -> Dict[str, Any]

Get vision cache statistics.

Source code in vllm_mlx/mllm_batch_generator.py

def get_vision_cache_stats(self) -> Dict[str, Any]:
    """Get vision cache statistics."""
    return self.vision_cache.get_stats()

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_prefix_cache_stats

get_prefix_cache_stats() -> Dict[str, Any]

Get KV prefix cache statistics.

Source code in vllm_mlx/mllm_batch_generator.py

def get_prefix_cache_stats(self) -> Dict[str, Any]:
    """Get KV prefix cache statistics."""
    if self.prefix_cache is not None:
        return self.prefix_cache.get_stats()
    return {
        "hits": 0,
        "misses": 0,
        "hit_rate": 0.0,
        "evictions": 0,
        "tokens_saved": 0,
        "current_memory_mb": 0.0,
        "max_memory_mb": 0.0,
        "memory_utilization": 0.0,
        "entry_count": 0,
    }

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.has_pending

has_pending() -> bool

Check if there are pending or active requests.

Source code in vllm_mlx/mllm_batch_generator.py

def has_pending(self) -> bool:
    """Check if there are pending or active requests."""
    return bool(self.unprocessed_requests or self.active_batch)

vllm_mlx.mllm_batch_generator._processors_can_retire

_processors_can_retire(processors: Optional[List[Callable]]) -> bool

True when any processor advertises a retire-to-content transition.

Source code in vllm_mlx/mllm_batch_generator.py

def _processors_can_retire(processors: Optional[List[Callable]]) -> bool:
    """True when any processor advertises a retire-to-content transition."""
    if os.getenv("VLLM_MLX_ENABLE_THINKING_RETIREMENT_RESUME") != "1":
        return False
    return bool(processors) and any(
        isinstance(getattr(p, "is_retired", None), bool) for p in processors
    )

vllm_mlx.mllm_batch_generator._mark_mtp_attempts_on_primary_responses

_mark_mtp_attempts_on_primary_responses(responses: List[MLLMBatchResponse], attempted_drafts_by_uid: Dict[int, int]) -> None

Mark only responses from steps that actually attempted MTP drafts.

Source code in vllm_mlx/mllm_batch_generator.py

def _mark_mtp_attempts_on_primary_responses(
    responses: List["MLLMBatchResponse"],
    attempted_drafts_by_uid: Dict[int, int],
) -> None:
    """Mark only responses from steps that actually attempted MTP drafts."""
    for response in responses:
        draft_count = attempted_drafts_by_uid.pop(response.uid, 0)
        if draft_count <= 0 or response.finish_reason is not None:
            continue
        response.mtp_attempted = True
        response.mtp_attempted_count = draft_count
    attempted_drafts_by_uid.clear()

vllm_mlx.mllm_batch_generator._drop_retired_processors

_drop_retired_processors(processors: Optional[List[Callable]]) -> tuple[Optional[List[Callable]], int]

Drop retire-capable processors that have completed their work.

Source code in vllm_mlx/mllm_batch_generator.py

def _drop_retired_processors(
    processors: Optional[List[Callable]],
) -> tuple[Optional[List[Callable]], int]:
    """Drop retire-capable processors that have completed their work."""
    if not processors:
        return processors, 0

    remaining = []
    retired_count = 0
    for processor in processors:
        if getattr(processor, "is_retired", False) is True:
            retired_count += 1
            continue
        remaining.append(processor)
    return (remaining or None), retired_count

vllm_mlx.mllm_batch_generator._request_uses_stochastic_sampling

_request_uses_stochastic_sampling(request: Any) -> bool

Return whether a request needs sampler-aware speculative verification.

Greedy (temperature 0) requests are excluded regardless of top_p/top_k/ min_p: _sampling_logprobs() collapses to an argmax delta distribution for temperature 0 and never applies those filters, so a greedy request left at a non-default top_p/top_k/min_p is not actually stochastic.

Source code in vllm_mlx/mllm_batch_generator.py

def _request_uses_stochastic_sampling(request: Any) -> bool:
    """Return whether a request needs sampler-aware speculative verification.

    Greedy (temperature 0) requests are excluded regardless of top_p/top_k/
    min_p: _sampling_logprobs() collapses to an argmax delta distribution for
    temperature 0 and never applies those filters, so a greedy request left at
    a non-default top_p/top_k/min_p is not actually stochastic.
    """
    temperature = getattr(request, "temperature", 0.0)
    if temperature in (0, 0.0):
        return False
    return (
        getattr(request, "top_p", 1.0) < 1.0
        or getattr(request, "top_k", 0) != 0
        or getattr(request, "min_p", 0.0) != 0.0
    )

vllm_mlx.mllm_batch_generator._sampling_logprobs

_sampling_logprobs(logits: array, request: Any) -> array

Match mlx-lm's request sampler in log-probability space.

Speculative decoding compares the post-filter distributions, not the raw target and draft logits. Keep this transformation here rather than reusing a greedy verifier for sampled requests.

Source code in vllm_mlx/mllm_batch_generator.py

def _sampling_logprobs(logits: mx.array, request: Any) -> mx.array:
    """Match mlx-lm's request sampler in log-probability space.

    Speculative decoding compares the post-filter distributions, not the raw
    target and draft logits. Keep this transformation here rather than reusing
    a greedy verifier for sampled requests.
    """
    from mlx_lm.sample_utils import apply_min_p, apply_top_k, apply_top_p

    temperature = getattr(request, "temperature", 0.0)
    top_p = getattr(request, "top_p", 1.0)
    top_k = getattr(request, "top_k", 0)
    min_p = getattr(request, "min_p", 0.0)

    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
    if temperature in (0, 0.0):
        token = mx.argmax(logprobs, axis=-1)
        result = mx.full(logprobs.shape, -float("inf"))
        return mx.put_along_axis(result, token[:, None], 0.0, axis=-1)

    if 0.0 < top_p < 1.0:
        logprobs = apply_top_p(logprobs, top_p)
    if min_p != 0.0:
        logprobs = apply_min_p(logprobs, min_p)
    if top_k > 0:
        logprobs = apply_top_k(logprobs, top_k)

    logprobs = logprobs * (1.0 / temperature)
    return logprobs - mx.logsumexp(logprobs, axis=-1, keepdims=True)

vllm_mlx.mllm_batch_generator._residual_logprobs

_residual_logprobs(target_logprobs: array, draft_logprobs: array) -> array

Return the normalized residual max(target - draft, 0) distribution.

Source code in vllm_mlx/mllm_batch_generator.py

def _residual_logprobs(
    target_logprobs: mx.array,
    draft_logprobs: mx.array,
) -> mx.array:
    """Return the normalized residual max(target - draft, 0) distribution."""
    residual = mx.maximum(mx.exp(target_logprobs) - mx.exp(draft_logprobs), 0.0)
    mass = mx.sum(residual, axis=-1, keepdims=True)
    fallback = target_logprobs
    normalized = mx.where(
        residual > 0,
        mx.log(residual) - mx.log(mass),
        -float("inf"),
    )
    return mx.where(mass > 1e-12, normalized, fallback)

vllm_mlx.mllm_batch_generator._accept_sampled_draft

_accept_sampled_draft(target_logprob: float, draft_logprob: float, uniform_draw: float) -> bool

Apply the exact min(1, p/q) stochastic speculative acceptance rule.

Source code in vllm_mlx/mllm_batch_generator.py

def _accept_sampled_draft(
    target_logprob: float,
    draft_logprob: float,
    uniform_draw: float,
) -> bool:
    """Apply the exact min(1, p/q) stochastic speculative acceptance rule."""
    log_acceptance = target_logprob - draft_logprob
    return log_acceptance >= 0.0 or math.log(max(uniform_draw, 1e-35)) < log_acceptance

vllm_mlx.mllm_batch_generator._cache_eval_tensors

_cache_eval_tensors(cache: List[Any]) -> List[Any]

Return realized tensors that break lazy cache graphs between chunks.

Source code in vllm_mlx/mllm_batch_generator.py

def _cache_eval_tensors(cache: List[Any]) -> List[Any]:
    """Return realized tensors that break lazy cache graphs between chunks."""
    tensors: List[Any] = []
    for c in cache:
        keys = getattr(c, "keys", None)
        values = getattr(c, "values", None)
        if keys is not None or values is not None:
            if keys is not None:
                tensors.append(keys)
            if values is not None:
                tensors.append(values)
            continue

        try:
            state = getattr(c, "state", None)
        except AttributeError:
            state = None
        if state is None:
            continue
        if isinstance(state, (list, tuple)):
            tensors.extend(s for s in state if s is not None)
        else:
            tensors.append(state)
    return tensors

vllm_mlx.mllm_batch_generator._eval_prompt_cache

_eval_prompt_cache(cache: List[Any]) -> None

Evaluate all cache tensors used by hybrid chunked prefill.

Source code in vllm_mlx/mllm_batch_generator.py

def _eval_prompt_cache(cache: List[Any]) -> None:
    """Evaluate all cache tensors used by hybrid chunked prefill."""
    tensors = _cache_eval_tensors(cache)
    if tensors:
        mx.eval(*tensors)

vllm_mlx.mllm_batch_generator._left_pad_prompts

_left_pad_prompts(prompts: List[List[int]], max_length: Optional[int] = None) -> array

Left-pad prompts to uniform length.

Parameters:

• prompts (List[List[int]]) –

  List of token lists

• max_length (Optional[int], default: None ) –

  Target length (computed if not provided)

Returns:

• array –

  Padded prompts as mx.array [batch_size, seq_len]

Source code in vllm_mlx/mllm_batch_generator.py

def _left_pad_prompts(
    prompts: List[List[int]], max_length: Optional[int] = None
) -> mx.array:
    """
    Left-pad prompts to uniform length.

    Args:
        prompts: List of token lists
        max_length: Target length (computed if not provided)

    Returns:
        Padded prompts as mx.array [batch_size, seq_len]
    """
    if max_length is None:
        max_length = max(len(p) for p in prompts)
    return mx.array([[0] * (max_length - len(p)) + list(p) for p in prompts])

vllm_mlx.mllm_batch_generator.install_mtp_mllm

install_mtp_mllm(batch_gen: MLLMBatchGenerator, language_model: Any, num_draft_tokens: int = 1) -> None

Install MTP (Multi-Token Prediction) on an MLLMBatchGenerator.

Adapts the always-advance MTP strategy from scheduler._install_mtp for the MLLM batched generation path. Handles hybrid model caches (BatchKVCache for attention + ArraysCache for recurrent layers).

Flow per generation step: 1. Use skip_state logits/hidden OR run model forward -> sample primary 2. MTP head drafts one token 3. Verify [primary, draft] in one model call (always advances cache) 4. Accept: skip_state from pos 1, defer draft for next step emission Reject: trim KV by 2 + restore RNN state + re-advance with primary 5. Draft is emitted in the NEXT generation step after primary

Source code in vllm_mlx/mllm_batch_generator.py

def install_mtp_mllm(
    batch_gen: "MLLMBatchGenerator",
    language_model: Any,
    num_draft_tokens: int = 1,
) -> None:
    """Install MTP (Multi-Token Prediction) on an MLLMBatchGenerator.

    Adapts the always-advance MTP strategy from scheduler._install_mtp
    for the MLLM batched generation path. Handles hybrid model caches
    (BatchKVCache for attention + ArraysCache for recurrent layers).

    Flow per generation step:
    1. Use skip_state logits/hidden OR run model forward -> sample primary
    2. MTP head drafts one token
    3. Verify [primary, draft] in one model call (always advances cache)
    4. Accept: skip_state from pos 1, defer draft for next step emission
       Reject: trim KV by 2 + restore RNN state + re-advance with primary
    5. Draft is emitted in the NEXT generation step after primary
    """
    from .scheduler import make_sampler

    _orig_step = batch_gen._step
    _draft_sampler = make_sampler(temp=0.0)

    # Skip state belongs to a request, not a batch position. Text-only work
    # may join/leave a continuous batch between decode steps; positional state
    # would then reuse one request's verified logits for another request.
    _skip_state_by_uid: Dict[int, dict] = {}

    # Deferred drafts keyed by UID
    _deferred_drafts: Dict[int, dict] = {}
    _attempted_drafts_by_uid: Dict[int, int] = {}

    # MTP stats. These are intentionally exposed through get_mtp_stats() so
    # /v1/status can distinguish "weights injected" from useful draft work.
    _mtp_stats_lock = threading.Lock()
    _mtp_stats = {"attempted": 0, "accepted": 0, "rejected": 0, "errors": 0}
    _bypass_counts = {
        "prefill": 0,
        "no_active_batch": 0,
        "concurrent_batch": 0,
        "logits_processors": 0,
    }

    def _get_mtp_stats() -> Dict[str, Any]:
        with _mtp_stats_lock:
            attempted = _mtp_stats["attempted"]
            accepted = _mtp_stats["accepted"]
            rejected = _mtp_stats["rejected"]
            errors = _mtp_stats["errors"]
            bypass_counts = dict(_bypass_counts)
        verified = accepted + rejected
        acceptance_rate = accepted / verified if verified > 0 else 0.0
        return {
            "enabled": True,
            "requested_draft_tokens": num_draft_tokens,
            "effective_draft_tokens": 1,
            "mode": "request_local_sampler_aware_verified",
            "attempted": attempted,
            "accepted": accepted,
            "rejected": rejected,
            "errors": errors,
            "acceptance_rate": acceptance_rate,
            "bypass_counts": bypass_counts,
            "bypass_counts_semantics": "per_condition_overlapping_not_total_steps",
        }

    batch_gen.get_mtp_stats = _get_mtp_stats

    def _mtp_step(
        input_tokens: mx.array,
        cache: List[Any],
        logits_processors: Optional[List[Optional[List[Callable]]]] = None,
        output_tokens: Optional[List[List[int]]] = None,
        samplers: Optional[List[Optional[Callable]]] = None,
    ) -> Tuple[mx.array, List[mx.array]]:
        """Extended _step with MTP always-advance strategy."""
        batch_size = input_tokens.shape[0]
        active_requests = (
            list(batch_gen.active_batch.requests)
            if batch_gen.active_batch is not None
            else []
        )
        # Prefill and request-local logits processors remain non-speculative.
        # Sampling and concurrent batches are supported below with per-request
        # distributions and UID-keyed verified state.
        prefill_bypass = input_tokens.shape[1] > 1
        no_active_batch_bypass = batch_gen.active_batch is None
        logits_processors_bypass = logits_processors is not None and any(
            logits_processors
        )
        if prefill_bypass or no_active_batch_bypass or logits_processors_bypass:
            # Keep the descriptions near the guards so operator-facing
            # telemetry stays dynamic instead of duplicating code predicates:
            # prefill=input_tokens.shape[1] > 1
            # no_active_batch=active_batch is None
            # logits_processors=request-local processors are active
            with _mtp_stats_lock:
                if prefill_bypass:
                    _bypass_counts["prefill"] += 1
                if no_active_batch_bypass:
                    _bypass_counts["no_active_batch"] += 1
                if logits_processors_bypass:
                    _bypass_counts["logits_processors"] += 1
            _skip_state_by_uid.clear()
            return _orig_step(
                input_tokens, cache, logits_processors, output_tokens, samplers
            )

        current_uids = list(batch_gen.active_batch.uids)
        skip_entries = [_skip_state_by_uid.pop(uid, None) for uid in current_uids]
        if any(skip_entries) and not all(skip_entries):
            # Batch membership changed since skip state was last populated
            # (e.g. a chunked-prefill request finalizing mid-batch). Skip
            # state cannot be partially reused, so discard it and fall back
            # to a full forward for the whole batch this step -- the same
            # tradeoff the concurrent-rejection path below makes (lose this
            # step's acceleration, keep cache/state correct) rather than
            # crash the batch.
            logger.debug(
                "[MTP-MLLM] batch membership changed since last verified "
                "step; discarding stale skip state and forcing a full forward"
            )
            skip_entries = []

        if skip_entries and all(skip_entries):
            logits = mx.concatenate([entry["logits"] for entry in skip_entries], axis=0)
            hidden_states = mx.concatenate(
                [entry["hidden"] for entry in skip_entries], axis=0
            )
        else:
            # Normal forward with return_hidden
            model_output = language_model(input_tokens, cache=cache, return_hidden=True)
            if isinstance(model_output, tuple):
                logits, hidden_states = model_output
            else:
                return _orig_step(
                    input_tokens, cache, logits_processors, output_tokens, samplers
                )
            logits = logits[:, -1, :]

        # Apply logits processors before sampling
        if logits_processors and output_tokens and any(logits_processors):
            processed_logits = []
            for e in range(batch_size):
                sample_logits = logits[e : e + 1]
                if logits_processors[e]:
                    for processor in logits_processors[e]:
                        sample_logits = processor(
                            mx.array(output_tokens[e]), sample_logits
                        )
                processed_logits.append(sample_logits)
            logits = mx.concatenate(processed_logits, axis=0)

        # Sample primary (use per-request sampler if available)
        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
        if samplers and any(samplers):
            sampled_list = []
            for e in range(logprobs.shape[0]):
                s = samplers[e] if samplers[e] else batch_gen.sampler
                sampled_list.append(s(logprobs[e : e + 1]))
            primary_tokens = mx.concatenate(sampled_list, axis=0)
        else:
            primary_tokens = batch_gen.sampler(logprobs)

        # MTP draft + always-advance verify
        try:
            with _mtp_stats_lock:
                _mtp_stats["attempted"] += 1
            draft_logits = language_model.mtp_forward(
                hidden_states[:, -1:, :],
                primary_tokens[:, None],
                mtp_cache=None,
            )
            draft_logits = draft_logits[:, -1, :]
            sampled_rows = [
                _request_uses_stochastic_sampling(request)
                for request in active_requests
            ]
            uses_stochastic_sampling = any(sampled_rows)
            if uses_stochastic_sampling:
                draft_distribution = mx.concatenate(
                    [
                        _sampling_logprobs(draft_logits[row : row + 1], request)
                        for row, request in enumerate(active_requests)
                    ],
                    axis=0,
                )
                draft_tokens = mx.random.categorical(draft_distribution)
            else:
                draft_logprobs = draft_logits - mx.logsumexp(
                    draft_logits, axis=-1, keepdims=True
                )
                draft_tokens = _draft_sampler(draft_logprobs)
            for uid in current_uids:
                # Current MLLM MTP drafts one token per primary step. Keep this
                # as a count so future multi-token drafters can report >1.
                _attempted_drafts_by_uid[uid] = 1

            # Snapshot RNN state for hybrid models
            _rnn_snapshots = {}
            for _ci, _c in enumerate(cache):
                if not (hasattr(_c, "is_trimmable") and _c.is_trimmable()):
                    if hasattr(_c, "state"):
                        _rnn_snapshots[_ci] = [
                            mx.array(s) if s is not None else None for s in _c.state
                        ]

            # Verify [primary, draft]
            verify_input = mx.concatenate(
                [primary_tokens[:, None], draft_tokens[:, None]], axis=1
            )
            verify_output = language_model(
                verify_input, cache=cache, return_hidden=True
            )
            if isinstance(verify_output, tuple):
                verify_logits, verify_hidden = verify_output
            else:
                verify_logits = verify_output
                verify_hidden = None

            # Verify in each request's sampler space. The old argmax equality
            # check was valid only for greedy decoding and silently bypassed
            # Qwen's normal temperature/top-p/top-k requests.
            draft_list = draft_tokens.tolist()
            residual_tokens_by_uid: Dict[int, int] = {}
            residual_logprobs_by_uid: Dict[int, mx.array] = {}
            if uses_stochastic_sampling:
                verify_distribution = mx.concatenate(
                    [
                        _sampling_logprobs(verify_logits[row : row + 1, 0, :], request)
                        for row, request in enumerate(active_requests)
                    ],
                    axis=0,
                )
                draws = mx.random.uniform(shape=(batch_size,))
                mx.eval(draft_tokens, verify_distribution, draft_distribution, draws)
                all_accepted = True
                for row, uid in enumerate(current_uids):
                    draft_token = int(draft_list[row])
                    accepted = _accept_sampled_draft(
                        float(verify_distribution[row, draft_token].item()),
                        float(draft_distribution[row, draft_token].item()),
                        float(draws[row].item()),
                    )
                    all_accepted = all_accepted and accepted
                    if not accepted and batch_size == 1:
                        residual = _residual_logprobs(
                            verify_distribution[row : row + 1],
                            draft_distribution[row : row + 1],
                        )
                        residual_token = mx.random.categorical(residual)
                        mx.eval(residual_token)
                        residual_tokens_by_uid[uid] = int(residual_token.item())
                        residual_logprobs_by_uid[uid] = residual[0]
            else:
                verify_pred = mx.argmax(verify_logits[:, 0, :], axis=-1)
                mx.eval(verify_pred, draft_tokens)
                all_accepted = verify_pred.tolist() == draft_list

            if all_accepted and verify_hidden is not None:
                # ACCEPT
                mx.async_eval(verify_logits[:, 1, :], verify_hidden[:, -1:, :])
                verify_lp = verify_logits[:, 0, :] - mx.logsumexp(
                    verify_logits[:, 0, :], axis=-1, keepdims=True
                )
                for e in range(batch_size):
                    uid = current_uids[e]
                    _skip_state_by_uid[uid] = {
                        "logits": verify_logits[e : e + 1, 1, :],
                        "hidden": verify_hidden[e : e + 1, -1:, :],
                    }
                    _deferred_drafts[uid] = {
                        "token": draft_list[e],
                        "logprobs": verify_lp[e],
                    }
                with _mtp_stats_lock:
                    _mtp_stats["accepted"] += 1

            else:
                # A batch cache cannot roll back an individual row. On a mixed
                # concurrent rejection, replay only the primary token for every
                # row; this preserves each target distribution and trades that
                # step's acceleration for exact cache state. A single sampled
                # rejection can retain its residual token and still advance.
                sampled_single_reject = (
                    uses_stochastic_sampling
                    and batch_size == 1
                    and bool(residual_tokens_by_uid)
                )
                replay_tokens = primary_tokens
                if sampled_single_reject:
                    residual_token = residual_tokens_by_uid[current_uids[0]]
                    replay_tokens = mx.concatenate(
                        [primary_tokens[:, None], mx.array([[residual_token]])],
                        axis=1,
                    )

                if _rnn_snapshots:
                    # Hybrid model: undo verify then replay the actual emitted
                    # suffix (primary only, or primary + sampled residual).
                    for c in cache:
                        if (
                            hasattr(c, "is_trimmable")
                            and c.is_trimmable()
                            and hasattr(c, "trim")
                        ):
                            c.trim(2)
                    for _ci, _snap in _rnn_snapshots.items():
                        cache[_ci].state = _snap
                    rerun_out = language_model(
                        (
                            replay_tokens
                            if sampled_single_reject
                            else primary_tokens[:, None]
                        ),
                        cache=cache,
                        return_hidden=True,
                    )
                    if isinstance(rerun_out, tuple):
                        rerun_logits, rerun_hidden = rerun_out
                    else:
                        rerun_logits = rerun_out
                        rerun_hidden = None
                    if rerun_hidden is not None:
                        mx.async_eval(rerun_logits[:, -1, :], rerun_hidden[:, -1:, :])
                        for row, uid in enumerate(current_uids):
                            _skip_state_by_uid[uid] = {
                                "logits": rerun_logits[row : row + 1, -1, :],
                                "hidden": rerun_hidden[row : row + 1, -1:, :],
                            }
                    else:
                        _skip_state_by_uid.clear()
                else:
                    # Pure attention caches retain primary after trimming the
                    # speculative draft. A sampled residual is then advanced
                    # explicitly; greedy and concurrent fallbacks reuse the
                    # verified primary state.
                    for c in cache:
                        if (
                            hasattr(c, "is_trimmable")
                            and c.is_trimmable()
                            and hasattr(c, "trim")
                        ):
                            c.trim(1)
                    if sampled_single_reject:
                        residual_token = residual_tokens_by_uid[current_uids[0]]
                        rerun_out = language_model(
                            mx.array([[residual_token]]),
                            cache=cache,
                            return_hidden=True,
                        )
                        if isinstance(rerun_out, tuple):
                            rerun_logits, rerun_hidden = rerun_out
                            # language_model(...) returns (batch, seq, vocab)/
                            # (batch, seq, hidden); reduce logits to the same
                            # 2-D (batch, vocab) convention every other
                            # _skip_state_by_uid write in this function uses.
                            rerun_logits = rerun_logits[:, -1, :]
                            rerun_hidden = rerun_hidden[:, -1:, :]
                        else:
                            rerun_logits, rerun_hidden = rerun_out[:, -1, :], None
                    else:
                        rerun_logits, rerun_hidden = verify_logits[:, 0, :], (
                            verify_hidden[:, 0:1, :]
                            if verify_hidden is not None
                            else None
                        )

                    if rerun_hidden is not None:
                        mx.async_eval(rerun_logits, rerun_hidden)
                        for row, uid in enumerate(current_uids):
                            _skip_state_by_uid[uid] = {
                                "logits": rerun_logits[row : row + 1],
                                "hidden": rerun_hidden[row : row + 1],
                            }
                    else:
                        _skip_state_by_uid.clear()
                for row, uid in enumerate(current_uids):
                    _deferred_drafts.pop(uid, None)
                    if sampled_single_reject:
                        # Report the logprob from the residual distribution the
                        # token was actually drawn from, not the raw unfiltered
                        # target distribution at this position.
                        _deferred_drafts[uid] = {
                            "token": residual_tokens_by_uid[uid],
                            "logprobs": residual_logprobs_by_uid[uid],
                        }
                with _mtp_stats_lock:
                    _mtp_stats["rejected"] += 1

        except Exception as e:
            logger.warning(f"[MTP-MLLM] draft/verify failed: {e}")
            _skip_state_by_uid.clear()
            with _mtp_stats_lock:
                _mtp_stats["errors"] += 1

        # Log MTP stats every 50 steps
        with _mtp_stats_lock:
            acc = _mtp_stats["accepted"]
            rej = _mtp_stats["rejected"]
            err = _mtp_stats["errors"]
        total = acc + rej + err
        if total > 0 and total % 50 == 0:
            rate = acc / (acc + rej) * 100 if (acc + rej) > 0 else 0
            logger.info(
                f"[MTP-MLLM] stats: accepted={acc} rejected={rej} "
                f"errors={err} acceptance={rate:.0f}%"
            )

        return primary_tokens, list(logprobs)

    # Wrap _next to emit deferred MTP drafts
    batch_gen._inner_next = batch_gen._next

    def _mtp_next() -> List[MLLMBatchResponse]:
        """Wrapper around _next that emits deferred MTP draft tokens."""
        if batch_gen.active_batch is None:
            _skip_state_by_uid.clear()
            _deferred_drafts.clear()
            _attempted_drafts_by_uid.clear()

        # `_inner_next` may extend a text-only request into an active batch
        # before the next `_mtp_step` call. That's fine: `_mtp_step` itself
        # tolerates a batch whose UIDs only partially match verified skip
        # state (it discards the stale entries and forces a full forward for
        # that step), so no request needs to be held back here.

        # Save deferred drafts from previous step. The base generator emits
        # its pending input token on this turn, so the verified suffix follows
        # that token in the response stream.
        prev_deferred: Dict[int, dict] = {}
        if batch_gen.active_batch is not None:
            for uid in batch_gen.active_batch.uids:
                if uid in _deferred_drafts:
                    prev_deferred[uid] = _deferred_drafts.pop(uid)

        responses = batch_gen._inner_next()

        if responses:
            _mark_mtp_attempts_on_primary_responses(responses, _attempted_drafts_by_uid)

        # Augment responses with deferred drafts. When there's nothing to
        # augment, `augmented` is just `responses` -- but the trailing
        # skip-state eviction sweep below still needs to run unconditionally
        # so a request that finishes on a step with no pending deferred draft
        # doesn't leave its skip-state entry lingering.
        augmented: List[MLLMBatchResponse] = responses
        draft_end_uids: set = set()

        if prev_deferred and responses:
            augmented = []
            for r in responses:
                uid = r.uid
                augmented.append(r)

                if r.finish_reason is not None:
                    _skip_state_by_uid.pop(uid, None)
                    _deferred_drafts.pop(uid, None)
                    prev_deferred.pop(uid, None)
                    continue

                if uid in prev_deferred:
                    draft_info = prev_deferred.pop(uid)
                    draft_t = draft_info["token"]
                    draft_lp = draft_info["logprobs"]

                    if draft_t in batch_gen.stop_tokens:
                        augmented.append(
                            MLLMBatchResponse(
                                uid=uid,
                                request_id=r.request_id,
                                token=draft_t,
                                logprobs=draft_lp,
                                finish_reason="stop",
                                from_draft=True,
                            )
                        )
                        draft_end_uids.add(uid)
                    else:
                        draft_finish = None
                        batch = batch_gen.active_batch
                        if batch is not None:
                            for e, bu in enumerate(batch.uids):
                                if bu == uid:
                                    batch.num_tokens[e] += 1
                                    batch.requests[e].output_tokens.append(draft_t)
                                    if batch.num_tokens[e] >= batch.max_tokens[e]:
                                        draft_finish = "length"
                                        draft_end_uids.add(uid)
                                    break

                        augmented.append(
                            MLLMBatchResponse(
                                uid=uid,
                                request_id=r.request_id,
                                token=draft_t,
                                logprobs=draft_lp,
                                finish_reason=draft_finish,
                                from_draft=True,
                            )
                        )

            # Store prefix caches for draft-ended sequences BEFORE filtering
            if draft_end_uids and batch_gen.active_batch is not None:
                end_indices = [
                    e
                    for e, u in enumerate(batch_gen.active_batch.uids)
                    if u in draft_end_uids
                ]
                batch_gen._maybe_store_prefix_cache(batch_gen.active_batch, end_indices)

                keep = [
                    e
                    for e, u in enumerate(batch_gen.active_batch.uids)
                    if u not in draft_end_uids
                ]
                if keep:
                    batch_gen.active_batch.filter(keep)
                else:
                    batch_gen.active_batch = None

        active_uids = (
            set(batch_gen.active_batch.uids)
            if batch_gen.active_batch is not None
            else set()
        )
        for uid in list(_skip_state_by_uid):
            if uid not in active_uids:
                _skip_state_by_uid.pop(uid, None)

        return augmented

    batch_gen._step = _mtp_step
    batch_gen._next = _mtp_next

    if num_draft_tokens != 1:
        logger.warning(
            "[MTP-MLLM] num_draft_tokens=%d requested, but the current batched "
            "MLLM MTP path drafts exactly one token per verify step",
            num_draft_tokens,
        )
    logger.info(
        f"[MTP-MLLM] installed with num_draft_tokens={num_draft_tokens}, "
        "effective_draft_tokens=1, request-local sampler-aware verified mode"
    )

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm

install_chunked_prefill_mllm(batch_gen: MLLMBatchGenerator, budget: int = 1024) -> None

Install interleaved prefill/decode on an MLLMBatchGenerator.

When a long text-only request arrives, instead of blocking the entire event loop for 20-60+ seconds during prefill, this processes ONE chunk of the new request's prefill per step() call. Between steps the scheduler yields to the event loop (await asyncio.sleep(0)), so health/status/metrics endpoints remain responsive.

When an active batch is generating, prefill chunks are interleaved with generation steps to keep throughput for existing requests at 30-50 tok/s.

Parameters:

• batch_gen (MLLMBatchGenerator) –

  The MLLMBatchGenerator to patch.

• budget (int, default: 1024 ) –

  Max tokens to prefill per step (chunk size).

Source code in vllm_mlx/mllm_batch_generator.py

def install_chunked_prefill_mllm(
    batch_gen: "MLLMBatchGenerator",
    budget: int = 1024,
) -> None:
    """Install interleaved prefill/decode on an MLLMBatchGenerator.

    When a long text-only request arrives, instead of blocking the entire
    event loop for 20-60+ seconds during prefill, this processes ONE chunk
    of the new request's prefill per ``step()`` call.  Between steps the
    scheduler yields to the event loop (``await asyncio.sleep(0)``), so
    health/status/metrics endpoints remain responsive.

    When an active batch is generating, prefill chunks are interleaved with
    generation steps to keep throughput for existing requests at 30-50 tok/s.

    Args:
        batch_gen: The MLLMBatchGenerator to patch.
        budget: Max tokens to prefill per step (chunk size).
    """
    from mlx_lm.models.cache import make_prompt_cache

    _orig_next = batch_gen._next
    batch_gen._partial = None
    batch_gen._chunked_prefill_budget = budget

    logger.info(
        f"[chunked-prefill-mllm] Installing interleaved prefill/decode "
        f"(budget={budget} tokens/step)"
    )

    def _generation_step() -> List[MLLMBatchResponse]:
        """Run one generation step for the active batch. Returns responses."""
        # Collect pending error responses
        error_responses = list(batch_gen._pending_error_responses)
        batch_gen._pending_error_responses.clear()

        batch = batch_gen.active_batch
        if batch is None:
            return error_responses

        tic = time.perf_counter()
        y, logprobs = batch.y, batch.logprobs
        output_tokens = (
            [req.output_tokens for req in batch.requests]
            if batch.logits_processors
            else None
        )
        batch.y, batch.logprobs = batch_gen._step(
            y[:, None],
            batch.cache,
            batch.logits_processors,
            output_tokens,
            batch.samplers,
        )
        # Synchronous eval — must have results before context switch
        mx.eval(batch.y, batch.logprobs)

        y = y.tolist()
        batch_gen._stats.generation_time += time.perf_counter() - tic

        # Build responses and track finished
        keep_idx = []
        end_idx = []
        responses = []

        for i, (token, uid, request_id, num_tok, max_tok, req) in enumerate(
            zip(
                y,
                batch.uids,
                batch.request_ids,
                batch.num_tokens,
                batch.max_tokens,
                batch.requests,
            )
        ):
            num_tok += 1
            batch.num_tokens[i] = num_tok
            req.num_tokens = num_tok
            req.output_tokens.append(token)

            finish_reason = None

            if token in batch_gen.stop_tokens:
                finish_reason = "stop"
                end_idx.append(i)
            elif num_tok >= max_tok:
                finish_reason = "length"
                end_idx.append(i)
            else:
                keep_idx.append(i)

            if finish_reason is not None:
                batch_gen._prefill_progress.pop(request_id, None)

            responses.append(
                MLLMBatchResponse(
                    uid=uid,
                    request_id=request_id,
                    token=token,
                    logprobs=logprobs[i],
                    finish_reason=finish_reason,
                    prompt_cache=(
                        (lambda idx=i: batch.extract_cache(idx))
                        if finish_reason is not None
                        else None
                    ),
                )
            )

        # Store caches for finished text-only requests BEFORE filtering
        batch_gen._maybe_store_prefix_cache(batch, end_idx)

        # Remove finished requests from batch
        if end_idx:
            if keep_idx:
                batch.filter(keep_idx)
            else:
                batch_gen.active_batch = None

        batch_gen._stats.generation_tokens += len(responses)
        return error_responses + responses

    def _chunked_next() -> List[MLLMBatchResponse]:
        """Interleaved prefill/decode: one prefill chunk + one gen step."""

        # === Phase 1: Continue partial prefill ===
        if batch_gen._partial is not None:
            partial = batch_gen._partial
            req = partial["request"]

            # Abort check
            if req.request_id in batch_gen._aborted_request_ids:
                batch_gen._aborted_request_ids.discard(req.request_id)
                batch_gen._partial = None
                mx.clear_cache()
                batch_gen._prefill_progress.pop(req.request_id, None)
                batch_gen._pending_error_responses.append(
                    MLLMBatchResponse(
                        uid=req.uid,
                        request_id=req.request_id,
                        token=0,
                        logprobs=mx.zeros(1),
                        finish_reason="abort",
                    )
                )
                return _generation_step()

            step = batch_gen._chunked_prefill_budget
            remaining = partial["remaining_ids"]
            remaining_count = remaining.shape[1]

            if remaining_count > step:
                # Process ONE chunk
                tic = time.perf_counter()
                batch_gen.language_model(remaining[:, :step], cache=partial["cache"])
                _eval_prompt_cache(partial["cache"])
                partial["remaining_ids"] = remaining[:, step:]
                partial["processed"] += step
                partial["chunk_count"] += 1
                batch_gen._prefill_progress[req.request_id] = (
                    partial["cached_count"] + partial["processed"],
                    partial["total"],
                )
                batch_gen._stats.prompt_time += time.perf_counter() - tic

                # Periodic memory cleanup
                if partial["chunk_count"] % 4 == 0:
                    mx.clear_cache()

                # Process any short pending requests inline so they
                # don't wait for the entire long prefill to finish.
                # IMPORTANT: Only inline requests whose prompt fits
                # within the chunk budget — longer requests must wait
                # for their own interleaved prefill (Phase 2).
                if batch_gen.unprocessed_requests:
                    _budget = batch_gen._chunked_prefill_budget
                    short_reqs = []
                    for r in batch_gen.unprocessed_requests:
                        if r.images or r.videos:
                            continue
                        if r.input_ids is None:
                            try:
                                batch_gen._preprocess_request(r)
                            except Exception:
                                continue
                        if r.input_ids is not None and r.input_ids.size <= _budget:
                            short_reqs.append(r)
                    if short_reqs:
                        try:
                            new_batch = batch_gen._process_prompts(short_reqs)
                            if new_batch is not None:
                                if batch_gen.active_batch is not None:
                                    batch_gen.active_batch.extend(new_batch)
                                else:
                                    batch_gen.active_batch = new_batch
                        except Exception as e:
                            logger.warning(
                                f"[chunked-prefill-mllm] Failed to process "
                                f"inline short requests: {e}"
                            )

                if batch_gen.active_batch is not None:
                    return _generation_step()
                else:
                    # Idle server — yield to event loop between chunks
                    return []
            else:
                # Last chunk — finalize prefill
                tic = time.perf_counter()
                logits = batch_gen.language_model(remaining, cache=partial["cache"])
                if hasattr(logits, "logits"):
                    logits = logits.logits
                last_logits = logits[:, -1, :]

                # Apply logits processors for first token
                if getattr(req, "logits_processors", None):
                    empty_tokens = mx.array([], dtype=mx.int32)
                    for processor in req.logits_processors:
                        last_logits = processor(empty_tokens, last_logits)

                logprobs = last_logits - mx.logsumexp(
                    last_logits, axis=-1, keepdims=True
                )
                sampled = batch_gen.sampler(logprobs)
                mx.eval(sampled, logprobs)

                batch_gen._prefill_progress[req.request_id] = (
                    partial["total"],
                    partial["total"],
                )
                batch_gen._stats.prompt_time += time.perf_counter() - tic

                # Build single-request batch
                from mlx_lm.sample_utils import make_logits_processors, make_sampler

                req_lp = []
                need_rep = req.repetition_penalty and req.repetition_penalty != 1.0
                need_pres = req.presence_penalty and req.presence_penalty != 0.0
                if need_rep or need_pres:
                    lp_kwargs = {}
                    if need_rep:
                        lp_kwargs["repetition_penalty"] = req.repetition_penalty
                    if need_pres:
                        lp_kwargs["presence_penalty"] = req.presence_penalty
                    req_lp.extend(make_logits_processors(**lp_kwargs))
                if req.logits_processors:
                    req_lp.extend(req.logits_processors)

                req_sampler = None
                if req.top_k != 0 or req.min_p != 0.0:
                    req_sampler = make_sampler(
                        temp=req.temperature,
                        top_p=req.top_p,
                        top_k=req.top_k,
                        min_p=req.min_p,
                    )

                new_batch = MLLMBatch(
                    uids=[req.uid],
                    request_ids=[req.request_id],
                    y=sampled,
                    logprobs=[logprobs.squeeze(0)],
                    max_tokens=[req.max_tokens],
                    num_tokens=[0],
                    cache=partial["cache"],
                    requests=[req],
                    logits_processors=[req_lp] if req_lp else None,
                    samplers=[req_sampler] if req_sampler else None,
                )

                # Extend active batch or set as new
                if batch_gen.active_batch is not None:
                    # Convert per-request cache to batch-compatible format
                    # via merge (same as _process_prompts does)
                    from mlx_lm.models.cache import RotatingKVCache

                    request_cache = partial["cache"]
                    batch_gen._trim_rotating_caches(request_cache)
                    for layer_cache in request_cache:
                        if isinstance(layer_cache, RotatingKVCache):
                            if layer_cache.keys is not None:
                                actual_buf = layer_cache.keys.shape[2]
                                if layer_cache._idx != actual_buf and actual_buf > 0:
                                    layer_cache.keys = layer_cache._temporal_order(
                                        layer_cache.keys
                                    )
                                    layer_cache.values = layer_cache._temporal_order(
                                        layer_cache.values
                                    )
                                    layer_cache._idx = actual_buf

                    # Convert single-request cache to B=1 batch format
                    # so it can be extended into the active batch.
                    merged_cache = [
                        request_cache[layer_idx].merge([request_cache[layer_idx]])
                        for layer_idx in range(len(request_cache))
                    ]
                    new_batch.cache = merged_cache
                    batch_gen.active_batch.extend(new_batch)
                else:
                    # No active batch — convert single-request cache
                    # to B=1 batch format for the new active batch.
                    request_cache = partial["cache"]
                    merged_cache = [
                        request_cache[layer_idx].merge([request_cache[layer_idx]])
                        for layer_idx in range(len(request_cache))
                    ]
                    new_batch.cache = merged_cache
                    batch_gen.active_batch = new_batch

                # Store in prefix cache (prompt-only)
                if batch_gen.prefix_cache is not None and req.input_ids is not None:
                    try:
                        input_ids_list = req.input_ids.reshape(-1).tolist()
                        S = batch_gen._think_suffix_len
                        cache_key = input_ids_list[:-S] if S > 0 else input_ids_list
                        # Trim output: at store time output_count=0, so
                        # trim by S only (matching canonical path's
                        # output_count + S invariant).
                        trim_amount = S
                        store_cache = _trim_cache_offset(partial["cache"], trim_amount)
                        batch_gen.prefix_cache.store(cache_key, store_cache)
                    except Exception as e:
                        logger.warning(
                            f"Failed to store prefix cache after chunked "
                            f"prefill for {req.request_id}: {e}"
                        )

                logger.info(
                    f"[chunked-prefill-mllm] Completed interleaved prefill "
                    f"for {req.request_id[:12]}: "
                    f"{partial['total']} tokens in {partial['chunk_count']} chunks"
                )
                batch_gen._partial = None
                mx.clear_cache()
                return _generation_step()

        # === Phase 2: No partial — check for new requests ===
        batch = batch_gen.active_batch
        num_active = len(batch) if batch else 0

        if batch_gen.unprocessed_requests:
            # Find first text-only request eligible for interleaving
            text_only_req = None
            for r in batch_gen.unprocessed_requests:
                if not r.images and not r.videos:
                    text_only_req = r
                    break

            if text_only_req is not None:
                try:
                    # Preprocess to get input_ids
                    batch_gen._preprocess_request(text_only_req)
                except Exception as e:
                    logger.error(
                        f"Failed to preprocess request "
                        f"{text_only_req.request_id}: {e}"
                    )
                    batch_gen.unprocessed_requests.remove(text_only_req)
                    batch_gen._pending_error_responses.append(
                        MLLMBatchResponse(
                            uid=text_only_req.uid,
                            request_id=text_only_req.request_id,
                            token=0,
                            logprobs=mx.zeros(1),
                            finish_reason="error",
                        )
                    )
                    return _generation_step()

                # Check prefix cache
                input_ids = text_only_req.input_ids
                if input_ids.ndim == 1:
                    input_ids = input_ids[None, :]

                cached_kv = None
                remaining_ids = None
                cached_count = 0
                total_tokens = input_ids.shape[1]

                if batch_gen.prefix_cache is not None:
                    input_ids_list = input_ids.reshape(-1).tolist()
                    S = batch_gen._think_suffix_len
                    lookup_ids = input_ids_list[:-S] if S > 0 else input_ids_list
                    cached_kv, remaining_ids = batch_gen.prefix_cache.fetch(lookup_ids)
                    if cached_kv is not None and S > 0:
                        remaining_ids = list(remaining_ids) + input_ids_list[-S:]

                    # Check for empty rotating cache
                    if cached_kv is not None and batch_gen._has_empty_rotating_cache(
                        cached_kv
                    ):
                        cached_kv = None
                        remaining_ids = None

                if cached_kv is not None and remaining_ids:
                    # Prefix cache hit
                    request_cache = batch_gen._copy_prefix_cache(cached_kv)
                    batch_gen._trim_rotating_caches(request_cache)
                    remaining = mx.array(remaining_ids)[None, :]
                    cached_count = total_tokens - len(remaining_ids)
                    remaining_count = len(remaining_ids)
                elif cached_kv is not None and not remaining_ids:
                    # Exact hit — trim cache by 1 so replaying the last token
                    # produces correct logits (same as _process_prompts path).
                    request_cache = _trim_cache_offset(cached_kv, 1)
                    remaining = input_ids[:, -1:]
                    cached_count = total_tokens - 1
                    remaining_count = 1
                else:
                    # Cache miss — full prefill
                    request_cache = make_prompt_cache(
                        batch_gen.language_model,
                        max_kv_size=batch_gen.max_kv_size or None,
                    )
                    remaining = input_ids
                    cached_count = 0
                    remaining_count = total_tokens

                # Decide: interleave or immediate
                if remaining_count > batch_gen._chunked_prefill_budget:
                    # LONG prompt — start partial (interleaved) prefill
                    logger.info(
                        f"[chunked-prefill-mllm] Starting interleaved prefill "
                        f"for {text_only_req.request_id[:12]}: "
                        f"{remaining_count} remaining tokens "
                        f"(cached={cached_count}, budget={batch_gen._chunked_prefill_budget})"
                    )
                    batch_gen._partial = {
                        "request": text_only_req,
                        "cache": request_cache,
                        "remaining_ids": remaining,
                        "processed": 0,
                        "total": total_tokens,
                        "cached_count": cached_count,
                        "chunk_count": 0,
                    }
                    batch_gen.unprocessed_requests.remove(text_only_req)
                    text_only_req.vision_encoded = True

                    # Process first chunk immediately
                    step = batch_gen._chunked_prefill_budget
                    tic = time.perf_counter()
                    batch_gen.language_model(remaining[:, :step], cache=request_cache)
                    _eval_prompt_cache(request_cache)
                    batch_gen._partial["remaining_ids"] = remaining[:, step:]
                    batch_gen._partial["processed"] = step
                    batch_gen._partial["chunk_count"] = 1
                    batch_gen._prefill_progress[text_only_req.request_id] = (
                        cached_count + step,
                        total_tokens,
                    )
                    batch_gen._stats.prompt_time += time.perf_counter() - tic

                    if num_active > 0:
                        return _generation_step()
                    else:
                        # Idle server — yield to event loop between chunks
                        return []
                # else: SHORT prompt — fall through to _orig_next.
                # _preprocess_request is idempotent for text-only (sets
                # input_ids if not already set); _process_prompts checks
                # input_ids and skips redundant preprocessing.

        # === Phase 3: No partial, no long prompt — original behavior ===
        return _orig_next()

    # Patch remove() to handle partial abort
    _orig_remove = batch_gen.remove

    def _patched_remove(uids: List[int]) -> None:
        if batch_gen._partial is not None:
            if batch_gen._partial["request"].uid in set(uids):
                batch_gen._partial = None
                mx.clear_cache()
        _orig_remove(uids)

    batch_gen.remove = _patched_remove
    batch_gen._next = _chunked_next

    logger.info(f"[chunked-prefill-mllm] Installed (budget={budget} tokens/step)")

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.mllm_batch_generator._processors_can_retire · function

vllm_mlx.mllm_batch_generator._processors_can_retire(processors: Optional[List[Callable]]) -> bool

True when any processor advertises a retire-to-content transition.

Parameters

Name	Type	Required	Default	Description
processors	Optional[List[Callable]]	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: False; bool(processors) and any((isinstance(getattr(p, 'is_retired', None), bool) for p in processors))

Exceptions and behavior

Function _processors_can_retire calls os.getenv, bool, any, isinstance; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L37-L43.

vllm_mlx.mllm_batch_generator._mark_mtp_attempts_on_primary_responses · function

vllm_mlx.mllm_batch_generator._mark_mtp_attempts_on_primary_responses(responses: List['MLLMBatchResponse'], attempted_drafts_by_uid: Dict[int, int]) -> None

Mark only responses from steps that actually attempted MTP drafts.

Parameters

Name	Type	Required	Default	Description
responses	List['MLLMBatchResponse']	yes	none	Required positional or keyword input.
attempted_drafts_by_uid	Dict[int, int]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Function _mark_mtp_attempts_on_primary_responses calls attempted_drafts_by_uid.pop, attempted_drafts_by_uid.clear. No direct raise statement appears in this definition.

View source #L46-L57.

vllm_mlx.mllm_batch_generator._drop_retired_processors · function

vllm_mlx.mllm_batch_generator._drop_retired_processors(processors: Optional[List[Callable]]) -> tuple[Optional[List[Callable]], int]

Drop retire-capable processors that have completed their work.

Parameters

Name	Type	Required	Default	Description
processors	Optional[List[Callable]]	yes	none	Required positional or keyword input.

Returns

• Type: tuple[Optional[List[Callable]], int]
• Direct return expressions: (processors, 0); (remaining or None, retired_count)

Exceptions and behavior

Function _drop_retired_processors calls getattr, remaining.append; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L60-L74.

vllm_mlx.mllm_batch_generator._request_uses_stochastic_sampling · function

vllm_mlx.mllm_batch_generator._request_uses_stochastic_sampling(request: Any) -> bool

Return whether a request needs sampler-aware speculative verification.

Parameters

Name	Type	Required	Default	Description
request	Any	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: False; getattr(request, 'top_p', 1.0) < 1.0 or getattr(request, 'top_k', 0) != 0 or getattr(request, 'min_p', 0.0) != 0.0

Exceptions and behavior

Function _request_uses_stochastic_sampling calls getattr; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L77-L92.

vllm_mlx.mllm_batch_generator._sampling_logprobs · function

vllm_mlx.mllm_batch_generator._sampling_logprobs(logits: mx.array, request: Any) -> mx.array

Match mlx-lm's request sampler in log-probability space.

Parameters

Name	Type	Required	Default	Description
logits	mx.array	yes	none	Required positional or keyword input.
request	Any	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: mx.put_along_axis(result, token[:, None], 0.0, axis=-1); logprobs - mx.logsumexp(logprobs, axis=-1, keepdims=True)

Exceptions and behavior

Function _sampling_logprobs calls getattr, mx.logsumexp, mx.argmax, mx.full; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L95-L123.

vllm_mlx.mllm_batch_generator._residual_logprobs · function

vllm_mlx.mllm_batch_generator._residual_logprobs(target_logprobs: mx.array, draft_logprobs: mx.array) -> mx.array

Return the normalized residual max(target - draft, 0) distribution.

Parameters

Name	Type	Required	Default	Description
target_logprobs	mx.array	yes	none	Required positional or keyword input.
draft_logprobs	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: mx.where(mass > 1e-12, normalized, fallback)

Exceptions and behavior

Function _residual_logprobs calls mx.maximum, mx.exp, mx.sum, mx.where; returns mx.where(mass > 1e-12, normalized, fallback). No direct raise statement appears in this definition.

View source #L126-L139.

vllm_mlx.mllm_batch_generator._accept_sampled_draft · function

vllm_mlx.mllm_batch_generator._accept_sampled_draft(target_logprob: float, draft_logprob: float, uniform_draw: float) -> bool

Apply the exact min(1, p/q) stochastic speculative acceptance rule.

Parameters

Name	Type	Required	Default	Description
target_logprob	float	yes	none	Required positional or keyword input.
draft_logprob	float	yes	none	Required positional or keyword input.
uniform_draw	float	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: log_acceptance >= 0.0 or math.log(max(uniform_draw, 1e-35)) < log_acceptance

Exceptions and behavior

Function _accept_sampled_draft calls math.log, max; returns log_acceptance >= 0.0 or math.log(max(uniform_draw, 1e-35)) < log_acceptance. No direct raise statement appears in this definition.

View source #L142-L149.

vllm_mlx.mllm_batch_generator.PrefillAbortedError · class

vllm_mlx.mllm_batch_generator.PrefillAbortedError(request_id: str)

Raised when a prefill is aborted due to client disconnect.

Parameters

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

Returns

• Constructs: vllm_mlx.mllm_batch_generator.PrefillAbortedError

Exceptions and behavior

Class PrefillAbortedError derives from Exception and declares 1 direct member(s). No direct raise statement appears in this definition.

View source #L152-L157.

vllm_mlx.mllm_batch_generator.PrefillAbortedError.__init__ · method

vllm_mlx.mllm_batch_generator.PrefillAbortedError.__init__(request_id: str) -> not annotated

Method PrefillAbortedError.__init__ updates self.request_id; calls super().__init__, super.

Parameters

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

Returns

• Type: not annotated

Exceptions and behavior

Method PrefillAbortedError.__init__ updates self.request_id; calls super().__init__, super. No direct raise statement appears in this definition.

View source #L155-L157.

vllm_mlx.mllm_batch_generator._cache_eval_tensors · function

vllm_mlx.mllm_batch_generator._cache_eval_tensors(cache: List[Any]) -> List[Any]

Return realized tensors that break lazy cache graphs between chunks.

Parameters

Name	Type	Required	Default	Description
cache	List[Any]	yes	none	Required positional or keyword input.

Returns

• Type: List[Any]
• Direct return expressions: tensors

Exceptions and behavior

Function _cache_eval_tensors calls getattr, tensors.append, isinstance, tensors.extend; returns tensors. No direct raise statement appears in this definition.

View source #L160-L183.

vllm_mlx.mllm_batch_generator._eval_prompt_cache · function

vllm_mlx.mllm_batch_generator._eval_prompt_cache(cache: List[Any]) -> None

Evaluate all cache tensors used by hybrid chunked prefill.

Parameters

Name	Type	Required	Default	Description
cache	List[Any]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Function _eval_prompt_cache calls _cache_eval_tensors, mx.eval. No direct raise statement appears in this definition.

View source #L186-L190.

vllm_mlx.mllm_batch_generator.MLLMBatchRequest · class

vllm_mlx.mllm_batch_generator.MLLMBatchRequest(uid: int, request_id: str, prompt: str, images: Optional[List[str]] = None, videos: Optional[List[str]] = None, audio: Optional[List[str]] = None, max_tokens: int = 256, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 0, min_p: float = 0.0, presence_penalty: float = 0.0, repetition_penalty: float = 1.0, logits_processors: Optional[List[Callable]] = None, input_ids: Optional[mx.array] = None, pixel_values: Optional[mx.array] = None, attention_mask: Optional[mx.array] = None, image_grid_thw: Optional[mx.array] = None, extra_kwargs: Dict[str, Any] = field(default_factory=dict), is_text_only: bool = False, num_tokens: int = 0, output_tokens: List[int] = field(default_factory=list), vision_encoded: bool = False, cross_attention_states: Optional[Any] = None, encoder_outputs: Optional[Any] = None)

Request data for MLLM batch processing.

Parameters

Name	Type	Required	Default	Description
uid	int	yes	none	Required constructor field.
request_id	str	yes	none	Required constructor field.
prompt	str	yes	none	Required constructor field.
images	Optional[List[str]]	no	None	Optional constructor field; defaults to None.
videos	Optional[List[str]]	no	None	Optional constructor field; defaults to None.
audio	Optional[List[str]]	no	None	Optional constructor field; defaults to None.
max_tokens	int	no	256	Optional constructor field; defaults to 256.
temperature	float	no	0.7	Optional constructor field; defaults to 0.7.
top_p	float	no	0.9	Optional constructor field; defaults to 0.9.
top_k	int	no	0	Optional constructor field; defaults to 0.
min_p	float	no	0.0	Optional constructor field; defaults to 0.0.
presence_penalty	float	no	0.0	Optional constructor field; defaults to 0.0.
repetition_penalty	float	no	1.0	Optional constructor field; defaults to 1.0.
logits_processors	Optional[List[Callable]]	no	None	Optional constructor field; defaults to None.
input_ids	Optional[mx.array]	no	None	Optional constructor field; defaults to None.
pixel_values	Optional[mx.array]	no	None	Optional constructor field; defaults to None.
attention_mask	Optional[mx.array]	no	None	Optional constructor field; defaults to None.
image_grid_thw	Optional[mx.array]	no	None	Optional constructor field; defaults to None.
extra_kwargs	Dict[str, Any]	no	field(default_factory=dict)	Optional constructor field; defaults to field(default_factory=dict).
is_text_only	bool	no	False	Optional constructor field; defaults to False.
num_tokens	int	no	0	Optional constructor field; defaults to 0.
output_tokens	List[int]	no	field(default_factory=list)	Optional constructor field; defaults to field(default_factory=list).
vision_encoded	bool	no	False	Optional constructor field; defaults to False.
cross_attention_states	Optional[Any]	no	None	Optional constructor field; defaults to None.
encoder_outputs	Optional[Any]	no	None	Optional constructor field; defaults to None.

Returns

• Constructs: vllm_mlx.mllm_batch_generator.MLLMBatchRequest

Exceptions and behavior

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

View source #L194-L237.

vllm_mlx.mllm_batch_generator.MLLMBatchResponse · class

vllm_mlx.mllm_batch_generator.MLLMBatchResponse(uid: int, request_id: str, token: int, logprobs: mx.array, finish_reason: Optional[str] = None, prompt_cache: Optional[Callable[[], List[Any]]] = None, from_draft: bool = False, mtp_attempted: bool = False, mtp_attempted_count: int = 0)

Response from a batch generation step.

Parameters

Name	Type	Required	Default	Description
uid	int	yes	none	Required constructor field.
request_id	str	yes	none	Required constructor field.
token	int	yes	none	Required constructor field.
logprobs	mx.array	yes	none	Required constructor field.
finish_reason	Optional[str]	no	None	Optional constructor field; defaults to None.
prompt_cache	Optional[Callable[[], List[Any]]]	no	None	Optional constructor field; defaults to None.
from_draft	bool	no	False	Optional constructor field; defaults to False.
mtp_attempted	bool	no	False	Optional constructor field; defaults to False.
mtp_attempted_count	int	no	0	Optional constructor field; defaults to 0.

Returns

• Constructs: vllm_mlx.mllm_batch_generator.MLLMBatchResponse

Exceptions and behavior

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

View source #L241-L256.

vllm_mlx.mllm_batch_generator.MLLMBatch · class

vllm_mlx.mllm_batch_generator.MLLMBatch(uids: List[int], request_ids: List[str], y: mx.array, logprobs: List[mx.array], max_tokens: List[int], num_tokens: List[int], cache: List[Any], requests: List[MLLMBatchRequest], logits_processors: Optional[List[Optional[List[Callable]]]] = None, samplers: Optional[List[Optional[Callable]]] = None)

Represents an active batch of MLLM requests.

Parameters

Name	Type	Required	Default	Description
uids	List[int]	yes	none	Required constructor field.
request_ids	List[str]	yes	none	Required constructor field.
y	mx.array	yes	none	Required constructor field.
logprobs	List[mx.array]	yes	none	Required constructor field.
max_tokens	List[int]	yes	none	Required constructor field.
num_tokens	List[int]	yes	none	Required constructor field.
cache	List[Any]	yes	none	Required constructor field.
requests	List[MLLMBatchRequest]	yes	none	Required constructor field.
logits_processors	Optional[List[Optional[List[Callable]]]]	no	None	Optional constructor field; defaults to None.
samplers	Optional[List[Optional[Callable]]]	no	None	Optional constructor field; defaults to None.

Returns

• Constructs: vllm_mlx.mllm_batch_generator.MLLMBatch

Exceptions and behavior

Class MLLMBatch declares 4 direct member(s). No direct raise statement appears in this definition.

View source #L260-L392.

vllm_mlx.mllm_batch_generator.MLLMBatch.__len__ · method

vllm_mlx.mllm_batch_generator.MLLMBatch.__len__() -> int

Method MLLMBatch.__len__ calls len; returns len(self.uids).

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: len(self.uids)

Exceptions and behavior

Method MLLMBatch.__len__ calls len; returns len(self.uids). No direct raise statement appears in this definition.

View source #L279-L280.

vllm_mlx.mllm_batch_generator.MLLMBatch.filter · method

vllm_mlx.mllm_batch_generator.MLLMBatch.filter(keep_idx: List[int]) -> None

Filter batch to keep only requests at specified indices.

Parameters

Name	Type	Required	Default	Description
keep_idx	List[int]	yes	none	Indices of requests to keep

Returns

• Type: None

Exceptions and behavior

Method MLLMBatch.filter updates self.uids, self.request_ids, self.logprobs, self.max_tokens; calls mx.array, hasattr, c.filter. No direct raise statement appears in this definition.

View source #L282-L306.

vllm_mlx.mllm_batch_generator.MLLMBatch.extend · method

vllm_mlx.mllm_batch_generator.MLLMBatch.extend(other: 'MLLMBatch') -> None

Extend this batch with another batch.

Parameters

Name	Type	Required	Default	Description
other	'MLLMBatch'	yes	none	Batch to merge into this one

Returns

• Type: None

Exceptions and behavior

Method MLLMBatch.extend updates self.y, self.logits_processors, self.samplers; calls self.uids.extend, self.request_ids.extend, mx.concatenate, self.logprobs.extend. No direct raise statement appears in this definition.

View source #L308-L351.

vllm_mlx.mllm_batch_generator.MLLMBatch.extract_cache · method

vllm_mlx.mllm_batch_generator.MLLMBatch.extract_cache(idx: int) -> List[Any]

Extract cache for a single request (for prefix caching).

Parameters

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

Returns

• Type: List[Any]
• Direct return expressions: result

Exceptions and behavior

Method MLLMBatch.extract_cache calls hasattr, result.append, isinstance, RotatingKVCache; returns result. No direct raise statement appears in this definition.

View source #L353-L392.

vllm_mlx.mllm_batch_generator.MLLMBatchStats · class

vllm_mlx.mllm_batch_generator.MLLMBatchStats()

Statistics for MLLM batch generation.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.mllm_batch_generator.MLLMBatchStats

Exceptions and behavior

Class MLLMBatchStats declares 4 direct member(s). No direct raise statement appears in this definition.

View source #L395-L436.

vllm_mlx.mllm_batch_generator.MLLMBatchStats.__init__ · method

vllm_mlx.mllm_batch_generator.MLLMBatchStats.__init__() -> not annotated

Method MLLMBatchStats.__init__ updates self.prompt_tokens, self.prompt_time, self.generation_tokens, self.generation_time.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Method MLLMBatchStats.__init__ updates self.prompt_tokens, self.prompt_time, self.generation_tokens, self.generation_time. No direct raise statement appears in this definition.

View source #L398-L405.

vllm_mlx.mllm_batch_generator.MLLMBatchStats.prompt_tps · method

vllm_mlx.mllm_batch_generator.MLLMBatchStats.prompt_tps() -> float

Return measured multimodal prompt throughput in tokens per second.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: 0; self.prompt_tokens / self.prompt_time

Exceptions and behavior

Method MLLMBatchStats.prompt_tps has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L408-L413.

vllm_mlx.mllm_batch_generator.MLLMBatchStats.generation_tps · method

vllm_mlx.mllm_batch_generator.MLLMBatchStats.generation_tps() -> float

Return measured decode throughput in tokens per second.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: 0; self.generation_tokens / self.generation_time

Exceptions and behavior

Method MLLMBatchStats.generation_tps has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L416-L421.

vllm_mlx.mllm_batch_generator.MLLMBatchStats.to_dict · method

vllm_mlx.mllm_batch_generator.MLLMBatchStats.to_dict() -> Dict[str, Any]

Return token, timing, vision, and peak-memory statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: Dict[str, Any]
• Direct return expressions: {'prompt_tokens': self.prompt_tokens, 'prompt_time': self.prompt_time, 'prompt_tps': self.prompt_tps, 'generation_token…

Exceptions and behavior

Method MLLMBatchStats.to_dict returns {'prompt_tokens': self.prompt_tokens, 'prompt_time': self.prompt_time, 'prompt_tps': self.prompt_tps, 'generation_token…. No direct raise statement appears in this definition.

View source #L423-L436.

vllm_mlx.mllm_batch_generator._left_pad_prompts · function

vllm_mlx.mllm_batch_generator._left_pad_prompts(prompts: List[List[int]], max_length: Optional[int] = None) -> mx.array

Left-pad prompts to uniform length.

Parameters

Name	Type	Required	Default	Description
prompts	List[List[int]]	yes	none	List of token lists
max_length	Optional[int]	no	None	Target length (computed if not provided)

Returns

• Type: mx.array
• Direct return expressions: mx.array([[0] * (max_length - len(p)) + list(p) for p in prompts])

Exceptions and behavior

Function _left_pad_prompts calls max, len, mx.array, list; returns mx.array([[0] * (max_length - len(p)) + list(p) for p in prompts]). No direct raise statement appears in this definition.

View source #L439-L454.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator · class

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator(model: nn.Module, processor: Any, mm_processor: Optional[MultimodalProcessor] = None, max_tokens: int = 256, stop_tokens: Optional[set] = None, sampler: Optional[Callable[[mx.array], mx.array]] = None, prefill_batch_size: int = 4, completion_batch_size: int = 16, prefill_step_size: int = 1024, enable_vision_cache: bool = True, vision_cache_size: int = 100, prefix_cache_config: Optional[MemoryCacheConfig] = None, max_kv_size: int = 0)

Batch generator for Vision Language Models.

Parameters

Name	Type	Required	Default	Description
model	nn.Module	yes	none	The VLM model (must have model.language_model)
processor	Any	yes	none	The VLM processor for tokenization and image processing
mm_processor	Optional[MultimodalProcessor]	no	None	Optional MultimodalProcessor for input preparation
max_tokens	int	no	256	Default max tokens per request
stop_tokens	Optional[set]	no	None	Set of stop token IDs
sampler	Optional[Callable[[mx.array], mx.array]]	no	None	Sampling function (default: argmax)
prefill_batch_size	int	no	4	Max requests to prefill together
completion_batch_size	int	no	16	Max requests for completion batching
prefill_step_size	int	no	1024	Tokens to process per prefill step
enable_vision_cache	bool	no	True	Enable vision embedding caching
vision_cache_size	int	no	100	Max entries in vision cache
prefix_cache_config	Optional[MemoryCacheConfig]	no	None	Config for KV prefix cache (text-only requests)
max_kv_size	int	no	0	Maximum KV cache size per sequence (0 = unbounded)

Returns

• Constructs: vllm_mlx.mllm_batch_generator.MLLMBatchGenerator

Exceptions and behavior

Class MLLMBatchGenerator declares 26 direct member(s). No direct raise statement appears in this definition.

View source #L457-L2042.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.__init__ · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.__init__(model: nn.Module, processor: Any, mm_processor: Optional[MultimodalProcessor] = None, max_tokens: int = 256, stop_tokens: Optional[set] = None, sampler: Optional[Callable[[mx.array], mx.array]] = None, prefill_batch_size: int = 4, completion_batch_size: int = 16, prefill_step_size: int = 1024, enable_vision_cache: bool = True, vision_cache_size: int = 100, prefix_cache_config: Optional[MemoryCacheConfig] = None, max_kv_size: int = 0) -> not annotated

Initialize MLLM batch generator.

Parameters

Name	Type	Required	Default	Description
model	nn.Module	yes	none	The VLM model (must have model.language_model)
processor	Any	yes	none	The VLM processor for tokenization and image processing
mm_processor	Optional[MultimodalProcessor]	no	None	Optional MultimodalProcessor for input preparation
max_tokens	int	no	256	Default max tokens per request
stop_tokens	Optional[set]	no	None	Set of stop token IDs
sampler	Optional[Callable[[mx.array], mx.array]]	no	None	Sampling function (default: argmax)
prefill_batch_size	int	no	4	Max requests to prefill together
completion_batch_size	int	no	16	Max requests for completion batching
prefill_step_size	int	no	1024	Tokens to process per prefill step
enable_vision_cache	bool	no	True	Enable vision embedding caching
vision_cache_size	int	no	100	Max entries in vision cache
prefix_cache_config	Optional[MemoryCacheConfig]	no	None	Config for KV prefix cache (text-only requests)
max_kv_size	int	no	0	Maximum KV cache size per sequence (0 = unbounded)

Returns

• Type: not annotated

Exceptions and behavior

Method MLLMBatchGenerator.__init__ updates self.model, self.processor, self.mm_processor, self.max_kv_size; calls getattr, hasattr, logger.info, logger.warning. No direct raise statement appears in this definition.

View source #L484-L632.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._normalize_chat_template_for_prefix_cache · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._normalize_chat_template_for_prefix_cache() -> None

Patch chat template so historical assistant turns are prefix-stable.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MLLMBatchGenerator._normalize_chat_template_for_prefix_cache updates self.processor.chat_template; calls getattr, re.sub, hasattr, logger.info; returns None. No direct raise statement appears in this definition.

View source #L634-L697.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._compute_think_suffix_len · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._compute_think_suffix_len() -> int

Compute how many extra tokens enable_thinking=True adds at the END.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: 0; max(0, suffix_len)

Exceptions and behavior

Method MLLMBatchGenerator._compute_think_suffix_len calls getattr, hasattr, applicator.apply_chat_template, text_with.endswith; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L699-L758.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.close · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.close() -> None

Release resources and reset wired limit.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method MLLMBatchGenerator.close updates self._old_wired_limit; calls mx.synchronize, mx.set_wired_limit. No direct raise statement appears in this definition.

View source #L760-L765.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.abort_prefill · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.abort_prefill(request_id: str) -> None

Signal that a request's prefill should be aborted.

Parameters

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

Returns

• Type: None

Exceptions and behavior

Method MLLMBatchGenerator.abort_prefill calls self._aborted_request_ids.add, logger.info. No direct raise statement appears in this definition.

View source #L767-L775.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.schedule_removal · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.schedule_removal(uids: List[int]) -> None

Thread-safe deferred removal of UIDs from the batch.

Parameters

Name	Type	Required	Default	Description
uids	List[int]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method MLLMBatchGenerator.schedule_removal calls self._pending_removal_uids.update. No direct raise statement appears in this definition.

View source #L777-L789.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.process_pending_removals · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.process_pending_removals() -> None

Remove any UIDs enqueued via :meth:schedule_removal.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MLLMBatchGenerator.process_pending_removals updates self._pending_removal_uids; calls set, list, self.remove; returns None. No direct raise statement appears in this definition.

View source #L791-L808.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.__del__ · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.__del__() -> not annotated

Method MLLMBatchGenerator.__del__ calls self.close.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Method MLLMBatchGenerator.__del__ calls self.close. No direct raise statement appears in this definition.

View source #L810-L814.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.insert · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.insert(requests: List[MLLMBatchRequest]) -> List[int]

Insert requests for batch processing.

Parameters

Name	Type	Required	Default	Description
requests	List[MLLMBatchRequest]	yes	none	List of MLLMBatchRequest to process

Returns

• Type: List[int]
• Direct return expressions: uids

Exceptions and behavior

Method MLLMBatchGenerator.insert updates self.uid_counter, self.unprocessed_requests; calls self.unprocessed_requests.append, uids.append, sorted, logger.debug; returns uids. No direct raise statement appears in this definition.

View source #L816-L846.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.remove · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.remove(uids: List[int]) -> None

Remove requests from processing.

Parameters

Name	Type	Required	Default	Description
uids	List[int]	yes	none	List of UIDs to remove

Returns

• Type: None

Exceptions and behavior

Method MLLMBatchGenerator.remove updates self.active_batch, self.unprocessed_requests; calls set, enumerate, self.active_batch.filter. No direct raise statement appears in this definition.

View source #L848-L870.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._preprocess_request · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._preprocess_request(request: MLLMBatchRequest) -> None

Preprocess a single MLLM request (vision encoding).

Parameters

Name	Type	Required	Default	Description
request	MLLMBatchRequest	yes	none	Request to preprocess

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MLLMBatchGenerator._preprocess_request updates self._stats.num_images_processed, self._stats.vision_encoding_time; calls time.perf_counter, process_image_input, all_images.append, logger.warning; returns None. No direct raise statement appears in this definition.

View source #L872-L1023.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._copy_prefix_cache · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._copy_prefix_cache(cache_list) -> not annotated

Create shallow copies of cache objects to prevent mutation of stored prefix cache.

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: copies

Exceptions and behavior

Method MLLMBatchGenerator._copy_prefix_cache calls isinstance, RotatingKVCache, copies.append, KVCache; returns copies. No direct raise statement appears in this definition.

View source #L1026-L1054.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._has_empty_rotating_cache · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._has_empty_rotating_cache(cache_list) -> not annotated

Check if any RotatingKVCache layer has no data (keys=None).

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: True; False

Exceptions and behavior

Method MLLMBatchGenerator._has_empty_rotating_cache calls isinstance; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1057-L1069.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._trim_rotating_caches · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._trim_rotating_caches(cache_list) -> not annotated

Trim RotatingKVCache buffers restored from prefix cache.

Parameters

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

Returns

• Type: not annotated

Exceptions and behavior

Method MLLMBatchGenerator._trim_rotating_caches calls isinstance, layer_cache._trim, min, logger.warning. No direct raise statement appears in this definition.

View source #L1072-L1106.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._run_chunked_text_prefill · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._run_chunked_text_prefill(request: MLLMBatchRequest, cache: List[Any]) -> mx.array

Run prefill in chunks for text-only requests, reporting real progress.

Parameters

Name	Type	Required	Default	Description
request	MLLMBatchRequest	yes	none	Required positional or keyword input.
cache	List[Any]	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: output.logits; output

Exceptions and behavior

Method MLLMBatchGenerator._run_chunked_text_prefill calls self.language_model, request.extra_kwargs.clear, hasattr, logger.info; can raise PrefillAbortedError; has 2 explicit return paths. Directly raised exceptions: PrefillAbortedError.

View source #L1108-L1206.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._run_vision_encoding · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._run_vision_encoding(request: MLLMBatchRequest, cache: Optional[List[Any]] = None) -> mx.array

Run the initial VLM forward pass to encode vision and get first logits.

Parameters

Name	Type	Required	Default	Description
request	MLLMBatchRequest	yes	none	Preprocessed request with input_ids and pixel_values
cache	Optional[List[Any]]	no	None	KV cache list for the language model. If provided, the language model writes its KV state directly into this cache during the forward pass.

Returns

• Type: mx.array
• Direct return expressions: output.logits; output

Exceptions and behavior

Method MLLMBatchGenerator._run_vision_encoding calls dict, self.model, request.extra_kwargs.clear, hasattr; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1208-L1258.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._process_prompts · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._process_prompts(requests: List[MLLMBatchRequest]) -> MLLMBatch

Process a batch of requests through vision encoding and initial prefill.

Parameters

Name	Type	Required	Default	Description
requests	List[MLLMBatchRequest]	yes	none	Requests to process

Returns

• Type: MLLMBatch
• Direct return expressions: None; MLLMBatch(uids=[req.uid for req in requests], request_ids=[req.request_id for req in requests], y=y, logprobs=all_logpr…

Exceptions and behavior

Method MLLMBatchGenerator._process_prompts updates self._stats.prompt_tokens, self._stats.prompt_time; calls time.perf_counter, self._preprocess_request, logger.error, type; can raise PrefillAbortedError; has 2 explicit return paths. Directly raised exceptions: PrefillAbortedError.

View source #L1260-L1682.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._process_prompts._sample_first_token · nested function

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._process_prompts._sample_first_token(req: MLLMBatchRequest, logits: mx.array) -> not annotated

Nested Function MLLMBatchGenerator._process_prompts._sample_first_token calls logits_processors_by_request.get, mx.array, processor, mx.logsumexp; returns (sampled, logprobs).

Parameters

Name	Type	Required	Default	Description
req	MLLMBatchRequest	yes	none	Required positional or keyword input.
logits	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: (sampled, logprobs)

Exceptions and behavior

Nested Function MLLMBatchGenerator._process_prompts._sample_first_token calls logits_processors_by_request.get, mx.array, processor, mx.logsumexp; returns (sampled, logprobs). No direct raise statement appears in this definition.

View source #L1348-L1362.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._step · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._step(input_tokens: mx.array, cache: List[Any], logits_processors: Optional[List[Optional[List[Callable]]]] = None, output_tokens: Optional[List[List[int]]] = None, samplers: Optional[List[Optional[Callable]]] = None) -> Tuple[mx.array, List[mx.array]]

Run one generation step through the language model.

Parameters

Name	Type	Required	Default	Description
input_tokens	mx.array	yes	none	Input tokens [batch_size, 1] or [batch_size]
cache	List[Any]	yes	none	BatchKVCache for the language model
logits_processors	Optional[List[Optional[List[Callable]]]]	no	None	Per-request logits processors (e.g. repetition penalty)
output_tokens	Optional[List[List[int]]]	no	None	Per-request generated tokens so far (needed by processors)
samplers	Optional[List[Optional[Callable]]]	no	None	Per-request sampler functions (for top_k/min_p)

Returns

• Type: Tuple[mx.array, List[mx.array]]
• Direct return expressions: (sampled, list(logprobs))

Exceptions and behavior

Method MLLMBatchGenerator._step calls self.language_model, hasattr, any, range; returns (sampled, list(logprobs)). No direct raise statement appears in this definition.

View source #L1684-L1746.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._next · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._next() -> List[MLLMBatchResponse]

Internal next() implementation.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[MLLMBatchResponse]
• Direct return expressions: []; error_responses; error_responses + responses

Exceptions and behavior

Method MLLMBatchGenerator._next updates self.active_batch, self.unprocessed_requests, self._stats.prompt_time, self._stats.generation_time; calls time.perf_counter, len, self._process_prompts, logger.error; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L1748-L1964.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.next · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.next() -> List[MLLMBatchResponse]

Generate next token for all requests in the batch.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[MLLMBatchResponse]
• Direct return expressions: self._next()

Exceptions and behavior

Method MLLMBatchGenerator.next calls mx.stream, self._next; returns self._next(). No direct raise statement appears in this definition.

View source #L1966-L1974.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.stats · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.stats() -> MLLMBatchStats

Get generation statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: MLLMBatchStats
• Direct return expressions: self._stats

Exceptions and behavior

Method MLLMBatchGenerator.stats updates self._stats.peak_memory; calls mx.get_peak_memory; returns self._stats. No direct raise statement appears in this definition.

View source #L1976-L1984.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._maybe_store_prefix_cache · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator._maybe_store_prefix_cache(batch: MLLMBatch, end_indices: List[int]) -> None

Store KV caches for finished text-only requests into prefix cache.

Parameters

Name	Type	Required	Default	Description
batch	MLLMBatch	yes	none	Required positional or keyword input.
end_indices	List[int]	yes	none	Required positional or keyword input.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MLLMBatchGenerator._maybe_store_prefix_cache calls batch.extract_cache, req.input_ids.reshape(-1).tolist, req.input_ids.reshape, _trim_cache_offset; returns None. No direct raise statement appears in this definition.

View source #L1986-L2014.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_prefill_progress · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_prefill_progress(request_id: str) -> Optional[Tuple[int, int]]

Return (processed_tokens, total_tokens) or None.

Parameters

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

Returns

• Type: Optional[Tuple[int, int]]
• Direct return expressions: self._prefill_progress.get(request_id)

Exceptions and behavior

Method MLLMBatchGenerator.get_prefill_progress calls self._prefill_progress.get; returns self._prefill_progress.get(request_id). No direct raise statement appears in this definition.

View source #L2016-L2018.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_vision_cache_stats · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_vision_cache_stats() -> Dict[str, Any]

Get vision cache statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: Dict[str, Any]
• Direct return expressions: self.vision_cache.get_stats()

Exceptions and behavior

Method MLLMBatchGenerator.get_vision_cache_stats calls self.vision_cache.get_stats; returns self.vision_cache.get_stats(). No direct raise statement appears in this definition.

View source #L2020-L2022.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_prefix_cache_stats · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.get_prefix_cache_stats() -> Dict[str, Any]

Get KV prefix cache statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: Dict[str, Any]
• Direct return expressions: self.prefix_cache.get_stats(); {'hits': 0, 'misses': 0, 'hit_rate': 0.0, 'evictions': 0, 'tokens_saved': 0, 'current_memory_mb': 0.0, 'max_memory_mb':…

Exceptions and behavior

Method MLLMBatchGenerator.get_prefix_cache_stats calls self.prefix_cache.get_stats; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2024-L2038.

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.has_pending · method

vllm_mlx.mllm_batch_generator.MLLMBatchGenerator.has_pending() -> bool

Check if there are pending or active requests.

Parameters

This callable has no explicit inputs.

Returns

• Type: bool
• Direct return expressions: bool(self.unprocessed_requests or self.active_batch)

Exceptions and behavior

Method MLLMBatchGenerator.has_pending calls bool; returns bool(self.unprocessed_requests or self.active_batch). No direct raise statement appears in this definition.

View source #L2040-L2042.

vllm_mlx.mllm_batch_generator.install_mtp_mllm · function

vllm_mlx.mllm_batch_generator.install_mtp_mllm(batch_gen: 'MLLMBatchGenerator', language_model: Any, num_draft_tokens: int = 1) -> None

Install MTP (Multi-Token Prediction) on an MLLMBatchGenerator.

Parameters

Name	Type	Required	Default	Description
batch_gen	'MLLMBatchGenerator'	yes	none	Required positional or keyword input.
language_model	Any	yes	none	Required positional or keyword input.
num_draft_tokens	int	no	1	Optional positional or keyword input; defaults to 1.

Returns

• Type: None

Exceptions and behavior

Function install_mtp_mllm calls make_sampler, threading.Lock, logger.warning, logger.info. No direct raise statement appears in this definition.

View source #L2045-L2590.

vllm_mlx.mllm_batch_generator.install_mtp_mllm._get_mtp_stats · nested function

vllm_mlx.mllm_batch_generator.install_mtp_mllm._get_mtp_stats() -> Dict[str, Any]

Nested Function install_mtp_mllm._get_mtp_stats calls dict; returns {'enabled': True, 'requested_draft_tokens': num_draft_tokens, 'effective_draft_tokens': 1, 'mode': 'request_local_sampl….

Parameters

This callable has no explicit inputs.

Returns

• Type: Dict[str, Any]
• Direct return expressions: {'enabled': True, 'requested_draft_tokens': num_draft_tokens, 'effective_draft_tokens': 1, 'mode': 'request_local_sampl…

Exceptions and behavior

Nested Function install_mtp_mllm._get_mtp_stats calls dict; returns {'enabled': True, 'requested_draft_tokens': num_draft_tokens, 'effective_draft_tokens': 1, 'mode': 'request_local_sampl…. No direct raise statement appears in this definition.

View source #L2089-L2110.

vllm_mlx.mllm_batch_generator.install_mtp_mllm._mtp_step · nested function

vllm_mlx.mllm_batch_generator.install_mtp_mllm._mtp_step(input_tokens: mx.array, cache: List[Any], logits_processors: Optional[List[Optional[List[Callable]]]] = None, output_tokens: Optional[List[List[int]]] = None, samplers: Optional[List[Optional[Callable]]] = None) -> Tuple[mx.array, List[mx.array]]

Extended _step with MTP always-advance strategy.

Parameters

Name	Type	Required	Default	Description
input_tokens	mx.array	yes	none	Required positional or keyword input.
cache	List[Any]	yes	none	Required positional or keyword input.
logits_processors	Optional[List[Optional[List[Callable]]]]	no	None	Optional positional or keyword input; defaults to None.
output_tokens	Optional[List[List[int]]]	no	None	Optional positional or keyword input; defaults to None.
samplers	Optional[List[Optional[Callable]]]	no	None	Optional positional or keyword input; defaults to None.

Returns

• Type: Tuple[mx.array, List[mx.array]]
• Direct return expressions: _orig_step(input_tokens, cache, logits_processors, output_tokens, samplers); (primary_tokens, list(logprobs))

Exceptions and behavior

Nested Function install_mtp_mllm._mtp_step calls list, any, _skip_state_by_uid.clear, _orig_step; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2114-L2455.

vllm_mlx.mllm_batch_generator.install_mtp_mllm._mtp_next · nested function

vllm_mlx.mllm_batch_generator.install_mtp_mllm._mtp_next() -> List[MLLMBatchResponse]

Wrapper around _next that emits deferred MTP draft tokens.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[MLLMBatchResponse]
• Direct return expressions: augmented

Exceptions and behavior

Nested Function install_mtp_mllm._mtp_next calls _skip_state_by_uid.clear, _deferred_drafts.clear, _attempted_drafts_by_uid.clear, _deferred_drafts.pop; returns augmented. No direct raise statement appears in this definition.

View source #L2460-L2576.

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm · function

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm(batch_gen: 'MLLMBatchGenerator', budget: int = 1024) -> None

Install interleaved prefill/decode on an MLLMBatchGenerator.

Parameters

Name	Type	Required	Default	Description
batch_gen	'MLLMBatchGenerator'	yes	none	The MLLMBatchGenerator to patch.
budget	int	no	1024	Max tokens to prefill per step (chunk size).

Returns

• Type: None

Exceptions and behavior

Function install_chunked_prefill_mllm calls logger.info. No direct raise statement appears in this definition.

View source #L2593-L3073.

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm._generation_step · nested function

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm._generation_step() -> List[MLLMBatchResponse]

Run one generation step for the active batch.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[MLLMBatchResponse]
• Direct return expressions: error_responses; error_responses + responses

Exceptions and behavior

Nested Function install_chunked_prefill_mllm._generation_step calls list, batch_gen._pending_error_responses.clear, time.perf_counter, batch_gen._step; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2623-L2713.

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm._chunked_next · nested function

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm._chunked_next() -> List[MLLMBatchResponse]

Interleaved prefill/decode: one prefill chunk + one gen step.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[MLLMBatchResponse]
• Direct return expressions: _generation_step(); []; _orig_next()

Exceptions and behavior

Nested Function install_chunked_prefill_mllm._chunked_next calls batch_gen._aborted_request_ids.discard, mx.clear_cache, batch_gen._prefill_progress.pop, batch_gen._pending_error_responses.append; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L2715-L3058.

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm._patched_remove · nested function

vllm_mlx.mllm_batch_generator.install_chunked_prefill_mllm._patched_remove(uids: List[int]) -> None

Nested Function install_chunked_prefill_mllm._patched_remove calls set, mx.clear_cache, _orig_remove.

Parameters

Name	Type	Required	Default	Description
uids	List[int]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Nested Function install_chunked_prefill_mllm._patched_remove calls set, mx.clear_cache, _orig_remove. No direct raise statement appears in this definition.

View source #L3063-L3068.

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
_processors_can_retire	function	_processors_can_retire(processors: Optional[List[Callable]]) -> bool	True when any processor advertises a retire-to-content transition.	#L37-L43
_mark_mtp_attempts_on_primary_responses	function	_mark_mtp_attempts_on_primary_responses(responses: List['MLLMBatchResponse'], attempted_drafts_by_uid: Dict[int, int]) -> None	Mark only responses from steps that actually attempted MTP drafts.	#L46-L57
_drop_retired_processors	function	_drop_retired_processors(processors: Optional[List[Callable]]) -> tuple[Optional[List[Callable]], int]	Drop retire-capable processors that have completed their work.	#L60-L74
_request_uses_stochastic_sampling	function	_request_uses_stochastic_sampling(request: Any) -> bool	Return whether a request needs sampler-aware speculative verification.	#L77-L92
_sampling_logprobs	function	_sampling_logprobs(logits: mx.array, request: Any) -> mx.array	Match mlx-lm's request sampler in log-probability space.	#L95-L123
_residual_logprobs	function	_residual_logprobs(target_logprobs: mx.array, draft_logprobs: mx.array) -> mx.array	Return the normalized residual max(target - draft, 0) distribution.	#L126-L139
_accept_sampled_draft	function	_accept_sampled_draft(target_logprob: float, draft_logprob: float, uniform_draw: float) -> bool	Apply the exact min(1, p/q) stochastic speculative acceptance rule.	#L142-L149
PrefillAbortedError	class	PrefillAbortedError(request_id: str)	Raised when a prefill is aborted due to client disconnect.	#L152-L157
PrefillAbortedError.__init__	method	PrefillAbortedError.__init__(request_id: str) -> not annotated	Method PrefillAbortedError.__init__ updates self.request_id; calls super().__init__, super.	#L155-L157
_cache_eval_tensors	function	_cache_eval_tensors(cache: List[Any]) -> List[Any]	Return realized tensors that break lazy cache graphs between chunks.	#L160-L183
_eval_prompt_cache	function	_eval_prompt_cache(cache: List[Any]) -> None	Evaluate all cache tensors used by hybrid chunked prefill.	#L186-L190
MLLMBatchRequest	class	MLLMBatchRequest(uid: int, request_id: str, prompt: str, images: Optional[List[str]] = None, videos: Optional[List[str]] = None, audio: Optional[List[str]] = None, max_tokens: int = 256, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 0, min_p: float = 0.0, presence_penalty: float = 0.0, repetition_penalty: float = 1.0, logits_processors: Optional[List[Callable]] = None, input_ids: Optional[mx.array] = None, pixel_values: Optional[mx.array] = None, attention_mask: Optional[mx.array] = None, image_grid_thw: Optional[mx.array] = None, extra_kwargs: Dict[str, Any] = field(default_factory=dict), is_text_only: bool = False, num_tokens: int = 0, output_tokens: List[int] = field(default_factory=list), vision_encoded: bool = False, cross_attention_states: Optional[Any] = None, encoder_outputs: Optional[Any] = None)	Request data for MLLM batch processing.	#L194-L237
MLLMBatchResponse	class	MLLMBatchResponse(uid: int, request_id: str, token: int, logprobs: mx.array, finish_reason: Optional[str] = None, prompt_cache: Optional[Callable[[], List[Any]]] = None, from_draft: bool = False, mtp_attempted: bool = False, mtp_attempted_count: int = 0)	Response from a batch generation step.	#L241-L256
MLLMBatch	class	MLLMBatch(uids: List[int], request_ids: List[str], y: mx.array, logprobs: List[mx.array], max_tokens: List[int], num_tokens: List[int], cache: List[Any], requests: List[MLLMBatchRequest], logits_processors: Optional[List[Optional[List[Callable]]]] = None, samplers: Optional[List[Optional[Callable]]] = None)	Represents an active batch of MLLM requests.	#L260-L392
MLLMBatch.__len__	method	MLLMBatch.__len__() -> int	Method MLLMBatch.__len__ calls len; returns len(self.uids).	#L279-L280
MLLMBatch.filter	method	MLLMBatch.filter(keep_idx: List[int]) -> None	Filter batch to keep only requests at specified indices.	#L282-L306
MLLMBatch.extend	method	MLLMBatch.extend(other: 'MLLMBatch') -> None	Extend this batch with another batch.	#L308-L351
MLLMBatch.extract_cache	method	MLLMBatch.extract_cache(idx: int) -> List[Any]	Extract cache for a single request (for prefix caching).	#L353-L392
MLLMBatchStats	class	MLLMBatchStats()	Statistics for MLLM batch generation.	#L395-L436
MLLMBatchStats.__init__	method	MLLMBatchStats.__init__() -> not annotated	Method MLLMBatchStats.__init__ updates self.prompt_tokens, self.prompt_time, self.generation_tokens, self.generation_time.	#L398-L405
MLLMBatchStats.prompt_tps	method	MLLMBatchStats.prompt_tps() -> float	Return measured multimodal prompt throughput in tokens per second.	#L408-L413
MLLMBatchStats.generation_tps	method	MLLMBatchStats.generation_tps() -> float	Return measured decode throughput in tokens per second.	#L416-L421
MLLMBatchStats.to_dict	method	MLLMBatchStats.to_dict() -> Dict[str, Any]	Return token, timing, vision, and peak-memory statistics.	#L423-L436
_left_pad_prompts	function	_left_pad_prompts(prompts: List[List[int]], max_length: Optional[int] = None) -> mx.array	Left-pad prompts to uniform length.	#L439-L454
MLLMBatchGenerator	class	MLLMBatchGenerator(model: nn.Module, processor: Any, mm_processor: Optional[MultimodalProcessor] = None, max_tokens: int = 256, stop_tokens: Optional[set] = None, sampler: Optional[Callable[[mx.array], mx.array]] = None, prefill_batch_size: int = 4, completion_batch_size: int = 16, prefill_step_size: int = 1024, enable_vision_cache: bool = True, vision_cache_size: int = 100, prefix_cache_config: Optional[MemoryCacheConfig] = None, max_kv_size: int = 0)	Batch generator for Vision Language Models.	#L457-L2042
MLLMBatchGenerator.__init__	method	MLLMBatchGenerator.__init__(model: nn.Module, processor: Any, mm_processor: Optional[MultimodalProcessor] = None, max_tokens: int = 256, stop_tokens: Optional[set] = None, sampler: Optional[Callable[[mx.array], mx.array]] = None, prefill_batch_size: int = 4, completion_batch_size: int = 16, prefill_step_size: int = 1024, enable_vision_cache: bool = True, vision_cache_size: int = 100, prefix_cache_config: Optional[MemoryCacheConfig] = None, max_kv_size: int = 0) -> not annotated	Initialize MLLM batch generator.	#L484-L632
MLLMBatchGenerator._normalize_chat_template_for_prefix_cache	method	MLLMBatchGenerator._normalize_chat_template_for_prefix_cache() -> None	Patch chat template so historical assistant turns are prefix-stable.	#L634-L697
MLLMBatchGenerator._compute_think_suffix_len	method	MLLMBatchGenerator._compute_think_suffix_len() -> int	Compute how many extra tokens enable_thinking=True adds at the END.	#L699-L758
MLLMBatchGenerator.close	method	MLLMBatchGenerator.close() -> None	Release resources and reset wired limit.	#L760-L765
MLLMBatchGenerator.abort_prefill	method	MLLMBatchGenerator.abort_prefill(request_id: str) -> None	Signal that a request's prefill should be aborted.	#L767-L775
MLLMBatchGenerator.schedule_removal	method	MLLMBatchGenerator.schedule_removal(uids: List[int]) -> None	Thread-safe deferred removal of UIDs from the batch.	#L777-L789
MLLMBatchGenerator.process_pending_removals	method	MLLMBatchGenerator.process_pending_removals() -> None	Remove any UIDs enqueued via :meth:schedule_removal.	#L791-L808
MLLMBatchGenerator.__del__	method	MLLMBatchGenerator.__del__() -> not annotated	Method MLLMBatchGenerator.__del__ calls self.close.	#L810-L814
MLLMBatchGenerator.insert	method	MLLMBatchGenerator.insert(requests: List[MLLMBatchRequest]) -> List[int]	Insert requests for batch processing.	#L816-L846
MLLMBatchGenerator.remove	method	MLLMBatchGenerator.remove(uids: List[int]) -> None	Remove requests from processing.	#L848-L870
MLLMBatchGenerator._preprocess_request	method	MLLMBatchGenerator._preprocess_request(request: MLLMBatchRequest) -> None	Preprocess a single MLLM request (vision encoding).	#L872-L1023
MLLMBatchGenerator._copy_prefix_cache	method	MLLMBatchGenerator._copy_prefix_cache(cache_list) -> not annotated	Create shallow copies of cache objects to prevent mutation of stored prefix cache.	#L1026-L1054
MLLMBatchGenerator._has_empty_rotating_cache	method	MLLMBatchGenerator._has_empty_rotating_cache(cache_list) -> not annotated	Check if any RotatingKVCache layer has no data (keys=None).	#L1057-L1069
MLLMBatchGenerator._trim_rotating_caches	method	MLLMBatchGenerator._trim_rotating_caches(cache_list) -> not annotated	Trim RotatingKVCache buffers restored from prefix cache.	#L1072-L1106
MLLMBatchGenerator._run_chunked_text_prefill	method	MLLMBatchGenerator._run_chunked_text_prefill(request: MLLMBatchRequest, cache: List[Any]) -> mx.array	Run prefill in chunks for text-only requests, reporting real progress.	#L1108-L1206
MLLMBatchGenerator._run_vision_encoding	method	MLLMBatchGenerator._run_vision_encoding(request: MLLMBatchRequest, cache: Optional[List[Any]] = None) -> mx.array	Run the initial VLM forward pass to encode vision and get first logits.	#L1208-L1258
MLLMBatchGenerator._process_prompts	method	MLLMBatchGenerator._process_prompts(requests: List[MLLMBatchRequest]) -> MLLMBatch	Process a batch of requests through vision encoding and initial prefill.	#L1260-L1682
MLLMBatchGenerator._process_prompts._sample_first_token	nested function	MLLMBatchGenerator._process_prompts._sample_first_token(req: MLLMBatchRequest, logits: mx.array) -> not annotated	Nested Function MLLMBatchGenerator._process_prompts._sample_first_token calls logits_processors_by_request.get, mx.array, processor, mx.logsumexp; returns (sampled, logprobs).	#L1348-L1362
MLLMBatchGenerator._step	method	MLLMBatchGenerator._step(input_tokens: mx.array, cache: List[Any], logits_processors: Optional[List[Optional[List[Callable]]]] = None, output_tokens: Optional[List[List[int]]] = None, samplers: Optional[List[Optional[Callable]]] = None) -> Tuple[mx.array, List[mx.array]]	Run one generation step through the language model.	#L1684-L1746
MLLMBatchGenerator._next	method	MLLMBatchGenerator._next() -> List[MLLMBatchResponse]	Internal next() implementation.	#L1748-L1964
MLLMBatchGenerator.next	method	MLLMBatchGenerator.next() -> List[MLLMBatchResponse]	Generate next token for all requests in the batch.	#L1966-L1974
MLLMBatchGenerator.stats	method	MLLMBatchGenerator.stats() -> MLLMBatchStats	Get generation statistics.	#L1976-L1984
MLLMBatchGenerator._maybe_store_prefix_cache	method	MLLMBatchGenerator._maybe_store_prefix_cache(batch: MLLMBatch, end_indices: List[int]) -> None	Store KV caches for finished text-only requests into prefix cache.	#L1986-L2014
MLLMBatchGenerator.get_prefill_progress	method	MLLMBatchGenerator.get_prefill_progress(request_id: str) -> Optional[Tuple[int, int]]	Return (processed_tokens, total_tokens) or None.	#L2016-L2018
MLLMBatchGenerator.get_vision_cache_stats	method	MLLMBatchGenerator.get_vision_cache_stats() -> Dict[str, Any]	Get vision cache statistics.	#L2020-L2022
MLLMBatchGenerator.get_prefix_cache_stats	method	MLLMBatchGenerator.get_prefix_cache_stats() -> Dict[str, Any]	Get KV prefix cache statistics.	#L2024-L2038
MLLMBatchGenerator.has_pending	method	MLLMBatchGenerator.has_pending() -> bool	Check if there are pending or active requests.	#L2040-L2042
install_mtp_mllm	function	install_mtp_mllm(batch_gen: 'MLLMBatchGenerator', language_model: Any, num_draft_tokens: int = 1) -> None	Install MTP (Multi-Token Prediction) on an MLLMBatchGenerator.	#L2045-L2590
install_mtp_mllm._get_mtp_stats	nested function	install_mtp_mllm._get_mtp_stats() -> Dict[str, Any]	Nested Function install_mtp_mllm._get_mtp_stats calls dict; returns {'enabled': True, 'requested_draft_tokens': num_draft_tokens, 'effective_draft_tokens': 1, 'mode': 'request_local_sampl….	#L2089-L2110
install_mtp_mllm._mtp_step	nested function	install_mtp_mllm._mtp_step(input_tokens: mx.array, cache: List[Any], logits_processors: Optional[List[Optional[List[Callable]]]] = None, output_tokens: Optional[List[List[int]]] = None, samplers: Optional[List[Optional[Callable]]] = None) -> Tuple[mx.array, List[mx.array]]	Extended _step with MTP always-advance strategy.	#L2114-L2455
install_mtp_mllm._mtp_next	nested function	install_mtp_mllm._mtp_next() -> List[MLLMBatchResponse]	Wrapper around _next that emits deferred MTP draft tokens.	#L2460-L2576
install_chunked_prefill_mllm	function	install_chunked_prefill_mllm(batch_gen: 'MLLMBatchGenerator', budget: int = 1024) -> None	Install interleaved prefill/decode on an MLLMBatchGenerator.	#L2593-L3073
install_chunked_prefill_mllm._generation_step	nested function	install_chunked_prefill_mllm._generation_step() -> List[MLLMBatchResponse]	Run one generation step for the active batch.	#L2623-L2713
install_chunked_prefill_mllm._chunked_next	nested function	install_chunked_prefill_mllm._chunked_next() -> List[MLLMBatchResponse]	Interleaved prefill/decode: one prefill chunk + one gen step.	#L2715-L3058
install_chunked_prefill_mllm._patched_remove	nested function	install_chunked_prefill_mllm._patched_remove(uids: List[int]) -> None	Nested Function install_chunked_prefill_mllm._patched_remove calls set, mx.clear_cache, _orig_remove.	#L3063-L3068