vllm_mlx.scheduler

Scheduler for vllm-mlx continuous batching.

View the complete module source at #L1-L3518.

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

Scheduler for vllm-mlx continuous batching.

This module provides a Scheduler class that manages request scheduling using mlx-lm's BatchGenerator for efficient continuous batching.

The scheduler follows vLLM's design with: - Waiting queue for pending requests - Running set for active requests - Continuous batching via BatchGenerator

vllm_mlx.scheduler.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.scheduler.CACHE_CORRUPTION_PATTERNS module-attribute

CACHE_CORRUPTION_PATTERNS = ["'NoneType' object is not subscriptable", 'cache', 'BatchKVCache']

vllm_mlx.scheduler.SchedulingPolicy

Bases: Enum

Scheduling policy for request ordering.

vllm_mlx.scheduler.SchedulingPolicy.FCFS class-attribute instance-attribute

FCFS = 'fcfs'

vllm_mlx.scheduler.SchedulingPolicy.PRIORITY class-attribute instance-attribute

PRIORITY = 'priority'

vllm_mlx.scheduler.SchedulerConfig dataclass

SchedulerConfig(max_num_seqs: int = 256, max_num_batched_tokens: int = 8192, policy: SchedulingPolicy = FCFS, prefill_batch_size: int = 8, completion_batch_size: int = 32, prefill_step_size: int = 2048, mllm_prefill_step_size: Optional[int] = None, enable_prefix_cache: bool = True, prefix_cache_size: int = 100, use_memory_aware_cache: bool = True, cache_memory_mb: Optional[int] = None, cache_memory_percent: float = 0.2, kv_cache_quantization: bool = False, kv_cache_quantization_bits: int = 8, kv_cache_quantization_group_size: int = 64, kv_cache_min_quantize_tokens: int = 256, use_paged_cache: bool = False, paged_cache_block_size: int = 64, max_cache_blocks: int = 1000, chunked_prefill_tokens: int = 0, mid_prefill_save_interval: int = 8192, ssd_cache_dir: Optional[str] = None, ssd_cache_max_gb: float = 10.0, max_kv_size: int = 0, enable_mtp: bool = False, mtp_num_draft_tokens: int = 1, mtp_optimistic: bool = False)

Configuration for the scheduler.

vllm_mlx.scheduler.SchedulerConfig.max_num_seqs class-attribute instance-attribute

max_num_seqs: int = 256

vllm_mlx.scheduler.SchedulerConfig.max_num_batched_tokens class-attribute instance-attribute

max_num_batched_tokens: int = 8192

vllm_mlx.scheduler.SchedulerConfig.policy class-attribute instance-attribute

policy: SchedulingPolicy = SchedulingPolicy.FCFS

vllm_mlx.scheduler.SchedulerConfig.prefill_batch_size class-attribute instance-attribute

prefill_batch_size: int = 8

vllm_mlx.scheduler.SchedulerConfig.completion_batch_size class-attribute instance-attribute

completion_batch_size: int = 32

vllm_mlx.scheduler.SchedulerConfig.prefill_step_size class-attribute instance-attribute

prefill_step_size: int = 2048

vllm_mlx.scheduler.SchedulerConfig.mllm_prefill_step_size class-attribute instance-attribute

mllm_prefill_step_size: Optional[int] = None

vllm_mlx.scheduler.SchedulerConfig.enable_prefix_cache class-attribute instance-attribute

enable_prefix_cache: bool = True

vllm_mlx.scheduler.SchedulerConfig.prefix_cache_size class-attribute instance-attribute

prefix_cache_size: int = 100

vllm_mlx.scheduler.SchedulerConfig.use_memory_aware_cache class-attribute instance-attribute

use_memory_aware_cache: bool = True

vllm_mlx.scheduler.SchedulerConfig.cache_memory_mb class-attribute instance-attribute

cache_memory_mb: Optional[int] = None

vllm_mlx.scheduler.SchedulerConfig.cache_memory_percent class-attribute instance-attribute

cache_memory_percent: float = 0.2

vllm_mlx.scheduler.SchedulerConfig.kv_cache_quantization class-attribute instance-attribute

kv_cache_quantization: bool = False

vllm_mlx.scheduler.SchedulerConfig.kv_cache_quantization_bits class-attribute instance-attribute

kv_cache_quantization_bits: int = 8

vllm_mlx.scheduler.SchedulerConfig.kv_cache_quantization_group_size class-attribute instance-attribute

kv_cache_quantization_group_size: int = 64

vllm_mlx.scheduler.SchedulerConfig.kv_cache_min_quantize_tokens class-attribute instance-attribute

kv_cache_min_quantize_tokens: int = 256

vllm_mlx.scheduler.SchedulerConfig.use_paged_cache class-attribute instance-attribute

use_paged_cache: bool = False

vllm_mlx.scheduler.SchedulerConfig.paged_cache_block_size class-attribute instance-attribute

paged_cache_block_size: int = 64

vllm_mlx.scheduler.SchedulerConfig.max_cache_blocks class-attribute instance-attribute

max_cache_blocks: int = 1000

vllm_mlx.scheduler.SchedulerConfig.chunked_prefill_tokens class-attribute instance-attribute

chunked_prefill_tokens: int = 0

vllm_mlx.scheduler.SchedulerConfig.mid_prefill_save_interval class-attribute instance-attribute

mid_prefill_save_interval: int = 8192

vllm_mlx.scheduler.SchedulerConfig.ssd_cache_dir class-attribute instance-attribute

ssd_cache_dir: Optional[str] = None

vllm_mlx.scheduler.SchedulerConfig.ssd_cache_max_gb class-attribute instance-attribute

ssd_cache_max_gb: float = 10.0

vllm_mlx.scheduler.SchedulerConfig.max_kv_size class-attribute instance-attribute

max_kv_size: int = 0

vllm_mlx.scheduler.SchedulerConfig.enable_mtp class-attribute instance-attribute

enable_mtp: bool = False

vllm_mlx.scheduler.SchedulerConfig.mtp_num_draft_tokens class-attribute instance-attribute

mtp_num_draft_tokens: int = 1

vllm_mlx.scheduler.SchedulerConfig.mtp_optimistic class-attribute instance-attribute

mtp_optimistic: bool = False

vllm_mlx.scheduler.SchedulerConfig.__post_init__

__post_init__() -> None

Source code in vllm_mlx/scheduler.py

def __post_init__(self) -> None:
    if self.mllm_prefill_step_size is not None and self.mllm_prefill_step_size <= 0:
        raise ValueError("mllm_prefill_step_size must be > 0 when provided")

vllm_mlx.scheduler.SchedulerOutput dataclass

SchedulerOutput(scheduled_request_ids: List[str] = list(), num_scheduled_tokens: int = 0, finished_request_ids: Set[str] = set(), outputs: List[RequestOutput] = list(), has_work: bool = False)

Output from a scheduling step.

Contains information about what was scheduled and results.

vllm_mlx.scheduler.SchedulerOutput.scheduled_request_ids class-attribute instance-attribute

scheduled_request_ids: List[str] = field(default_factory=list)

vllm_mlx.scheduler.SchedulerOutput.num_scheduled_tokens class-attribute instance-attribute

num_scheduled_tokens: int = 0

vllm_mlx.scheduler.SchedulerOutput.finished_request_ids class-attribute instance-attribute

finished_request_ids: Set[str] = field(default_factory=set)

vllm_mlx.scheduler.SchedulerOutput.outputs class-attribute instance-attribute

outputs: List[RequestOutput] = field(default_factory=list)

vllm_mlx.scheduler.SchedulerOutput.has_work class-attribute instance-attribute

has_work: bool = False

vllm_mlx.scheduler._MTPStatsState dataclass

_MTPStatsState(counters: Dict[str, int] = (lambda: {'attempted': 0, 'accepted': 0, 'rejected': 0, 'errors': 0})(), bypass_counts: Dict[str, int] = (lambda: {'prefill': 0, 'no_active_batch': 0, 'cache_mismatch': 0})(), lock: Any = Lock())

Cumulative native-MTP counters shared across generator instances.

vllm_mlx.scheduler._MTPStatsState.counters class-attribute instance-attribute

counters: Dict[str, int] = field(default_factory=lambda: {'attempted': 0, 'accepted': 0, 'rejected': 0, 'errors': 0})

vllm_mlx.scheduler._MTPStatsState.bypass_counts class-attribute instance-attribute

bypass_counts: Dict[str, int] = field(default_factory=lambda: {'prefill': 0, 'no_active_batch': 0, 'cache_mismatch': 0})

vllm_mlx.scheduler._MTPStatsState.lock class-attribute instance-attribute

lock: Any = field(default_factory=Lock)

vllm_mlx.scheduler.Scheduler

Scheduler(model: Any, tokenizer: Any, config: Optional[SchedulerConfig] = None)

Scheduler for continuous batching using mlx-lm BatchGenerator.

This scheduler manages the lifecycle of requests: 1. Requests arrive and are added to the waiting queue 2. Scheduler moves requests from waiting to running (via BatchGenerator) 3. BatchGenerator processes all running requests together 4. Finished requests are removed and outputs returned

The key insight is that mlx-lm's BatchGenerator already implements continuous batching at the token level, so we use it as the backend.

Initialize the scheduler.

Parameters:

• model (Any) –

  The MLX model

• tokenizer (Any) –

  The tokenizer

• config (Optional[SchedulerConfig], default: None ) –

  Scheduler configuration

Source code in vllm_mlx/scheduler.py

def __init__(
    self,
    model: Any,
    tokenizer: Any,
    config: Optional[SchedulerConfig] = None,
):
    """
    Initialize the scheduler.

    Args:
        model: The MLX model
        tokenizer: The tokenizer
        config: Scheduler configuration
    """
    self.model = model
    self.tokenizer = tokenizer
    self.config = config or SchedulerConfig()

    # Detect if tokenizer is a processor (MLLM) and get the actual tokenizer
    self._actual_tokenizer = self._get_actual_tokenizer(tokenizer)

    # Per-request streaming detokenizers for UTF-8-safe incremental decode
    self._detokenizer_pool: Dict[str, Any] = {}

    # Request management - following vLLM's design
    self.waiting: deque[Request] = deque()  # Waiting queue (FCFS)
    self.running: Dict[str, Request] = {}  # Running requests by ID
    self.requests: Dict[str, Request] = {}  # All requests by ID
    self.finished_req_ids: Set[str] = set()  # Recently finished

    # Mapping between our request IDs and BatchGenerator UIDs
    self.request_id_to_uid: Dict[str, int] = {}
    self.uid_to_request_id: Dict[int, str] = {}

    # BatchGenerator - the actual batching engine
    self.batch_generator: Optional[BatchGenerator] = None
    self._current_sampler_params: Optional[Tuple] = None

    # Prefix cache for KV state reuse
    self.prefix_cache: Optional[PrefixCacheManager] = None
    self.memory_aware_cache: Optional[MemoryAwarePrefixCache] = None
    self.paged_cache_manager: Optional[PagedCacheManager] = None
    self.block_aware_cache: Optional[BlockAwarePrefixCache] = None
    self._ssd_tier: Optional[SSDCacheTier] = None

    if self.config.enable_prefix_cache:
        if self.config.use_paged_cache:
            # Use paged cache for memory efficiency
            self.paged_cache_manager = PagedCacheManager(
                block_size=self.config.paged_cache_block_size,
                max_blocks=self.config.max_cache_blocks,
            )
            self.block_aware_cache = BlockAwarePrefixCache(
                model=model,
                paged_cache_manager=self.paged_cache_manager,
            )
            logger.info(
                f"Paged cache enabled: block_size={self.config.paged_cache_block_size}, "
                f"max_blocks={self.config.max_cache_blocks}"
            )
        elif self.config.use_memory_aware_cache:
            # Use memory-aware cache (recommended for large models)
            cache_config = MemoryCacheConfig(
                max_memory_mb=self.config.cache_memory_mb,
                max_memory_percent=self.config.cache_memory_percent,
                kv_quantize=self.config.kv_cache_quantization,
                kv_bits=self.config.kv_cache_quantization_bits,
                kv_group_size=self.config.kv_cache_quantization_group_size,
                kv_min_quantize_tokens=self.config.kv_cache_min_quantize_tokens,
            )
            self.memory_aware_cache = MemoryAwarePrefixCache(
                model=model,
                config=cache_config,
            )
            logger.info(
                f"Memory-aware cache enabled: "
                f"limit={self.memory_aware_cache.memory_limit_mb:.1f}MB"
            )

            if self.config.ssd_cache_dir is not None:
                ssd_config = SSDCacheConfig(
                    cache_dir=self.config.ssd_cache_dir,
                    max_size_gb=self.config.ssd_cache_max_gb,
                )
                self._ssd_tier = SSDCacheTier(ssd_config)
                self._ssd_tier.start_writer()
                self._ssd_tier.reconcile()
                self.memory_aware_cache.set_ssd_tier(self._ssd_tier)
                logger.info(
                    f"SSD cache tier enabled: dir={self.config.ssd_cache_dir}, "
                    f"max={self.config.ssd_cache_max_gb}GB"
                )
        else:
            # Use legacy entry-count based prefix cache
            self.prefix_cache = PrefixCacheManager(
                model=model,
                max_entries=self.config.prefix_cache_size,
            )
            logger.info(
                f"Prefix cache enabled with max_entries={self.config.prefix_cache_size}"
            )

    # Thread-safe set for deferred aborts (main thread → executor thread)
    # CPython GIL guarantees set.add() and `x in set` are atomic.
    self._pending_abort_ids: Set[str] = set()

    # Statistics
    self.num_requests_processed = 0
    self.total_prompt_tokens = 0
    self.total_completion_tokens = 0
    self._mtp_stats_state = _MTPStatsState()

    # Memory management: periodic mx.clear_cache() to free Metal command buffers
    # Lower interval = less VRAM spike during generation but slight throughput cost
    self._step_count = 0
    self._clear_cache_interval = 32
    self._memory_log_interval = 256

vllm_mlx.scheduler.Scheduler.model instance-attribute

model = model

vllm_mlx.scheduler.Scheduler.tokenizer instance-attribute

tokenizer = tokenizer

vllm_mlx.scheduler.Scheduler.config instance-attribute

config = config or SchedulerConfig()

vllm_mlx.scheduler.Scheduler._actual_tokenizer instance-attribute

_actual_tokenizer = self._get_actual_tokenizer(tokenizer)

vllm_mlx.scheduler.Scheduler._detokenizer_pool instance-attribute

_detokenizer_pool: Dict[str, Any] = {}

vllm_mlx.scheduler.Scheduler.waiting instance-attribute

waiting: deque[Request] = deque()

vllm_mlx.scheduler.Scheduler.running instance-attribute

running: Dict[str, Request] = {}

vllm_mlx.scheduler.Scheduler.requests instance-attribute

requests: Dict[str, Request] = {}

vllm_mlx.scheduler.Scheduler.finished_req_ids instance-attribute

finished_req_ids: Set[str] = set()

vllm_mlx.scheduler.Scheduler.request_id_to_uid instance-attribute

request_id_to_uid: Dict[str, int] = {}

vllm_mlx.scheduler.Scheduler.uid_to_request_id instance-attribute

uid_to_request_id: Dict[int, str] = {}

vllm_mlx.scheduler.Scheduler.batch_generator instance-attribute

batch_generator: Optional[BatchGenerator] = None

vllm_mlx.scheduler.Scheduler._current_sampler_params instance-attribute

_current_sampler_params: Optional[Tuple] = None

vllm_mlx.scheduler.Scheduler.prefix_cache instance-attribute

prefix_cache: Optional[PrefixCacheManager] = None

vllm_mlx.scheduler.Scheduler.memory_aware_cache instance-attribute

memory_aware_cache: Optional[MemoryAwarePrefixCache] = None

vllm_mlx.scheduler.Scheduler.paged_cache_manager instance-attribute

paged_cache_manager: Optional[PagedCacheManager] = None

vllm_mlx.scheduler.Scheduler.block_aware_cache instance-attribute

block_aware_cache: Optional[BlockAwarePrefixCache] = None

vllm_mlx.scheduler.Scheduler._ssd_tier instance-attribute

_ssd_tier: Optional[SSDCacheTier] = None

vllm_mlx.scheduler.Scheduler._pending_abort_ids instance-attribute

_pending_abort_ids: Set[str] = set()

vllm_mlx.scheduler.Scheduler.num_requests_processed instance-attribute

num_requests_processed = 0

vllm_mlx.scheduler.Scheduler.total_prompt_tokens instance-attribute

total_prompt_tokens = 0

vllm_mlx.scheduler.Scheduler.total_completion_tokens instance-attribute

total_completion_tokens = 0

vllm_mlx.scheduler.Scheduler._mtp_stats_state instance-attribute

_mtp_stats_state = _MTPStatsState()

vllm_mlx.scheduler.Scheduler._step_count instance-attribute

_step_count = 0

vllm_mlx.scheduler.Scheduler._clear_cache_interval instance-attribute

_clear_cache_interval = 32

vllm_mlx.scheduler.Scheduler._memory_log_interval instance-attribute

_memory_log_interval = 256

vllm_mlx.scheduler.Scheduler.SNAPSHOT_REFRESH_TOKENS class-attribute instance-attribute

SNAPSHOT_REFRESH_TOKENS = 4096

vllm_mlx.scheduler.Scheduler._get_actual_tokenizer

_get_actual_tokenizer(tokenizer: Any) -> Any

Get the actual tokenizer from a processor or tokenizer.

MLLM models use processors (e.g., Qwen3VLProcessor) which wrap the tokenizer. This method extracts the actual tokenizer.

Source code in vllm_mlx/scheduler.py

def _get_actual_tokenizer(self, tokenizer: Any) -> Any:
    """
    Get the actual tokenizer from a processor or tokenizer.

    MLLM models use processors (e.g., Qwen3VLProcessor) which wrap
    the tokenizer. This method extracts the actual tokenizer.
    """
    # If it has encode method, it's already a tokenizer
    if hasattr(tokenizer, "encode") and callable(tokenizer.encode):
        return tokenizer
    # If it's a processor, get the wrapped tokenizer
    if hasattr(tokenizer, "tokenizer"):
        return tokenizer.tokenizer
    # Fallback to the original
    return tokenizer

vllm_mlx.scheduler.Scheduler._decode_tokens

_decode_tokens(token_ids: List[int]) -> str

Decode token IDs to text, handling both tokenizers and processors.

Source code in vllm_mlx/scheduler.py

def _decode_tokens(self, token_ids: List[int]) -> str:
    """
    Decode token IDs to text, handling both tokenizers and processors.
    """
    return self._actual_tokenizer.decode(token_ids)

vllm_mlx.scheduler.Scheduler._get_detokenizer

_get_detokenizer(request_id: str) -> Any

Get or create a streaming detokenizer for a request.

Source code in vllm_mlx/scheduler.py

def _get_detokenizer(self, request_id: str) -> Any:
    """Get or create a streaming detokenizer for a request."""
    if request_id not in self._detokenizer_pool:
        detok = NaiveStreamingDetokenizer(self._actual_tokenizer)
        self._detokenizer_pool[request_id] = detok
    return self._detokenizer_pool[request_id]

vllm_mlx.scheduler.Scheduler._cleanup_detokenizer

_cleanup_detokenizer(request_id: str) -> None

Remove the streaming detokenizer for a finished request.

Source code in vllm_mlx/scheduler.py

def _cleanup_detokenizer(self, request_id: str) -> None:
    """Remove the streaming detokenizer for a finished request."""
    self._detokenizer_pool.pop(request_id, None)

vllm_mlx.scheduler.Scheduler._get_stop_tokens

_get_stop_tokens() -> Set[int]

Get stop token IDs from tokenizer or processor.

Source code in vllm_mlx/scheduler.py

def _get_stop_tokens(self) -> Set[int]:
    """Get stop token IDs from tokenizer or processor."""
    stop_tokens = set()
    # Check both the processor/tokenizer and the actual tokenizer
    for tok in [self.tokenizer, self._actual_tokenizer]:
        if tok is None:
            continue
        if hasattr(tok, "eos_token_id") and tok.eos_token_id is not None:
            if isinstance(tok.eos_token_id, list):
                stop_tokens.update(tok.eos_token_id)
            else:
                stop_tokens.add(tok.eos_token_id)
        if hasattr(tok, "eos_token_ids") and tok.eos_token_ids is not None:
            if isinstance(tok.eos_token_ids, (list, set, tuple)):
                stop_tokens.update(tok.eos_token_ids)
            else:
                # Handle case where eos_token_ids is a single int
                stop_tokens.add(tok.eos_token_ids)
    return stop_tokens

vllm_mlx.scheduler.Scheduler._create_batch_generator

_create_batch_generator(sampling_params: SamplingParams) -> BatchGenerator

Create a BatchGenerator with the given sampling parameters.

Source code in vllm_mlx/scheduler.py

def _create_batch_generator(
    self, sampling_params: SamplingParams
) -> BatchGenerator:
    """Create a BatchGenerator with the given sampling parameters."""
    sampler = make_sampler(
        temp=sampling_params.temperature,
        top_p=sampling_params.top_p,
        min_p=sampling_params.min_p,
    )

    stop_tokens = self._get_stop_tokens()
    # Add custom stop token IDs
    if sampling_params.stop_token_ids:
        stop_tokens.update(sampling_params.stop_token_ids)

    def _prefill_progress(progress_list):
        """Log prefill progress for each uid chunk."""
        for uid, processed, total in progress_list:
            rid = self.uid_to_request_id.get(uid, "?")
            logger.info(
                f"[prefill] request={rid[:12] if isinstance(rid, str) else rid} "
                f"tokens={processed}/{total}"
            )

    bg = BatchGenerator(
        model=self.model,
        max_tokens=sampling_params.max_tokens,
        stop_tokens=stop_tokens,
        sampler=sampler,
        prefill_batch_size=self.config.prefill_batch_size,
        completion_batch_size=self.config.completion_batch_size,
        prefill_step_size=self.config.prefill_step_size,
    )
    # Set callback as attribute — used by _install_chunked_prefill
    # monkey-patch. Not a BatchGenerator constructor parameter.
    bg.prompt_progress_callback = _prefill_progress

    # Install chunked prefill only when explicitly configured.
    # memory_aware_cache fetch/store works independently; the mid-prefill
    # save callback is an optimisation, not a requirement.
    # When chunked_prefill_tokens == 0 (the default), honour the user's
    # intent — do NOT silently re-enable chunked prefill just because
    # memory_aware_cache is active (see #178).
    chunked_budget = self.config.chunked_prefill_tokens
    need_chunked = chunked_budget > 0

    prompt_cache_cb = None
    if self.memory_aware_cache is not None:
        prompt_cache_cb = self._make_prompt_cache_save_callback()

    if need_chunked:
        _configure_chunked_prefill(
            self,
            bg,
            chunked_budget,
            prompt_cache_cb,
        )

    # When chunked prefill is off but memory_aware_cache is active,
    # install the lightweight _process_prompts hook so prompt-only
    # cache entries are still captured.  This is the only safe capture
    # point for hybrid Mamba+Transformer models (#178).
    if not need_chunked and prompt_cache_cb is not None:
        if hasattr(bg, "_process_prompts"):
            _install_prompt_cache_save(bg, prompt_cache_cb)

    # Install MTP if the model supports it
    if self.config.enable_mtp:
        if hasattr(self.model, "mtp") and self.model.mtp is not None:
            _install_mtp(
                bg,
                model=self.model,
                num_draft_tokens=self.config.mtp_num_draft_tokens,
                optimistic=self.config.mtp_optimistic,
                stats_state=self._mtp_stats_state,
            )
        else:
            logger.warning(
                "[MTP] --enable-mtp is set but model has no MTP head "
                "(model.mtp is None). MTP will be disabled."
            )

    return bg

vllm_mlx.scheduler.Scheduler._make_prompt_cache_save_callback

_make_prompt_cache_save_callback()

Create a callback that stores prompt-only KV/Mamba cache.

Called from _generation_step right before the first output token is fed into the model. At that point num_tokens == 0 and the batch cache contains the exact prompt-only state (correct for both KVCache and MambaCache/ArraysCache layers).

The cache is stored with key = prompt_token_ids so that a future request with the identical prompt gets an exact hit.

Source code in vllm_mlx/scheduler.py

def _make_prompt_cache_save_callback(self):
    """Create a callback that stores prompt-only KV/Mamba cache.

    Called from ``_generation_step`` right before the first output token
    is fed into the model.  At that point ``num_tokens == 0`` and the
    batch cache contains the exact prompt-only state (correct for both
    KVCache and MambaCache/ArraysCache layers).

    The cache is stored with key = prompt_token_ids so that a future
    request with the identical prompt gets an exact hit.
    """
    import time as _time

    def _prompt_cache_save(uid, extracted_cache):
        request_id = self.uid_to_request_id.get(uid)
        if not request_id:
            return
        request = self.requests.get(request_id)
        if not request or not request.prompt_token_ids:
            return

        prompt_tokens = list(request.prompt_token_ids)
        # Trim cache by 1 so the stored KV has offset = N-1.
        # On exact fetch the scheduler sends the last prompt token
        # for reprocessing (lines 1872-1877).  Without this trim
        # the last token would be placed at position N instead of N-1.
        from .memory_cache import _trim_cache_offset

        trimmed_cache = _trim_cache_offset(extracted_cache, 1)
        _t0 = _time.monotonic()
        # evict_prefixes=False: keep mid-prefill boundary entries so
        # that future requests with the same prefix but different
        # suffix get a prefix cache hit (critical for agentic multi-turn).
        stored = self.memory_aware_cache.store(
            prompt_tokens, trimmed_cache, evict_prefixes=False
        )
        _dt = _time.monotonic() - _t0
        if stored:
            logger.info(
                f"[prompt_cache_save] request={request_id[:12]} "
                f"prompt_tokens={len(prompt_tokens)} "
                f"store_time={_dt:.3f}s"
            )

    return _prompt_cache_save

vllm_mlx.scheduler.Scheduler._make_mid_prefill_save_callback

_make_mid_prefill_save_callback(save_interval: int)

Create a callback for saving intermediate KV cache during chunked prefill.

The callback is called after each chunk with (uid, processed_tokens, prompt_cache). It extracts the cache state (immutable MLX array snapshots), reconstructs KVCache objects, and stores them in the memory-aware prefix cache so that a subsequent request with the same prompt prefix can skip the already-computed tokens.

Source code in vllm_mlx/scheduler.py

def _make_mid_prefill_save_callback(self, save_interval: int):
    """Create a callback for saving intermediate KV cache during chunked prefill.

    The callback is called after each chunk with (uid, processed_tokens,
    prompt_cache).  It extracts the cache state (immutable MLX array
    snapshots), reconstructs KVCache objects, and stores them in the
    memory-aware prefix cache so that a subsequent request with the same
    prompt prefix can skip the already-computed tokens.
    """
    import time as _time

    def _mid_prefill_save(uid, processed_tokens, prompt_cache):
        request_id = self.uid_to_request_id.get(uid)
        if not request_id:
            return
        request = self.requests.get(request_id)
        if not request or not request.prompt_token_ids:
            return

        total_cached = (request.cached_tokens or 0) + processed_tokens

        # Always save at prefix_boundary (message boundary for cache
        # reuse with different final user messages).
        prefix_boundary = getattr(request, "prefix_boundary", 0)
        at_prefix_boundary = prefix_boundary > 0 and total_cached == prefix_boundary

        # Throttle: only save every save_interval tokens,
        # unless we're at the prefix boundary.
        last_save = getattr(request, "_mid_prefill_last_save", 0)
        if not at_prefix_boundary and total_cached - last_save < save_interval:
            return

        # Extract immutable state snapshots
        extracted = self._extract_cache_states(prompt_cache)
        if not extracted:
            return

        # Reconstruct cache objects (directly usable by BatchGenerator)
        reconstructed = self._reconstruct_cache_from_states(extracted)
        if not reconstructed:
            return

        prefix_tokens = list(request.prompt_token_ids[:total_cached])

        # Remove previous intermediate entry to avoid memory waste
        old_key = getattr(request, "_mid_prefill_cache_key", None)
        if old_key is not None:
            self.memory_aware_cache.remove(list(old_key))

        _t0 = _time.monotonic()
        stored = self.memory_aware_cache.store(prefix_tokens, reconstructed)
        _dt = _time.monotonic() - _t0

        if stored:
            request._mid_prefill_last_save = total_cached
            request._mid_prefill_cache_key = tuple(prefix_tokens)
            logger.info(
                f"[mid_prefill_cache] request={request_id[:12]} "
                f"saved {total_cached}/{len(request.prompt_token_ids)} tokens "
                f"({total_cached * 100 // len(request.prompt_token_ids)}%) "
                f"store_time={_dt:.3f}s"
            )
        else:
            logger.debug(
                f"[mid_prefill_cache] request={request_id[:12]} "
                f"store rejected for {total_cached} tokens"
            )

    return _mid_prefill_save

vllm_mlx.scheduler.Scheduler._close_batch_generator

_close_batch_generator() -> None

Properly close BatchGenerator to restore wired_limit.

Source code in vllm_mlx/scheduler.py

def _close_batch_generator(self) -> None:
    """Properly close BatchGenerator to restore wired_limit."""
    if self.batch_generator is not None:
        try:
            if hasattr(self.batch_generator, "close"):
                self.batch_generator.close()
        except Exception as e:
            logger.debug(f"Error closing BatchGenerator: {e}")
        self.batch_generator = None

vllm_mlx.scheduler.Scheduler._ensure_batch_generator

_ensure_batch_generator(sampling_params: SamplingParams) -> None

Ensure BatchGenerator exists with compatible settings.

Source code in vllm_mlx/scheduler.py

def _ensure_batch_generator(self, sampling_params: SamplingParams) -> None:
    """Ensure BatchGenerator exists with compatible settings."""
    sampler_params = (
        sampling_params.temperature,
        sampling_params.top_p,
        sampling_params.min_p,
    )

    # Create new generator if needed or if sampling params changed
    if (
        self.batch_generator is None
        or self._current_sampler_params != sampler_params
    ):
        # If we have an existing generator with requests, we need to drain it first
        if self.batch_generator is not None and self.running:
            logger.warning(
                "Sampling parameters changed with active requests. "
                "New requests will use new parameters after current batch completes."
            )
            return

        # Keep prefix cache across BatchGenerator recreations.
        # KV cache entries depend only on the input tokens, not on
        # sampling params (temperature, top_p, min_p).  Since the
        # server runs a single model, the cache is always valid.
        if self.batch_generator is not None:
            n_entries = 0
            if self.memory_aware_cache is not None:
                n_entries = len(self.memory_aware_cache._entries)
            elif self.prefix_cache is not None:
                n_entries = (
                    len(self.prefix_cache)
                    if hasattr(self.prefix_cache, "__len__")
                    else 0
                )
            logger.info(
                f"[batch_generator] recreating (sampler params changed), "
                f"keeping {n_entries} cache entries"
            )

        self._close_batch_generator()
        self.batch_generator = self._create_batch_generator(sampling_params)
        self._current_sampler_params = sampler_params

vllm_mlx.scheduler.Scheduler._validate_cache

_validate_cache(cache: Any) -> bool

Validate that a cache object is usable.

Checks for None references AND shape compatibility. Restored cache entries must have batch_size == 1 (single sequence) so they can be merged into the running batch by _merge_caches. A shape mismatch here (e.g. batch=2 from a stale entry) would cause a concatenation crash inside _merge_caches.

Parameters:

• cache (Any) –

  The cache object to validate

Returns:

• bool –

  True if cache is valid and usable

Source code in vllm_mlx/scheduler.py

def _validate_cache(self, cache: Any) -> bool:
    """
    Validate that a cache object is usable.

    Checks for None references AND shape compatibility.  Restored
    cache entries must have batch_size == 1 (single sequence) so
    they can be merged into the running batch by _merge_caches.
    A shape mismatch here (e.g. batch=2 from a stale entry) would
    cause a concatenation crash inside _merge_caches.

    Args:
        cache: The cache object to validate

    Returns:
        True if cache is valid and usable
    """
    if cache is None:
        return False

    # Check if it's a list of cache layers
    if isinstance(cache, list):
        if len(cache) == 0:
            return False
        # Check each layer
        for layer_cache in cache:
            if layer_cache is None:
                return False
            # Check if layer has expected structure
            if hasattr(layer_cache, "keys") and layer_cache.keys is None:
                return False
            if hasattr(layer_cache, "values") and layer_cache.values is None:
                return False
            # Validate batch dimension == 1 for KVCache layers
            if hasattr(layer_cache, "keys") and layer_cache.keys is not None:
                if layer_cache.keys.shape[0] != 1:
                    logger.debug(
                        f"Cache layer invalid: keys batch={layer_cache.keys.shape[0]}, expected 1"
                    )
                    return False
            # Validate batch dimension for MambaCache layers
            if hasattr(layer_cache, "cache") and isinstance(
                layer_cache.cache, list
            ):
                for arr in layer_cache.cache:
                    if arr is not None and arr.shape[0] != 1:
                        logger.debug(
                            f"Cache layer invalid: mamba batch={arr.shape[0]}, expected 1"
                        )
                        return False

    # Check BatchKVCache structure
    if hasattr(cache, "caches"):
        if cache.caches is None:
            return False
        for c in cache.caches:
            if c is None:
                return False

    return True

vllm_mlx.scheduler.Scheduler._extract_cache_states

_extract_cache_states(raw_cache: List[Any]) -> List[Dict[str, Any]]

Extract actual tensor state from each layer cache.

This extracts the real KV data using mlx-lm's cache.state property, allowing the data to be stored and reconstructed later even after the BatchGenerator is recreated.

Parameters:

• raw_cache (List[Any]) –

  List of KVCache objects from mlx-lm

Returns:

• List[Dict[str, Any]] –

  List of dicts with {state: (keys, values), meta_state: (offset,), class_name: str}

Source code in vllm_mlx/scheduler.py

def _extract_cache_states(self, raw_cache: List[Any]) -> List[Dict[str, Any]]:
    """
    Extract actual tensor state from each layer cache.

    This extracts the real KV data using mlx-lm's cache.state property,
    allowing the data to be stored and reconstructed later even after
    the BatchGenerator is recreated.

    Args:
        raw_cache: List of KVCache objects from mlx-lm

    Returns:
        List of dicts with {state: (keys, values), meta_state: (offset,), class_name: str}
    """
    if not raw_cache:
        return []

    extracted = []
    for layer_cache in raw_cache:
        try:
            if hasattr(layer_cache, "state") and hasattr(layer_cache, "meta_state"):
                state = layer_cache.state  # (keys, values) or more for Mamba
                meta = layer_cache.meta_state  # (offset,) as strings
                extracted.append(
                    {
                        "state": state,
                        "meta_state": meta,
                        "class_name": type(layer_cache).__name__,
                        "class_ref": type(layer_cache),
                    }
                )
        except Exception as e:
            logger.debug(f"Failed to extract state from cache layer: {e}")
            continue

    return extracted if len(extracted) == len(raw_cache) else []

vllm_mlx.scheduler.Scheduler._reconstruct_cache_from_states

_reconstruct_cache_from_states(extracted_states: List[Dict[str, Any]]) -> Optional[List[Any]]

Reconstruct cache objects from extracted cache states.

This is the inverse of _extract_cache_states(). Uses mlx-lm's _BaseCache.from_state() to reconstruct any cache type (KVCache, MambaCache, etc.) from its state/meta_state.

Parameters:

• extracted_states (List[Dict[str, Any]]) –

  List of dicts from _extract_cache_states()

Returns:

• Optional[List[Any]] –

  List of cache objects, or None if reconstruction fails

Source code in vllm_mlx/scheduler.py

def _reconstruct_cache_from_states(
    self, extracted_states: List[Dict[str, Any]]
) -> Optional[List[Any]]:
    """
    Reconstruct cache objects from extracted cache states.

    This is the inverse of _extract_cache_states(). Uses mlx-lm's
    _BaseCache.from_state() to reconstruct any cache type (KVCache,
    MambaCache, etc.) from its state/meta_state.

    Args:
        extracted_states: List of dicts from _extract_cache_states()

    Returns:
        List of cache objects, or None if reconstruction fails
    """
    if not extracted_states:
        return None

    try:
        caches = []
        for layer_state in extracted_states:
            state = layer_state.get("state")
            meta_state = layer_state.get("meta_state")
            cache_cls = layer_state.get("class_ref")
            if state is None:
                return None

            if cache_cls is not None and hasattr(cache_cls, "from_state"):
                # BatchKVCache doesn't inherit from KVCache, so
                # _merge_caches can't handle it. Convert to KVCache
                # (safe because mid-prefill save is always batch_size=1).
                from mlx_lm.models.cache import (
                    BatchKVCache as _BatchKVCache,
                    KVCache as _KVCache,
                )

                if cache_cls is _BatchKVCache:
                    # BatchKVCache.state = (keys, values, offset, left_padding)
                    keys, values = state[0], state[1]
                    cache = _KVCache()
                    cache.keys = keys
                    cache.values = values
                    cache.offset = keys.shape[2]
                else:
                    cache = cache_cls.from_state(state, meta_state)
            else:
                # Fallback: try KVCache manual reconstruction
                from mlx_lm.models.cache import KVCache

                if len(state) != 2:
                    return None
                cache = KVCache()
                cache.keys, cache.values = state
                cache.offset = (
                    int(meta_state[0]) if meta_state else cache.keys.shape[2]
                )

            caches.append(cache)

        return caches

    except Exception as e:
        logger.info(f"[mid_prefill_cache] reconstruct EXCEPTION: {e}")
        return None

vllm_mlx.scheduler.Scheduler.add_request

add_request(request: Request) -> None

Add a new request to the scheduler.

Parameters:

• request (Request) –

  The request to add

Source code in vllm_mlx/scheduler.py

def add_request(self, request: Request) -> None:
    """
    Add a new request to the scheduler.

    Args:
        request: The request to add
    """
    if request.request_id in self.requests:
        raise ValueError(f"Request {request.request_id} already exists")

    # Tokenize if needed
    if request.prompt_token_ids is None:
        if isinstance(request.prompt, str):
            # Handle both tokenizers and processors (for MLLM models)
            if hasattr(self.tokenizer, "encode"):
                request.prompt_token_ids = self.tokenizer.encode(request.prompt)
            elif hasattr(self.tokenizer, "tokenizer") and hasattr(
                self.tokenizer.tokenizer, "encode"
            ):
                # Processor wraps tokenizer (e.g., Qwen3VLProcessor)
                request.prompt_token_ids = self.tokenizer.tokenizer.encode(
                    request.prompt
                )
            else:
                raise AttributeError(
                    f"Tokenizer {type(self.tokenizer)} has no 'encode' method. "
                    "Continuous batching requires a tokenizer with encode support."
                )
        else:
            request.prompt_token_ids = list(request.prompt)
        request.num_prompt_tokens = len(request.prompt_token_ids)

    # Check prefix cache for cached KV state
    if self.block_aware_cache is not None:
        # Use paged cache
        block_table, remaining = self.block_aware_cache.fetch_cache(
            request.request_id,
            request.prompt_token_ids,
        )
        if block_table and block_table.num_tokens > 0:
            request.cache_hit_type = "hit"
            # Reconstruct actual KVCache objects from stored tensor data
            reconstructed = self.block_aware_cache.reconstruct_cache(block_table)
            if reconstructed:
                request.prompt_cache = reconstructed
                request.block_table = block_table
                request.cached_tokens = block_table.num_tokens
                request.shared_prefix_blocks = len(block_table.block_ids)
                request.remaining_tokens = remaining
                logger.debug(
                    f"Request {request.request_id}: paged cache hit, "
                    f"{request.cached_tokens} tokens in {request.shared_prefix_blocks} blocks, "
                    f"{len(remaining)} tokens remaining, cache reconstructed"
                )
            else:
                # Reconstruction failed, treat as cache miss
                request.cache_hit_type = "miss"
                request.remaining_tokens = request.prompt_token_ids
                logger.debug(
                    f"Request {request.request_id}: paged cache reconstruction failed"
                )
        else:
            request.cache_hit_type = "miss"
            request.remaining_tokens = request.prompt_token_ids
    elif self.memory_aware_cache is not None:
        # Use memory-aware prefix cache
        import time as _time

        _fetch_t0 = _time.monotonic()
        cache, remaining = self.memory_aware_cache.fetch(request.prompt_token_ids)
        _fetch_dt = _time.monotonic() - _fetch_t0
        request.cache_hit_type = self.memory_aware_cache._last_match_type
        if cache:
            request.prompt_cache = cache
            request.cached_tokens = len(request.prompt_token_ids) - len(remaining)
            request.remaining_tokens = remaining
            logger.info(
                f"[cache_fetch] request={request.request_id[:12]} HIT "
                f"prompt_tokens={len(request.prompt_token_ids)} "
                f"cached={request.cached_tokens} remaining={len(remaining)} "
                f"time={_fetch_dt:.3f}s"
            )
        else:
            request.remaining_tokens = request.prompt_token_ids
            logger.info(
                f"[cache_fetch] request={request.request_id[:12]} MISS "
                f"prompt_tokens={len(request.prompt_token_ids)} "
                f"time={_fetch_dt:.3f}s entries={len(self.memory_aware_cache._entries)}"
            )
            # Check SSD tier for cold-tier hit
            if self._ssd_tier is not None:
                ssd_candidate = self.memory_aware_cache.check_ssd(
                    request.prompt_token_ids
                )
                if ssd_candidate is not None:
                    request.cache_hit_type = "ssd_pending"
                    request._ssd_candidate = ssd_candidate
    elif self.prefix_cache is not None:
        # Use legacy prefix cache
        cache, remaining = self.prefix_cache.fetch_cache(request.prompt_token_ids)
        if cache:
            request.cache_hit_type = "hit"
            request.prompt_cache = cache
            request.cached_tokens = len(request.prompt_token_ids) - len(remaining)
            request.remaining_tokens = remaining
            logger.debug(
                f"Request {request.request_id}: cache hit, "
                f"{request.cached_tokens} tokens cached, "
                f"{len(remaining)} tokens remaining"
            )
        else:
            request.cache_hit_type = "miss"
            request.remaining_tokens = request.prompt_token_ids
    else:
        request.cache_hit_type = "miss"
        request.remaining_tokens = request.prompt_token_ids

    # Add to tracking
    self.requests[request.request_id] = request
    self.waiting.append(request)

    logger.debug(
        f"Added request {request.request_id} with {request.num_prompt_tokens} prompt tokens"
    )

vllm_mlx.scheduler.Scheduler.abort_request

abort_request(request_id: str) -> bool

Queue request for abort. Thread-safe, called from any thread.

The actual abort is deferred to the executor thread (inside step()) to avoid race conditions with in-flight Metal GPU operations.

Parameters:

• request_id (str) –

  The request ID to abort

Returns:

• bool –

  True (abort is always enqueued)

Source code in vllm_mlx/scheduler.py

def abort_request(self, request_id: str) -> bool:
    """
    Queue request for abort. Thread-safe, called from any thread.

    The actual abort is deferred to the executor thread (inside step())
    to avoid race conditions with in-flight Metal GPU operations.

    Args:
        request_id: The request ID to abort

    Returns:
        True (abort is always enqueued)
    """
    self._pending_abort_ids.add(request_id)
    logger.info(f"[abort_request] {request_id[:12]} enqueued for deferred abort")
    return True

vllm_mlx.scheduler.Scheduler._process_pending_aborts

_process_pending_aborts() -> None

Drain and process pending abort requests. Called from executor thread.

Source code in vllm_mlx/scheduler.py

def _process_pending_aborts(self) -> None:
    """Drain and process pending abort requests. Called from executor thread."""
    while self._pending_abort_ids:
        request_id = self._pending_abort_ids.pop()
        self._do_abort_request(request_id)

vllm_mlx.scheduler.Scheduler._do_abort_request

_do_abort_request(request_id: str) -> bool

Actually abort a request. Must be called from the executor thread.

Handles the case where the request was already removed from self.requests by _cleanup_request() but still lives in the BatchGenerator (e.g. in _partial or active_batch).

Parameters:

• request_id (str) –

  The request ID to abort

Returns:

• bool –

  True if any cleanup was performed, False otherwise

Source code in vllm_mlx/scheduler.py

def _do_abort_request(self, request_id: str) -> bool:
    """
    Actually abort a request. Must be called from the executor thread.

    Handles the case where the request was already removed from
    self.requests by _cleanup_request() but still lives in the
    BatchGenerator (e.g. in _partial or active_batch).

    Args:
        request_id: The request ID to abort

    Returns:
        True if any cleanup was performed, False otherwise
    """
    request = self.requests.get(request_id)
    was_waiting = False
    was_running = False
    removed_from_batch = False

    # Remove from waiting queue
    if request is not None and request.status == RequestStatus.WAITING:
        was_waiting = True
        try:
            self.waiting.remove(request)
        except ValueError:
            pass

    # Remove from running (BatchGenerator) — do this even if request
    # was already cleaned up from self.requests, because the UID may
    # still be live inside the BatchGenerator (_partial / active_batch).
    if request_id in self.request_id_to_uid:
        was_running = True
        uid = self.request_id_to_uid[request_id]
        if self.batch_generator is not None:
            self.batch_generator.remove([uid])
            removed_from_batch = True
        del self.uid_to_request_id[uid]
        del self.request_id_to_uid[request_id]

    if request_id in self.running:
        del self.running[request_id]

    # Credit in-flight tokens so dashboard metrics stay accurate
    # (without this, aborted requests' tokens vanish from /v1/status).
    if request is not None and request.num_output_tokens > 0:
        self.total_completion_tokens += request.num_output_tokens
        self.total_prompt_tokens += request.num_prompt_tokens

    if request is not None:
        request.set_finished(RequestStatus.FINISHED_ABORTED)
        # Release cache references so Metal buffers can be freed
        request.prompt_cache = None
        request._extracted_cache = None
    self.finished_req_ids.add(request_id)
    self._cleanup_detokenizer(request_id)

    # Flush Metal encoders after removing arrays from batch
    mx.clear_cache()

    logger.info(
        f"[abort_request] {request_id[:12]} ABORTED "
        f"was_waiting={was_waiting} was_running={was_running} "
        f"removed_from_batch={removed_from_batch} "
        f"remaining_running={len(self.running)} remaining_waiting={len(self.waiting)}"
    )
    return True

vllm_mlx.scheduler.Scheduler.has_requests

has_requests() -> bool

Check if there are any pending or running requests.

Source code in vllm_mlx/scheduler.py

def has_requests(self) -> bool:
    """Check if there are any pending or running requests."""
    return bool(self.waiting or self.running)

vllm_mlx.scheduler.Scheduler.get_num_waiting

get_num_waiting() -> int

Get number of waiting requests.

Source code in vllm_mlx/scheduler.py

def get_num_waiting(self) -> int:
    """Get number of waiting requests."""
    return len(self.waiting)

vllm_mlx.scheduler.Scheduler.get_num_running

get_num_running() -> int

Get number of running requests.

Source code in vllm_mlx/scheduler.py

def get_num_running(self) -> int:
    """Get number of running requests."""
    return len(self.running)

vllm_mlx.scheduler.Scheduler._schedule_waiting

_schedule_waiting() -> List[Request]

Move requests from waiting queue to running.

Returns:

• List[Request] –

  List of requests that were scheduled

Source code in vllm_mlx/scheduler.py

def _schedule_waiting(self) -> List[Request]:
    """
    Move requests from waiting queue to running.

    Returns:
        List of requests that were scheduled
    """
    # Attempt synchronous SSD promotion for any ssd_pending requests
    # before scheduling. This keeps SSD I/O out of fetch() while
    # avoiding engine modifications.
    if self._ssd_tier is not None:
        self._try_promote_ssd_pending()

    scheduled = []

    while self.waiting and len(self.running) < self.config.max_num_seqs:
        request = self.waiting.popleft()

        # Ensure we have a batch generator
        self._ensure_batch_generator(request.sampling_params)

        if self.batch_generator is None:
            # Put back and try again later
            self.waiting.appendleft(request)
            break

        # Determine tokens to process and cache to use
        # Note: Don't use `remaining_tokens or prompt_token_ids` because empty list
        # is falsy in Python. For exact cache match, remaining_tokens=[] but we should
        # pass just the last token so BatchGenerator can start generation.
        if (
            request.remaining_tokens is not None
            and len(request.remaining_tokens) == 0
        ):
            # Exact cache match. Re-feeding the last token is only correct
            # when the cached state stops one token short of the key; for a
            # state that already covers the whole key it duplicates that
            # token in the KV cache and the model then answers from a
            # corrupted context (measured: same prompt, different output).
            # Entries stored post-prefill cover the full key, so drop the
            # cache and prefill instead of guessing which kind this is.
            if getattr(request, "cache_hit_type", None) in {
                "exact",
                "supersequence",
            }:
                logger.debug(
                    "[cache] %s match on a full-coverage entry; "
                    "prefilling to avoid duplicating the last token",
                    request.cache_hit_type,
                )
                cache_to_use = None
                request.prompt_cache = None
                request.cached_tokens = 0
                request.remaining_tokens = request.prompt_token_ids
                tokens_to_process = request.prompt_token_ids
            else:
                tokens_to_process = request.prompt_token_ids[-1:]
        elif request.remaining_tokens:
            tokens_to_process = request.remaining_tokens
        else:
            tokens_to_process = request.prompt_token_ids
        cache_to_use = request.prompt_cache  # May be None

        # Create bounded cache when max_kv_size is configured and no cache exists
        if cache_to_use is None and self.config.max_kv_size > 0:
            from mlx_lm.models.cache import make_prompt_cache

            cache_to_use = make_prompt_cache(
                self.model, max_kv_size=self.config.max_kv_size
            )

        # Validate cache before using it
        if cache_to_use is not None and not self._validate_cache(cache_to_use):
            logger.debug(
                f"Request {request.request_id}: invalid cache detected, "
                f"proceeding without cache"
            )
            cache_to_use = None
            request.prompt_cache = None
            request.cached_tokens = 0
            request.remaining_tokens = request.prompt_token_ids
            tokens_to_process = request.prompt_token_ids

        # Build per-request logits_processors from repetition_penalty and
        # any caller-supplied extras (e.g. JSON schema constrained
        # decoding).
        rep_penalty = request.sampling_params.repetition_penalty
        extra_lp = request.sampling_params.logits_processors or []
        combined_lp: list = []
        if rep_penalty and rep_penalty != 1.0:
            combined_lp.extend(
                make_logits_processors(repetition_penalty=rep_penalty)
            )
            logger.info(
                f"[rep_penalty] request={request.request_id[:12]} "
                f"penalty={rep_penalty}"
            )
        if extra_lp:
            combined_lp.extend(extra_lp)
            logger.info(
                f"[logits_proc] request={request.request_id[:12]} "
                f"extra_processors={len(extra_lp)}"
            )
        lp = combined_lp

        # Insert into BatchGenerator with optional cache.
        # Wrap in try/except: if cache shapes are incompatible
        # (e.g. stale entry after BatchGenerator recreation),
        # fall back to no-cache insert instead of crashing.
        insert_kwargs = {
            "max_tokens": [request.sampling_params.max_tokens],
            "caches": [cache_to_use] if cache_to_use else None,
            # Always pass logits_processors (even empty list) so that
            # mlx_lm BatchGenerator never stores None per-sequence.
            "logits_processors": [lp] if lp else [[]],
        }
        try:
            uids = self.batch_generator.insert(
                [tokens_to_process],
                **insert_kwargs,
            )
        except Exception as e:
            if cache_to_use is not None:
                logger.warning(
                    f"[cache_insert_error] request={request.request_id[:12]} "
                    f"cache insert failed ({e}), retrying without cache"
                )
                cache_to_use = None
                request.prompt_cache = None
                request.cached_tokens = 0
                request.remaining_tokens = request.prompt_token_ids
                tokens_to_process = request.prompt_token_ids
                insert_kwargs["caches"] = None
                uids = self.batch_generator.insert(
                    [tokens_to_process],
                    **insert_kwargs,
                )
            else:
                raise

        if uids:
            uid = uids[0]
            self.request_id_to_uid[request.request_id] = uid
            self.uid_to_request_id[uid] = request.request_id
            request.batch_uid = uid
            request.status = RequestStatus.RUNNING
            # Release the prompt cache reference now that BatchGenerator
            # has its own copy.  Holding this reference prevents MLX from
            # freeing the Metal buffers until the request object is GC'd,
            # which under sustained traffic can accumulate hundreds of GB
            # of wired memory (issue #442).
            request.prompt_cache = None
            self.running[request.request_id] = request
            scheduled.append(request)

            self.total_prompt_tokens += request.num_prompt_tokens
            cache_info = (
                f", {request.cached_tokens} cached"
                if request.cached_tokens > 0
                else ""
            )
            tokens_to_prefill = len(tokens_to_process)
            rep_info = (
                f" rep_penalty={rep_penalty}"
                if rep_penalty and rep_penalty != 1.0
                else ""
            )
            logger.info(
                f"[schedule] request={request.request_id[:12]} uid={uid} "
                f"prompt_tokens={request.num_prompt_tokens} "
                f"tokens_to_prefill={tokens_to_prefill}{cache_info} "
                f"max_tokens={request.sampling_params.max_tokens}{rep_info} "
                f"running={len(self.running)} waiting={len(self.waiting)}"
            )

    return scheduled

vllm_mlx.scheduler.Scheduler._copy_cache_state staticmethod

_copy_cache_state(value: Any) -> Any

Deep-copy a cache state payload.

Sharing the arrays is not safe: RotatingKVCache writes into its ring buffer and PoolingCache writes into its remainder buffer, both in place, so a snapshot that aliases them would be rewritten by the very generation it is supposed to predate. x + 0 forces a fresh array while staying on the GPU.

Source code in vllm_mlx/scheduler.py

@staticmethod
def _copy_cache_state(value: Any) -> Any:
    """Deep-copy a cache ``state`` payload.

    Sharing the arrays is not safe: RotatingKVCache writes into its ring
    buffer and PoolingCache writes into its remainder buffer, both in
    place, so a snapshot that aliases them would be rewritten by the very
    generation it is supposed to predate. ``x + 0`` forces a fresh array
    while staying on the GPU.
    """
    import mlx.core as mx

    if isinstance(value, mx.array):
        return value + 0
    if isinstance(value, (list, tuple)):
        copied = [Scheduler._copy_cache_state(v) for v in value]
        return type(value)(copied) if isinstance(value, tuple) else copied
    return value

vllm_mlx.scheduler.Scheduler._prompt_output_entry_is_useless staticmethod

_prompt_output_entry_is_useless(cache: Any) -> bool

Would a prompt+output entry built from this cache ever be reusable?

Only via a trim: any later query is shorter than a prompt+output key, so the generated tail has to come off first. When the cache cannot be trimmed the entry is dead weight — and far from free, since each one holds a full-length KV copy and Metal runs out of buffers long before the byte budget is reached.

Source code in vllm_mlx/scheduler.py

@staticmethod
def _prompt_output_entry_is_useless(cache: Any) -> bool:
    """Would a prompt+output entry built from this cache ever be reusable?

    Only via a trim: any later query is shorter than a prompt+output key, so
    the generated tail has to come off first. When the cache cannot be
    trimmed the entry is dead weight — and far from free, since each one
    holds a full-length KV copy and Metal runs out of buffers long before
    the byte budget is reached.
    """
    try:
        from mlx_lm.models.cache import can_trim_prompt_cache

        return not can_trim_prompt_cache(cache)
    except Exception:
        return False

vllm_mlx.scheduler.Scheduler._extract_cache_for_uid

_extract_cache_for_uid(uid: int) -> Any

Pull one sequence's cache out of the live BatchGenerator batch.

Source code in vllm_mlx/scheduler.py

def _extract_cache_for_uid(self, uid: int) -> Any:
    """Pull one sequence's cache out of the live BatchGenerator batch."""
    bg = self.batch_generator
    if bg is None:
        return None
    for attr in ("_generation_batch", "_prompt_batch"):
        batch = getattr(bg, attr, None)
        uids = getattr(batch, "uids", None)
        if not uids or uid not in uids:
            continue
        extract = getattr(batch, "extract_cache", None)
        if extract is None:
            continue
        try:
            return extract(uids.index(uid))
        except Exception as e:
            logger.debug("extract_cache(%s) on %s failed: %s", uid, attr, e)
    return None

vllm_mlx.scheduler.Scheduler._make_snapshot_destination

_make_snapshot_destination(live_cache: Any) -> Any

Build a destination cache with the same topology as the live one.

make_prompt_cache(model) is not a safe source for this. A plain KVCache destination cannot take a RotatingKVCache's state or meta_state; the assignment raises, the broad handler below logs a warning, and the snapshot is silently never stored — on exactly the sliding-window configurations this feature exists for.

Deriving it from config.max_kv_size instead is also wrong, which I only found by measuring: _create_batch_generator does not pass max_kv_size to BatchGenerator, so with max_kv_size=512 configured the live layers were still plain KVCache and a config-derived destination mismatched in the opposite direction.

So mirror the live objects themselves. A shallow copy keeps the class and every scalar attribute (max_size, keep, step, _idx) and the caller overwrites the arrays, which is the only part that must not be shared.

Source code in vllm_mlx/scheduler.py

def _make_snapshot_destination(self, live_cache: Any) -> Any:
    """Build a destination cache with the same topology as the live one.

    ``make_prompt_cache(model)`` is not a safe source for this. A plain
    ``KVCache`` destination cannot take a ``RotatingKVCache``'s state or
    meta_state; the assignment raises, the broad handler below logs a
    warning, and the snapshot is silently never stored — on exactly the
    sliding-window configurations this feature exists for.

    Deriving it from ``config.max_kv_size`` instead is also wrong, which I
    only found by measuring: ``_create_batch_generator`` does not pass
    ``max_kv_size`` to ``BatchGenerator``, so with ``max_kv_size=512``
    configured the live layers were still plain ``KVCache`` and a
    config-derived destination mismatched in the opposite direction.

    So mirror the live objects themselves. A shallow copy keeps the class
    and every scalar attribute (``max_size``, ``keep``, ``step``, ``_idx``)
    and the caller overwrites the arrays, which is the only part that must
    not be shared.
    """
    import copy

    def _mirror(layer: Any) -> Any:
        children = getattr(layer, "caches", None)
        if children:
            # copy.copy on a container shares the child cache objects, so
            # the "snapshot" would follow live generation. Rebuild it from
            # mirrored children instead.
            mirrored = [_mirror(child) for child in children]
            container = copy.copy(layer)
            container.caches = type(children)(mirrored)
            return container
        return copy.copy(layer)

    try:
        return [_mirror(layer) for layer in live_cache]
    except Exception:
        logger.warning(
            "[cache_store_prompt] could not mirror live cache topology; "
            "not storing",
            exc_info=True,
        )
        return None

vllm_mlx.scheduler.Scheduler._cache_coverage staticmethod

_cache_coverage(cache: Any) -> int | None

How many tokens the live cache actually holds.

Containers have to be descended into: CacheList carries no offset of its own, so reading the attribute off the layer returns None and the caller silently falls back to a prompt-only key — the misalignment this is here to prevent, on exactly the architectures (DeepSeek-V4) that group several caches per layer.

Source code in vllm_mlx/scheduler.py

@staticmethod
def _cache_coverage(cache: Any) -> int | None:
    """How many tokens the live cache actually holds.

    Containers have to be descended into: ``CacheList`` carries no
    ``offset`` of its own, so reading the attribute off the layer returns
    None and the caller silently falls back to a prompt-only key — the
    misalignment this is here to prevent, on exactly the architectures
    (DeepSeek-V4) that group several caches per layer.
    """

    def _offset_of(layer: Any) -> int | None:
        offset = getattr(layer, "offset", None)
        if isinstance(offset, int):
            return offset
        children = getattr(layer, "caches", None)
        if children:
            for child in children:
                found = _offset_of(child)
                if found is not None:
                    return found
        return None

    for layer in cache:
        found = _offset_of(layer)
        if found is not None:
            return found
    return None

vllm_mlx.scheduler.Scheduler._cache_key_for_snapshot

_cache_key_for_snapshot(request: Any, response: Any, raw_cache: Any) -> list[int] | None

Key the entry by the tokens the cache covers, not by the prompt.

The snapshot is taken while processing the response that carries the first generated token, and by then the batch has already fed that token through the cache: measured prompt_len=5, cache_offset=6. Storing that under prompt_token_ids leaves every warm reuse one token ahead of its key.

Trimming the overshoot off is not available here — these are precisely the caches that cannot be trimmed — so the key is extended instead. The extra token is the first token of the reply, which the next turn's prompt also contains, so the entry still matches by strict prefix.

Returns None rather than storing a misaligned entry.

Source code in vllm_mlx/scheduler.py

def _cache_key_for_snapshot(
    self, request: Any, response: Any, raw_cache: Any
) -> list[int] | None:
    """Key the entry by the tokens the cache covers, not by the prompt.

    The snapshot is taken while processing the response that carries the
    first generated token, and by then the batch has already fed that token
    through the cache: measured ``prompt_len=5, cache_offset=6``. Storing
    that under ``prompt_token_ids`` leaves every warm reuse one token ahead
    of its key.

    Trimming the overshoot off is not available here — these are precisely
    the caches that cannot be trimmed — so the key is extended instead. The
    extra token is the first token of the reply, which the next turn's
    prompt also contains, so the entry still matches by strict prefix.

    Returns None rather than storing a misaligned entry.
    """
    covered = self._cache_coverage(raw_cache)
    prompt_ids = list(request.prompt_token_ids)
    if covered is None:
        # Fail closed. A cache that exposes no offset — a pure ArraysCache,
        # for instance — still has the first generated token folded into it
        # by the time the first response arrives, so assuming prompt-only
        # coverage stores state under a key one token short. The next turn
        # then replays that token into cumulative recurrent state and
        # corrupts it, and for these models this is the only entry that
        # ever gets stored. Skipping costs a prefill; guessing costs
        # correctness.
        logger.debug(
            "[cache_store_prompt] coverage unknown for %s; not storing",
            ", ".join(sorted({type(layer).__name__ for layer in raw_cache})),
        )
        return None

    overshoot = covered - len(prompt_ids)
    if overshoot == 0:
        return prompt_ids
    if overshoot < 0:
        logger.debug(
            "[cache_store_prompt] cache covers %d of %d prompt tokens; "
            "not storing",
            covered,
            len(prompt_ids),
        )
        return None

    token = getattr(response, "token", None)
    generated = [] if token is None else [int(token)]
    if overshoot > len(generated):
        logger.debug(
            "[cache_store_prompt] cache is %d tokens past the prompt but "
            "only %d are known; not storing",
            overshoot,
            len(generated),
        )
        return None
    return prompt_ids + generated[:overshoot]

vllm_mlx.scheduler.Scheduler._store_prompt_only_cache

_store_prompt_only_cache(request: Any, response: Any) -> None

Store the post-prefill cache under the prompt tokens alone.

Called once per request, at the point where the cache covers exactly the prompt. Entries keyed this way are reusable without any trimming, which is what models with sliding-window or pooled KV need.

Source code in vllm_mlx/scheduler.py

def _store_prompt_only_cache(self, request: Any, response: Any) -> None:
    """Store the post-prefill cache under the prompt tokens alone.

    Called once per request, at the point where the cache covers exactly
    the prompt. Entries keyed this way are reusable without any trimming,
    which is what models with sliding-window or pooled KV need.
    """
    if self.memory_aware_cache is None:
        return

    # Do not refresh an entry that already covers nearly all of this
    # prompt: the older one still gives a prefix hit next turn, only a few
    # tokens shorter, so the refresh buys almost nothing. The copy itself is
    # cheap (measured make/copy/eval at 0.00/0.00/0.01s for a 43-layer,
    # 11k-token cache), but it allocates a fresh set of per-layer arrays
    # every turn, and buffer count — not bytes — is what Metal runs out of.
    # One copy per SNAPSHOT_REFRESH_TOKENS of growth instead of one per turn.
    # Only throttle REFRESHES. covered > 0 means an existing entry served
    # this prompt as a prefix hit; if it already covers all but a small
    # tail, re-copying the whole cache buys a few tokens at the cost of a
    # fresh set of per-layer arrays every turn. A cold prompt (covered ==
    # 0) must always be stored — gating it on the same threshold silently
    # disabled caching for every conversation shorter than the threshold.
    covered = getattr(request, "cached_tokens", 0) or 0
    if (
        covered > 0
        and len(request.prompt_token_ids) - covered <= self.SNAPSHOT_REFRESH_TOKENS
    ):
        return

    try:
        raw_cache = getattr(response, "prompt_cache", None)
        if callable(raw_cache):
            raw_cache = raw_cache()
        if not raw_cache:
            # mlx-lm only attaches prompt_cache to the response that
            # carries a finish_reason; mid-generation it is None. Pull the
            # per-sequence cache out of the live batch instead, which is
            # what that attribute is built from anyway.
            raw_cache = self._extract_cache_for_uid(response.uid)
        if not raw_cache:
            return

        # Only topologies whose completion-time entry is unusable need this.
        # A trimmable cache already gets a correct entry from the normal
        # path; adding an N+1 snapshot here would evict it and leave an
        # identical N-token prompt matching a supersequence, where the
        # scheduler replays prompt[-1] and duplicates that token. It would
        # also copy and evaluate the whole context before the first token
        # goes out, for no benefit.
        if not self._prompt_output_entry_is_useless(raw_cache):
            return

        cache_key = self._cache_key_for_snapshot(request, response, raw_cache)
        if cache_key is None:
            return

        import mlx.core as mx

        import time as _t

        _t0 = _t.monotonic()
        snapshot = self._make_snapshot_destination(raw_cache)
        _t1 = _t.monotonic()
        if snapshot is None:
            return
        states = []
        for dst, src in zip(snapshot, raw_cache):
            state = self._copy_cache_state(src.state)
            meta = getattr(src, "meta_state", None)
            if meta is not None:
                dst.meta_state = meta
            dst.state = state
            states.append(state)
        _t2 = _t.monotonic()
        mx.eval(states)
        _t3 = _t.monotonic()
        logger.debug(
            "[snapshot_timing] make=%.2fs copy=%.2fs eval=%.2fs layers=%d",
            _t1 - _t0,
            _t2 - _t1,
            _t3 - _t2,
            len(snapshot),
        )

        # evict_prefixes=True is essential here, not cosmetic. In an
        # agentic loop each turn's prompt extends the previous one, so
        # without it every turn adds another full-length KV copy: measured
        # 45 entries of a 46k-token cache, which exhausted Metal's buffer
        # count ("[metal::malloc] Resource limit (499000) exceeded") and
        # aborted generation mid-request. Evicting the superseded prefix
        # keeps one entry per conversation.
        stored = self.memory_aware_cache.store(
            cache_key,
            snapshot,
            evict_prefixes=True,
        )
        logger.info(
            "[cache_store_prompt] request=%s key_tokens=%d prompt_tokens=%d "
            "stored=%s entries=%d",
            request.request_id[:12],
            len(cache_key),
            len(request.prompt_token_ids),
            stored,
            len(self.memory_aware_cache._entries),
        )
    except Exception as e:
        logger.warning(
            "[cache_store_prompt] request=%s snapshot failed: %s",
            request.request_id[:12],
            e,
        )

vllm_mlx.scheduler.Scheduler._process_batch_responses

_process_batch_responses(responses: List[Any]) -> Tuple[List[RequestOutput], Set[str]]

Process responses from BatchGenerator.

Parameters:

• responses (List[Any]) –

  List of BatchGenerator.Response objects

Returns:

• Tuple[List[RequestOutput], Set[str]] –

  Tuple of (outputs, finished_request_ids)

Source code in vllm_mlx/scheduler.py

def _process_batch_responses(
    self, responses: List[Any]
) -> Tuple[List[RequestOutput], Set[str]]:
    """
    Process responses from BatchGenerator.

    Args:
        responses: List of BatchGenerator.Response objects

    Returns:
        Tuple of (outputs, finished_request_ids)
    """
    outputs = []
    finished_ids = set()

    for response in responses:
        request_id = self.uid_to_request_id.get(response.uid)
        if request_id is None:
            continue

        request = self.running.get(request_id)
        if request is None:
            continue

        # Snapshot the cache while it still covers exactly the prompt, i.e.
        # before the first generated token is appended. Storing that under
        # the prompt tokens is the only reuse path open to caches that
        # cannot be trimmed: a later request whose prompt repeats or
        # extends this one then gets an exact or strict-prefix match, and
        # neither needs a trim.
        #
        # The prompt+output entry stored at completion can never be reused
        # by such models. Any future query is shorter than that key, so it
        # would have to trim the generated tail away — and DeepSeek-V4's
        # sliding-window layers physically overwrite older KV once the
        # window wraps (RotatingKVCache.is_trimmable() is offset<max_size),
        # while its PoolingCache cannot split a pooled window. That data is
        # gone, so no trim can recover it.
        if request.num_output_tokens == 0:
            self._store_prompt_only_cache(request, response)

        # Append token to request
        request.append_output_token(response.token)

        # Record first token time for TTFT metric
        if request.first_token_time is None and request.num_output_tokens > 0:
            import time as _time

            request.first_token_time = _time.time()

        # Decode the new token using streaming detokenizer (UTF-8 safe)
        if response.finish_reason == "stop":
            new_text = ""
        else:
            detok = self._get_detokenizer(request_id)
            detok.add_token(response.token)
            new_text = detok.last_segment

        # Create output
        output = RequestOutput(
            request_id=request_id,
            new_token_ids=[response.token],
            new_text=new_text,
            output_token_ids=request.output_token_ids,
            prompt_tokens=request.num_prompt_tokens,
            completion_tokens=request.num_output_tokens,
        )

        # Check if finished
        if response.finish_reason is not None:
            if response.finish_reason == "stop":
                request.set_finished(RequestStatus.FINISHED_STOPPED)
            elif response.finish_reason == "length":
                request.set_finished(RequestStatus.FINISHED_LENGTH_CAPPED)

            output.finished = True
            output.finish_reason = response.finish_reason
            finished_ids.add(request_id)

            # Finalize streaming detokenizer and get full output
            detok = self._detokenizer_pool.get(request_id)
            if detok is not None:
                detok.finalize()
                output.output_text = detok.text
            else:
                output.output_text = self._decode_tokens(request.output_token_ids)
            request.output_text = output.output_text
            self._cleanup_detokenizer(request_id)

            # Extract cache for future reuse (critical for agentic multi-turn)
            if hasattr(response, "prompt_cache"):
                try:
                    # prompt_cache may be callable or direct attribute
                    if callable(response.prompt_cache):
                        raw_cache = response.prompt_cache()
                    else:
                        raw_cache = response.prompt_cache

                    if raw_cache and not self._prompt_output_entry_is_useless(
                        raw_cache
                    ):
                        # For paged cache, extract actual tensor states
                        # This allows cache to survive BatchGenerator recreation
                        if self.block_aware_cache is not None:
                            extracted_cache = self._extract_cache_states(raw_cache)
                            if extracted_cache:
                                request._extracted_cache = extracted_cache
                                logger.debug(
                                    f"Extracted {len(extracted_cache)} layer states "
                                    f"for request {request_id}"
                                )
                        else:
                            # Standard cache stores object references
                            request._extracted_cache = raw_cache
                except Exception as e:
                    logger.debug(f"Failed to extract cache for {request_id}: {e}")

            self.total_completion_tokens += request.num_output_tokens
            self.num_requests_processed += 1

            logger.debug(
                f"Request {request_id} finished: {response.finish_reason}, "
                f"{request.num_output_tokens} tokens"
            )

        outputs.append(output)

    return outputs, finished_ids

vllm_mlx.scheduler.Scheduler._cleanup_finished

_cleanup_finished(finished_ids: Set[str]) -> None

Clean up finished requests and store caches for reuse.

Source code in vllm_mlx/scheduler.py

def _cleanup_finished(self, finished_ids: Set[str]) -> None:
    """Clean up finished requests and store caches for reuse."""
    for request_id in finished_ids:
        request = self.running.get(request_id)

        # Store cache for future reuse
        if request is not None and request.prompt_token_ids:
            if self.block_aware_cache is not None:
                # Store in paged cache
                # Key includes both prompt and output tokens for multi-turn chat caching
                if (
                    hasattr(request, "_extracted_cache")
                    and request._extracted_cache is not None
                ):
                    try:
                        full_token_sequence = list(request.prompt_token_ids) + list(
                            request.output_token_ids
                        )
                        self.block_aware_cache.store_cache(
                            request_id,
                            full_token_sequence,
                            request._extracted_cache,
                        )
                        logger.debug(
                            f"Stored paged cache for request {request_id} "
                            f"({len(full_token_sequence)} tokens: {len(request.prompt_token_ids)} prompt + {len(request.output_token_ids)} output)"
                        )
                    except Exception as e:
                        logger.debug(
                            f"Failed to store paged cache for {request_id}: {e}"
                        )
                # NOTE: Do NOT call release_cache here - blocks should persist
                # for future requests to share. The LRU eviction will clean up
                # unused blocks when under memory pressure.

            elif self.memory_aware_cache is not None:
                # Keep mid-prefill entry as prefix cache for future
                # requests that share a common prefix (e.g. same system
                # prompt + tools but different user message).  LRU
                # eviction handles memory pressure.

                # Store in memory-aware prefix cache
                # Key includes both prompt and output tokens for multi-turn chat caching
                if (
                    hasattr(request, "_extracted_cache")
                    and request._extracted_cache is not None
                ):
                    try:
                        full_token_sequence = list(request.prompt_token_ids) + list(
                            request.output_token_ids
                        )
                        import time as _time

                        _store_t0 = _time.monotonic()
                        stored = self.memory_aware_cache.store(
                            full_token_sequence,
                            request._extracted_cache,
                            evict_prefixes=False,
                        )
                        _store_dt = _time.monotonic() - _store_t0
                        # NOTE: We intentionally do NOT store a prompt-only
                        # cache entry.  Hybrid Mamba+Transformer models
                        # (like Qwen3-Coder-Next) have MambaCache layers
                        # whose state is cumulative and cannot be trimmed
                        # back to "prompt only".  Reusing such state causes
                        # the model to immediately produce EOS.
                        # The full prompt+output entry is stored above; a
                        # future request with the same prompt will hit the
                        # supersequence match path in the fetch, which is
                        # now disabled for safety (see memory_cache.py).

                        logger.info(
                            f"[cache_store] request={request_id[:12]} "
                            f"tokens={len(full_token_sequence)} "
                            f"({len(request.prompt_token_ids)} prompt + {len(request.output_token_ids)} output) "
                            f"stored={stored} time={_store_dt:.3f}s "
                            f"cache_entries={len(self.memory_aware_cache._entries)} "
                            f"cache_mem={self.memory_aware_cache._current_memory / 1e6:.0f}MB"
                        )
                        # Release the original FP16 cache reference so
                        # memory can be reclaimed (the quantized copy
                        # lives inside the prefix cache now).
                        request._extracted_cache = None
                    except Exception as e:
                        logger.debug(
                            f"Failed to store memory-aware cache for {request_id}: {e}"
                        )

            elif self.prefix_cache is not None:
                # Store in legacy prefix cache
                # Key includes both prompt and output tokens for multi-turn chat caching
                # The next turn's prompt will include the previous response
                if (
                    hasattr(request, "_extracted_cache")
                    and request._extracted_cache is not None
                ):
                    try:
                        full_token_sequence = list(request.prompt_token_ids) + list(
                            request.output_token_ids
                        )
                        self.prefix_cache.store_cache(
                            full_token_sequence,
                            request._extracted_cache,
                        )
                        logger.debug(
                            f"Stored cache for request {request_id} "
                            f"({len(full_token_sequence)} tokens: {len(request.prompt_token_ids)} prompt + {len(request.output_token_ids)} output)"
                        )
                    except Exception as e:
                        logger.debug(f"Failed to store cache for {request_id}: {e}")

        # Evaluate stored cache tensors incrementally (per-layer) to prevent
        # a deferred batch evaluation spike when all lazy ops resolve at once.
        # This spreads the VRAM cost across smaller per-layer evaluations.
        if (
            request is not None
            and hasattr(request, "_extracted_cache")
            and request._extracted_cache
        ):
            for layer in request._extracted_cache:
                if isinstance(layer, dict) and "state" in layer:
                    keys, values = layer["state"]
                    mx.eval(keys, values)
                elif hasattr(layer, "keys") and hasattr(layer, "values"):
                    keys_attr = layer.keys
                    values_attr = layer.values
                    if not callable(keys_attr) and not callable(values_attr):
                        mx.eval(keys_attr, values_attr)

        # Release all cache references on the request so Metal buffers
        # can be freed.  The prefix cache (if any) holds its own copy;
        # keeping a second reference here pins the buffers in wired memory
        # until the request object is GC'd (issue #442).
        if request is not None:
            request.prompt_cache = None
            request._extracted_cache = None

        # Remove from running
        if request_id in self.running:
            del self.running[request_id]

        # Remove UID mappings
        if request_id in self.request_id_to_uid:
            uid = self.request_id_to_uid[request_id]
            if uid in self.uid_to_request_id:
                del self.uid_to_request_id[uid]
            del self.request_id_to_uid[request_id]

        # Track as finished
        self.finished_req_ids.add(request_id)

    # Free Metal command buffers after cleanup (prevents end-of-generation spike)
    if finished_ids:
        mx.clear_cache()

vllm_mlx.scheduler.Scheduler._is_cache_corruption_error

_is_cache_corruption_error(error: Exception) -> bool

Check if an error indicates cache corruption.

Source code in vllm_mlx/scheduler.py

def _is_cache_corruption_error(self, error: Exception) -> bool:
    """Check if an error indicates cache corruption."""
    error_str = str(error)
    return any(pattern in error_str for pattern in CACHE_CORRUPTION_PATTERNS)

vllm_mlx.scheduler.Scheduler._is_stream_thread_error

_is_stream_thread_error(error: Exception) -> bool

Check if an error indicates MLX stream/thread ownership mismatch.

Source code in vllm_mlx/scheduler.py

def _is_stream_thread_error(self, error: Exception) -> bool:
    """Check if an error indicates MLX stream/thread ownership mismatch."""
    error_str = str(error)
    return "no Stream(" in error_str or "no Stream(gpu" in error_str

vllm_mlx.scheduler.Scheduler._recover_from_cache_error

_recover_from_cache_error() -> None

Recover from cache corruption error.

Source code in vllm_mlx/scheduler.py

def _recover_from_cache_error(self) -> None:
    """Recover from cache corruption error."""
    # Properly close batch generator (this is the source of the corruption)
    self._close_batch_generator()
    self._current_sampler_params = None

    # Clear caches
    if self.block_aware_cache is not None:
        self.block_aware_cache.clear()
    if self.memory_aware_cache is not None:
        self.memory_aware_cache.clear()
    if self.prefix_cache is not None:
        self.prefix_cache.clear()

    # Clear UID mappings
    self.request_id_to_uid.clear()
    self.uid_to_request_id.clear()

    logger.info("Cache recovery completed")

vllm_mlx.scheduler.Scheduler._recover_from_generation_error

_recover_from_generation_error() -> Set[str]

Recover from fatal generation error (OOM, Metal crash).

Aborts all running requests and resets batch state. Unlike cache corruption recovery, does NOT reschedule — the request that OOMed would just OOM again.

Returns:

• Set[str] –

  Set of aborted request IDs.

Source code in vllm_mlx/scheduler.py

def _recover_from_generation_error(self) -> Set[str]:
    """Recover from fatal generation error (OOM, Metal crash).

    Aborts all running requests and resets batch state.
    Unlike cache corruption recovery, does NOT reschedule —
    the request that OOMed would just OOM again.

    Returns:
        Set of aborted request IDs.
    """
    # Close batch generator (clears _partial state, active_batch)
    self._close_batch_generator()
    self._current_sampler_params = None

    # Abort all running requests
    aborted_ids: Set[str] = set()
    for request_id in list(self.running):
        request = self.running.get(request_id)
        if request is not None:
            request.set_finished(RequestStatus.FINISHED_ABORTED)
        aborted_ids.add(request_id)
        self.finished_req_ids.add(request_id)
    self.running.clear()
    self._detokenizer_pool.clear()

    # Clear UID mappings (batch generator is gone)
    self.request_id_to_uid.clear()
    self.uid_to_request_id.clear()

    # Release Metal memory
    mx.clear_cache()

    logger.warning(
        f"[generation_error_recovery] aborted {len(aborted_ids)} running requests, "
        f"batch generator closed, Metal cache cleared"
    )
    return aborted_ids

vllm_mlx.scheduler.Scheduler._reschedule_running_requests

_reschedule_running_requests() -> None

Move running requests back to waiting queue for retry.

Source code in vllm_mlx/scheduler.py

def _reschedule_running_requests(self) -> None:
    """Move running requests back to waiting queue for retry."""
    count = len(self.running)
    for request_id, request in list(self.running.items()):
        # Reset request state
        request.status = RequestStatus.WAITING
        request.batch_uid = None
        request.prompt_cache = None
        request.cached_tokens = 0
        request.remaining_tokens = request.prompt_token_ids

        # Move to waiting queue (at front for priority)
        self.waiting.appendleft(request)
        del self.running[request_id]

    if count > 0:
        logger.info(f"Rescheduled {count} requests for retry")

vllm_mlx.scheduler.Scheduler.step

step(max_retries: int = 1) -> SchedulerOutput

Execute one scheduling step with automatic error recovery.

This method: 1. Schedules waiting requests into the batch 2. Runs one generation step via BatchGenerator 3. Processes outputs and handles finished requests 4. Automatically recovers from cache corruption errors

Parameters:

• max_retries (int, default: 1 ) –

  Number of times to retry on cache errors (default 1)

Returns:

• SchedulerOutput –

  SchedulerOutput with results of this step

Source code in vllm_mlx/scheduler.py

def step(self, max_retries: int = 1) -> SchedulerOutput:
    """
    Execute one scheduling step with automatic error recovery.

    This method:
    1. Schedules waiting requests into the batch
    2. Runs one generation step via BatchGenerator
    3. Processes outputs and handles finished requests
    4. Automatically recovers from cache corruption errors

    Args:
        max_retries: Number of times to retry on cache errors (default 1)

    Returns:
        SchedulerOutput with results of this step
    """
    output = SchedulerOutput()

    # Process pending aborts FIRST (in executor thread, safe for MLX)
    self._process_pending_aborts()

    for attempt in range(max_retries + 1):
        try:
            # Schedule waiting requests
            scheduled = self._schedule_waiting()
            output.scheduled_request_ids = [r.request_id for r in scheduled]
            output.num_scheduled_tokens = sum(
                r.num_prompt_tokens for r in scheduled
            )

            # Run generation step if we have running requests
            if self.batch_generator is not None and self.running:
                _sanitize_batch_generator_logits_processors(self.batch_generator)
                result = self.batch_generator.next()
                output.has_work = True

                # mlx-lm >=0.31.x returns (prompt_responses, generation_responses);
                # older versions returned a flat list.
                if isinstance(result, tuple):
                    responses = result[1]  # generation_responses only
                else:
                    responses = result

                if responses:
                    outputs, finished_ids = self._process_batch_responses(responses)
                    output.outputs = outputs
                    output.finished_request_ids = finished_ids
                    self._cleanup_finished(finished_ids)

            # Success - break out of retry loop
            break

        except TypeError as e:
            # Catch the NoneType error specifically
            if self._is_cache_corruption_error(e):
                if attempt < max_retries:
                    logger.warning(
                        f"Cache corruption detected (attempt {attempt + 1}), "
                        f"performing recovery and retry..."
                    )
                    # Deep reset to recover
                    self._recover_from_cache_error()
                    # Re-add any running requests back to waiting
                    self._reschedule_running_requests()
                else:
                    logger.error(
                        f"Cache corruption not recoverable after "
                        f"{max_retries + 1} attempts"
                    )
                    raise
            else:
                raise
        except Exception as e:
            if self._is_stream_thread_error(e):
                raise
            import traceback

            logger.error(
                f"Error in batch generation step: {e}\n{traceback.format_exc()}"
            )
            # Recover from fatal errors (OOM, Metal crash) instead of
            # re-raising, which would cause infinite loop in engine_core.
            aborted_ids = self._recover_from_generation_error()
            for rid in aborted_ids:
                output.outputs.append(
                    RequestOutput(
                        request_id=rid,
                        finished=True,
                        finish_reason="error",
                    )
                )
            output.finished_request_ids = aborted_ids
            break

    # Clear finished tracking for next step
    old_finished = self.finished_req_ids
    self.finished_req_ids = set()

    # Adaptive interval: scale inversely with concurrency to prevent
    # Metal resource handle exhaustion under high-concurrency workloads.
    active_seqs = len(self.running)
    min_interval = max(4, self._clear_cache_interval // 4)
    effective_interval = max(
        min_interval, self._clear_cache_interval // max(1, active_seqs // 8)
    )

    self._step_count += 1
    if self._step_count % effective_interval == 0:
        # Evaluate batch tokens to collapse lazy concatenation chains
        if (
            self.batch_generator is not None
            and hasattr(self.batch_generator, "active_batch")
            and self.batch_generator.active_batch is not None
            and hasattr(self.batch_generator.active_batch, "tokens")
        ):
            tokens = self.batch_generator.active_batch.tokens
            if tokens:
                mx.eval(*tokens)
        mx.clear_cache()

    # Periodically log memory stats for monitoring
    if self._step_count % self._memory_log_interval == 0:
        try:
            if mx.metal.is_available():
                active_gb = mx.get_active_memory() / 1e9
                peak_gb = mx.get_peak_memory() / 1e9
                cache_gb = mx.get_cache_memory() / 1e9
                logger.info(
                    f"[Metal memory] active={active_gb:.1f}GB "
                    f"peak={peak_gb:.1f}GB cache={cache_gb:.1f}GB "
                    f"step={self._step_count} "
                    f"running={len(self.running)} waiting={len(self.waiting)}"
                )
        except Exception:
            pass

    return output

vllm_mlx.scheduler.Scheduler.get_request

get_request(request_id: str) -> Optional[Request]

Get a request by ID.

Source code in vllm_mlx/scheduler.py

def get_request(self, request_id: str) -> Optional[Request]:
    """Get a request by ID."""
    return self.requests.get(request_id)

vllm_mlx.scheduler.Scheduler.remove_finished_request

remove_finished_request(request_id: str) -> Optional[Request]

Remove a finished request from tracking.

Source code in vllm_mlx/scheduler.py

def remove_finished_request(self, request_id: str) -> Optional[Request]:
    """Remove a finished request from tracking."""
    return self.requests.pop(request_id, None)

vllm_mlx.scheduler.Scheduler.get_running_requests_info

get_running_requests_info() -> List[Dict[str, Any]]

Per-request details for status endpoint.

Source code in vllm_mlx/scheduler.py

def get_running_requests_info(self) -> List[Dict[str, Any]]:
    """Per-request details for status endpoint."""
    import time as _time

    now = _time.time()
    result = []

    # Waiting requests
    for req in self.waiting:
        result.append(
            {
                "request_id": req.request_id,
                "status": "waiting",
                "phase": "queued",
                "elapsed_s": round(now - req.arrival_time, 2),
                "prompt_tokens": req.num_prompt_tokens,
                "completion_tokens": 0,
                "max_tokens": req.max_tokens,
                "progress": 0.0,
                "tokens_per_second": None,
                "ttft_s": None,
                "cache_hit_type": req.cache_hit_type,
                "cached_tokens": req.cached_tokens,
            }
        )

    # Running requests
    for req in self.running.values():
        n_out = req.num_output_tokens
        elapsed = now - req.arrival_time

        # Phase detection
        if n_out == 0:
            phase = "prefill"
        else:
            phase = "generation"

        # Tokens per second (generation phase only)
        tok_s = None
        ttft = None
        if req.first_token_time is not None:
            ttft = round(req.first_token_time - req.arrival_time, 3)
            gen_elapsed = now - req.first_token_time
            if gen_elapsed > 0 and n_out > 0:
                tok_s = round(n_out / gen_elapsed, 1)

        # Progress: completion_tokens / max_tokens
        progress = round(n_out / req.max_tokens, 3) if req.max_tokens > 0 else 0.0

        result.append(
            {
                "request_id": req.request_id,
                "status": "running",
                "phase": phase,
                "elapsed_s": round(elapsed, 2),
                "prompt_tokens": req.num_prompt_tokens,
                "completion_tokens": n_out,
                "max_tokens": req.max_tokens,
                "progress": min(progress, 1.0),
                "tokens_per_second": tok_s,
                "ttft_s": ttft,
                "cache_hit_type": req.cache_hit_type,
                "cached_tokens": req.cached_tokens,
            }
        )

    return result

vllm_mlx.scheduler.Scheduler.get_stats

get_stats() -> Dict[str, Any]

Get scheduler statistics.

Source code in vllm_mlx/scheduler.py

def get_stats(self) -> Dict[str, Any]:
    """Get scheduler statistics."""
    stats = {
        "num_waiting": len(self.waiting),
        "num_running": len(self.running),
        "num_requests_processed": self.num_requests_processed,
        "total_prompt_tokens": self.total_prompt_tokens,
        "total_completion_tokens": self.total_completion_tokens,
    }
    stats.update(_mtp_status_snapshot(self.batch_generator))
    # Include Metal memory stats
    try:
        if mx.metal.is_available():
            stats["metal_active_memory_gb"] = round(mx.get_active_memory() / 1e9, 2)
            stats["metal_peak_memory_gb"] = round(mx.get_peak_memory() / 1e9, 2)
            stats["metal_cache_memory_gb"] = round(mx.get_cache_memory() / 1e9, 2)
    except Exception:
        pass

    # Include cache stats
    if self.block_aware_cache is not None:
        stats["paged_cache"] = self.block_aware_cache.get_stats()
    elif self.memory_aware_cache is not None:
        stats["memory_aware_cache"] = self.memory_aware_cache.get_stats()
    elif self.prefix_cache is not None:
        stats["prefix_cache"] = self.prefix_cache.get_stats()
    return stats

vllm_mlx.scheduler.Scheduler.get_cache_stats

get_cache_stats() -> Optional[Dict[str, Any]]

Get cache statistics.

Source code in vllm_mlx/scheduler.py

def get_cache_stats(self) -> Optional[Dict[str, Any]]:
    """Get cache statistics."""
    if self.block_aware_cache is not None:
        return self.block_aware_cache.get_stats()
    elif self.memory_aware_cache is not None:
        return self.memory_aware_cache.get_stats()
    elif self.prefix_cache is not None:
        return self.prefix_cache.get_stats()
    return None

vllm_mlx.scheduler.Scheduler.clear_runtime_caches

clear_runtime_caches() -> Dict[str, bool]

Clear prefix-cache state without resetting scheduler/request state.

Source code in vllm_mlx/scheduler.py

def clear_runtime_caches(self) -> Dict[str, bool]:
    """Clear prefix-cache state without resetting scheduler/request state."""
    cleared = {
        "paged_cache": False,
        "memory_aware_cache": False,
        "prefix_cache": False,
    }
    if self.block_aware_cache is not None:
        self.block_aware_cache.clear()
        cleared["paged_cache"] = True
    if self.memory_aware_cache is not None:
        self.memory_aware_cache.clear()
        cleared["memory_aware_cache"] = True
    if self.prefix_cache is not None:
        self.prefix_cache.clear()
        cleared["prefix_cache"] = True
    return cleared

vllm_mlx.scheduler.Scheduler.reset

reset() -> None

Reset the scheduler state.

Source code in vllm_mlx/scheduler.py

def reset(self) -> None:
    """Reset the scheduler state."""
    # Drain any pending deferred aborts
    self._pending_abort_ids.clear()

    # Abort all requests directly (reset is synchronous)
    for request_id in list(self.requests.keys()):
        self._do_abort_request(request_id)

    self.waiting.clear()
    self.running.clear()
    self.requests.clear()
    self.finished_req_ids.clear()
    self.request_id_to_uid.clear()
    self.uid_to_request_id.clear()
    self._detokenizer_pool.clear()
    self._close_batch_generator()
    self._current_sampler_params = None

    # Clear caches
    self.clear_runtime_caches()

    # Close SSD tier on reset
    self.close_ssd_tier()

vllm_mlx.scheduler.Scheduler.deep_reset

deep_reset() -> None

Deep reset that clears ALL cache state including model-level caches.

This is more aggressive than reset() and should be used when switching engines or recovering from errors.

Source code in vllm_mlx/scheduler.py

def deep_reset(self) -> None:
    """
    Deep reset that clears ALL cache state including model-level caches.

    This is more aggressive than reset() and should be used when
    switching engines or recovering from errors.
    """
    # Standard reset first
    self.reset()

    # Clear any model-level cache state
    # MLX models may have internal cache references
    if hasattr(self.model, "cache"):
        self.model.cache = None

    # Some MLX models store cache in layers
    if hasattr(self.model, "layers"):
        for layer in self.model.layers:
            if hasattr(layer, "cache"):
                layer.cache = None
            if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "cache"):
                layer.self_attn.cache = None

    # Force garbage collection of any lingering cache objects
    import gc

    gc.collect()

    logger.info("Deep reset completed - all caches cleared")

vllm_mlx.scheduler.Scheduler.save_cache_to_disk

save_cache_to_disk(cache_dir: str) -> bool

Save prefix cache to disk for persistence across restarts.

Source code in vllm_mlx/scheduler.py

def save_cache_to_disk(self, cache_dir: str) -> bool:
    """Save prefix cache to disk for persistence across restarts."""
    if self.memory_aware_cache is not None:
        return self.memory_aware_cache.save_to_disk(cache_dir)
    logger.info("[cache_persist] no memory-aware cache to save")
    return False

vllm_mlx.scheduler.Scheduler.load_cache_from_disk

load_cache_from_disk(cache_dir: str) -> int

Load prefix cache from disk. Returns number of entries loaded.

Source code in vllm_mlx/scheduler.py

def load_cache_from_disk(self, cache_dir: str) -> int:
    """Load prefix cache from disk. Returns number of entries loaded."""
    if self.memory_aware_cache is not None:
        return self.memory_aware_cache.load_from_disk(cache_dir)
    logger.info("[cache_persist] no memory-aware cache to load into")
    return 0

vllm_mlx.scheduler.Scheduler.clear_prefix_cache

clear_prefix_cache() -> None

Clear the in-memory prefix cache (keeps disk cache untouched).

Source code in vllm_mlx/scheduler.py

def clear_prefix_cache(self) -> None:
    """Clear the in-memory prefix cache (keeps disk cache untouched)."""
    if self.memory_aware_cache is not None and hasattr(
        self.memory_aware_cache, "clear"
    ):
        self.memory_aware_cache.clear()
        logger.info("[clear_prefix_cache] memory-aware cache cleared")
        return
    if self.prefix_cache is not None and hasattr(self.prefix_cache, "clear"):
        self.prefix_cache.clear()
        logger.info("[clear_prefix_cache] prefix cache cleared")

vllm_mlx.scheduler.Scheduler.close_ssd_tier

close_ssd_tier() -> None

Shut down the SSD cache tier if present.

Source code in vllm_mlx/scheduler.py

def close_ssd_tier(self) -> None:
    """Shut down the SSD cache tier if present."""
    if self._ssd_tier is not None:
        self._ssd_tier.close()
        self._ssd_tier = None
        logger.info("SSD cache tier closed")

vllm_mlx.scheduler.Scheduler._try_promote_ssd_pending

_try_promote_ssd_pending() -> None

Attempt synchronous SSD promotion for waiting requests tagged ssd_pending.

Called from _schedule_waiting() before requests are moved to running. Reads SSD entries synchronously (disk I/O stays out of fetch() per spec).

Source code in vllm_mlx/scheduler.py

def _try_promote_ssd_pending(self) -> None:
    """Attempt synchronous SSD promotion for waiting requests tagged ssd_pending.

    Called from _schedule_waiting() before requests are moved to running.
    Reads SSD entries synchronously (disk I/O stays out of fetch() per spec).
    """
    for request in self.waiting:
        if getattr(request, "cache_hit_type", None) != "ssd_pending":
            continue

        candidate = getattr(request, "_ssd_candidate", None)
        if candidate is None:
            continue

        memory_bytes = candidate["memory_bytes"]

        # Check RAM budget availability
        if self.memory_aware_cache is None:
            request.cache_hit_type = "miss"
            continue

        if not self.memory_aware_cache.try_reserve_memory(memory_bytes):
            self._ssd_tier._stats.promotion_failures += 1
            request.cache_hit_type = "miss"
            logger.info(
                f"[ssd_promote] request={request.request_id[:12]} "
                f"budget denied ({memory_bytes} bytes)"
            )
            continue

        # Use the SSD entry's actual token count for read and store,
        # NOT the full prompt tokens. For prefix hits these differ.
        matched_count = candidate["matched_tokens"]
        matched_tokens = tuple(request.prompt_token_ids[:matched_count])

        try:
            cache_layers = self._ssd_tier._read_entry(
                matched_tokens, candidate["file_path"]
            )
        except Exception:
            self.memory_aware_cache.release_reserved_memory(memory_bytes)
            self._ssd_tier._stats.promotion_failures += 1
            request.cache_hit_type = "miss"
            logger.exception(
                f"[ssd_promote] request={request.request_id[:12]} "
                f"disk read failed"
            )
            continue

        if cache_layers is None:
            self.memory_aware_cache.release_reserved_memory(memory_bytes)
            self._ssd_tier._stats.promotion_failures += 1
            request.cache_hit_type = "miss"
            continue

        # Release tentative budget (store() will account properly)
        self.memory_aware_cache.release_reserved_memory(memory_bytes)

        # Reconstruct and store under the matched prefix tokens
        reconstructed = self._reconstruct_ssd_layers(cache_layers)
        if reconstructed is None:
            request.cache_hit_type = "miss"
            continue

        self.memory_aware_cache.store(
            list(matched_tokens), reconstructed, evict_prefixes=False
        )

        request.prompt_cache = reconstructed
        request.cached_tokens = matched_count
        request.remaining_tokens = request.prompt_token_ids[matched_count:]
        request.cache_hit_type = "ssd_hit"

        self._ssd_tier._stats.ssd_hits += 1
        self._ssd_tier._index.touch(matched_tokens)

        logger.info(
            f"[ssd_promote] request={request.request_id[:12]} "
            f"{candidate['match_type']} promote: {matched_count}/{len(request.prompt_token_ids)} tokens from SSD, "
            f"{len(request.remaining_tokens)} remaining"
        )

vllm_mlx.scheduler.Scheduler.promote_from_ssd async

promote_from_ssd(request) -> bool

Promote a cold-tier cache entry for a request (async version).

Alternative to _try_promote_ssd_pending() for callers with an async event loop. Uses asyncio.to_thread for non-blocking disk I/O.

Returns True if promotion succeeded and request was updated.

Source code in vllm_mlx/scheduler.py

async def promote_from_ssd(self, request) -> bool:
    """Promote a cold-tier cache entry for a request (async version).

    Alternative to _try_promote_ssd_pending() for callers with an
    async event loop. Uses asyncio.to_thread for non-blocking disk I/O.

    Returns True if promotion succeeded and request was updated.
    """
    if self._ssd_tier is None:
        return False

    candidate = getattr(request, "_ssd_candidate", None)
    if candidate is None:
        return False

    def reserve_budget(nbytes: int) -> bool:
        """Tentatively reserve RAM budget for promotion."""
        if self.memory_aware_cache is None:
            return False
        return self.memory_aware_cache.try_reserve_memory(nbytes)

    def release_budget(nbytes: int) -> None:
        """Release tentatively reserved budget on failure."""
        if self.memory_aware_cache is not None:
            self.memory_aware_cache.release_reserved_memory(nbytes)

    # Use matched token count, not full prompt, for prefix hits
    matched_count = candidate.get("matched_tokens", len(request.prompt_token_ids))
    matched_tokens = tuple(request.prompt_token_ids[:matched_count])

    cache_layers = await self._ssd_tier.async_promote(
        matched_tokens, reserve_budget, release_budget
    )

    if cache_layers is None:
        request.cache_hit_type = "miss"
        return False

    # Release tentative budget — store() will account properly
    release_budget(candidate["memory_bytes"])

    # Reconstruct cache objects from deserialized layer dicts
    reconstructed = self._reconstruct_ssd_layers(cache_layers)
    if reconstructed is None:
        request.cache_hit_type = "miss"
        return False

    # Store in RAM cache under the matched prefix tokens
    self.memory_aware_cache.store(
        list(matched_tokens), reconstructed, evict_prefixes=False
    )

    request.prompt_cache = reconstructed
    request.cached_tokens = matched_count
    request.remaining_tokens = request.prompt_token_ids[matched_count:]
    request.cache_hit_type = "ssd_hit"

    logger.info(
        f"[ssd_promote] request={request.request_id[:12]} "
        f"{candidate.get('match_type', 'exact')} promote: "
        f"{matched_count}/{len(request.prompt_token_ids)} tokens from SSD, "
        f"{len(request.remaining_tokens)} remaining"
    )
    return True

vllm_mlx.scheduler.Scheduler._reconstruct_ssd_layers

_reconstruct_ssd_layers(layer_dicts: list[dict]) -> list | None

Reconstruct cache objects from deserialized layer dicts.

Converts numpy arrays back to MLX arrays and creates KVCache objects.

Source code in vllm_mlx/scheduler.py

def _reconstruct_ssd_layers(self, layer_dicts: list[dict]) -> list | None:
    """Reconstruct cache objects from deserialized layer dicts.

    Converts numpy arrays back to MLX arrays and creates KVCache objects.
    """
    try:
        from mlx_lm.models.cache import ArraysCache, KVCache

        # Cast restored arrays back to their original dtype if the spill
        # path upcast for numpy (bf16 → fp32). None = mlx lacks the named
        # dtype on this version; accept default from mx.array(np_fp32).
        def _mx_dtype_from_name(name: str):
            return getattr(mx, name, None)

        result = []
        for ld in layer_dicts:
            if "keys" in ld and "values" in ld:
                kv = KVCache()
                kv.keys = mx.array(ld["keys"])
                kv.values = mx.array(ld["values"])
                keys_orig = ld.get("keys_original_dtype")
                if keys_orig is not None:
                    dt = _mx_dtype_from_name(keys_orig)
                    if dt is not None:
                        kv.keys = kv.keys.astype(dt)
                values_orig = ld.get("values_original_dtype")
                if values_orig is not None:
                    dt = _mx_dtype_from_name(values_orig)
                    if dt is not None:
                        kv.values = kv.values.astype(dt)
                kv.offset = ld["offset"]
                for attr in ("max_size", "keep", "step", "_idx"):
                    if attr in ld:
                        setattr(kv, attr, ld[attr])
                result.append(kv)
            elif "state" in ld:
                state_arrays = [mx.array(a) for a in ld["state"]]
                state_dtypes = ld.get("state_original_dtypes")
                if state_dtypes is not None:
                    for i, dtype_name in enumerate(state_dtypes):
                        if dtype_name is None:
                            continue
                        dt = _mx_dtype_from_name(dtype_name)
                        if dt is not None:
                            state_arrays[i] = state_arrays[i].astype(dt)
                layer_obj = ArraysCache(len(state_arrays))
                layer_obj.state = state_arrays
                result.append(layer_obj)
            else:
                logger.warning(
                    f"[ssd_promote] unknown layer dict format: {list(ld.keys())}"
                )
                return None
        return result
    except Exception as e:
        logger.warning(f"[ssd_promote] reconstruction failed: {e}")
        return None

vllm_mlx.scheduler._normalize_logits_processors

_normalize_logits_processors(logits_processors)

Normalize empty per-sequence processor slots to lists.

Source code in vllm_mlx/scheduler.py

def _normalize_logits_processors(logits_processors):
    """Normalize empty per-sequence processor slots to lists."""
    if logits_processors is None:
        return None
    return [processors or [] for processors in logits_processors]

vllm_mlx.scheduler._sanitize_batch_generator_logits_processors

_sanitize_batch_generator_logits_processors(batch_generator) -> None

Sanitize stale BatchGenerator processor state before decode.

Source code in vllm_mlx/scheduler.py

def _sanitize_batch_generator_logits_processors(batch_generator) -> None:
    """Sanitize stale BatchGenerator processor state before decode."""
    active_batch = getattr(batch_generator, "active_batch", None)
    if active_batch is not None and hasattr(active_batch, "logits_processors"):
        active_batch.logits_processors = _normalize_logits_processors(
            active_batch.logits_processors
        )

    partial = getattr(batch_generator, "_partial", None)
    if isinstance(partial, dict) and "logits_processors" in partial:
        partial["logits_processors"] = _normalize_logits_processors(
            partial["logits_processors"]
        )

vllm_mlx.scheduler._install_prompt_cache_save

_install_prompt_cache_save(batch_gen: BatchGenerator, prompt_cache_save) -> None

Monkey-patch _process_prompts to capture prompt-only cache state.

Can be installed independently of chunked prefill. If chunked prefill is also installed, it takes over _process_prompts and invokes the callback itself, so call this before _install_chunked_prefill.

Source code in vllm_mlx/scheduler.py

def _install_prompt_cache_save(batch_gen: "BatchGenerator", prompt_cache_save) -> None:
    """Monkey-patch ``_process_prompts`` to capture prompt-only cache state.

    Can be installed independently of chunked prefill.  If chunked prefill is
    also installed, *it* takes over ``_process_prompts`` and invokes the
    callback itself, so call this **before** ``_install_chunked_prefill``.
    """
    _orig_process_prompts = batch_gen._process_prompts

    try:
        from mlx_lm.generate import Batch as _batch_cls
    except ImportError:
        _batch_cls = None  # extract_cache fallback handled in patched fn

    def _patched_process_prompts(prompts, _self=batch_gen):
        batch = _orig_process_prompts(prompts)
        for e, uid in enumerate(batch.uids):
            if batch.num_tokens[e] == 0:
                try:
                    prompt_cache_save(uid, batch.extract_cache(e))
                except Exception:
                    pass
        return batch

    batch_gen._process_prompts = _patched_process_prompts

vllm_mlx.scheduler._install_chunked_prefill

_install_chunked_prefill(batch_gen: BatchGenerator, budget: int, mid_prefill_save=None, prompt_cache_save=None, pending_abort_ids: Optional[Set[str]] = None, uid_to_request_id: Optional[Dict[int, str]] = None, requests: Optional[Dict[str, Any]] = None) -> None

Monkey-patch a BatchGenerator instance so that large prefills are broken into chunks of at most budget tokens each.

Between chunks the generation loop gets a chance to produce one token for every active request, preventing starvation during long prefills.

Parameters:

• batch_gen (BatchGenerator) –

  The BatchGenerator to patch.

• budget (int) –

  Max tokens per prefill chunk.

• mid_prefill_save –

  Optional callback(uid, processed, prompt_cache) called after each chunk to save intermediate KV cache state.

Source code in vllm_mlx/scheduler.py

def _install_chunked_prefill(
    batch_gen: "BatchGenerator",
    budget: int,
    mid_prefill_save=None,
    prompt_cache_save=None,
    pending_abort_ids: Optional[Set[str]] = None,
    uid_to_request_id: Optional[Dict[int, str]] = None,
    requests: Optional[Dict[str, Any]] = None,
) -> None:
    """
    Monkey-patch a BatchGenerator instance so that large prefills are
    broken into chunks of at most *budget* tokens each.

    Between chunks the generation loop gets a chance to produce one token
    for every active request, preventing starvation during long prefills.

    Args:
        batch_gen: The BatchGenerator to patch.
        budget: Max tokens per prefill chunk.
        mid_prefill_save: Optional callback(uid, processed, prompt_cache)
            called after each chunk to save intermediate KV cache state.
    """
    import time as _time

    from mlx_lm.generate import (
        _left_pad_prompts,
        _make_cache,
        _merge_caches,
        _right_pad_prompts,
    )

    try:
        from mlx_lm.generate import _lazy_extract_cache
    except ImportError:

        def _lazy_extract_cache(cache, idx):
            return (c.extract(idx) for c in cache)

    try:
        from mlx_lm.generate import Batch as _batch_cls
    except ImportError:

        @dataclass
        class _batch_cls:
            uids: List[int]
            y: Any
            logprobs: List[Any]
            max_tokens: List[int]
            num_tokens: List[int]
            cache: List[Any]
            samplers: List[Any]
            logits_processors: List[Any]
            tokens: List[Any]

            def __len__(self):
                return len(self.uids)

            def filter(self, keep_idx: List[int]):
                self.uids = [self.uids[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.samplers = [self.samplers[k] for k in keep_idx]
                self.logits_processors = [self.logits_processors[k] for k in keep_idx]
                self.tokens = [self.tokens[k] for k in keep_idx]
                keep_idx_mx = mx.array(keep_idx, mx.int32)
                self.y = self.y[keep_idx_mx]
                for c in self.cache:
                    c.filter(keep_idx_mx)

            def extend(self, other):
                self.uids.extend(other.uids)
                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.samplers.extend(other.samplers)
                self.logits_processors.extend(other.logits_processors)
                self.tokens.extend(other.tokens)
                for c, o in zip(self.cache, other.cache):
                    c.extend(o)

            def extract_cache(self, idx):
                return [c.extract(idx) for c in self.cache]

    # Keep references to originals
    _orig_next = batch_gen._next
    _orig_remove = batch_gen.remove
    _orig_process_prompts = batch_gen._process_prompts

    # Partial prefill state (None when no prefill in progress)
    batch_gen._partial = None

    # Monkey-patch _process_prompts to capture prompt-only cache state.
    # At the point where _process_prompts returns, the Batch cache contains
    # the exact prompt-only state: all prompt tokens have been processed
    # through the model, but no output token has been fed back yet.
    # This is the only safe capture point for hybrid Mamba+Transformer
    # models whose MambaCache state is cumulative.
    if prompt_cache_save is not None:

        def _patched_process_prompts(prompts, _self=batch_gen):
            batch = _orig_process_prompts(prompts)
            for e, uid in enumerate(batch.uids):
                if batch.num_tokens[e] == 0:
                    try:
                        prompt_cache_save(uid, batch.extract_cache(e))
                    except Exception:
                        pass
            return batch

        batch_gen._process_prompts = _patched_process_prompts

    def _generation_step(self=batch_gen):
        """Run one generation step on the active batch. Returns responses."""
        batch = self.active_batch
        if batch is None or len(batch) == 0:
            return []

        tic_gen = _time.perf_counter()
        y, logprobs = batch.y, batch.logprobs
        for i, toks in enumerate(batch.tokens):
            batch.tokens[i] = mx.concatenate((toks, y[i : i + 1]))
        batch.y, batch.logprobs = self._step(
            y[:, None],
            batch.cache,
            batch.samplers,
            batch.logits_processors,
            batch.tokens,
        )
        mx.async_eval(batch.y, batch.logprobs)
        # Evaluate accumulated tokens to prevent Metal buffer buildup
        # from lazy mx.concatenate() chains holding AGXAllocation handles
        if batch.tokens:
            mx.async_eval(*batch.tokens)

        y = y.tolist()
        self._stats.generation_time += _time.perf_counter() - tic_gen

        keep_idx = []
        end_idx = []
        responses = []
        for e, (t, uid, num_tok, max_tok) in enumerate(
            zip(y, batch.uids, batch.num_tokens, batch.max_tokens)
        ):
            cache_out = None
            num_tok += 1
            batch.num_tokens[e] = num_tok
            if t in self.stop_tokens:
                finish_reason = "stop"
                end_idx.append(e)
            elif num_tok >= max_tok:
                finish_reason = "length"
                end_idx.append(e)
            else:
                finish_reason = None
                keep_idx.append(e)
            if finish_reason is not None:
                cache_out = batch.extract_cache(e)
            responses.append(
                self.Response(uid, t, logprobs[e], finish_reason, cache_out)
            )

        if len(end_idx):
            if len(keep_idx) > 0:
                batch.filter(keep_idx)
            else:
                self.active_batch = None

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

    def _chunked_next(self=batch_gen):  # noqa: C901
        """
        Replacement for _next() that chunks large prefills.

        Only intercepts when:
        1. A partial prefill is in progress (_partial is not None)
        2. The next prompt batch exceeds the budget

        Everything else delegates to the original _next().
        """
        # ----- Continue a partial prefill -----
        if self._partial is not None:
            # Check for pending aborts BEFORE processing next chunk
            if pending_abort_ids is not None and uid_to_request_id is not None:
                partial_rids = {uid_to_request_id.get(u) for u in self._partial["uids"]}
                aborted_rids = partial_rids & pending_abort_ids
                if aborted_rids:
                    logger.info(
                        f"[chunked_prefill] abort detected mid-prefill, "
                        f"clearing partial for: {aborted_rids}"
                    )
                    self._partial = None
                    mx.clear_cache()
                    return self._generation_step()

            tic = _time.perf_counter()
            partial = self._partial
            inputs = partial["inputs"]
            prompt_cache = partial["cache"]
            remaining = inputs.shape[1]
            prompt_checkpoint = max(1, int(partial.get("prompt_checkpoint", 1)))

            n_to_process = (
                min(budget, remaining - prompt_checkpoint)
                if remaining > prompt_checkpoint
                else 0
            )

            if n_to_process > 0:
                self.model(mx.contiguous(inputs[:, :n_to_process]), cache=prompt_cache)
                mx.eval([c.state for c in prompt_cache])
                inputs = inputs[:, n_to_process:]
                partial["inputs"] = inputs
                partial["processed"] += n_to_process

                self.prompt_progress_callback(
                    [
                        (uid, partial["processed"], partial["total"])
                        for uid in partial["uids"]
                    ]
                )

                # Save intermediate cache for disconnect resilience
                if mid_prefill_save is not None and len(partial["uids"]) == 1:
                    mid_prefill_save(
                        partial["uids"][0], partial["processed"], prompt_cache
                    )

                if partial.get("is_cached"):
                    mx.clear_cache()

            # Check if prefill is done once only the checkpoint tail remains.
            if inputs.shape[1] <= prompt_checkpoint:
                # Finalize
                if partial.get("is_cached"):
                    mx.eval([c.state for c in prompt_cache])
                    inputs = partial["last_inputs"]

                for c in prompt_cache:
                    c.finalize()

                if self.prompt_checkpoint_callback is not None:
                    self.prompt_checkpoint_callback(
                        [
                            (
                                uid,
                                prompt_checkpoint,
                                _lazy_extract_cache(prompt_cache, i),
                            )
                            for i, uid in enumerate(partial["uids"])
                        ]
                    )
                mx.clear_cache()

                # Mirror upstream BatchGenerator semantics: after finalize() and
                # the checkpoint callback, replay the remaining checkpoint tail
                # except for the final token, which _step() consumes.
                if prompt_checkpoint > 1:
                    self.model(
                        mx.contiguous(inputs[:, : prompt_checkpoint - 1]),
                        cache=prompt_cache,
                    )
                    mx.eval([c.state for c in prompt_cache])
                    mx.clear_cache()

                y, logprobs = self._step(
                    inputs,
                    prompt_cache,
                    partial["samplers"],
                    partial["logits_processors"],
                    partial["tokens"],
                )
                mx.async_eval(y, logprobs)

                new_batch = _batch_cls(
                    list(partial["uids"]),
                    y,
                    list(logprobs),
                    list(partial["max_tokens"]),
                    [0] * len(partial["uids"]),
                    prompt_cache,
                    list(partial["samplers"]),
                    list(partial["logits_processors"]),
                    partial["tokens"],
                )

                # Save prompt-only cache BEFORE merging into active batch.
                # This is the chunked-prefill equivalent of the
                # _patched_process_prompts hook — at this point the cache
                # contains the exact prompt-only state (num_tokens == 0).
                if prompt_cache_save is not None and len(partial["uids"]) == 1:
                    uid = partial["uids"][0]
                    try:
                        prompt_cache_save(uid, new_batch.extract_cache(0))
                    except Exception:
                        pass

                if self.active_batch is None:
                    self.active_batch = new_batch
                else:
                    self.active_batch.extend(new_batch)

                self._partial = None
                self._stats.prompt_time += _time.perf_counter() - tic
            else:
                # Not done yet — record prompt time for this chunk
                self._stats.prompt_time += _time.perf_counter() - tic

            # Generation step for active requests between chunks
            return self._generation_step()

        # ----- No partial — check if next prompt batch needs chunking -----
        num_active = len(self.active_batch) if self.active_batch else 0
        num_to_add = self.completion_batch_size - num_active

        if num_to_add >= self.prefill_batch_size and self.unprocessed_prompts:
            batch_prompts = self.unprocessed_prompts[: self.prefill_batch_size]
            if batch_prompts:
                total_tokens = sum(len(p[1]) for p in batch_prompts)

                # Check if any prompt has a prefix_boundary that
                # requires two-phase prefill for cache save at that boundary.
                _needs_boundary_split = False
                if requests is not None and uid_to_request_id is not None:
                    for _uid, _toks, *_ in batch_prompts:
                        _rid = uid_to_request_id.get(_uid)
                        _req = requests.get(_rid) if _rid else None
                        if _req and getattr(_req, "prefix_boundary", 0) > 0:
                            _needs_boundary_split = True
                            break

                if total_tokens > budget or _needs_boundary_split:
                    # Large prompt batch or prefix boundary — start partial prefill
                    tic = _time.perf_counter()

                    # Eval outstanding generation tokens before switching.
                    # Also drain pending async_eval when active_batch is None
                    # (previous request finished) — stale async_eval work on
                    # generation_stream can block subsequent model forwards.
                    if self.active_batch is not None:
                        mx.eval(self.active_batch.y, self.active_batch.logprobs)
                        self._stats.generation_time += _time.perf_counter() - tic
                        tic = _time.perf_counter()
                    else:
                        mx.clear_cache()

                    (
                        uids,
                        inputs_raw,
                        max_tokens_list,
                        caches,
                        samplers,
                        logits_processors,
                        prompt_checkpoints,
                    ) = zip(*batch_prompts)
                    lengths = [len(p) for p in inputs_raw]
                    max_length = max(lengths)
                    padding = [max_length - ln for ln in lengths]
                    tokens = [mx.array(inp) for inp in inputs_raw]
                    # Match mlx-lm's prompt_checkpoint contract: positive values
                    # name the checkpoint token position in the prompt, while
                    # non-positive values already encode an offset from the end.
                    checkpoint_offsets = [
                        (ln - pc if pc > 0 else -pc)
                        for ln, pc in zip(lengths, prompt_checkpoints)
                    ]
                    prompt_checkpoint = max(1, max(checkpoint_offsets))
                    is_cached = not all(c[0].empty() for c in caches)

                    self._stats.prompt_tokens += sum(lengths)

                    if not is_cached:
                        padded = _left_pad_prompts(inputs_raw, max_length=max_length)
                        prompt_cache = _make_cache(
                            self.model, padding, self.max_kv_size
                        )
                    else:
                        last_inputs = mx.array(
                            [p[-prompt_checkpoint:] for p in inputs_raw]
                        )
                        padded = _right_pad_prompts(inputs_raw, max_length=max_length)
                        prompt_cache = _merge_caches(caches)
                        for c in prompt_cache:
                            c.prepare(
                                lengths=[ln - prompt_checkpoint for ln in lengths],
                                right_padding=padding,
                            )

                    # Remove from unprocessed
                    self.unprocessed_prompts = self.unprocessed_prompts[
                        self.prefill_batch_size :
                    ]

                    # Process first chunk — if prefix_boundary is set,
                    # use it as the first chunk size so that mid_prefill_save
                    # can capture the exact prefix cache state (critical for
                    # hybrid Mamba+Transformer models where trim is unsafe).
                    # When the request already has cached tokens (cache hit),
                    # adjust the boundary relative to the remaining tokens.
                    _first_chunk = budget
                    if _needs_boundary_split and len(batch_prompts) == 1:
                        _uid0 = uids[0]
                        _rid0 = uid_to_request_id.get(_uid0)
                        _req0 = requests.get(_rid0) if _rid0 else None
                        _pb = getattr(_req0, "prefix_boundary", 0) if _req0 else 0
                        _cached = getattr(_req0, "cached_tokens", 0) if _req0 else 0
                        _adjusted_pb = _pb - _cached
                        if 0 < _adjusted_pb < padded.shape[1] - prompt_checkpoint + 1:
                            _first_chunk = _adjusted_pb
                    n_to_process = min(
                        _first_chunk, padded.shape[1] - prompt_checkpoint
                    )
                    if n_to_process > 0:
                        self.model(
                            mx.contiguous(padded[:, :n_to_process]),
                            cache=prompt_cache,
                        )
                        mx.eval([c.state for c in prompt_cache])
                        padded = padded[:, n_to_process:]
                        if is_cached:
                            mx.clear_cache()

                    self._partial = {
                        "uids": list(uids),
                        "inputs": padded,
                        "cache": prompt_cache,
                        "tokens": tokens,
                        "max_tokens": list(max_tokens_list),
                        "samplers": list(samplers),
                        "logits_processors": list(logits_processors),
                        "prompt_checkpoint": prompt_checkpoint,
                        "processed": n_to_process,
                        "total": max_length,
                        "is_cached": is_cached,
                    }
                    if is_cached:
                        self._partial["last_inputs"] = last_inputs

                    self.prompt_progress_callback(
                        [
                            (uid, n_to_process, max_length)
                            for uid in self._partial["uids"]
                        ]
                    )

                    # Save intermediate cache for disconnect resilience
                    if mid_prefill_save is not None and len(uids) == 1:
                        mid_prefill_save(uids[0], n_to_process, prompt_cache)

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

                    # Generation step for active requests
                    return self._generation_step()

                else:
                    # Small prompt batch — process directly without _orig_next.
                    # _orig_next's while loop processes multiple batches per call
                    # which causes batch-dimension mismatches in DeltaRNN conv_state
                    # when mixing prefix-cached and fresh prompts.
                    # Processing one batch per _next call avoids this.
                    tic = _time.perf_counter()

                    # Eval outstanding generation tokens before prefill.
                    # Also drain when active_batch is None to clear stale
                    # async_eval work from the previous request.
                    if self.active_batch is not None:
                        mx.eval(self.active_batch.y, self.active_batch.logprobs)
                        self._stats.generation_time += _time.perf_counter() - tic
                        tic = _time.perf_counter()
                    else:
                        mx.clear_cache()

                    new_batch = self._process_prompts(batch_prompts)
                    self.unprocessed_prompts = self.unprocessed_prompts[
                        self.prefill_batch_size :
                    ]

                    if self.active_batch is None:
                        self.active_batch = new_batch
                    else:
                        self.active_batch.extend(new_batch)

                    self._stats.prompt_time += _time.perf_counter() - tic
                    return self._generation_step()

        # Pure generation or no work — run generation step directly
        return self._generation_step()

    def _patched_remove(uids_to_remove, _self=batch_gen):
        """Clear partial state if aborted request is being prefilled."""
        if _self._partial is not None:
            partial_uids = set(_self._partial["uids"])
            if partial_uids & set(uids_to_remove):
                logger.info(
                    f"[chunked_prefill] clearing partial state for aborted uids: "
                    f"{partial_uids & set(uids_to_remove)}"
                )
                _self._partial = None
                mx.clear_cache()  # flush Metal encoders after dropping partial state
        _orig_remove(uids_to_remove)

    batch_gen._next = _chunked_next
    batch_gen._generation_step = _generation_step
    batch_gen.remove = _patched_remove

    logger.info(f"[chunked_prefill] installed with budget={budget} tokens per step")

vllm_mlx.scheduler._configure_chunked_prefill

_configure_chunked_prefill(scheduler: Scheduler, batch_gen: BatchGenerator, budget: int, prompt_cache_save) -> None

Enable the matching legacy or native mlx-lm chunked-prefill API.

Source code in vllm_mlx/scheduler.py

def _configure_chunked_prefill(
    scheduler: "Scheduler",
    batch_gen: "BatchGenerator",
    budget: int,
    prompt_cache_save,
) -> None:
    """Enable the matching legacy or native mlx-lm chunked-prefill API."""
    legacy_api = hasattr(batch_gen, "_process_prompts") and hasattr(
        batch_gen, "active_batch"
    )
    if legacy_api:
        save_interval = scheduler.config.mid_prefill_save_interval
        mid_prefill_save = None
        if save_interval > 0 and scheduler.memory_aware_cache is not None:
            mid_prefill_save = scheduler._make_mid_prefill_save_callback(save_interval)
            logger.info(
                "[mid_prefill_cache] enabled, interval=%s",
                save_interval,
            )
        _install_chunked_prefill(
            batch_gen,
            budget,
            mid_prefill_save,
            prompt_cache_save=prompt_cache_save,
            pending_abort_ids=scheduler._pending_abort_ids,
            uid_to_request_id=scheduler.uid_to_request_id,
            requests=scheduler.requests,
        )
        return

    native_api = all(
        hasattr(batch_gen, attribute)
        for attribute in (
            "_prompt_batch",
            "_generation_batch",
            "_unprocessed_sequences",
            "_next",
        )
    )
    if native_api:
        # Native mlx-lm chunking processes at most this many prompt tokens per
        # scheduler turn and returns to generation between turns. Its internal
        # API has no safe extension point for the legacy prompt-cache and
        # mid-prefill callbacks, which were already unavailable on this layout.
        batch_gen.prefill_step_size = budget
        logger.info(
            "Chunked prefill enabled through native mlx-lm BatchGenerator: "
            "budget=%s tokens per step",
            budget,
        )
        return

    logger.warning(
        "Chunked prefill disabled: mlx-lm BatchGenerator matches neither "
        "the legacy nor native chunked-prefill API."
    )

vllm_mlx.scheduler._install_mtp

_install_mtp(batch_gen: BatchGenerator, model: Any, num_draft_tokens: int = 1, optimistic: bool = False, stats_state: Optional[_MTPStatsState] = None) -> None

Monkey-patch a BatchGenerator to use MTP (Multi-Token Prediction) with always-advance strategy for hybrid MambaCache + KVCache.

Flow per generation step: 1. Use skip_state logits/hidden OR run model forward -> sample primary 2. MTP head drafts one token after primary 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 KVCache by 1, skip_state from pos 0 (no cold start) 5. Draft is emitted in the NEXT generation step after primary

Source code in vllm_mlx/scheduler.py

def _install_mtp(
    batch_gen: "BatchGenerator",
    model: Any,
    num_draft_tokens: int = 1,
    optimistic: bool = False,
    stats_state: Optional["_MTPStatsState"] = None,
) -> None:
    """
    Monkey-patch a BatchGenerator to use MTP (Multi-Token Prediction)
    with always-advance strategy for hybrid MambaCache + KVCache.

    Flow per generation step:
    1. Use skip_state logits/hidden OR run model forward -> sample primary
    2. MTP head drafts one token after primary
    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 KVCache by 1, skip_state from pos 0 (no cold start)
    5. Draft is emitted in the NEXT generation step after primary
    """
    _orig_step = batch_gen._step

    # Greedy sampler for MTP draft tokens
    _draft_sampler = make_sampler(temp=0.0)

    # Skip state: when MTP accepts, the cache already consumed [primary, draft].
    # Next _step call receives primary as input but must NOT re-feed it.
    # Instead, use stored logits from the verify pass.
    # Format: {'logits': (B, V), 'hidden': (B, 1, H)}
    _skip_state = [None]

    # Deferred drafts: draft tokens to emit in the NEXT generation step,
    # keyed by UID for stability across batch changes.
    # Format: {uid: {'token': int, 'logprobs': mx.array}}
    _deferred_drafts = {}

    # Scheduler-created generators share one state so sampler-driven generator
    # replacement does not reset the operator-facing counters.
    if stats_state is None:
        stats_state = _MTPStatsState()
    _mtp_stats = stats_state.counters
    _mtp_bypass_counts = stats_state.bypass_counts
    _mtp_stats_lock = stats_state.lock

    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(_mtp_bypass_counts)
        verified = accepted + rejected
        return {
            "enabled": True,
            "requested_draft_tokens": num_draft_tokens,
            "effective_draft_tokens": 1,
            "mode": (
                "always_advance_optimistic" if optimistic else "always_advance_verified"
            ),
            "attempted": attempted,
            "accepted": accepted,
            "rejected": rejected,
            "errors": errors,
            "acceptance_rate": accepted / verified if verified else 0.0,
            "bypass_counts": bypass_counts,
            "bypass_counts_semantics": "per_condition_overlapping_not_total_steps",
        }

    batch_gen.get_mtp_stats = _get_mtp_stats

    def _mtp_bypass_reasons(input_tokens, prompt_cache):
        reasons = []
        if input_tokens.shape[1] > 1:
            reasons.append("prefill")
        if batch_gen.active_batch is None:
            reasons.append("no_active_batch")
        elif prompt_cache is not batch_gen.active_batch.cache:
            reasons.append("cache_mismatch")
        return reasons

    def _record_mtp_bypass(reasons) -> None:
        with _mtp_stats_lock:
            for reason in reasons:
                _mtp_bypass_counts[reason] += 1

    def _mtp_step(
        input_tokens,
        prompt_cache,
        samplers,
        logits_processors,
        tokens,
    ):
        """
        Extended _step with MTP always-advance strategy.

        Every step (after skip):
        1. Use skip_state logits/hidden OR run model forward
        2. Sample primary token P
        3. MTP head drafts token D
        4. Verify [P, D] in one model call (always advances cache)
        5. Accept: skip_state from position 1 (after D), defer D
           Reject: trim KVCache by 1, skip_state from position 0 (after P)

        No snapshot/restore — eliminates cold starts after rejection.
        MambaCache layers accept minor pollution on reject (exponential decay).

        During prefill (multi-token input), MTP is skipped entirely.
        """
        batch_size = input_tokens.shape[0]

        # --- Prefill guard: skip MTP for multi-token input,
        # during _process_prompts (active_batch not yet set), or when
        # the cache doesn't belong to the active batch (e.g. during
        # _process_prompts in the 2nd+ iteration of _orig_next's loop
        # or during _chunked_next partial prefill finalization).
        bypass_reasons = _mtp_bypass_reasons(input_tokens, prompt_cache)
        if bypass_reasons:
            _record_mtp_bypass(bypass_reasons)
            _skip_state[0] = None
            return _orig_step(
                input_tokens,
                prompt_cache,
                samplers,
                logits_processors,
                tokens,
            )

        # --- Check skip state from previous MTP step ---
        skip = _skip_state[0]
        if skip is not None:
            if skip["logits"].shape[0] != batch_size:
                # Batch size changed since skip was stored — invalidate
                skip = None
                _skip_state[0] = None

        if skip is not None:
            # Skip mode: model already processed input_tokens during
            # previous verify. Use stored logits + hidden instead.
            logits = skip["logits"]
            hidden_states = skip["hidden"]
            _skip_state[0] = None
        else:
            # Normal model forward
            model_output = model(input_tokens, cache=prompt_cache, return_hidden=True)
            if isinstance(model_output, tuple):
                logits, hidden_states = model_output
            else:
                # Model doesn't support return_hidden — fall back
                return _orig_step(
                    input_tokens,
                    prompt_cache,
                    samplers,
                    logits_processors,
                    tokens,
                )
            logits = logits[:, -1, :]

        # --- Apply logits processors + sample primary ---
        logits_processors = _normalize_logits_processors(logits_processors) or []
        if any(logits_processors):
            logger.debug(
                f"[logits_proc] applying {sum(len(lp) for lp in logits_processors)} "
                f"processors to batch_size={batch_size}"
            )
            processed_logits = []
            for e in range(batch_size):
                sample_logits = logits[e : e + 1]
                for processor in logits_processors[e]:
                    sample_logits = processor(tokens[e], sample_logits)
                processed_logits.append(sample_logits)
            logits = mx.concatenate(processed_logits, axis=0)

        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
        if any(samplers):
            all_samples = []
            for e in range(batch_size):
                sample_sampler = samplers[e] or batch_gen.sampler
                sampled = sample_sampler(logprobs[e : e + 1])
                all_samples.append(sampled)
            primary_tokens = mx.concatenate(all_samples, axis=0)
        else:
            primary_tokens = batch_gen.sampler(logprobs)

        # Get current UIDs (guaranteed non-empty: prefill guard above
        # prevents MTP from running when active_batch is None).
        current_uids = list(batch_gen.active_batch.uids)

        # --- MTP draft + always-advance verify ---
        try:
            with _mtp_stats_lock:
                _mtp_stats["attempted"] += 1
            # Draft: predict token n+2 from hidden states + primary (n+1)
            draft_logits = model.mtp_forward(
                hidden_states[:, -1:, :],
                primary_tokens[:, None],
                mtp_cache=None,
            )
            draft_logits = draft_logits[:, -1, :]
            draft_logprobs = draft_logits - mx.logsumexp(
                draft_logits, axis=-1, keepdims=True
            )
            draft_tokens = _draft_sampler(draft_logprobs)

            # Always-advance: feed [primary, draft] and let cache advance.
            #
            # Hybrid models (e.g. Qwen3-Next) mix attention (KVCache) and
            # recurrent layers (MambaCache/DeltaRNN).  KVCache supports
            # trim(1) to undo the draft token on reject, but recurrent
            # state is irreversible — rejected drafts permanently pollute
            # the RNN state, causing progressive output corruption.
            #
            # For hybrid models we snapshot recurrent state before verify
            # and on reject: trim KV by 2 (remove both P and D), restore
            # RNN snapshot, then re-advance with just P so both cache
            # types end up consistent at [..., P].
            # Skip RNN snapshots in optimistic mode — it never rejects,
            # so the copies are wasted (~147 MB/step of lazy graph nodes
            # that prevent pre-verify Metal buffers from being freed).
            _rnn_snapshots = {}
            if not optimistic:
                for _ci, _c in enumerate(prompt_cache):
                    if not (hasattr(_c, "is_trimmable") and _c.is_trimmable()):
                        if hasattr(_c, "state"):
                            _rnn_snapshots[_ci] = [
                                s.copy() if s is not None else None for s in _c.state
                            ]

            verify_input = mx.concatenate(
                [primary_tokens[:, None], draft_tokens[:, None]], axis=1
            )
            verify_output = model(verify_input, cache=prompt_cache, return_hidden=True)
            if isinstance(verify_output, tuple):
                verify_logits, verify_hidden = verify_output
            else:
                verify_logits = verify_output
                verify_hidden = None

            if optimistic:
                # --- OPTIMISTIC: always accept, zero sync ---
                if verify_hidden is not None:
                    _skip_state[0] = {
                        "logits": verify_logits[:, 1, :],
                        "hidden": verify_hidden[:, -1:, :],
                    }
                    verify_lp = verify_logits[:, 0, :] - mx.logsumexp(
                        verify_logits[:, 0, :], axis=-1, keepdims=True
                    )
                    mx.async_eval(
                        _skip_state[0]["logits"],
                        _skip_state[0]["hidden"],
                        draft_tokens,
                        verify_lp,
                    )
                    for e in range(batch_size):
                        uid = current_uids[e]
                        _deferred_drafts[uid] = {
                            "token_array": draft_tokens[e : e + 1],
                            "logprobs": verify_lp[e],
                        }
                else:
                    _skip_state[0] = None
                with _mtp_stats_lock:
                    _mtp_stats["accepted"] += 1
            else:
                # --- VERIFIED MODE: single eval + Python comparison ---
                verify_pred = mx.argmax(verify_logits[:, 0, :], axis=-1)
                mx.eval(verify_pred, draft_tokens)
                pred_list = verify_pred.tolist()
                draft_list = draft_tokens.tolist()
                all_accepted = pred_list == draft_list

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

                else:
                    # --- REJECT (always-advance) ---
                    if _rnn_snapshots:
                        # Hybrid model: undo the entire verify pass
                        # (both P and D) for all cache types, then
                        # re-advance with just P for a consistent state.
                        for c in prompt_cache:
                            if (
                                hasattr(c, "is_trimmable")
                                and c.is_trimmable()
                                and hasattr(c, "trim")
                            ):
                                c.trim(2)
                        for _ci, _snap in _rnn_snapshots.items():
                            prompt_cache[_ci].state = _snap
                        # Re-advance with primary only — both KV and RNN
                        # now advance by exactly 1 (the primary token).
                        rerun_out = model(
                            primary_tokens[:, None],
                            cache=prompt_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:
                            _skip_state[0] = {
                                "logits": rerun_logits[:, -1, :],
                                "hidden": rerun_hidden[:, -1:, :],
                            }
                            mx.async_eval(
                                _skip_state[0]["logits"],
                                _skip_state[0]["hidden"],
                            )
                        else:
                            _skip_state[0] = None
                    else:
                        # Pure attention model: simple trim(1) is enough.
                        for c in prompt_cache:
                            if (
                                hasattr(c, "is_trimmable")
                                and c.is_trimmable()
                                and hasattr(c, "trim")
                            ):
                                c.trim(1)
                        if verify_hidden is not None:
                            _skip_state[0] = {
                                "logits": verify_logits[:, 0, :],
                                "hidden": verify_hidden[:, 0:1, :],
                            }
                            mx.async_eval(
                                _skip_state[0]["logits"],
                                _skip_state[0]["hidden"],
                            )
                        else:
                            _skip_state[0] = None
                    for uid in current_uids:
                        _deferred_drafts.pop(uid, None)
                    with _mtp_stats_lock:
                        _mtp_stats["rejected"] += 1

        except Exception as e:
            logger.debug(f"[MTP] draft/verify failed: {e}")
            _skip_state[0] = None
            with _mtp_stats_lock:
                _mtp_stats["errors"] += 1

        return primary_tokens, list(logprobs)

    # Wrap _next() to emit deferred MTP drafts after each primary token.
    # This works regardless of whether _chunked_next or original _next is
    # the current _next implementation, because it sits at the top level.
    # Store as attribute so it's always the correct reference, even after
    # BatchGenerator recreation.
    batch_gen._inner_next = batch_gen._next

    def _mtp_next(self=batch_gen):
        """Wrapper around _next that emits deferred MTP draft tokens.

        After each primary token, if the previous step's MTP draft was
        accepted, it is emitted as an additional response.
        """
        # Clear stale MTP state when no batch is active.
        # This prevents skip_state/deferred_drafts from a finished request
        # from leaking into the next request and causing stale computation
        # graph references on generation_stream.
        if self.active_batch is None:
            _skip_state[0] = None
            _deferred_drafts.clear()

        # Save deferred drafts from PREVIOUS step before _inner_next
        # runs _mtp_step, which may store NEW deferred drafts.
        prev_deferred = {}
        if self.active_batch is not None:
            for uid in self.active_batch.uids:
                if uid in _deferred_drafts:
                    prev_deferred[uid] = _deferred_drafts.pop(uid)

        # Run the inner _next (original or chunked) — calls _mtp_step
        responses = self._inner_next()

        if not prev_deferred or not responses:
            return responses

        # Augment responses with deferred drafts from the previous step.
        # The Response from _next reports the OLD batch.y (the primary
        # from the *previous* _step call). The deferred draft follows
        # that primary in the token stream, so emit it AFTER the primary.
        augmented = []
        draft_end_uids = set()
        for r in responses:
            uid = r.uid

            # Emit the primary response first
            augmented.append(r)

            if r.finish_reason is not None:
                # Sequence ended with primary — discard any pending draft
                _deferred_drafts.pop(uid, None)
                prev_deferred.pop(uid, None)
                continue

            # Emit deferred draft AFTER its primary
            if uid in prev_deferred:
                draft_info = prev_deferred.pop(uid)
                if "token" in draft_info:
                    draft_t = draft_info["token"]
                else:
                    draft_t = draft_info["token_array"].item()
                draft_lp = draft_info["logprobs"]

                if draft_t in self.stop_tokens:
                    augmented.append(
                        self.Response(uid, draft_t, draft_lp, "stop", None)
                    )
                    draft_end_uids.add(uid)
                else:
                    draft_finish = None
                    batch = self.active_batch
                    if batch is not None:
                        for e, bu in enumerate(batch.uids):
                            if bu == uid:
                                batch.num_tokens[e] += 1
                                batch.tokens[e] = mx.concatenate(
                                    (batch.tokens[e], mx.array([draft_t]))
                                )
                                if batch.num_tokens[e] >= batch.max_tokens[e]:
                                    draft_finish = "length"
                                    draft_end_uids.add(uid)
                                break

                    draft_cache_out = None
                    if draft_finish is not None and batch is not None:
                        for e, bu in enumerate(batch.uids):
                            if bu == uid:
                                draft_cache_out = batch.extract_cache(e)
                                break

                    augmented.append(
                        self.Response(
                            uid, draft_t, draft_lp, draft_finish, draft_cache_out
                        )
                    )

        # Remove sequences that finished due to draft tokens
        if draft_end_uids and self.active_batch is not None:
            keep = [
                e
                for e, u in enumerate(self.active_batch.uids)
                if u not in draft_end_uids
            ]
            if keep:
                self.active_batch.filter(keep)
            else:
                self.active_batch = None

        return augmented

    batch_gen._step = _mtp_step
    batch_gen._next = _mtp_next

    if num_draft_tokens != 1:
        logger.warning(
            "[MTP] num_draft_tokens=%d requested, but the current batched MTP "
            "path drafts exactly one token per verify step",
            num_draft_tokens,
        )
    mode_str = "optimistic (no verify)" if optimistic else "always-advance"
    logger.info(
        f"[MTP] installed with num_draft_tokens={num_draft_tokens}, "
        f"effective_draft_tokens=1, {mode_str} mode"
    )

vllm_mlx.scheduler._mtp_status_snapshot

_mtp_status_snapshot(batch_generator) -> Dict[str, Any]

Source code in vllm_mlx/scheduler.py

def _mtp_status_snapshot(batch_generator) -> Dict[str, Any]:
    get_mtp_stats = getattr(batch_generator, "get_mtp_stats", None)
    if callable(get_mtp_stats):
        return {"mtp": get_mtp_stats()}
    return {}

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.scheduler._normalize_logits_processors · function

vllm_mlx.scheduler._normalize_logits_processors(logits_processors) -> not annotated

Normalize empty per-sequence processor slots to lists.

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: None; [processors or [] for processors in logits_processors]

Exceptions and behavior

Function _normalize_logits_processors has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L46-L50.

vllm_mlx.scheduler._sanitize_batch_generator_logits_processors · function

vllm_mlx.scheduler._sanitize_batch_generator_logits_processors(batch_generator) -> None

Sanitize stale BatchGenerator processor state before decode.

Parameters

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

Returns

• Type: None

Exceptions and behavior

Function _sanitize_batch_generator_logits_processors calls getattr, hasattr, _normalize_logits_processors, isinstance. No direct raise statement appears in this definition.

View source #L53-L65.

vllm_mlx.scheduler.SchedulingPolicy · class

vllm_mlx.scheduler.SchedulingPolicy()

Scheduling policy for request ordering.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.scheduler.SchedulingPolicy

Exceptions and behavior

Class SchedulingPolicy derives from Enum and declares 0 direct member(s). No direct raise statement appears in this definition.

View source #L68-L72.

vllm_mlx.scheduler.SchedulerConfig · class

vllm_mlx.scheduler.SchedulerConfig(max_num_seqs: int = 256, max_num_batched_tokens: int = 8192, policy: SchedulingPolicy = SchedulingPolicy.FCFS, prefill_batch_size: int = 8, completion_batch_size: int = 32, prefill_step_size: int = 2048, mllm_prefill_step_size: Optional[int] = None, enable_prefix_cache: bool = True, prefix_cache_size: int = 100, use_memory_aware_cache: bool = True, cache_memory_mb: Optional[int] = None, cache_memory_percent: float = 0.2, kv_cache_quantization: bool = False, kv_cache_quantization_bits: int = 8, kv_cache_quantization_group_size: int = 64, kv_cache_min_quantize_tokens: int = 256, use_paged_cache: bool = False, paged_cache_block_size: int = 64, max_cache_blocks: int = 1000, chunked_prefill_tokens: int = 0, mid_prefill_save_interval: int = 8192, ssd_cache_dir: Optional[str] = None, ssd_cache_max_gb: float = 10.0, max_kv_size: int = 0, enable_mtp: bool = False, mtp_num_draft_tokens: int = 1, mtp_optimistic: bool = False)

Configuration for the scheduler.

Parameters

Name	Type	Required	Default	Description
max_num_seqs	int	no	256	Optional constructor field; defaults to 256.
max_num_batched_tokens	int	no	8192	Optional constructor field; defaults to 8192.
policy	SchedulingPolicy	no	SchedulingPolicy.FCFS	Optional constructor field; defaults to SchedulingPolicy.FCFS.
prefill_batch_size	int	no	8	Optional constructor field; defaults to 8.
completion_batch_size	int	no	32	Optional constructor field; defaults to 32.
prefill_step_size	int	no	2048	Optional constructor field; defaults to 2048.
mllm_prefill_step_size	Optional[int]	no	None	Optional constructor field; defaults to None.
enable_prefix_cache	bool	no	True	Optional constructor field; defaults to True.
prefix_cache_size	int	no	100	Optional constructor field; defaults to 100.
use_memory_aware_cache	bool	no	True	Optional constructor field; defaults to True.
cache_memory_mb	Optional[int]	no	None	Optional constructor field; defaults to None.
cache_memory_percent	float	no	0.2	Optional constructor field; defaults to 0.2.
kv_cache_quantization	bool	no	False	Optional constructor field; defaults to False.
kv_cache_quantization_bits	int	no	8	Optional constructor field; defaults to 8.
kv_cache_quantization_group_size	int	no	64	Optional constructor field; defaults to 64.
kv_cache_min_quantize_tokens	int	no	256	Optional constructor field; defaults to 256.
use_paged_cache	bool	no	False	Optional constructor field; defaults to False.
paged_cache_block_size	int	no	64	Optional constructor field; defaults to 64.
max_cache_blocks	int	no	1000	Optional constructor field; defaults to 1000.
chunked_prefill_tokens	int	no	0	Optional constructor field; defaults to 0.
mid_prefill_save_interval	int	no	8192	Optional constructor field; defaults to 8192.
ssd_cache_dir	Optional[str]	no	None	Optional constructor field; defaults to None.
ssd_cache_max_gb	float	no	10.0	Optional constructor field; defaults to 10.0.
max_kv_size	int	no	0	Optional constructor field; defaults to 0.
enable_mtp	bool	no	False	Optional constructor field; defaults to False.
mtp_num_draft_tokens	int	no	1	Optional constructor field; defaults to 1.
mtp_optimistic	bool	no	False	Optional constructor field; defaults to False.

Returns

• Constructs: vllm_mlx.scheduler.SchedulerConfig

Exceptions and behavior

Class SchedulerConfig declares 1 direct member(s). No direct raise statement appears in this definition.

View source #L76-L140.

vllm_mlx.scheduler.SchedulerConfig.__post_init__ · method

vllm_mlx.scheduler.SchedulerConfig.__post_init__() -> None

Method SchedulerConfig.__post_init__ calls ValueError; can raise ValueError.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SchedulerConfig.__post_init__ calls ValueError; can raise ValueError. Directly raised exceptions: ValueError.

View source #L138-L140.

vllm_mlx.scheduler.SchedulerOutput · class

vllm_mlx.scheduler.SchedulerOutput(scheduled_request_ids: List[str] = field(default_factory=list), num_scheduled_tokens: int = 0, finished_request_ids: Set[str] = field(default_factory=set), outputs: List[RequestOutput] = field(default_factory=list), has_work: bool = False)

Output from a scheduling step.

Parameters

Name	Type	Required	Default	Description
scheduled_request_ids	List[str]	no	field(default_factory=list)	Optional constructor field; defaults to field(default_factory=list).
num_scheduled_tokens	int	no	0	Optional constructor field; defaults to 0.
finished_request_ids	Set[str]	no	field(default_factory=set)	Optional constructor field; defaults to field(default_factory=set).
outputs	List[RequestOutput]	no	field(default_factory=list)	Optional constructor field; defaults to field(default_factory=list).
has_work	bool	no	False	Optional constructor field; defaults to False.

Returns

• Constructs: vllm_mlx.scheduler.SchedulerOutput

Exceptions and behavior

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

View source #L144-L160.

vllm_mlx.scheduler._install_prompt_cache_save · function

vllm_mlx.scheduler._install_prompt_cache_save(batch_gen: 'BatchGenerator', prompt_cache_save) -> None

Monkey-patch _process_prompts to capture prompt-only cache state.

Parameters

Name	Type	Required	Default	Description
batch_gen	'BatchGenerator'	yes	none	Required positional or keyword input.
prompt_cache_save	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Function _install_prompt_cache_save contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L163-L187.

vllm_mlx.scheduler._install_prompt_cache_save._patched_process_prompts · nested function

vllm_mlx.scheduler._install_prompt_cache_save._patched_process_prompts(prompts, _self = batch_gen) -> not annotated

Nested Function _install_prompt_cache_save._patched_process_prompts calls _orig_process_prompts, enumerate, prompt_cache_save, batch.extract_cache; returns batch.

Parameters

Name	Type	Required	Default	Description
prompts	not annotated	yes	none	Required positional or keyword input.
_self	not annotated	no	batch_gen	Optional positional or keyword input; defaults to batch_gen.

Returns

• Type: not annotated
• Direct return expressions: batch

Exceptions and behavior

Nested Function _install_prompt_cache_save._patched_process_prompts calls _orig_process_prompts, enumerate, prompt_cache_save, batch.extract_cache; returns batch. No direct raise statement appears in this definition.

View source #L177-L185.

vllm_mlx.scheduler._install_chunked_prefill · function

vllm_mlx.scheduler._install_chunked_prefill(batch_gen: 'BatchGenerator', budget: int, mid_prefill_save = None, prompt_cache_save = None, pending_abort_ids: Optional[Set[str]] = None, uid_to_request_id: Optional[Dict[int, str]] = None, requests: Optional[Dict[str, Any]] = None) -> None

Monkey-patch a BatchGenerator instance so that large prefills are broken into chunks of at most budget tokens each.

Parameters

Name	Type	Required	Default	Description
batch_gen	'BatchGenerator'	yes	none	The BatchGenerator to patch.
budget	int	yes	none	Max tokens per prefill chunk.
mid_prefill_save	not annotated	no	None	Optional callback(uid, processed, prompt_cache) called after each chunk to save intermediate KV cache state.
prompt_cache_save	not annotated	no	None	Optional positional or keyword input; defaults to None.
pending_abort_ids	Optional[Set[str]]	no	None	Optional positional or keyword input; defaults to None.
uid_to_request_id	Optional[Dict[int, str]]	no	None	Optional positional or keyword input; defaults to None.
requests	Optional[Dict[str, Any]]	no	None	Optional positional or keyword input; defaults to None.

Returns

• Type: None

Exceptions and behavior

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

View source #L190-L697.

vllm_mlx.scheduler._install_chunked_prefill._lazy_extract_cache · nested function

vllm_mlx.scheduler._install_chunked_prefill._lazy_extract_cache(cache, idx) -> not annotated

Nested Function _install_chunked_prefill._lazy_extract_cache calls c.extract; returns (c.extract(idx) for c in cache).

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: (c.extract(idx) for c in cache)

Exceptions and behavior

Nested Function _install_chunked_prefill._lazy_extract_cache calls c.extract; returns (c.extract(idx) for c in cache). No direct raise statement appears in this definition.

View source #L225-L226.

vllm_mlx.scheduler._install_chunked_prefill._batch_cls · nested class

vllm_mlx.scheduler._install_chunked_prefill._batch_cls(uids: List[int], y: Any, logprobs: List[Any], max_tokens: List[int], num_tokens: List[int], cache: List[Any], samplers: List[Any], logits_processors: List[Any], tokens: List[Any])

Nested Class _install_chunked_prefill._batch_cls declares 4 direct member(s).

Parameters

Name	Type	Required	Default	Description
uids	List[int]	yes	none	Required constructor field.
y	Any	yes	none	Required constructor field.
logprobs	List[Any]	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.
samplers	List[Any]	yes	none	Required constructor field.
logits_processors	List[Any]	yes	none	Required constructor field.
tokens	List[Any]	yes	none	Required constructor field.

Returns

• Constructs: vllm_mlx.scheduler._install_chunked_prefill._batch_cls

Exceptions and behavior

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

View source #L233-L273.

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.__len__ · nested function

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.__len__() -> not annotated

Nested Function _install_chunked_prefill._batch_cls.__len__ calls len; returns len(self.uids).

Parameters

This callable has no explicit inputs.

Returns

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

Exceptions and behavior

Nested Function _install_chunked_prefill._batch_cls.__len__ calls len; returns len(self.uids). No direct raise statement appears in this definition.

View source #L244-L245.

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.filter · nested function

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.filter(keep_idx: List[int]) -> not annotated

Nested Function _install_chunked_prefill._batch_cls.filter updates self.uids, self.logprobs, self.max_tokens, self.num_tokens; calls mx.array, c.filter.

Parameters

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

Returns

• Type: not annotated

Exceptions and behavior

Nested Function _install_chunked_prefill._batch_cls.filter updates self.uids, self.logprobs, self.max_tokens, self.num_tokens; calls mx.array, c.filter. No direct raise statement appears in this definition.

View source #L247-L258.

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.extend · nested function

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.extend(other) -> not annotated

Nested Function _install_chunked_prefill._batch_cls.extend updates self.y; calls self.uids.extend, mx.concatenate, self.logprobs.extend, self.num_tokens.extend.

Parameters

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

Returns

• Type: not annotated

Exceptions and behavior

Nested Function _install_chunked_prefill._batch_cls.extend updates self.y; calls self.uids.extend, mx.concatenate, self.logprobs.extend, self.num_tokens.extend. No direct raise statement appears in this definition.

View source #L260-L270.

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.extract_cache · nested function

vllm_mlx.scheduler._install_chunked_prefill._batch_cls.extract_cache(idx) -> not annotated

Nested Function _install_chunked_prefill._batch_cls.extract_cache calls c.extract; returns [c.extract(idx) for c in self.cache].

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: [c.extract(idx) for c in self.cache]

Exceptions and behavior

Nested Function _install_chunked_prefill._batch_cls.extract_cache calls c.extract; returns [c.extract(idx) for c in self.cache]. No direct raise statement appears in this definition.

View source #L272-L273.

vllm_mlx.scheduler._install_chunked_prefill._patched_process_prompts · nested function

vllm_mlx.scheduler._install_chunked_prefill._patched_process_prompts(prompts, _self = batch_gen) -> not annotated

Nested Function _install_chunked_prefill._patched_process_prompts calls _orig_process_prompts, enumerate, prompt_cache_save, batch.extract_cache; returns batch.

Parameters

Name	Type	Required	Default	Description
prompts	not annotated	yes	none	Required positional or keyword input.
_self	not annotated	no	batch_gen	Optional positional or keyword input; defaults to batch_gen.

Returns

• Type: not annotated
• Direct return expressions: batch

Exceptions and behavior

Nested Function _install_chunked_prefill._patched_process_prompts calls _orig_process_prompts, enumerate, prompt_cache_save, batch.extract_cache; returns batch. No direct raise statement appears in this definition.

View source #L291-L299.

vllm_mlx.scheduler._install_chunked_prefill._generation_step · nested function

vllm_mlx.scheduler._install_chunked_prefill._generation_step() -> not annotated

Run one generation step on the active batch.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: []; responses

Exceptions and behavior

Nested Function _install_chunked_prefill._generation_step updates self._stats.generation_time, self.active_batch, self._stats.generation_tokens; calls len, _time.perf_counter, enumerate, mx.concatenate; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L303-L360.

vllm_mlx.scheduler._install_chunked_prefill._chunked_next · nested function

vllm_mlx.scheduler._install_chunked_prefill._chunked_next() -> not annotated

Replacement for _next() that chunks large prefills.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: self._generation_step()

Exceptions and behavior

Nested Function _install_chunked_prefill._chunked_next updates self._partial, self.active_batch, self._stats.prompt_time, self._stats.generation_time; calls uid_to_request_id.get, logger.info, mx.clear_cache, self._generation_step; returns self._generation_step(). No direct raise statement appears in this definition.

View source #L362-L678.

vllm_mlx.scheduler._install_chunked_prefill._patched_remove · nested function

vllm_mlx.scheduler._install_chunked_prefill._patched_remove(uids_to_remove, _self = batch_gen) -> not annotated

Clear partial state if aborted request is being prefilled.

Parameters

Name	Type	Required	Default	Description
uids_to_remove	not annotated	yes	none	Required positional or keyword input.
_self	not annotated	no	batch_gen	Optional positional or keyword input; defaults to batch_gen.

Returns

• Type: not annotated

Exceptions and behavior

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

View source #L680-L691.

vllm_mlx.scheduler._MTPStatsState · class

vllm_mlx.scheduler._MTPStatsState(counters: Dict[str, int] = field(default_factory=lambda: {'attempted': 0, 'accepted': 0, 'rejected': 0, 'errors': 0}), bypass_counts: Dict[str, int] = field(default_factory=lambda: {'prefill': 0, 'no_active_batch': 0, 'cache_mismatch': 0}), lock: Any = field(default_factory=Lock))

Cumulative native-MTP counters shared across generator instances.

Parameters

Name	Type	Required	Default	Description
counters	Dict[str, int]	no	field(default_factory=lambda: {'attempted': 0, 'accepted': 0, 'rejected': 0, 'errors': 0})	Optional constructor field; defaults to field(default_factory=lambda: {'attempted': 0, 'accepted': 0, 'rejected': 0, 'errors': 0}).
bypass_counts	Dict[str, int]	no	field(default_factory=lambda: {'prefill': 0, 'no_active_batch': 0, 'cache_mismatch': 0})	Optional constructor field; defaults to field(default_factory=lambda: {'prefill': 0, 'no_active_batch': 0, 'cache_mismatch': 0}).
lock	Any	no	field(default_factory=Lock)	Optional constructor field; defaults to field(default_factory=Lock).

Returns

• Constructs: vllm_mlx.scheduler._MTPStatsState

Exceptions and behavior

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

View source #L701-L719.

vllm_mlx.scheduler._configure_chunked_prefill · function

vllm_mlx.scheduler._configure_chunked_prefill(scheduler: 'Scheduler', batch_gen: 'BatchGenerator', budget: int, prompt_cache_save) -> None

Enable the matching legacy or native mlx-lm chunked-prefill API.

Parameters

Name	Type	Required	Default	Description
scheduler	'Scheduler'	yes	none	Required positional or keyword input.
batch_gen	'BatchGenerator'	yes	none	Required positional or keyword input.
budget	int	yes	none	Required positional or keyword input.
prompt_cache_save	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Function _configure_chunked_prefill calls hasattr, scheduler._make_mid_prefill_save_callback, logger.info, _install_chunked_prefill; returns None. No direct raise statement appears in this definition.

View source #L722-L777.

vllm_mlx.scheduler._install_mtp · function

vllm_mlx.scheduler._install_mtp(batch_gen: 'BatchGenerator', model: Any, num_draft_tokens: int = 1, optimistic: bool = False, stats_state: Optional['_MTPStatsState'] = None) -> None

Monkey-patch a BatchGenerator to use MTP (Multi-Token Prediction) with always-advance strategy for hybrid MambaCache + KVCache.

Parameters

Name	Type	Required	Default	Description
batch_gen	'BatchGenerator'	yes	none	Required positional or keyword input.
model	Any	yes	none	Required positional or keyword input.
num_draft_tokens	int	no	1	Optional positional or keyword input; defaults to 1.
optimistic	bool	no	False	Optional positional or keyword input; defaults to False.
stats_state	Optional['_MTPStatsState']	no	None	Optional positional or keyword input; defaults to None.

Returns

• Type: None

Exceptions and behavior

Function _install_mtp calls make_sampler, _MTPStatsState, logger.warning, logger.info. No direct raise statement appears in this definition.

View source #L780-L1262.

vllm_mlx.scheduler._install_mtp._get_mtp_stats · nested function

vllm_mlx.scheduler._install_mtp._get_mtp_stats() -> Dict[str, Any]

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

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': 'always_advance_opti…

Exceptions and behavior

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

View source #L823-L845.

vllm_mlx.scheduler._install_mtp._mtp_bypass_reasons · nested function

vllm_mlx.scheduler._install_mtp._mtp_bypass_reasons(input_tokens, prompt_cache) -> not annotated

Nested Function _install_mtp._mtp_bypass_reasons calls reasons.append; returns reasons.

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: reasons

Exceptions and behavior

Nested Function _install_mtp._mtp_bypass_reasons calls reasons.append; returns reasons. No direct raise statement appears in this definition.

View source #L849-L857.

vllm_mlx.scheduler._install_mtp._record_mtp_bypass · nested function

vllm_mlx.scheduler._install_mtp._record_mtp_bypass(reasons) -> None

Nested Function _install_mtp._record_mtp_bypass contains no state mutation, call, raise, return, await, or yield.

Parameters

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

Returns

• Type: None

Exceptions and behavior

Nested Function _install_mtp._record_mtp_bypass contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L859-L862.

vllm_mlx.scheduler._install_mtp._mtp_step · nested function

vllm_mlx.scheduler._install_mtp._mtp_step(input_tokens, prompt_cache, samplers, logits_processors, tokens) -> not annotated

Extended _step with MTP always-advance strategy.

Parameters

Name	Type	Required	Default	Description
input_tokens	not annotated	yes	none	Required positional or keyword input.
prompt_cache	not annotated	yes	none	Required positional or keyword input.
samplers	not annotated	yes	none	Required positional or keyword input.
logits_processors	not annotated	yes	none	Required positional or keyword input.
tokens	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: _orig_step(input_tokens, prompt_cache, samplers, logits_processors, tokens); (primary_tokens, list(logprobs))

Exceptions and behavior

Nested Function _install_mtp._mtp_step calls _mtp_bypass_reasons, _record_mtp_bypass, _orig_step, model; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L864-L1138.

vllm_mlx.scheduler._install_mtp._mtp_next · nested function

vllm_mlx.scheduler._install_mtp._mtp_next() -> not annotated

Wrapper around _next that emits deferred MTP draft tokens.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: responses; augmented

Exceptions and behavior

Nested Function _install_mtp._mtp_next updates self.active_batch; calls _deferred_drafts.clear, _deferred_drafts.pop, self._inner_next, set; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1147-L1247.

vllm_mlx.scheduler._mtp_status_snapshot · function

vllm_mlx.scheduler._mtp_status_snapshot(batch_generator) -> Dict[str, Any]

Function _mtp_status_snapshot calls getattr, callable, get_mtp_stats; has 2 explicit return paths.

Parameters

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

Returns

• Type: Dict[str, Any]
• Direct return expressions: {'mtp': get_mtp_stats()}; {}

Exceptions and behavior

Function _mtp_status_snapshot calls getattr, callable, get_mtp_stats; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1265-L1269.

vllm_mlx.scheduler.Scheduler · class

vllm_mlx.scheduler.Scheduler(model: Any, tokenizer: Any, config: Optional[SchedulerConfig] = None)

Scheduler for continuous batching using mlx-lm BatchGenerator.

Parameters

Name	Type	Required	Default	Description
model	Any	yes	none	The MLX model
tokenizer	Any	yes	none	The tokenizer
config	Optional[SchedulerConfig]	no	None	Scheduler configuration

Returns

• Constructs: vllm_mlx.scheduler.Scheduler

Exceptions and behavior

Class Scheduler declares 52 direct member(s). No direct raise statement appears in this definition.

View source #L1272-L3518.

vllm_mlx.scheduler.Scheduler.__init__ · method

vllm_mlx.scheduler.Scheduler.__init__(model: Any, tokenizer: Any, config: Optional[SchedulerConfig] = None) -> not annotated

Initialize the scheduler.

Parameters

Name	Type	Required	Default	Description
model	Any	yes	none	The MLX model
tokenizer	Any	yes	none	The tokenizer
config	Optional[SchedulerConfig]	no	None	Scheduler configuration

Returns

• Type: not annotated

Exceptions and behavior

Method Scheduler.__init__ updates self.model, self.tokenizer, self.config, self._actual_tokenizer; calls SchedulerConfig, self._get_actual_tokenizer, deque, set. No direct raise statement appears in this definition.

View source #L1286-L1402.

vllm_mlx.scheduler.Scheduler._get_actual_tokenizer · method

vllm_mlx.scheduler.Scheduler._get_actual_tokenizer(tokenizer: Any) -> Any

Get the actual tokenizer from a processor or tokenizer.

Parameters

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

Returns

• Type: Any
• Direct return expressions: tokenizer; tokenizer.tokenizer

Exceptions and behavior

Method Scheduler._get_actual_tokenizer calls hasattr, callable; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1404-L1418.

vllm_mlx.scheduler.Scheduler._decode_tokens · method

vllm_mlx.scheduler.Scheduler._decode_tokens(token_ids: List[int]) -> str

Decode token IDs to text, handling both tokenizers and processors.

Parameters

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

Returns

• Type: str
• Direct return expressions: self._actual_tokenizer.decode(token_ids)

Exceptions and behavior

Method Scheduler._decode_tokens calls self._actual_tokenizer.decode; returns self._actual_tokenizer.decode(token_ids). No direct raise statement appears in this definition.

View source #L1420-L1424.

vllm_mlx.scheduler.Scheduler._get_detokenizer · method

vllm_mlx.scheduler.Scheduler._get_detokenizer(request_id: str) -> Any

Get or create a streaming detokenizer for a request.

Parameters

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

Returns

• Type: Any
• Direct return expressions: self._detokenizer_pool[request_id]

Exceptions and behavior

Method Scheduler._get_detokenizer calls NaiveStreamingDetokenizer; returns self._detokenizer_pool[request_id]. No direct raise statement appears in this definition.

View source #L1426-L1431.

vllm_mlx.scheduler.Scheduler._cleanup_detokenizer · method

vllm_mlx.scheduler.Scheduler._cleanup_detokenizer(request_id: str) -> None

Remove the streaming detokenizer for a finished request.

Parameters

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

Returns

• Type: None

Exceptions and behavior

Method Scheduler._cleanup_detokenizer calls self._detokenizer_pool.pop. No direct raise statement appears in this definition.

View source #L1433-L1435.

vllm_mlx.scheduler.Scheduler._get_stop_tokens · method

vllm_mlx.scheduler.Scheduler._get_stop_tokens() -> Set[int]

Get stop token IDs from tokenizer or processor.

Parameters

This callable has no explicit inputs.

Returns

• Type: Set[int]
• Direct return expressions: stop_tokens

Exceptions and behavior

Method Scheduler._get_stop_tokens calls set, hasattr, isinstance, stop_tokens.update; returns stop_tokens. No direct raise statement appears in this definition.

View source #L1437-L1455.

vllm_mlx.scheduler.Scheduler._create_batch_generator · method

vllm_mlx.scheduler.Scheduler._create_batch_generator(sampling_params: SamplingParams) -> BatchGenerator

Create a BatchGenerator with the given sampling parameters.

Parameters

Name	Type	Required	Default	Description
sampling_params	SamplingParams	yes	none	Required positional or keyword input.

Returns

• Type: BatchGenerator
• Direct return expressions: bg

Exceptions and behavior

Method Scheduler._create_batch_generator calls make_sampler, self._get_stop_tokens, stop_tokens.update, BatchGenerator; returns bg. No direct raise statement appears in this definition.

View source #L1457-L1539.

vllm_mlx.scheduler.Scheduler._create_batch_generator._prefill_progress · nested function

vllm_mlx.scheduler.Scheduler._create_batch_generator._prefill_progress(progress_list) -> not annotated

Log prefill progress for each uid chunk.

Parameters

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

Returns

• Type: not annotated

Exceptions and behavior

Nested Function Scheduler._create_batch_generator._prefill_progress calls self.uid_to_request_id.get, logger.info, isinstance. No direct raise statement appears in this definition.

View source #L1472-L1479.

vllm_mlx.scheduler.Scheduler._make_prompt_cache_save_callback · method

vllm_mlx.scheduler.Scheduler._make_prompt_cache_save_callback() -> not annotated

Create a callback that stores prompt-only KV/Mamba cache.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: _prompt_cache_save

Exceptions and behavior

Method Scheduler._make_prompt_cache_save_callback returns _prompt_cache_save. No direct raise statement appears in this definition.

View source #L1541-L1585.

vllm_mlx.scheduler.Scheduler._make_prompt_cache_save_callback._prompt_cache_save · nested function

vllm_mlx.scheduler.Scheduler._make_prompt_cache_save_callback._prompt_cache_save(uid, extracted_cache) -> not annotated

Nested Function Scheduler._make_prompt_cache_save_callback._prompt_cache_save calls self.uid_to_request_id.get, self.requests.get, list, _trim_cache_offset; returns None.

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: None

Exceptions and behavior

Nested Function Scheduler._make_prompt_cache_save_callback._prompt_cache_save calls self.uid_to_request_id.get, self.requests.get, list, _trim_cache_offset; returns None. No direct raise statement appears in this definition.

View source #L1554-L1583.

vllm_mlx.scheduler.Scheduler._make_mid_prefill_save_callback · method

vllm_mlx.scheduler.Scheduler._make_mid_prefill_save_callback(save_interval: int) -> not annotated

Create a callback for saving intermediate KV cache during chunked prefill.

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: _mid_prefill_save

Exceptions and behavior

Method Scheduler._make_mid_prefill_save_callback returns _mid_prefill_save. No direct raise statement appears in this definition.

View source #L1587-L1655.

vllm_mlx.scheduler.Scheduler._make_mid_prefill_save_callback._mid_prefill_save · nested function

vllm_mlx.scheduler.Scheduler._make_mid_prefill_save_callback._mid_prefill_save(uid, processed_tokens, prompt_cache) -> not annotated

Nested Function Scheduler._make_mid_prefill_save_callback._mid_prefill_save calls self.uid_to_request_id.get, self.requests.get, getattr, self._extract_cache_states; returns None.

Parameters

Name	Type	Required	Default	Description
uid	not annotated	yes	none	Required positional or keyword input.
processed_tokens	not annotated	yes	none	Required positional or keyword input.
prompt_cache	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: None

Exceptions and behavior

Nested Function Scheduler._make_mid_prefill_save_callback._mid_prefill_save calls self.uid_to_request_id.get, self.requests.get, getattr, self._extract_cache_states; returns None. No direct raise statement appears in this definition.

View source #L1598-L1653.

vllm_mlx.scheduler.Scheduler._close_batch_generator · method

vllm_mlx.scheduler.Scheduler._close_batch_generator() -> None

Properly close BatchGenerator to restore wired_limit.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler._close_batch_generator updates self.batch_generator; calls hasattr, self.batch_generator.close, logger.debug. No direct raise statement appears in this definition.

View source #L1657-L1665.

vllm_mlx.scheduler.Scheduler._ensure_batch_generator · method

vllm_mlx.scheduler.Scheduler._ensure_batch_generator(sampling_params: SamplingParams) -> None

Ensure BatchGenerator exists with compatible settings.

Parameters

Name	Type	Required	Default	Description
sampling_params	SamplingParams	yes	none	Required positional or keyword input.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method Scheduler._ensure_batch_generator updates self.batch_generator, self._current_sampler_params; calls logger.warning, len, hasattr, logger.info; returns None. No direct raise statement appears in this definition.

View source #L1667-L1709.

vllm_mlx.scheduler.Scheduler._validate_cache · method

vllm_mlx.scheduler.Scheduler._validate_cache(cache: Any) -> bool

Validate that a cache object is usable.

Parameters

Name	Type	Required	Default	Description
cache	Any	yes	none	The cache object to validate

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Method Scheduler._validate_cache calls isinstance, len, hasattr, logger.debug; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1711-L1769.

vllm_mlx.scheduler.Scheduler._extract_cache_states · method

vllm_mlx.scheduler.Scheduler._extract_cache_states(raw_cache: List[Any]) -> List[Dict[str, Any]]

Extract actual tensor state from each layer cache.

Parameters

Name	Type	Required	Default	Description
raw_cache	List[Any]	yes	none	List of KVCache objects from mlx-lm

Returns

• Type: List[Dict[str, Any]]
• Direct return expressions: []; extracted if len(extracted) == len(raw_cache) else []

Exceptions and behavior

Method Scheduler._extract_cache_states calls hasattr, extracted.append, type, logger.debug; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1771-L1806.

vllm_mlx.scheduler.Scheduler._reconstruct_cache_from_states · method

vllm_mlx.scheduler.Scheduler._reconstruct_cache_from_states(extracted_states: List[Dict[str, Any]]) -> Optional[List[Any]]

Reconstruct cache objects from extracted cache states.

Parameters

Name	Type	Required	Default	Description
extracted_states	List[Dict[str, Any]]	yes	none	List of dicts from _extract_cache_states()

Returns

• Type: Optional[List[Any]]
• Direct return expressions: None; caches

Exceptions and behavior

Method Scheduler._reconstruct_cache_from_states calls layer_state.get, hasattr, _KVCache, cache_cls.from_state; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1808-L1872.

vllm_mlx.scheduler.Scheduler.add_request · method

vllm_mlx.scheduler.Scheduler.add_request(request: Request) -> None

Add a new request to the scheduler.

Parameters

Name	Type	Required	Default	Description
request	Request	yes	none	The request to add

Returns

• Type: None

Exceptions and behavior

Method Scheduler.add_request calls ValueError, isinstance, hasattr, self.tokenizer.encode; can raise ValueError, AttributeError. Directly raised exceptions: ValueError, AttributeError.

View source #L1874-L1997.

vllm_mlx.scheduler.Scheduler.abort_request · method

vllm_mlx.scheduler.Scheduler.abort_request(request_id: str) -> bool

Queue request for abort.

Parameters

Name	Type	Required	Default	Description
request_id	str	yes	none	The request ID to abort

Returns

• Type: bool
• Direct return expressions: True

Exceptions and behavior

Method Scheduler.abort_request calls self._pending_abort_ids.add, logger.info; returns True. No direct raise statement appears in this definition.

View source #L1999-L2014.

vllm_mlx.scheduler.Scheduler._process_pending_aborts · method

vllm_mlx.scheduler.Scheduler._process_pending_aborts() -> None

Drain and process pending abort requests.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler._process_pending_aborts calls self._pending_abort_ids.pop, self._do_abort_request. No direct raise statement appears in this definition.

View source #L2016-L2020.

vllm_mlx.scheduler.Scheduler._do_abort_request · method

vllm_mlx.scheduler.Scheduler._do_abort_request(request_id: str) -> bool

Actually abort a request.

Parameters

Name	Type	Required	Default	Description
request_id	str	yes	none	The request ID to abort

Returns

• Type: bool
• Direct return expressions: True

Exceptions and behavior

Method Scheduler._do_abort_request updates self.total_completion_tokens, self.total_prompt_tokens; calls self.requests.get, self.waiting.remove, self.batch_generator.remove, request.set_finished; returns True. No direct raise statement appears in this definition.

View source #L2022-L2087.

vllm_mlx.scheduler.Scheduler.has_requests · method

vllm_mlx.scheduler.Scheduler.has_requests() -> bool

Check if there are any pending or running requests.

Parameters

This callable has no explicit inputs.

Returns

• Type: bool
• Direct return expressions: bool(self.waiting or self.running)

Exceptions and behavior

Method Scheduler.has_requests calls bool; returns bool(self.waiting or self.running). No direct raise statement appears in this definition.

View source #L2089-L2091.

vllm_mlx.scheduler.Scheduler.get_num_waiting · method

vllm_mlx.scheduler.Scheduler.get_num_waiting() -> int

Get number of waiting requests.

Parameters

This callable has no explicit inputs.

Returns

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

Exceptions and behavior

Method Scheduler.get_num_waiting calls len; returns len(self.waiting). No direct raise statement appears in this definition.

View source #L2093-L2095.

vllm_mlx.scheduler.Scheduler.get_num_running · method

vllm_mlx.scheduler.Scheduler.get_num_running() -> int

Get number of running requests.

Parameters

This callable has no explicit inputs.

Returns

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

Exceptions and behavior

Method Scheduler.get_num_running calls len; returns len(self.running). No direct raise statement appears in this definition.

View source #L2097-L2099.

vllm_mlx.scheduler.Scheduler._schedule_waiting · method

vllm_mlx.scheduler.Scheduler._schedule_waiting() -> List[Request]

Move requests from waiting queue to running.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[Request]
• Direct return expressions: scheduled

Exceptions and behavior

Method Scheduler._schedule_waiting updates self.total_prompt_tokens; calls self._try_promote_ssd_pending, len, self.waiting.popleft, self._ensure_batch_generator; returns scheduled. No direct raise statement appears in this definition.

View source #L2101-L2276.

vllm_mlx.scheduler.Scheduler._copy_cache_state · method

vllm_mlx.scheduler.Scheduler._copy_cache_state(value: Any) -> Any

Deep-copy a cache state payload.

Parameters

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

Returns

• Type: Any
• Direct return expressions: value + 0; type(value)(copied) if isinstance(value, tuple) else copied; value

Exceptions and behavior

Method Scheduler._copy_cache_state calls isinstance, Scheduler._copy_cache_state, type(value), type; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L2279-L2295.

vllm_mlx.scheduler.Scheduler._prompt_output_entry_is_useless · method

vllm_mlx.scheduler.Scheduler._prompt_output_entry_is_useless(cache: Any) -> bool

Would a prompt+output entry built from this cache ever be reusable?

Parameters

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

Returns

• Type: bool
• Direct return expressions: not can_trim_prompt_cache(cache); False

Exceptions and behavior

Method Scheduler._prompt_output_entry_is_useless calls can_trim_prompt_cache; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2303-L2317.

vllm_mlx.scheduler.Scheduler._extract_cache_for_uid · method

vllm_mlx.scheduler.Scheduler._extract_cache_for_uid(uid: int) -> Any

Pull one sequence's cache out of the live BatchGenerator batch.

Parameters

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

Returns

• Type: Any
• Direct return expressions: None; extract(uids.index(uid))

Exceptions and behavior

Method Scheduler._extract_cache_for_uid calls getattr, extract, uids.index, logger.debug; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2319-L2336.

vllm_mlx.scheduler.Scheduler._make_snapshot_destination · method

vllm_mlx.scheduler.Scheduler._make_snapshot_destination(live_cache: Any) -> Any

Build a destination cache with the same topology as the live one.

Parameters

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

Returns

• Type: Any
• Direct return expressions: [_mirror(layer) for layer in live_cache]; None

Exceptions and behavior

Method Scheduler._make_snapshot_destination calls _mirror, logger.warning; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2338-L2380.

vllm_mlx.scheduler.Scheduler._make_snapshot_destination._mirror · nested function

vllm_mlx.scheduler.Scheduler._make_snapshot_destination._mirror(layer: Any) -> Any

Nested Function Scheduler._make_snapshot_destination._mirror calls getattr, _mirror, copy.copy, type(children); has 2 explicit return paths.

Parameters

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

Returns

• Type: Any
• Direct return expressions: container; copy.copy(layer)

Exceptions and behavior

Nested Function Scheduler._make_snapshot_destination._mirror calls getattr, _mirror, copy.copy, type(children); has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2360-L2370.

vllm_mlx.scheduler.Scheduler._cache_coverage · method

vllm_mlx.scheduler.Scheduler._cache_coverage(cache: Any) -> int | None

How many tokens the live cache actually holds.

Parameters

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

Returns

• Type: int | None
• Direct return expressions: found; None

Exceptions and behavior

Method Scheduler._cache_coverage calls _offset_of; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L2383-L2409.

vllm_mlx.scheduler.Scheduler._cache_coverage._offset_of · nested function

vllm_mlx.scheduler.Scheduler._cache_coverage._offset_of(layer: Any) -> int | None

Nested Function Scheduler._cache_coverage._offset_of calls getattr, isinstance, _offset_of; has 3 explicit return paths.

Parameters

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

Returns

• Type: int | None
• Direct return expressions: offset; found; None

Exceptions and behavior

Nested Function Scheduler._cache_coverage._offset_of calls getattr, isinstance, _offset_of; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L2393-L2403.

vllm_mlx.scheduler.Scheduler._cache_key_for_snapshot · method

vllm_mlx.scheduler.Scheduler._cache_key_for_snapshot(request: Any, response: Any, raw_cache: Any) -> list[int] | None

Key the entry by the tokens the cache covers, not by the prompt.

Parameters

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

Returns

• Type: list[int] | None
• Direct return expressions: None; prompt_ids; prompt_ids + generated[:overshoot]

Exceptions and behavior

Method Scheduler._cache_key_for_snapshot calls self._cache_coverage, list, logger.debug, ', '.join; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L2411-L2468.

vllm_mlx.scheduler.Scheduler._store_prompt_only_cache · method

vllm_mlx.scheduler.Scheduler._store_prompt_only_cache(request: Any, response: Any) -> None

Store the post-prefill cache under the prompt tokens alone.

Parameters

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

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method Scheduler._store_prompt_only_cache calls getattr, len, callable, raw_cache; returns None. No direct raise statement appears in this definition.

View source #L2470-L2581.

vllm_mlx.scheduler.Scheduler._process_batch_responses · method

vllm_mlx.scheduler.Scheduler._process_batch_responses(responses: List[Any]) -> Tuple[List[RequestOutput], Set[str]]

Process responses from BatchGenerator.

Parameters

Name	Type	Required	Default	Description
responses	List[Any]	yes	none	List of BatchGenerator.Response objects

Returns

• Type: Tuple[List[RequestOutput], Set[str]]
• Direct return expressions: (outputs, finished_ids)

Exceptions and behavior

Method Scheduler._process_batch_responses updates self.total_completion_tokens, self.num_requests_processed; calls set, self.uid_to_request_id.get, self.running.get, self._store_prompt_only_cache; returns (outputs, finished_ids). No direct raise statement appears in this definition.

View source #L2583-L2710.

vllm_mlx.scheduler.Scheduler._cleanup_finished · method

vllm_mlx.scheduler.Scheduler._cleanup_finished(finished_ids: Set[str]) -> None

Clean up finished requests and store caches for reuse.

Parameters

Name	Type	Required	Default	Description
finished_ids	Set[str]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method Scheduler._cleanup_finished calls self.running.get, hasattr, list, self.block_aware_cache.store_cache. No direct raise statement appears in this definition.

View source #L2712-L2865.

vllm_mlx.scheduler.Scheduler._is_cache_corruption_error · method

vllm_mlx.scheduler.Scheduler._is_cache_corruption_error(error: Exception) -> bool

Check if an error indicates cache corruption.

Parameters

Name	Type	Required	Default	Description
error	Exception	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: any((pattern in error_str for pattern in CACHE_CORRUPTION_PATTERNS))

Exceptions and behavior

Method Scheduler._is_cache_corruption_error calls str, any; returns any((pattern in error_str for pattern in CACHE_CORRUPTION_PATTERNS)). No direct raise statement appears in this definition.

View source #L2867-L2870.

vllm_mlx.scheduler.Scheduler._is_stream_thread_error · method

vllm_mlx.scheduler.Scheduler._is_stream_thread_error(error: Exception) -> bool

Check if an error indicates MLX stream/thread ownership mismatch.

Parameters

Name	Type	Required	Default	Description
error	Exception	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: 'no Stream(' in error_str or 'no Stream(gpu' in error_str

Exceptions and behavior

Method Scheduler._is_stream_thread_error calls str; returns 'no Stream(' in error_str or 'no Stream(gpu' in error_str. No direct raise statement appears in this definition.

View source #L2872-L2875.

vllm_mlx.scheduler.Scheduler._recover_from_cache_error · method

vllm_mlx.scheduler.Scheduler._recover_from_cache_error() -> None

Recover from cache corruption error.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler._recover_from_cache_error updates self._current_sampler_params; calls self._close_batch_generator, self.block_aware_cache.clear, self.memory_aware_cache.clear, self.prefix_cache.clear. No direct raise statement appears in this definition.

View source #L2877-L2895.

vllm_mlx.scheduler.Scheduler._recover_from_generation_error · method

vllm_mlx.scheduler.Scheduler._recover_from_generation_error() -> Set[str]

Recover from fatal generation error (OOM, Metal crash).

Parameters

This callable has no explicit inputs.

Returns

• Type: Set[str]
• Direct return expressions: aborted_ids

Exceptions and behavior

Method Scheduler._recover_from_generation_error updates self._current_sampler_params; calls self._close_batch_generator, set, list, self.running.get; returns aborted_ids. No direct raise statement appears in this definition.

View source #L2897-L2933.

vllm_mlx.scheduler.Scheduler._reschedule_running_requests · method

vllm_mlx.scheduler.Scheduler._reschedule_running_requests() -> None

Move running requests back to waiting queue for retry.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler._reschedule_running_requests calls len, list, self.running.items, self.waiting.appendleft. No direct raise statement appears in this definition.

View source #L2935-L2951.

vllm_mlx.scheduler.Scheduler.step · method

vllm_mlx.scheduler.Scheduler.step(max_retries: int = 1) -> SchedulerOutput

Execute one scheduling step with automatic error recovery.

Parameters

Name	Type	Required	Default	Description
max_retries	int	no	1	Number of times to retry on cache errors (default 1)

Returns

• Type: SchedulerOutput
• Direct return expressions: output

Exceptions and behavior

Method Scheduler.step updates self.finished_req_ids, self._step_count; calls SchedulerOutput, self._process_pending_aborts, range, self._schedule_waiting; returns output. No direct raise statement appears in this definition.

View source #L2953-L3089.

vllm_mlx.scheduler.Scheduler.get_request · method

vllm_mlx.scheduler.Scheduler.get_request(request_id: str) -> Optional[Request]

Get a request by ID.

Parameters

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

Returns

• Type: Optional[Request]
• Direct return expressions: self.requests.get(request_id)

Exceptions and behavior

Method Scheduler.get_request calls self.requests.get; returns self.requests.get(request_id). No direct raise statement appears in this definition.

View source #L3091-L3093.

vllm_mlx.scheduler.Scheduler.remove_finished_request · method

vllm_mlx.scheduler.Scheduler.remove_finished_request(request_id: str) -> Optional[Request]

Remove a finished request from tracking.

Parameters

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

Returns

• Type: Optional[Request]
• Direct return expressions: self.requests.pop(request_id, None)

Exceptions and behavior

Method Scheduler.remove_finished_request calls self.requests.pop; returns self.requests.pop(request_id, None). No direct raise statement appears in this definition.

View source #L3095-L3097.

vllm_mlx.scheduler.Scheduler.get_running_requests_info · method

vllm_mlx.scheduler.Scheduler.get_running_requests_info() -> List[Dict[str, Any]]

Per-request details for status endpoint.

Parameters

This callable has no explicit inputs.

Returns

• Type: List[Dict[str, Any]]
• Direct return expressions: result

Exceptions and behavior

Method Scheduler.get_running_requests_info calls _time.time, result.append, round, self.running.values; returns result. No direct raise statement appears in this definition.

View source #L3099-L3165.

vllm_mlx.scheduler.Scheduler.get_stats · method

vllm_mlx.scheduler.Scheduler.get_stats() -> Dict[str, Any]

Get scheduler statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: Dict[str, Any]
• Direct return expressions: stats

Exceptions and behavior

Method Scheduler.get_stats calls len, stats.update, _mtp_status_snapshot, mx.metal.is_available; returns stats. No direct raise statement appears in this definition.

View source #L3167-L3193.

vllm_mlx.scheduler.Scheduler.get_cache_stats · method

vllm_mlx.scheduler.Scheduler.get_cache_stats() -> Optional[Dict[str, Any]]

Get cache statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: Optional[Dict[str, Any]]
• Direct return expressions: self.block_aware_cache.get_stats(); self.memory_aware_cache.get_stats(); self.prefix_cache.get_stats(); None

Exceptions and behavior

Method Scheduler.get_cache_stats calls self.block_aware_cache.get_stats, self.memory_aware_cache.get_stats, self.prefix_cache.get_stats; has 4 explicit return paths. No direct raise statement appears in this definition.

View source #L3195-L3203.

vllm_mlx.scheduler.Scheduler.clear_runtime_caches · method

vllm_mlx.scheduler.Scheduler.clear_runtime_caches() -> Dict[str, bool]

Clear prefix-cache state without resetting scheduler/request state.

Parameters

This callable has no explicit inputs.

Returns

• Type: Dict[str, bool]
• Direct return expressions: cleared

Exceptions and behavior

Method Scheduler.clear_runtime_caches calls self.block_aware_cache.clear, self.memory_aware_cache.clear, self.prefix_cache.clear; returns cleared. No direct raise statement appears in this definition.

View source #L3205-L3221.

vllm_mlx.scheduler.Scheduler.reset · method

vllm_mlx.scheduler.Scheduler.reset() -> None

Reset the scheduler state.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler.reset updates self._current_sampler_params; calls self._pending_abort_ids.clear, list, self.requests.keys, self._do_abort_request. No direct raise statement appears in this definition.

View source #L3223-L3246.

vllm_mlx.scheduler.Scheduler.deep_reset · method

vllm_mlx.scheduler.Scheduler.deep_reset() -> None

Deep reset that clears ALL cache state including model-level caches.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler.deep_reset updates self.model.cache; calls self.reset, hasattr, gc.collect, logger.info. No direct raise statement appears in this definition.

View source #L3248-L3276.

vllm_mlx.scheduler.Scheduler.save_cache_to_disk · method

vllm_mlx.scheduler.Scheduler.save_cache_to_disk(cache_dir: str) -> bool

Save prefix cache to disk for persistence across restarts.

Parameters

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

Returns

• Type: bool
• Direct return expressions: self.memory_aware_cache.save_to_disk(cache_dir); False

Exceptions and behavior

Method Scheduler.save_cache_to_disk calls self.memory_aware_cache.save_to_disk, logger.info; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L3282-L3287.

vllm_mlx.scheduler.Scheduler.load_cache_from_disk · method

vllm_mlx.scheduler.Scheduler.load_cache_from_disk(cache_dir: str) -> int

Load prefix cache from disk.

Parameters

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

Returns

• Type: int
• Direct return expressions: self.memory_aware_cache.load_from_disk(cache_dir); 0

Exceptions and behavior

Method Scheduler.load_cache_from_disk calls self.memory_aware_cache.load_from_disk, logger.info; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L3289-L3294.

vllm_mlx.scheduler.Scheduler.clear_prefix_cache · method

vllm_mlx.scheduler.Scheduler.clear_prefix_cache() -> None

Clear the in-memory prefix cache (keeps disk cache untouched).

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method Scheduler.clear_prefix_cache calls hasattr, self.memory_aware_cache.clear, logger.info, self.prefix_cache.clear; returns None. No direct raise statement appears in this definition.

View source #L3296-L3306.

vllm_mlx.scheduler.Scheduler.close_ssd_tier · method

vllm_mlx.scheduler.Scheduler.close_ssd_tier() -> None

Shut down the SSD cache tier if present.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler.close_ssd_tier updates self._ssd_tier; calls self._ssd_tier.close, logger.info. No direct raise statement appears in this definition.

View source #L3308-L3313.

vllm_mlx.scheduler.Scheduler._try_promote_ssd_pending · method

vllm_mlx.scheduler.Scheduler._try_promote_ssd_pending() -> None

Attempt synchronous SSD promotion for waiting requests tagged ssd_pending.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method Scheduler._try_promote_ssd_pending updates self._ssd_tier._stats.promotion_failures, self._ssd_tier._stats.ssd_hits; calls getattr, self.memory_aware_cache.try_reserve_memory, logger.info, tuple. No direct raise statement appears in this definition.

View source #L3315-L3395.

vllm_mlx.scheduler.Scheduler.promote_from_ssd · method

async vllm_mlx.scheduler.Scheduler.promote_from_ssd(request) -> bool

Promote a cold-tier cache entry for a request (async version).

Parameters

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

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Method Scheduler.promote_from_ssd calls getattr, candidate.get, len, tuple; awaits asynchronous work; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L3397-L3460.

vllm_mlx.scheduler.Scheduler.promote_from_ssd.reserve_budget · nested function

vllm_mlx.scheduler.Scheduler.promote_from_ssd.reserve_budget(nbytes: int) -> bool

Tentatively reserve RAM budget for promotion.

Parameters

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

Returns

• Type: bool
• Direct return expressions: False; self.memory_aware_cache.try_reserve_memory(nbytes)

Exceptions and behavior

Nested Function Scheduler.promote_from_ssd.reserve_budget calls self.memory_aware_cache.try_reserve_memory; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L3412-L3416.

vllm_mlx.scheduler.Scheduler.promote_from_ssd.release_budget · nested function

vllm_mlx.scheduler.Scheduler.promote_from_ssd.release_budget(nbytes: int) -> None

Release tentatively reserved budget on failure.

Parameters

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

Returns

• Type: None

Exceptions and behavior

Nested Function Scheduler.promote_from_ssd.release_budget calls self.memory_aware_cache.release_reserved_memory. No direct raise statement appears in this definition.

View source #L3418-L3421.

vllm_mlx.scheduler.Scheduler._reconstruct_ssd_layers · method

vllm_mlx.scheduler.Scheduler._reconstruct_ssd_layers(layer_dicts: list[dict]) -> list | None

Reconstruct cache objects from deserialized layer dicts.

Parameters

Name	Type	Required	Default	Description
layer_dicts	list[dict]	yes	none	Required positional or keyword input.

Returns

• Type: list | None
• Direct return expressions: None; result

Exceptions and behavior

Method Scheduler._reconstruct_ssd_layers calls KVCache, mx.array, ld.get, _mx_dtype_from_name; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L3462-L3518.

vllm_mlx.scheduler.Scheduler._reconstruct_ssd_layers._mx_dtype_from_name · nested function

vllm_mlx.scheduler.Scheduler._reconstruct_ssd_layers._mx_dtype_from_name(name: str) -> not annotated

Nested Function Scheduler._reconstruct_ssd_layers._mx_dtype_from_name calls getattr; returns getattr(mx, name, None).

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: getattr(mx, name, None)

Exceptions and behavior

Nested Function Scheduler._reconstruct_ssd_layers._mx_dtype_from_name calls getattr; returns getattr(mx, name, None). No direct raise statement appears in this definition.

View source #L3473-L3474.

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
_normalize_logits_processors	function	_normalize_logits_processors(logits_processors) -> not annotated	Normalize empty per-sequence processor slots to lists.	#L46-L50
_sanitize_batch_generator_logits_processors	function	_sanitize_batch_generator_logits_processors(batch_generator) -> None	Sanitize stale BatchGenerator processor state before decode.	#L53-L65
SchedulingPolicy	class	SchedulingPolicy()	Scheduling policy for request ordering.	#L68-L72
SchedulerConfig	class	SchedulerConfig(max_num_seqs: int = 256, max_num_batched_tokens: int = 8192, policy: SchedulingPolicy = SchedulingPolicy.FCFS, prefill_batch_size: int = 8, completion_batch_size: int = 32, prefill_step_size: int = 2048, mllm_prefill_step_size: Optional[int] = None, enable_prefix_cache: bool = True, prefix_cache_size: int = 100, use_memory_aware_cache: bool = True, cache_memory_mb: Optional[int] = None, cache_memory_percent: float = 0.2, kv_cache_quantization: bool = False, kv_cache_quantization_bits: int = 8, kv_cache_quantization_group_size: int = 64, kv_cache_min_quantize_tokens: int = 256, use_paged_cache: bool = False, paged_cache_block_size: int = 64, max_cache_blocks: int = 1000, chunked_prefill_tokens: int = 0, mid_prefill_save_interval: int = 8192, ssd_cache_dir: Optional[str] = None, ssd_cache_max_gb: float = 10.0, max_kv_size: int = 0, enable_mtp: bool = False, mtp_num_draft_tokens: int = 1, mtp_optimistic: bool = False)	Configuration for the scheduler.	#L76-L140
SchedulerConfig.__post_init__	method	SchedulerConfig.__post_init__() -> None	Method SchedulerConfig.__post_init__ calls ValueError; can raise ValueError.	#L138-L140
SchedulerOutput	class	SchedulerOutput(scheduled_request_ids: List[str] = field(default_factory=list), num_scheduled_tokens: int = 0, finished_request_ids: Set[str] = field(default_factory=set), outputs: List[RequestOutput] = field(default_factory=list), has_work: bool = False)	Output from a scheduling step.	#L144-L160
_install_prompt_cache_save	function	_install_prompt_cache_save(batch_gen: 'BatchGenerator', prompt_cache_save) -> None	Monkey-patch _process_prompts to capture prompt-only cache state.	#L163-L187
_install_prompt_cache_save._patched_process_prompts	nested function	_install_prompt_cache_save._patched_process_prompts(prompts, _self = batch_gen) -> not annotated	Nested Function _install_prompt_cache_save._patched_process_prompts calls _orig_process_prompts, enumerate, prompt_cache_save, batch.extract_cache; returns batch.	#L177-L185
_install_chunked_prefill	function	_install_chunked_prefill(batch_gen: 'BatchGenerator', budget: int, mid_prefill_save = None, prompt_cache_save = None, pending_abort_ids: Optional[Set[str]] = None, uid_to_request_id: Optional[Dict[int, str]] = None, requests: Optional[Dict[str, Any]] = None) -> None	Monkey-patch a BatchGenerator instance so that large prefills are broken into chunks of at most budget tokens each.	#L190-L697
_install_chunked_prefill._lazy_extract_cache	nested function	_install_chunked_prefill._lazy_extract_cache(cache, idx) -> not annotated	Nested Function _install_chunked_prefill._lazy_extract_cache calls c.extract; returns (c.extract(idx) for c in cache).	#L225-L226
_install_chunked_prefill._batch_cls	nested class	_install_chunked_prefill._batch_cls(uids: List[int], y: Any, logprobs: List[Any], max_tokens: List[int], num_tokens: List[int], cache: List[Any], samplers: List[Any], logits_processors: List[Any], tokens: List[Any])	Nested Class _install_chunked_prefill._batch_cls declares 4 direct member(s).	#L233-L273
_install_chunked_prefill._batch_cls.__len__	nested function	_install_chunked_prefill._batch_cls.__len__() -> not annotated	Nested Function _install_chunked_prefill._batch_cls.__len__ calls len; returns len(self.uids).	#L244-L245
_install_chunked_prefill._batch_cls.filter	nested function	_install_chunked_prefill._batch_cls.filter(keep_idx: List[int]) -> not annotated	Nested Function _install_chunked_prefill._batch_cls.filter updates self.uids, self.logprobs, self.max_tokens, self.num_tokens; calls mx.array, c.filter.	#L247-L258
_install_chunked_prefill._batch_cls.extend	nested function	_install_chunked_prefill._batch_cls.extend(other) -> not annotated	Nested Function _install_chunked_prefill._batch_cls.extend updates self.y; calls self.uids.extend, mx.concatenate, self.logprobs.extend, self.num_tokens.extend.	#L260-L270
_install_chunked_prefill._batch_cls.extract_cache	nested function	_install_chunked_prefill._batch_cls.extract_cache(idx) -> not annotated	Nested Function _install_chunked_prefill._batch_cls.extract_cache calls c.extract; returns [c.extract(idx) for c in self.cache].	#L272-L273
_install_chunked_prefill._patched_process_prompts	nested function	_install_chunked_prefill._patched_process_prompts(prompts, _self = batch_gen) -> not annotated	Nested Function _install_chunked_prefill._patched_process_prompts calls _orig_process_prompts, enumerate, prompt_cache_save, batch.extract_cache; returns batch.	#L291-L299
_install_chunked_prefill._generation_step	nested function	_install_chunked_prefill._generation_step() -> not annotated	Run one generation step on the active batch.	#L303-L360
_install_chunked_prefill._chunked_next	nested function	_install_chunked_prefill._chunked_next() -> not annotated	Replacement for _next() that chunks large prefills.	#L362-L678
_install_chunked_prefill._patched_remove	nested function	_install_chunked_prefill._patched_remove(uids_to_remove, _self = batch_gen) -> not annotated	Clear partial state if aborted request is being prefilled.	#L680-L691
_MTPStatsState	class	_MTPStatsState(counters: Dict[str, int] = field(default_factory=lambda: {'attempted': 0, 'accepted': 0, 'rejected': 0, 'errors': 0}), bypass_counts: Dict[str, int] = field(default_factory=lambda: {'prefill': 0, 'no_active_batch': 0, 'cache_mismatch': 0}), lock: Any = field(default_factory=Lock))	Cumulative native-MTP counters shared across generator instances.	#L701-L719
_configure_chunked_prefill	function	_configure_chunked_prefill(scheduler: 'Scheduler', batch_gen: 'BatchGenerator', budget: int, prompt_cache_save) -> None	Enable the matching legacy or native mlx-lm chunked-prefill API.	#L722-L777
_install_mtp	function	_install_mtp(batch_gen: 'BatchGenerator', model: Any, num_draft_tokens: int = 1, optimistic: bool = False, stats_state: Optional['_MTPStatsState'] = None) -> None	Monkey-patch a BatchGenerator to use MTP (Multi-Token Prediction) with always-advance strategy for hybrid MambaCache + KVCache.	#L780-L1262
_install_mtp._get_mtp_stats	nested function	_install_mtp._get_mtp_stats() -> Dict[str, Any]	Nested Function _install_mtp._get_mtp_stats calls dict; returns {'enabled': True, 'requested_draft_tokens': num_draft_tokens, 'effective_draft_tokens': 1, 'mode': 'always_advance_opti….	#L823-L845
_install_mtp._mtp_bypass_reasons	nested function	_install_mtp._mtp_bypass_reasons(input_tokens, prompt_cache) -> not annotated	Nested Function _install_mtp._mtp_bypass_reasons calls reasons.append; returns reasons.	#L849-L857
_install_mtp._record_mtp_bypass	nested function	_install_mtp._record_mtp_bypass(reasons) -> None	Nested Function _install_mtp._record_mtp_bypass contains no state mutation, call, raise, return, await, or yield.	#L859-L862
_install_mtp._mtp_step	nested function	_install_mtp._mtp_step(input_tokens, prompt_cache, samplers, logits_processors, tokens) -> not annotated	Extended _step with MTP always-advance strategy.	#L864-L1138
_install_mtp._mtp_next	nested function	_install_mtp._mtp_next() -> not annotated	Wrapper around _next that emits deferred MTP draft tokens.	#L1147-L1247
_mtp_status_snapshot	function	_mtp_status_snapshot(batch_generator) -> Dict[str, Any]	Function _mtp_status_snapshot calls getattr, callable, get_mtp_stats; has 2 explicit return paths.	#L1265-L1269
Scheduler	class	Scheduler(model: Any, tokenizer: Any, config: Optional[SchedulerConfig] = None)	Scheduler for continuous batching using mlx-lm BatchGenerator.	#L1272-L3518
Scheduler.__init__	method	Scheduler.__init__(model: Any, tokenizer: Any, config: Optional[SchedulerConfig] = None) -> not annotated	Initialize the scheduler.	#L1286-L1402
Scheduler._get_actual_tokenizer	method	Scheduler._get_actual_tokenizer(tokenizer: Any) -> Any	Get the actual tokenizer from a processor or tokenizer.	#L1404-L1418
Scheduler._decode_tokens	method	Scheduler._decode_tokens(token_ids: List[int]) -> str	Decode token IDs to text, handling both tokenizers and processors.	#L1420-L1424
Scheduler._get_detokenizer	method	Scheduler._get_detokenizer(request_id: str) -> Any	Get or create a streaming detokenizer for a request.	#L1426-L1431
Scheduler._cleanup_detokenizer	method	Scheduler._cleanup_detokenizer(request_id: str) -> None	Remove the streaming detokenizer for a finished request.	#L1433-L1435
Scheduler._get_stop_tokens	method	Scheduler._get_stop_tokens() -> Set[int]	Get stop token IDs from tokenizer or processor.	#L1437-L1455
Scheduler._create_batch_generator	method	Scheduler._create_batch_generator(sampling_params: SamplingParams) -> BatchGenerator	Create a BatchGenerator with the given sampling parameters.	#L1457-L1539
Scheduler._create_batch_generator._prefill_progress	nested function	Scheduler._create_batch_generator._prefill_progress(progress_list) -> not annotated	Log prefill progress for each uid chunk.	#L1472-L1479
Scheduler._make_prompt_cache_save_callback	method	Scheduler._make_prompt_cache_save_callback() -> not annotated	Create a callback that stores prompt-only KV/Mamba cache.	#L1541-L1585
Scheduler._make_prompt_cache_save_callback._prompt_cache_save	nested function	Scheduler._make_prompt_cache_save_callback._prompt_cache_save(uid, extracted_cache) -> not annotated	Nested Function Scheduler._make_prompt_cache_save_callback._prompt_cache_save calls self.uid_to_request_id.get, self.requests.get, list, _trim_cache_offset; returns None.	#L1554-L1583
Scheduler._make_mid_prefill_save_callback	method	Scheduler._make_mid_prefill_save_callback(save_interval: int) -> not annotated	Create a callback for saving intermediate KV cache during chunked prefill.	#L1587-L1655
Scheduler._make_mid_prefill_save_callback._mid_prefill_save	nested function	Scheduler._make_mid_prefill_save_callback._mid_prefill_save(uid, processed_tokens, prompt_cache) -> not annotated	Nested Function Scheduler._make_mid_prefill_save_callback._mid_prefill_save calls self.uid_to_request_id.get, self.requests.get, getattr, self._extract_cache_states; returns None.	#L1598-L1653
Scheduler._close_batch_generator	method	Scheduler._close_batch_generator() -> None	Properly close BatchGenerator to restore wired_limit.	#L1657-L1665
Scheduler._ensure_batch_generator	method	Scheduler._ensure_batch_generator(sampling_params: SamplingParams) -> None	Ensure BatchGenerator exists with compatible settings.	#L1667-L1709
Scheduler._validate_cache	method	Scheduler._validate_cache(cache: Any) -> bool	Validate that a cache object is usable.	#L1711-L1769
Scheduler._extract_cache_states	method	Scheduler._extract_cache_states(raw_cache: List[Any]) -> List[Dict[str, Any]]	Extract actual tensor state from each layer cache.	#L1771-L1806
Scheduler._reconstruct_cache_from_states	method	Scheduler._reconstruct_cache_from_states(extracted_states: List[Dict[str, Any]]) -> Optional[List[Any]]	Reconstruct cache objects from extracted cache states.	#L1808-L1872
Scheduler.add_request	method	Scheduler.add_request(request: Request) -> None	Add a new request to the scheduler.	#L1874-L1997
Scheduler.abort_request	method	Scheduler.abort_request(request_id: str) -> bool	Queue request for abort.	#L1999-L2014
Scheduler._process_pending_aborts	method	Scheduler._process_pending_aborts() -> None	Drain and process pending abort requests.	#L2016-L2020
Scheduler._do_abort_request	method	Scheduler._do_abort_request(request_id: str) -> bool	Actually abort a request.	#L2022-L2087
Scheduler.has_requests	method	Scheduler.has_requests() -> bool	Check if there are any pending or running requests.	#L2089-L2091
Scheduler.get_num_waiting	method	Scheduler.get_num_waiting() -> int	Get number of waiting requests.	#L2093-L2095
Scheduler.get_num_running	method	Scheduler.get_num_running() -> int	Get number of running requests.	#L2097-L2099
Scheduler._schedule_waiting	method	Scheduler._schedule_waiting() -> List[Request]	Move requests from waiting queue to running.	#L2101-L2276
Scheduler._copy_cache_state	method	Scheduler._copy_cache_state(value: Any) -> Any	Deep-copy a cache state payload.	#L2279-L2295
Scheduler._prompt_output_entry_is_useless	method	Scheduler._prompt_output_entry_is_useless(cache: Any) -> bool	Would a prompt+output entry built from this cache ever be reusable?	#L2303-L2317
Scheduler._extract_cache_for_uid	method	Scheduler._extract_cache_for_uid(uid: int) -> Any	Pull one sequence's cache out of the live BatchGenerator batch.	#L2319-L2336
Scheduler._make_snapshot_destination	method	Scheduler._make_snapshot_destination(live_cache: Any) -> Any	Build a destination cache with the same topology as the live one.	#L2338-L2380
Scheduler._make_snapshot_destination._mirror	nested function	Scheduler._make_snapshot_destination._mirror(layer: Any) -> Any	Nested Function Scheduler._make_snapshot_destination._mirror calls getattr, _mirror, copy.copy, type(children); has 2 explicit return paths.	#L2360-L2370
Scheduler._cache_coverage	method	Scheduler._cache_coverage(cache: Any) -> int \| None	How many tokens the live cache actually holds.	#L2383-L2409
Scheduler._cache_coverage._offset_of	nested function	Scheduler._cache_coverage._offset_of(layer: Any) -> int \| None	Nested Function Scheduler._cache_coverage._offset_of calls getattr, isinstance, _offset_of; has 3 explicit return paths.	#L2393-L2403
Scheduler._cache_key_for_snapshot	method	Scheduler._cache_key_for_snapshot(request: Any, response: Any, raw_cache: Any) -> list[int] \| None	Key the entry by the tokens the cache covers, not by the prompt.	#L2411-L2468
Scheduler._store_prompt_only_cache	method	Scheduler._store_prompt_only_cache(request: Any, response: Any) -> None	Store the post-prefill cache under the prompt tokens alone.	#L2470-L2581
Scheduler._process_batch_responses	method	Scheduler._process_batch_responses(responses: List[Any]) -> Tuple[List[RequestOutput], Set[str]]	Process responses from BatchGenerator.	#L2583-L2710
Scheduler._cleanup_finished	method	Scheduler._cleanup_finished(finished_ids: Set[str]) -> None	Clean up finished requests and store caches for reuse.	#L2712-L2865
Scheduler._is_cache_corruption_error	method	Scheduler._is_cache_corruption_error(error: Exception) -> bool	Check if an error indicates cache corruption.	#L2867-L2870
Scheduler._is_stream_thread_error	method	Scheduler._is_stream_thread_error(error: Exception) -> bool	Check if an error indicates MLX stream/thread ownership mismatch.	#L2872-L2875
Scheduler._recover_from_cache_error	method	Scheduler._recover_from_cache_error() -> None	Recover from cache corruption error.	#L2877-L2895
Scheduler._recover_from_generation_error	method	Scheduler._recover_from_generation_error() -> Set[str]	Recover from fatal generation error (OOM, Metal crash).	#L2897-L2933
Scheduler._reschedule_running_requests	method	Scheduler._reschedule_running_requests() -> None	Move running requests back to waiting queue for retry.	#L2935-L2951
Scheduler.step	method	Scheduler.step(max_retries: int = 1) -> SchedulerOutput	Execute one scheduling step with automatic error recovery.	#L2953-L3089
Scheduler.get_request	method	Scheduler.get_request(request_id: str) -> Optional[Request]	Get a request by ID.	#L3091-L3093
Scheduler.remove_finished_request	method	Scheduler.remove_finished_request(request_id: str) -> Optional[Request]	Remove a finished request from tracking.	#L3095-L3097
Scheduler.get_running_requests_info	method	Scheduler.get_running_requests_info() -> List[Dict[str, Any]]	Per-request details for status endpoint.	#L3099-L3165
Scheduler.get_stats	method	Scheduler.get_stats() -> Dict[str, Any]	Get scheduler statistics.	#L3167-L3193
Scheduler.get_cache_stats	method	Scheduler.get_cache_stats() -> Optional[Dict[str, Any]]	Get cache statistics.	#L3195-L3203
Scheduler.clear_runtime_caches	method	Scheduler.clear_runtime_caches() -> Dict[str, bool]	Clear prefix-cache state without resetting scheduler/request state.	#L3205-L3221
Scheduler.reset	method	Scheduler.reset() -> None	Reset the scheduler state.	#L3223-L3246
Scheduler.deep_reset	method	Scheduler.deep_reset() -> None	Deep reset that clears ALL cache state including model-level caches.	#L3248-L3276
Scheduler.save_cache_to_disk	method	Scheduler.save_cache_to_disk(cache_dir: str) -> bool	Save prefix cache to disk for persistence across restarts.	#L3282-L3287
Scheduler.load_cache_from_disk	method	Scheduler.load_cache_from_disk(cache_dir: str) -> int	Load prefix cache from disk.	#L3289-L3294
Scheduler.clear_prefix_cache	method	Scheduler.clear_prefix_cache() -> None	Clear the in-memory prefix cache (keeps disk cache untouched).	#L3296-L3306
Scheduler.close_ssd_tier	method	Scheduler.close_ssd_tier() -> None	Shut down the SSD cache tier if present.	#L3308-L3313
Scheduler._try_promote_ssd_pending	method	Scheduler._try_promote_ssd_pending() -> None	Attempt synchronous SSD promotion for waiting requests tagged ssd_pending.	#L3315-L3395
Scheduler.promote_from_ssd	method	async Scheduler.promote_from_ssd(request) -> bool	Promote a cold-tier cache entry for a request (async version).	#L3397-L3460
Scheduler.promote_from_ssd.reserve_budget	nested function	Scheduler.promote_from_ssd.reserve_budget(nbytes: int) -> bool	Tentatively reserve RAM budget for promotion.	#L3412-L3416
Scheduler.promote_from_ssd.release_budget	nested function	Scheduler.promote_from_ssd.release_budget(nbytes: int) -> None	Release tentatively reserved budget on failure.	#L3418-L3421
Scheduler._reconstruct_ssd_layers	method	Scheduler._reconstruct_ssd_layers(layer_dicts: list[dict]) -> list \| None	Reconstruct cache objects from deserialized layer dicts.	#L3462-L3518
Scheduler._reconstruct_ssd_layers._mx_dtype_from_name	nested function	Scheduler._reconstruct_ssd_layers._mx_dtype_from_name(name: str) -> not annotated	Nested Function Scheduler._reconstruct_ssd_layers._mx_dtype_from_name calls getattr; returns getattr(mx, name, None).	#L3473-L3474