wehub-resource-sync
10686 lines
457 KiB
Python
10686 lines
457 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""
|
||
Scheduler for oMLX 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
|
||
"""
|
||
|
||
import concurrent.futures
|
||
import copy
|
||
import gc
|
||
import importlib
|
||
import logging
|
||
import os
|
||
import threading
|
||
import time
|
||
from collections import OrderedDict, defaultdict, deque
|
||
from collections.abc import Callable
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass, field
|
||
from enum import Enum
|
||
from pathlib import Path
|
||
from typing import Any, NamedTuple, Optional
|
||
|
||
import mlx.core as mx
|
||
from mlx_lm.generate import (
|
||
BatchGenerator,
|
||
GenerationBatch,
|
||
PromptProcessingBatch,
|
||
SequenceStateMachine,
|
||
generation_stream,
|
||
)
|
||
from mlx_lm.models.cache import (
|
||
KVCache as _MLXKVCache,
|
||
)
|
||
from mlx_lm.models.cache import (
|
||
RotatingKVCache as _MLXRotatingKVCache,
|
||
)
|
||
from mlx_lm.models.cache import (
|
||
make_prompt_cache,
|
||
)
|
||
from mlx_lm.sample_utils import make_logits_processors
|
||
|
||
from .cache.observability import CacheRateTracker
|
||
from .cache.paged_cache import PagedCacheManager
|
||
from .cache.prefix_cache import BlockAwarePrefixCache
|
||
from .exceptions import PrefillMemoryExceededError, is_cache_corruption_error
|
||
from .patches.sdpa256_attention import set_unfused_headroom_provider
|
||
from .prefill_progress import get_prefill_tracker
|
||
from .prefill_transient_tracker import PrefillTransientTracker
|
||
from .request import Request, RequestOutput, RequestStatus, SamplingParams
|
||
from .speculative.vlm_mtp import VLMMTPDrafter, run_vlm_mtp_decode
|
||
from .utils.fatal import FATAL_TEARDOWN_TIMEOUT_S, fatal_exit
|
||
from .utils.generation_config import load_generation_config_token_ids
|
||
from .utils.proc_memory import get_phys_footprint
|
||
from .utils.sampling import make_sampler as omlx_make_sampler
|
||
from .utils.tokenizer import create_streaming_detokenizer
|
||
|
||
# Module-level alias so Scheduler.__init__ can fall back to mlx-lm's default
|
||
# stream when no per-engine stream is provided.
|
||
_default_generation_stream = generation_stream
|
||
|
||
|
||
def _apply_suppress_token_ids(logits: Any, suppress_token_ids: tuple[int, ...]) -> Any:
|
||
if suppress_token_ids:
|
||
logits[..., list(suppress_token_ids)] = mx.array(float("-inf"))
|
||
return logits
|
||
|
||
|
||
def _make_suppress_logits_processor(
|
||
suppress_token_ids: set[int],
|
||
) -> Callable[[Any, Any], Any] | None:
|
||
suppress_tuple = tuple(sorted(int(t) for t in suppress_token_ids))
|
||
if not suppress_tuple:
|
||
return None
|
||
|
||
def _suppress_logits(tokens: Any, logits: Any) -> Any:
|
||
return _apply_suppress_token_ids(logits, suppress_tuple)
|
||
|
||
return _suppress_logits
|
||
|
||
|
||
def _make_suppressing_sampler(
|
||
sampler: Callable[[Any], Any],
|
||
suppress_token_ids: set[int],
|
||
) -> Callable[[Any], Any]:
|
||
suppress_tuple = tuple(sorted(int(t) for t in suppress_token_ids))
|
||
if not suppress_tuple:
|
||
return sampler
|
||
|
||
def _sample(logits: Any) -> Any:
|
||
return sampler(_apply_suppress_token_ids(logits, suppress_tuple))
|
||
|
||
return _sample
|
||
|
||
|
||
@dataclass
|
||
class _PreflightRejection:
|
||
"""Typed return for ``_preflight_memory_check`` / its token-count
|
||
helper. Carries the human-readable diagnostic plus the numeric
|
||
estimated / limit bytes so callers can populate
|
||
``PrefillMemoryExceededError`` without parsing the string.
|
||
|
||
Same shim as ``preflight_or_raise`` (restored on main after the
|
||
upstream merge dropped it); the typed shape is what
|
||
``tests/test_scheduler_prefill_memory_guard.py`` asserts against,
|
||
and what PR #1452 carries upstream.
|
||
"""
|
||
|
||
message: str
|
||
estimated_bytes: int
|
||
limit_bytes: int
|
||
|
||
|
||
@dataclass
|
||
class _VLMMTPDecodeState:
|
||
"""Per-request state for vlm_mtp decode that bypasses BatchGenerator.
|
||
|
||
The wrapper generator yields plain Python ints (single-request mode).
|
||
Scheduler iterates it one token per ``step()`` and feeds each token
|
||
into ``_process_batch_responses`` via a synthesized ``_VLMMTPResponse``.
|
||
"""
|
||
|
||
generator: Any # Generator[int, None, None] from run_vlm_mtp_decode
|
||
request: Request
|
||
prompt_cache: list[Any]
|
||
sampler: Callable[[Any], Any]
|
||
state_machine: Any
|
||
max_tokens: int
|
||
# Plain stop-token set (EOS + request-specific) for direct membership
|
||
# check; mlx-lm's SequenceStateMachine doesn't expose a "did the last
|
||
# token finish" helper, so we keep a copy.
|
||
stop_token_ids: set[int] = field(default_factory=set)
|
||
emitted: int = 0
|
||
finished: bool = False
|
||
|
||
|
||
@dataclass
|
||
class _VLMMTPResponse:
|
||
"""BatchGenerator.Response shim emitted by the vlm_mtp decode loop.
|
||
|
||
Same field surface used by ``_process_batch_responses``: ``uid``,
|
||
``token``, ``finish_reason``, ``logprobs``, and an optional
|
||
``prompt_cache`` returned on the terminal yield so paged-cache reuse
|
||
keeps working.
|
||
"""
|
||
|
||
uid: int
|
||
token: int
|
||
finish_reason: Optional[str] = None
|
||
logprobs: Any = None
|
||
prompt_cache: Any = None
|
||
|
||
|
||
# Serializes Metal buffer-protocol access from the async store-cache worker
|
||
# against inference-thread mx.clear_cache / mx.synchronize calls that can
|
||
# invalidate the underlying buffer pool. Closes a SIGABRT path where
|
||
# _async_store_cache_worker reads tensor bytes via memoryview while the
|
||
# inference thread concurrently issues a reclaim-triggering mx op.
|
||
# See: https://github.com/jundot/omlx/issues/1106
|
||
_mx_buffer_access_lock = threading.RLock()
|
||
|
||
|
||
def _sync_and_clear_cache(stream=None):
|
||
"""Synchronize in-flight GPU work before clearing the Metal buffer cache.
|
||
|
||
Without synchronization, mx.clear_cache() can release Metal buffers that
|
||
are still referenced by in-flight command buffers submitted via
|
||
mx.async_eval(). This causes the GPU driver to hit a
|
||
'completeMemory() prepare count underflow' kernel panic on M4 hardware
|
||
(and SIGSEGV/SIGABRT on M3).
|
||
|
||
Held under _mx_buffer_access_lock so the async store-cache worker cannot
|
||
observe a half-reclaimed Metal buffer pool while it is in the middle of
|
||
reading tensor bytes via the Python buffer protocol (#1106).
|
||
|
||
See: https://github.com/jundot/omlx/issues/300, #888, #1106
|
||
"""
|
||
with _mx_buffer_access_lock:
|
||
# The engine stream may not have in-flight work on the current thread
|
||
# (for example, during teardown before that thread submits work). On
|
||
# some MLX builds mx.synchronize raises "There is no Stream(gpu, 0) in
|
||
# current thread" in that case; swallow it since there is nothing to
|
||
# drain.
|
||
target = stream if stream is not None else _default_generation_stream
|
||
try:
|
||
mx.synchronize(target)
|
||
except RuntimeError:
|
||
pass
|
||
mx.synchronize() # default stream
|
||
mx.clear_cache()
|
||
|
||
|
||
def _safe_sync_stream(stream=None):
|
||
"""mx.synchronize(stream) that tolerates cross-thread calls.
|
||
|
||
The per-engine stream is owned by the engine's executor thread. Teardown
|
||
paths that run on the main thread (via EngineCore.close) hit "no Stream in
|
||
current thread" RuntimeError. Swallow that specific case so cleanup can
|
||
proceed; re-raise anything else so real GPU errors stay visible.
|
||
"""
|
||
target = stream if stream is not None else _default_generation_stream
|
||
try:
|
||
mx.synchronize(target)
|
||
except RuntimeError as e:
|
||
if "no Stream" not in str(e):
|
||
raise
|
||
|
||
|
||
def _env_int(name: str, default: int = 0) -> int:
|
||
value = os.environ.get(name)
|
||
if value is None or value == "":
|
||
return default
|
||
try:
|
||
return int(value)
|
||
except ValueError:
|
||
logger.warning("Ignoring invalid integer env %s=%r", name, value)
|
||
return default
|
||
|
||
|
||
def _collect_mx_arrays(value, out: list[mx.array]) -> None:
|
||
if isinstance(value, mx.array):
|
||
out.append(value)
|
||
elif isinstance(value, dict):
|
||
for item in value.values():
|
||
_collect_mx_arrays(item, out)
|
||
elif isinstance(value, (list, tuple)):
|
||
for item in value:
|
||
_collect_mx_arrays(item, out)
|
||
|
||
|
||
def _eval_generation_batch_cache(batch_generator) -> int:
|
||
generation_batch = getattr(batch_generator, "_generation_batch", None)
|
||
prompt_cache = getattr(generation_batch, "prompt_cache", None)
|
||
if not prompt_cache:
|
||
return 0
|
||
arrays: list[mx.array] = []
|
||
for cache in prompt_cache:
|
||
state = getattr(cache, "state", None)
|
||
if state is not None:
|
||
_collect_mx_arrays(state, arrays)
|
||
if arrays:
|
||
mx.eval(*arrays)
|
||
return len(arrays)
|
||
|
||
|
||
class _StoreCacheGate:
|
||
"""Non-blocking counter that bounds in-flight store-cache submissions.
|
||
|
||
Tracks how many KV caches are alive in the post-completion store-cache
|
||
pipeline. _cleanup_finished records each submission with note_submitted()
|
||
and _drain_pending_async_removes clears it with note_done() after the
|
||
deferred batch_generator.remove() releases the request cache references;
|
||
neither blocks the generation step. Backpressure is applied at admission
|
||
instead — _schedule_waiting declines to admit new prefills while
|
||
in_flight >= cap (see has_capacity), so token generation never stalls
|
||
waiting for an SSD write (#1496).
|
||
|
||
cap still bounds the concurrent extracted-KV count, which is the OOM
|
||
guard for the burst-finish RAM growth reported in #1383. It is adjusted
|
||
at runtime from ProcessMemoryEnforcer so the pipeline shrinks under
|
||
memory pressure on smaller systems.
|
||
"""
|
||
|
||
def __init__(self, cap: int) -> None:
|
||
self._cap = max(1, cap)
|
||
self._in_flight = 0
|
||
self._lock = threading.Lock()
|
||
|
||
def note_submitted(self) -> None:
|
||
"""Record a store-cache job handed to the executor (never blocks)."""
|
||
with self._lock:
|
||
self._in_flight += 1
|
||
|
||
def note_done(self) -> None:
|
||
"""Record a store-cache job finished (future done callback)."""
|
||
with self._lock:
|
||
if self._in_flight > 0:
|
||
self._in_flight -= 1
|
||
|
||
def set_cap(self, cap: int) -> None:
|
||
with self._lock:
|
||
self._cap = max(1, cap)
|
||
|
||
@property
|
||
def cap(self) -> int:
|
||
with self._lock:
|
||
return self._cap
|
||
|
||
@property
|
||
def in_flight(self) -> int:
|
||
with self._lock:
|
||
return self._in_flight
|
||
|
||
@property
|
||
def has_capacity(self) -> bool:
|
||
"""True when another submission would stay within cap.
|
||
|
||
Read by _schedule_waiting to decide whether to admit a new prefill.
|
||
"""
|
||
with self._lock:
|
||
return self._in_flight < self._cap
|
||
|
||
|
||
# Import tiered cache components
|
||
try:
|
||
from .cache.boundary_snapshot_store import BoundarySnapshotSSDStore
|
||
from .cache.paged_ssd_cache import PagedSSDCacheManager
|
||
from .memory_monitor import MemoryMonitor, estimate_mla_kv_bytes_per_token
|
||
|
||
HAS_TIERED_CACHE = True
|
||
except ImportError:
|
||
PagedSSDCacheManager = None
|
||
BoundarySnapshotSSDStore = None
|
||
MemoryMonitor = None
|
||
estimate_mla_kv_bytes_per_token = None
|
||
HAS_TIERED_CACHE = False
|
||
|
||
# Import cache type handlers for hybrid cache support
|
||
try:
|
||
from .cache.hybrid_cache import ModelCacheConfig
|
||
from .cache.type_registry import CacheTypeRegistry
|
||
|
||
HAS_CACHE_TYPE_HANDLERS = True
|
||
except ImportError:
|
||
CacheTypeRegistry = None
|
||
ModelCacheConfig = None
|
||
HAS_CACHE_TYPE_HANDLERS = False
|
||
|
||
# Import protocol-specific output parser support
|
||
try:
|
||
from .adapter.output_parser import (
|
||
OutputParserFactory,
|
||
OutputParserSession,
|
||
detect_output_parser,
|
||
)
|
||
|
||
HAS_OUTPUT_PARSER = True
|
||
except ImportError:
|
||
OutputParserFactory = None
|
||
OutputParserSession = None
|
||
detect_output_parser = None
|
||
HAS_OUTPUT_PARSER = False
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class _PrefillAbortedError(Exception):
|
||
"""Raised when prefill is interrupted by a pending abort."""
|
||
|
||
def __init__(self, aborted_uids: list[int], processed_tokens: int):
|
||
self.aborted_uids = aborted_uids
|
||
self.processed_tokens = processed_tokens
|
||
super().__init__(
|
||
f"Prefill aborted for UIDs {aborted_uids} " f"at {processed_tokens} tokens"
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class PrefillEvictionRequest:
|
||
"""Internal request for async LRU model eviction before prefill."""
|
||
|
||
request_id: str
|
||
model_id: str
|
||
current_bytes: int
|
||
target_cap_bytes: int
|
||
predicted_transient_bytes: int
|
||
requested_tokens: int
|
||
reason: str
|
||
|
||
|
||
class _PrefillEvictionNeeded(Exception):
|
||
"""Raised inside scheduler.step() to pause and request async eviction."""
|
||
|
||
def __init__(self, request: PrefillEvictionRequest):
|
||
super().__init__(request.reason)
|
||
self.request = request
|
||
|
||
|
||
def _prefill_memory_error_output(
|
||
request_id: str,
|
||
message: str,
|
||
*,
|
||
estimated_bytes: int | None = None,
|
||
limit_bytes: int | None = None,
|
||
) -> RequestOutput:
|
||
metadata: dict[str, int | str] = {"request_id": request_id}
|
||
if estimated_bytes is not None:
|
||
metadata["estimated_bytes"] = estimated_bytes
|
||
if limit_bytes is not None:
|
||
metadata["limit_bytes"] = limit_bytes
|
||
return RequestOutput(
|
||
request_id=request_id,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=message,
|
||
error_code="prefill_memory_exceeded",
|
||
error_metadata=metadata,
|
||
)
|
||
|
||
|
||
def _prefill_memory_exception_output(
|
||
request_id: str,
|
||
exc: PrefillMemoryExceededError,
|
||
) -> RequestOutput:
|
||
return _prefill_memory_error_output(
|
||
request_id,
|
||
str(exc),
|
||
estimated_bytes=exc.estimated_bytes,
|
||
limit_bytes=exc.limit_bytes,
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class _PrefillState:
|
||
"""Intermediate state for a request undergoing chunked prefill.
|
||
|
||
When chunked_prefill=True, a long prefill is spread across multiple
|
||
step() calls (one prefill_step_size chunk per step). This dataclass
|
||
holds all the state needed to resume prefill between steps.
|
||
"""
|
||
|
||
request: Any
|
||
cache: list # Accumulated prompt_cache (mutated in-place by each chunk)
|
||
tokens_remaining: Any # mx.array shape (1, N) — tokens not yet prefilled
|
||
last_token: list # tokens[-1:] — passed to batch_generator.insert()
|
||
tokens_processed: int # Cumulative count for boundary snapshot math
|
||
base_size: int # Prefix cache offset at prefill start (for alignment)
|
||
emitted_boundaries: dict # {request_id: int} — last emitted boundary count
|
||
boundary_enabled: bool # Whether boundary snapshots are active
|
||
block_size: int # Copied from config.paged_cache_block_size
|
||
total_length: int # len(original tokens) for completeness
|
||
# Pre-built insert-time params (set by _schedule_waiting before enqueuing)
|
||
sampler: Any = None
|
||
sm: Any = None
|
||
per_row_lps: Any = None
|
||
|
||
|
||
@dataclass
|
||
class _InflightStoreInfo:
|
||
tokens: list[int]
|
||
extra_keys: tuple[Any, ...] | None = None
|
||
extra_key_token_start: int | None = None
|
||
extra_key_ranges: list[tuple[int, tuple[Any, ...]]] | None = None
|
||
|
||
|
||
@dataclass
|
||
class _CacheFreshnessWait:
|
||
store_request_id: str
|
||
future: concurrent.futures.Future
|
||
common_prefix: int
|
||
prompt_len: int
|
||
deadline_s: float
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Monkey-patch GenerationBatch._step to call grammar accept_token() after
|
||
# sampling. In the pipelined _step(), logits processors fill the bitmask
|
||
# (constrain NEXT token) but can't know which token was just sampled.
|
||
# After _original_step returns, self._next_tokens holds the freshly sampled
|
||
# tokens. We eval them synchronously and accept in grammar processors.
|
||
# ---------------------------------------------------------------------------
|
||
# Authoritative per-uid row state for the generation batch.
|
||
#
|
||
# mlx-lm keeps ``samplers`` / ``logits_processors`` as positional lists that
|
||
# must stay aligned with ``uids``. Heterogeneous continuous batching
|
||
# (extend/filter/split across prompt and generation batches) can leave stale
|
||
# or offset row slots behind; #1799 made the step crash-safe by normalising
|
||
# ``None`` slots, but a misaligned row silently runs the WRONG sampler and
|
||
# logits processors (e.g. a grammar/thinking-budget request decoding with no
|
||
# constraints at all). The registry below records, at insert time, what each
|
||
# uid is supposed to run; the step chokepoint realigns the positional lists
|
||
# from it. Bounded so a missing cleanup can never grow it unbounded.
|
||
class _RegisteredRow(NamedTuple):
|
||
"""What a uid is supposed to run, recorded at request insert."""
|
||
|
||
sampler: Any
|
||
logits_processors: list
|
||
|
||
|
||
_UID_ROW_REGISTRY_MAX = 4096
|
||
# Keyed by (id(model), uid): mlx-lm's BatchGenerator numbers uids per
|
||
# instance starting at 0, so two engines serving concurrently (or an engine
|
||
# reload) produce colliding uid sequences. The model object is the one
|
||
# identity both the insert sites and the step chokepoint can see.
|
||
_uid_row_registry: "OrderedDict[tuple[int, int], _RegisteredRow]" = OrderedDict()
|
||
# Engines run on separate executor threads and share this module-level
|
||
# registry; a plain OrderedDict is not safe under concurrent mutation.
|
||
_uid_row_registry_lock = threading.Lock()
|
||
# Drift corrections are worth one log line each, but a pathological batching
|
||
# pattern could correct on every merge; cap the WARNING rate and route the
|
||
# rest to DEBUG so the signal survives without flooding the logs.
|
||
_UID_ROW_DRIFT_WARNING_INTERVAL_S = 60.0
|
||
_uid_row_drift_last_warning = float("-inf")
|
||
|
||
|
||
def _register_uid_rows(model, uids, samplers, lps_rows) -> None:
|
||
"""Record the sampler and logits processors each freshly-inserted uid must run.
|
||
|
||
Each (model, uid) key is inserted exactly once per request, so plain
|
||
insertion order is enough for the oldest-first backstop eviction.
|
||
"""
|
||
with _uid_row_registry_lock:
|
||
for uid, sampler, lps in zip(uids, samplers, lps_rows):
|
||
_uid_row_registry[(id(model), uid)] = _RegisteredRow(
|
||
sampler, list(lps or ())
|
||
)
|
||
while len(_uid_row_registry) > _UID_ROW_REGISTRY_MAX:
|
||
_uid_row_registry.popitem(last=False)
|
||
|
||
|
||
def _unregister_uid_row(model, uid) -> None:
|
||
"""Drop a finished request's row so heavy processors are not pinned
|
||
until FIFO eviction; the bounded size stays as the backstop."""
|
||
with _uid_row_registry_lock:
|
||
_uid_row_registry.pop((id(model), uid), None)
|
||
|
||
|
||
def _unregister_uid_rows_for_model(model) -> None:
|
||
"""Drop every registry row for a model (generator reset, recovery, shutdown).
|
||
|
||
The recovery and reset paths clear the uid maps wholesale instead of
|
||
finishing requests one by one; releasing by model covers them, and leaves
|
||
nothing behind that a later engine load could match if ``id(model)`` were
|
||
recycled.
|
||
"""
|
||
model_id = id(model)
|
||
with _uid_row_registry_lock:
|
||
for key in [key for key in _uid_row_registry if key[0] == model_id]:
|
||
del _uid_row_registry[key]
|
||
|
||
|
||
def _row_drifted(current_lps, expected_lps) -> bool:
|
||
"""True when a slot's processors genuinely differ from the registered row.
|
||
|
||
Two distinct empty lists are equivalent — the #1799 normalisation mints
|
||
fresh ``[]`` objects every step — so only differing content counts. The
|
||
caller's identity check is the steady-state fast path; this only runs
|
||
past it.
|
||
"""
|
||
if not current_lps and not expected_lps:
|
||
return False
|
||
return current_lps != expected_lps
|
||
|
||
|
||
def _log_drift_correction(uids, slot_count) -> None:
|
||
"""Log a corrected drift: one WARNING per window, the rest at DEBUG."""
|
||
global _uid_row_drift_last_warning
|
||
now = time.monotonic()
|
||
rate_limited = now - _uid_row_drift_last_warning < _UID_ROW_DRIFT_WARNING_INTERVAL_S
|
||
if not rate_limited:
|
||
_uid_row_drift_last_warning = now
|
||
(logger.debug if rate_limited else logger.warning)(
|
||
"Realigned generation-batch row state from the uid registry "
|
||
f"(uids={list(uids)}, had {slot_count} processor slots); "
|
||
"stale or offset slots would have run the wrong sampler/processors."
|
||
)
|
||
|
||
|
||
def _realigned_rows(model, uids, cur_samplers, cur_lps):
|
||
"""Rebuild the positional row lists in uid order from the registry.
|
||
|
||
Registered uids take their recorded row; unregistered uids keep their
|
||
current slot (the #1799 fallback), padding when the lists are shorter
|
||
than ``uids``. Returns ``(samplers, logits_processors, drift)`` — drift
|
||
only drives logging, the rebuilt lists are always installed. In steady
|
||
state the slots already are the registry lists, so the identity check
|
||
skips any comparison work.
|
||
"""
|
||
model_id = id(model)
|
||
with _uid_row_registry_lock:
|
||
rows = [_uid_row_registry.get((model_id, uid)) for uid in uids]
|
||
|
||
drift = len(cur_lps) != len(uids)
|
||
samplers, lps = [], []
|
||
for i, row in enumerate(rows):
|
||
if row is not None:
|
||
if not drift:
|
||
if i >= len(cur_samplers):
|
||
drift = row.sampler is not None
|
||
elif cur_samplers[i] is not row.sampler:
|
||
drift = True
|
||
if (
|
||
not drift
|
||
and i < len(cur_lps)
|
||
and cur_lps[i] is not row.logits_processors
|
||
):
|
||
drift = _row_drifted(cur_lps[i], row.logits_processors)
|
||
samplers.append(row.sampler)
|
||
lps.append(row.logits_processors)
|
||
else:
|
||
samplers.append(cur_samplers[i] if i < len(cur_samplers) else None)
|
||
lps.append(cur_lps[i] if i < len(cur_lps) else [])
|
||
return samplers, lps, drift
|
||
|
||
|
||
def _omlx_realign_generation_batch_rows(self) -> None:
|
||
"""Realign positional row state with ``uids`` before any decode path reads it."""
|
||
if self.logits_processors is None:
|
||
self.logits_processors = []
|
||
else:
|
||
self.logits_processors = [
|
||
procs if procs is not None else [] for procs in self.logits_processors
|
||
]
|
||
|
||
uids = getattr(self, "uids", None) or []
|
||
if not uids:
|
||
return
|
||
|
||
new_samplers, new_lps, drift = _realigned_rows(
|
||
getattr(self, "model", None),
|
||
uids,
|
||
getattr(self, "samplers", None) or [],
|
||
self.logits_processors,
|
||
)
|
||
if drift:
|
||
_log_drift_correction(uids, len(self.logits_processors))
|
||
self.logits_processors = new_lps
|
||
self.samplers = new_samplers
|
||
|
||
|
||
_original_generation_batch_step = GenerationBatch._step
|
||
|
||
|
||
def _patched_generation_batch_step(self):
|
||
# Build per-batch mRoPE deltas from UID mapping before each step.
|
||
# This handles batch size changes during prompt split/generate.
|
||
model = self.model
|
||
if (
|
||
getattr(model, "_uses_mrope", False)
|
||
and getattr(model, "_uid_rope_deltas", None)
|
||
and self.uids
|
||
):
|
||
deltas = [model._uid_rope_deltas.get(uid, 0.0) for uid in self.uids]
|
||
model.set_batch_rope_deltas(mx.array(deltas))
|
||
|
||
# Defensive: mlx-lm's GenerationBatch._step does `any(self.logits_processors)`
|
||
# and `for p in self.logits_processors[e]`, both of which crash when a row
|
||
# slot is None. Normalise the whole list AND every per-row slot to [] here,
|
||
# at the single consumption chokepoint, so the original step and the
|
||
# grammar-accept loop below are both safe regardless of slot origin.
|
||
#
|
||
# The insert call sites already wrap each request's processors as a list,
|
||
# but that is not enough: on a heterogeneous continuous-batch merge,
|
||
# mlx-lm's GenerationBatch.extend() re-introduces None slots via
|
||
# `if not any(self.logits_processors): self.logits_processors =
|
||
# [None] * len(self.uids)`. `any([[], []])` is False, so empty-list slots
|
||
# collapse back to None whenever a batch with no *active* processor merges
|
||
# with a grammar-constrained one (e.g. a plain chat request joining a batch
|
||
# that is serving a structured json_schema request). Per-row normalisation
|
||
# at this chokepoint is the only place that covers both insert and merge.
|
||
# See #934 / #1747.
|
||
_omlx_realign_generation_batch_rows(self)
|
||
|
||
result = _original_generation_batch_step(self)
|
||
|
||
# self._next_tokens contains the just-sampled tokens (async eval pending).
|
||
# We need to accept them NOW so the next __call__ fills the correct bitmask.
|
||
if any(self.logits_processors):
|
||
from .api.grammar import GrammarConstraintProcessor
|
||
|
||
has_grammar = any(
|
||
isinstance(p, GrammarConstraintProcessor)
|
||
for procs in self.logits_processors
|
||
for p in procs
|
||
)
|
||
if has_grammar:
|
||
# Force eval of the sampled tokens so we can read them.
|
||
mx.eval(self._next_tokens)
|
||
sampled = self._next_tokens.tolist()
|
||
for e in range(len(self.uids)):
|
||
for proc in self.logits_processors[e]:
|
||
if isinstance(proc, GrammarConstraintProcessor):
|
||
proc.accept_token(sampled[e])
|
||
|
||
return result
|
||
|
||
|
||
GenerationBatch._omlx_realign_rows = _omlx_realign_generation_batch_rows
|
||
GenerationBatch._step = _patched_generation_batch_step
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Monkey-patch GenerationBatch.filter to keep logits_processors aligned with
|
||
# uids. mlx-lm's filter only reindexes the processor list when at least one
|
||
# row has an active processor:
|
||
#
|
||
# if any(self.logits_processors):
|
||
# self.logits_processors = [self.logits_processors[idx] for idx in keep]
|
||
#
|
||
# There is no else branch (unlike the prompt-batch class, which resets to
|
||
# ``[[]] * len(keep)``), so when every slot is empty — the normal state after
|
||
# serving requests without per-request processors — the stale list survives
|
||
# while uids/tokens shrink. A later extend() then appends the next request's
|
||
# processors BEHIND its own row index: the row reads a leftover empty slot and
|
||
# the real processor (thinking budget, grammar constraint) is silently never
|
||
# applied. Which requests are affected depends on insertion/removal order,
|
||
# and alignment self-heals once the broken request finishes, so the symptom
|
||
# is an intermittently ignored thinking_budget or grammar. See #934/#1747
|
||
# for the sibling None-slot collapse handled in _patched_generation_batch_step.
|
||
_original_generation_batch_filter = GenerationBatch.filter
|
||
|
||
|
||
def _patched_generation_batch_filter(self, keep):
|
||
lps = self.logits_processors
|
||
lps_inert = not lps or not any(lps)
|
||
if lps is None:
|
||
# ``any(None)`` inside the original filter raises TypeError.
|
||
self.logits_processors = []
|
||
_original_generation_batch_filter(self, keep)
|
||
if lps_inert:
|
||
# Original filter skipped the reindex; reset to one empty slot per
|
||
# surviving row so extend() appends at the correct indices.
|
||
self.logits_processors = [[] for _ in keep]
|
||
|
||
|
||
GenerationBatch.filter = _patched_generation_batch_filter
|
||
|
||
|
||
_TQ_SINGLETON_CACHE_TYPE: type[Any] | None = None
|
||
|
||
# Monkey-patch TurboQuantKVCache.merge so _merge_caches() works
|
||
try:
|
||
from mlx_vlm.turboquant import TurboQuantKVCache as _TQCache
|
||
|
||
from .turboquant_kv import BatchTurboQuantKVCache as _BTQCache
|
||
|
||
_TQ_SINGLETON_CACHE_TYPE = _TQCache
|
||
if not hasattr(_TQCache, "merge"):
|
||
_TQCache.merge = _BTQCache.merge
|
||
except ImportError:
|
||
pass
|
||
|
||
|
||
# Regular singleton KV caches are already the fastest decode representation.
|
||
# mlx-lm's default _merge_caches([cache]) turns them into BatchKVCache even
|
||
# when there is only one active row, which slows text-only VLM decode. Install
|
||
# the minimal BatchGenerator methods needed while the row count remains one;
|
||
# _patched_extend_cache converts them back to batched caches before a second
|
||
# row is appended.
|
||
def _batch_indices_len(batch_indices: Any) -> int:
|
||
try:
|
||
return len(batch_indices)
|
||
except TypeError:
|
||
return int(getattr(batch_indices, "shape", (0,))[0] or 0)
|
||
|
||
|
||
def _regular_kv_filter_singleton(self, batch_indices):
|
||
n = _batch_indices_len(batch_indices)
|
||
if n == 0:
|
||
self.keys = None
|
||
self.values = None
|
||
self.offset = 0
|
||
return
|
||
if n == 1:
|
||
return
|
||
raise NotImplementedError(
|
||
f"{type(self).__name__}.filter only supports singleton pass-through; "
|
||
"convert to a batched cache before keeping multiple rows."
|
||
)
|
||
|
||
|
||
def _regular_rotating_kv_filter_singleton(self, batch_indices):
|
||
n = _batch_indices_len(batch_indices)
|
||
if n == 0:
|
||
self.keys = None
|
||
self.values = None
|
||
self.offset = 0
|
||
self._idx = 0
|
||
return
|
||
if n == 1:
|
||
return
|
||
raise NotImplementedError(
|
||
f"{type(self).__name__}.filter only supports singleton pass-through; "
|
||
"convert to a batched cache before keeping multiple rows."
|
||
)
|
||
|
||
|
||
def _regular_cache_extract_singleton(self, idx: int):
|
||
if int(idx) != 0:
|
||
raise IndexError(f"{type(self).__name__} singleton cache only has row 0")
|
||
return self
|
||
|
||
|
||
def _regular_cache_extend_singleton(self, other):
|
||
raise NotImplementedError(
|
||
f"{type(self).__name__}.extend requires batched conversion first"
|
||
)
|
||
|
||
|
||
def _turboquant_filter_singleton(self, batch_indices):
|
||
n = _batch_indices_len(batch_indices)
|
||
if n == 0:
|
||
self.keys = None
|
||
self.values = None
|
||
self.offset = 0
|
||
self._cached_state = None
|
||
self._cached_state_offset = -1
|
||
if hasattr(self, "_shadow_keys"):
|
||
self._shadow_keys = None
|
||
if hasattr(self, "_shadow_values"):
|
||
self._shadow_values = None
|
||
return
|
||
if n == 1:
|
||
return
|
||
raise NotImplementedError(
|
||
f"{type(self).__name__}.filter only supports singleton pass-through; "
|
||
"convert to a batched cache before keeping multiple rows."
|
||
)
|
||
|
||
|
||
if not hasattr(_MLXKVCache, "filter"):
|
||
_MLXKVCache.filter = _regular_kv_filter_singleton
|
||
if not hasattr(_MLXKVCache, "extract"):
|
||
_MLXKVCache.extract = _regular_cache_extract_singleton
|
||
if not hasattr(_MLXKVCache, "extend"):
|
||
_MLXKVCache.extend = _regular_cache_extend_singleton
|
||
|
||
if not hasattr(_MLXRotatingKVCache, "filter"):
|
||
_MLXRotatingKVCache.filter = _regular_rotating_kv_filter_singleton
|
||
if not hasattr(_MLXRotatingKVCache, "extract"):
|
||
_MLXRotatingKVCache.extract = _regular_cache_extract_singleton
|
||
if not hasattr(_MLXRotatingKVCache, "extend"):
|
||
_MLXRotatingKVCache.extend = _regular_cache_extend_singleton
|
||
|
||
if _TQ_SINGLETON_CACHE_TYPE is not None:
|
||
if not hasattr(_TQ_SINGLETON_CACHE_TYPE, "filter"):
|
||
_TQ_SINGLETON_CACHE_TYPE.filter = _turboquant_filter_singleton
|
||
if not hasattr(_TQ_SINGLETON_CACHE_TYPE, "extract"):
|
||
_TQ_SINGLETON_CACHE_TYPE.extract = _regular_cache_extract_singleton
|
||
if not hasattr(_TQ_SINGLETON_CACHE_TYPE, "extend"):
|
||
_TQ_SINGLETON_CACHE_TYPE.extend = _regular_cache_extend_singleton
|
||
|
||
_mlx_lm_generate_module = importlib.import_module("mlx_lm.generate")
|
||
_original_merge_caches = _mlx_lm_generate_module._merge_caches
|
||
_original_ppb_split = PromptProcessingBatch.split
|
||
_REGULAR_SINGLETON_CACHE_TYPES = (_MLXKVCache, _MLXRotatingKVCache)
|
||
|
||
|
||
def _cache_layer_supports_singleton_passthrough(cache_obj: Any) -> bool:
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
return all(_cache_layer_supports_singleton_passthrough(c) for c in sub_caches)
|
||
return hasattr(cache_obj, "filter") and hasattr(cache_obj, "extract")
|
||
|
||
|
||
def _to_batched_cache_layer(cache_obj: Any) -> Any:
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
converted = tuple(_to_batched_cache_layer(c) for c in sub_caches)
|
||
if all(a is b for a, b in zip(sub_caches, converted)):
|
||
return cache_obj
|
||
return type(cache_obj)(*converted)
|
||
if isinstance(cache_obj, _REGULAR_SINGLETON_CACHE_TYPES):
|
||
return cache_obj.merge([cache_obj])
|
||
if (
|
||
_TQ_SINGLETON_CACHE_TYPE is not None
|
||
and type(cache_obj) is _TQ_SINGLETON_CACHE_TYPE
|
||
):
|
||
return cache_obj.merge([cache_obj])
|
||
return cache_obj
|
||
|
||
|
||
def _extend_cache_layer(cache_a: Any, cache_b: Any) -> Any:
|
||
sub_a = getattr(cache_a, "caches", None)
|
||
sub_b = getattr(cache_b, "caches", None)
|
||
if isinstance(sub_a, (list, tuple)) and isinstance(sub_b, (list, tuple)):
|
||
cache_a.caches = tuple(
|
||
_extend_cache_layer(ca, cb) for ca, cb in zip(sub_a, sub_b)
|
||
)
|
||
return cache_a
|
||
|
||
cache_a = _to_batched_cache_layer(cache_a)
|
||
cache_b = _to_batched_cache_layer(cache_b)
|
||
cache_a.extend(cache_b)
|
||
return cache_a
|
||
|
||
|
||
def _patched_merge_caches(caches):
|
||
if not caches:
|
||
return []
|
||
if len(caches) == 1:
|
||
merged = []
|
||
for layer_cache in caches[0]:
|
||
if _cache_layer_supports_singleton_passthrough(layer_cache):
|
||
merged.append(layer_cache)
|
||
elif hasattr(layer_cache, "merge"):
|
||
merged.append(layer_cache.merge([layer_cache]))
|
||
else:
|
||
raise ValueError(
|
||
f"{type(layer_cache)} does not yet support batching with history"
|
||
)
|
||
return merged
|
||
return _original_merge_caches(caches)
|
||
|
||
|
||
def _patched_extend_cache(cache_a, cache_b):
|
||
if not cache_a:
|
||
return cache_b
|
||
if not cache_b:
|
||
return cache_a
|
||
return [_extend_cache_layer(ca, cb) for ca, cb in zip(cache_a, cache_b)]
|
||
|
||
|
||
def _patched_ppb_split(self, indices):
|
||
sorted_indices = sorted(indices)
|
||
if sorted_indices and sorted_indices == list(range(len(self.uids))):
|
||
new_batch = self.__class__.__new__(self.__class__)
|
||
new_batch.model = self.model
|
||
new_batch.uids = self.uids
|
||
new_batch.prompt_cache = self.prompt_cache
|
||
new_batch.tokens = self.tokens
|
||
new_batch.prefill_step_size = self.prefill_step_size
|
||
new_batch.samplers = self.samplers
|
||
new_batch.fallback_sampler = self.fallback_sampler
|
||
# Defensive: normalise None → [] to avoid mlx-lm crash in _step
|
||
lps = self.logits_processors if self.logits_processors is not None else []
|
||
new_batch.logits_processors = lps
|
||
new_batch.state_machines = self.state_machines
|
||
new_batch.max_tokens = self.max_tokens
|
||
if hasattr(self, "_omlx_glm_dsa_adaptive_prefill"):
|
||
new_batch._omlx_glm_dsa_adaptive_prefill = (
|
||
self._omlx_glm_dsa_adaptive_prefill
|
||
)
|
||
|
||
self.uids = []
|
||
self.prompt_cache = []
|
||
self.tokens = []
|
||
self.samplers = []
|
||
self.logits_processors = []
|
||
self.state_machines = []
|
||
self.max_tokens = []
|
||
return new_batch
|
||
return _original_ppb_split(self, indices)
|
||
|
||
|
||
_mlx_lm_generate_module._merge_caches = _patched_merge_caches
|
||
_mlx_lm_generate_module._extend_cache = _patched_extend_cache
|
||
PromptProcessingBatch.split = _patched_ppb_split
|
||
|
||
|
||
# Monkey-patch ChunkedKVCache for Llama-4 (Scout / Maverick): mlx_lm's
|
||
# ChunkedKVCache lacks the batch-aware methods (`merge`, `filter`, `extract`,
|
||
# `size`, `extend`) that BatchGenerator's continuous-batching code path
|
||
# expects, so any chat completion targeting a Llama-4 model raises
|
||
# `Cache corruption not recoverable: <ChunkedKVCache> does not yet support
|
||
# batching with history` and returns 500.
|
||
#
|
||
# Real continuous batching with chunked attention is unimplemented upstream;
|
||
# this patch installs batch=1 pass-throughs so serialized requests work.
|
||
# Run the server with `--max-concurrent-requests 1` to honor the assumption.
|
||
try:
|
||
from mlx_lm.models.cache import ChunkedKVCache as _CKVCache
|
||
|
||
_ckvcache_methods_skipped: list[str] = []
|
||
|
||
if not hasattr(_CKVCache, "merge"):
|
||
|
||
@classmethod
|
||
def _ckvcache_merge_passthrough(cls, caches):
|
||
if len(caches) == 1:
|
||
return caches[0]
|
||
raise NotImplementedError(
|
||
"ChunkedKVCache.merge for batch_size > 1 is not implemented. "
|
||
"Run with --max-concurrent-requests 1 when serving Llama-4."
|
||
)
|
||
|
||
_CKVCache.merge = _ckvcache_merge_passthrough
|
||
else:
|
||
_ckvcache_methods_skipped.append("merge")
|
||
|
||
if not hasattr(_CKVCache, "filter"):
|
||
|
||
def _ckvcache_filter_passthrough(self, batch_indices):
|
||
try:
|
||
n = len(batch_indices)
|
||
except TypeError:
|
||
n = int(getattr(batch_indices, "shape", (0,))[0] or 0)
|
||
if n == 0:
|
||
self.keys = None
|
||
self.values = None
|
||
self.offset = 0
|
||
self.start_position = 0
|
||
return
|
||
if n == 1:
|
||
return
|
||
raise NotImplementedError(
|
||
f"ChunkedKVCache.filter with batch_size={n} > 1 is not "
|
||
"implemented. Run with --max-concurrent-requests 1 when "
|
||
"serving Llama-4."
|
||
)
|
||
|
||
_CKVCache.filter = _ckvcache_filter_passthrough
|
||
else:
|
||
_ckvcache_methods_skipped.append("filter")
|
||
|
||
if not hasattr(_CKVCache, "extract"):
|
||
|
||
def _ckvcache_extract_passthrough(self, idx):
|
||
return self
|
||
|
||
_CKVCache.extract = _ckvcache_extract_passthrough
|
||
else:
|
||
_ckvcache_methods_skipped.append("extract")
|
||
|
||
if not hasattr(_CKVCache, "size"):
|
||
|
||
def _ckvcache_size(self):
|
||
return max(0, self.offset - self.start_position)
|
||
|
||
_CKVCache.size = _ckvcache_size
|
||
else:
|
||
_ckvcache_methods_skipped.append("size")
|
||
|
||
if not hasattr(_CKVCache, "extend"):
|
||
|
||
def _ckvcache_extend_passthrough(self, other):
|
||
if other is None or other.empty():
|
||
return
|
||
if self.empty():
|
||
self.keys = other.keys
|
||
self.values = other.values
|
||
self.offset = other.offset
|
||
self.start_position = other.start_position
|
||
return
|
||
raise NotImplementedError(
|
||
"ChunkedKVCache.extend across non-empty caches is not "
|
||
"supported. Run with --max-concurrent-requests 1."
|
||
)
|
||
|
||
_CKVCache.extend = _ckvcache_extend_passthrough
|
||
else:
|
||
_ckvcache_methods_skipped.append("extend")
|
||
|
||
if _ckvcache_methods_skipped:
|
||
# Upstream may have landed implementations between mlx_lm upgrades.
|
||
# Surface which ones so a regression in Llama-4 batching is visible
|
||
# to operators without diffing the patch against installed mlx_lm.
|
||
logger.info(
|
||
"ChunkedKVCache patch: methods already present upstream, " "skipped: %s",
|
||
", ".join(_ckvcache_methods_skipped),
|
||
)
|
||
except ImportError:
|
||
pass
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Monkey-patch PromptProcessingBatch.prompt to set mRoPE deltas before the
|
||
# prompt processing loop. Without this, batched VLM prompt processing
|
||
# (e.g. the 1-token final prompt after external prefill) would use
|
||
# per-request offsets without rope_deltas, corrupting attention masks
|
||
# for concurrent VLM requests.
|
||
# ---------------------------------------------------------------------------
|
||
_original_ppb_prompt = PromptProcessingBatch.prompt
|
||
|
||
|
||
def _patched_ppb_prompt(self, tokens):
|
||
model = self.model
|
||
if (
|
||
getattr(model, "_uses_mrope", False)
|
||
and getattr(model, "_uid_rope_deltas", None)
|
||
and self.uids
|
||
):
|
||
deltas = [model._uid_rope_deltas.get(uid, 0.0) for uid in self.uids]
|
||
model.set_batch_rope_deltas(mx.array(deltas))
|
||
return _original_ppb_prompt(self, tokens)
|
||
|
||
|
||
PromptProcessingBatch.prompt = _patched_ppb_prompt
|
||
|
||
|
||
# Cache class names known to be sliceable (no boundary snapshots needed).
|
||
# ChunkedKVCache is included once the batch=1 patch above installs its
|
||
# extract/filter/size pass-throughs; without it Llama-4 requests fall
|
||
# back to the snapshot path unnecessarily.
|
||
_KNOWN_SLICEABLE_CACHE_TYPES = frozenset(
|
||
{
|
||
"KVCache",
|
||
"BatchKVCache",
|
||
"QuantizedKVCache",
|
||
"TurboQuantKVCache",
|
||
"BatchTurboQuantKVCache",
|
||
"ChunkedKVCache",
|
||
"MiniMaxM3KVCache",
|
||
}
|
||
)
|
||
|
||
|
||
_TURBOQUANT_KV_CACHE_TYPES = frozenset(
|
||
{
|
||
"TurboQuantKVCache",
|
||
"BatchTurboQuantKVCache",
|
||
}
|
||
)
|
||
|
||
|
||
def _is_turboquant_kv_cache(cache_obj: Any) -> bool:
|
||
return type(cache_obj).__name__ in _TURBOQUANT_KV_CACHE_TYPES
|
||
|
||
|
||
def _is_turboquant_kv_family_cache(cache_obj: Any) -> bool:
|
||
"""Cache layer counted by TurboQuant's skip-last full-attention rule."""
|
||
return isinstance(cache_obj, _MLXKVCache) or _is_turboquant_kv_cache(cache_obj)
|
||
|
||
|
||
class _BoundaryStoreUnavailable(Exception):
|
||
"""No boundary-aligned snapshot exists for non-sliceable cache state.
|
||
|
||
Raised inside the final-response store path when the model needs
|
||
boundary snapshots but none were captured (e.g. every capture was
|
||
skipped by the speculative-decode skew guard). Storing the live
|
||
state instead would label off-boundary recurrent/rotating state
|
||
with a block-boundary token count and corrupt later prefix hits;
|
||
the handler skips the store and releases the block references.
|
||
"""
|
||
|
||
|
||
def _first_leaf_cache_offset(cache_obj: Any) -> int | None:
|
||
"""First integer ``offset`` found walking into composite caches.
|
||
|
||
CacheList-style wrappers (DeepSeek-V4 / GLM per-layer caches) expose
|
||
no ``offset`` themselves — the token position lives on their leading
|
||
sub-cache (RotatingKVCache / KVCache, whose ``offset`` counts
|
||
forwarded tokens). Walk depth-first so that leading leaf decides;
|
||
later sub-caches may count something else entirely (PoolingCache's
|
||
``offset`` is the pooled-window count).
|
||
"""
|
||
subs = getattr(cache_obj, "caches", None)
|
||
if isinstance(subs, (list, tuple)):
|
||
for sub in subs:
|
||
offset = _first_leaf_cache_offset(sub)
|
||
if offset is not None:
|
||
return offset
|
||
return None
|
||
offset = getattr(cache_obj, "offset", None)
|
||
if offset is None:
|
||
return None
|
||
try:
|
||
return int(offset)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _prompt_cache_needs_snapshots(prompt_cache: list[Any]) -> bool:
|
||
"""Return True if any layer cache is non-sliceable (needs snapshots).
|
||
|
||
Checks the cache objects created during prefill. If all layers
|
||
are known-sliceable types (e.g. KVCache), boundary snapshots
|
||
are unnecessary and can be skipped entirely.
|
||
"""
|
||
for cache_obj in prompt_cache:
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
for sub in sub_caches:
|
||
if type(sub).__name__ not in _KNOWN_SLICEABLE_CACHE_TYPES:
|
||
return True
|
||
elif type(cache_obj).__name__ not in _KNOWN_SLICEABLE_CACHE_TYPES:
|
||
return True
|
||
return False
|
||
|
||
|
||
def _batch_generator_all_tokens(request: Any) -> list[int]:
|
||
"""Seed tokens for mlx-lm's TokenBuffer before the kickoff token."""
|
||
token_ids = getattr(request, "prompt_token_ids", None)
|
||
if token_ids is None:
|
||
return []
|
||
return list(token_ids[:-1])
|
||
|
||
|
||
def _cache_layer_token_count(cache_obj: Any) -> int:
|
||
"""Return the number of tokens stored in a single cache layer."""
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)) and sub_caches:
|
||
return max(_cache_layer_token_count(sub_cache) for sub_cache in sub_caches)
|
||
|
||
offset = getattr(cache_obj, "offset", None)
|
||
if isinstance(offset, (int, float)):
|
||
return int(offset)
|
||
|
||
size_fn = getattr(cache_obj, "size", None)
|
||
if callable(size_fn):
|
||
try:
|
||
return int(size_fn())
|
||
except Exception:
|
||
return 0
|
||
|
||
return 0
|
||
|
||
|
||
def _cache_base_sizes(caches: list[Any]) -> int:
|
||
"""Return the base token count of a single-request cache list."""
|
||
if not caches:
|
||
return 0
|
||
try:
|
||
return max(_cache_layer_token_count(c) for c in caches)
|
||
except Exception:
|
||
return 0
|
||
|
||
|
||
def _collect_cache_storage_arrays(cache_obj: Any) -> list[mx.array]:
|
||
"""Collect concrete backing arrays from cache objects, not state slices."""
|
||
arrays: list[mx.array] = []
|
||
|
||
if isinstance(cache_obj, mx.array):
|
||
return [cache_obj]
|
||
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
for sub_cache in sub_caches:
|
||
arrays.extend(_collect_cache_storage_arrays(sub_cache))
|
||
|
||
array_cache = getattr(cache_obj, "cache", None)
|
||
if isinstance(array_cache, (list, tuple)):
|
||
for item in array_cache:
|
||
arrays.extend(_collect_cache_storage_arrays(item))
|
||
|
||
for attr in ("keys", "values", "left_padding", "lengths"):
|
||
value = getattr(cache_obj, attr, None)
|
||
if isinstance(value, mx.array):
|
||
arrays.append(value)
|
||
|
||
return arrays
|
||
|
||
|
||
def _materialize_cache_storage(cache_list: list[Any]) -> None:
|
||
"""Force restored cache backing arrays concrete before decode begins."""
|
||
arrays: list[mx.array] = []
|
||
for cache_obj in cache_list:
|
||
arrays.extend(_collect_cache_storage_arrays(cache_obj))
|
||
if arrays:
|
||
with _mx_buffer_access_lock:
|
||
mx.eval(*arrays)
|
||
|
||
|
||
def _seed_text_only_mrope_delta_for_cached_prefill(model: Any, request: Any) -> None:
|
||
"""Seed zero mRoPE delta after clearing text-only cached-prefix state."""
|
||
if getattr(request, "cached_tokens", 0) <= 0:
|
||
return
|
||
lm = getattr(model, "_language_model", None)
|
||
if lm is None or not hasattr(lm, "_rope_deltas"):
|
||
return
|
||
lm._rope_deltas = mx.zeros((1, 1), dtype=mx.int64)
|
||
|
||
|
||
def _vlm_extra_seq_slice(val: mx.array, s: slice) -> mx.array:
|
||
"""Slice a VLM extra tensor along its seq dimension.
|
||
|
||
Standard layout (batch=1, seq, ...): seq at dim 1.
|
||
Special layout (e.g. mRoPE (3, batch, seq)): seq at last dim.
|
||
"""
|
||
if val.ndim >= 3 and val.shape[0] == 1:
|
||
return val[:, s]
|
||
if val.ndim >= 3:
|
||
return val[..., s]
|
||
return val[:, s]
|
||
|
||
|
||
def _slice_vlm_extra(extra: dict[str, Any], n: int) -> dict[str, Any]:
|
||
"""Slice VLM extra kwargs to first n tokens along seq dimension."""
|
||
sliced: dict[str, Any] = {}
|
||
for key, val in extra.items():
|
||
if isinstance(val, mx.array) and val.ndim >= 2:
|
||
sliced[key] = _vlm_extra_seq_slice(val, slice(None, n))
|
||
else:
|
||
sliced[key] = val
|
||
return sliced
|
||
|
||
|
||
def _advance_vlm_extra(extra: dict[str, Any], n: int) -> dict[str, Any]:
|
||
"""Advance VLM extra kwargs past first n tokens along seq dimension."""
|
||
advanced: dict[str, Any] = {}
|
||
for key, val in extra.items():
|
||
if isinstance(val, mx.array) and val.ndim >= 2:
|
||
advanced[key] = _vlm_extra_seq_slice(val, slice(n, None))
|
||
else:
|
||
advanced[key] = val
|
||
return advanced
|
||
|
||
|
||
def _get_attr_or_key(obj: Any, name: str) -> Any:
|
||
if isinstance(obj, dict):
|
||
return obj.get(name)
|
||
try:
|
||
value = getattr(obj, name)
|
||
except Exception:
|
||
return None
|
||
if type(value).__module__.startswith("unittest.mock"):
|
||
return None
|
||
return value
|
||
|
||
|
||
def _model_declares_llama4(model: Any) -> bool:
|
||
"""Return True if the loaded model/config tree declares Llama 4."""
|
||
seen: set[int] = set()
|
||
stack = [model]
|
||
while stack:
|
||
obj = stack.pop()
|
||
if obj is None:
|
||
continue
|
||
obj_id = id(obj)
|
||
if obj_id in seen:
|
||
continue
|
||
seen.add(obj_id)
|
||
|
||
if _get_attr_or_key(obj, "model_type") == "llama4":
|
||
return True
|
||
|
||
for attr in ("config", "args", "text_config", "language_config", "llm_config"):
|
||
child = _get_attr_or_key(obj, attr)
|
||
if child is not None and not isinstance(
|
||
child, (str, bytes, int, float, bool)
|
||
):
|
||
stack.append(child)
|
||
return False
|
||
|
||
|
||
class SchedulingPolicy(Enum):
|
||
"""Scheduling policy for request ordering."""
|
||
|
||
FCFS = "fcfs" # First-Come-First-Served
|
||
PRIORITY = "priority" # Priority-based
|
||
|
||
|
||
@dataclass
|
||
class SchedulerConfig:
|
||
"""Configuration for the scheduler."""
|
||
|
||
# Maximum number of concurrent requests in the batch
|
||
max_num_seqs: int = 256
|
||
# Maximum tokens to process per step (for prefill chunking)
|
||
max_num_batched_tokens: int = 8192
|
||
# Scheduling policy
|
||
policy: SchedulingPolicy = SchedulingPolicy.FCFS
|
||
# BatchGenerator settings (passed directly to mlx-lm)
|
||
completion_batch_size: int = 32
|
||
# Per-forward embedding input chunk size
|
||
embedding_batch_size: int = 32
|
||
prefill_step_size: int = 2048
|
||
# When True, long prefills are processed one chunk per step() call,
|
||
# interleaved with decode steps for already-running requests. This
|
||
# reduces TTFT for concurrent requests but adds per-step overhead.
|
||
chunked_prefill: bool = False
|
||
|
||
# Paged cache settings (internal defaults)
|
||
paged_cache_block_size: int = 256 # Tokens per block
|
||
max_cache_blocks: int | None = (
|
||
None # Auto-calculated from available KV cache memory
|
||
)
|
||
initial_cache_blocks: int = (
|
||
256 # Starting blocks (grows dynamically to max_cache_blocks)
|
||
)
|
||
|
||
# paged SSD cache settings (oMLX only supports paged SSD-based caching)
|
||
# When paged_ssd_cache_dir is set, oMLX stores KV cache on paged SSD for prefix reuse.
|
||
# When None, no oMLX caching (mlx-lm BatchGenerator manages KV internally).
|
||
paged_ssd_cache_dir: str | None = (
|
||
None # Path for paged SSD cache storage (None = disabled)
|
||
)
|
||
hot_cache_only: bool = False
|
||
paged_ssd_cache_max_size: int = 100 * 1024 * 1024 * 1024 # 100GB default
|
||
hot_cache_max_size: int = 0 # In-memory hot cache size in bytes (0 = disabled)
|
||
hot_cache_budget: Any | None = None # Shared process-wide hot cache budget
|
||
|
||
# Model identification (for cache isolation between different models)
|
||
model_name: str = "" # OpenAI API model name (e.g., "mlx-community/Llama-3.2-3B")
|
||
model_path: str = "" # Filesystem path to the model (e.g., "/cache/models--Org--Name/snapshots/abc123")
|
||
|
||
# GC/cleanup settings (memory optimization)
|
||
gc_cleanup_interval: int = 0 # Steps between gc.collect() calls (0=disabled)
|
||
mlx_cache_cleanup_interval: int = 512 # Steps between mx.clear_cache() calls
|
||
|
||
|
||
@dataclass
|
||
class SchedulerOutput:
|
||
"""
|
||
Output from a scheduling step.
|
||
|
||
Contains information about what was scheduled and results.
|
||
"""
|
||
|
||
# Requests scheduled in this step
|
||
scheduled_request_ids: list[str] = field(default_factory=list)
|
||
# Total tokens scheduled
|
||
num_scheduled_tokens: int = 0
|
||
# Requests that finished in this step
|
||
finished_request_ids: set[str] = field(default_factory=set)
|
||
# Request outputs (tokens generated)
|
||
outputs: list[RequestOutput] = field(default_factory=list)
|
||
# Internal signal consumed by EngineCore; not part of any API response.
|
||
prefill_eviction_request: PrefillEvictionRequest | None = None
|
||
# Whether any work was done
|
||
has_work: bool = False
|
||
|
||
|
||
class _BoundarySnapshotProvider:
|
||
"""Dict-like loader for extracted boundary snapshots.
|
||
|
||
Used by ``store_cache()`` to load snapshots from SSD one block at a time
|
||
or serve pre-extracted in-memory snapshots. In-memory snapshots must already
|
||
be in the ``_extract_cache_states`` dict format so this provider can be used
|
||
safely from the async store-cache worker without touching raw MLX cache
|
||
objects on the wrong thread.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
store: Any, # Optional[BoundarySnapshotSSDStore]
|
||
request_id: str,
|
||
valid_tcs: list[int],
|
||
in_memory_snapshots: dict[int, Any],
|
||
) -> None:
|
||
self._store = store
|
||
self._request_id = request_id
|
||
self._valid_tcs = set(valid_tcs)
|
||
self._in_memory = in_memory_snapshots
|
||
|
||
def __contains__(self, tc: int) -> bool:
|
||
return tc in self._valid_tcs
|
||
|
||
def __getitem__(self, tc: int) -> Any:
|
||
snap = self._in_memory.get(tc)
|
||
if snap is not None:
|
||
return snap
|
||
if self._store is not None:
|
||
return self._store.load(self._request_id, tc)
|
||
return None
|
||
|
||
def __len__(self) -> int:
|
||
return len(self._valid_tcs)
|
||
|
||
def __bool__(self) -> bool:
|
||
return bool(self._valid_tcs)
|
||
|
||
def iter_in_memory_extracted(self):
|
||
"""Yield pre-extracted in-memory snapshots for pre-evaluation."""
|
||
for tc in sorted(self._valid_tcs):
|
||
snap = self._in_memory.get(tc)
|
||
if snap is not None:
|
||
yield snap
|
||
|
||
|
||
class Scheduler:
|
||
"""
|
||
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
|
||
|
||
.. note::
|
||
|
||
``_DEFERRED_CLEAR_DELAY`` controls how many generation steps to wait
|
||
after the last request completion before calling ``mx.clear_cache()``.
|
||
Immediate clearing races with IOKit's asynchronous ``completeMemory()``
|
||
callbacks, causing 'prepare count underflow' kernel panics (#435).
|
||
8 steps (~10-40 ms at typical generation speeds) gives IOKit ample
|
||
time to process those callbacks while still reclaiming Metal buffers
|
||
fast enough to prevent TTFT spikes (#411).
|
||
|
||
The key insight is that mlx-lm's BatchGenerator already implements
|
||
continuous batching at the token level, so we use it as the backend.
|
||
"""
|
||
|
||
_DEFERRED_CLEAR_DELAY: int = 8
|
||
_GENERATION_OVERFLOW_PATTERN = "__next_prime overflow"
|
||
_MAX_GENERATION_OVERFLOW_RETRIES = 1
|
||
|
||
def __init__(
|
||
self,
|
||
model: Any,
|
||
tokenizer: Any,
|
||
config: SchedulerConfig | None = None,
|
||
stream: Any | None = None,
|
||
):
|
||
"""
|
||
Initialize the scheduler.
|
||
|
||
Args:
|
||
model: The MLX model
|
||
tokenizer: The tokenizer
|
||
config: Scheduler configuration
|
||
stream: Optional mx.Stream for this engine. Falls back to the
|
||
module-level _default_generation_stream when not provided.
|
||
"""
|
||
self.model = model
|
||
# Deep-copy the tokenizer so the scheduler owns an independent Rust
|
||
# tokenizer backend. Without this, concurrent access from the asyncio
|
||
# event loop (encode/apply_chat_template in engine handlers) and the
|
||
# MLX executor thread (scheduler.step) causes
|
||
# "RuntimeError: Already borrowed" from the HuggingFace tokenizers
|
||
# Rust RefCell. See: https://github.com/huggingface/tokenizers/issues/537
|
||
self.tokenizer = copy.deepcopy(tokenizer)
|
||
self.config = copy.copy(config) if config else SchedulerConfig()
|
||
self._stream = stream if stream is not None else _default_generation_stream
|
||
self._serialize_llama4_requests = _model_declares_llama4(model)
|
||
if self._serialize_llama4_requests and self.config.max_num_seqs > 1:
|
||
logger.info(
|
||
"Llama 4 detected; serializing requests because ChunkedKVCache "
|
||
"does not support multi-row batching yet"
|
||
)
|
||
|
||
# Load additional EOS tokens from generation_config.json.
|
||
# Some models (e.g. GLM-4.6V) define multiple EOS tokens there
|
||
# that are not in tokenizer.eos_token_id.
|
||
self._generation_config_eos: set[int] | None = (
|
||
self._load_generation_config_eos()
|
||
)
|
||
|
||
# Load generation_config.suppress_tokens once and apply them on every
|
||
# sampling path. Gemma 4 uses this to suppress multimodal close markers.
|
||
self._model_suppress_tokens: set[int] = self._load_model_suppress_tokens()
|
||
|
||
# For strict RotatingKVCache reuse, align paged cache block size to
|
||
# the model's rotating window size when paged cache is enabled.
|
||
self._align_block_size_with_rotating_window()
|
||
# For ArraysCache-only models (no RotatingKVCache), use a larger block
|
||
# size to reduce boundary snapshot overhead during prefill.
|
||
self._enlarge_block_size_for_arrays_cache()
|
||
|
||
# TurboQuant KV cache (set by engine if model_settings has it enabled)
|
||
self._turboquant_kv_bits: float | None = None
|
||
self._turboquant_skip_last: bool = True
|
||
# Memoized MLA-architecture detection (see _model_uses_mla / #1613).
|
||
self._mla_model: bool | None = None
|
||
self._glm_dsa_adaptive_prefill = None
|
||
try:
|
||
from .patches.glm_moe_dsa.generate_patch import (
|
||
_glm_dsa_adaptive_prefill_config,
|
||
)
|
||
|
||
self._glm_dsa_adaptive_prefill = _glm_dsa_adaptive_prefill_config(
|
||
model, self.config.prefill_step_size
|
||
)
|
||
except Exception:
|
||
logger.debug("GLM DSA adaptive prefill config unavailable", exc_info=True)
|
||
if self._glm_dsa_adaptive_prefill is not None:
|
||
logger.info(
|
||
"GLM DSA adaptive scheduler prefill enabled: step=%d after=%d "
|
||
"min_remaining=%d",
|
||
self._glm_dsa_adaptive_prefill.step_size,
|
||
self._glm_dsa_adaptive_prefill.after,
|
||
self._glm_dsa_adaptive_prefill.min_remaining,
|
||
)
|
||
self._minimax_m3_adaptive_prefill = None
|
||
try:
|
||
from .patches.minimax_m3.generate_patch import (
|
||
_minimax_m3_adaptive_prefill_config,
|
||
)
|
||
|
||
self._minimax_m3_adaptive_prefill = _minimax_m3_adaptive_prefill_config(
|
||
model,
|
||
self.config.prefill_step_size,
|
||
getattr(self.config, "model_name", None),
|
||
)
|
||
except Exception:
|
||
logger.debug(
|
||
"MiniMax M3 adaptive prefill config unavailable", exc_info=True
|
||
)
|
||
if self._minimax_m3_adaptive_prefill is not None:
|
||
logger.info(
|
||
"MiniMax M3 adaptive scheduler prefill enabled: step=%d after=%d "
|
||
"min_remaining=%d",
|
||
self._minimax_m3_adaptive_prefill.step_size,
|
||
self._minimax_m3_adaptive_prefill.after,
|
||
self._minimax_m3_adaptive_prefill.min_remaining,
|
||
)
|
||
|
||
# Request management - following vLLM's design
|
||
self.waiting: deque[Request] = deque() # Waiting queue (FCFS)
|
||
self.running: dict[str, Request] = {} # Running requests by ID
|
||
# Chunked prefill queue: requests whose prefill spans multiple steps.
|
||
# Populated when chunked_prefill=True and prompt exceeds prefill_step_size.
|
||
self.prefilling: deque[Request] = deque()
|
||
self._prefill_states: dict[str, _PrefillState] = {}
|
||
self.requests: dict[str, Request] = {} # All requests by ID
|
||
self.finished_req_ids: set[str] = set() # Recently finished
|
||
self._generation_overflow_recovery_ids: set[str] = set()
|
||
|
||
# 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()
|
||
|
||
# Deferred between-turn Metal reclaim, requested by the (asyncio-thread)
|
||
# ProcessMemoryEnforcer under pinned-model memory pressure. A bare bool
|
||
# is GIL-atomic to set; it is drained on the inference thread at the top
|
||
# of step() (same cross-thread idiom as _pending_abort_ids) because the
|
||
# enforcer must never touch Metal directly.
|
||
self._pending_reclaim_request: bool = False
|
||
|
||
# Lock-free admin snapshot. Published at the end of each step() while
|
||
# the engine thread is the sole writer of running/waiting; the admin
|
||
# endpoint reads the dict reference atomically (GIL) and never iterates
|
||
# the live mutable structures.
|
||
self._admin_snapshot: dict[str, Any] = {
|
||
"running_by_id": {},
|
||
"waiting": [],
|
||
}
|
||
|
||
# Memory limits for inline prefill checking.
|
||
# Set by ProcessMemoryEnforcer; propagated to BatchGenerator.
|
||
self._memory_limit_bytes: int = 0 # soft limit (dynamic, jittery)
|
||
self._memory_hard_limit_bytes: int = 0 # dynamic ceiling (throttle target)
|
||
# Stable physical cap = min(static_ceiling, metal_cap). Used ONLY to
|
||
# abort an in-flight prefill, so a transient dynamic-ceiling dip can't
|
||
# kill a near-complete request that actually fits. 0 => fall back to
|
||
# _memory_hard_limit_bytes (pre-propagation / old enforcer).
|
||
self._memory_abort_limit_bytes: int = 0
|
||
# Last mx.get_active_memory() sample taken on this scheduler's MLX
|
||
# executor thread. The background memory enforcer reads this cached
|
||
# value during active decode instead of touching MLX/Metal directly.
|
||
self._last_mlx_active_memory_bytes: int = 0
|
||
# Component ceilings — propagated alongside the hard limit so the
|
||
# rejection-path error message can identify which constraint is
|
||
# binding and suggest the right remedy (close apps / raise tier /
|
||
# raise iogpu.wired_limit_mb / reduce context). 0 = not set yet.
|
||
self._memory_static_ceiling_bytes: int = 0
|
||
self._memory_dynamic_ceiling_bytes: int = 0
|
||
self._memory_metal_cap_bytes: int = 0
|
||
self._memory_hot_cache_reserved_bytes: int = 0
|
||
# Tier name propagated alongside the breakdown. For ``custom`` the
|
||
# "dynamic" ceiling is the user-pinned ``custom_ceiling_bytes``
|
||
# rather than computed reclaimable memory, so the advice ladder
|
||
# must steer the user to that knob instead of "close other apps".
|
||
self._memory_guard_tier: str = "balanced"
|
||
self._prefill_memory_guard: bool = False # set by ProcessMemoryEnforcer
|
||
# Set to True by ProcessMemoryEnforcer when phys_footprint crosses
|
||
# soft_threshold. Schedulers stop admitting new prefills while this is
|
||
# set; in-flight requests proceed.
|
||
self._admission_paused: bool = False
|
||
# Adaptive prefill throttle params, propagated from enforcer.
|
||
# Until set, _adaptive_chunk_size is a no-op (returns requested as-is).
|
||
self._prefill_headroom_safety: float = self._PREFILL_HEADROOM_SAFETY
|
||
self._prefill_safe_zone_ratio: float = 0.80
|
||
self._prefill_min_chunk_tokens: int = 256
|
||
self._prefill_abort_margin: float = self._PREFILL_ABORT_MARGIN
|
||
self._pending_prefill_eviction_request: PrefillEvictionRequest | None = None
|
||
self._memory_admission_blocked_request_id: str | None = None
|
||
self._memory_admission_blocked_since: float = 0.0
|
||
self._store_cache_admission_blocked_request_id: str | None = None
|
||
self._store_cache_admission_blocked_since: float = 0.0
|
||
# EWMA estimator of per-token chunk transient bytes, used by
|
||
# _adaptive_chunk_size in the caution zone. Owned per-scheduler.
|
||
_tracker_model_id = ""
|
||
if config is not None and config.model_name:
|
||
_tracker_model_id = config.model_name
|
||
self._prefill_transient_tracker = PrefillTransientTracker(
|
||
model_id=_tracker_model_id
|
||
)
|
||
# Let the sdpa256 head_dim-256 prefill route ask for live guard
|
||
# headroom so it only takes the slow O(L) tiled pass when the faster
|
||
# unfused fallback would not fit (issue #2204). Weakly held; harmless
|
||
# when the patch is never applied.
|
||
set_unfused_headroom_provider(self._sdpa256_unfused_headroom)
|
||
|
||
# SpecPrefill: draft model for attention-based sparse prefill
|
||
self._specprefill_draft_model: Any | None = None
|
||
self._draft_paged_ssd_cache_manager: Any | None = None
|
||
# Track active specprefill request for RoPE cleanup
|
||
self._specprefill_active_request_id: str | None = None
|
||
|
||
# DEBUG-only prefix-cache divergence probe (issue #1003): recent
|
||
# stored cache sequences, so a miss can be traced to the exact
|
||
# token where the new prompt diverges from what was cached.
|
||
# Populated only when debug logging is enabled — zero cost otherwise.
|
||
self._cache_probe_seqs: deque[tuple[str, list[int]]] = deque(maxlen=4)
|
||
|
||
model_name_lower = (self.config.model_name or "").lower()
|
||
default_kv_eval_interval = 256 if "minimax" in model_name_lower else 0
|
||
self._decode_eval_kv_cache_interval: int = max(
|
||
0,
|
||
_env_int(
|
||
"OMLX_DECODE_EVAL_KV_CACHE_INTERVAL",
|
||
default_kv_eval_interval,
|
||
),
|
||
)
|
||
self._tokens_since_kv_cache_eval: int = 0
|
||
if self._decode_eval_kv_cache_interval > 0:
|
||
logger.info(
|
||
"Decode KV cache materialization interval set to %d tokens",
|
||
self._decode_eval_kv_cache_interval,
|
||
)
|
||
|
||
# VLM MTP: gemma4_assistant drafter attached by VLMBatchedEngine.
|
||
# When set, eligible requests bypass mlx-lm BatchGenerator for decode
|
||
# and run through mlx-vlm's _mtp_rounds round loop instead.
|
||
self._vlm_mtp_drafter: VLMMTPDrafter | None = None
|
||
# Active vlm_mtp decode generators keyed by synthesized negative uid
|
||
# (negative to make collision with BatchGenerator uids impossible).
|
||
self._vlm_mtp_active: dict[int, _VLMMTPDecodeState] = {}
|
||
self._vlm_mtp_next_uid: int = -1
|
||
# Per-request settings snapshot for vlm_mtp routing (block size etc.).
|
||
# Injected by VLMBatchedEngine.set_vlm_mtp_drafter alongside the drafter.
|
||
self._vlm_mtp_draft_block_size: int | None = None
|
||
|
||
# Phase timing instrumentation for cache-on overhead diagnostics.
|
||
# Accumulated wall-time per phase + invocation count, dumped at request
|
||
# end or via get_phase_stats(). Adds ~100ns per measurement.
|
||
self._phase_total_ms: dict[str, float] = defaultdict(float)
|
||
self._phase_count: dict[str, int] = defaultdict(int)
|
||
|
||
# Async store_cache executor (G2-async). Offloads the post-finish
|
||
# bulk memcpy (28GB+ per 32k request) off the inference thread so
|
||
# response streaming isn't blocked by it.
|
||
self._store_cache_executor: concurrent.futures.ThreadPoolExecutor | None = None
|
||
# Gate that caps in-flight store-cache submissions. Set only when
|
||
# tiered cache is enabled (alongside _store_cache_executor).
|
||
self._store_cache_gate: _StoreCacheGate | None = None
|
||
# Pending (uid, request_id, future) entries waiting for async store
|
||
# to finish before batch_generator.remove() can safely run. Drained
|
||
# at the start of every step.
|
||
self._pending_async_removes: deque = deque()
|
||
# Track in-flight store futures per request_id for lookup wait /
|
||
# shutdown wait.
|
||
self._inflight_store_futures: dict[str, concurrent.futures.Future] = {}
|
||
self._inflight_store_info: dict[str, _InflightStoreInfo] = {}
|
||
# Admission-only cache freshness waits. A waiting request can pause at
|
||
# the front of the queue for a relevant in-flight store without
|
||
# blocking the scheduler step that continues existing decode/prefill.
|
||
self._cache_freshness_waits: dict[str, _CacheFreshnessWait] = {}
|
||
self._prefix_cache_prepared: set[str] = set()
|
||
|
||
# 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: BatchGenerator | None = None
|
||
self._current_sampler_params: tuple | None = None
|
||
# Boundary cache snapshots for stateful non-sliceable caches (e.g., ArraysCache).
|
||
# request_id -> {token_count -> snapshot_cache_or_None}
|
||
# Multiple snapshots per request to support per-block ArraysCache state storage.
|
||
# Values are None when offloaded to SSD via _boundary_snapshot_store.
|
||
self._boundary_cache_snapshots: dict[str, dict[int, Any]] = {}
|
||
# Lazy detection flag: True/False once determined, None before first check.
|
||
self._boundary_snapshot_required: bool | None = None
|
||
# SSD store for offloading boundary snapshots (initialized in _init_tiered_cache).
|
||
self._boundary_snapshot_store: BoundarySnapshotSSDStore | None = None
|
||
|
||
# paged SSD cache for KV state persistence (oMLX only supports paged SSD-based caching)
|
||
self.paged_cache_manager: PagedCacheManager | None = None
|
||
self.block_aware_cache: BlockAwarePrefixCache | None = None
|
||
self.paged_ssd_cache_manager: PagedSSDCacheManager | None = None
|
||
self._cache_rate_tracker = CacheRateTracker()
|
||
# Prefill-peak estimator used by ``_preflight_memory_check`` /
|
||
# ``preflight_or_raise``. Only the estimator path is exercised
|
||
# here (it reads head_dim / num_layers / num_kv_heads via
|
||
# ``set_model_info`` below); ``eviction_enabled=False`` so the
|
||
# monitor does not gate on ``max_kv_cache_memory`` — paged SSD
|
||
# mode never wants this monitor making eviction decisions, and
|
||
# we have no real value to pass for that field at this point.
|
||
#
|
||
# This auto-init was wired up in b6a69c4 then silently dropped
|
||
# by an upstream merge (same pattern as ``preflight_or_raise``
|
||
# in d40ab80). Without it the guard short-circuits at the
|
||
# ``memory_monitor is None`` gate for every request — Pi prompts
|
||
# that should be rejected by the configured hard limit instead
|
||
# sail straight into chunked prefill and OOM at the Metal cap.
|
||
if MemoryMonitor is not None:
|
||
self.memory_monitor: MemoryMonitor | None = MemoryMonitor(
|
||
max_kv_cache_memory=None,
|
||
eviction_enabled=False,
|
||
)
|
||
self._set_model_info_for_monitor()
|
||
else:
|
||
self.memory_monitor: MemoryMonitor | None = None
|
||
|
||
# Initialize paged SSD cache if paged_ssd_cache_dir is specified
|
||
if self.config.paged_ssd_cache_dir:
|
||
# Calculate max_blocks automatically if not specified
|
||
if self.config.max_cache_blocks is not None:
|
||
max_blocks = self.config.max_cache_blocks
|
||
else:
|
||
max_blocks = self._calculate_max_blocks()
|
||
|
||
# Initialize paged cache manager for block metadata
|
||
self.paged_cache_manager = PagedCacheManager(
|
||
block_size=self.config.paged_cache_block_size,
|
||
max_blocks=max_blocks,
|
||
model_name=self.config.model_name,
|
||
initial_blocks=self.config.initial_cache_blocks,
|
||
)
|
||
self.block_aware_cache = BlockAwarePrefixCache(
|
||
model=model,
|
||
paged_cache_manager=self.paged_cache_manager,
|
||
)
|
||
|
||
# Initialize paged SSD cache. If the backing directory is not
|
||
# usable (for example, an external cache drive is disconnected),
|
||
# continue with cache disabled instead of leaving partial state.
|
||
cache_initialized = self._init_tiered_cache()
|
||
|
||
# Set cold restore callback for prefix cache
|
||
if cache_initialized and self.paged_ssd_cache_manager is not None:
|
||
self.block_aware_cache.set_cold_restore_callback(
|
||
self._restore_block_from_cold
|
||
)
|
||
if self.config.hot_cache_only:
|
||
logger.info(
|
||
f"hot-cache-only mode enabled: "
|
||
f"block_size={self.config.paged_cache_block_size}, "
|
||
f"max_blocks={max_blocks}"
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"paged SSD cache enabled: {self.config.paged_ssd_cache_dir}, "
|
||
f"block_size={self.config.paged_cache_block_size}, "
|
||
f"max_blocks={max_blocks}"
|
||
)
|
||
|
||
# Async store_cache executor: single worker so submissions are
|
||
# serialized (matches the original synchronous order) and we
|
||
# never have two stores racing on the same paged_ssd index.
|
||
self._store_cache_executor = concurrent.futures.ThreadPoolExecutor(
|
||
max_workers=1,
|
||
thread_name_prefix="omlx-store-cache",
|
||
)
|
||
# Gate caps the post-completion store-cache pipeline so a burst
|
||
# of finishes cannot pile up unbounded KV caches in memory while
|
||
# the single writer drains. Cap starts at max_concurrent_requests
|
||
# and is shrunk by ProcessMemoryEnforcer under pressure (#1383).
|
||
self._store_cache_gate = _StoreCacheGate(cap=self.config.max_num_seqs)
|
||
else:
|
||
self._disable_paged_cache_components()
|
||
logger.info(
|
||
"oMLX cache disabled after paged SSD cache initialization failed"
|
||
)
|
||
else:
|
||
logger.info(
|
||
"oMLX cache disabled (mlx-lm BatchGenerator manages KV internally)"
|
||
)
|
||
|
||
# Streaming detokenizers for proper UTF-8 handling (one per active request)
|
||
# NOTE: No pooling - each request gets a fresh instance to prevent state contamination
|
||
self._request_detokenizers: dict[str, Any] = (
|
||
{}
|
||
) # request_id → active detokenizer
|
||
|
||
# Protocol-specific output parser support (e.g. Harmony, Gemma 4)
|
||
self._output_parser_factory: OutputParserFactory | None = None
|
||
self._output_parser_kind: str | None = None
|
||
self._output_parser_sessions: dict[str, OutputParserSession] = {}
|
||
self._is_harmony_model: bool = False
|
||
if HAS_OUTPUT_PARSER and detect_output_parser is not None:
|
||
try:
|
||
model_config = None
|
||
if hasattr(model, "config"):
|
||
# model.config may be a Pydantic model or dict
|
||
try:
|
||
if hasattr(model.config, "model_dump"):
|
||
model_config = model.config.model_dump()
|
||
elif hasattr(model.config, "dict"):
|
||
model_config = model.config.dict()
|
||
elif isinstance(model.config, dict):
|
||
model_config = model.config
|
||
else:
|
||
# Try to convert to dict via __dict__
|
||
model_config = getattr(model.config, "__dict__", None)
|
||
except Exception as e:
|
||
logger.debug(f"Failed to extract model.config: {e}")
|
||
elif hasattr(model, "args"):
|
||
try:
|
||
if hasattr(model.args, "model_dump"):
|
||
model_config = model.args.model_dump()
|
||
elif hasattr(model.args, "__dict__"):
|
||
model_config = model.args.__dict__
|
||
except Exception as e:
|
||
logger.debug(f"Failed to extract model.args: {e}")
|
||
|
||
self._output_parser_factory = detect_output_parser(
|
||
self.config.model_name,
|
||
self.tokenizer,
|
||
model_config,
|
||
model_path=self.config.model_path,
|
||
)
|
||
if self._output_parser_factory is not None:
|
||
self._output_parser_kind = self._output_parser_factory.kind
|
||
self._is_harmony_model = self._output_parser_kind == "harmony"
|
||
logger.info(
|
||
"Output parser detected: %s for %s, stop_tokens=%s",
|
||
self._output_parser_kind,
|
||
self.config.model_name,
|
||
sorted(self._output_parser_factory.stop_token_ids),
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Error detecting output parser: {e}, assuming none")
|
||
self._output_parser_factory = None
|
||
self._output_parser_kind = None
|
||
self._is_harmony_model = False
|
||
|
||
# Statistics
|
||
self.num_requests_processed = 0
|
||
self.total_prompt_tokens = 0
|
||
self.total_completion_tokens = 0
|
||
|
||
# Step counter for periodic cleanup
|
||
self._step_counter = 0
|
||
# Deferred Metal cache cleanup after request completion.
|
||
# Immediate mx.clear_cache() after request completion races with
|
||
# IOKit's asynchronous completeMemory() callbacks, causing
|
||
# 'prepare count underflow' kernel panics. Deferring the clear
|
||
# by a few generation steps gives IOKit time to process callbacks.
|
||
#
|
||
# Stored as the absolute step number at which the clear should fire,
|
||
# rather than a countdown integer. This avoids the burst-completion
|
||
# bug (#557): with max_num_seqs > 1 two requests can finish in the
|
||
# same batch. The old "only set if None" guard meant the second
|
||
# completion never extended the window, so the first request's KV
|
||
# cache blocks could be re-allocated before IOKit finished its
|
||
# completeMemory() callbacks. Using max() ensures the window always
|
||
# covers the *latest* completion.
|
||
# None = no deferred clear pending; int = step at which to fire.
|
||
self._deferred_clear_at: int | None = None
|
||
|
||
# Cache XTC special tokens (newline + EOS) — stable per tokenizer.
|
||
# Must be after _is_harmony_model / _generation_config_eos init
|
||
# since _get_xtc_special_tokens() delegates to _get_stop_tokens().
|
||
self._xtc_special_tokens: list[int] = self._get_xtc_special_tokens()
|
||
|
||
# Drop transient aliases after ownership moves to scheduler fields;
|
||
# close()/deep_reset() clear those fields during teardown.
|
||
model = None
|
||
tokenizer = None
|
||
|
||
@contextmanager
|
||
def _phase_timer(self, phase: str):
|
||
"""Lightweight wall-time accumulator for cache-on overhead diagnostics.
|
||
|
||
Tracks total ms and invocation count per named phase. Intended for
|
||
boundary capture / store_cache / hot cache eviction hot paths.
|
||
"""
|
||
t0 = time.perf_counter()
|
||
try:
|
||
yield
|
||
finally:
|
||
self._phase_total_ms[phase] += (time.perf_counter() - t0) * 1000.0
|
||
self._phase_count[phase] += 1
|
||
|
||
def get_phase_stats(self) -> dict[str, dict[str, float]]:
|
||
"""Return accumulated phase timings for diagnostics.
|
||
|
||
Returns dict of phase -> {total_ms, count, avg_ms}.
|
||
"""
|
||
result = {}
|
||
for phase, total in self._phase_total_ms.items():
|
||
count = self._phase_count.get(phase, 0)
|
||
result[phase] = {
|
||
"total_ms": total,
|
||
"count": count,
|
||
"avg_ms": total / count if count else 0.0,
|
||
}
|
||
return result
|
||
|
||
def _periodic_clear_threshold_bytes(self) -> int:
|
||
"""Cache-bytes threshold above which the periodic clear runs.
|
||
|
||
Defaults to memory_limit/3 when a process memory limit is set,
|
||
otherwise an absolute 2 GiB floor. Each periodic clear releases
|
||
the entire MLX buffer pool in one batch; gating it on accumulated
|
||
bytes avoids producing IOGPUFamily refcount bursts when the pool
|
||
is already small.
|
||
"""
|
||
if self._memory_limit_bytes > 0:
|
||
return max(self._memory_limit_bytes // 3, 2 * 1024**3)
|
||
return 2 * 1024**3
|
||
|
||
def _should_periodic_clear_cache(self) -> bool:
|
||
"""Decide whether the per-step periodic clear should fire.
|
||
|
||
Returns False unless ``mlx_cache_cleanup_interval`` is configured,
|
||
the step counter just landed on the interval boundary, AND the
|
||
MLX buffer pool exceeds the threshold. See #978 / #1040 for the
|
||
kernel panic class this gating is meant to mitigate.
|
||
"""
|
||
interval = self.config.mlx_cache_cleanup_interval
|
||
if interval <= 0 or self._step_counter % interval != 0:
|
||
return False
|
||
return mx.get_cache_memory() > self._periodic_clear_threshold_bytes()
|
||
|
||
@staticmethod
|
||
def _collect_arrays_from_extracted_cache(
|
||
extracted_cache: list[Any],
|
||
) -> list[Any]:
|
||
"""Collect lazy mx.array references from an _extracted_cache payload.
|
||
|
||
Used by G2-async to force a single batched mx.eval on the inference
|
||
thread before handing the cache off to the store_cache worker. The
|
||
worker can then call _extract_tensor_bytes safely (no further Metal
|
||
graph evaluation needed for non-bfloat16, no-op for already-evaluated).
|
||
|
||
Walks the per-layer dict format produced by _extract_cache_states:
|
||
each layer is {state, meta_state, class_name, cache_type}, where
|
||
state is a tuple of mx.arrays (or nested for CacheList / TurboQuant).
|
||
"""
|
||
arrays: list[Any] = []
|
||
for layer in extracted_cache or []:
|
||
if not isinstance(layer, dict):
|
||
continue
|
||
state = layer.get("state", ())
|
||
if isinstance(state, mx.array):
|
||
arrays.append(state)
|
||
continue
|
||
if not isinstance(state, (list, tuple)):
|
||
continue
|
||
for item in state:
|
||
if isinstance(item, mx.array):
|
||
arrays.append(item)
|
||
elif isinstance(item, (list, tuple)):
|
||
for sub in item:
|
||
if isinstance(sub, mx.array):
|
||
arrays.append(sub)
|
||
elif hasattr(sub, "_fields"):
|
||
# NamedTuple state (TurboQuant). Walk fields.
|
||
for fname in sub._fields:
|
||
val = getattr(sub, fname, None)
|
||
if isinstance(val, mx.array):
|
||
arrays.append(val)
|
||
elif hasattr(item, "_fields"):
|
||
for fname in item._fields:
|
||
val = getattr(item, fname, None)
|
||
if isinstance(val, mx.array):
|
||
arrays.append(val)
|
||
return arrays
|
||
|
||
def _async_store_cache_worker(
|
||
self,
|
||
request_id: str,
|
||
token_sequence_to_store: list[int],
|
||
cache_to_store: list[Any],
|
||
model_cache_config: Any | None,
|
||
intermediate_snapshots: dict[int, list[Any]] | None,
|
||
extra_keys: tuple[Any, ...] | None,
|
||
extra_key_token_start: int | None,
|
||
extra_key_ranges: list[tuple[int, tuple[Any, ...]]] | None,
|
||
hot_cache_write_back: bool = True,
|
||
) -> None:
|
||
"""Run store_cache + paged_cache cleanup off the inference thread.
|
||
|
||
Pre-conditions enforced by the caller (_cleanup_finished):
|
||
- mx.eval() (a FULL, blocking eval — NOT async_eval) was called on
|
||
the inference thread for all KV cache arrays in cache_to_store, so
|
||
they are fully materialized concrete buffers before this worker
|
||
runs. This is the load-bearing invariant: MLX streams ARE
|
||
thread-local (each engine's generation stream is created via
|
||
mx.new_thread_local_stream on the inference thread). A KV array
|
||
left lazy and bound to self._stream cannot be materialized on this
|
||
worker thread — _extract_block_tensor_slice (slices) and
|
||
_extract_tensor_bytes (bf16 -> uint16 view) would re-dispatch the
|
||
source op to self._stream's index, which does NOT exist on this
|
||
thread, aborting the process ("There is no Stream(gpu, N) in
|
||
current thread"). Boundary snapshots get the same treatment via
|
||
_eval_snapshot_cache at capture time.
|
||
- Because the sources are concrete, the worker's own slice/view ops
|
||
consume materialized buffers and bind their new ops to the
|
||
always-present default stream (gpu,0). The _safe_sync_stream call
|
||
below is now belt-and-suspenders (the owner already drained the
|
||
work); correctness no longer depends on its cross-thread
|
||
mx.synchronize, whose "no Stream" RuntimeError it tolerantly
|
||
swallows.
|
||
- All mx-buffer access here is held under _mx_buffer_access_lock,
|
||
serializing the bf16 view+eval and the buffer-protocol reads
|
||
against inference-thread _sync_and_clear_cache (which also takes
|
||
that lock), so mx.clear_cache cannot reclaim a buffer mid-read.
|
||
- batch_generator.remove(uid) is deferred until this worker
|
||
completes (handled by _drain_pending_async_removes).
|
||
|
||
paged_cache_manager and block_aware_cache rely on
|
||
threading.RLock so concurrent access from main and worker is safe.
|
||
"""
|
||
try:
|
||
# Hold _mx_buffer_access_lock across the worker's mx-buffer
|
||
# access. store_cache eventually drives _extract_tensor_bytes,
|
||
# which reads raw bytes via the buffer protocol; serializing
|
||
# against inference-thread mx.clear_cache / mx.synchronize calls
|
||
# prevents a SIGABRT when those reclaim the underlying Metal
|
||
# buffer pool mid-read (#1106).
|
||
with _mx_buffer_access_lock:
|
||
with self._phase_timer("store_cache_worker_sync"):
|
||
_safe_sync_stream(self._stream)
|
||
if hot_cache_write_back:
|
||
block_table = self.block_aware_cache.store_cache(
|
||
request_id,
|
||
token_sequence_to_store,
|
||
cache_to_store,
|
||
model_cache_config=model_cache_config,
|
||
boundary_snapshots=intermediate_snapshots,
|
||
extra_keys=extra_keys,
|
||
extra_key_token_start=extra_key_token_start,
|
||
extra_key_ranges=extra_key_ranges,
|
||
)
|
||
else:
|
||
block_table = self.block_aware_cache.store_cache(
|
||
request_id,
|
||
token_sequence_to_store,
|
||
cache_to_store,
|
||
model_cache_config=model_cache_config,
|
||
boundary_snapshots=intermediate_snapshots,
|
||
extra_keys=extra_keys,
|
||
extra_key_token_start=extra_key_token_start,
|
||
extra_key_ranges=extra_key_ranges,
|
||
hot_cache_write_back=False,
|
||
)
|
||
if block_table is None and self.paged_cache_manager is not None:
|
||
block_table = self.paged_cache_manager.get_block_table(request_id)
|
||
if block_table and self.paged_cache_manager is not None:
|
||
self.paged_cache_manager.release_for_eviction(block_table.block_ids)
|
||
if self.block_aware_cache is not None:
|
||
self.block_aware_cache.clear_request_entry(request_id)
|
||
except Exception as e:
|
||
logger.warning("Async store_cache failed for %s: %s", request_id, e)
|
||
|
||
def _drain_pending_async_removes(self) -> bool:
|
||
"""Process deferred batch_generator.remove() calls from prior steps.
|
||
|
||
Called at the start of every step. For each pending entry whose async
|
||
store_cache future has finished, perform batch_generator.remove() on
|
||
the inference thread (Metal-safe) and finalize cleanup state. Entries
|
||
whose futures are still in flight are kept for a later step, but they
|
||
do not block later completed entries from releasing cache references.
|
||
"""
|
||
if not self._pending_async_removes:
|
||
return False
|
||
drained = False
|
||
pending: deque = deque()
|
||
while self._pending_async_removes:
|
||
uid, request_id, future = self._pending_async_removes.popleft()
|
||
if future is not None and not future.done():
|
||
# Worker still busy. Keep it for the next step, but continue
|
||
# scanning so later completed futures can release memory now.
|
||
pending.append((uid, request_id, future))
|
||
continue
|
||
# Surface worker exceptions for visibility (don't crash step loop).
|
||
if future is not None:
|
||
try:
|
||
exc = future.exception()
|
||
except concurrent.futures.CancelledError:
|
||
logger.warning("Async store_cache for %s was cancelled", request_id)
|
||
else:
|
||
if exc is not None:
|
||
logger.warning(
|
||
"Async store_cache for %s raised: %s", request_id, exc
|
||
)
|
||
try:
|
||
# Run batch_generator.remove on the inference thread.
|
||
try:
|
||
_safe_sync_stream(self._stream)
|
||
self._remove_uid_from_active_batch(uid)
|
||
if hasattr(self.model, "unregister_rope_delta"):
|
||
self.model.unregister_rope_delta(uid)
|
||
except Exception as e:
|
||
logger.warning(
|
||
"Deferred batch_generator.remove(uid=%s) failed: %s",
|
||
uid,
|
||
e,
|
||
)
|
||
# Cleanup uid maps now that the slot is reclaimable.
|
||
_unregister_uid_row(self.model, uid)
|
||
if uid in self.uid_to_request_id:
|
||
del self.uid_to_request_id[uid]
|
||
if request_id in self.request_id_to_uid:
|
||
del self.request_id_to_uid[request_id]
|
||
self._inflight_store_futures.pop(request_id, None)
|
||
self._inflight_store_info.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
# Boundary snapshots were kept on disk for the worker; safe to
|
||
# delete now that the future has completed. Cleanup was
|
||
# deferred from _cleanup_finished to avoid racing the worker's
|
||
# boundary_snapshot_store.load() calls with rmtree.
|
||
if self._boundary_snapshot_store is not None:
|
||
self._boundary_snapshot_store.cleanup_request(request_id)
|
||
# Worker no longer holds extracted_cache — pop request from
|
||
# self.requests and drop the cache buffer references so MLX
|
||
# arrays can be freed.
|
||
req_to_remove = self.requests.pop(request_id, None)
|
||
if req_to_remove is not None:
|
||
req_to_remove._extracted_cache = None
|
||
req_to_remove.prompt_cache = None
|
||
finally:
|
||
gate = self._store_cache_gate
|
||
if gate is not None:
|
||
gate.note_done()
|
||
drained = True
|
||
self._pending_async_removes = pending
|
||
return drained
|
||
|
||
def _calculate_max_blocks(self) -> int:
|
||
"""
|
||
Calculate maximum cache blocks for paged SSD-only mode.
|
||
|
||
In paged SSD-only mode, blocks don't consume GPU memory (data is on paged SSD),
|
||
so we use a large default that can be limited by SSD capacity.
|
||
|
||
Returns:
|
||
Maximum number of cache blocks to allocate.
|
||
"""
|
||
# In paged SSD-only mode, use a large default since blocks don't consume GPU memory
|
||
# The actual limit is SSD capacity (paged_ssd_cache_max_size)
|
||
max_blocks = 100000 # Large default for paged SSD-only mode
|
||
|
||
block_size = self.config.paged_cache_block_size
|
||
logger.info(
|
||
f"paged SSD-only mode: max_blocks={max_blocks}, block_size={block_size} tokens"
|
||
)
|
||
|
||
return max_blocks
|
||
|
||
def _collect_rotating_window_sizes(
|
||
self,
|
||
cache_obj: Any,
|
||
window_sizes: set[int],
|
||
) -> None:
|
||
"""Collect rotating window sizes recursively from cache objects."""
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
for sub_cache in sub_caches:
|
||
self._collect_rotating_window_sizes(sub_cache, window_sizes)
|
||
|
||
class_name = type(cache_obj).__name__
|
||
is_rotating_cache = class_name in ("RotatingKVCache", "BatchRotatingKVCache")
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None:
|
||
is_rotating_cache = (
|
||
is_rotating_cache or CacheTypeRegistry.is_rotating_family(class_name)
|
||
)
|
||
if is_rotating_cache:
|
||
max_size = getattr(cache_obj, "max_size", 0)
|
||
if isinstance(max_size, int) and max_size > 0:
|
||
window_sizes.add(max_size)
|
||
|
||
def _detect_rotating_window_sizes(self) -> set[int]:
|
||
"""Detect rotating window sizes from model.make_cache() if available."""
|
||
if not hasattr(self.model, "make_cache"):
|
||
return set()
|
||
|
||
try:
|
||
cache_list = self.model.make_cache()
|
||
except Exception as e:
|
||
logger.debug(f"Failed to inspect model rotating window sizes: {e}")
|
||
return set()
|
||
|
||
if cache_list is None:
|
||
return set()
|
||
|
||
window_sizes: set[int] = set()
|
||
for cache_obj in cache_list:
|
||
self._collect_rotating_window_sizes(cache_obj, window_sizes)
|
||
|
||
return window_sizes
|
||
|
||
# Target range for RotatingKVCache block size alignment.
|
||
# Using a multiple of window_size within this range reduces SSD I/O
|
||
# overhead (fewer, larger block files) while keeping cache restore
|
||
# reprocessing reasonable.
|
||
_ROTATING_BLOCK_SIZE_MIN = 512
|
||
_ROTATING_BLOCK_SIZE_MAX = 1024
|
||
|
||
def _align_block_size_with_rotating_window(self) -> None:
|
||
"""
|
||
Align paged cache block size to a multiple of RotatingKVCache
|
||
window size, targeting 512-1024 tokens per block.
|
||
|
||
Block size must be a multiple of window_size so that block
|
||
boundaries align with rotation boundaries. When window_size is
|
||
small (e.g. 128), using it directly as block_size creates too
|
||
many small files. Instead we pick the smallest multiple of
|
||
window_size that falls within [_ROTATING_BLOCK_SIZE_MIN,
|
||
_ROTATING_BLOCK_SIZE_MAX].
|
||
"""
|
||
if not self.config.paged_ssd_cache_dir:
|
||
return
|
||
|
||
window_sizes = self._detect_rotating_window_sizes()
|
||
if not window_sizes:
|
||
return
|
||
|
||
if len(window_sizes) > 1:
|
||
raise ValueError(
|
||
"Multiple RotatingKVCache window sizes detected "
|
||
f"({sorted(window_sizes)}). Set a single aligned block size or "
|
||
"disable paged cache for this model."
|
||
)
|
||
|
||
window_size = next(iter(window_sizes))
|
||
|
||
# Find the smallest multiple of window_size >= _ROTATING_BLOCK_SIZE_MIN.
|
||
# If window_size itself is already >= max, just use window_size.
|
||
lo = self._ROTATING_BLOCK_SIZE_MIN
|
||
hi = self._ROTATING_BLOCK_SIZE_MAX
|
||
|
||
if window_size >= hi or window_size >= lo:
|
||
target_block_size = window_size
|
||
else:
|
||
# window_size < lo: pick smallest multiple in [lo, hi]
|
||
multiplier = (lo + window_size - 1) // window_size # ceil(lo / ws)
|
||
target_block_size = multiplier * window_size
|
||
if target_block_size > hi:
|
||
# Fall back to largest multiple <= hi
|
||
target_block_size = (hi // window_size) * window_size
|
||
if target_block_size < window_size:
|
||
target_block_size = window_size
|
||
|
||
if self.config.paged_cache_block_size != target_block_size:
|
||
logger.info(
|
||
"Aligning paged cache block_size=%s to %s "
|
||
"(RotatingKVCache window_size=%s, multiplier=%sx)",
|
||
self.config.paged_cache_block_size,
|
||
target_block_size,
|
||
window_size,
|
||
target_block_size // window_size,
|
||
)
|
||
self.config.paged_cache_block_size = target_block_size
|
||
|
||
# Default block size for ArraysCache-only hybrid models.
|
||
# Match prefill_step_size (2048) so that boundary caching ON/OFF
|
||
# produces identical prefill chunk sizes, eliminating float32↔dtype
|
||
# roundtrip differences in GatedDeltaNet recurrent state.
|
||
_ARRAYS_CACHE_BLOCK_SIZE = 2048
|
||
|
||
def _enlarge_block_size_for_arrays_cache(self) -> None:
|
||
"""Enlarge block size for ArraysCache-only hybrid models.
|
||
|
||
When a model uses ArraysCache (GatedDeltaNet) but not RotatingKVCache,
|
||
a larger block size reduces the number of boundary snapshot stops during
|
||
prefill while still storing valid per-block recurrent state.
|
||
|
||
This is skipped if RotatingKVCache was already detected (block size was
|
||
aligned to its window size) or if the user explicitly set a block size
|
||
larger than the default.
|
||
"""
|
||
if not self.config.paged_ssd_cache_dir:
|
||
return
|
||
|
||
# Skip if RotatingKVCache already adjusted block size.
|
||
rotating_sizes = self._detect_rotating_window_sizes()
|
||
if rotating_sizes:
|
||
return
|
||
|
||
# Detect ArraysCache from model.make_cache()
|
||
if not hasattr(self.model, "make_cache"):
|
||
return
|
||
|
||
try:
|
||
cache_list = self.model.make_cache()
|
||
except Exception:
|
||
return
|
||
|
||
if cache_list is None:
|
||
return
|
||
|
||
has_arrays_cache = any(
|
||
self._cache_tree_has_arrays_cache(cache_obj) for cache_obj in cache_list
|
||
)
|
||
if not has_arrays_cache:
|
||
return
|
||
|
||
target = self._ARRAYS_CACHE_BLOCK_SIZE
|
||
if self.config.paged_cache_block_size >= target:
|
||
return
|
||
|
||
logger.info(
|
||
"Enlarging paged cache block_size=%s to %s for "
|
||
"ArraysCache hybrid model (reduces boundary snapshot overhead)",
|
||
self.config.paged_cache_block_size,
|
||
target,
|
||
)
|
||
self.config.paged_cache_block_size = target
|
||
|
||
@staticmethod
|
||
def _cache_tree_has_arrays_cache(cache_obj: Any) -> bool:
|
||
"""Return True if cache_obj contains ArraysCache (recursively)."""
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
return any(
|
||
Scheduler._cache_tree_has_arrays_cache(sub) for sub in sub_caches
|
||
)
|
||
return type(cache_obj).__name__ in ("ArraysCache", "SizedArraysCache")
|
||
|
||
def _load_generation_config_eos(self) -> set[int] | None:
|
||
"""Load EOS token IDs from generation_config.json if available."""
|
||
try:
|
||
model_ref = getattr(self.tokenizer, "name_or_path", None) or getattr(
|
||
self.config, "model_name", None
|
||
)
|
||
if not model_ref:
|
||
return None
|
||
|
||
result = load_generation_config_token_ids(model_ref, "eos_token_id")
|
||
if result is None:
|
||
return None
|
||
|
||
# Only return if there are tokens beyond what tokenizer already provides
|
||
tokenizer_eos = getattr(self.tokenizer, "eos_token_id", None)
|
||
if tokenizer_eos is not None:
|
||
existing = (
|
||
{tokenizer_eos}
|
||
if isinstance(tokenizer_eos, int)
|
||
else set(tokenizer_eos)
|
||
)
|
||
extra = result - existing
|
||
if extra:
|
||
logger.info(
|
||
f"Loaded {len(extra)} additional EOS token(s) from "
|
||
f"generation_config.json: {extra}"
|
||
)
|
||
return result
|
||
return result
|
||
except Exception as e:
|
||
logger.debug(f"Could not load generation_config.json: {e}")
|
||
return None
|
||
|
||
def _load_model_suppress_tokens(self) -> set[int]:
|
||
"""Load suppress_tokens from generation_config.json if available.
|
||
|
||
These tokens are set to -inf during generation. For Gemma 4 unified,
|
||
generation_config marks the multimodal close markers (<image|>,
|
||
<audio|>) this way.
|
||
"""
|
||
try:
|
||
model_ref = getattr(self.tokenizer, "name_or_path", None) or getattr(
|
||
self.config, "model_name", None
|
||
)
|
||
if not model_ref:
|
||
return set()
|
||
|
||
result = load_generation_config_token_ids(model_ref, "suppress_tokens")
|
||
if not result:
|
||
return set()
|
||
logger.info(
|
||
f"Loaded {len(result)} suppress token(s) from "
|
||
f"generation_config.json: {result}"
|
||
)
|
||
return result
|
||
except Exception as e:
|
||
logger.debug(f"Could not load suppress_tokens from generation_config: {e}")
|
||
return set()
|
||
|
||
def _get_stop_tokens(self) -> set[int]:
|
||
"""Get stop token IDs from tokenizer and generation_config."""
|
||
stop_tokens = set()
|
||
if (
|
||
hasattr(self.tokenizer, "eos_token_id")
|
||
and self.tokenizer.eos_token_id is not None
|
||
):
|
||
if isinstance(self.tokenizer.eos_token_id, list):
|
||
stop_tokens.update(self.tokenizer.eos_token_id)
|
||
else:
|
||
stop_tokens.add(self.tokenizer.eos_token_id)
|
||
if (
|
||
hasattr(self.tokenizer, "eos_token_ids")
|
||
and self.tokenizer.eos_token_ids is not None
|
||
):
|
||
eos_ids = self.tokenizer.eos_token_ids
|
||
if isinstance(eos_ids, int):
|
||
stop_tokens.add(eos_ids)
|
||
else:
|
||
stop_tokens.update(eos_ids)
|
||
|
||
# Include end-of-turn token for models that use turn-based
|
||
# conversation delimiters (e.g. Gemma 4 with <turn|>). Without
|
||
# this the model generates the full next turn after its response.
|
||
eot_token_id = getattr(self.tokenizer, "eot_token_id", None)
|
||
if eot_token_id is not None:
|
||
if isinstance(eot_token_id, list):
|
||
stop_tokens.update(eot_token_id)
|
||
else:
|
||
stop_tokens.add(eot_token_id)
|
||
elif hasattr(self.tokenizer, "eot_token") and self.tokenizer.eot_token:
|
||
# Encode the string value if eot_token_id isn't directly exposed
|
||
try:
|
||
encoded = self.tokenizer.encode(
|
||
self.tokenizer.eot_token, add_special_tokens=False
|
||
)
|
||
if encoded:
|
||
stop_tokens.update(encoded)
|
||
except Exception:
|
||
pass
|
||
|
||
# Read additional EOS tokens from generation_config.json.
|
||
# Some models (e.g. GLM-4.6V) define multiple EOS tokens there
|
||
# that are not reflected in tokenizer.eos_token_id.
|
||
if self._generation_config_eos is not None:
|
||
stop_tokens.update(self._generation_config_eos)
|
||
|
||
# Protocol parsers need to observe their own stop tokens so they can
|
||
# apply channel-aware handling (for example, Harmony analysis end
|
||
# should continue into the final channel).
|
||
|
||
return stop_tokens
|
||
|
||
# _update_stop_tokens deleted — per-request stop tokens are now
|
||
# handled via SequenceStateMachine passed to insert().
|
||
|
||
def _get_detokenizer(self, request_id: str):
|
||
"""Get or create a streaming detokenizer for a request.
|
||
|
||
This enables proper UTF-8 handling for multi-byte characters
|
||
(Korean, Chinese, Japanese, etc.) during streaming.
|
||
|
||
NOTE: Each request gets a fresh detokenizer instance. Pooling was removed
|
||
because internal state (byte buffers) can leak between requests even after
|
||
finalize()/reset(), causing text corruption (e.g., spaces inserted in paths,
|
||
character swaps like 'features' -> 'featurse').
|
||
"""
|
||
if request_id not in self._request_detokenizers:
|
||
# Always create a fresh detokenizer - no pooling to prevent state contamination
|
||
detok = create_streaming_detokenizer(
|
||
self.tokenizer,
|
||
model_path=self.config.model_path,
|
||
)
|
||
if detok is None:
|
||
# Fallback: return None, we'll use decode([token])
|
||
return None
|
||
detok.reset()
|
||
self._request_detokenizers[request_id] = detok
|
||
return self._request_detokenizers[request_id]
|
||
|
||
def _cleanup_detokenizer(self, request_id: str):
|
||
"""Clean up detokenizer for a finished request.
|
||
|
||
NOTE: Detokenizers are NOT pooled - each request gets a fresh instance
|
||
to prevent state contamination that causes text corruption.
|
||
"""
|
||
self._request_detokenizers.pop(request_id, None)
|
||
# Let GC collect - no pooling to prevent state contamination
|
||
|
||
def _get_output_parser_session(
|
||
self, request_id: str
|
||
) -> Optional["OutputParserSession"]:
|
||
"""Get or create a protocol-specific output parser session."""
|
||
if self._output_parser_factory is None:
|
||
return None
|
||
|
||
if request_id not in self._output_parser_sessions:
|
||
self._output_parser_sessions[request_id] = (
|
||
self._output_parser_factory.create_session(self.tokenizer)
|
||
)
|
||
return self._output_parser_sessions[request_id]
|
||
|
||
def _cleanup_output_parser_session(self, request_id: str):
|
||
"""Remove any per-request protocol parser session."""
|
||
self._output_parser_sessions.pop(request_id, None)
|
||
|
||
def _get_xtc_special_tokens(self) -> list[int]:
|
||
"""Get special tokens to exclude from XTC sampling.
|
||
|
||
Parser-owned stop tokens stay out of BatchGenerator stop-token matching
|
||
so protocol parsers can handle them channel-aware, but XTC must still
|
||
protect them from diversity masking.
|
||
"""
|
||
tokens = self.tokenizer.encode("\n")
|
||
tokens.extend(self._get_stop_tokens())
|
||
if self._output_parser_factory is not None:
|
||
tokens.extend(self._output_parser_factory.stop_token_ids)
|
||
return tokens
|
||
|
||
def _create_batch_generator(
|
||
self, sampling_params: SamplingParams
|
||
) -> BatchGenerator:
|
||
"""Create a BatchGenerator with the given sampling parameters."""
|
||
sampler = omlx_make_sampler(
|
||
temp=sampling_params.temperature,
|
||
top_p=sampling_params.top_p,
|
||
min_p=sampling_params.min_p,
|
||
top_k=sampling_params.top_k,
|
||
xtc_probability=sampling_params.xtc_probability,
|
||
xtc_threshold=sampling_params.xtc_threshold,
|
||
xtc_special_tokens=self._xtc_special_tokens,
|
||
)
|
||
|
||
# Create logits processors for repetition/presence/frequency penalties
|
||
logits_processors = make_logits_processors(
|
||
repetition_penalty=(
|
||
sampling_params.repetition_penalty
|
||
if sampling_params.repetition_penalty != 1.0
|
||
else None
|
||
),
|
||
presence_penalty=(
|
||
sampling_params.presence_penalty
|
||
if sampling_params.presence_penalty != 0.0
|
||
else None
|
||
),
|
||
frequency_penalty=(
|
||
sampling_params.frequency_penalty
|
||
if sampling_params.frequency_penalty != 0.0
|
||
else None
|
||
),
|
||
)
|
||
|
||
suppress_processor = _make_suppress_logits_processor(
|
||
self._model_suppress_tokens
|
||
)
|
||
if suppress_processor is not None:
|
||
logits_processors.append(suppress_processor)
|
||
|
||
# Convert stop tokens from Set[int] to Sequence[Sequence[int]]
|
||
# for the new BatchGenerator API (each stop token is a sequence).
|
||
stop_tokens_set = self._get_stop_tokens()
|
||
if sampling_params.stop_token_ids:
|
||
stop_tokens_set.update(sampling_params.stop_token_ids)
|
||
stop_tokens_seq = [[t] for t in stop_tokens_set] if stop_tokens_set else None
|
||
|
||
bg = BatchGenerator(
|
||
model=self.model,
|
||
max_tokens=sampling_params.max_tokens,
|
||
stop_tokens=stop_tokens_seq,
|
||
sampler=sampler,
|
||
logits_processors=logits_processors if logits_processors else [],
|
||
prefill_batch_size=1,
|
||
completion_batch_size=self.config.completion_batch_size,
|
||
prefill_step_size=self.config.prefill_step_size,
|
||
stream=self._stream,
|
||
)
|
||
|
||
return bg
|
||
|
||
def _on_prompt_progress(self, updates: list[tuple[int, int, int]]) -> None:
|
||
"""Callback from BatchGenerator's prefill loop.
|
||
|
||
Called once per prefill chunk (default 2048 tokens) with a list of
|
||
(uid, processed_tokens, total_tokens) tuples. Updates the global
|
||
PrefillProgressTracker so the admin dashboard can display per-request
|
||
prefill progress. Only touches CPU counters — zero GPU overhead.
|
||
"""
|
||
tracker = get_prefill_tracker()
|
||
model_id = self.config.model_name
|
||
for uid, processed, total in updates:
|
||
request_id = self.uid_to_request_id.get(uid)
|
||
if request_id is None:
|
||
continue
|
||
tracker.update(
|
||
request_id=request_id,
|
||
processed=processed,
|
||
total=total,
|
||
model_id=model_id,
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# External prefill (composition pattern — replaces _process_prompts)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _model_uses_mla(self) -> bool:
|
||
"""Detect Multi-head Latent Attention models (DeepSeek-V2/V3/V4,
|
||
GLM-4-MoE / GLM-4.7-Flash, Kimi-K2, ...).
|
||
|
||
MLA compresses K/V into a low-rank latent plus a separate rope key and
|
||
reads the *fetched* cache tensors directly — e.g.
|
||
``kv_latent, k_pe = cache.update_and_fetch(...)`` then
|
||
``k_pe.swapaxes(-1, -2)`` (mlx_lm/models/glm4_moe_lite.py). TurboQuant
|
||
replaces the cache state with quantized NamedTuples that have no array
|
||
methods, so that ``.swapaxes`` raises ``AttributeError`` (#1613). MLA
|
||
also stores keys/values with mismatched head dims, which the codec does
|
||
not support. Such models stay fp16 — no crash, no TurboQuant.
|
||
|
||
Result is memoized: the model never changes for a scheduler instance.
|
||
"""
|
||
cached = getattr(self, "_mla_model", None)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
detected = False
|
||
model = getattr(self, "model", None)
|
||
|
||
# kv_lora_rank is the defining MLA hyperparameter and is an int on real
|
||
# models. It may sit on the top-level config or be nested under a
|
||
# text/LM sub-config (VLM MLA, e.g. kimi_vl -> text_config). The
|
||
# isinstance(int) check guards against mocks where it is a sentinel.
|
||
def _cfg_has_kv_lora(cfg: Any, depth: int = 0) -> bool:
|
||
if cfg is None or depth > 3:
|
||
return False
|
||
if isinstance(getattr(cfg, "kv_lora_rank", None), int):
|
||
return True
|
||
return any(
|
||
_cfg_has_kv_lora(getattr(cfg, sub, None), depth + 1)
|
||
for sub in (
|
||
"text_config",
|
||
"llm_config",
|
||
"language_config",
|
||
"thinker_config",
|
||
)
|
||
)
|
||
|
||
# Config signal. For VLMs the scheduler sees VLMModelAdapter, whose
|
||
# .args delegates to the language model; also probe (_)language_model.
|
||
for holder in (
|
||
model,
|
||
getattr(model, "_language_model", None),
|
||
getattr(model, "language_model", None),
|
||
):
|
||
if holder is None:
|
||
continue
|
||
if _cfg_has_kv_lora(getattr(holder, "args", None)) or _cfg_has_kv_lora(
|
||
getattr(holder, "config", None)
|
||
):
|
||
detected = True
|
||
break
|
||
|
||
# Architecture signal: an attention submodule carrying the MLA
|
||
# down-projection, latent layernorm, or latent rank. Covers models
|
||
# whose config does not surface kv_lora_rank where the scheduler can
|
||
# see it (e.g. a directly-loaded VLM with a nested text config).
|
||
if not detected and model is not None and hasattr(model, "modules"):
|
||
try:
|
||
for m in model.modules():
|
||
if (
|
||
hasattr(m, "kv_a_proj_with_mqa")
|
||
or hasattr(m, "kv_a_layernorm")
|
||
or isinstance(getattr(m, "kv_lora_rank", None), int)
|
||
):
|
||
detected = True
|
||
break
|
||
except Exception:
|
||
pass
|
||
|
||
if detected:
|
||
logger.info(
|
||
"TurboQuant disabled: model uses Multi-head Latent Attention "
|
||
"(MLA), which is incompatible with quantized KV cache states; "
|
||
"keeping fp16 KV cache (#1613)."
|
||
)
|
||
self._mla_model = detected
|
||
return detected
|
||
|
||
def _model_uses_attention_sinks(self) -> bool:
|
||
"""Detect models whose attention path passes sink logits to SDPA.
|
||
|
||
TurboQuant's quantized attention kernels currently do not implement the
|
||
sink term used by attention-sink models. Ignoring it silently changes
|
||
the model's attention distribution, so these models must keep fp16 KV
|
||
unless the attention patch falls back to dequantized sink-aware SDPA.
|
||
"""
|
||
cached = getattr(self, "_attention_sink_model", None)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
detected = False
|
||
model = getattr(self, "model", None)
|
||
|
||
def _has_real_sink_attr(obj: Any) -> bool:
|
||
for name in ("sinks", "attention_sink_bias", "attn_sink"):
|
||
value = None
|
||
if isinstance(obj, dict):
|
||
value = obj.get(name)
|
||
if value is None:
|
||
data = getattr(obj, "__dict__", {})
|
||
if isinstance(data, dict):
|
||
value = data.get(name)
|
||
if isinstance(value, mx.array):
|
||
return True
|
||
if value is not None and isinstance(value, (int, float, list, tuple)):
|
||
return True
|
||
return False
|
||
|
||
try:
|
||
modules = getattr(model, "modules", None)
|
||
except Exception:
|
||
modules = None
|
||
if type(modules).__module__.startswith("unittest.mock"):
|
||
modules = None
|
||
if not detected and callable(modules):
|
||
try:
|
||
for m in modules():
|
||
if _has_real_sink_attr(m):
|
||
detected = True
|
||
break
|
||
except Exception:
|
||
pass
|
||
|
||
if detected:
|
||
logger.info(
|
||
"TurboQuant disabled: model uses attention sinks, which are "
|
||
"not supported by TurboQuant's quantized attention kernels; "
|
||
"keeping fp16 KV cache."
|
||
)
|
||
self._attention_sink_model = detected
|
||
return detected
|
||
|
||
def _turboquant_eligible(self, prompt_cache: list[Any]) -> bool:
|
||
"""True if this cache layout can safely mix TQ and pass-through caches.
|
||
|
||
Plain KVCache layers are TurboQuant-convertible. State-array caches and
|
||
rotating/sliding-window caches are pass-through: they stay in their
|
||
native form while adjacent full-attention KVCache layers are converted.
|
||
|
||
MLA models (DeepSeek / GLM-4.7-Flash) and attention-sink models are
|
||
excluded because their attention paths need semantics TurboQuant's
|
||
quantized cache states/kernels do not currently provide.
|
||
"""
|
||
from mlx_lm.models.cache import ArraysCache, CacheList, KVCache
|
||
|
||
if self._model_uses_mla():
|
||
return False
|
||
if self._model_uses_attention_sinks():
|
||
return False
|
||
|
||
def _ok(c: Any) -> bool:
|
||
if isinstance(c, KVCache):
|
||
return True
|
||
if isinstance(c, ArraysCache):
|
||
return True
|
||
class_name = type(c).__name__
|
||
if class_name in (
|
||
"SizedArraysCache",
|
||
"RotatingKVCache",
|
||
"BatchRotatingKVCache",
|
||
"PrefillReadyRotatingKVCache",
|
||
"BufferedRotatingKVCache",
|
||
"TurboQuantKVCache",
|
||
"BatchTurboQuantKVCache",
|
||
):
|
||
return True
|
||
if class_name in ("MiniMaxM3KVCache", "MiniMaxM3BatchKVCache"):
|
||
return False
|
||
if isinstance(c, CacheList):
|
||
return all(_ok(inner) for inner in c.caches)
|
||
return False
|
||
|
||
return bool(prompt_cache) and all(_ok(c) for c in prompt_cache)
|
||
|
||
def _apply_turboquant_kv_empty(self, prompt_cache: list[Any]) -> None:
|
||
"""Replace empty KVCache layers with empty TurboQuantKVCache.
|
||
|
||
Tokens are quantized on the fly during update_and_fetch, avoiding
|
||
the peak memory spike from storing full-precision KV then converting.
|
||
Used only when there is no prefill history to preserve (the single
|
||
last token is quantized during insert()'s prompt step). Skips the
|
||
last KVCache layer if turboquant_skip_last is set.
|
||
"""
|
||
from mlx_lm.models.cache import CacheList, KVCache
|
||
from mlx_vlm.turboquant import TurboQuantKVCache
|
||
|
||
kv_indices = [
|
||
i for i, c in enumerate(prompt_cache) if _is_turboquant_kv_family_cache(c)
|
||
]
|
||
skip_last = self._turboquant_skip_last and len(kv_indices) > 1
|
||
last_kv_idx = kv_indices[-1] if skip_last else -1
|
||
|
||
converted = 0
|
||
bits = float(self._turboquant_kv_bits)
|
||
for i, cache_obj in enumerate(prompt_cache):
|
||
if isinstance(cache_obj, KVCache):
|
||
if i == last_kv_idx:
|
||
continue
|
||
prompt_cache[i] = TurboQuantKVCache(bits=bits)
|
||
converted += 1
|
||
elif isinstance(cache_obj, CacheList):
|
||
new_caches = []
|
||
for c in cache_obj.caches:
|
||
if isinstance(c, KVCache):
|
||
new_caches.append(TurboQuantKVCache(bits=bits))
|
||
converted += 1
|
||
else:
|
||
new_caches.append(c)
|
||
cache_obj.caches = tuple(new_caches)
|
||
if converted > 0:
|
||
skip_msg = ", skipped last KVCache layer" if skip_last else ""
|
||
logger.info(
|
||
f"TurboQuant: {converted}/{len(prompt_cache)} "
|
||
f"cache layers set to {bits}-bit{skip_msg}"
|
||
)
|
||
|
||
def _apply_turboquant_kv_convert(self, prompt_cache: list[Any]) -> None:
|
||
"""Convert populated KVCache data to TurboQuantKVCache via from_cache().
|
||
|
||
Called AFTER fp16 prefill completes (or on an SSD-restored fp16
|
||
cache): the completed full-precision KV is quantized once, so prefill
|
||
hidden states stay exact and quantization error only enters at
|
||
decode-time reads. This is the key difference from #717/#771, which
|
||
quantized on the fly during prefill and corrupted hidden states.
|
||
"""
|
||
from mlx_lm.models.cache import CacheList, KVCache
|
||
from mlx_vlm.turboquant import TurboQuantKVCache
|
||
|
||
kv_indices = [
|
||
i for i, c in enumerate(prompt_cache) if _is_turboquant_kv_family_cache(c)
|
||
]
|
||
skip_last = self._turboquant_skip_last and len(kv_indices) > 1
|
||
last_kv_idx = kv_indices[-1] if skip_last else -1
|
||
|
||
converted = 0
|
||
bits = float(self._turboquant_kv_bits)
|
||
for i, cache_obj in enumerate(prompt_cache):
|
||
if isinstance(cache_obj, KVCache):
|
||
if i == last_kv_idx:
|
||
continue
|
||
prompt_cache[i] = TurboQuantKVCache.from_cache(cache_obj, bits=bits)
|
||
converted += 1
|
||
elif isinstance(cache_obj, CacheList):
|
||
new_caches = []
|
||
for c in cache_obj.caches:
|
||
if isinstance(c, KVCache):
|
||
new_caches.append(TurboQuantKVCache.from_cache(c, bits=bits))
|
||
converted += 1
|
||
else:
|
||
new_caches.append(c)
|
||
cache_obj.caches = tuple(new_caches)
|
||
if converted > 0:
|
||
skip_msg = ", skipped last KVCache layer" if skip_last else ""
|
||
logger.info(
|
||
f"TurboQuant: converted {converted}/{len(prompt_cache)} "
|
||
f"cache layers to {bits}-bit{skip_msg}"
|
||
)
|
||
|
||
def _do_external_prefill(
|
||
self,
|
||
request: "Request",
|
||
tokens: list[int],
|
||
existing_cache: list[Any] | None,
|
||
vlm_embeds: tuple[mx.array, dict[str, Any], int] | None = None,
|
||
) -> tuple[list[Any], list[int]]:
|
||
"""Run prefill externally (outside BatchGenerator) for a single request.
|
||
|
||
Processes tokens[0:N-1] through the model. The last token tokens[N-1]
|
||
is NOT processed here — it will be passed to BatchGenerator.insert()
|
||
so that the first decode step produces the correct logit.
|
||
|
||
Args:
|
||
request: The request being prefilled.
|
||
tokens: Full token list to prefill.
|
||
existing_cache: Restored cache from paged SSD (or None).
|
||
vlm_embeds: Optional (inputs_embeds, extra_kwargs, start_offset)
|
||
tuple for VLM requests.
|
||
|
||
Returns:
|
||
(prefilled_cache, last_token_list) where last_token_list contains
|
||
the single last token to pass to insert().
|
||
|
||
Raises:
|
||
_PrefillAbortedError: If prefill is interrupted by a pending abort.
|
||
RuntimeError: If memory limit exceeded during prefill.
|
||
"""
|
||
n_tokens = len(tokens)
|
||
if n_tokens <= 1:
|
||
# Nothing to prefill, return cache + tokens as-is.
|
||
cache = existing_cache or make_prompt_cache(self.model)
|
||
# TurboQuant: a TQ cache here makes _merge_caches() build a
|
||
# BatchTurboQuantKVCache (via the monkey-patched merge), so the
|
||
# one decode token quantizes against TQ history. An empty fresh
|
||
# cache gets empty TQ layers; a restored cache preserves its data.
|
||
if self._turboquant_kv_bits is not None and self._turboquant_eligible(
|
||
cache
|
||
):
|
||
if existing_cache is None:
|
||
self._apply_turboquant_kv_empty(cache)
|
||
else:
|
||
self._apply_turboquant_kv_convert(cache)
|
||
return cache, tokens
|
||
|
||
# Create or reuse cache
|
||
if existing_cache is not None:
|
||
prompt_cache = existing_cache
|
||
else:
|
||
prompt_cache = make_prompt_cache(self.model)
|
||
|
||
# Fresh TurboQuant requests run fp16 during the cold prefill loop and
|
||
# are quantized once at the end. Restored TurboQuant prefix caches stay
|
||
# quantized while pre-filling the uncached suffix, then keep using TQ for
|
||
# decode. Rotating/sliding-window layers remain native pass-through
|
||
# caches; only full-attention KVCache layers are converted.
|
||
|
||
# Clear stale mRoPE position state for text-only requests.
|
||
if vlm_embeds is None and hasattr(self.model, "clear_vlm_position_state"):
|
||
self.model.clear_vlm_position_state()
|
||
_seed_text_only_mrope_delta_for_cached_prefill(self.model, request)
|
||
|
||
# Boundary snapshot setup
|
||
block_size = self.config.paged_cache_block_size
|
||
boundary_enabled = (
|
||
block_size > 0
|
||
and self.block_aware_cache is not None
|
||
and _prompt_cache_needs_snapshots(prompt_cache)
|
||
)
|
||
base_size = _cache_base_sizes(prompt_cache) if boundary_enabled else 0
|
||
# Sanity check: base_size from cache offsets should match the number
|
||
# of tokens actually cached. A mismatch indicates stale meta_state
|
||
# in a restored RotatingKVCache (e.g. shared layer_meta_states from
|
||
# an earlier store_cache bug). Use cached_tokens which is always
|
||
# derived from block_table.num_tokens and therefore trustworthy.
|
||
if (
|
||
boundary_enabled
|
||
and hasattr(request, "cached_tokens")
|
||
and request.cached_tokens > 0
|
||
):
|
||
if base_size != request.cached_tokens:
|
||
logger.debug(
|
||
"Cache base_size mismatch: computed %d, expected %d "
|
||
"(cached_tokens). Using cached_tokens for boundary "
|
||
"alignment.",
|
||
base_size,
|
||
request.cached_tokens,
|
||
)
|
||
base_size = request.cached_tokens
|
||
|
||
# Prepare VLM embeddings for prefill
|
||
embeds_array: mx.array | None = None
|
||
extra_kwargs: dict[str, Any] | None = None
|
||
if vlm_embeds is not None:
|
||
embeds_array, extra_kwargs, start_offset = vlm_embeds
|
||
embeds_array = embeds_array[:, start_offset:] # skip cached portion
|
||
if start_offset > 0 and extra_kwargs:
|
||
extra_kwargs = _advance_vlm_extra(extra_kwargs, start_offset)
|
||
# Force _position_ids path in language model for cached VLM
|
||
# prefill. Without this, the delta approach gives sequential
|
||
# positions to image tokens that need 3D mRoPE positions.
|
||
# Setting _rope_deltas=None makes the language model use
|
||
# _position_ids (set by get_input_embeddings) instead.
|
||
# Saved and restored after prefill for decode rope_deltas capture.
|
||
# Only applies to mRoPE VLMs (Qwen2-VL, Qwen2.5-VL, GLM-4V, etc.);
|
||
# non-mRoPE VLMs like Gemma 4 have no _rope_deltas attribute.
|
||
_saved_rope_deltas = None
|
||
if start_offset > 0:
|
||
lm = getattr(self.model, "_language_model", None)
|
||
if lm is not None and hasattr(lm, "_rope_deltas"):
|
||
_saved_rope_deltas = lm._rope_deltas
|
||
lm._rope_deltas = None
|
||
# Stash so the #1405 requeue path can restore it if this prefill
|
||
# raises before the normal restore below runs.
|
||
request._prefill_saved_rope_deltas = _saved_rope_deltas
|
||
|
||
# Prefill tokens[0:N-1] (leave last token for insert())
|
||
prefill_tokens = tokens[:-1]
|
||
last_token = tokens[-1:]
|
||
total_length = len(tokens)
|
||
|
||
# Build the input row on the engine stream: the chunk forwards below
|
||
# run inside mx.stream(self._stream), and a worker-default-stream
|
||
# view would split the chunk eval graph across two streams (see the
|
||
# loop comment below).
|
||
with mx.stream(self._stream):
|
||
input_arr = mx.array(prefill_tokens)[None] # (1, seq_len)
|
||
processed_tokens = 0
|
||
uid = self.request_id_to_uid.get(request.request_id)
|
||
|
||
emitted_boundaries: dict[int, int] = {}
|
||
|
||
while input_arr.shape[1] > 0:
|
||
remaining = input_arr.shape[1]
|
||
prefill_step_size = self._prefill_step_size_for_progress(
|
||
processed_tokens, remaining
|
||
)
|
||
n_to_process = min(prefill_step_size, remaining)
|
||
|
||
if processed_tokens == 0:
|
||
_sync_and_clear_cache(self._stream)
|
||
|
||
# Boundary-limited step size
|
||
if boundary_enabled and block_size > 0:
|
||
current_total = base_size + processed_tokens
|
||
next_boundary = ((current_total // block_size) + 1) * block_size
|
||
target_boundary_prefill = next_boundary - base_size
|
||
delta = target_boundary_prefill - processed_tokens
|
||
if delta > 0:
|
||
n_to_process = min(n_to_process, delta)
|
||
n_to_process = max(1, n_to_process)
|
||
|
||
# Adaptive throttle: shrink chunk when entering the caution zone
|
||
# so the hard cap is honored before the chunk-end check. Raises
|
||
# RuntimeError if the min chunk would exceed the cap — the
|
||
# #1405 cleanup path catches it and emits an error to the client.
|
||
n_to_process = self._adaptive_chunk_size(
|
||
n_to_process,
|
||
request_id=request.request_id,
|
||
loop_label="external",
|
||
kv_len=base_size + processed_tokens,
|
||
)
|
||
|
||
# Pre-chunk safety guard: NEVER submit a chunk whose predicted peak
|
||
# would breach the prefill safety cap. The Metal command-buffer
|
||
# OOM is an async, uncatchable SIGABRT, so it must be prevented
|
||
# before submission — a post-chunk check is too late. Falls back to
|
||
# min_chunk after a reclaim; raises gracefully only if even the
|
||
# floor can't fit (caught by the #1405 path → requeue/clean error).
|
||
n_to_process = self._guard_prefill_chunk(
|
||
n_to_process,
|
||
kv_len=base_size + processed_tokens,
|
||
progress=processed_tokens,
|
||
loop_label="external",
|
||
request_id=request.request_id,
|
||
)
|
||
|
||
_throttle_pre = get_phys_footprint()
|
||
# External prefill bypasses BatchGenerator, so it must establish
|
||
# the per-engine stream context itself. Native lazy primitives
|
||
# otherwise bind to the worker's unrelated default stream and can
|
||
# fail at mx.eval with "There is no Stream(gpu, X) in current
|
||
# thread" (issue #2170). The chunk views (input slices, VLM embed
|
||
# slices, and the advance views for the next chunk) stay inside
|
||
# the same context: a worker-default-stream view splits the chunk
|
||
# eval graph across two streams and adds a per-chunk cross-stream
|
||
# fence, the synchronization pattern implicated in the #2197 and
|
||
# #2183 engine hangs on macOS 26.
|
||
with mx.stream(self._stream):
|
||
model_kwargs: dict[str, Any] = {}
|
||
if embeds_array is not None and embeds_array.shape[1] > 0:
|
||
model_kwargs["inputs_embeds"] = embeds_array[:, :n_to_process]
|
||
if extra_kwargs:
|
||
model_kwargs["vlm_extra_kwargs"] = _slice_vlm_extra(
|
||
extra_kwargs, n_to_process
|
||
)
|
||
self.model(
|
||
input_arr[:, :n_to_process],
|
||
cache=prompt_cache,
|
||
**model_kwargs,
|
||
)
|
||
mx.eval([c.state for c in prompt_cache])
|
||
input_arr = input_arr[:, n_to_process:]
|
||
if embeds_array is not None:
|
||
embeds_array = embeds_array[:, n_to_process:]
|
||
if extra_kwargs:
|
||
extra_kwargs = _advance_vlm_extra(extra_kwargs, n_to_process)
|
||
_throttle_post = get_phys_footprint()
|
||
self._record_chunk_transient(
|
||
n_to_process,
|
||
_throttle_pre,
|
||
_throttle_post,
|
||
request_id=request.request_id,
|
||
loop_label="external",
|
||
)
|
||
|
||
processed_tokens += n_to_process
|
||
|
||
# Progress callback
|
||
if uid is not None:
|
||
self._on_prompt_progress([(uid, processed_tokens, total_length)])
|
||
|
||
# Boundary snapshot emission
|
||
if boundary_enabled:
|
||
total_tokens = base_size + processed_tokens
|
||
if (
|
||
total_tokens > 0
|
||
and total_tokens % block_size == 0
|
||
and emitted_boundaries.get(request.request_id, -1) < total_tokens
|
||
):
|
||
self._emit_prefill_boundary_snapshot(
|
||
request, prompt_cache, total_tokens
|
||
)
|
||
emitted_boundaries[request.request_id] = total_tokens
|
||
|
||
# Memory monitoring — use max(active, phys_footprint) so MLX
|
||
# cache pool and IOAccelerator-backed allocations that don't
|
||
# show in mx.get_active_memory() still trigger the guard.
|
||
# See utils/proc_memory.py for why phys_footprint matters.
|
||
if self._memory_limit_bytes > 0:
|
||
current = self._current_usage_bytes()
|
||
_hard = self._memory_hard_limit_bytes
|
||
_soft = self._memory_limit_bytes
|
||
# Only log when crossing the soft watermark — that's the
|
||
# caution zone where adaptive throttle decisions matter.
|
||
# Skipped on healthy traffic to keep the log quiet.
|
||
if current > _soft:
|
||
logger.debug(
|
||
"[memcheck:external] rid=%s n=%d processed=%d "
|
||
"current=%.3fGB soft=%.3fGB hard=%.3fGB %s",
|
||
request.request_id,
|
||
n_to_process,
|
||
processed_tokens,
|
||
current / 1024**3,
|
||
_soft / 1024**3,
|
||
_hard / 1024**3,
|
||
"OVER_HARD" if _hard > 0 and current > _hard else "OVER_SOFT",
|
||
)
|
||
# Abort decision uses the STABLE physical cap, not the jittery
|
||
# dynamic ceiling: only kill an in-flight prefill if it would
|
||
# breach what Metal actually allows. Throttling above still
|
||
# targets the dynamic ceiling. Falls back to the dynamic hard
|
||
# limit if the abort limit hasn't been propagated yet.
|
||
_abort = self._memory_abort_limit_bytes or self._memory_hard_limit_bytes
|
||
if _abort > 0 and current > _abort:
|
||
# Reclaim the just-computed chunk's Metal transients before
|
||
# giving up — they are still resident at this pre-clear
|
||
# check and are usually what tipped us over the cap.
|
||
current = self._reclaim_prefill_headroom()
|
||
if current > _abort:
|
||
logger.warning(
|
||
f"Prefill force-stopped at {processed_tokens} "
|
||
f"tokens: memory {current / 1024**3:.1f}GB "
|
||
f"exceeds physical cap "
|
||
f"{_abort / 1024**3:.1f}GB (after reclaim)"
|
||
)
|
||
raise RuntimeError("Memory limit exceeded during prefill")
|
||
logger.info(
|
||
"Prefill recovered after reclaim at %d tokens "
|
||
"(%.1fGB <= cap %.1fGB)",
|
||
processed_tokens,
|
||
current / 1024**3,
|
||
_abort / 1024**3,
|
||
)
|
||
elif current > self._memory_limit_bytes:
|
||
logger.warning(
|
||
f"Prefill above max_bytes at "
|
||
f"{processed_tokens} tokens: "
|
||
f"{current / 1024**3:.1f}GB > "
|
||
f"{self._memory_limit_bytes / 1024**3:.1f}GB "
|
||
f"(ceiling: "
|
||
f"{self._memory_hard_limit_bytes / 1024**3:.1f}GB)"
|
||
)
|
||
|
||
# Check for pending aborts between prefill chunks.
|
||
abort_uids = self._check_pending_aborts_for_uids(
|
||
[uid] if uid is not None else []
|
||
)
|
||
if abort_uids:
|
||
logger.info(
|
||
f"Prefill interrupted at {processed_tokens}/"
|
||
f"{total_length} tokens: "
|
||
f"{len(abort_uids)} request(s) aborted"
|
||
)
|
||
if vlm_embeds is not None and _saved_rope_deltas is not None:
|
||
self.model._language_model._rope_deltas = _saved_rope_deltas
|
||
request._prefill_saved_rope_deltas = None
|
||
|
||
# Drop partial-prefill references before clearing the Metal pool.
|
||
# Otherwise the traceback frame can keep large KV/cache arrays
|
||
# alive until after the abort handler returns.
|
||
input_arr = None
|
||
embeds_array = None
|
||
extra_kwargs = None
|
||
model_kwargs = {}
|
||
prompt_cache = None
|
||
_sync_and_clear_cache(self._stream)
|
||
raise _PrefillAbortedError(abort_uids, processed_tokens)
|
||
|
||
# Reclaim Metal intermediates between prefill chunks.
|
||
_sync_and_clear_cache(self._stream)
|
||
|
||
# Emit final boundary snapshot if prompt lands exactly on boundary.
|
||
if boundary_enabled:
|
||
total_tokens = base_size + processed_tokens
|
||
if (
|
||
total_tokens > 0
|
||
and total_tokens % block_size == 0
|
||
and emitted_boundaries.get(request.request_id, -1) < total_tokens
|
||
):
|
||
self._emit_prefill_boundary_snapshot(
|
||
request, prompt_cache, total_tokens
|
||
)
|
||
|
||
_sync_and_clear_cache(self._stream)
|
||
|
||
# Restore _rope_deltas after cached VLM prefill (for decode capture)
|
||
if vlm_embeds is not None and _saved_rope_deltas is not None:
|
||
self.model._language_model._rope_deltas = _saved_rope_deltas
|
||
request._prefill_saved_rope_deltas = None
|
||
|
||
# Quantize the completed fp16 KV cache to TurboQuant for decode.
|
||
# Done here (after the prefill loop, after boundary snapshots are
|
||
# captured fp16) so prefill hidden states stay exact and the paged-SSD
|
||
# format is unchanged. _merge_caches() then builds a
|
||
# BatchTurboQuantKVCache when this request is inserted. Gated to dense
|
||
# KVCache models — chunked/rotating caches stay fp16.
|
||
if self._turboquant_kv_bits is not None and self._turboquant_eligible(
|
||
prompt_cache
|
||
):
|
||
self._apply_turboquant_kv_convert(prompt_cache)
|
||
|
||
if getattr(request, "cached_tokens", 0) > 0:
|
||
with mx.stream(self._stream):
|
||
_materialize_cache_storage(prompt_cache)
|
||
|
||
return prompt_cache, last_token
|
||
|
||
# ------------------------------------------------------------------
|
||
# Adaptive prefill throttle
|
||
# ------------------------------------------------------------------
|
||
|
||
# Discrete step sizes used by the watermark-based throttle. Each tier
|
||
# halves SDPA-fallback transient (∝ query_len × kv_len), so crossing
|
||
# one tier under memory pressure roughly doubles the available
|
||
# headroom for the next chunk's intermediates.
|
||
_PREFILL_STEP_TIERS: tuple[int, ...] = (1024, 512)
|
||
|
||
# Safety margin applied to the headroom (hard_cap - current) when sizing
|
||
# a chunk predictively. The remaining 10% absorbs estimator error and
|
||
# Metal command-buffer overhead above the modeled SDPA + KV growth.
|
||
_PREFILL_HEADROOM_SAFETY: float = 0.90
|
||
|
||
# Default fraction of the physical abort cap we allow a chunk's predicted
|
||
# PEAK to reach. ProcessMemoryEnforcer can override this per tier. The
|
||
# remaining headroom is reserved for Metal command-buffer overhead: a chunk
|
||
# whose peak lands on the wired limit can make Metal abort the command
|
||
# buffer asynchronously (kIOGPUCommandBufferCallbackError OutOfMemory) —
|
||
# an uncatchable SIGABRT — so we keep a hard margin below it.
|
||
_PREFILL_ABORT_MARGIN: float = 0.90
|
||
|
||
# Safety multiplier on the predicted per-chunk transient. The transient
|
||
# scales with query_len * kv_len, so per-token cost grows with context
|
||
# length; this covers one chunk's worth of growth + measurement noise.
|
||
_PREFILL_TRANSIENT_SAFETY: float = 1.3
|
||
_MEMORY_ADMISSION_STALL_TIMEOUT_S: float = 60.0
|
||
_STORE_CACHE_ADMISSION_STALL_TIMEOUT_S: float = 60.0
|
||
|
||
def _predicted_chunk_transient(self, n_tokens: int, kv_len: int) -> float:
|
||
"""Conservative predicted Metal peak growth for one prefill chunk.
|
||
|
||
The per-chunk SDPA/MoE transient scales with ``query_len * kv_len``, so
|
||
the per-token cost GROWS with context length. A long-run EWMA average
|
||
lags that growth and underestimates the next chunk — the cause of the
|
||
Metal command-buffer OOM crash at large kv_len. We therefore take the
|
||
MAX of three signals and apply a safety factor:
|
||
- the most recently MEASURED per-token growth (last_delta / last_n)
|
||
— anchored on reality at the current kv_len regime,
|
||
- the long-run EWMA (model-specific constants the static misses),
|
||
- the kv_len-aware static estimate (SDPA transient + this chunk's
|
||
newly allocated KV).
|
||
Returns 0 only when nothing is known (first chunk, no model info).
|
||
"""
|
||
if n_tokens <= 0:
|
||
return 0.0
|
||
per_token = 0.0
|
||
tracker = self._prefill_transient_tracker
|
||
if tracker is not None:
|
||
if tracker.last_n_tokens > 0 and tracker.last_delta_bytes > 0:
|
||
per_token = max(
|
||
per_token, tracker.last_delta_bytes / tracker.last_n_tokens
|
||
)
|
||
if tracker.bytes_per_token > 0:
|
||
per_token = max(per_token, tracker.bytes_per_token)
|
||
if self.memory_monitor is not None:
|
||
static = self.memory_monitor.estimate_chunk_transient_bytes(
|
||
n_tokens, kv_len + n_tokens
|
||
)
|
||
static += self.memory_monitor.estimate_prompt_kv_bytes(n_tokens)
|
||
per_token = max(per_token, float(static) / n_tokens)
|
||
return per_token * n_tokens * self._PREFILL_TRANSIENT_SAFETY
|
||
|
||
def _prefill_abort_cap(self) -> int:
|
||
"""Safety cap a chunk's predicted peak must stay under.
|
||
|
||
Uses the stable abort limit (min(static, metal_cap)) with a margin so
|
||
we never submit a chunk that could trip the async Metal OOM. Falls back
|
||
to the dynamic hard limit before the abort limit is propagated.
|
||
"""
|
||
cap = self._memory_abort_limit_bytes or self._memory_hard_limit_bytes
|
||
return int(cap * self._prefill_abort_margin) if cap > 0 else 0
|
||
|
||
def _prefill_abort_description(self) -> tuple[int, int, float]:
|
||
"""Return (base cap, safety cap, margin) for diagnostics."""
|
||
base_cap = self._memory_abort_limit_bytes or self._memory_hard_limit_bytes
|
||
safety_cap = self._prefill_abort_cap()
|
||
return base_cap, safety_cap, self._prefill_abort_margin
|
||
|
||
def _sdpa256_unfused_headroom(self) -> int:
|
||
"""Live headroom (bytes) for one unfused SDPA transient, under the
|
||
same target the adaptive prefill throttle enforces (hard ceiling x
|
||
headroom safety, clamped by the abort cap). Negative when no ceiling
|
||
is active (enforcer not propagated yet / guard disabled), which tells
|
||
the sdpa256 route to keep its memory-safe tiled default. Called from
|
||
the route gate on the MLX step thread mid-prefill, where refreshing
|
||
the active-memory sample is safe (issue #2204)."""
|
||
hard_cap = self._memory_hard_limit_bytes
|
||
if hard_cap <= 0:
|
||
return -1
|
||
headroom_safety = getattr(
|
||
self, "_prefill_headroom_safety", self._PREFILL_HEADROOM_SAFETY
|
||
)
|
||
target = int(hard_cap * headroom_safety)
|
||
abort_cap = self._prefill_abort_cap()
|
||
if abort_cap > 0:
|
||
target = min(target, abort_cap)
|
||
return target - self._current_usage_bytes()
|
||
|
||
_MAX_PREFILL_EVICTION_RETRIES = 1
|
||
|
||
def _raise_prefill_eviction_if_available(
|
||
self,
|
||
*,
|
||
request_id: str,
|
||
current: int,
|
||
target_cap: int,
|
||
predicted_transient: int,
|
||
requested_tokens: int,
|
||
reason: str,
|
||
) -> None:
|
||
"""Pause a request once so EngineCore can evict idle LRU models."""
|
||
request = self.requests.get(request_id)
|
||
if request is None:
|
||
return
|
||
max_retries = getattr(
|
||
self,
|
||
"_MAX_PREFILL_EVICTION_RETRIES",
|
||
Scheduler._MAX_PREFILL_EVICTION_RETRIES,
|
||
)
|
||
if request.prefill_eviction_retries >= max_retries:
|
||
return
|
||
if target_cap <= 0 or predicted_transient <= 0:
|
||
return
|
||
|
||
request.prefill_eviction_retries += 1
|
||
config = getattr(self, "config", None)
|
||
eviction_request = PrefillEvictionRequest(
|
||
request_id=request_id,
|
||
model_id=getattr(config, "model_name", ""),
|
||
current_bytes=int(current),
|
||
target_cap_bytes=int(target_cap),
|
||
predicted_transient_bytes=int(predicted_transient),
|
||
requested_tokens=int(requested_tokens),
|
||
reason=reason,
|
||
)
|
||
logger.info(
|
||
"Request %s needs prefill headroom before throttling "
|
||
"(reason=%s, current=%.2fGB, predicted=%.2fGB, target=%.2fGB)",
|
||
request_id,
|
||
reason,
|
||
current / 1024**3,
|
||
predicted_transient / 1024**3,
|
||
target_cap / 1024**3,
|
||
)
|
||
raise _PrefillEvictionNeeded(eviction_request)
|
||
|
||
def _guard_prefill_chunk(
|
||
self,
|
||
n_tokens: int,
|
||
*,
|
||
kv_len: int,
|
||
progress: int,
|
||
loop_label: str,
|
||
request_id: str | None = None,
|
||
) -> int:
|
||
"""Clamp/abort a prefill chunk so its predicted peak can never reach
|
||
the physical Metal cap (the uncatchable async OOM crash).
|
||
|
||
Returns a chunk size whose predicted peak fits under the margined cap
|
||
(possibly shrunk from ``n_tokens``). If even the minimum chunk would
|
||
not fit after a reclaim, raises a clean RuntimeError — the context is
|
||
genuinely too large for available memory. That message intentionally
|
||
does NOT contain "Memory limit exceeded", so ``_requeue_or_fail_prefill``
|
||
fails it fast with a clear error rather than looping a doomed retry.
|
||
"""
|
||
base_cap, cap, margin = self._prefill_abort_description()
|
||
if cap <= 0:
|
||
return n_tokens
|
||
min_chunk = max(1, self._prefill_min_chunk_tokens)
|
||
current = self._current_usage_bytes()
|
||
if current + self._predicted_chunk_transient(n_tokens, kv_len) <= cap:
|
||
return n_tokens
|
||
|
||
# Predicted to breach — reclaim transients and re-measure once.
|
||
current = self._reclaim_prefill_headroom()
|
||
min_transient = self._predicted_chunk_transient(min_chunk, kv_len)
|
||
if current + min_transient > cap:
|
||
maybe_raise_eviction = getattr(
|
||
self, "_raise_prefill_eviction_if_available", None
|
||
)
|
||
if request_id is not None and callable(maybe_raise_eviction):
|
||
maybe_raise_eviction(
|
||
request_id=request_id,
|
||
current=current,
|
||
target_cap=cap,
|
||
predicted_transient=int(min_transient),
|
||
requested_tokens=min_chunk,
|
||
reason="prefill_safety_cap",
|
||
)
|
||
logger.warning(
|
||
"[guard:%s] context too large at progress=%d kv_len=%d: "
|
||
"%.2fGB + min-chunk transient exceeds prefill safety cap "
|
||
"%.2fGB (%d%% of effective ceiling %.2fGB)",
|
||
loop_label,
|
||
progress,
|
||
kv_len,
|
||
current / 1024**3,
|
||
cap / 1024**3,
|
||
round(margin * 100),
|
||
base_cap / 1024**3,
|
||
)
|
||
message = (
|
||
"Prefill context too large for available memory "
|
||
f"(pre-chunk guard at {progress} tokens, kv_len={kv_len}): "
|
||
"predicted peak would exceed prefill safety cap "
|
||
f"{cap / 1024**3:.1f}GB "
|
||
f"({round(margin * 100)}% of effective ceiling "
|
||
f"{base_cap / 1024**3:.1f}GB)"
|
||
)
|
||
raise PrefillMemoryExceededError(
|
||
message=message,
|
||
request_id=request_id,
|
||
estimated_bytes=int(current + min_transient),
|
||
limit_bytes=int(cap),
|
||
)
|
||
|
||
# The floor fits — pick the largest chunk that still fits under the cap.
|
||
per_token = self._predicted_chunk_transient(n_tokens, kv_len) / n_tokens
|
||
safe_n = int((cap - current) / per_token) if per_token > 0 else n_tokens
|
||
n_fit = max(min_chunk, min(n_tokens, safe_n))
|
||
if n_fit < n_tokens:
|
||
logger.debug(
|
||
"[guard:%s] shrink %d -> %d at progress=%d kv_len=%d "
|
||
"(current=%.2fGB cap=%.2fGB)",
|
||
loop_label,
|
||
n_tokens,
|
||
n_fit,
|
||
progress,
|
||
kv_len,
|
||
current / 1024**3,
|
||
cap / 1024**3,
|
||
)
|
||
return n_fit
|
||
|
||
def _adaptive_chunk_size(
|
||
self,
|
||
requested: int,
|
||
*,
|
||
request_id: str,
|
||
loop_label: str,
|
||
kv_len: int = 0,
|
||
) -> int:
|
||
"""Size the next prefill chunk so its predicted peak stays under a
|
||
safety margin below the hard cap.
|
||
|
||
The chunk is sized so that ``current + predicted_transient(n) <=
|
||
hard_cap * safety``. If the full requested chunk already fits, it runs
|
||
unchanged — no behavior change on healthy traffic. Crucially the gate
|
||
is on the *predicted peak*, not on current memory crossing the soft
|
||
watermark: a single large chunk's transient (e.g. MoE prefill at tens
|
||
of MB/token) can blow the ceiling from a low baseline before current
|
||
ever reaches the watermark, which is the failure this prevents.
|
||
|
||
Two predictors feed the sizing:
|
||
- Measured: once the per-scheduler EWMA has samples, use its
|
||
``bytes_per_token`` (× the same 1.2 safety factor ``predict()``
|
||
applies) — this is measurement-based and model-agnostic.
|
||
- First chunk (no samples yet): fall back to the static SDPA + KV
|
||
growth estimate for the requested candidate chunk. ``kv_len`` is
|
||
the current context span (cached prefix + already-prefilled
|
||
tokens), so a large prefix-cache hit with a small suffix is
|
||
throttled correctly without classifying large prefill chunks as
|
||
vector-path traffic.
|
||
|
||
The discrete watermark tiers are retained as a *secondary clamp* —
|
||
they only ever shrink further, never enlarge the predicted size.
|
||
|
||
The chunk-end memory check (``self._memory_hard_limit_bytes``
|
||
comparison in the prefill loops) remains the safety net: if memory
|
||
still exceeds the cap after this shrink, the loop attempts reclaim
|
||
(``_reclaim_prefill_headroom``) and, failing that, raises so the
|
||
#1405 cleanup path can requeue or emit ``finish_reason="error"``.
|
||
|
||
Args:
|
||
requested: The chunk size the caller would have used without
|
||
throttle (already clamped by boundary alignment).
|
||
request_id: For debug log correlation.
|
||
loop_label: "external" or "chunked_step", used only for debug
|
||
log identification.
|
||
kv_len: Current context span (base/cached + processed tokens)
|
||
used for the first-chunk static peak-growth estimate.
|
||
|
||
Returns:
|
||
The chunk size to actually process (>= 1, <= requested).
|
||
"""
|
||
soft_base = self._memory_limit_bytes
|
||
hard_cap = self._memory_hard_limit_bytes
|
||
if soft_base <= 0 or hard_cap <= 0 or requested <= 0:
|
||
return requested
|
||
|
||
current = self._current_usage_bytes()
|
||
min_chunk = max(1, self._prefill_min_chunk_tokens)
|
||
|
||
# Conservative per-token peak growth (measured-last / EWMA / static, ×
|
||
# safety) — see _predicted_chunk_transient. Anchored on the most recent
|
||
# measurement so it tracks growth with kv_len instead of lagging behind
|
||
# a long-run average.
|
||
per_token = self._predicted_chunk_transient(requested, kv_len) / requested
|
||
predictor = "measured" if per_token > 0 else "none"
|
||
|
||
# Keep each chunk's predicted peak under the LOWER of the dynamic
|
||
# throttle target and the prefill safety cap, so the peak can never
|
||
# reach the Metal wall (the uncatchable async OOM).
|
||
headroom_safety = getattr(
|
||
self, "_prefill_headroom_safety", self._PREFILL_HEADROOM_SAFETY
|
||
)
|
||
safe_target = int(hard_cap * headroom_safety)
|
||
abort_cap = self._prefill_abort_cap()
|
||
target = min(safe_target, abort_cap) if abort_cap > 0 else safe_target
|
||
soft_watermark = int(soft_base * self._prefill_safe_zone_ratio)
|
||
|
||
if per_token <= 0:
|
||
# No usable predictor (e.g. model info unavailable). Fall back to
|
||
# the legacy watermark gate so we never run unbounded.
|
||
if current < soft_watermark:
|
||
return requested
|
||
n_fit = requested
|
||
else:
|
||
# Predicted-peak gate: if the FULL requested chunk fits under the
|
||
# target it runs unchanged (covers all healthy traffic). Gated on
|
||
# the predicted peak, not on current crossing the soft watermark —
|
||
# a single big chunk's transient can blow the cap from a low
|
||
# baseline (MoE prefill at tens of MB/token), the failure this
|
||
# prevents.
|
||
if current + per_token * requested <= target:
|
||
return requested
|
||
maybe_raise_eviction = getattr(
|
||
self, "_raise_prefill_eviction_if_available", None
|
||
)
|
||
if callable(maybe_raise_eviction):
|
||
maybe_raise_eviction(
|
||
request_id=request_id,
|
||
current=current,
|
||
target_cap=target,
|
||
predicted_transient=int(per_token * requested),
|
||
requested_tokens=requested,
|
||
reason="adaptive_prefill_throttle",
|
||
)
|
||
headroom = max(target - current, 0)
|
||
n_fit = int(headroom / per_token)
|
||
|
||
n = max(min_chunk, min(requested, n_fit))
|
||
|
||
# Secondary clamp: once in the watermark caution zone, cap by the
|
||
# discrete tiers so a mispredicting EWMA can't run an oversized chunk
|
||
# in deep pressure. Skipped below the watermark so a low-baseline chunk
|
||
# with ample headroom isn't needlessly shrunk.
|
||
band_ratio = -1.0
|
||
if current >= soft_watermark and hard_cap > soft_watermark:
|
||
band = hard_cap - soft_watermark
|
||
band_ratio = max(0.0, min(1.0, (current - soft_watermark) / band))
|
||
if band_ratio < 0.50:
|
||
bucket = self._PREFILL_STEP_TIERS[0] # 1024
|
||
else:
|
||
bucket = self._PREFILL_STEP_TIERS[1] # 512
|
||
n = max(min_chunk, min(n, bucket))
|
||
|
||
if n < requested:
|
||
logger.debug(
|
||
"[throttle:%s] shrink rid=%s chunk %d -> %d "
|
||
"(predictor=%s per_token=%.1fKB current=%.2fGB "
|
||
"safe_target=%.2fGB ceiling=%.2fGB kv_len=%d band_ratio=%.2f)",
|
||
loop_label,
|
||
request_id,
|
||
requested,
|
||
n,
|
||
predictor,
|
||
per_token / 1024,
|
||
current / 1024**3,
|
||
safe_target / 1024**3,
|
||
hard_cap / 1024**3,
|
||
kv_len,
|
||
band_ratio,
|
||
)
|
||
return n
|
||
|
||
def get_cached_mlx_active_memory_bytes(self) -> int:
|
||
"""Return the last MLX active-memory sample taken on the executor."""
|
||
return self._last_mlx_active_memory_bytes
|
||
|
||
def _hot_cache_cpu_bytes(self) -> int:
|
||
"""Return serialized hot-cache bytes safe to exclude from phys guard."""
|
||
config = getattr(self, "config", None)
|
||
budget = getattr(config, "hot_cache_budget", None)
|
||
if budget is not None:
|
||
try:
|
||
return max(0, int(getattr(budget, "total_bytes", 0)))
|
||
except Exception:
|
||
logger.debug("Failed to read shared hot-cache byte budget")
|
||
return 0
|
||
|
||
manager = getattr(self, "paged_ssd_cache_manager", None)
|
||
if manager is None:
|
||
return 0
|
||
|
||
try:
|
||
stats = manager.get_stats()
|
||
return max(0, int(getattr(stats, "hot_cache_size_bytes", 0)))
|
||
except Exception:
|
||
try:
|
||
return max(0, int(getattr(manager, "_hot_cache_total_bytes", 0)))
|
||
except Exception:
|
||
logger.debug("Failed to read local hot-cache byte counter")
|
||
return 0
|
||
|
||
def _current_usage_bytes(self, *, refresh_mlx_active: bool = True) -> int:
|
||
"""Current memory usage for scheduler-side guard checks.
|
||
|
||
Scheduler steps run on the MLX executor thread, so they can refresh
|
||
mx.get_active_memory() safely. Event-loop callers such as early
|
||
preflight use the cached executor sample and phys_footprint instead.
|
||
"""
|
||
active = self._last_mlx_active_memory_bytes
|
||
if refresh_mlx_active:
|
||
active = max(0, int(mx.get_active_memory()))
|
||
self._last_mlx_active_memory_bytes = active
|
||
hot_cache_cpu_bytes = getattr(self, "_hot_cache_cpu_bytes", None)
|
||
if callable(hot_cache_cpu_bytes):
|
||
hot_cache_bytes = hot_cache_cpu_bytes()
|
||
else:
|
||
hot_cache_bytes = Scheduler._hot_cache_cpu_bytes(self)
|
||
phys = max(0, int(get_phys_footprint()) - hot_cache_bytes)
|
||
return max(active, phys)
|
||
|
||
def get_active_hot_cache_block_hashes(self) -> set[bytes]:
|
||
"""Return hot-cache block hashes owned by active in-flight requests."""
|
||
manager = getattr(self, "paged_cache_manager", None)
|
||
if manager is None:
|
||
return set()
|
||
|
||
hashes: set[bytes] = set()
|
||
active_requests = list(self.running.values()) + list(self.prefilling)
|
||
for request in active_requests:
|
||
block_table = getattr(request, "block_table", None)
|
||
if block_table is None:
|
||
continue
|
||
for block_id in getattr(block_table, "block_ids", []) or []:
|
||
try:
|
||
block = manager.blocks[block_id]
|
||
block_hash = getattr(block, "block_hash", None)
|
||
except Exception:
|
||
continue
|
||
if block_hash is not None:
|
||
hashes.add(bytes(block_hash))
|
||
return hashes
|
||
|
||
def _clear_memory_admission_blocker(self, request_id: str | None = None) -> None:
|
||
if (
|
||
request_id is not None
|
||
and request_id != self._memory_admission_blocked_request_id
|
||
):
|
||
return
|
||
self._memory_admission_blocked_request_id = None
|
||
self._memory_admission_blocked_since = 0.0
|
||
|
||
def _clear_store_cache_admission_blocker(
|
||
self, request_id: str | None = None
|
||
) -> None:
|
||
if (
|
||
request_id is not None
|
||
and request_id != self._store_cache_admission_blocked_request_id
|
||
):
|
||
return
|
||
self._store_cache_admission_blocked_request_id = None
|
||
self._store_cache_admission_blocked_since = 0.0
|
||
|
||
def _clear_request_admission_bookkeeping(self, request_id: str) -> None:
|
||
self._cache_freshness_waits.pop(request_id, None)
|
||
self._prefix_cache_prepared.discard(request_id)
|
||
self._clear_memory_admission_blocker(request_id)
|
||
self._clear_store_cache_admission_blocker(request_id)
|
||
|
||
def _memory_admission_stall_output(self, reason: str) -> RequestOutput | None:
|
||
"""Fail one head-of-line request after persistent memory admission stall."""
|
||
if not self.waiting:
|
||
self._clear_memory_admission_blocker()
|
||
return None
|
||
|
||
request = self.waiting[0]
|
||
request_id = request.request_id
|
||
now = time.monotonic()
|
||
if request_id != self._memory_admission_blocked_request_id:
|
||
self._memory_admission_blocked_request_id = request_id
|
||
self._memory_admission_blocked_since = now
|
||
return None
|
||
|
||
timeout = getattr(
|
||
self,
|
||
"_MEMORY_ADMISSION_STALL_TIMEOUT_S",
|
||
Scheduler._MEMORY_ADMISSION_STALL_TIMEOUT_S,
|
||
)
|
||
if now - self._memory_admission_blocked_since < timeout:
|
||
return None
|
||
|
||
stalled_for = now - self._memory_admission_blocked_since
|
||
self.waiting.popleft()
|
||
self._release_paged_cache_for_request(request_id)
|
||
self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
get_prefill_tracker().remove(request_id)
|
||
self._clear_memory_admission_blocker(request_id)
|
||
|
||
message = (
|
||
"Request could not be admitted because memory pressure persisted "
|
||
f"for {stalled_for:.1f}s ({reason}). Reduce context length, free "
|
||
"memory, lower hot_cache_max_size, or loosen memory_guard_tier."
|
||
)
|
||
logger.warning("Memory admission stalled for %s: %s", request_id, message)
|
||
return RequestOutput(
|
||
request_id=request_id,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=message,
|
||
error_code="memory_admission_stalled",
|
||
error_metadata={
|
||
"request_id": request_id,
|
||
"reason": reason,
|
||
"stalled_seconds": int(stalled_for),
|
||
},
|
||
)
|
||
|
||
def _store_cache_admission_stall_output(
|
||
self,
|
||
reason: str,
|
||
*,
|
||
gate_in_flight: int,
|
||
gate_cap: int,
|
||
pending_cleanups: int,
|
||
) -> RequestOutput | None:
|
||
"""Fail one head-of-line request after persistent store-cache stall."""
|
||
if not self.waiting:
|
||
self._clear_store_cache_admission_blocker()
|
||
return None
|
||
|
||
request = self.waiting[0]
|
||
request_id = request.request_id
|
||
now = time.monotonic()
|
||
if request_id != self._store_cache_admission_blocked_request_id:
|
||
self._store_cache_admission_blocked_request_id = request_id
|
||
self._store_cache_admission_blocked_since = now
|
||
return None
|
||
|
||
timeout = getattr(
|
||
self,
|
||
"_STORE_CACHE_ADMISSION_STALL_TIMEOUT_S",
|
||
Scheduler._STORE_CACHE_ADMISSION_STALL_TIMEOUT_S,
|
||
)
|
||
if now - self._store_cache_admission_blocked_since < timeout:
|
||
return None
|
||
|
||
stalled_for = now - self._store_cache_admission_blocked_since
|
||
self.waiting.popleft()
|
||
self._release_paged_cache_for_request(request_id)
|
||
self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
get_prefill_tracker().remove(request_id)
|
||
|
||
message = (
|
||
"Request could not be admitted because store-cache cleanup stayed "
|
||
f"full for {stalled_for:.1f}s ({reason}). The previous response "
|
||
"cache is still being persisted; retry after the cache writer drains "
|
||
"or reduce cache/write pressure."
|
||
)
|
||
logger.warning(
|
||
"Store-cache admission stalled for %s: %s "
|
||
"(in_flight=%d pending_cleanups=%d cap=%d)",
|
||
request_id,
|
||
message,
|
||
gate_in_flight,
|
||
pending_cleanups,
|
||
gate_cap,
|
||
)
|
||
return RequestOutput(
|
||
request_id=request_id,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=message,
|
||
error_code="store_cache_admission_stalled",
|
||
error_metadata={
|
||
"request_id": request_id,
|
||
"reason": reason,
|
||
"stalled_seconds": int(stalled_for),
|
||
"store_cache_in_flight": gate_in_flight,
|
||
"pending_store_cleanups": pending_cleanups,
|
||
"store_cache_cap": gate_cap,
|
||
},
|
||
)
|
||
|
||
def _bypass_hot_cache_under_pressure(self) -> bool:
|
||
"""Return True when SSD-backed hot-cache acceleration should be bypassed."""
|
||
if not self._prefill_memory_guard:
|
||
return False
|
||
if self._memory_limit_bytes <= 0:
|
||
return False
|
||
config = getattr(self, "config", None)
|
||
if config is None:
|
||
return False
|
||
if getattr(config, "hot_cache_only", False):
|
||
return False
|
||
if int(getattr(config, "hot_cache_max_size", 0) or 0) <= 0:
|
||
return False
|
||
if getattr(self, "paged_ssd_cache_manager", None) is None:
|
||
return False
|
||
try:
|
||
current = self._current_usage_bytes()
|
||
except Exception:
|
||
logger.debug("Failed to sample memory for hot-cache pressure bypass")
|
||
return False
|
||
return current >= self._memory_limit_bytes
|
||
|
||
def _record_chunk_transient(
|
||
self,
|
||
n_tokens: int,
|
||
pre_bytes: int,
|
||
post_bytes: int,
|
||
*,
|
||
request_id: str,
|
||
loop_label: str,
|
||
) -> None:
|
||
"""Feed one chunk's measured transient into the EWMA tracker."""
|
||
delta = post_bytes - pre_bytes
|
||
min_chunk = max(1, self._prefill_min_chunk_tokens)
|
||
if n_tokens < min_chunk:
|
||
logger.debug(
|
||
"[throttle:%s] measure rid=%s n=%d delta=%.2fMB "
|
||
"(skipped: tail < min_chunk=%d)",
|
||
loop_label,
|
||
request_id,
|
||
n_tokens,
|
||
delta / 1024**2,
|
||
min_chunk,
|
||
)
|
||
return
|
||
if delta <= 0:
|
||
logger.debug(
|
||
"[throttle:%s] measure rid=%s n=%d delta=%dB (skipped: <=0)",
|
||
loop_label,
|
||
request_id,
|
||
n_tokens,
|
||
delta,
|
||
)
|
||
return
|
||
self._prefill_transient_tracker.update(n_tokens, delta)
|
||
logger.debug(
|
||
"[throttle:%s] measure rid=%s n=%d transient=%.2fMB "
|
||
"per_token=%.1fKB ewma=%.1fKB samples=%d",
|
||
loop_label,
|
||
request_id,
|
||
n_tokens,
|
||
delta / 1024**2,
|
||
(delta / max(n_tokens, 1)) / 1024,
|
||
self._prefill_transient_tracker.bytes_per_token / 1024,
|
||
self._prefill_transient_tracker.samples,
|
||
)
|
||
|
||
def _reclaim_prefill_headroom(self) -> int:
|
||
"""Reclaim Metal headroom mid-prefill and return the re-measured usage.
|
||
|
||
The prefill loops measure the hard-limit at the chunk boundary, which
|
||
is *before* the per-chunk ``_sync_and_clear_cache`` runs — so the
|
||
just-completed forward pass's SDPA intermediates are still resident
|
||
when the limit is checked. Synchronizing and clearing the Metal buffer
|
||
cache here releases those transients, which is exactly the spike that
|
||
drives prefill OOM (observed: 42.8GB at the check → 24.6GB after the
|
||
buffers are reclaimed). This is the only lever that actually lowers
|
||
the physical footprint: paged-cache block eviction merely recycles
|
||
``CacheBlock`` metadata back into the free queue (the pool never
|
||
shrinks, see ``PagedCacheManager._grow_blocks``), so it is deliberately
|
||
not attempted here — it would drop reusable prefix-cache entries for no
|
||
memory benefit.
|
||
|
||
Returns:
|
||
``max(active, phys_footprint)`` after reclaim.
|
||
"""
|
||
_sync_and_clear_cache(self._stream)
|
||
return self._current_usage_bytes()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Chunked prefill helpers (used when config.chunked_prefill=True)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _prefill_step_size_for_progress(
|
||
self, processed_tokens: int, remaining_tokens: int
|
||
) -> int:
|
||
"""Return the scheduler prefill chunk size for the current progress."""
|
||
adaptive_prefill = self._glm_dsa_adaptive_prefill
|
||
if adaptive_prefill is not None:
|
||
from .patches.glm_moe_dsa.generate_patch import (
|
||
_prefill_step_size_for_progress,
|
||
)
|
||
|
||
return _prefill_step_size_for_progress(
|
||
self.config.prefill_step_size,
|
||
processed_tokens,
|
||
remaining_tokens,
|
||
adaptive_prefill,
|
||
)
|
||
|
||
adaptive_prefill = getattr(self, "_minimax_m3_adaptive_prefill", None)
|
||
if adaptive_prefill is None:
|
||
return self.config.prefill_step_size
|
||
from .patches.minimax_m3.generate_patch import (
|
||
_prefill_step_size_for_progress as _minimax_prefill_step_size,
|
||
)
|
||
|
||
return _minimax_prefill_step_size(
|
||
self.config.prefill_step_size,
|
||
processed_tokens,
|
||
remaining_tokens,
|
||
adaptive_prefill,
|
||
)
|
||
|
||
def _begin_prefill(
|
||
self,
|
||
request: "Request",
|
||
tokens: list[int],
|
||
existing_cache: "list[Any] | None",
|
||
) -> _PrefillState:
|
||
"""Initialise a _PrefillState for a non-VLM request.
|
||
|
||
Performs all once-per-request setup (cache creation, boundary config,
|
||
token splitting) without running any model forward passes.
|
||
"""
|
||
if hasattr(self.model, "clear_vlm_position_state"):
|
||
self.model.clear_vlm_position_state()
|
||
_seed_text_only_mrope_delta_for_cached_prefill(self.model, request)
|
||
|
||
prompt_cache = (
|
||
existing_cache
|
||
if existing_cache is not None
|
||
else make_prompt_cache(self.model)
|
||
)
|
||
|
||
block_size = self.config.paged_cache_block_size
|
||
boundary_enabled = (
|
||
block_size > 0
|
||
and self.block_aware_cache is not None
|
||
and _prompt_cache_needs_snapshots(prompt_cache)
|
||
)
|
||
base_size = _cache_base_sizes(prompt_cache) if boundary_enabled else 0
|
||
if (
|
||
boundary_enabled
|
||
and hasattr(request, "cached_tokens")
|
||
and request.cached_tokens > 0
|
||
and base_size != request.cached_tokens
|
||
):
|
||
logger.debug(
|
||
"Cache base_size mismatch: computed %d, expected %d "
|
||
"(cached_tokens). Using cached_tokens for boundary alignment.",
|
||
base_size,
|
||
request.cached_tokens,
|
||
)
|
||
base_size = request.cached_tokens
|
||
|
||
prefill_tokens = tokens[:-1]
|
||
last_token = tokens[-1:]
|
||
# Build the input row on the engine stream so chunk eval graphs stay
|
||
# single-stream (see _do_external_prefill, #2197/#2183).
|
||
with mx.stream(self._stream):
|
||
input_arr = mx.array(prefill_tokens)[None] # (1, N-1)
|
||
|
||
return _PrefillState(
|
||
request=request,
|
||
cache=prompt_cache,
|
||
tokens_remaining=input_arr,
|
||
last_token=last_token,
|
||
tokens_processed=0,
|
||
base_size=base_size,
|
||
emitted_boundaries={},
|
||
boundary_enabled=boundary_enabled,
|
||
block_size=block_size,
|
||
total_length=len(tokens),
|
||
)
|
||
|
||
def _step_prefill_chunk(self, state: _PrefillState) -> bool:
|
||
"""Process one prefill chunk from *state*.
|
||
|
||
Runs the model on at most prefill_step_size tokens, evals the cache,
|
||
emits any due boundary snapshot, updates the prefill progress
|
||
tracker, and clears Metal intermediates.
|
||
|
||
Returns:
|
||
True when all tokens_remaining have been consumed (prefill done).
|
||
|
||
Raises:
|
||
RuntimeError: If the hard memory limit is exceeded.
|
||
"""
|
||
if state.tokens_remaining.shape[1] == 0:
|
||
return True
|
||
|
||
remaining = state.tokens_remaining.shape[1]
|
||
prefill_step_size = self._prefill_step_size_for_progress(
|
||
state.tokens_processed, remaining
|
||
)
|
||
n = min(prefill_step_size, remaining)
|
||
|
||
if state.tokens_processed == 0:
|
||
_sync_and_clear_cache(self._stream)
|
||
|
||
# Clamp to the next block boundary so boundary snapshots fire exactly.
|
||
if state.boundary_enabled and state.block_size > 0:
|
||
current_total = state.base_size + state.tokens_processed
|
||
next_boundary = ((current_total // state.block_size) + 1) * state.block_size
|
||
delta = (next_boundary - state.base_size) - state.tokens_processed
|
||
if delta > 0:
|
||
n = min(n, delta)
|
||
n = max(1, n)
|
||
|
||
# Adaptive throttle — see _adaptive_chunk_size docstring. Raises
|
||
# if even prefill_min_chunk_tokens would exceed the cap; #1405
|
||
# cleanup paths in _schedule_waiting / _advance_chunked_prefills
|
||
# convert that into a finish_reason="error" output for the client.
|
||
n = self._adaptive_chunk_size(
|
||
n,
|
||
request_id=state.request.request_id,
|
||
loop_label="chunked_step",
|
||
kv_len=state.base_size + state.tokens_processed,
|
||
)
|
||
|
||
# Pre-chunk safety guard (mirrors the external loop): never submit a
|
||
# chunk whose predicted peak would trip the uncatchable async Metal OOM.
|
||
n = self._guard_prefill_chunk(
|
||
n,
|
||
kv_len=state.base_size + state.tokens_processed,
|
||
progress=state.tokens_processed,
|
||
loop_label="chunked_step",
|
||
request_id=state.request.request_id,
|
||
)
|
||
|
||
_throttle_pre = get_phys_footprint()
|
||
# Chunked prefill also bypasses BatchGenerator and must establish the
|
||
# same per-engine stream context as the regular external prefill path.
|
||
# The chunk views stay inside it for the same reason (single-stream
|
||
# chunk eval graph, #2197/#2183).
|
||
with mx.stream(self._stream):
|
||
chunk = state.tokens_remaining[:, :n]
|
||
state.tokens_remaining = state.tokens_remaining[:, n:]
|
||
self.model(chunk, cache=state.cache)
|
||
mx.eval([c.state for c in state.cache])
|
||
_throttle_post = get_phys_footprint()
|
||
self._record_chunk_transient(
|
||
n,
|
||
_throttle_pre,
|
||
_throttle_post,
|
||
request_id=state.request.request_id,
|
||
loop_label="chunked_step",
|
||
)
|
||
state.tokens_processed += n
|
||
|
||
# Boundary snapshot
|
||
if state.boundary_enabled:
|
||
total_tokens = state.base_size + state.tokens_processed
|
||
rid = state.request.request_id
|
||
if (
|
||
total_tokens > 0
|
||
and total_tokens % state.block_size == 0
|
||
and state.emitted_boundaries.get(rid, -1) < total_tokens
|
||
):
|
||
self._emit_prefill_boundary_snapshot(
|
||
state.request, state.cache, total_tokens
|
||
)
|
||
state.emitted_boundaries[rid] = total_tokens
|
||
|
||
# Progress callback so the admin UI prefilling list advances during
|
||
# chunked prefill. _do_external_prefill calls _on_prompt_progress
|
||
# via the temp_uid mapping; the chunked path has no temp uid so we
|
||
# talk to the tracker directly with the request_id.
|
||
get_prefill_tracker().update(
|
||
state.request.request_id,
|
||
state.tokens_processed,
|
||
state.total_length - 1,
|
||
(
|
||
self.config.model_name
|
||
if self.config.model_name
|
||
else ""
|
||
),
|
||
)
|
||
|
||
# Memory monitoring — use max(active, phys_footprint) so MLX cache
|
||
# pool and IOAccelerator-backed allocations that don't show up in
|
||
# mx.get_active_memory() still trigger the guard. Matches the
|
||
# _do_external_prefill check; on macOS jetsam watches
|
||
# phys_footprint, so the active-only check could miss the page
|
||
# before the kernel kills us.
|
||
if self._memory_limit_bytes > 0:
|
||
current = self._current_usage_bytes()
|
||
_hard = self._memory_hard_limit_bytes
|
||
_soft = self._memory_limit_bytes
|
||
# Caution-zone-only memcheck log (see external loop counterpart).
|
||
if current > _soft:
|
||
logger.debug(
|
||
"[memcheck:chunked_step] rid=%s n=%d processed=%d/%d "
|
||
"current=%.3fGB soft=%.3fGB hard=%.3fGB %s",
|
||
state.request.request_id,
|
||
n,
|
||
state.tokens_processed,
|
||
state.total_length - 1,
|
||
current / 1024**3,
|
||
_soft / 1024**3,
|
||
_hard / 1024**3,
|
||
"OVER_HARD" if _hard > 0 and current > _hard else "OVER_SOFT",
|
||
)
|
||
# Abort on the stable physical cap, not the jittery dynamic ceiling
|
||
# (mirrors the external prefill loop).
|
||
_abort = self._memory_abort_limit_bytes or self._memory_hard_limit_bytes
|
||
if _abort > 0 and current > _abort:
|
||
# Reclaim the just-computed chunk's Metal transients before
|
||
# giving up (mirrors the external prefill loop).
|
||
current = self._reclaim_prefill_headroom()
|
||
if current > _abort:
|
||
raise RuntimeError(
|
||
f"Memory limit exceeded during chunked prefill at "
|
||
f"{state.tokens_processed}/{state.total_length - 1} tokens: "
|
||
f"{current / 1024**3:.1f}GB exceeds physical cap "
|
||
f"{_abort / 1024**3:.1f}GB (after reclaim)"
|
||
)
|
||
logger.info(
|
||
"Chunked prefill recovered after reclaim at %d/%d tokens "
|
||
"(%.1fGB <= cap %.1fGB)",
|
||
state.tokens_processed,
|
||
state.total_length - 1,
|
||
current / 1024**3,
|
||
_abort / 1024**3,
|
||
)
|
||
elif current > self._memory_limit_bytes:
|
||
logger.warning(
|
||
f"Chunked prefill above max_bytes at "
|
||
f"{state.tokens_processed} tokens: "
|
||
f"{current / 1024**3:.1f}GB > "
|
||
f"{self._memory_limit_bytes / 1024**3:.1f}GB "
|
||
f"(ceiling: "
|
||
f"{self._memory_hard_limit_bytes / 1024**3:.1f}GB)"
|
||
)
|
||
|
||
_sync_and_clear_cache(self._stream)
|
||
return state.tokens_remaining.shape[1] == 0
|
||
|
||
def _emit_final_boundary_if_needed(self, state: _PrefillState) -> None:
|
||
"""Emit a final boundary snapshot if the prefill landed on a boundary."""
|
||
if not state.boundary_enabled:
|
||
return
|
||
total_tokens = state.base_size + state.tokens_processed
|
||
rid = state.request.request_id
|
||
if (
|
||
total_tokens > 0
|
||
and total_tokens % state.block_size == 0
|
||
and state.emitted_boundaries.get(rid, -1) < total_tokens
|
||
):
|
||
self._emit_prefill_boundary_snapshot(
|
||
state.request, state.cache, total_tokens
|
||
)
|
||
|
||
def _finalize_chunked_prefill_cache_for_insert(
|
||
self, request: "Request", prompt_cache: list[Any] | None
|
||
) -> None:
|
||
"""Mirror external prefill's post-prefill cache epilogue."""
|
||
if not prompt_cache or self._turboquant_kv_bits is None:
|
||
return
|
||
if not self._turboquant_eligible(prompt_cache):
|
||
return
|
||
|
||
self._apply_turboquant_kv_convert(prompt_cache)
|
||
if getattr(request, "cached_tokens", 0) > 0:
|
||
with mx.stream(self._stream):
|
||
_materialize_cache_storage(prompt_cache)
|
||
_sync_and_clear_cache(self._stream)
|
||
|
||
def _insert_prefilled_request(
|
||
self,
|
||
request: "Request",
|
||
state: _PrefillState,
|
||
scheduled: "list[Request]",
|
||
) -> None:
|
||
"""Insert a fully-prefilled request into BatchGenerator.
|
||
|
||
Handles the batch_generator.insert() call, uid bookkeeping, and moving
|
||
the request to self.running. Called from both the inline chunked path
|
||
(first chunk completed immediately) and _advance_chunked_prefills()
|
||
(last chunk completed across steps).
|
||
|
||
Precondition: state.sampler, state.sm, state.per_row_lps are set.
|
||
"""
|
||
self._finalize_chunked_prefill_cache_for_insert(request, state.cache)
|
||
|
||
if request.sampling_params.seed is not None:
|
||
mx.random.seed(request.sampling_params.seed)
|
||
|
||
per_row_lps = state.per_row_lps if state.per_row_lps is not None else []
|
||
uids = self.batch_generator.insert(
|
||
[state.last_token],
|
||
max_tokens=[request.sampling_params.max_tokens],
|
||
caches=[state.cache] if state.cache else None,
|
||
all_tokens=[_batch_generator_all_tokens(request)],
|
||
samplers=[state.sampler],
|
||
logits_processors=[per_row_lps],
|
||
state_machines=[state.sm],
|
||
)
|
||
if uids:
|
||
_register_uid_rows(self.model, uids, [state.sampler], [per_row_lps])
|
||
uid = uids[0]
|
||
self.request_id_to_uid[request.request_id] = uid
|
||
self.uid_to_request_id[uid] = request.request_id
|
||
now = time.monotonic()
|
||
request.batch_uid = uid
|
||
request.status = RequestStatus.RUNNING
|
||
request.generation_started_at = now
|
||
request.last_activity_at = now
|
||
self.running[request.request_id] = request
|
||
scheduled.append(request)
|
||
|
||
if hasattr(self.model, "register_rope_delta"):
|
||
self.model.register_rope_delta(uid, request.rope_deltas)
|
||
|
||
self.total_prompt_tokens += request.num_prompt_tokens
|
||
cache_info = (
|
||
f", {request.cached_tokens} cached" if request.cached_tokens > 0 else ""
|
||
)
|
||
logger.debug(
|
||
"Scheduled chunked-prefill request %s (uid=%d) "
|
||
"with %d tokens (%d total)%s",
|
||
request.request_id,
|
||
uid,
|
||
len(state.last_token),
|
||
request.num_prompt_tokens,
|
||
cache_info,
|
||
)
|
||
|
||
def _advance_chunked_prefills(
|
||
self,
|
||
scheduled: "list[Request]",
|
||
rejected: "list[RequestOutput]",
|
||
) -> None:
|
||
"""Process one prefill chunk per in-flight chunked-prefill request.
|
||
|
||
Called at the start of each step() before _schedule_waiting(). Each
|
||
call advances every request in self.prefilling by one prefill_step_size
|
||
chunk. When a request's prefill completes it is inserted into
|
||
BatchGenerator and moved to self.running.
|
||
|
||
Args:
|
||
scheduled: The step's running list of newly-scheduled requests;
|
||
completed chunked-prefill requests are appended here.
|
||
rejected: Per-step rejected outputs. A chunked prefill that hits
|
||
the memory hard limit emits a finish_reason="error" entry
|
||
here so the engine can surface the failure to the client.
|
||
"""
|
||
if not self.prefilling:
|
||
return
|
||
|
||
pending_prefills = list(self.prefilling)
|
||
still_prefilling: deque[Request] = deque()
|
||
|
||
for index, request in enumerate(pending_prefills):
|
||
rid = request.request_id
|
||
state = self._prefill_states.get(rid)
|
||
|
||
# State missing means the request was aborted and cleaned up by
|
||
# _do_abort_request() between steps — just skip it.
|
||
if state is None:
|
||
continue
|
||
|
||
try:
|
||
done = self._step_prefill_chunk(state)
|
||
except _PrefillAbortedError:
|
||
# Request aborted mid-chunk. Discard state; the abort will
|
||
# be fully processed by _process_pending_aborts() next step.
|
||
self._prefill_states.pop(rid, None)
|
||
_sync_and_clear_cache(self._stream)
|
||
continue
|
||
except _PrefillEvictionNeeded as e:
|
||
self._pending_prefill_eviction_request = e.request
|
||
still_prefilling.append(request)
|
||
still_prefilling.extend(pending_prefills[index + 1 :])
|
||
logger.info(
|
||
"Paused chunked prefill request %s for LRU eviction " "(reason=%s)",
|
||
rid,
|
||
e.request.reason,
|
||
)
|
||
break
|
||
except PrefillMemoryExceededError as e:
|
||
logger.error("Chunked prefill capacity rejected for %s: %s", rid, e)
|
||
self._prefill_states.pop(rid, None)
|
||
self._release_paged_cache_for_request(rid)
|
||
self.requests.pop(rid, None)
|
||
self._clear_request_admission_bookkeeping(rid)
|
||
get_prefill_tracker().remove(rid)
|
||
_sync_and_clear_cache()
|
||
rejected.append(_prefill_memory_exception_output(rid, e))
|
||
continue
|
||
except RuntimeError as e:
|
||
logger.error("Chunked prefill failed for %s: %s", rid, e)
|
||
self._prefill_states.pop(rid, None)
|
||
self._release_paged_cache_for_request(rid)
|
||
self.requests.pop(rid, None)
|
||
self._clear_request_admission_bookkeeping(rid)
|
||
get_prefill_tracker().remove(rid)
|
||
# Drop Metal cache pool buffers held by the aborted chunk's
|
||
# forward / mx.eval transients. Without this, enforcer keeps
|
||
# seeing the burst footprint until the next mx.clear_cache().
|
||
_sync_and_clear_cache()
|
||
# Try a bounded requeue before surfacing the failure: a
|
||
# memory-pressure prefill gets a fresh, better-throttled
|
||
# attempt. Only after the retry budget is exhausted (or for
|
||
# non-memory errors) do we emit the client-facing error.
|
||
if self._requeue_or_fail_prefill(request, e):
|
||
continue
|
||
# Surface the failure to the engine. Without this, the
|
||
# request is silently dropped and the client hangs.
|
||
rejected.append(
|
||
RequestOutput(
|
||
request_id=rid,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=str(e),
|
||
)
|
||
)
|
||
continue
|
||
|
||
if not done:
|
||
still_prefilling.append(request)
|
||
continue
|
||
|
||
# Prefill complete — emit final boundary snapshot and insert.
|
||
self._prefill_states.pop(rid, None)
|
||
self._emit_final_boundary_if_needed(state)
|
||
_sync_and_clear_cache(self._stream)
|
||
|
||
# Ensure a BatchGenerator exists (may not if all requests were
|
||
# previously in chunked prefill with no running decode).
|
||
self._ensure_batch_generator(request.sampling_params)
|
||
if self.batch_generator is None:
|
||
# Unlikely, but if BG creation fails put request back.
|
||
logger.error(
|
||
"BatchGenerator unavailable at chunked-prefill completion "
|
||
"for %s; requeueing.",
|
||
rid,
|
||
)
|
||
still_prefilling.append(request)
|
||
self._prefill_states[rid] = state
|
||
continue
|
||
|
||
# Clean up the prefill-progress tracker entry.
|
||
get_prefill_tracker().remove(rid)
|
||
|
||
self._insert_prefilled_request(request, state, scheduled)
|
||
|
||
self.prefilling = still_prefilling
|
||
|
||
def _build_state_machine(self, request: "Request") -> SequenceStateMachine:
|
||
"""Build a SequenceStateMachine for per-request stop tokens.
|
||
|
||
Combines base stop tokens (EOS, Harmony) with request-specific
|
||
stop_token_ids and tokenized stop strings into a single state
|
||
machine that tells BatchGenerator when to stop generating for
|
||
this request.
|
||
"""
|
||
stop_tokens_set = self._get_stop_tokens()
|
||
if request.sampling_params.stop_token_ids:
|
||
stop_tokens_set.update(request.sampling_params.stop_token_ids)
|
||
|
||
transitions: dict[str, list] = {
|
||
"normal": [([t], None) for t in stop_tokens_set]
|
||
}
|
||
|
||
# Tokenize stop strings into token sequences. mlx-lm's
|
||
# SequenceStateMachine uses Aho-Corasick, so per-token match
|
||
# cost stays O(1) regardless of how many sequences are added.
|
||
# BPE merge edge cases (where a stop string boundary lands
|
||
# mid-token) may miss; that is a known limitation.
|
||
for stop_str in request.sampling_params.stop or []:
|
||
if not isinstance(stop_str, str) or not stop_str:
|
||
continue
|
||
try:
|
||
seq = self.tokenizer.encode(stop_str, add_special_tokens=False)
|
||
except TypeError:
|
||
seq = self.tokenizer.encode(stop_str)
|
||
if seq:
|
||
transitions["normal"].append((list(seq), None))
|
||
|
||
if transitions["normal"]:
|
||
return SequenceStateMachine(transitions, initial="normal")
|
||
return SequenceStateMachine({}, initial="normal")
|
||
|
||
def _emit_prefill_boundary_snapshot(
|
||
self,
|
||
request: "Request",
|
||
prompt_cache: list[Any],
|
||
total_tokens: int,
|
||
) -> None:
|
||
"""Capture boundary snapshot from individual (non-batch) cache.
|
||
|
||
During external prefill we have direct access to per-layer cache
|
||
objects (not BatchKVCache). Extract non-sliceable layers for
|
||
boundary snapshot storage.
|
||
|
||
Pass ``request_id`` directly. The request is mid-prefill and has
|
||
not been inserted into ``BatchGenerator`` yet, so
|
||
``request_id_to_uid`` has no entry for it. The earlier shape
|
||
routed through ``self.request_id_to_uid.get(request_id, -1)`` →
|
||
``uid_to_request_id.get(-1)`` → ``None`` → silent return,
|
||
dropping every snapshot. For ArraysCache / GDN / hybrid models
|
||
that made every non-last block store a placeholder, and the
|
||
next identical-prompt request rejected the cache and re-
|
||
prefilled from scratch.
|
||
"""
|
||
snapshot_cache = [
|
||
c if type(c).__name__ not in _KNOWN_SLICEABLE_CACHE_TYPES else None
|
||
for c in prompt_cache
|
||
]
|
||
self._on_prefill_boundary_snapshot(
|
||
request.request_id,
|
||
snapshot_cache,
|
||
total_tokens,
|
||
)
|
||
|
||
def _build_sampler_and_processors(
|
||
self, sampling_params: SamplingParams, request: Any = None
|
||
) -> tuple[Callable[[mx.array], mx.array], list[Callable]]:
|
||
"""Build per-request sampler and logits processors."""
|
||
# Use omlx.utils.sampling.make_sampler instead of mlx_lm.sample_utils.
|
||
# The mlx-lm version decorates categorical_sampling and apply_* with
|
||
# @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state),
|
||
# which fails to advance the RNG state after the first call in this
|
||
# server environment. Identical prompts then produce identical output
|
||
# even at temperature > 1.
|
||
sampler = omlx_make_sampler(
|
||
temp=sampling_params.temperature,
|
||
top_p=sampling_params.top_p,
|
||
min_p=sampling_params.min_p,
|
||
top_k=sampling_params.top_k,
|
||
xtc_probability=sampling_params.xtc_probability,
|
||
xtc_threshold=sampling_params.xtc_threshold,
|
||
xtc_special_tokens=self._xtc_special_tokens,
|
||
)
|
||
logits_processors = make_logits_processors(
|
||
repetition_penalty=(
|
||
sampling_params.repetition_penalty
|
||
if sampling_params.repetition_penalty != 1.0
|
||
else None
|
||
),
|
||
presence_penalty=(
|
||
sampling_params.presence_penalty
|
||
if sampling_params.presence_penalty != 0.0
|
||
else None
|
||
),
|
||
frequency_penalty=(
|
||
sampling_params.frequency_penalty
|
||
if sampling_params.frequency_penalty != 0.0
|
||
else None
|
||
),
|
||
)
|
||
|
||
suppress_processor = _make_suppress_logits_processor(
|
||
self._model_suppress_tokens
|
||
)
|
||
if suppress_processor is not None:
|
||
logits_processors.append(suppress_processor)
|
||
|
||
# Add thinking budget processor for reasoning models
|
||
if (
|
||
sampling_params.thinking_budget is not None
|
||
and request is not None
|
||
and (
|
||
getattr(request, "needs_think_prefix", False)
|
||
or self._get_output_parser_thinking_end_text() is not None
|
||
)
|
||
):
|
||
think_end_ids = self._resolve_think_end_token_ids()
|
||
if think_end_ids:
|
||
from .api.thinking import ThinkingBudgetProcessor
|
||
|
||
think_start_id = self._get_think_token_id("think_start_id")
|
||
leading_ids, trailing_ids = self._resolve_think_close_pattern(
|
||
self._get_output_parser_thinking_end_text()
|
||
)
|
||
parser_trailing_ids = (
|
||
self._resolve_output_parser_thinking_trailing_ids()
|
||
)
|
||
if parser_trailing_ids is not None:
|
||
trailing_ids = parser_trailing_ids
|
||
processor = ThinkingBudgetProcessor(
|
||
think_end_token_ids=think_end_ids,
|
||
budget=sampling_params.thinking_budget,
|
||
think_start_token_id=think_start_id,
|
||
leading_token_ids=leading_ids,
|
||
trailing_token_ids=trailing_ids,
|
||
token_to_piece=self._thinking_budget_token_to_piece,
|
||
)
|
||
logits_processors.append(processor)
|
||
|
||
# Add grammar constraint processor for structured output.
|
||
# Phase awareness (thinking vs output) is handled by the compiled
|
||
# grammar itself via xgrammar structural tags, so we don't need
|
||
# think_end_ids here.
|
||
if sampling_params.compiled_grammar is not None:
|
||
try:
|
||
from .api.grammar import GrammarConstraintProcessor
|
||
|
||
vocab_size = self._get_model_vocab_size()
|
||
if vocab_size is not None:
|
||
processor = GrammarConstraintProcessor(
|
||
compiled_grammar=sampling_params.compiled_grammar,
|
||
vocab_size=vocab_size,
|
||
)
|
||
logits_processors.append(processor)
|
||
else:
|
||
logger.warning(
|
||
"Cannot determine vocab_size; skipping grammar constraint"
|
||
)
|
||
except ImportError:
|
||
logger.warning("xgrammar not installed; skipping grammar constraint")
|
||
|
||
return sampler, logits_processors
|
||
|
||
def _get_model_vocab_size(self) -> int | None:
|
||
"""Return vocab_size from model config, or None if unavailable."""
|
||
from .utils.tokenizer import resolve_vocab_size
|
||
|
||
return resolve_vocab_size(self.model)
|
||
|
||
def _get_think_token_id(self, attr: str) -> int | None:
|
||
"""Safely read a think token id from the tokenizer.
|
||
|
||
mlx-lm tokenizers expose ``think_start_id`` / ``think_end_id`` as
|
||
properties that may raise ``ValueError`` (multi-token sequence) or
|
||
``TypeError`` (``_think_start_tokens`` is ``None`` for models without
|
||
thinking support, e.g. context-1 / harmony parser).
|
||
|
||
Returns the token id, or ``None`` when unavailable.
|
||
"""
|
||
try:
|
||
return getattr(self.tokenizer, attr, None)
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
def _get_output_parser_thinking_end_text(self) -> str | None:
|
||
"""Return parser-provided thinking close text, if the parser has one."""
|
||
factory = getattr(self, "_output_parser_factory", None)
|
||
if factory is None:
|
||
return None
|
||
return getattr(factory, "thinking_end_text", None)
|
||
|
||
def _get_output_parser_thinking_start_text(self) -> str | None:
|
||
"""Return parser-provided thinking open text, if the parser has one."""
|
||
factory = getattr(self, "_output_parser_factory", None)
|
||
if factory is None:
|
||
return None
|
||
return getattr(factory, "thinking_start_text", None)
|
||
|
||
def _get_output_parser_thinking_start_output_text(self) -> str | None:
|
||
"""Return normalized text to prepend when parser thinking starts in prompt."""
|
||
factory = getattr(self, "_output_parser_factory", None)
|
||
if factory is None:
|
||
return None
|
||
return getattr(factory, "thinking_start_output_text", None)
|
||
|
||
def _encode_thinking_marker(self, text: str) -> list[int] | None:
|
||
"""Encode a parser/tokenizer thinking marker into token IDs."""
|
||
try:
|
||
ids = self.tokenizer.encode(text, add_special_tokens=False)
|
||
except TypeError:
|
||
try:
|
||
ids = self.tokenizer.encode(text)
|
||
except Exception:
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
if ids:
|
||
return list(ids)
|
||
return None
|
||
|
||
def _thinking_budget_token_to_piece(self, token_id: int) -> str | bytes | None:
|
||
"""Best-effort token piece lookup for UTF-8-safe budget forcing."""
|
||
try:
|
||
token = self.tokenizer.convert_ids_to_tokens(token_id)
|
||
if token is not None:
|
||
byte_piece = self._token_piece_to_bytes(token)
|
||
return byte_piece if byte_piece is not None else token
|
||
except (AttributeError, KeyError, TypeError, ValueError):
|
||
pass
|
||
|
||
try:
|
||
return self.tokenizer.decode([token_id], skip_special_tokens=False)
|
||
except TypeError:
|
||
try:
|
||
return self.tokenizer.decode([token_id])
|
||
except Exception:
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
def _token_piece_to_bytes(self, token: str) -> bytes | None:
|
||
"""Convert byte-fallback tokenizer pieces to raw bytes when possible."""
|
||
import re
|
||
|
||
byte_fallback = re.fullmatch(r"(?:<0x[0-9A-Fa-f]{2}>)+", token)
|
||
if byte_fallback is not None:
|
||
return bytes(
|
||
int(match.group(1), 16)
|
||
for match in re.finditer(r"<0x([0-9A-Fa-f]{2})>", token)
|
||
)
|
||
|
||
byte_decoder = getattr(self.tokenizer, "byte_decoder", None)
|
||
if isinstance(byte_decoder, dict) and token:
|
||
try:
|
||
return bytes(byte_decoder[ch] for ch in token)
|
||
except (KeyError, TypeError, ValueError):
|
||
pass
|
||
|
||
return None
|
||
|
||
def _resolve_output_parser_thinking_trailing_ids(self) -> list[int] | None:
|
||
"""Resolve parser-provided tokens that should follow a forced close."""
|
||
factory = getattr(self, "_output_parser_factory", None)
|
||
if factory is None:
|
||
return None
|
||
|
||
trailing_text = getattr(factory, "thinking_end_trailing_text", None)
|
||
if not trailing_text:
|
||
return None
|
||
|
||
return self._encode_thinking_marker(trailing_text)
|
||
|
||
def _resolve_think_end_token_ids(self) -> list[int] | None:
|
||
"""Resolve token ID(s) for the close-think tag.
|
||
|
||
Uses mlx-lm's built-in think_end_id which supports both
|
||
</think> and </longcat_think> automatically.
|
||
"""
|
||
parser_think_end = self._get_output_parser_thinking_end_text()
|
||
if parser_think_end is not None:
|
||
return self._encode_thinking_marker(parser_think_end)
|
||
|
||
# Tier 1: mlx-lm tokenizer attribute (covers all known think variants)
|
||
think_end_id = self._get_think_token_id("think_end_id")
|
||
if think_end_id is not None:
|
||
return [think_end_id]
|
||
|
||
# Tier 2: encode the think_end string
|
||
think_end_str = getattr(self.tokenizer, "think_end", "</think>")
|
||
try:
|
||
ids = self.tokenizer.encode(think_end_str, add_special_tokens=False)
|
||
if ids:
|
||
return list(ids)
|
||
except Exception:
|
||
pass
|
||
|
||
# Tier 3: direct token lookup
|
||
try:
|
||
tid = self.tokenizer.convert_tokens_to_ids("</think>")
|
||
if tid != getattr(self.tokenizer, "unk_token_id", None):
|
||
return [tid]
|
||
except (AttributeError, KeyError, TypeError):
|
||
pass
|
||
|
||
return None
|
||
|
||
def _resolve_think_close_pattern(
|
||
self, think_end_str: str | None = None
|
||
) -> tuple[list[int] | None, list[int] | None]:
|
||
"""Detect leading/trailing tokens around </think> from the chat template.
|
||
|
||
Different models use different patterns:
|
||
- Qwen3/3.5, MiniMax: ``\\n</think>\\n\\n``
|
||
- DeepSeek V3.2, GLM-5: ``</think>`` (no newlines)
|
||
- GLM-4.6V: ``</think>\\n``
|
||
- Step-3.5-Flash: ``\\n</think>\\n``
|
||
|
||
Returns (leading_token_ids, trailing_token_ids) or (None, None).
|
||
"""
|
||
import re
|
||
|
||
if think_end_str is None:
|
||
think_end_str = getattr(self.tokenizer, "think_end", None) or "</think>"
|
||
|
||
# Try to get the chat template text
|
||
template_text = self._get_chat_template_text()
|
||
if not template_text:
|
||
return None, None
|
||
|
||
# Find the close pattern in the template, e.g. \n</think>\n\n
|
||
# Look for the think_end_str surrounded by whitespace/newlines in string literals
|
||
escaped = re.escape(think_end_str)
|
||
# Match patterns like: \n</think>\n\n or </think> in template strings
|
||
match = re.search(
|
||
r"(\\n|\\r|[\n\r])*" + escaped + r"((?:\\n|\\r|[\n\r])*)",
|
||
template_text,
|
||
)
|
||
if not match:
|
||
return None, None
|
||
|
||
# Extract raw leading/trailing whitespace, converting \n escapes to actual newlines
|
||
raw_leading = (
|
||
match.group(0)
|
||
.split(think_end_str)[0]
|
||
.replace("\\n", "\n")
|
||
.replace("\\r", "\r")
|
||
)
|
||
raw_trailing = (
|
||
match.group(0)
|
||
.split(think_end_str)[1]
|
||
.replace("\\n", "\n")
|
||
.replace("\\r", "\r")
|
||
)
|
||
|
||
# Encode to token IDs
|
||
leading_ids = None
|
||
trailing_ids = None
|
||
if raw_leading:
|
||
try:
|
||
ids = self.tokenizer.encode(raw_leading, add_special_tokens=False)
|
||
if ids:
|
||
leading_ids = list(ids)
|
||
except Exception:
|
||
pass
|
||
if raw_trailing:
|
||
try:
|
||
ids = self.tokenizer.encode(raw_trailing, add_special_tokens=False)
|
||
if ids:
|
||
trailing_ids = list(ids)
|
||
except Exception:
|
||
pass
|
||
|
||
return leading_ids, trailing_ids
|
||
|
||
def _get_chat_template_text(self) -> str | None:
|
||
"""Get chat template text from the tokenizer or model directory."""
|
||
# Try tokenizer's chat_template attribute (Jinja string)
|
||
ct = getattr(self.tokenizer, "_chat_template", None)
|
||
if ct:
|
||
return ct if isinstance(ct, str) else str(ct)
|
||
ct = getattr(self.tokenizer, "chat_template", None)
|
||
if ct:
|
||
return ct if isinstance(ct, str) else str(ct)
|
||
|
||
# Try reading the .jinja file from model directory
|
||
import os
|
||
|
||
model_path = getattr(self.config, "model_path", None) or ""
|
||
jinja_path = os.path.join(model_path, "chat_template.jinja")
|
||
if os.path.isfile(jinja_path):
|
||
try:
|
||
with open(jinja_path, encoding="utf-8") as f:
|
||
return f.read()
|
||
except Exception:
|
||
pass
|
||
|
||
return None
|
||
|
||
def _detect_needs_think_prefix(self, request: "Request") -> bool:
|
||
"""Detect if prompt ends with an open <think> tag (thinking enabled).
|
||
|
||
Returns False for disabled-thinking patterns like <think></think>
|
||
where </think> immediately follows <think> in the prompt tail.
|
||
"""
|
||
think_start_ids = None
|
||
think_start_id = self._get_think_token_id("think_start_id")
|
||
if think_start_id is not None:
|
||
think_start_ids = [think_start_id]
|
||
else:
|
||
think_start_text = (
|
||
self._get_output_parser_thinking_start_text() or "<think>"
|
||
)
|
||
try:
|
||
token_id = self.tokenizer.convert_tokens_to_ids(think_start_text)
|
||
if isinstance(token_id, int) and token_id != getattr(
|
||
self.tokenizer, "unk_token_id", None
|
||
):
|
||
think_start_ids = [token_id]
|
||
except (AttributeError, KeyError, TypeError):
|
||
think_start_ids = None
|
||
|
||
if think_start_ids is None:
|
||
think_start_ids = self._encode_thinking_marker(think_start_text)
|
||
|
||
if not think_start_ids or not request.prompt_token_ids:
|
||
return False
|
||
|
||
lookback = max(3, len(think_start_ids) + 2)
|
||
last_tokens = list(request.prompt_token_ids[-lookback:])
|
||
last_idx = None
|
||
for idx in range(len(last_tokens) - len(think_start_ids), -1, -1):
|
||
if last_tokens[idx : idx + len(think_start_ids)] == think_start_ids:
|
||
last_idx = idx
|
||
break
|
||
if last_idx is None:
|
||
return False
|
||
|
||
# <think> found. Check if </think> follows it (disabled thinking pattern).
|
||
after_start = last_tokens[last_idx + len(think_start_ids) :]
|
||
|
||
if after_start:
|
||
think_end_ids = self._resolve_think_end_token_ids()
|
||
if think_end_ids and len(after_start) >= len(think_end_ids):
|
||
for idx in range(len(after_start) - len(think_end_ids) + 1):
|
||
if after_start[idx : idx + len(think_end_ids)] == think_end_ids:
|
||
return False
|
||
elif think_end_ids and think_end_ids[0] in after_start:
|
||
return False
|
||
|
||
return True
|
||
|
||
def _ensure_batch_generator(self, sampling_params: SamplingParams) -> None:
|
||
"""Ensure BatchGenerator exists with compatible settings."""
|
||
# Only create once; per-request samplers are passed at insert time.
|
||
if self.batch_generator is None:
|
||
self.batch_generator = self._create_batch_generator(sampling_params)
|
||
|
||
# Track latest params for debugging/metrics.
|
||
self._current_sampler_params = (
|
||
sampling_params.temperature,
|
||
sampling_params.top_p,
|
||
sampling_params.min_p,
|
||
sampling_params.top_k,
|
||
sampling_params.repetition_penalty,
|
||
)
|
||
|
||
def _cache_tree_has_stateful_non_sliceable(self, cache_obj: Any) -> bool:
|
||
"""Detect non-sliceable recurrent cache layers requiring snapshots."""
|
||
# None placeholders from boundary snapshots (sliceable layers replaced).
|
||
if cache_obj is None:
|
||
return False
|
||
|
||
# CacheList nests multiple cache objects.
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
return any(
|
||
self._cache_tree_has_stateful_non_sliceable(sub_cache)
|
||
for sub_cache in sub_caches
|
||
)
|
||
|
||
class_name = type(cache_obj).__name__
|
||
|
||
# Known sliceable cache types — no boundary snapshots needed.
|
||
if class_name in (
|
||
"KVCache",
|
||
"BatchKVCache",
|
||
"QuantizedKVCache",
|
||
):
|
||
return False
|
||
|
||
# Stateful non-sliceable caches require boundary-safe snapshots.
|
||
is_rotating_cache = class_name in (
|
||
"RotatingKVCache",
|
||
"BatchRotatingKVCache",
|
||
)
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None:
|
||
is_rotating_cache = (
|
||
is_rotating_cache or CacheTypeRegistry.is_rotating_family(class_name)
|
||
)
|
||
|
||
if is_rotating_cache or class_name in ("ArraysCache", "SizedArraysCache"):
|
||
return True
|
||
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None:
|
||
handler = CacheTypeRegistry.get_handler_by_class_name(class_name)
|
||
if not handler.supports_block_slicing:
|
||
return True
|
||
|
||
# Best-effort fallback for unknown recurrent cache structures.
|
||
state_list = getattr(cache_obj, "cache", None)
|
||
if isinstance(state_list, list):
|
||
return True
|
||
|
||
return False
|
||
|
||
def _cache_list_needs_boundary_snapshot(self, cache_list: list[Any]) -> bool:
|
||
"""Return True if any layer cache requires boundary snapshots."""
|
||
if not cache_list:
|
||
return False
|
||
return any(
|
||
self._cache_tree_has_stateful_non_sliceable(layer_cache)
|
||
for layer_cache in cache_list
|
||
)
|
||
|
||
def _eval_snapshot_cache(self, snapshot_cache: list[Any]) -> None:
|
||
"""Force the leaf KV tensors of an in-memory boundary snapshot concrete.
|
||
|
||
Runs on the capturing (owner/inference) thread. The store-cache worker
|
||
later re-extracts and slices these via _BoundarySnapshotProvider; MLX
|
||
streams are thread-local, so any leaf op still lazy at that point would
|
||
re-dispatch to this thread's stream index, which does not exist on the
|
||
worker -> SIGABRT. Extracting + eval'ing the leaves here (under this
|
||
thread's stream, where that stream lives) makes the worker's slicing
|
||
operate exclusively on already-materialized buffers, which bind their
|
||
new ops to the always-present default stream (gpu,0).
|
||
"""
|
||
if not snapshot_cache:
|
||
return
|
||
extracted, _ = self._extract_cache_states(snapshot_cache)
|
||
leaves = self._collect_arrays_from_extracted_cache(extracted)
|
||
if leaves:
|
||
with mx.stream(self._stream):
|
||
mx.eval(*leaves)
|
||
|
||
def _on_prefill_boundary_snapshot(
|
||
self,
|
||
request_id: str,
|
||
snapshot_cache: list[Any],
|
||
token_count: int,
|
||
) -> None:
|
||
"""Record boundary snapshots captured during prefill processing.
|
||
|
||
Called from ``_emit_prefill_boundary_snapshot`` at each block
|
||
boundary crossed during prefill. Keyed by ``request_id`` rather
|
||
than ``uid`` because the request has not been inserted into
|
||
``BatchGenerator`` yet and the uid mapping does not exist —
|
||
routing through it dropped every snapshot silently (#TBD).
|
||
"""
|
||
if self.block_aware_cache is None:
|
||
return
|
||
|
||
block_size = self.config.paged_cache_block_size
|
||
if block_size <= 0 or token_count <= 0 or token_count % block_size != 0:
|
||
return
|
||
|
||
if not self._cache_list_needs_boundary_snapshot(snapshot_cache):
|
||
return
|
||
|
||
if request_id not in self._boundary_cache_snapshots:
|
||
self._boundary_cache_snapshots[request_id] = {}
|
||
|
||
# Skip if we already have a snapshot at this token count
|
||
if token_count in self._boundary_cache_snapshots[request_id]:
|
||
return
|
||
|
||
# Offload snapshot to SSD if store is available, keeping only a
|
||
# None marker in the dict. Falls back to in-memory storage when
|
||
# the SSD store is unavailable or the write fails.
|
||
if self._boundary_snapshot_store is not None:
|
||
saved = self._boundary_snapshot_store.save(
|
||
request_id,
|
||
token_count,
|
||
snapshot_cache,
|
||
self._extract_cache_states,
|
||
)
|
||
if saved:
|
||
self._boundary_cache_snapshots[request_id][token_count] = None
|
||
else:
|
||
# In-memory fallback: this snapshot will be sliced later on the
|
||
# store-cache worker thread (via _BoundarySnapshotProvider ->
|
||
# _extract_cache_states). MLX streams are thread-local, so the
|
||
# worker cannot materialize a lazy op bound to THIS (owner) thread's
|
||
# stream. Force it concrete now, on the capturing thread, so the
|
||
# worker only ever slices already-evaluated buffers.
|
||
self._eval_snapshot_cache(snapshot_cache)
|
||
self._boundary_cache_snapshots[request_id][token_count] = snapshot_cache
|
||
else:
|
||
self._eval_snapshot_cache(snapshot_cache)
|
||
self._boundary_cache_snapshots[request_id][token_count] = snapshot_cache
|
||
|
||
self._boundary_snapshot_required = True
|
||
logger.debug(
|
||
"Captured prefill boundary cache snapshot for %s at %s tokens",
|
||
request_id,
|
||
token_count,
|
||
)
|
||
|
||
def _detect_boundary_snapshot_need(self) -> bool:
|
||
"""
|
||
Determine whether boundary snapshots are needed for the current model.
|
||
|
||
Evaluated lazily by inspecting model.make_cache() output instead of
|
||
the active batch (which no longer exists in the new API).
|
||
"""
|
||
if self._boundary_snapshot_required is not None:
|
||
return self._boundary_snapshot_required
|
||
|
||
if not hasattr(self.model, "make_cache"):
|
||
self._boundary_snapshot_required = False
|
||
return False
|
||
|
||
try:
|
||
cache_list = self.model.make_cache()
|
||
except Exception:
|
||
self._boundary_snapshot_required = False
|
||
return False
|
||
|
||
if not cache_list:
|
||
self._boundary_snapshot_required = False
|
||
return False
|
||
|
||
try:
|
||
self._boundary_snapshot_required = any(
|
||
self._cache_tree_has_stateful_non_sliceable(layer_cache)
|
||
for layer_cache in cache_list
|
||
)
|
||
except TypeError:
|
||
# make_cache() returned something non-iterable (stub models in
|
||
# tests); treat as not snapshot-needing.
|
||
self._boundary_snapshot_required = False
|
||
return False
|
||
|
||
if self._boundary_snapshot_required:
|
||
logger.info(
|
||
"Enabled boundary cache snapshots for stateful non-sliceable "
|
||
"cache layers"
|
||
)
|
||
# Ask the speculative (MTP) decode path to land its commits on
|
||
# block boundaries: its emit queue otherwise leaves the cache a
|
||
# few tokens ahead of the emitted count when a boundary token
|
||
# surfaces, which forces the consistency guard in
|
||
# _extract_boundary_snapshot to skip most captures.
|
||
block = int(self.config.paged_cache_block_size or 0)
|
||
if block > 0:
|
||
try:
|
||
self.model._omlx_mtp_commit_align = block
|
||
except Exception:
|
||
pass
|
||
else:
|
||
logger.debug(
|
||
"Boundary cache snapshots disabled (no stateful non-sliceable "
|
||
"cache layers detected)"
|
||
)
|
||
|
||
return self._boundary_snapshot_required
|
||
|
||
def _extract_boundary_snapshot(
|
||
self, uid: int, expected_tokens: int | None = None
|
||
) -> list[Any] | None:
|
||
"""Extract a per-request prompt cache snapshot via extract_cache().
|
||
|
||
Uses BatchGenerator.extract_cache() which returns
|
||
Dict[uid, (cache_list, tokens_list)].
|
||
|
||
``expected_tokens`` guards positional consistency: a snapshot labeled
|
||
"state at N tokens" must be extracted while the cache holds exactly N
|
||
forwarded tokens. The standard decode step always satisfies this at
|
||
emit time, but speculative (MTP) decode advances the cache in bursts
|
||
and emits from a queue, so the cache can be a few tokens ahead of —
|
||
or one behind — the emitted count when the boundary token surfaces.
|
||
A skewed snapshot would pair block-aligned KV with recurrent (SSM)
|
||
state from a different position and corrupt later prefix-cache hits
|
||
on hybrid models; skipping the capture merely costs a reuse
|
||
opportunity.
|
||
"""
|
||
if self.batch_generator is None:
|
||
return None
|
||
|
||
try:
|
||
# Synchronize pending engine stream operations before
|
||
# accessing batch cache tensors.
|
||
with self._phase_timer("boundary_capture_sync"):
|
||
_safe_sync_stream(self._stream)
|
||
with self._phase_timer("boundary_capture_extract"):
|
||
with mx.stream(self._stream):
|
||
result = self.batch_generator.extract_cache([uid])
|
||
if uid not in result:
|
||
return None
|
||
cache_list, _tokens = result[uid]
|
||
if expected_tokens is not None:
|
||
# Walk into CacheList wrappers: DeepSeek-V4/GLM layer
|
||
# caches carry their token offset on a sub-cache, and
|
||
# checking only the wrapper would silently pass a
|
||
# skewed capture through.
|
||
for c in cache_list:
|
||
offset = _first_leaf_cache_offset(c)
|
||
if offset is None:
|
||
continue
|
||
if offset != expected_tokens:
|
||
logger.debug(
|
||
"Skipping boundary snapshot for uid=%s: "
|
||
"cache offset %d != boundary %d "
|
||
"(speculative decode skew)",
|
||
uid,
|
||
offset,
|
||
expected_tokens,
|
||
)
|
||
return None
|
||
break
|
||
# Only extract non-sliceable layers to avoid costly
|
||
# deep-copy accumulation (same rationale as prefill path).
|
||
return [
|
||
(
|
||
c
|
||
if type(c).__name__ not in _KNOWN_SLICEABLE_CACHE_TYPES
|
||
else None
|
||
)
|
||
for c in cache_list
|
||
]
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"Failed to extract boundary cache snapshot for uid={uid}: {e}"
|
||
)
|
||
return None
|
||
|
||
def _maybe_capture_boundary_snapshot(self, request: Request, uid: int) -> None:
|
||
"""Capture cache snapshot exactly at block boundaries for safe reuse."""
|
||
if self.block_aware_cache is None:
|
||
return
|
||
|
||
block_size = self.config.paged_cache_block_size
|
||
if block_size <= 0:
|
||
return
|
||
|
||
total_tokens = request.num_tokens
|
||
if total_tokens <= 0 or total_tokens % block_size != 0:
|
||
return
|
||
|
||
if not self._detect_boundary_snapshot_need():
|
||
return
|
||
|
||
snapshot_cache = self._extract_boundary_snapshot(
|
||
uid, expected_tokens=total_tokens
|
||
)
|
||
if not snapshot_cache:
|
||
return
|
||
|
||
if request.request_id not in self._boundary_cache_snapshots:
|
||
self._boundary_cache_snapshots[request.request_id] = {}
|
||
|
||
# Offload to SSD with in-memory fallback.
|
||
if self._boundary_snapshot_store is not None:
|
||
with self._phase_timer("boundary_snapshot_save"):
|
||
saved = self._boundary_snapshot_store.save(
|
||
request.request_id,
|
||
total_tokens,
|
||
snapshot_cache,
|
||
self._extract_cache_states,
|
||
)
|
||
if saved:
|
||
self._boundary_cache_snapshots[request.request_id][total_tokens] = None
|
||
else:
|
||
# In-memory fallback: the store-cache worker slices this snapshot
|
||
# off-thread (via _BoundarySnapshotProvider -> _extract_cache_states).
|
||
# MLX streams are thread-local, so force the leaves concrete now on
|
||
# the capturing thread; otherwise the worker re-dispatches a lazy op
|
||
# to this thread's stream index -> "no Stream(gpu, N)" -> SIGABRT.
|
||
# Mirrors _on_prefill_boundary_snapshot's in-memory fallback.
|
||
self._eval_snapshot_cache(snapshot_cache)
|
||
self._boundary_cache_snapshots[request.request_id][
|
||
total_tokens
|
||
] = snapshot_cache
|
||
else:
|
||
self._eval_snapshot_cache(snapshot_cache)
|
||
self._boundary_cache_snapshots[request.request_id][
|
||
total_tokens
|
||
] = snapshot_cache
|
||
|
||
logger.debug(
|
||
f"Captured boundary cache snapshot for {request.request_id} at "
|
||
f"{total_tokens} tokens"
|
||
)
|
||
|
||
def _get_boundary_store_override(
|
||
self,
|
||
request_id: str,
|
||
full_token_sequence: list[int],
|
||
) -> (
|
||
tuple[
|
||
list[int],
|
||
list[dict[str, Any]],
|
||
Optional["ModelCacheConfig"],
|
||
dict[int, list[dict[str, Any]]],
|
||
]
|
||
| None
|
||
):
|
||
"""
|
||
Return boundary-aligned cache payload when final request ends on partial block.
|
||
|
||
Returns:
|
||
Tuple of (truncated_tokens, extracted_cache, model_cache_config,
|
||
intermediate_snapshots) where intermediate_snapshots maps
|
||
token_count -> extracted cache states for per-block storage.
|
||
"""
|
||
snapshots = self._boundary_cache_snapshots.get(request_id)
|
||
if not snapshots:
|
||
return None
|
||
|
||
total_tokens = len(full_token_sequence)
|
||
block_size = self.config.paged_cache_block_size
|
||
|
||
# Find all valid boundary-aligned snapshot token counts
|
||
valid_counts = sorted(
|
||
tc
|
||
for tc in snapshots.keys()
|
||
if 0 < tc <= total_tokens and tc % block_size == 0
|
||
)
|
||
if not valid_counts:
|
||
return None
|
||
|
||
# Find the latest snapshot that leaves trailing partial tokens
|
||
# (or equals total if it's block-aligned).
|
||
latest_tc = valid_counts[-1]
|
||
if latest_tc < total_tokens:
|
||
# Trailing partial tokens exist — use this snapshot for truncation
|
||
pass
|
||
elif latest_tc == total_tokens and total_tokens % block_size == 0:
|
||
# Exactly block-aligned — no truncation needed but we still
|
||
# provide intermediate snapshots for per-block storage.
|
||
latest_tc = total_tokens
|
||
else:
|
||
return None
|
||
|
||
# Load latest snapshot — may be on SSD (None marker) or in memory.
|
||
#
|
||
# In-memory snapshots are raw mlx-lm cache objects. They must be
|
||
# converted to the extracted dict format here on the engine MLX thread.
|
||
# The async store-cache worker does not own the generation stream; if it
|
||
# touches raw Rotating/Arrays cache state, MLX can abort with
|
||
# "There is no Stream(gpu, X) in current thread" (#1568).
|
||
latest_snapshot = snapshots[latest_tc]
|
||
if latest_snapshot is None and self._boundary_snapshot_store is not None:
|
||
# Offloaded to SSD — load back.
|
||
extracted_cache = self._boundary_snapshot_store.load(request_id, latest_tc)
|
||
if not extracted_cache:
|
||
return None
|
||
# Build model_cache_config from the main request cache config
|
||
# since the SSD snapshot doesn't carry it.
|
||
model_cache_config = getattr(
|
||
self.requests.get(request_id), "_model_cache_config", None
|
||
)
|
||
elif latest_snapshot is not None:
|
||
extracted_cache, model_cache_config = self._extract_cache_states(
|
||
latest_snapshot
|
||
)
|
||
if not extracted_cache:
|
||
return None
|
||
else:
|
||
return None
|
||
|
||
# Build provider for intermediate snapshots. SSD-backed snapshots remain
|
||
# lazy-loaded, but in-memory snapshots are extracted eagerly on this
|
||
# engine thread before the provider is handed to the async worker.
|
||
intermediate_tcs = [tc for tc in valid_counts if tc != latest_tc]
|
||
provider_tcs: list[int] = []
|
||
extracted_in_memory: dict[int, list[dict[str, Any]]] = {}
|
||
for tc in intermediate_tcs:
|
||
snap = snapshots.get(tc)
|
||
if snap is None:
|
||
if self._boundary_snapshot_store is not None:
|
||
provider_tcs.append(tc)
|
||
continue
|
||
|
||
extracted_snapshot, _ = self._extract_cache_states(snap)
|
||
if extracted_snapshot:
|
||
extracted_in_memory[tc] = extracted_snapshot
|
||
provider_tcs.append(tc)
|
||
|
||
intermediate_snapshots = _BoundarySnapshotProvider(
|
||
store=self._boundary_snapshot_store,
|
||
request_id=request_id,
|
||
valid_tcs=provider_tcs,
|
||
in_memory_snapshots=extracted_in_memory,
|
||
)
|
||
|
||
token_sequence = (
|
||
full_token_sequence[:latest_tc]
|
||
if latest_tc < total_tokens
|
||
else full_token_sequence
|
||
)
|
||
|
||
return (
|
||
token_sequence,
|
||
extracted_cache,
|
||
model_cache_config,
|
||
intermediate_snapshots,
|
||
)
|
||
|
||
@staticmethod
|
||
def _merge_boundary_with_full_cache(
|
||
boundary_cache: list[dict[str, Any]],
|
||
full_cache: list[dict[str, Any]],
|
||
) -> list[dict[str, Any]]:
|
||
"""Fill placeholder layers in boundary cache from full extracted cache.
|
||
|
||
Boundary snapshots skip sliceable (KVCache) layers to save memory,
|
||
leaving them as ``{'state': (), ...}`` placeholders. For block
|
||
storage the KV tensors are needed, so we copy them from the full
|
||
extracted cache (which contains the complete sequence).
|
||
"""
|
||
if not full_cache or len(boundary_cache) != len(full_cache):
|
||
return boundary_cache
|
||
|
||
merged = []
|
||
for bc, fc in zip(boundary_cache, full_cache):
|
||
state = bc.get("state", ())
|
||
# Placeholder layers have state == () (empty tuple).
|
||
if isinstance(state, tuple) and len(state) == 0:
|
||
# Take full cache layer instead.
|
||
merged.append(fc)
|
||
else:
|
||
merged.append(bc)
|
||
return merged
|
||
|
||
@staticmethod
|
||
def _is_empty_boundary_placeholder(layer_state: Any) -> bool:
|
||
if not isinstance(layer_state, dict):
|
||
return False
|
||
state = layer_state.get("state", ())
|
||
return isinstance(state, tuple) and len(state) == 0
|
||
|
||
@staticmethod
|
||
def _extracted_layer_type_name(layer_state: dict[str, Any]) -> str:
|
||
class_name = str(layer_state.get("class_name") or "")
|
||
if class_name:
|
||
return class_name
|
||
return str(layer_state.get("cache_type") or "")
|
||
|
||
@staticmethod
|
||
def _is_sliceable_extracted_layer(layer_state: dict[str, Any]) -> bool:
|
||
type_name = Scheduler._extracted_layer_type_name(layer_state)
|
||
if type_name in _KNOWN_SLICEABLE_CACHE_TYPES:
|
||
return True
|
||
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None and type_name:
|
||
try:
|
||
if type_name not in CacheTypeRegistry.list_known_class_names():
|
||
return False
|
||
handler = CacheTypeRegistry.get_handler_by_class_name(type_name)
|
||
return bool(handler.supports_block_slicing)
|
||
except Exception:
|
||
return False
|
||
|
||
return False
|
||
|
||
def _fill_boundary_placeholders_from_live_cache(
|
||
self,
|
||
boundary_cache: list[dict[str, Any]],
|
||
live_cache: list[dict[str, Any]],
|
||
) -> list[dict[str, Any]] | None:
|
||
"""Fill boundary placeholders only with proven sliceable live layers."""
|
||
if (
|
||
not boundary_cache
|
||
or not live_cache
|
||
or len(boundary_cache) != len(live_cache)
|
||
):
|
||
return None
|
||
|
||
merged: list[dict[str, Any]] = []
|
||
for layer_idx, (boundary_layer, live_layer) in enumerate(
|
||
zip(boundary_cache, live_cache)
|
||
):
|
||
if not self._is_empty_boundary_placeholder(boundary_layer):
|
||
merged.append(boundary_layer)
|
||
continue
|
||
|
||
if not isinstance(
|
||
live_layer, dict
|
||
) or not self._is_sliceable_extracted_layer(live_layer):
|
||
logger.debug(
|
||
"Cannot fill boundary placeholder for layer %s from non-sliceable "
|
||
"live cache type %s",
|
||
layer_idx,
|
||
(
|
||
self._extracted_layer_type_name(live_layer)
|
||
if isinstance(live_layer, dict)
|
||
else type(live_layer).__name__
|
||
),
|
||
)
|
||
return None
|
||
|
||
merged.append(live_layer)
|
||
|
||
if any(self._is_empty_boundary_placeholder(layer) for layer in merged):
|
||
return None
|
||
return merged
|
||
|
||
def _extract_live_request_cache_for_store(
|
||
self,
|
||
request_id: str,
|
||
uid: int,
|
||
expected_tokens: list[int],
|
||
) -> tuple[list[dict[str, Any]], Optional["ModelCacheConfig"]] | None:
|
||
"""Extract live cache only when its token prefix matches exactly."""
|
||
if self.batch_generator is None or uid is None or uid < 0:
|
||
return None
|
||
|
||
try:
|
||
_safe_sync_stream(self._stream)
|
||
with mx.stream(self._stream):
|
||
result = self.batch_generator.extract_cache([uid])
|
||
if uid not in result:
|
||
logger.debug(
|
||
"Cannot extract live cache for %s: uid %s not present",
|
||
request_id,
|
||
uid,
|
||
)
|
||
return None
|
||
|
||
live_cache, live_tokens = result[uid]
|
||
live_tokens_list = list(live_tokens) if live_tokens is not None else []
|
||
if len(live_tokens_list) < len(expected_tokens) or (
|
||
live_tokens_list[: len(expected_tokens)] != expected_tokens
|
||
):
|
||
logger.debug(
|
||
"Skipping parser-stop cache store for %s: live cache tokens do "
|
||
"not match prompt boundary prefix (%s/%s tokens)",
|
||
request_id,
|
||
min(len(live_tokens_list), len(expected_tokens)),
|
||
len(expected_tokens),
|
||
)
|
||
return None
|
||
|
||
extracted_cache, model_cache_config = self._extract_cache_states(live_cache)
|
||
if not extracted_cache:
|
||
return None
|
||
return extracted_cache, model_cache_config
|
||
except Exception as e:
|
||
logger.debug(
|
||
"Failed to extract live cache for parser-stop cache store %s: %s",
|
||
request_id,
|
||
e,
|
||
)
|
||
return None
|
||
|
||
def _prepare_prompt_boundary_cache_store(
|
||
self,
|
||
request_id: str,
|
||
request: Request,
|
||
uid: int,
|
||
) -> (
|
||
tuple[
|
||
list[int],
|
||
list[dict[str, Any]],
|
||
Optional["ModelCacheConfig"],
|
||
Any | None,
|
||
]
|
||
| None
|
||
):
|
||
"""Prepare a prompt-only cache payload when a scheduler stop skipped it.
|
||
|
||
Parser-side stops (for example tool-call end markers) can finish a request
|
||
after a normal streaming token response. That response has no
|
||
``prompt_cache``, so the usual final-response cache extraction never runs.
|
||
This fallback stores only prompt tokens up to a prefill block boundary,
|
||
never generated output tokens.
|
||
"""
|
||
if self.block_aware_cache is None:
|
||
return None
|
||
if request.specprefill_indices is not None:
|
||
return None
|
||
|
||
block_size = self.config.paged_cache_block_size
|
||
if block_size <= 0:
|
||
return None
|
||
|
||
prompt_tokens = list(request.prompt_token_ids or [])
|
||
boundary_len = (len(prompt_tokens) // block_size) * block_size
|
||
if boundary_len <= 0:
|
||
return None
|
||
|
||
token_sequence = prompt_tokens[:boundary_len]
|
||
|
||
boundary_override = self._get_boundary_store_override(request_id, prompt_tokens)
|
||
if boundary_override is not None:
|
||
(
|
||
token_sequence,
|
||
boundary_cache,
|
||
boundary_model_config,
|
||
intermediate_snapshots,
|
||
) = boundary_override
|
||
|
||
live_payload = None
|
||
if any(
|
||
self._is_empty_boundary_placeholder(layer) for layer in boundary_cache
|
||
):
|
||
live_payload = self._extract_live_request_cache_for_store(
|
||
request_id,
|
||
uid,
|
||
token_sequence,
|
||
)
|
||
if live_payload is None:
|
||
return None
|
||
live_cache, live_model_config = live_payload
|
||
cache_to_store = self._fill_boundary_placeholders_from_live_cache(
|
||
boundary_cache,
|
||
live_cache,
|
||
)
|
||
if cache_to_store is None:
|
||
return None
|
||
model_cache_config = boundary_model_config or live_model_config
|
||
else:
|
||
cache_to_store = boundary_cache
|
||
model_cache_config = boundary_model_config
|
||
|
||
logger.info(
|
||
"Using prompt boundary cache snapshot for %s: storing %s/%s prompt "
|
||
"tokens after scheduler-side stop (skipping output tokens, %s "
|
||
"intermediate snapshots)",
|
||
request_id,
|
||
len(token_sequence),
|
||
len(prompt_tokens),
|
||
len(intermediate_snapshots) if intermediate_snapshots else 0,
|
||
)
|
||
return (
|
||
token_sequence,
|
||
cache_to_store,
|
||
model_cache_config,
|
||
intermediate_snapshots,
|
||
)
|
||
|
||
# Pure sliceable cache models do not need boundary snapshots. For models
|
||
# that do need snapshots, a missing snapshot means the non-sliceable state
|
||
# at this boundary is unavailable, so skip rather than storing an unsafe
|
||
# live decode tail.
|
||
if self._detect_boundary_snapshot_need():
|
||
return None
|
||
|
||
live_payload = self._extract_live_request_cache_for_store(
|
||
request_id,
|
||
uid,
|
||
token_sequence,
|
||
)
|
||
if live_payload is None:
|
||
return None
|
||
|
||
live_cache, live_model_config = live_payload
|
||
if not all(self._is_sliceable_extracted_layer(layer) for layer in live_cache):
|
||
logger.debug(
|
||
"Skipping parser-stop cache store for %s: live cache has "
|
||
"non-sliceable layers but no boundary snapshot",
|
||
request_id,
|
||
)
|
||
return None
|
||
|
||
logger.info(
|
||
"Using live prompt cache for %s: storing %s/%s prompt tokens after "
|
||
"scheduler-side stop (skipping output tokens)",
|
||
request_id,
|
||
len(token_sequence),
|
||
len(prompt_tokens),
|
||
)
|
||
return token_sequence, live_cache, live_model_config, None
|
||
|
||
def _validate_cache(self, cache: Any) -> bool:
|
||
"""
|
||
Validate that a cache object is usable.
|
||
|
||
This prevents NoneType errors when mlx-lm's BatchKVCache
|
||
contains invalid/stale references.
|
||
|
||
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
|
||
# RotatingKVCache may have keys=None (legacy) or zero-length
|
||
# keys (hybrid window padding). Both are valid empty states
|
||
# that will be filled during padding reprocessing.
|
||
if hasattr(layer_cache, "keys") and layer_cache.keys is None:
|
||
if hasattr(layer_cache, "max_size"):
|
||
continue # Valid empty RotatingKVCache (keys=None)
|
||
return False
|
||
if hasattr(layer_cache, "values") and layer_cache.values is None:
|
||
if hasattr(layer_cache, "max_size"):
|
||
continue # Valid empty RotatingKVCache (values=None)
|
||
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
|
||
|
||
def _normalize_rotating_snapshot_state(
|
||
self,
|
||
layer_cache: Any,
|
||
state: tuple[Any, Any],
|
||
meta_state: Any,
|
||
layer_idx: int | None = None,
|
||
) -> tuple[tuple[Any, Any], tuple[str, str, str, str]]:
|
||
"""
|
||
Normalize RotatingKVCache state into merge-safe canonical form.
|
||
|
||
Boundary snapshots captured mid-prefill can expose oversized rotating
|
||
buffers (e.g., max_size + chunk_size - 1). Those states are valid for
|
||
in-flight prefill but break BatchRotatingKVCache.merge() after SSD
|
||
restore because merge expects per-request rotating buffers capped to
|
||
max_size. This method canonicalizes to the latest max_size tokens.
|
||
"""
|
||
if not isinstance(state, (list, tuple)) or len(state) < 2:
|
||
return state, (
|
||
tuple(meta_state) if isinstance(meta_state, (list, tuple)) else ()
|
||
)
|
||
|
||
keys = state[0]
|
||
values = state[1]
|
||
if keys is None or values is None or not hasattr(keys, "shape"):
|
||
return state, (
|
||
tuple(meta_state) if isinstance(meta_state, (list, tuple)) else ()
|
||
)
|
||
|
||
try:
|
||
keep = (
|
||
int(meta_state[0])
|
||
if meta_state and len(meta_state) >= 1
|
||
else int(getattr(layer_cache, "keep", 0))
|
||
)
|
||
max_size = (
|
||
int(meta_state[1])
|
||
if meta_state and len(meta_state) >= 2
|
||
else int(getattr(layer_cache, "max_size", keys.shape[2]))
|
||
)
|
||
offset = (
|
||
int(meta_state[2])
|
||
if meta_state and len(meta_state) >= 3
|
||
else int(getattr(layer_cache, "offset", keys.shape[2]))
|
||
)
|
||
idx = (
|
||
int(meta_state[3])
|
||
if meta_state and len(meta_state) >= 4
|
||
else int(getattr(layer_cache, "_idx", keys.shape[2]))
|
||
)
|
||
except Exception:
|
||
return state, (
|
||
tuple(meta_state) if isinstance(meta_state, (list, tuple)) else ()
|
||
)
|
||
|
||
ordered_keys = keys
|
||
ordered_values = values
|
||
temporal_order = getattr(layer_cache, "_temporal_order", None)
|
||
if callable(temporal_order):
|
||
try:
|
||
ordered_keys = temporal_order(keys)
|
||
ordered_values = temporal_order(values)
|
||
except Exception:
|
||
ordered_keys = keys
|
||
ordered_values = values
|
||
|
||
original_len = int(ordered_keys.shape[2]) if len(ordered_keys.shape) >= 3 else 0
|
||
normalized_keys = ordered_keys
|
||
normalized_values = ordered_values
|
||
|
||
if max_size > 0 and original_len > max_size:
|
||
if keep > 0 and keep < max_size:
|
||
tail_len = max_size - keep
|
||
normalized_keys = mx.concatenate(
|
||
[
|
||
ordered_keys[..., :keep, :],
|
||
ordered_keys[..., -tail_len:, :],
|
||
],
|
||
axis=2,
|
||
)
|
||
normalized_values = mx.concatenate(
|
||
[
|
||
ordered_values[..., :keep, :],
|
||
ordered_values[..., -tail_len:, :],
|
||
],
|
||
axis=2,
|
||
)
|
||
else:
|
||
normalized_keys = ordered_keys[..., -max_size:, :]
|
||
normalized_values = ordered_values[..., -max_size:, :]
|
||
|
||
try:
|
||
normalized_keys = mx.contiguous(normalized_keys)
|
||
normalized_values = mx.contiguous(normalized_values)
|
||
except Exception:
|
||
pass
|
||
|
||
normalized_len = (
|
||
int(normalized_keys.shape[2]) if len(normalized_keys.shape) >= 3 else 0
|
||
)
|
||
# Force case 1 of _temporal_order: _idx == keys.shape[2] means the
|
||
# buffer is already in temporal order (which is exactly what the
|
||
# oversized trim above produces — the contiguous tail of the most
|
||
# recent tokens). Anything else lets _temporal_order re-slice the
|
||
# buffer in the rotated branch (case 2), which is wasted work and
|
||
# obscures the merge contract. See cache.py:431-447 for the branches.
|
||
normalized_idx = normalized_len
|
||
|
||
normalized_meta = (
|
||
str(keep),
|
||
str(max_size),
|
||
str(offset),
|
||
str(normalized_idx),
|
||
)
|
||
|
||
if original_len != normalized_len or idx != normalized_idx:
|
||
layer_tag = f"layer {layer_idx}: " if layer_idx is not None else ""
|
||
logger.debug(
|
||
"%sNormalized RotatingKVCache snapshot: len %s->%s, idx %s->%s, "
|
||
"offset=%s, max_size=%s",
|
||
layer_tag,
|
||
original_len,
|
||
normalized_len,
|
||
idx,
|
||
normalized_idx,
|
||
offset,
|
||
max_size,
|
||
)
|
||
|
||
return (normalized_keys, normalized_values), normalized_meta
|
||
|
||
def _extract_cache_states(
|
||
self,
|
||
raw_cache: list[Any],
|
||
) -> tuple[list[dict[str, Any]], Optional["ModelCacheConfig"]]:
|
||
"""
|
||
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.
|
||
|
||
Also creates a ModelCacheConfig with per-layer type information to
|
||
support hybrid cache models (e.g., KVCache + ArraysCache).
|
||
|
||
Args:
|
||
raw_cache: List of cache objects from mlx-lm (KVCache, ArraysCache, etc.)
|
||
|
||
Returns:
|
||
Tuple of:
|
||
- List of dicts with {state, meta_state, class_name, cache_type}
|
||
- ModelCacheConfig with per-layer type information (or None)
|
||
"""
|
||
if not raw_cache:
|
||
return [], None
|
||
|
||
# Build ModelCacheConfig for type information.
|
||
# Skip if raw_cache contains None entries (boundary snapshots with
|
||
# sliceable layers replaced by None) — from_cache_list expects real
|
||
# cache objects and would log noisy NoneType warnings.
|
||
model_cache_config = None
|
||
has_none_layers = any(c is None for c in raw_cache)
|
||
if (
|
||
HAS_CACHE_TYPE_HANDLERS
|
||
and ModelCacheConfig is not None
|
||
and not has_none_layers
|
||
):
|
||
try:
|
||
model_cache_config = ModelCacheConfig.from_cache_list(
|
||
raw_cache,
|
||
model_name=self.model_name if hasattr(self, "model_name") else "",
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"Failed to build ModelCacheConfig: {e}")
|
||
|
||
extracted = []
|
||
for layer_idx, layer_cache in enumerate(raw_cache):
|
||
# Boundary snapshots may contain None for sliceable layers
|
||
# (KVCache) that were skipped during capture to save memory.
|
||
# Insert a placeholder to preserve layer index alignment.
|
||
if layer_cache is None:
|
||
extracted.append(
|
||
{
|
||
"state": (),
|
||
"meta_state": (),
|
||
"class_name": "KVCache",
|
||
"cache_type": "KVCache",
|
||
}
|
||
)
|
||
continue
|
||
try:
|
||
class_name = type(layer_cache).__name__
|
||
|
||
# Determine cache type using registry if available
|
||
cache_type_name = class_name
|
||
handler = None
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None:
|
||
try:
|
||
cache_type = CacheTypeRegistry.detect_cache_type(layer_cache)
|
||
cache_type_name = cache_type.value
|
||
handler = CacheTypeRegistry.get_handler(cache_type)
|
||
except Exception:
|
||
pass
|
||
|
||
# CacheList: composite cache with multiple sub-caches
|
||
if cache_type_name == "CacheList" or class_name == "CacheList":
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None:
|
||
try:
|
||
handler = CacheTypeRegistry.get_handler_by_class_name(
|
||
"CacheList"
|
||
)
|
||
state_dict = handler.extract_state(layer_cache)
|
||
sub_states = list(state_dict.get("sub_states", []))
|
||
sub_class_names = list(
|
||
state_dict.get("sub_class_names", [])
|
||
)
|
||
sub_meta_states = list(
|
||
state_dict.get("sub_meta_states", [])
|
||
)
|
||
|
||
sub_caches = getattr(layer_cache, "caches", ())
|
||
for sub_idx, sub_cache in enumerate(sub_caches):
|
||
if sub_idx >= len(sub_states):
|
||
break
|
||
|
||
sub_class_name = type(sub_cache).__name__
|
||
if sub_class_name in (
|
||
"RotatingKVCache",
|
||
"BatchRotatingKVCache",
|
||
"PrefillReadyRotatingKVCache",
|
||
):
|
||
normalized_state, normalized_meta = (
|
||
self._normalize_rotating_snapshot_state(
|
||
sub_cache,
|
||
sub_states[sub_idx],
|
||
(
|
||
sub_meta_states[sub_idx]
|
||
if sub_idx < len(sub_meta_states)
|
||
else getattr(
|
||
sub_cache, "meta_state", ()
|
||
)
|
||
),
|
||
layer_idx=layer_idx,
|
||
)
|
||
)
|
||
sub_states[sub_idx] = normalized_state
|
||
if sub_idx < len(sub_meta_states):
|
||
sub_meta_states[sub_idx] = normalized_meta
|
||
|
||
extracted.append(
|
||
{
|
||
"state": sub_states,
|
||
"meta_state": (
|
||
sub_class_names,
|
||
sub_meta_states,
|
||
),
|
||
"class_name": "CacheList",
|
||
"cache_type": "CacheList",
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"CacheList handler extraction failed: {e}")
|
||
extracted.append(
|
||
{
|
||
"state": [],
|
||
"meta_state": ([], []),
|
||
"class_name": "CacheList",
|
||
"cache_type": "CacheList",
|
||
}
|
||
)
|
||
else:
|
||
# Fallback: extract sub-cache state/meta without handlers
|
||
# MUST append to extracted to prevent layer count mismatch (Issue #1)
|
||
sub_caches = getattr(layer_cache, "caches", ())
|
||
sub_states = []
|
||
sub_class_names = []
|
||
sub_meta_states = []
|
||
for sc in sub_caches:
|
||
sub_states.append(sc.state if hasattr(sc, "state") else ())
|
||
sub_class_names.append(type(sc).__name__)
|
||
sub_meta_states.append(getattr(sc, "meta_state", ()))
|
||
extracted.append(
|
||
{
|
||
"state": sub_states,
|
||
"meta_state": (sub_class_names, sub_meta_states),
|
||
"class_name": "CacheList",
|
||
"cache_type": "CacheList",
|
||
}
|
||
)
|
||
continue
|
||
|
||
if hasattr(layer_cache, "state"):
|
||
if handler is not None and class_name in (
|
||
"MiniMaxM3KVCache",
|
||
"MiniMaxM3BatchKVCache",
|
||
):
|
||
state = handler.serialize_state(layer_cache)
|
||
meta = handler.serialize_meta_state(layer_cache)
|
||
else:
|
||
state = layer_cache.state
|
||
meta = getattr(layer_cache, "meta_state", ())
|
||
|
||
is_rotating_cache = class_name in (
|
||
"RotatingKVCache",
|
||
"BatchRotatingKVCache",
|
||
"PrefillReadyRotatingKVCache",
|
||
"BufferedRotatingKVCache",
|
||
)
|
||
if HAS_CACHE_TYPE_HANDLERS and CacheTypeRegistry is not None:
|
||
is_rotating_cache = (
|
||
is_rotating_cache
|
||
or CacheTypeRegistry.is_rotating_family(class_name)
|
||
)
|
||
if is_rotating_cache:
|
||
state, meta = self._normalize_rotating_snapshot_state(
|
||
layer_cache,
|
||
state,
|
||
meta,
|
||
layer_idx=layer_idx,
|
||
)
|
||
|
||
# Preserve the full state tuple regardless of length.
|
||
# Legacy 2-tuple caches (KVCache, RotatingKVCache, ...)
|
||
# surface as (keys, values); 3-tuple caches like
|
||
# PoolingCache surface as (buf_kv, buf_gate, pooled);
|
||
# 4-tuple caches like BatchKVCache surface with the
|
||
# extra offset/padding metadata. Downstream
|
||
# serialization (paged_ssd_cache, boundary_snapshot)
|
||
# is N-tuple aware after the cache architecture
|
||
# refactor — see Section 6 of the implementation
|
||
# plan.
|
||
if isinstance(state, (list, tuple)) and len(state) >= 1:
|
||
# Validate non-None for legacy KV-style caches only.
|
||
# PoolingCache's buf_kv may legitimately be None
|
||
# (fresh cache before any update), so skip the
|
||
# null guard for non-KV cache classes.
|
||
if (
|
||
class_name in ("KVCache", "RotatingKVCache", "BatchKVCache")
|
||
or (
|
||
HAS_CACHE_TYPE_HANDLERS
|
||
and CacheTypeRegistry is not None
|
||
and CacheTypeRegistry.is_rotating_family(class_name)
|
||
)
|
||
) and len(state) >= 2:
|
||
if state[0] is None or state[1] is None:
|
||
logger.debug(
|
||
f"Layer {layer_idx} ({class_name}) has None keys/values, "
|
||
f"skipping cache extraction"
|
||
)
|
||
return [], None # Return empty - cache is corrupted
|
||
|
||
extracted.append(
|
||
{
|
||
"state": tuple(state),
|
||
"meta_state": meta,
|
||
"class_name": class_name,
|
||
"cache_type": cache_type_name,
|
||
}
|
||
)
|
||
else:
|
||
# Unexpected state format (e.g. a non-tuple scalar).
|
||
logger.debug(
|
||
f"Layer {layer_idx} ({class_name}) has unexpected state format"
|
||
)
|
||
meta = getattr(layer_cache, "meta_state", ())
|
||
# Wrap the scalar so downstream code still gets a
|
||
# tuple-shaped state. This path is essentially dead
|
||
# in practice — kept defensive only.
|
||
extracted.append(
|
||
{
|
||
"state": (state,),
|
||
"meta_state": meta,
|
||
"class_name": class_name,
|
||
"cache_type": cache_type_name,
|
||
}
|
||
)
|
||
elif hasattr(layer_cache, "cache"):
|
||
# ArraysCache style: state stored in .cache list
|
||
cache_list = layer_cache.cache
|
||
if isinstance(cache_list, list) and len(cache_list) >= 2:
|
||
state = (cache_list[0], cache_list[1])
|
||
meta = getattr(layer_cache, "meta_state", ())
|
||
extracted.append(
|
||
{
|
||
"state": state,
|
||
"meta_state": meta,
|
||
"class_name": class_name,
|
||
"cache_type": cache_type_name,
|
||
}
|
||
)
|
||
else:
|
||
logger.debug(
|
||
f"Layer {layer_idx} ({class_name}) has invalid cache list"
|
||
)
|
||
continue
|
||
else:
|
||
logger.debug(
|
||
f"Layer {layer_idx} ({class_name}) has no state or cache attribute"
|
||
)
|
||
continue
|
||
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"Failed to extract state from cache layer {layer_idx}: {e}"
|
||
)
|
||
continue
|
||
|
||
if len(extracted) != len(raw_cache):
|
||
logger.debug(
|
||
f"Incomplete cache extraction: {len(extracted)}/{len(raw_cache)} layers"
|
||
)
|
||
return [], None
|
||
|
||
return extracted, model_cache_config
|
||
|
||
@staticmethod
|
||
def _common_prefix_len(a: list[int], b: list[int]) -> int:
|
||
n = min(len(a), len(b))
|
||
for i in range(n):
|
||
if a[i] != b[i]:
|
||
return i
|
||
return n
|
||
|
||
def _log_prefix_divergence(self, request: Request) -> None:
|
||
"""DEBUG-only prefix-cache miss diagnostic (issue #1003).
|
||
|
||
Compares the new prompt against recently stored cache sequences and
|
||
logs the first divergent token offset with decoded context on both
|
||
sides, so an always-miss report can be traced to the exact prompt
|
||
position (template re-render drift, client echo changes, eviction)
|
||
instead of guessing from hit counters.
|
||
"""
|
||
prompt = request.prompt_token_ids or []
|
||
if not prompt or not self._cache_probe_seqs:
|
||
return
|
||
best_id, best_seq, best_p = None, None, -1
|
||
for ref_id, seq in list(self._cache_probe_seqs):
|
||
p = self._common_prefix_len(prompt, seq)
|
||
if p > best_p:
|
||
best_id, best_seq, best_p = ref_id, seq, p
|
||
if best_seq is None:
|
||
return
|
||
cached = request.cached_tokens or 0
|
||
reusable = min(len(prompt), len(best_seq))
|
||
block = self.config.paged_cache_block_size
|
||
logger.debug(
|
||
f"Request {request.request_id}: prefix probe vs stored {best_id}: "
|
||
f"common_prefix={best_p}/{reusable} tokens "
|
||
f"(~{best_p // max(1, block)} blocks of {block}), "
|
||
f"served cached_tokens={cached}, prompt={len(prompt)}"
|
||
)
|
||
if best_p < reusable:
|
||
lo = max(0, best_p - 12)
|
||
hi = best_p + 12
|
||
try:
|
||
stored_ctx = self.tokenizer.decode(best_seq[lo:hi])
|
||
prompt_ctx = self.tokenizer.decode(prompt[lo:hi])
|
||
except Exception:
|
||
stored_ctx = prompt_ctx = "<decode failed>"
|
||
logger.debug(
|
||
f"Request {request.request_id}: first divergence at token "
|
||
f"{best_p}: stored=...{stored_ctx!r} vs prompt=...{prompt_ctx!r}"
|
||
)
|
||
|
||
_CACHE_FRESHNESS_WAIT_MIN_PROMPT_TOKENS = 8192
|
||
_CACHE_FRESHNESS_WAIT_MIN_COMMON_TOKENS = 8192
|
||
_CACHE_FRESHNESS_WAIT_MIN_PROMPT_RATIO = 0.30
|
||
_CACHE_FRESHNESS_WAIT_TIMEOUT_S = 4.0
|
||
|
||
@staticmethod
|
||
def _store_extra_keys_match(
|
||
info: _InflightStoreInfo,
|
||
request: Request,
|
||
) -> bool:
|
||
return (
|
||
info.extra_keys == request.vlm_extra_keys_for_cache
|
||
and info.extra_key_token_start
|
||
== request.vlm_extra_key_token_start_for_cache
|
||
and info.extra_key_ranges == request.vlm_extra_key_ranges_for_cache
|
||
)
|
||
|
||
def _find_relevant_inflight_store(
|
||
self,
|
||
request: Request,
|
||
) -> tuple[str, concurrent.futures.Future, int] | None:
|
||
"""Find a pending store_cache job worth waiting for before lookup."""
|
||
if not self._inflight_store_futures:
|
||
return None
|
||
|
||
prompt = request.prompt_token_ids or []
|
||
if len(prompt) < self._CACHE_FRESHNESS_WAIT_MIN_PROMPT_TOKENS:
|
||
return None
|
||
|
||
best_rid: str | None = None
|
||
best_future: concurrent.futures.Future | None = None
|
||
best_common = 0
|
||
for rid, future in list(self._inflight_store_futures.items()):
|
||
if future.done():
|
||
continue
|
||
info = self._inflight_store_info.get(rid)
|
||
if info is None or not self._store_extra_keys_match(info, request):
|
||
continue
|
||
|
||
common = self._common_prefix_len(prompt, info.tokens)
|
||
if common > best_common:
|
||
best_rid = rid
|
||
best_future = future
|
||
best_common = common
|
||
|
||
if best_future is None or best_rid is None:
|
||
return None
|
||
if not (
|
||
best_common >= self._CACHE_FRESHNESS_WAIT_MIN_COMMON_TOKENS
|
||
or best_common / len(prompt) >= self._CACHE_FRESHNESS_WAIT_MIN_PROMPT_RATIO
|
||
):
|
||
return None
|
||
|
||
return best_rid, best_future, best_common
|
||
|
||
def _should_defer_for_cache_freshness(self, request: Request) -> bool:
|
||
"""Defer only this waiting request until a relevant store is visible.
|
||
|
||
This intentionally does not call Future.result(). The scheduler step can
|
||
return immediately, so running decode rows and chunked prefills continue
|
||
on subsequent steps while the waiting head holds admission order.
|
||
"""
|
||
if self.block_aware_cache is None:
|
||
return False
|
||
|
||
now = time.monotonic()
|
||
wait = self._cache_freshness_waits.get(request.request_id)
|
||
if wait is not None:
|
||
if wait.future.done():
|
||
self._cache_freshness_waits.pop(request.request_id, None)
|
||
try:
|
||
exc = wait.future.exception()
|
||
except concurrent.futures.CancelledError:
|
||
logger.debug(
|
||
"Cache freshness deferral saw cancelled store_cache %s "
|
||
"before prefix lookup for %s",
|
||
wait.store_request_id,
|
||
request.request_id,
|
||
)
|
||
else:
|
||
if exc is None:
|
||
logger.debug(
|
||
"Completed cache freshness deferral for store_cache %s "
|
||
"before prefix lookup for %s",
|
||
wait.store_request_id,
|
||
request.request_id,
|
||
)
|
||
else:
|
||
logger.debug(
|
||
"Cache freshness deferral saw failed store_cache %s "
|
||
"before prefix lookup for %s: %s",
|
||
wait.store_request_id,
|
||
request.request_id,
|
||
exc,
|
||
)
|
||
return False
|
||
if now >= wait.deadline_s:
|
||
self._cache_freshness_waits.pop(request.request_id, None)
|
||
logger.debug(
|
||
"Timed out cache freshness deferral for store_cache %s before "
|
||
"prefix lookup for %s (common_prefix=%d/%d)",
|
||
wait.store_request_id,
|
||
request.request_id,
|
||
wait.common_prefix,
|
||
wait.prompt_len,
|
||
)
|
||
return False
|
||
return True
|
||
|
||
match = self._find_relevant_inflight_store(request)
|
||
if match is None:
|
||
return False
|
||
|
||
store_request_id, future, common_prefix = match
|
||
prompt_len = len(request.prompt_token_ids or [])
|
||
timeout = self._CACHE_FRESHNESS_WAIT_TIMEOUT_S
|
||
self._cache_freshness_waits[request.request_id] = _CacheFreshnessWait(
|
||
store_request_id=store_request_id,
|
||
future=future,
|
||
common_prefix=common_prefix,
|
||
prompt_len=prompt_len,
|
||
deadline_s=now + timeout,
|
||
)
|
||
|
||
logger.debug(
|
||
"Deferring admission up to %.1fs for in-flight store_cache %s before "
|
||
"prefix lookup for %s (common_prefix=%d/%d running=%d prefilling=%d)",
|
||
timeout,
|
||
store_request_id,
|
||
request.request_id,
|
||
common_prefix,
|
||
prompt_len,
|
||
len(self.running),
|
||
len(self.prefilling),
|
||
)
|
||
return True
|
||
|
||
def _prepare_prefix_cache_for_request(self, request: Request) -> None:
|
||
if request.request_id in self._prefix_cache_prepared:
|
||
return
|
||
|
||
# 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,
|
||
extra_keys=request.vlm_extra_keys_for_cache,
|
||
extra_key_token_start=request.vlm_extra_key_token_start_for_cache,
|
||
extra_key_ranges=request.vlm_extra_key_ranges_for_cache,
|
||
)
|
||
if block_table and block_table.num_tokens > 0:
|
||
bypass_hot_cache = self._bypass_hot_cache_under_pressure()
|
||
if bypass_hot_cache:
|
||
logger.info(
|
||
"Skipping hot-cache preload for %s under memory pressure",
|
||
request.request_id,
|
||
)
|
||
else:
|
||
self.block_aware_cache.preload_blocks(block_table)
|
||
# Reconstruct actual KVCache objects from stored tensor data
|
||
# Note: reconstruct_cache may modify block_table in-place if
|
||
# partial reconstruction occurs (some blocks invalid)
|
||
original_tokens = block_table.num_tokens
|
||
if bypass_hot_cache:
|
||
reconstructed = self.block_aware_cache.reconstruct_cache(
|
||
block_table,
|
||
promote_to_hot_cache=False,
|
||
)
|
||
else:
|
||
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)
|
||
# Recalculate remaining_tokens in case block_table was truncated
|
||
request.remaining_tokens = request.prompt_token_ids[
|
||
block_table.num_tokens :
|
||
]
|
||
if self._align_minimax_m3_partial_cache_to_prefill_step(request):
|
||
request.cached_tokens = block_table.num_tokens
|
||
request.shared_prefix_blocks = len(block_table.block_ids)
|
||
request.remaining_tokens = request.prompt_token_ids[
|
||
block_table.num_tokens :
|
||
]
|
||
# For exact prefix hits we need cache state at (N-1) and the
|
||
# last prompt token as input to produce the first decode logit.
|
||
# Reusing cache state at N and feeding the last token again
|
||
# shifts the model state and can change greedy output.
|
||
if len(request.remaining_tokens) == 0 and request.cached_tokens > 0:
|
||
if self._cache_list_needs_boundary_snapshot(
|
||
request.prompt_cache
|
||
):
|
||
# Stateful non-sliceable caches (Rotating/Arrays)
|
||
# cannot be safely converted from N to N-1 state
|
||
# without cache-type-specific logic.
|
||
if self.paged_cache_manager is not None:
|
||
self.paged_cache_manager.delete_block_table(
|
||
request.request_id
|
||
)
|
||
request.prompt_cache = None
|
||
request.block_table = None
|
||
request.cached_tokens = 0
|
||
request.shared_prefix_blocks = 0
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
logger.debug(
|
||
f"Request {request.request_id}: exact cache hit with "
|
||
f"stateful cache type, falling back to full prefill "
|
||
f"for deterministic kickoff"
|
||
)
|
||
elif self._trim_prompt_cache_for_generation(
|
||
request.prompt_cache
|
||
):
|
||
request.cached_tokens = max(0, request.cached_tokens - 1)
|
||
request.remaining_tokens = request.prompt_token_ids[-1:]
|
||
logger.debug(
|
||
f"Request {request.request_id}: exact cache hit adjusted "
|
||
f"to N-1 state for generation kickoff "
|
||
f"(cached_tokens={request.cached_tokens}, "
|
||
f"remaining={len(request.remaining_tokens)})"
|
||
)
|
||
else:
|
||
# Fallback to full recompute when cache layers cannot
|
||
# be safely trimmed by one token (e.g., non-trimmable
|
||
# recurrent state caches).
|
||
if self.paged_cache_manager is not None:
|
||
self.paged_cache_manager.delete_block_table(
|
||
request.request_id
|
||
)
|
||
request.prompt_cache = None
|
||
request.block_table = None
|
||
request.cached_tokens = 0
|
||
request.shared_prefix_blocks = 0
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
logger.debug(
|
||
f"Request {request.request_id}: exact cache hit could "
|
||
f"not be trimmed safely, falling back to full prefill"
|
||
)
|
||
if block_table.num_tokens < original_tokens:
|
||
logger.debug(
|
||
f"Request {request.request_id}: partial cache hit, "
|
||
f"{request.cached_tokens} tokens in {request.shared_prefix_blocks} blocks "
|
||
f"(originally {original_tokens} tokens), "
|
||
f"{len(request.remaining_tokens)} tokens remaining"
|
||
)
|
||
else:
|
||
logger.debug(
|
||
f"Request {request.request_id}: paged cache hit, "
|
||
f"{request.cached_tokens} tokens in {request.shared_prefix_blocks} blocks, "
|
||
f"{len(request.remaining_tokens)} tokens remaining, cache reconstructed"
|
||
)
|
||
else:
|
||
# Reconstruction failed, treat as cache miss
|
||
if self.paged_cache_manager is not None:
|
||
self.paged_cache_manager.delete_block_table(request.request_id)
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
logger.debug(
|
||
f"Request {request.request_id}: paged cache reconstruction failed, "
|
||
"released shared blocks"
|
||
)
|
||
else:
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
else:
|
||
# No paged SSD cache configured - process all tokens
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
|
||
# DEBUG-only: trace where this prompt diverges from recently stored
|
||
# cache sequences (issue #1003 always-miss diagnosis).
|
||
if logger.isEnabledFor(logging.DEBUG):
|
||
self._log_prefix_divergence(request)
|
||
|
||
# SpecPrefill: score remaining tokens with draft model if applicable.
|
||
# Must run AFTER prefix cache check (scoring applies only to uncached suffix).
|
||
self._try_specprefill_scoring(request)
|
||
self._prefix_cache_prepared.add(request.request_id)
|
||
|
||
def add_request(self, request: Request) -> None:
|
||
"""
|
||
Add a new request to the scheduler.
|
||
|
||
Raises SchedulerQueueFullError when the waiting queue is at or above
|
||
the configured cap (max(max_num_seqs * 4, 32)). Server layer maps
|
||
this to HTTP 503 + Retry-After.
|
||
|
||
Args:
|
||
request: The request to add
|
||
"""
|
||
if request.request_id in self.requests:
|
||
raise ValueError(f"Request {request.request_id} already exists")
|
||
|
||
# Cap the waiting queue so client-side polling can't accumulate
|
||
# unbounded work and the scheduler can apply backpressure via 503.
|
||
max_waiting = max(self.config.max_num_seqs * 4, 32)
|
||
if len(self.waiting) >= max_waiting:
|
||
from .exceptions import SchedulerQueueFullError
|
||
|
||
raise SchedulerQueueFullError(
|
||
current_depth=len(self.waiting),
|
||
max_depth=max_waiting,
|
||
)
|
||
|
||
# Tokenize if needed
|
||
if request.prompt_token_ids is None:
|
||
if isinstance(request.prompt, str):
|
||
request.prompt_token_ids = self.tokenizer.encode(request.prompt)
|
||
else:
|
||
request.prompt_token_ids = list(request.prompt)
|
||
request.num_prompt_tokens = len(request.prompt_token_ids)
|
||
|
||
# Prefix-cache lookup is intentionally delayed until admission. That
|
||
# lets a same-prefix request wait for a relevant in-flight store_cache
|
||
# without blocking the scheduler lane that continues decode/prefill.
|
||
#
|
||
# Keep the immediate preflight only when no prefix cache lookup can
|
||
# change cached_tokens. With a block-aware cache, the in-stream
|
||
# _preflight_memory_check runs after lookup with the final cache state.
|
||
if self.block_aware_cache is None:
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
try:
|
||
self.preflight_or_raise(
|
||
num_prompt_tokens=request.num_prompt_tokens,
|
||
cached_tokens=request.cached_tokens or 0,
|
||
request_id=request.request_id,
|
||
)
|
||
except Exception:
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
raise
|
||
|
||
# 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"
|
||
)
|
||
|
||
def set_specprefill_draft_model(
|
||
self, draft_model: Any, draft_model_name: str | None = None
|
||
) -> None:
|
||
"""Set the draft model for SpecPrefill scoring.
|
||
|
||
Creates separate block and SSD cache managers for the draft model so
|
||
target TurboQuant signatures cannot invalidate draft cache blocks.
|
||
"""
|
||
if not self._close_specprefill_draft_cache_manager():
|
||
raise RuntimeError(
|
||
"Could not close the previous SpecPrefill draft SSD cache manager"
|
||
)
|
||
self._specprefill_draft_model = draft_model
|
||
self._draft_prefix_cache: Any | None = None
|
||
if not draft_model_name:
|
||
logger.info(
|
||
"SpecPrefill: draft model set without a stable model name "
|
||
"(no SSD cache)"
|
||
)
|
||
return
|
||
|
||
if (
|
||
self.paged_cache_manager is not None
|
||
and self.paged_ssd_cache_manager is not None
|
||
):
|
||
try:
|
||
from .cache.paged_cache import PagedCacheManager
|
||
from .cache.prefix_cache import BlockAwarePrefixCache
|
||
|
||
name = draft_model_name
|
||
draft_cache_list = make_prompt_cache(draft_model)
|
||
draft_layer_cache_types = None
|
||
if HAS_CACHE_TYPE_HANDLERS and ModelCacheConfig is not None:
|
||
try:
|
||
draft_model_cache_config = ModelCacheConfig.from_cache_list(
|
||
draft_cache_list,
|
||
model_name=name,
|
||
)
|
||
draft_layer_cache_types = draft_model_cache_config.get_type_names()
|
||
except Exception as e:
|
||
logger.debug(
|
||
"Could not infer SpecPrefill draft cache layout: %s", e
|
||
)
|
||
draft_paged = PagedCacheManager(
|
||
block_size=self.config.paged_cache_block_size,
|
||
max_blocks=self.paged_cache_manager.max_blocks,
|
||
model_name=name,
|
||
)
|
||
draft_ssd = PagedSSDCacheManager(
|
||
cache_dir=Path(self.config.paged_ssd_cache_dir),
|
||
max_size_bytes=self.config.paged_ssd_cache_max_size,
|
||
hot_cache_max_bytes=self.config.hot_cache_max_size,
|
||
hot_cache_only=self.config.hot_cache_only,
|
||
hot_cache_budget=self.config.hot_cache_budget,
|
||
expected_model_name=name,
|
||
expected_num_layers=len(draft_cache_list),
|
||
expected_block_size=self.config.paged_cache_block_size,
|
||
expected_block_size_tokens=self.config.paged_cache_block_size,
|
||
expected_kv_bytes_per_token=getattr(
|
||
self.paged_ssd_cache_manager,
|
||
"_expected_kv_bytes_per_token",
|
||
200_000,
|
||
),
|
||
expected_layer_cache_types=draft_layer_cache_types,
|
||
)
|
||
self._draft_paged_ssd_cache_manager = draft_ssd
|
||
draft_paged.set_paged_ssd_cache_manager(draft_ssd)
|
||
self._draft_prefix_cache = BlockAwarePrefixCache(
|
||
model=draft_model,
|
||
paged_cache_manager=draft_paged,
|
||
paged_ssd_cache_manager=draft_ssd,
|
||
)
|
||
self._draft_prefix_cache.set_cold_restore_callback(
|
||
self._restore_block_from_cold
|
||
)
|
||
logger.info(
|
||
f"SpecPrefill: draft model set with SSD cache (model_name={name})"
|
||
)
|
||
except Exception as e:
|
||
self._draft_prefix_cache = None
|
||
self._close_specprefill_draft_cache_manager()
|
||
logger.warning(f"SpecPrefill: draft SSD cache setup failed: {e}")
|
||
logger.info("SpecPrefill: draft model set (no SSD cache)")
|
||
else:
|
||
logger.info("SpecPrefill: draft model set (no SSD cache)")
|
||
|
||
def _close_specprefill_draft_cache_manager(self) -> bool:
|
||
manager = self._draft_paged_ssd_cache_manager
|
||
if manager is None:
|
||
return True
|
||
try:
|
||
manager.close()
|
||
except Exception as e:
|
||
logger.warning("SpecPrefill draft SSD cache shutdown error: %s", e)
|
||
return False
|
||
writer_thread = getattr(manager, "_writer_thread", None)
|
||
if writer_thread is not None and writer_thread.is_alive():
|
||
logger.warning(
|
||
"SpecPrefill draft SSD cache writer remains active after shutdown"
|
||
)
|
||
return False
|
||
|
||
self._draft_paged_ssd_cache_manager = None
|
||
return True
|
||
|
||
def set_vlm_mtp_drafter(
|
||
self,
|
||
drafter: VLMMTPDrafter | None,
|
||
draft_block_size: int | None = None,
|
||
) -> None:
|
||
"""Attach an MTP drafter for VLM MTP speculative decode.
|
||
|
||
Called by ``VLMBatchedEngine.set_vlm_mtp_drafter`` once the drafter
|
||
artifact is loaded. Supports any drafter that mlx-vlm's
|
||
``load_drafter()`` resolves to ``kind="mtp"`` (gemma4_assistant,
|
||
qwen3_5_mtp, etc.). ``None`` clears the toggle.
|
||
"""
|
||
self._vlm_mtp_drafter = drafter
|
||
self._vlm_mtp_draft_block_size = draft_block_size
|
||
if drafter is not None:
|
||
logger.info(
|
||
"VLM MTP drafter attached to scheduler (block_size=%s)",
|
||
draft_block_size,
|
||
)
|
||
|
||
def _route_to_vlm_mtp(
|
||
self,
|
||
request: Request,
|
||
prefilled_cache: list[Any],
|
||
last_tokens: list[int],
|
||
sampler: Callable[[Any], Any],
|
||
state_machine: Any,
|
||
) -> int | None:
|
||
"""Bypass BatchGenerator and stand up a vlm_mtp generator instead.
|
||
|
||
Runs the final forward on ``last_tokens`` with ``return_hidden=True``
|
||
and ``return_shared_kv=True`` so the drafter has the targets it
|
||
needs, samples the first bonus token from the resulting logits, and
|
||
returns a synthesized uid that ``step()`` will drive.
|
||
|
||
Returns ``None`` if the eligibility check fails at the last second
|
||
(drafter missing, language model lacks rollback hook, etc.) so the
|
||
caller can fall back to the normal BatchGenerator path.
|
||
"""
|
||
drafter = self._vlm_mtp_drafter
|
||
if drafter is None:
|
||
return None
|
||
|
||
# Gemma4AssistantDraftModel keeps ``_shared_kv`` / ``_input_embed`` on
|
||
# the module instance, so multiple in-flight ``_mtp_rounds`` generators
|
||
# share one drafter and effectively serialize on it: each round has
|
||
# to ``set_shared_kv`` for its own request before ``draft_block`` runs.
|
||
# Output stays correct because target-side verify is the source of
|
||
# truth in speculative decoding (a stale-drafter round just rejects
|
||
# everything and falls back to a target-only step), but the
|
||
# per-request tok/s is roughly halved under concurrency. Empirically
|
||
# at 4 concurrent, vlm_mtp gives ~14 tok/s each vs BatchGenerator's
|
||
# ~27 tok/s each — BG's batched matmul beats serialized speculative
|
||
# rounds. So we route only the first eligible request through
|
||
# vlm_mtp and let subsequent concurrent requests fall back. A future
|
||
# commit can swap this gate for true batched MTP via
|
||
# ``_mtp_rounds_batch`` if and when omlx prefill exposes batched
|
||
# hidden/shared_kv outputs.
|
||
if self._vlm_mtp_active:
|
||
logger.info(
|
||
"vlm_mtp routing skipped for %s: drafter is busy with %d "
|
||
"request(s); falling back to BatchGenerator",
|
||
request.request_id,
|
||
len(self._vlm_mtp_active),
|
||
)
|
||
return None
|
||
|
||
lm = getattr(self.model, "_language_model", None)
|
||
if lm is None or not hasattr(lm, "rollback_speculative_cache"):
|
||
logger.warning(
|
||
"vlm_mtp toggle on but model lacks _language_model with "
|
||
"rollback_speculative_cache (model=%s); falling back to "
|
||
"standard decode for request %s",
|
||
type(self.model).__name__,
|
||
request.request_id,
|
||
)
|
||
return None
|
||
target_model = self.model
|
||
|
||
if not last_tokens:
|
||
logger.warning(
|
||
"vlm_mtp routing skipped: last_tokens empty for request %s",
|
||
request.request_id,
|
||
)
|
||
return None
|
||
|
||
mtp_sampler = _make_suppressing_sampler(sampler, self._model_suppress_tokens)
|
||
last_arr = mx.array(last_tokens)[None] # (1, len_last)
|
||
try:
|
||
with mx.stream(self._stream):
|
||
set_batch_rope = getattr(target_model, "set_batch_rope_deltas", None)
|
||
if callable(set_batch_rope):
|
||
set_batch_rope(mx.array([request.rope_deltas]))
|
||
out = target_model(
|
||
last_arr,
|
||
cache=prefilled_cache,
|
||
return_hidden=True,
|
||
return_shared_kv=True,
|
||
)
|
||
mx.eval([c.state for c in prefilled_cache])
|
||
except Exception as e:
|
||
logger.warning(
|
||
"vlm_mtp final-prefill forward failed (%s); falling back "
|
||
"to standard decode for request %s",
|
||
e,
|
||
request.request_id,
|
||
)
|
||
return None
|
||
|
||
# Handle current LanguageModelOutput and legacy tuple
|
||
# (logits, hidden, gdn_states) MTP runtime patch returns.
|
||
if isinstance(out, tuple):
|
||
logits = out[0][:, -1, :]
|
||
hidden_raw = out[1]
|
||
else:
|
||
logits = out.logits[:, -1, :]
|
||
hidden_raw = out.hidden_states
|
||
|
||
first_bonus_arr = mtp_sampler(logits) # mx.array shape [1]
|
||
mx.eval(first_bonus_arr)
|
||
|
||
if isinstance(hidden_raw, list):
|
||
hidden = hidden_raw[-1]
|
||
else:
|
||
hidden = hidden_raw
|
||
# Slice to last position so the drafter sees a [B, 1, H] tensor
|
||
# regardless of how many tokens this forward processed.
|
||
if hidden.shape[1] > 1:
|
||
hidden = hidden[:, -1:, :]
|
||
|
||
# Combine base stop tokens (EOS, Harmony, generation_config) with
|
||
# request-specific stop_token_ids — same shape as _build_state_machine.
|
||
eos_ids: set[int] = self._get_stop_tokens()
|
||
if request.sampling_params.stop_token_ids:
|
||
eos_ids.update(request.sampling_params.stop_token_ids)
|
||
|
||
try:
|
||
generator = run_vlm_mtp_decode(
|
||
target_language_model=target_model,
|
||
drafter=drafter,
|
||
prompt_cache=prefilled_cache,
|
||
hidden=hidden,
|
||
shared_kv_states=(
|
||
getattr(out, "shared_kv_states", {})
|
||
if not isinstance(out, tuple)
|
||
else {}
|
||
),
|
||
first_bonus=int(first_bonus_arr.item()),
|
||
max_tokens=request.sampling_params.max_tokens,
|
||
sampler=mtp_sampler,
|
||
draft_block_size=self._vlm_mtp_draft_block_size,
|
||
token_dtype=mx.int32,
|
||
eos_token_ids=eos_ids or None,
|
||
)
|
||
except Exception as e:
|
||
logger.warning(
|
||
"vlm_mtp generator setup failed (%s); falling back for %s",
|
||
e,
|
||
request.request_id,
|
||
)
|
||
return None
|
||
|
||
uid = self._vlm_mtp_next_uid
|
||
self._vlm_mtp_next_uid -= 1
|
||
self._vlm_mtp_active[uid] = _VLMMTPDecodeState(
|
||
generator=generator,
|
||
request=request,
|
||
prompt_cache=prefilled_cache,
|
||
sampler=mtp_sampler,
|
||
state_machine=state_machine,
|
||
max_tokens=request.sampling_params.max_tokens,
|
||
stop_token_ids=set(eos_ids),
|
||
)
|
||
logger.info(
|
||
"vlm_mtp decode started: request=%s uid=%d block_size=%s",
|
||
request.request_id,
|
||
uid,
|
||
self._vlm_mtp_draft_block_size,
|
||
)
|
||
return uid
|
||
|
||
def _log_vlm_mtp_stats(
|
||
self, state: "_VLMMTPDecodeState", finish_reason: str
|
||
) -> None:
|
||
"""Emit one INFO line per finished vlm_mtp request with the drafter
|
||
acceptance rate measured for that request.
|
||
|
||
Reads ``Gemma4AssistantDraftModel.accept_lens`` — a list of accepted
|
||
draft counts per round, populated inside mlx-vlm's ``_mtp_rounds``.
|
||
The drafter mutates this in place and ``reset()`` (called at the
|
||
start of every new round-loop entry) clears it, so we have to read
|
||
before the next eligible request lands. The serialized routing in
|
||
``_route_to_vlm_mtp`` guarantees one in-flight vlm_mtp generator
|
||
at a time, so the value we read here belongs to ``state.request``.
|
||
"""
|
||
drafter = self._vlm_mtp_drafter
|
||
if drafter is None:
|
||
return
|
||
accept_lens = getattr(drafter.model, "accept_lens", None)
|
||
if not accept_lens:
|
||
return
|
||
try:
|
||
lens = [int(x) for x in accept_lens]
|
||
except Exception:
|
||
return
|
||
rounds = len(lens)
|
||
if rounds == 0:
|
||
return
|
||
total_accepted = sum(lens)
|
||
block_size = self._vlm_mtp_draft_block_size or int(
|
||
getattr(drafter.model.config, "block_size", 4)
|
||
)
|
||
max_per_round = max(1, block_size - 1)
|
||
acceptance_rate = total_accepted / (rounds * max_per_round)
|
||
avg_tokens_per_round = (total_accepted + rounds) / rounds
|
||
logger.info(
|
||
"vlm_mtp stats: request=%s finish=%s rounds=%d "
|
||
"accepted=%d/%d (%.1f%%) tokens_per_round=%.2f "
|
||
"emitted=%d block_size=%d",
|
||
state.request.request_id,
|
||
finish_reason,
|
||
rounds,
|
||
total_accepted,
|
||
rounds * max_per_round,
|
||
acceptance_rate * 100,
|
||
avg_tokens_per_round,
|
||
state.emitted,
|
||
block_size,
|
||
)
|
||
|
||
def _step_vlm_mtp(self) -> list[_VLMMTPResponse]:
|
||
"""Advance every active vlm_mtp generator by one yield.
|
||
|
||
Returns the synthesized responses for ``_process_batch_responses``.
|
||
Mirrors mlx-lm BatchGenerator's per-step contract: one
|
||
``GenerationBatch.Response``-shaped object per active uid.
|
||
"""
|
||
if not self._vlm_mtp_active:
|
||
return []
|
||
|
||
responses: list[_VLMMTPResponse] = []
|
||
for uid, state in list(self._vlm_mtp_active.items()):
|
||
try:
|
||
with mx.stream(self._stream):
|
||
token_val = next(state.generator)
|
||
except StopIteration:
|
||
# Round loop exited naturally — terminate with prompt cache
|
||
# so the prefix-cache layer can keep using it.
|
||
self._log_vlm_mtp_stats(state, "length")
|
||
responses.append(
|
||
_VLMMTPResponse(
|
||
uid=uid,
|
||
token=0,
|
||
finish_reason="length",
|
||
prompt_cache=state.prompt_cache,
|
||
)
|
||
)
|
||
state.finished = True
|
||
continue
|
||
|
||
# Single-request mode yields ints; batch mode (not yet routed
|
||
# by omlx) would yield a list. Guard so the path stays robust
|
||
# if we widen routing later.
|
||
if isinstance(token_val, list):
|
||
# Take the first row (we only route singles for now).
|
||
tok = next((t for t in token_val if t is not None), None)
|
||
if tok is None:
|
||
responses.append(
|
||
_VLMMTPResponse(
|
||
uid=uid,
|
||
token=0,
|
||
finish_reason="length",
|
||
prompt_cache=state.prompt_cache,
|
||
)
|
||
)
|
||
state.finished = True
|
||
continue
|
||
token = int(tok)
|
||
else:
|
||
token = int(token_val)
|
||
|
||
state.emitted += 1
|
||
finish_reason: str | None = None
|
||
if state.stop_token_ids and token in state.stop_token_ids:
|
||
finish_reason = "stop"
|
||
elif state.emitted >= state.max_tokens:
|
||
finish_reason = "length"
|
||
|
||
if finish_reason is not None:
|
||
self._log_vlm_mtp_stats(state, finish_reason)
|
||
|
||
responses.append(
|
||
_VLMMTPResponse(
|
||
uid=uid,
|
||
token=token,
|
||
finish_reason=finish_reason,
|
||
prompt_cache=(
|
||
state.prompt_cache if finish_reason is not None else None
|
||
),
|
||
)
|
||
)
|
||
if finish_reason is not None:
|
||
state.finished = True
|
||
|
||
# Drop finished entries.
|
||
for uid in [u for u, s in self._vlm_mtp_active.items() if s.finished]:
|
||
del self._vlm_mtp_active[uid]
|
||
|
||
return responses
|
||
|
||
def _try_specprefill_scoring(self, request: Request) -> None:
|
||
"""Score tokens with draft model if SpecPrefill is applicable.
|
||
|
||
Uses paged SSD cache for the draft model: if the prompt prefix
|
||
was already scored in a previous turn, the draft cache is restored
|
||
and only the new suffix is prefilled through the draft model.
|
||
"""
|
||
if self._specprefill_draft_model is None:
|
||
return
|
||
|
||
specprefill_enabled = getattr(request, "_specprefill_enabled", False)
|
||
if not specprefill_enabled:
|
||
return
|
||
|
||
if request.vlm_inputs_embeds is not None:
|
||
return
|
||
|
||
remaining = request.remaining_tokens or request.prompt_token_ids
|
||
if remaining is None:
|
||
return
|
||
|
||
from .patches.specprefill import DEFAULT_KEEP_RATE, DEFAULT_THRESHOLD
|
||
from .specprefill.policy import plan_specprefill_scoring
|
||
|
||
# Apply deterministic admission before the draft workflow.
|
||
plan = plan_specprefill_scoring(
|
||
remaining_tokens=remaining,
|
||
system_prompt_end=request.specprefill_system_end,
|
||
cached_tokens=request.cached_tokens,
|
||
requested_threshold=getattr(request, "_specprefill_threshold", None),
|
||
requested_keep_pct=getattr(request, "_specprefill_keep_pct", None),
|
||
default_threshold=DEFAULT_THRESHOLD,
|
||
default_keep_pct=DEFAULT_KEEP_RATE,
|
||
)
|
||
if plan is None:
|
||
return
|
||
|
||
from .specprefill.draft import run_specprefill_draft_scoring
|
||
|
||
run_specprefill_draft_scoring(
|
||
request=request,
|
||
plan=plan,
|
||
draft_model=self._specprefill_draft_model,
|
||
draft_prefix_cache=self._draft_prefix_cache,
|
||
model_id=self.config.model_name,
|
||
prefill_step_size=self.config.prefill_step_size,
|
||
stream=self._stream,
|
||
extract_cache_states=self._extract_cache_states,
|
||
sync_and_clear_cache=lambda: _sync_and_clear_cache(self._stream),
|
||
log=logger,
|
||
)
|
||
|
||
def _cleanup_specprefill(self, request_id: str) -> None:
|
||
"""Clean up SpecPrefill RoPE patches when a request finishes."""
|
||
if self._specprefill_active_request_id == request_id:
|
||
from .patches.specprefill import cleanup_rope
|
||
|
||
cleanup_rope(self.model)
|
||
self._specprefill_active_request_id = None
|
||
logger.debug(
|
||
f"SpecPrefill: RoPE restored for finished request {request_id}"
|
||
)
|
||
|
||
def _trim_prompt_cache_for_generation(self, cache_list: list[Any]) -> bool:
|
||
"""Trim each cache layer by one token for exact-hit generation kickoff."""
|
||
return self._trim_prompt_cache_by_tokens(cache_list, 1)
|
||
|
||
def _trim_prompt_cache_by_tokens(self, cache_list: list[Any], n: int) -> bool:
|
||
"""Trim each cache layer by n tokens."""
|
||
if not cache_list:
|
||
return False
|
||
if n <= 0:
|
||
return True
|
||
|
||
for cache_obj in cache_list:
|
||
if not self._trim_cache_tree_by_tokens(cache_obj, n):
|
||
return False
|
||
return True
|
||
|
||
def _trim_cache_tree_by_one(self, cache_obj: Any) -> bool:
|
||
"""Trim one token from cache object (recursively for CacheList)."""
|
||
return self._trim_cache_tree_by_tokens(cache_obj, 1)
|
||
|
||
def _trim_cache_tree_by_tokens(self, cache_obj: Any, n: int) -> bool:
|
||
"""Trim n tokens from cache object (recursively for CacheList)."""
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
return all(
|
||
self._trim_cache_tree_by_tokens(sub_cache, n)
|
||
for sub_cache in sub_caches
|
||
)
|
||
|
||
trim_fn = getattr(cache_obj, "trim", None)
|
||
if not callable(trim_fn):
|
||
return False
|
||
|
||
try:
|
||
trimmed = trim_fn(n)
|
||
if trimmed is None:
|
||
return True
|
||
return int(trimmed) >= n
|
||
except Exception:
|
||
return False
|
||
|
||
def _cache_tree_has_class_name(
|
||
self,
|
||
cache_obj: Any,
|
||
class_names: frozenset[str],
|
||
) -> bool:
|
||
"""Return True when a cache tree contains one of the named cache classes."""
|
||
if type(cache_obj).__name__ in class_names:
|
||
return True
|
||
sub_caches = getattr(cache_obj, "caches", None)
|
||
if isinstance(sub_caches, (list, tuple)):
|
||
return any(
|
||
self._cache_tree_has_class_name(sub_cache, class_names)
|
||
for sub_cache in sub_caches
|
||
)
|
||
return False
|
||
|
||
def _align_minimax_m3_partial_cache_to_prefill_step(
|
||
self,
|
||
request: "Request",
|
||
) -> bool:
|
||
"""Align MiniMax M3 partial hits to external prefill chunk boundaries."""
|
||
cache_list = request.prompt_cache
|
||
block_table = request.block_table
|
||
prompt_tokens = request.prompt_token_ids or []
|
||
if not cache_list or block_table is None or not block_table.block_ids:
|
||
return False
|
||
if block_table.num_tokens <= 0 or block_table.num_tokens >= len(prompt_tokens):
|
||
return False
|
||
|
||
minimax_m3_names = frozenset({"MiniMaxM3KVCache"})
|
||
has_minimax_m3 = any(
|
||
self._cache_tree_has_class_name(cache_obj, minimax_m3_names)
|
||
for cache_obj in cache_list
|
||
)
|
||
if not has_minimax_m3:
|
||
return False
|
||
|
||
block_size = int(getattr(self.config, "paged_cache_block_size", 0) or 0)
|
||
prefill_step = int(getattr(self.config, "prefill_step_size", 0) or 0)
|
||
if block_size <= 0 or prefill_step <= block_size:
|
||
return False
|
||
|
||
aligned_tokens = (block_table.num_tokens // prefill_step) * prefill_step
|
||
aligned_tokens = (aligned_tokens // block_size) * block_size
|
||
if aligned_tokens <= 0 or aligned_tokens >= block_table.num_tokens:
|
||
return False
|
||
|
||
target_block_count = 0
|
||
target_tokens = 0
|
||
for block_id in block_table.block_ids:
|
||
block = (
|
||
self.paged_cache_manager.allocated_blocks.get(block_id)
|
||
if self.paged_cache_manager is not None
|
||
else None
|
||
)
|
||
token_count = int(getattr(block, "token_count", block_size) or block_size)
|
||
if target_tokens + token_count > aligned_tokens:
|
||
break
|
||
target_tokens += token_count
|
||
target_block_count += 1
|
||
|
||
if target_tokens != aligned_tokens:
|
||
logger.debug(
|
||
"MiniMax M3 partial cache alignment skipped for %s: cannot align "
|
||
"block table from %d to %d tokens",
|
||
request.request_id,
|
||
block_table.num_tokens,
|
||
aligned_tokens,
|
||
)
|
||
return False
|
||
|
||
trim_tokens = block_table.num_tokens - aligned_tokens
|
||
if not self._trim_prompt_cache_by_tokens(cache_list, trim_tokens):
|
||
logger.debug(
|
||
"MiniMax M3 partial cache alignment skipped for %s: cache trim "
|
||
"by %d tokens failed",
|
||
request.request_id,
|
||
trim_tokens,
|
||
)
|
||
return False
|
||
|
||
dropped_block_ids = block_table.block_ids[target_block_count:]
|
||
if self.paged_cache_manager is not None:
|
||
for block_id in dropped_block_ids:
|
||
self.paged_cache_manager.free_block(block_id)
|
||
block_table.block_ids = block_table.block_ids[:target_block_count]
|
||
block_table.num_tokens = aligned_tokens
|
||
|
||
logger.info(
|
||
"MiniMax M3 partial cache aligned to prefill step for %s: "
|
||
"%d -> %d tokens, dropped %d block(s)",
|
||
request.request_id,
|
||
aligned_tokens + trim_tokens,
|
||
aligned_tokens,
|
||
len(dropped_block_ids),
|
||
)
|
||
return True
|
||
|
||
def _remove_uid_from_active_batch(self, uid: int) -> None:
|
||
"""Remove UID from BatchGenerator safely.
|
||
|
||
vlm_mtp uses negative uids that BatchGenerator never sees; the
|
||
per-uid generator state is owned by ``_vlm_mtp_active`` and gets
|
||
dropped when ``_step_vlm_mtp`` marks the entry finished.
|
||
"""
|
||
if uid < 0:
|
||
return
|
||
if self.batch_generator is None:
|
||
return
|
||
|
||
self.batch_generator.remove([uid])
|
||
|
||
def _check_pending_aborts_for_uids(self, uids: list[int]) -> list[int]:
|
||
"""Return UIDs that have pending aborts.
|
||
|
||
Called during prefill to detect aborted
|
||
requests between chunks. GIL guarantees thread-safe reads of
|
||
_pending_abort_ids from the executor thread.
|
||
"""
|
||
if not self._pending_abort_ids:
|
||
return []
|
||
aborted = []
|
||
for uid in uids:
|
||
request_id = self.uid_to_request_id.get(uid)
|
||
if request_id and request_id in self._pending_abort_ids:
|
||
aborted.append(uid)
|
||
return aborted
|
||
|
||
def abort_request(self, request_id: str) -> bool:
|
||
"""
|
||
Enqueue a request for deferred abort.
|
||
|
||
The actual abort is processed at the start of the next step() call,
|
||
ensuring thread safety with the hybrid executor pattern. CPython GIL
|
||
guarantees set.add() is atomic.
|
||
|
||
Args:
|
||
request_id: The request ID to abort
|
||
|
||
Returns:
|
||
True (abort is always enqueued)
|
||
"""
|
||
self._pending_abort_ids.add(request_id)
|
||
logger.debug(f"Enqueued deferred abort for request {request_id}")
|
||
return True
|
||
|
||
def _process_pending_aborts(self) -> None:
|
||
"""Drain and process pending abort requests.
|
||
|
||
Called from step() to ensure aborts are processed in the same
|
||
execution context as generation (thread-safe).
|
||
"""
|
||
while self._pending_abort_ids:
|
||
request_id = self._pending_abort_ids.pop()
|
||
self._do_abort_request(request_id)
|
||
|
||
def _cleanup_prefill_abort_request(
|
||
self, request: "Request", temp_uid: int | None = None
|
||
) -> None:
|
||
"""Finish cleanup for a request aborted while it was being prefetched.
|
||
|
||
External prefill removes the request from ``waiting`` before it has a
|
||
real BatchGenerator UID. If a client abort arrives at that point, the
|
||
normal next-step deferred abort can be stranded because ``has_requests``
|
||
no longer sees queued work. Finish it synchronously on the scheduler
|
||
thread instead.
|
||
"""
|
||
if temp_uid is not None:
|
||
self.uid_to_request_id.pop(temp_uid, None)
|
||
self.request_id_to_uid.pop(request.request_id, None)
|
||
|
||
self._pending_abort_ids.discard(request.request_id)
|
||
self._do_abort_request(request.request_id)
|
||
|
||
def request_idle_reclaim(self) -> None:
|
||
"""Enqueue a between-turn Metal reclaim (thread-safe, no Metal touch).
|
||
|
||
Called by ProcessMemoryEnforcer (asyncio thread) when memory pressure
|
||
is hard but every loaded model is pinned and no load is in progress —
|
||
the case where there is nothing to evict. Setting the flag is
|
||
GIL-atomic; the actual ``_sync_and_clear_cache`` runs on the inference
|
||
thread when step() drains it, and only when the scheduler is idle.
|
||
"""
|
||
self._pending_reclaim_request = True
|
||
|
||
def _process_pending_reclaim(self) -> None:
|
||
"""Drain a deferred idle reclaim request (inference-thread side).
|
||
|
||
Only reclaims when truly idle (no running / prefilling / waiting work)
|
||
so we never clear Metal buffers an in-flight decode or prefill still
|
||
references.
|
||
"""
|
||
if not self._pending_reclaim_request:
|
||
return
|
||
self._pending_reclaim_request = False
|
||
if self.running or self.prefilling or self.waiting:
|
||
return
|
||
before = self._current_usage_bytes()
|
||
after = self._reclaim_prefill_headroom()
|
||
logger.info(
|
||
"Idle reclaim: trimmed Metal transients between turns "
|
||
"(%.1fGB -> %.1fGB)",
|
||
before / 1024**3,
|
||
after / 1024**3,
|
||
)
|
||
|
||
def _do_abort_request(self, request_id: str) -> bool:
|
||
"""
|
||
Actually abort a request. Must be called from the step() context.
|
||
|
||
Args:
|
||
request_id: The request ID to abort
|
||
|
||
Returns:
|
||
True if request was found and aborted, False otherwise
|
||
"""
|
||
request = self.requests.get(request_id)
|
||
if request is None:
|
||
return False
|
||
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
|
||
# Remove from waiting queue
|
||
if request.status == RequestStatus.WAITING:
|
||
try:
|
||
self.waiting.remove(request)
|
||
except ValueError:
|
||
pass
|
||
|
||
# Remove from chunked-prefill queue (if mid-prefill)
|
||
if request_id in self._prefill_states:
|
||
self._prefill_states.pop(request_id, None)
|
||
self.prefilling = deque(
|
||
r for r in self.prefilling if r.request_id != request_id
|
||
)
|
||
|
||
# Remove from running (BatchGenerator)
|
||
if request.request_id in self.request_id_to_uid:
|
||
uid = self.request_id_to_uid[request.request_id]
|
||
# Synchronize in-flight GPU work before modifying batch state.
|
||
# batch_generator.remove() triggers lazy KV cache array slicing
|
||
# that replaces references to arrays still used by in-flight
|
||
# Metal command buffers. Without this barrier the Metal driver
|
||
# can hit 'completeMemory() prepare count underflow'.
|
||
_safe_sync_stream(self._stream)
|
||
self._remove_uid_from_active_batch(uid)
|
||
if hasattr(self.model, "unregister_rope_delta"):
|
||
self.model.unregister_rope_delta(uid)
|
||
if uid < 0:
|
||
mtp_state = self._vlm_mtp_active.pop(uid, None)
|
||
if mtp_state is not None:
|
||
close = getattr(mtp_state.generator, "close", None)
|
||
if callable(close):
|
||
close()
|
||
_unregister_uid_row(self.model, uid)
|
||
del self.uid_to_request_id[uid]
|
||
del self.request_id_to_uid[request.request_id]
|
||
|
||
if request_id in self.running:
|
||
del self.running[request_id]
|
||
|
||
# Release blocks for eviction (same as _cleanup_finished)
|
||
if self.paged_cache_manager is not None:
|
||
block_table = self.paged_cache_manager.get_block_table(request_id)
|
||
if block_table is None and hasattr(request, "block_table"):
|
||
block_table = request.block_table
|
||
if block_table:
|
||
released = self.paged_cache_manager.release_for_eviction(
|
||
block_table.block_ids
|
||
)
|
||
if released > 0:
|
||
logger.debug(
|
||
f"Released {released} blocks for eviction on abort "
|
||
f"(request {request_id})"
|
||
)
|
||
|
||
# Clear request entry from block_aware_cache
|
||
if self.block_aware_cache is not None:
|
||
self.block_aware_cache.clear_request_entry(request_id)
|
||
|
||
# Clean up streaming detokenizer to prevent state contamination
|
||
self._cleanup_detokenizer(request_id)
|
||
|
||
# Clean up protocol-specific output parser session
|
||
self._cleanup_output_parser_session(request_id)
|
||
|
||
# Clean up VLM adapter state to prevent contamination
|
||
if hasattr(self.model, "clear_vlm_position_state"):
|
||
self.model.clear_vlm_position_state()
|
||
if hasattr(self.model, "clear_pending_embeddings"):
|
||
self.model.clear_pending_embeddings()
|
||
|
||
# Drop any boundary snapshot for this request.
|
||
self._boundary_cache_snapshots.pop(request_id, None)
|
||
if self._boundary_snapshot_store is not None:
|
||
self._boundary_snapshot_store.cleanup_request(request_id)
|
||
|
||
# Remove from prefill progress tracker.
|
||
get_prefill_tracker().remove(request_id)
|
||
|
||
# Mark as aborted
|
||
request.set_finished(RequestStatus.FINISHED_ABORTED)
|
||
self.finished_req_ids.add(request_id)
|
||
|
||
# Remove from requests dict and clear cache references to release
|
||
# MLX arrays promptly (mirrors _cleanup_finished behavior).
|
||
# _cleanup_request (engine_core) no longer calls remove_finished_request,
|
||
# so this is the single cleanup point for aborted requests.
|
||
req_to_remove = self.requests.pop(request_id, None)
|
||
if req_to_remove is not None:
|
||
req_to_remove._extracted_cache = None
|
||
req_to_remove.prompt_cache = None
|
||
|
||
logger.debug(f"Aborted request {request_id}")
|
||
return True
|
||
|
||
def has_requests(self) -> bool:
|
||
"""Check if there are any pending or running requests.
|
||
|
||
Also returns True when a deferred Metal cache clear is pending,
|
||
so that the engine loop keeps calling step() until the clear fires.
|
||
Without this, an idle server would never reach the target step and
|
||
stale buffers would accumulate indefinitely.
|
||
"""
|
||
return bool(
|
||
self.waiting
|
||
or self.prefilling
|
||
or self.running
|
||
or self._pending_async_removes
|
||
or self._deferred_clear_at is not None
|
||
)
|
||
|
||
def _refresh_generation_overflow_recovery_ids(self) -> None:
|
||
"""Drop serial-retry markers once the affected requests leave the scheduler."""
|
||
if not self._generation_overflow_recovery_ids:
|
||
return
|
||
active_ids = set(self.running)
|
||
active_ids.update(request.request_id for request in self.waiting)
|
||
active_ids.update(request.request_id for request in self.prefilling)
|
||
self._generation_overflow_recovery_ids.intersection_update(active_ids)
|
||
|
||
def _effective_max_num_seqs(self) -> int:
|
||
"""Current admission cap, narrowed for models that require serial decode."""
|
||
self._refresh_generation_overflow_recovery_ids()
|
||
if self._serialize_llama4_requests or self._generation_overflow_recovery_ids:
|
||
return 1
|
||
return max(1, self.config.max_num_seqs)
|
||
|
||
def fail_all_requests(self) -> list[str]:
|
||
"""Remove all running and waiting requests after unrecoverable error.
|
||
|
||
Used as a safety net by engine_core when step() raises an
|
||
unexpected exception, to prevent infinite loops.
|
||
|
||
Only resets batch_generator (not full cache) because this method
|
||
is called for non-corruption errors — corruption is already
|
||
handled inside step().
|
||
|
||
Returns:
|
||
List of failed request IDs.
|
||
"""
|
||
failed_ids: list[str] = []
|
||
for request_id in list(self.running):
|
||
failed_ids.append(request_id)
|
||
req = self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
if req is not None:
|
||
req._extracted_cache = None
|
||
req.prompt_cache = None
|
||
self.running.clear()
|
||
for request in list(self.prefilling):
|
||
failed_ids.append(request.request_id)
|
||
req = self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
if req is not None:
|
||
req._extracted_cache = None
|
||
req.prompt_cache = None
|
||
self.prefilling.clear()
|
||
self._prefill_states.clear()
|
||
for request in list(self.waiting):
|
||
failed_ids.append(request.request_id)
|
||
req = self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
if req is not None:
|
||
req._extracted_cache = None
|
||
req.prompt_cache = None
|
||
self.waiting.clear()
|
||
# Catch in-flight orphans: a request popped from self.waiting but
|
||
# not yet added to self.running (or self.prefilling) sits as a
|
||
# local in _schedule_waiting. If _do_external_prefill raises, the
|
||
# request is unreachable through the three queues but still lives
|
||
# in self.requests (and the engine_core collector / finished_event
|
||
# for its id is still waiting). Without this pass, fail_all_requests
|
||
# returns an incomplete list and the HTTP request hangs forever.
|
||
#
|
||
# Exclude finished requests still awaiting async cache-store cleanup
|
||
# (those have an entry in ``_inflight_store_futures`` — see
|
||
# ``_cleanup_finished`` line ~5267). They have already emitted a
|
||
# ``finished=True`` output to their collector; ``_drain_pending_async_removes``
|
||
# pops them from ``self.requests`` after the store future completes.
|
||
# Failing them here would append an error output that wins over the
|
||
# success for non-streaming ``generate()`` callers (engine_core
|
||
# returns the last queued output).
|
||
for request_id in list(self.requests):
|
||
if request_id in self._inflight_store_futures:
|
||
continue
|
||
failed_ids.append(request_id)
|
||
req = self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
if req is not None:
|
||
req._extracted_cache = None
|
||
req.prompt_cache = None
|
||
# Clear stale uid mappings for every failed id. Running requests hold
|
||
# real uids; the in-flight orphan above holds the temp_uid assigned at
|
||
# _schedule_waiting (id(request)) that its success-path cleanup never
|
||
# reached. batch_generator is reset below, so these mappings are dead
|
||
# either way. failed_ids excludes _inflight_store_futures ids, so the
|
||
# async-cleanup uids that _drain_pending_async_removes still needs are
|
||
# left intact.
|
||
for rid in failed_ids:
|
||
uid = self.request_id_to_uid.pop(rid, None)
|
||
if uid is not None:
|
||
self.uid_to_request_id.pop(uid, None)
|
||
self._generation_overflow_recovery_ids.difference_update(failed_ids)
|
||
# Reset batch generator only (cache is not corrupted). Every row dies
|
||
# with it; survivors re-register at re-insert.
|
||
_unregister_uid_rows_for_model(self.model)
|
||
self.batch_generator = None
|
||
self._current_sampler_params = None
|
||
# Reclaim fragmented Metal buffers after generation failure.
|
||
# Without this, subsequent requests may hit the same resource
|
||
# limit even though Python references have been cleared.
|
||
# Wrapped in try-except because Metal may already be in an error
|
||
# state — mx.synchronize() or mx.clear_cache() can throw a C++
|
||
# exception that causes SIGABRT if uncaught (#435).
|
||
try:
|
||
_sync_and_clear_cache(self._stream)
|
||
except Exception as e:
|
||
logger.warning(f"Metal cache clear failed during error recovery: {e}")
|
||
# Requests failed mid-prefill leave PrefillProgressTracker entries
|
||
# behind (auto-removal only fires at processed >= total). The local
|
||
# RuntimeError handlers in the prefill paths cover the common memory
|
||
# errors (#1405), but any other exception type bubbles up here and
|
||
# would leak a phantom "PP" row on the dashboard.
|
||
tracker = get_prefill_tracker()
|
||
for rid in failed_ids:
|
||
tracker.remove(rid)
|
||
# Republish the admin snapshot now that the queues are empty. The
|
||
# snapshot is normally published at the end of a successful step();
|
||
# when step() raises, the last published snapshot still lists the
|
||
# failed requests, so the dashboard and the macOS app keep showing
|
||
# them as "generating" until the next successful step (#2126).
|
||
self._publish_admin_snapshot()
|
||
return failed_ids
|
||
|
||
def get_num_waiting(self) -> int:
|
||
"""Get number of waiting requests."""
|
||
return len(self.waiting)
|
||
|
||
def get_num_running(self) -> int:
|
||
"""Get number of running requests."""
|
||
return len(self.running)
|
||
|
||
def _num_admitted_requests(self) -> int:
|
||
"""Return requests already occupying scheduler capacity."""
|
||
return len(self.running) + len(self.prefilling)
|
||
|
||
def _preflight_memory_check(
|
||
self, request: "Request"
|
||
) -> "_PreflightRejection | None":
|
||
"""
|
||
Estimate whether prefill would exceed memory limits.
|
||
|
||
Computes worst-case peak memory for the last prefill chunk
|
||
(model weights + KV cache + SDPA activation/scratch) and rejects
|
||
if it would exceed the hard limit.
|
||
|
||
Mirrors MLX SDPA dispatch closely enough that unsupported prefill
|
||
head dimensions are charged for the unfused fp32 score matrix.
|
||
|
||
Returns:
|
||
``_PreflightRejection`` carrying the message + numeric
|
||
estimated / limit bytes if the request should be rejected,
|
||
otherwise ``None``. The structured return lets the server
|
||
layer populate ``PrefillMemoryExceededError.estimated_bytes``
|
||
/ ``limit_bytes`` without parsing the human string.
|
||
"""
|
||
if not self._prefill_memory_guard:
|
||
return None
|
||
if self._memory_hard_limit_bytes <= 0:
|
||
return None
|
||
if self.memory_monitor is None:
|
||
return None
|
||
|
||
prompt_tokens = request.num_prompt_tokens
|
||
cached_tokens = request.cached_tokens or 0
|
||
new_tokens = max(prompt_tokens - cached_tokens, 0)
|
||
|
||
if new_tokens == 0:
|
||
return None
|
||
|
||
peak = self.memory_monitor.estimate_prefill_peak_bytes(
|
||
new_tokens, self.config.prefill_step_size, cached_tokens=cached_tokens
|
||
)
|
||
if peak == 0:
|
||
return None # can't estimate, skip
|
||
|
||
current = self._current_usage_bytes()
|
||
estimated = current + peak
|
||
hard_limit = self._memory_hard_limit_bytes
|
||
|
||
if estimated > hard_limit:
|
||
# Try LRU eviction first (upstream's predictive-throttle
|
||
# path): if eviction can free enough headroom this raises
|
||
# ``_PrefillEvictionNeeded`` and the request is paused for
|
||
# retry. If eviction can't help (already retried, no idle
|
||
# models), the call is a no-op and we fall through to the
|
||
# typed rejection.
|
||
self._raise_prefill_eviction_if_available(
|
||
request_id=request.request_id,
|
||
current=current,
|
||
target_cap=hard_limit,
|
||
predicted_transient=peak,
|
||
requested_tokens=min(new_tokens, self.config.prefill_step_size),
|
||
reason="prefill_preflight",
|
||
)
|
||
|
||
message = self._format_rejection_message(
|
||
estimated=estimated,
|
||
current=current,
|
||
peak=peak,
|
||
hard_limit=hard_limit,
|
||
)
|
||
return _PreflightRejection(
|
||
message=message,
|
||
estimated_bytes=int(estimated),
|
||
limit_bytes=int(hard_limit),
|
||
)
|
||
safety_rejection = self._preflight_safety_rejection(
|
||
num_prompt_tokens=prompt_tokens,
|
||
cached_tokens=cached_tokens,
|
||
current_usage_bytes=current,
|
||
)
|
||
if safety_rejection is not None:
|
||
requested_tokens = min(max(1, self._prefill_min_chunk_tokens), new_tokens)
|
||
self._raise_prefill_eviction_if_available(
|
||
request_id=request.request_id,
|
||
current=current,
|
||
target_cap=safety_rejection.limit_bytes,
|
||
predicted_transient=max(
|
||
0, int(safety_rejection.estimated_bytes) - int(current)
|
||
),
|
||
requested_tokens=requested_tokens,
|
||
reason="prefill_safety_cap",
|
||
)
|
||
return safety_rejection
|
||
return None
|
||
|
||
def _memory_component_limit_for_rejection(self, component_limit: int) -> int:
|
||
if component_limit <= 0:
|
||
return 0
|
||
hot_reserved = max(0, int(self._memory_hot_cache_reserved_bytes or 0))
|
||
if hot_reserved <= 0:
|
||
return component_limit
|
||
return max(1, component_limit - hot_reserved)
|
||
|
||
def _format_rejection_message(
|
||
self,
|
||
*,
|
||
estimated: int,
|
||
current: int,
|
||
peak: int,
|
||
hard_limit: int,
|
||
) -> str:
|
||
"""Build the prefill-rejection diagnostic.
|
||
|
||
Identifies which of static / dynamic / metal_cap is binding so the
|
||
message can steer the user to the right remedy (close apps for
|
||
dynamic, raise sysctl for metal_cap, raise tier or reduce context
|
||
for static). Component ceilings are propagated by
|
||
``ProcessMemoryEnforcer._propagate_memory_limit``; if a caller
|
||
wired this scheduler outside that path the components stay 0 and
|
||
we fall back to a generic message.
|
||
"""
|
||
from .utils.hardware import format_bytes
|
||
|
||
static = self._memory_static_ceiling_bytes
|
||
dynamic = self._memory_dynamic_ceiling_bytes
|
||
metal_cap = self._memory_metal_cap_bytes
|
||
|
||
binding: list[str] = []
|
||
if static and self._memory_component_limit_for_rejection(static) == hard_limit:
|
||
binding.append("static")
|
||
if (
|
||
dynamic
|
||
and self._memory_component_limit_for_rejection(dynamic) == hard_limit
|
||
):
|
||
binding.append("dynamic")
|
||
if (
|
||
metal_cap
|
||
and self._memory_component_limit_for_rejection(metal_cap) == hard_limit
|
||
):
|
||
binding.append("metal_cap")
|
||
binding_str = "/".join(binding) if binding else "effective"
|
||
|
||
# Order remedies by likelihood of helping for the binding cause.
|
||
# Dynamic-bound on a reclaim tier (safe/balanced/aggressive) means
|
||
# reclaimable memory is low right now even though the static cap
|
||
# has room — closing apps raises ``free`` / ``inactive`` and a
|
||
# more aggressive ``memory_guard_tier`` raises the active-reclaim
|
||
# ratio. Dynamic-bound under ``custom`` means the user pinned the
|
||
# ceiling there; the only knob that helps is raising
|
||
# ``custom_ceiling_bytes`` itself. Metal-cap bound means the
|
||
# kernel sysctl is the ceiling, so raising ``iogpu.wired_limit_mb``
|
||
# is the only knob that helps. Static-bound (or no breakdown
|
||
# available) leaves ``memory_guard_tier`` / context length as the
|
||
# levers.
|
||
is_custom = self._memory_guard_tier == "custom"
|
||
if "dynamic" in binding and is_custom:
|
||
advice = (
|
||
f"raise custom_ceiling_bytes in admin Memory settings "
|
||
f"(currently pinned at {format_bytes(dynamic)}), "
|
||
f"or reduce context length"
|
||
)
|
||
elif "dynamic" in binding and static and static > dynamic:
|
||
headroom = max(0, dynamic - current)
|
||
advice = (
|
||
f"close other apps to free RAM "
|
||
f"(static cap is {format_bytes(static)} but only "
|
||
f"{format_bytes(headroom)} is reclaimable right now), "
|
||
f"raise memory_guard_tier (safe → balanced → aggressive), "
|
||
f"or reduce context length"
|
||
)
|
||
elif "metal_cap" in binding:
|
||
advice = (
|
||
f"raise kernel iogpu.wired_limit_mb in Terminal "
|
||
f"(currently caps Metal at {format_bytes(metal_cap)}), "
|
||
f"or reduce context length"
|
||
)
|
||
else:
|
||
advice = (
|
||
"reduce context length or raise memory_guard_tier "
|
||
"(safe → balanced → aggressive)"
|
||
)
|
||
advice = advice[:1].upper() + advice[1:]
|
||
|
||
return (
|
||
f"Prefill would require ~{format_bytes(estimated)} peak "
|
||
f"(current {format_bytes(current)} + KV+SDPA {format_bytes(peak)}) "
|
||
f"but {binding_str} ceiling is {format_bytes(hard_limit)}. "
|
||
f"{advice}."
|
||
)
|
||
|
||
def preflight_or_raise(
|
||
self,
|
||
*,
|
||
num_prompt_tokens: int,
|
||
cached_tokens: int = 0,
|
||
request_id: str | None = None,
|
||
) -> None:
|
||
"""Pre-StreamingResponse prefill memory check.
|
||
|
||
Called from the engine's ``preflight_chat`` / ``preflight_completion``
|
||
before the FastAPI route wraps the body in a ``StreamingResponse``,
|
||
so the typed exception can be mapped to HTTP 400 by the registered
|
||
handler. A no-op when the guard is disabled or the request fits.
|
||
|
||
Mirrors the ``_preflight_memory_check`` math but takes token counts
|
||
directly (no Request object) and raises instead of returning a
|
||
message — the in-stream re-check inside ``_schedule_waiting``
|
||
remains as defense-in-depth.
|
||
"""
|
||
if not self._prefill_memory_guard:
|
||
return
|
||
if self._memory_hard_limit_bytes <= 0:
|
||
return
|
||
if self.memory_monitor is None:
|
||
return
|
||
|
||
new_tokens = max(int(num_prompt_tokens) - max(int(cached_tokens), 0), 0)
|
||
if new_tokens == 0:
|
||
return
|
||
|
||
peak = self.memory_monitor.estimate_prefill_peak_bytes(
|
||
new_tokens, self.config.prefill_step_size, cached_tokens=cached_tokens
|
||
)
|
||
if peak == 0:
|
||
return
|
||
|
||
current = self._current_usage_bytes(refresh_mlx_active=False)
|
||
if not request_id:
|
||
import uuid as _uuid
|
||
|
||
request_id = f"preflight-{_uuid.uuid4().hex[:8]}"
|
||
|
||
if current + peak > self._memory_hard_limit_bytes:
|
||
message = self._format_rejection_message(
|
||
estimated=current + peak,
|
||
current=current,
|
||
peak=peak,
|
||
hard_limit=self._memory_hard_limit_bytes,
|
||
)
|
||
|
||
logger.warning(
|
||
"Preflight rejected (%d tokens, cached=%d, request_id=%s): %s",
|
||
num_prompt_tokens,
|
||
cached_tokens,
|
||
request_id,
|
||
message,
|
||
)
|
||
raise PrefillMemoryExceededError(
|
||
message=message,
|
||
request_id=request_id,
|
||
estimated_bytes=int(current + peak),
|
||
limit_bytes=int(self._memory_hard_limit_bytes),
|
||
)
|
||
|
||
safety_rejection = self._preflight_safety_rejection(
|
||
num_prompt_tokens=num_prompt_tokens,
|
||
cached_tokens=cached_tokens,
|
||
current_usage_bytes=current,
|
||
)
|
||
if safety_rejection is None:
|
||
return
|
||
|
||
logger.warning(
|
||
"Preflight safety-cap rejected (%d tokens, cached=%d, "
|
||
"request_id=%s): %s",
|
||
num_prompt_tokens,
|
||
cached_tokens,
|
||
request_id,
|
||
safety_rejection.message,
|
||
)
|
||
raise PrefillMemoryExceededError(
|
||
message=safety_rejection.message,
|
||
request_id=request_id,
|
||
estimated_bytes=safety_rejection.estimated_bytes,
|
||
limit_bytes=safety_rejection.limit_bytes,
|
||
)
|
||
|
||
def preflight_eviction_request(
|
||
self,
|
||
*,
|
||
num_prompt_tokens: int,
|
||
cached_tokens: int = 0,
|
||
request_id: str | None = None,
|
||
) -> PrefillEvictionRequest | None:
|
||
"""Return an idle-model eviction request for route-level preflight.
|
||
|
||
``preflight_or_raise`` runs before a ``Request`` is admitted, so it
|
||
cannot use the request-bound ``_raise_prefill_eviction_if_available``.
|
||
The API-facing engines call this first, run the async pool callback if
|
||
needed, then call ``preflight_or_raise`` to re-measure and reject only
|
||
if eviction did not create enough headroom.
|
||
"""
|
||
if not self._prefill_memory_guard:
|
||
return None
|
||
if self._memory_hard_limit_bytes <= 0:
|
||
return None
|
||
if self.memory_monitor is None:
|
||
return None
|
||
|
||
new_tokens = max(int(num_prompt_tokens) - max(int(cached_tokens), 0), 0)
|
||
if new_tokens == 0:
|
||
return None
|
||
|
||
current = self._current_usage_bytes(refresh_mlx_active=False)
|
||
request_id = request_id or "preflight"
|
||
|
||
peak = self.memory_monitor.estimate_prefill_peak_bytes(
|
||
new_tokens, self.config.prefill_step_size, cached_tokens=cached_tokens
|
||
)
|
||
if peak and current + peak > self._memory_hard_limit_bytes:
|
||
return PrefillEvictionRequest(
|
||
request_id=request_id,
|
||
model_id=getattr(self.config, "model_name", ""),
|
||
current_bytes=int(current),
|
||
target_cap_bytes=int(self._memory_hard_limit_bytes),
|
||
predicted_transient_bytes=int(peak),
|
||
requested_tokens=int(min(new_tokens, self.config.prefill_step_size)),
|
||
reason="prefill_preflight",
|
||
)
|
||
|
||
safety_rejection = self._preflight_safety_rejection(
|
||
num_prompt_tokens=num_prompt_tokens,
|
||
cached_tokens=cached_tokens,
|
||
current_usage_bytes=current,
|
||
)
|
||
if safety_rejection is None:
|
||
return None
|
||
|
||
requested_tokens = min(max(1, self._prefill_min_chunk_tokens), new_tokens)
|
||
return PrefillEvictionRequest(
|
||
request_id=request_id,
|
||
model_id=getattr(self.config, "model_name", ""),
|
||
current_bytes=int(current),
|
||
target_cap_bytes=int(safety_rejection.limit_bytes),
|
||
predicted_transient_bytes=max(
|
||
0, int(safety_rejection.estimated_bytes) - int(current)
|
||
),
|
||
requested_tokens=int(requested_tokens),
|
||
reason="prefill_safety_cap",
|
||
)
|
||
|
||
def _preflight_safety_rejection(
|
||
self,
|
||
*,
|
||
num_prompt_tokens: int,
|
||
cached_tokens: int = 0,
|
||
current_usage_bytes: int,
|
||
) -> _PreflightRejection | None:
|
||
"""Predict whether even the safety floor chunk cannot fit.
|
||
|
||
This mirrors the mid-prefill ``_guard_prefill_chunk`` rejection, but
|
||
runs before the route returns a ``StreamingResponse``. It charges the
|
||
resident KV that will be allocated by the prompt plus the minimum
|
||
chunk transient at the full prompt context length.
|
||
"""
|
||
if self.memory_monitor is None:
|
||
return None
|
||
base_cap, cap, margin = self._prefill_abort_description()
|
||
if cap <= 0:
|
||
return None
|
||
|
||
new_tokens = max(int(num_prompt_tokens) - max(int(cached_tokens), 0), 0)
|
||
if new_tokens == 0:
|
||
return None
|
||
|
||
floor_chunk = min(max(1, self._prefill_min_chunk_tokens), new_tokens)
|
||
kv_len = max(int(num_prompt_tokens) - 1, 1)
|
||
kv_growth = self.memory_monitor.estimate_prompt_kv_bytes(new_tokens)
|
||
min_transient = self._predicted_chunk_transient(floor_chunk, kv_len)
|
||
if kv_growth <= 0 and min_transient <= 0:
|
||
return None
|
||
|
||
estimated = int(current_usage_bytes + kv_growth + min_transient)
|
||
if estimated <= cap:
|
||
return None
|
||
|
||
from .utils.hardware import format_bytes
|
||
|
||
message = (
|
||
"Prefill context too large for available memory "
|
||
f"(preflight safety guard, kv_len={kv_len}, "
|
||
f"min_chunk={floor_chunk}): predicted peak would require "
|
||
f"~{format_bytes(estimated)} "
|
||
f"(current {format_bytes(current_usage_bytes)} + "
|
||
f"KV {format_bytes(kv_growth)} + "
|
||
f"min-chunk transient {format_bytes(min_transient)}) "
|
||
f"but prefill safety cap is {format_bytes(cap)} "
|
||
f"({round(margin * 100)}% of effective ceiling "
|
||
f"{format_bytes(base_cap)}). Reduce context length, free system "
|
||
"memory, or loosen memory_guard_tier (safe → balanced → aggressive)."
|
||
)
|
||
return _PreflightRejection(
|
||
message=message,
|
||
estimated_bytes=estimated,
|
||
limit_bytes=int(cap),
|
||
)
|
||
|
||
def _schedule_waiting(
|
||
self,
|
||
) -> tuple[list["Request"], list[RequestOutput]]:
|
||
"""
|
||
Move requests from waiting queue to running.
|
||
|
||
Each request is prefilled externally before being inserted into
|
||
BatchGenerator, so prefill_batch_size=1 is always used. Cache
|
||
status homogeneity tracking is kept for safety since it affects
|
||
how we handle the existing_cache argument.
|
||
|
||
Returns:
|
||
Tuple of (scheduled requests, rejected error outputs)
|
||
"""
|
||
scheduled = []
|
||
rejected_outputs: list[RequestOutput] = []
|
||
|
||
# Track cache status of first scheduled request to ensure homogeneity
|
||
# None = not determined yet, True = has cache, False = no cache
|
||
batch_cache_status: bool | None = None
|
||
# Track VLM status: VLM and text-only requests cannot be in the same prefill batch
|
||
# None = not determined yet, True = VLM request, False = text-only request
|
||
batch_vlm_status: bool | None = None
|
||
# Track SpecPrefill: these requests must be alone (RoPE patching affects whole model)
|
||
batch_specprefill_status: bool | None = None
|
||
|
||
while (
|
||
self.waiting
|
||
and self._num_admitted_requests() < self._effective_max_num_seqs()
|
||
):
|
||
# Admission pause: set by ProcessMemoryEnforcer when phys
|
||
# crosses soft_threshold. New prefills wait; in-flight requests
|
||
# continue. First request always passes (no admitted work yet)
|
||
# so admission can recover by completing the current generation.
|
||
admitted = self._num_admitted_requests()
|
||
if self._admission_paused and admitted:
|
||
logger.debug(
|
||
"Admission paused by memory pressure, %d admitted",
|
||
admitted,
|
||
)
|
||
stalled = self._memory_admission_stall_output("admission_paused")
|
||
if stalled is not None:
|
||
rejected_outputs.append(stalled)
|
||
break
|
||
|
||
# Store-cache backpressure: when the post-completion pipeline is
|
||
# at its cleanup cap, defer admitting new prefills instead of
|
||
# blocking the generation step on the store-cache write (#1496).
|
||
# The cap bounds concurrent extracted-KV copies (the #1383 OOM
|
||
# guard) and shrinks under memory pressure via
|
||
# adjust_store_cache_cap. This also applies between sequential
|
||
# turns: a new prefill must not start while async store-cache
|
||
# cleanup still owns too many large cache payloads (#1684).
|
||
gate = self._store_cache_gate
|
||
pending_store_cleanups = len(self._pending_async_removes)
|
||
if gate is not None and (
|
||
not gate.has_capacity or pending_store_cleanups >= gate.cap
|
||
):
|
||
logger.debug(
|
||
"Admission deferred: store-cache pipeline full "
|
||
"(in_flight=%d pending_cleanups=%d cap=%d), %d running",
|
||
gate.in_flight,
|
||
pending_store_cleanups,
|
||
gate.cap,
|
||
len(self.running),
|
||
)
|
||
memory_related_gate = self._admission_paused
|
||
if (
|
||
not memory_related_gate
|
||
and self._prefill_memory_guard
|
||
and self._memory_limit_bytes > 0
|
||
):
|
||
try:
|
||
memory_related_gate = (
|
||
self._current_usage_bytes() >= self._memory_limit_bytes
|
||
)
|
||
except Exception:
|
||
memory_related_gate = False
|
||
if memory_related_gate:
|
||
if self.waiting:
|
||
self._clear_store_cache_admission_blocker(
|
||
self.waiting[0].request_id
|
||
)
|
||
stalled = self._memory_admission_stall_output(
|
||
"store_cache_backpressure"
|
||
)
|
||
if stalled is not None:
|
||
rejected_outputs.append(stalled)
|
||
else:
|
||
if self.waiting:
|
||
self._clear_memory_admission_blocker(self.waiting[0].request_id)
|
||
stalled = self._store_cache_admission_stall_output(
|
||
"store_cache_backpressure",
|
||
gate_in_flight=gate.in_flight,
|
||
gate_cap=gate.cap,
|
||
pending_cleanups=pending_store_cleanups,
|
||
)
|
||
if stalled is not None:
|
||
rejected_outputs.append(stalled)
|
||
break
|
||
|
||
# Generation memory guard: when requests are already admitted,
|
||
# defer scheduling if memory pressure is high to prevent
|
||
# Metal allocation failures during batch_generator.next().
|
||
# First request always passes (no admitted work yet).
|
||
if self._prefill_memory_guard and self._memory_limit_bytes > 0 and admitted:
|
||
current = self._current_usage_bytes()
|
||
if current > self._memory_limit_bytes:
|
||
logger.debug(
|
||
"Generation memory guard: deferring scheduling "
|
||
"(%s > %s), %d admitted",
|
||
current,
|
||
self._memory_limit_bytes,
|
||
admitted,
|
||
)
|
||
stalled = self._memory_admission_stall_output(
|
||
"generation_memory_guard"
|
||
)
|
||
if stalled is not None:
|
||
rejected_outputs.append(stalled)
|
||
break
|
||
|
||
request = self.waiting[0]
|
||
self._clear_memory_admission_blocker(request.request_id)
|
||
self._clear_store_cache_admission_blocker(request.request_id)
|
||
if self._should_defer_for_cache_freshness(request):
|
||
break
|
||
|
||
request = self.waiting.popleft()
|
||
self._cache_freshness_waits.pop(request.request_id, None)
|
||
self._clear_memory_admission_blocker(request.request_id)
|
||
self._clear_store_cache_admission_blocker(request.request_id)
|
||
|
||
# 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
|
||
|
||
self._prepare_prefix_cache_for_request(request)
|
||
|
||
# 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 - pass only last token for generation kickoff
|
||
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
|
||
|
||
# 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
|
||
|
||
# SpecPrefill requests must be alone in the batch (RoPE patching
|
||
# affects the entire model). Also block scheduling if another
|
||
# specprefill request is already running (offset RoPE active).
|
||
request_is_specprefill = request.specprefill_indices is not None
|
||
if (
|
||
self._specprefill_active_request_id is not None
|
||
and not request_is_specprefill
|
||
):
|
||
# A specprefill request is running — defer all others until it finishes
|
||
self.waiting.appendleft(request)
|
||
break
|
||
if batch_specprefill_status is None:
|
||
batch_specprefill_status = request_is_specprefill
|
||
elif batch_specprefill_status != request_is_specprefill:
|
||
self.waiting.appendleft(request)
|
||
break
|
||
if request_is_specprefill and len(scheduled) > 0:
|
||
# SpecPrefill request must be alone
|
||
self.waiting.appendleft(request)
|
||
break
|
||
|
||
# Check VLM status homogeneity: VLM and text-only requests use
|
||
# different prefill paths (embeddings vs token IDs)
|
||
request_is_vlm = request.vlm_inputs_embeds is not None
|
||
if batch_vlm_status is None:
|
||
batch_vlm_status = request_is_vlm
|
||
elif batch_vlm_status != request_is_vlm:
|
||
# VLM status mismatch - defer this request to next batch
|
||
self.waiting.appendleft(request)
|
||
logger.debug(
|
||
f"Deferring request {request.request_id} to next batch "
|
||
f"(VLM status mismatch: batch={batch_vlm_status}, request={request_is_vlm})"
|
||
)
|
||
break
|
||
|
||
# Check cache status homogeneity (kept for consistent prefill behavior)
|
||
request_has_cache = cache_to_use is not None
|
||
if batch_cache_status is None:
|
||
batch_cache_status = request_has_cache
|
||
elif batch_cache_status != request_has_cache:
|
||
# Cache status mismatch - defer this request to next batch
|
||
self.waiting.appendleft(request)
|
||
logger.debug(
|
||
f"Deferring request {request.request_id} to next batch "
|
||
f"(cache status mismatch: batch={batch_cache_status}, request={request_has_cache})"
|
||
)
|
||
break
|
||
|
||
# Mark as Harmony model if applicable (before think detection)
|
||
if self._is_harmony_model:
|
||
request.is_harmony_model = True
|
||
|
||
# Check if prompt ends with <think> token for reasoning models.
|
||
# Must happen before _build_sampler_and_processors so the thinking
|
||
# budget processor can check needs_think_prefix.
|
||
if self._detect_needs_think_prefix(request):
|
||
request.needs_think_prefix = True
|
||
|
||
# Per-request sampler/logits processors to avoid BatchGenerator recreation.
|
||
sampler, logits_processors = self._build_sampler_and_processors(
|
||
request.sampling_params, request
|
||
)
|
||
|
||
# Pre-flight memory guard: estimate peak memory for this request
|
||
# and reject if it would exceed the hard limit. The check
|
||
# may raise ``_PrefillEvictionNeeded`` (upstream's
|
||
# predictive-throttle path) to pause and retry under
|
||
# eviction headroom; only if eviction can't help does the
|
||
# typed rejection propagate.
|
||
try:
|
||
preflight_rejection = self._preflight_memory_check(request)
|
||
except _PrefillEvictionNeeded as e:
|
||
self._pause_for_prefill_eviction(request, e.request)
|
||
break
|
||
if preflight_rejection is not None:
|
||
logger.warning(
|
||
f"Request {request.request_id} rejected by prefill "
|
||
f"memory guard: {preflight_rejection.message}"
|
||
)
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
rejected_outputs.append(
|
||
_prefill_memory_error_output(
|
||
request.request_id,
|
||
preflight_rejection.message,
|
||
estimated_bytes=preflight_rejection.estimated_bytes,
|
||
limit_bytes=preflight_rejection.limit_bytes,
|
||
)
|
||
)
|
||
continue
|
||
|
||
# SpecPrefill: replace tokens with selected subset and pre-fill
|
||
# cache via sparse_prefill before inserting into BatchGenerator.
|
||
#
|
||
# Key design: sparse_prefill processes selected tokens (excluding
|
||
# the last prompt token). BatchGenerator then processes the last
|
||
# prompt token to produce generation logits. This avoids:
|
||
# - Double-processing the last token (Bug #2)
|
||
# - Off-by-one RoPE positions (Bug #1)
|
||
#
|
||
# Position math:
|
||
# sparse_prefill: N' tokens, adjustment = M - N'
|
||
# We subtract 1: adjustment = M - N' - 1
|
||
# BatchGenerator last token: pos = N' + (M - N' - 1) = M - 1
|
||
# First gen token: pos = (N'+1) + (M - N' - 1) = M
|
||
if request.specprefill_indices is not None:
|
||
tracker = get_prefill_tracker()
|
||
model_id = self.config.model_name
|
||
total_pp = 0
|
||
try:
|
||
sys_count = getattr(request, "_specprefill_system_tokens", 0)
|
||
all_tokens = tokens_to_process
|
||
|
||
from .specprefill.planning import plan_specprefill_target
|
||
|
||
target_plan = plan_specprefill_target(
|
||
all_tokens=all_tokens,
|
||
system_token_count=sys_count,
|
||
selected_indices=request.specprefill_indices.tolist(),
|
||
position_offset=request.specprefill_position_offset,
|
||
)
|
||
m_pre = target_plan.conversation_token_count
|
||
n_eff = target_plan.sparse_selected_token_count
|
||
total_pp = target_plan.total_tracker_prefill_count
|
||
tracker.update(request.request_id, 0, total_pp, model_id)
|
||
|
||
spec_sparse_extra = {
|
||
"prompt_tokens": request.num_prompt_tokens,
|
||
"system_tokens": request.specprefill_system_end,
|
||
"conversation_tokens": request.num_prompt_tokens
|
||
- request.specprefill_system_end,
|
||
"cached_tokens": request.cached_tokens,
|
||
"scored_tokens": m_pre,
|
||
"selected_tokens": n_eff,
|
||
"keep_percent": (
|
||
round(n_eff / m_pre * 100) if m_pre > 0 else 0
|
||
),
|
||
}
|
||
|
||
def _check_specprefill_abort(processed: int) -> None:
|
||
if request.request_id in self._pending_abort_ids:
|
||
logger.info(
|
||
f"SpecPrefill interrupted at {processed}/{total_pp} "
|
||
f"tokens: request aborted"
|
||
)
|
||
tracker.remove(request.request_id)
|
||
self.waiting.appendleft(request)
|
||
raise _PrefillAbortedError([], processed)
|
||
|
||
def _report_system_progress(processed: int, total: int) -> None:
|
||
tracker.update(
|
||
request.request_id,
|
||
min(processed, total_pp - 1),
|
||
total_pp,
|
||
model_id,
|
||
phase="specprefill_system",
|
||
detail="system prompt prefill",
|
||
extra=spec_sparse_extra,
|
||
)
|
||
|
||
def _report_sparse_progress(processed: int, total: int) -> None:
|
||
_check_specprefill_abort(sys_count + processed)
|
||
tracker.update(
|
||
request.request_id,
|
||
min(sys_count + processed, total_pp - 1),
|
||
total_pp,
|
||
model_id,
|
||
phase="specprefill_sparse",
|
||
detail="sparse target prefill",
|
||
extra={
|
||
"scored_tokens": m_pre,
|
||
"selected_tokens": n_eff,
|
||
"keep_percent": (
|
||
round(n_eff / m_pre * 100) if m_pre > 0 else 0
|
||
),
|
||
"prompt_tokens": request.num_prompt_tokens,
|
||
"system_tokens": request.specprefill_system_end,
|
||
"conversation_tokens": request.num_prompt_tokens
|
||
- request.specprefill_system_end,
|
||
"cached_tokens": request.cached_tokens,
|
||
},
|
||
)
|
||
|
||
from .specprefill.target import run_specprefill_target_prefill
|
||
|
||
target_result = run_specprefill_target_prefill(
|
||
target_model=self.model,
|
||
request=request,
|
||
plan=target_plan,
|
||
all_tokens=all_tokens,
|
||
selected_indices=request.specprefill_indices,
|
||
prefill_step_size=self.config.prefill_step_size,
|
||
stream=self._stream,
|
||
check_abort=_check_specprefill_abort,
|
||
report_system_progress=_report_system_progress,
|
||
report_sparse_progress=_report_sparse_progress,
|
||
sync_and_clear_cache=lambda: _sync_and_clear_cache(self._stream),
|
||
log=logger,
|
||
)
|
||
|
||
cache_to_use = target_result.prompt_cache
|
||
tokens_to_process = target_result.tokens_to_process
|
||
self._specprefill_active_request_id = request.request_id
|
||
|
||
# Mark spec-prefill complete (auto-removes tracker entry).
|
||
tracker.update(request.request_id, total_pp, total_pp, model_id)
|
||
|
||
except _PrefillAbortedError:
|
||
from .patches.specprefill import cleanup_rope
|
||
|
||
cleanup_rope(self.model)
|
||
request.specprefill_indices = None
|
||
tracker.remove(request.request_id)
|
||
_sync_and_clear_cache(self._stream)
|
||
self._cleanup_prefill_abort_request(request)
|
||
continue
|
||
except Exception as e:
|
||
from .patches.specprefill import cleanup_rope
|
||
|
||
logger.error(f"SpecPrefill sparse prefill failed: {e}")
|
||
cleanup_rope(self.model)
|
||
request.specprefill_indices = None
|
||
tracker.remove(request.request_id)
|
||
# Fall through to normal prefill
|
||
# External prefill: process tokens[0:N-1] outside BatchGenerator.
|
||
# Only the last token goes to insert() for the first decode step.
|
||
# SpecPrefill already handled its own prefill above, so skip for those.
|
||
if request.specprefill_indices is None and len(tokens_to_process) > 1:
|
||
vlm_embeds = None
|
||
if request.vlm_inputs_embeds is not None:
|
||
vlm_embeds = (
|
||
request.vlm_inputs_embeds,
|
||
request.vlm_extra_kwargs or {},
|
||
request.cached_tokens,
|
||
)
|
||
|
||
# Chunked prefill: non-VLM prompts longer than one step are
|
||
# spread across multiple step() calls. The first chunk is run
|
||
# here; subsequent chunks run in _advance_chunked_prefills().
|
||
if (
|
||
self.config.chunked_prefill
|
||
and vlm_embeds is None
|
||
and len(tokens_to_process) > self.config.prefill_step_size + 1
|
||
):
|
||
sm = self._build_state_machine(request)
|
||
per_row_lps = list(logits_processors) if logits_processors else []
|
||
state = self._begin_prefill(
|
||
request, tokens_to_process, cache_to_use
|
||
)
|
||
state.sampler = sampler
|
||
state.sm = sm
|
||
state.per_row_lps = per_row_lps
|
||
|
||
try:
|
||
done = self._step_prefill_chunk(state)
|
||
except _PrefillAbortedError:
|
||
_sync_and_clear_cache(self._stream)
|
||
self._cleanup_prefill_abort_request(request)
|
||
continue
|
||
except _PrefillEvictionNeeded as e:
|
||
# Raised by the adaptive throttle before the first
|
||
# chunk's forward pass, so a reconstructed prefix
|
||
# (prompt_cache / block_table / cached_tokens) is
|
||
# still valid. Keep it attached across the pause:
|
||
# if no idle model gets evicted, the retry prefills
|
||
# only the uncached suffix under the adaptive
|
||
# throttle instead of recomputing the whole prompt
|
||
# cold (#2180). Mirrors the in-flight pause in
|
||
# _advance_chunked_prefills, which also keeps state.
|
||
self._pause_for_prefill_eviction(request, e.request)
|
||
break
|
||
except PrefillMemoryExceededError as e:
|
||
logger.error(
|
||
"Chunked prefill (first chunk) capacity rejected "
|
||
"for %s: %s",
|
||
request.request_id,
|
||
e,
|
||
)
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
get_prefill_tracker().remove(request.request_id)
|
||
_sync_and_clear_cache()
|
||
rejected_outputs.append(
|
||
_prefill_memory_exception_output(request.request_id, e)
|
||
)
|
||
continue
|
||
except RuntimeError as e:
|
||
# Hard memory limit hit on the first chunk.
|
||
# _step_prefill_chunk updates the PrefillProgressTracker
|
||
# before the limit check, so without this catch the
|
||
# tracker entry leaks and stays in the dashboard
|
||
# forever (#1405). Mirrors the cleanup in
|
||
# _advance_chunked_prefills (d736bfd).
|
||
logger.error(
|
||
"Chunked prefill (first chunk) failed for %s: %s",
|
||
request.request_id,
|
||
e,
|
||
)
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
get_prefill_tracker().remove(request.request_id)
|
||
# Drop Metal cache pool buffers held by the aborted
|
||
# first chunk's forward / mx.eval transients.
|
||
_sync_and_clear_cache()
|
||
if self._requeue_or_fail_prefill(request, e):
|
||
continue
|
||
rejected_outputs.append(
|
||
RequestOutput(
|
||
request_id=request.request_id,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=str(e),
|
||
)
|
||
)
|
||
continue
|
||
|
||
if done:
|
||
self._emit_final_boundary_if_needed(state)
|
||
_sync_and_clear_cache(self._stream)
|
||
get_prefill_tracker().remove(request.request_id)
|
||
self._insert_prefilled_request(request, state, scheduled)
|
||
else:
|
||
self.prefilling.append(request)
|
||
self._prefill_states[request.request_id] = state
|
||
continue # Skip normal prefill + insert path
|
||
|
||
# Normal (non-chunked) full prefill path.
|
||
# Assign a temporary UID so progress callbacks can map
|
||
# uid→request_id during external prefill. Replaced by the
|
||
# real UID returned from insert().
|
||
temp_uid = id(request) # unique, won't collide with BatchGenerator UIDs
|
||
self.request_id_to_uid[request.request_id] = temp_uid
|
||
self.uid_to_request_id[temp_uid] = request.request_id
|
||
|
||
try:
|
||
prefilled_cache, last_token = self._do_external_prefill(
|
||
request,
|
||
tokens_to_process,
|
||
cache_to_use,
|
||
vlm_embeds=vlm_embeds,
|
||
)
|
||
except _PrefillAbortedError:
|
||
self._cleanup_prefill_abort_request(request, temp_uid=temp_uid)
|
||
continue
|
||
except _PrefillEvictionNeeded as e:
|
||
self.uid_to_request_id.pop(temp_uid, None)
|
||
self.request_id_to_uid.pop(request.request_id, None)
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
get_prefill_tracker().remove(request.request_id)
|
||
self._pause_for_prefill_eviction(request, e.request)
|
||
break
|
||
except PrefillMemoryExceededError as e:
|
||
logger.error(
|
||
"Prefill capacity rejected for %s: %s",
|
||
request.request_id,
|
||
e,
|
||
)
|
||
self.uid_to_request_id.pop(temp_uid, None)
|
||
self.request_id_to_uid.pop(request.request_id, None)
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
get_prefill_tracker().remove(request.request_id)
|
||
_sync_and_clear_cache()
|
||
rejected_outputs.append(
|
||
_prefill_memory_exception_output(request.request_id, e)
|
||
)
|
||
continue
|
||
except RuntimeError as e:
|
||
# Hard memory limit hit during external prefill. Without
|
||
# this catch, the exception bubbles up to step() and then
|
||
# engine_core's fail_all_requests(), which pops
|
||
# self.requests but cannot reach the PrefillProgressTracker
|
||
# singleton, so the dashboard entry leaks across model
|
||
# reload (#1405). Mirrors the cleanup in
|
||
# _advance_chunked_prefills (d736bfd).
|
||
logger.error("Prefill failed for %s: %s", request.request_id, e)
|
||
self.uid_to_request_id.pop(temp_uid, None)
|
||
self.request_id_to_uid.pop(request.request_id, None)
|
||
self._release_paged_cache_for_request(request.request_id)
|
||
self.requests.pop(request.request_id, None)
|
||
self._clear_request_admission_bookkeeping(request.request_id)
|
||
get_prefill_tracker().remove(request.request_id)
|
||
# Drop Metal cache pool buffers held by the aborted
|
||
# chunk's forward / mx.eval transients.
|
||
_sync_and_clear_cache()
|
||
if self._requeue_or_fail_prefill(request, e):
|
||
continue
|
||
rejected_outputs.append(
|
||
RequestOutput(
|
||
request_id=request.request_id,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=str(e),
|
||
)
|
||
)
|
||
continue
|
||
|
||
# Clean up temp UID mapping
|
||
del self.uid_to_request_id[temp_uid]
|
||
del self.request_id_to_uid[request.request_id]
|
||
|
||
# Prefill complete: remove from progress tracker so dashboard
|
||
# shows "generating" instead of "PP" during decode.
|
||
get_prefill_tracker().remove(request.request_id)
|
||
|
||
cache_to_use = prefilled_cache
|
||
tokens_to_process = last_token
|
||
|
||
# Capture per-request mRoPE rope_deltas for decode.
|
||
# Prefer _captured_rope_deltas from per-request extra_kwargs
|
||
# (set during get_input_embeddings), since the global
|
||
# _rope_deltas may be stale when explicit position_ids are used.
|
||
if request.vlm_inputs_embeds is not None:
|
||
extra = request.vlm_extra_kwargs or {}
|
||
captured = extra.get("_captured_rope_deltas")
|
||
if captured is not None:
|
||
if hasattr(captured, "item"):
|
||
request.rope_deltas = float(captured.item())
|
||
else:
|
||
request.rope_deltas = float(captured)
|
||
elif hasattr(self.model, "get_last_rope_deltas"):
|
||
request.rope_deltas = self.model.get_last_rope_deltas()
|
||
|
||
# Build per-request state machine for stop tokens
|
||
sm = self._build_state_machine(request)
|
||
|
||
# Set random seed for reproducible generation (best-effort).
|
||
# This affects global MLX random state, so concurrent requests
|
||
# may interfere. Matches OpenAI's best-effort seed semantics.
|
||
if request.sampling_params.seed is not None:
|
||
mx.random.seed(request.sampling_params.seed)
|
||
|
||
# TurboQuant KV is quantized at the end of _do_external_prefill
|
||
# (fp16 prefill → quantize once); _merge_caches() turns the per
|
||
# request TQ cache into a BatchTurboQuantKVCache on insert.
|
||
|
||
# VLM MTP routing: if a gemma4_assistant drafter is attached, run
|
||
# an extra last-token forward to capture hidden + shared_kv_states,
|
||
# sample the first bonus, and hand the request to a vlm_mtp
|
||
# generator instead of BatchGenerator. Falls through on any
|
||
# eligibility issue so other speculative paths stay intact.
|
||
if self._vlm_mtp_drafter is not None and cache_to_use is not None:
|
||
vlm_mtp_uid = self._route_to_vlm_mtp(
|
||
request, cache_to_use, tokens_to_process, sampler, sm
|
||
)
|
||
if vlm_mtp_uid is not None:
|
||
self.request_id_to_uid[request.request_id] = vlm_mtp_uid
|
||
self.uid_to_request_id[vlm_mtp_uid] = request.request_id
|
||
now = time.monotonic()
|
||
request.batch_uid = vlm_mtp_uid
|
||
request.status = RequestStatus.RUNNING
|
||
request.generation_started_at = now
|
||
request.last_activity_at = now
|
||
self.running[request.request_id] = request
|
||
scheduled.append(request)
|
||
self.total_prompt_tokens += request.num_prompt_tokens
|
||
logger.debug(
|
||
f"Scheduled request {request.request_id} via vlm_mtp "
|
||
f"(uid={vlm_mtp_uid}, {request.num_prompt_tokens} prompt tokens)"
|
||
)
|
||
continue
|
||
|
||
# Insert into BatchGenerator with pre-filled cache + last token.
|
||
# BatchGenerator only handles decode from here.
|
||
#
|
||
# IMPORTANT: ``logits_processors`` MUST be passed as a per-row
|
||
# list (possibly empty), never None. mlx-lm's
|
||
# GenerationBatch._step does ``for p in self.logits_processors[e]``
|
||
# in any branch where ``any(self.logits_processors)`` is True
|
||
# (e.g., heterogeneous merge with another row that has a
|
||
# processor). A None slot crashes that loop with
|
||
# ``TypeError: 'NoneType' object is not iterable``, which then
|
||
# bubbles into the engine retry loop and presents as a hang.
|
||
# See vllm-mlx-patched commit 8d4052b for the same root cause
|
||
# in a sibling project, and #934 for the user-visible symptom.
|
||
per_row_lps = list(logits_processors) if logits_processors else []
|
||
uids = self.batch_generator.insert(
|
||
[tokens_to_process],
|
||
max_tokens=[request.sampling_params.max_tokens],
|
||
caches=[cache_to_use] if cache_to_use else None,
|
||
all_tokens=[_batch_generator_all_tokens(request)],
|
||
samplers=[sampler],
|
||
logits_processors=[per_row_lps],
|
||
state_machines=[sm],
|
||
)
|
||
if uids:
|
||
_register_uid_rows(self.model, uids, [sampler], [per_row_lps])
|
||
uid = uids[0]
|
||
self.request_id_to_uid[request.request_id] = uid
|
||
self.uid_to_request_id[uid] = request.request_id
|
||
now = time.monotonic()
|
||
request.batch_uid = uid
|
||
request.status = RequestStatus.RUNNING
|
||
request.generation_started_at = now
|
||
request.last_activity_at = now
|
||
self.running[request.request_id] = request
|
||
scheduled.append(request)
|
||
|
||
# Register per-UID rope_delta for mRoPE decode.
|
||
if hasattr(self.model, "register_rope_delta"):
|
||
self.model.register_rope_delta(uid, request.rope_deltas)
|
||
|
||
self.total_prompt_tokens += request.num_prompt_tokens
|
||
cache_info = (
|
||
f", {request.cached_tokens} cached"
|
||
if request.cached_tokens > 0
|
||
else ""
|
||
)
|
||
cache_used = "with cache" if cache_to_use else "no cache"
|
||
logger.debug(
|
||
f"Scheduled request {request.request_id} (uid={uid}) "
|
||
f"with {len(tokens_to_process)} tokens to process "
|
||
f"({request.num_prompt_tokens} total){cache_info}, {cache_used}"
|
||
)
|
||
|
||
return scheduled, rejected_outputs
|
||
|
||
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()
|
||
|
||
step_now = time.monotonic()
|
||
generated_at = time.perf_counter()
|
||
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
|
||
|
||
request.last_activity_at = step_now
|
||
completion_tokens_before = request.num_output_tokens
|
||
|
||
# Release VLM embeddings after first decode token (prefill is done)
|
||
if request.vlm_inputs_embeds is not None:
|
||
request.vlm_inputs_embeds = None
|
||
request.vlm_extra_kwargs = None
|
||
|
||
# Check finish reason first - don't include EOS token in output
|
||
# (following mlx-lm's batch_generate behavior)
|
||
is_stop = response.finish_reason == "stop"
|
||
is_length = response.finish_reason == "length"
|
||
is_finished = response.finish_reason is not None
|
||
|
||
# Only append token if not stopping due to EOS token
|
||
new_text = ""
|
||
|
||
# Check if this request uses a protocol-specific output parser
|
||
parser_session = self._get_output_parser_session(request_id)
|
||
|
||
if parser_session is not None and not is_stop:
|
||
parser_result = parser_session.process_token(response.token)
|
||
new_text = parser_result.stream_text
|
||
if parser_result.visible_text:
|
||
request.output_text += parser_result.visible_text
|
||
|
||
# Parser-defined stop token can override finish reason
|
||
if parser_result.is_stop and not is_finished:
|
||
is_finished = True
|
||
is_stop = True
|
||
response.finish_reason = "stop"
|
||
|
||
should_record_token = (
|
||
parser_result.record_token
|
||
if parser_result.record_token is not None
|
||
else not is_stop
|
||
)
|
||
if should_record_token:
|
||
request.append_output_token(response.token)
|
||
|
||
elif not is_stop:
|
||
# Standard processing without a protocol parser
|
||
request.append_output_token(response.token)
|
||
|
||
# Decode the new token using streaming detokenizer for proper UTF-8 handling
|
||
detokenizer = self._get_detokenizer(request_id)
|
||
if detokenizer is not None:
|
||
detokenizer.add_token(response.token)
|
||
new_text = detokenizer.last_segment
|
||
else:
|
||
# Fallback to single-token decode
|
||
new_text = self.tokenizer.decode([response.token])
|
||
|
||
# Text-level stop-string fallback. Catches BPE edge cases
|
||
# where the tokenized stop sequence does not match the
|
||
# model's actual output tokens (e.g. " delta" vs "delta").
|
||
# Only scans the tail to keep cost O(stop_len) per step.
|
||
stop_strs = request.sampling_params.stop or []
|
||
if stop_strs and not is_finished and detokenizer is not None:
|
||
full_text = detokenizer.text
|
||
prev_len = len(full_text) - len(new_text)
|
||
for ss in stop_strs:
|
||
if not ss:
|
||
continue
|
||
scan_start = max(0, prev_len - len(ss) + 1)
|
||
idx_in_tail = full_text.find(ss, scan_start)
|
||
if idx_in_tail < 0:
|
||
continue
|
||
is_finished = True
|
||
is_stop = True
|
||
response.finish_reason = "stop"
|
||
if idx_in_tail >= prev_len:
|
||
new_text = new_text[: idx_in_tail - prev_len]
|
||
else:
|
||
new_text = ""
|
||
break
|
||
|
||
# Prepend <think> tag for first chunk if this is a reasoning model.
|
||
# Protocol parsers may expose a normalized prefix when their prompt
|
||
# uses a model-specific open-think marker (e.g. MiniMax <mm:think>).
|
||
if getattr(request, "needs_think_prefix", False):
|
||
if not getattr(request, "think_prefix_sent", False):
|
||
if parser_session is None:
|
||
think_tag = getattr(self.tokenizer, "think_start", "<think>")
|
||
prefix_text = think_tag + "\n"
|
||
else:
|
||
prefix_text = (
|
||
self._get_output_parser_thinking_start_output_text() or ""
|
||
)
|
||
if prefix_text:
|
||
new_text = prefix_text + new_text
|
||
if parser_session is not None:
|
||
request.output_text = prefix_text + request.output_text
|
||
request.think_prefix_sent = True
|
||
|
||
# Immediately discard logprobs if not requested to free memory (~800KB per response)
|
||
# This prevents accumulation of large MLX arrays during streaming
|
||
if (
|
||
hasattr(response, "logprobs")
|
||
and response.logprobs is not None
|
||
and not request.sampling_params.logprobs
|
||
):
|
||
response.logprobs = None
|
||
|
||
# Create output
|
||
output_generated_at = (
|
||
generated_at
|
||
if request.num_output_tokens > completion_tokens_before
|
||
else None
|
||
)
|
||
output = RequestOutput(
|
||
request_id=request_id,
|
||
new_token_ids=[response.token] if not is_stop else [],
|
||
new_text=new_text,
|
||
output_token_ids=list(request.output_token_ids),
|
||
prompt_tokens=request.num_prompt_tokens,
|
||
completion_tokens=request.num_output_tokens,
|
||
generated_at=output_generated_at,
|
||
generated_until=output_generated_at,
|
||
cached_tokens=request.cached_tokens,
|
||
)
|
||
|
||
if not is_finished:
|
||
self._maybe_capture_boundary_snapshot(request, response.uid)
|
||
|
||
# Handle finished requests
|
||
if is_finished:
|
||
if is_stop:
|
||
request.set_finished(RequestStatus.FINISHED_STOPPED)
|
||
elif is_length:
|
||
request.set_finished(RequestStatus.FINISHED_LENGTH_CAPPED)
|
||
|
||
output.finished = True
|
||
output.finish_reason = response.finish_reason
|
||
finished_ids.add(request_id)
|
||
|
||
if parser_session is not None:
|
||
final_result = parser_session.finalize()
|
||
if final_result.stream_text:
|
||
output.new_text += final_result.stream_text
|
||
if final_result.visible_text:
|
||
request.output_text += final_result.visible_text
|
||
if final_result.output_text_prefix:
|
||
request.output_text = (
|
||
final_result.output_text_prefix + request.output_text
|
||
)
|
||
if final_result.tool_calls:
|
||
output.tool_calls = final_result.tool_calls
|
||
if final_result.finish_reason:
|
||
output.finish_reason = final_result.finish_reason
|
||
output.output_text = request.output_text
|
||
else:
|
||
# Standard finalization without a protocol parser
|
||
# Finalize detokenizer to flush any remaining bytes
|
||
detokenizer = self._get_detokenizer(request_id)
|
||
if detokenizer is not None:
|
||
detokenizer.finalize()
|
||
final_segment = detokenizer.last_segment
|
||
if final_segment:
|
||
output.new_text += final_segment
|
||
|
||
# Decode full output
|
||
output.output_text = self.tokenizer.decode(request.output_token_ids)
|
||
request.output_text = output.output_text
|
||
|
||
# Trim accumulated output text at the first stop string
|
||
# match so non-streaming responses do not include the
|
||
# stop sequence itself (matches OpenAI semantics).
|
||
if is_stop:
|
||
stop_strs = request.sampling_params.stop or []
|
||
for ss in stop_strs:
|
||
if not ss:
|
||
continue
|
||
cut = output.output_text.find(ss)
|
||
if cut >= 0:
|
||
output.output_text = output.output_text[:cut]
|
||
request.output_text = output.output_text
|
||
break
|
||
|
||
# Extract cache for future reuse.
|
||
# In the new API, prompt_cache is a direct value (not callable).
|
||
raw_cache = getattr(response, "prompt_cache", None)
|
||
if raw_cache is not None:
|
||
try:
|
||
# SpecPrefill: sparse KV data can't be stored in
|
||
# paged cache (hash mismatch with full token IDs).
|
||
if request.specprefill_indices is not None:
|
||
raw_cache = None
|
||
|
||
# For paged cache, extract actual tensor states
|
||
# This allows cache to survive BatchGenerator recreation
|
||
elif self.block_aware_cache is not None:
|
||
extracted_cache, model_cache_config = (
|
||
self._extract_cache_states(raw_cache)
|
||
)
|
||
if extracted_cache:
|
||
request._extracted_cache = extracted_cache
|
||
request._model_cache_config = model_cache_config
|
||
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
|
||
request._model_cache_config = None
|
||
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"
|
||
)
|
||
logger.log(
|
||
5, "Request %s generated text:\n%s", request_id, output.output_text
|
||
)
|
||
|
||
outputs.append(output)
|
||
|
||
return outputs, finished_ids
|
||
|
||
def _release_paged_cache_for_request(self, request_id: str) -> None:
|
||
"""Drop a request's paged-cache footprint on rejection paths.
|
||
|
||
``add_request`` routes through ``block_aware_cache.fetch_cache``
|
||
which records the request in ``_request_tables`` and increments
|
||
ref counts on every prefix-matched paged-cache block. The normal
|
||
completion path releases that state in ``_cleanup_finished``;
|
||
the prefill-rejection paths in ``_advance_chunked_prefills`` /
|
||
``_schedule_waiting`` must do the same or rejected requests
|
||
leak block refs (pinning the paged cache and compounding the
|
||
very memory pressure that triggered the rejection) and orphan
|
||
``_request_tables`` entries.
|
||
"""
|
||
if self.block_aware_cache is not None:
|
||
self.block_aware_cache.release_cache(request_id)
|
||
elif self.paged_cache_manager is not None:
|
||
self.paged_cache_manager.delete_block_table(request_id)
|
||
|
||
# SpecPrefill primes an independent ``_draft_prefix_cache`` in
|
||
# ``_try_specprefill_scoring`` whose block refs are tracked
|
||
# separately from the target ``block_aware_cache``. Without
|
||
# releasing it on the rejection path a rejected SpecPrefill
|
||
# request leaks every draft-block ref symmetric to the
|
||
# target-cache leak the main branch above guards against.
|
||
draft_cache = getattr(self, "_draft_prefix_cache", None)
|
||
if draft_cache is not None:
|
||
try:
|
||
draft_cache.release_cache(request_id)
|
||
except Exception: # noqa: BLE001
|
||
logger.debug(
|
||
"Draft prefix cache release_cache(%s) raised; ignoring",
|
||
request_id,
|
||
exc_info=True,
|
||
)
|
||
|
||
def _cleanup_finished(self, finished_ids: set[str]) -> None:
|
||
"""Clean up finished requests and store caches for reuse."""
|
||
# Synchronize pending engine stream operations before cache storage.
|
||
# store_cache -> mx.save_safetensors triggers implicit mx.eval() which
|
||
# can conflict with async Metal operations on the generation stream.
|
||
if finished_ids:
|
||
with self._phase_timer("cleanup_finished_sync"):
|
||
_safe_sync_stream(self._stream)
|
||
|
||
# SpecPrefill: restore original RoPE if active request finished
|
||
for rid in finished_ids:
|
||
self._cleanup_specprefill(rid)
|
||
|
||
# Remove finished requests from prefill progress tracker.
|
||
tracker = get_prefill_tracker()
|
||
for rid in finished_ids:
|
||
tracker.remove(rid)
|
||
|
||
for request_id in finished_ids:
|
||
request = self.running.get(request_id)
|
||
|
||
# Store cache for future reuse (G2-async): submit to background
|
||
# executor so the post-finish 28GB+ memcpy doesn't block response
|
||
# streaming. The inference thread does mx.synchronize +
|
||
# boundary merge + a single batched mx.eval here; the worker
|
||
# handles _extract_tensor_bytes (CPU memcpy) + index/queue
|
||
# registration. batch_generator.remove(uid) is deferred and
|
||
# picked up at the next step's _drain_pending_async_removes.
|
||
store_future = None
|
||
if request is not None and request.prompt_token_ids:
|
||
if self.block_aware_cache is not None:
|
||
if (
|
||
hasattr(request, "_extracted_cache")
|
||
and request._extracted_cache is not None
|
||
):
|
||
prompt_boundary_store = None
|
||
else:
|
||
uid_for_store = self.request_id_to_uid.get(request_id, -1)
|
||
prompt_boundary_store = (
|
||
self._prepare_prompt_boundary_cache_store(
|
||
request_id,
|
||
request,
|
||
uid_for_store,
|
||
)
|
||
)
|
||
|
||
if (
|
||
hasattr(request, "_extracted_cache")
|
||
and request._extracted_cache is not None
|
||
) or prompt_boundary_store is not None:
|
||
try:
|
||
full_token_sequence = list(request.prompt_token_ids) + list(
|
||
request.output_token_ids
|
||
)
|
||
|
||
if prompt_boundary_store is not None:
|
||
(
|
||
token_sequence_to_store,
|
||
cache_to_store,
|
||
model_cache_config,
|
||
intermediate_snapshots,
|
||
) = prompt_boundary_store
|
||
cacheable_sequence = list(token_sequence_to_store)
|
||
else:
|
||
# For reasoning models, only cache prompt tokens.
|
||
# Output contains <think> tokens that the API layer
|
||
# strips before the next turn, so they never match.
|
||
if getattr(request, "needs_think_prefix", False):
|
||
cacheable_sequence = list(request.prompt_token_ids)
|
||
else:
|
||
cacheable_sequence = full_token_sequence
|
||
token_sequence_to_store = cacheable_sequence
|
||
cache_to_store = request._extracted_cache
|
||
model_cache_config = getattr(
|
||
request, "_model_cache_config", None
|
||
)
|
||
intermediate_snapshots = None
|
||
|
||
# Inference-thread store_cache prep, timed as
|
||
# three sub-phases (boundary / collect / dispatch)
|
||
# mirroring boundary_capture_* granularity.
|
||
# The dispatch phase does a FULL mx.eval (not
|
||
# async_eval) so the KV arrays are concrete on THIS
|
||
# thread before the store-cache worker slices/views
|
||
# them. MLX streams are thread-local; a lazy op left
|
||
# for the worker to materialize re-dispatches to this
|
||
# thread's stream index, which is absent on the worker
|
||
# -> SIGABRT. See the dispatch-phase comment below.
|
||
with mx.stream(self._stream):
|
||
with self._phase_timer("store_cache_main_boundary"):
|
||
if prompt_boundary_store is None:
|
||
boundary_override = (
|
||
self._get_boundary_store_override(
|
||
request_id,
|
||
cacheable_sequence,
|
||
)
|
||
)
|
||
if (
|
||
boundary_override is None
|
||
and self._detect_boundary_snapshot_need()
|
||
):
|
||
# Non-sliceable cache state is only
|
||
# storable from boundary-aligned
|
||
# snapshots; the live state sits at
|
||
# the current decode offset (which
|
||
# speculative decode can leave off
|
||
# the emitted count entirely).
|
||
raise _BoundaryStoreUnavailable()
|
||
if boundary_override is not None:
|
||
(
|
||
token_sequence_to_store,
|
||
boundary_cache,
|
||
boundary_model_config,
|
||
intermediate_snapshots,
|
||
) = boundary_override
|
||
cache_to_store = (
|
||
self._merge_boundary_with_full_cache(
|
||
boundary_cache,
|
||
request._extracted_cache,
|
||
)
|
||
)
|
||
if boundary_model_config is not None:
|
||
model_cache_config = (
|
||
boundary_model_config
|
||
)
|
||
logger.info(
|
||
f"Using boundary cache snapshot for {request_id}: "
|
||
f"storing {len(token_sequence_to_store)}/"
|
||
f"{len(full_token_sequence)} tokens "
|
||
f"(skipping trailing partial block, "
|
||
f"{len(intermediate_snapshots) if intermediate_snapshots else 0} "
|
||
f"intermediate snapshots)"
|
||
)
|
||
# DEBUG-only divergence probe (issue #1003).
|
||
# Record the exact token sequence submitted to
|
||
# store_cache, including boundary truncation.
|
||
if logger.isEnabledFor(logging.DEBUG):
|
||
self._cache_probe_seqs.append(
|
||
(
|
||
request.request_id,
|
||
list(token_sequence_to_store),
|
||
)
|
||
)
|
||
with self._phase_timer("store_cache_main_collect"):
|
||
pre_eval_arrays = (
|
||
self._collect_arrays_from_extracted_cache(
|
||
cache_to_store
|
||
)
|
||
)
|
||
if intermediate_snapshots is not None:
|
||
for (
|
||
snapshot_cache
|
||
) in (
|
||
intermediate_snapshots.iter_in_memory_extracted()
|
||
):
|
||
pre_eval_arrays.extend(
|
||
self._collect_arrays_from_extracted_cache(
|
||
snapshot_cache
|
||
)
|
||
)
|
||
with self._phase_timer("store_cache_main_dispatch"):
|
||
if pre_eval_arrays:
|
||
# FULL eval (not async_eval) on the owner
|
||
# thread. MLX streams are thread-local:
|
||
# these KV arrays carry self._stream
|
||
# (a per-engine ThreadLocalStream created
|
||
# on THIS thread). The store-cache worker
|
||
# later slices them (_extract_block_tensor_slice)
|
||
# and views bf16->uint16 (_extract_tensor_bytes);
|
||
# if the source op is still LAZY at that
|
||
# point, materializing it on the worker
|
||
# re-dispatches to self._stream's index,
|
||
# which does not exist on the worker thread
|
||
# -> "There is no Stream(gpu, N) in current
|
||
# thread" -> std::terminate -> SIGABRT.
|
||
# Forcing concrete materialization here
|
||
# means every downstream worker op consumes
|
||
# an already-evaluated buffer and binds its
|
||
# own new ops to the always-present default
|
||
# stream (gpu,0). The big host memcpy
|
||
# (bytes(memoryview(...))) and the disk
|
||
# write stay on the worker — only the GPU
|
||
# completion fence moves onto this thread.
|
||
mx.eval(*pre_eval_arrays)
|
||
|
||
hot_cache_write_back = (
|
||
not self._bypass_hot_cache_under_pressure()
|
||
)
|
||
if not hot_cache_write_back:
|
||
logger.info(
|
||
"Using SSD write-through for %s "
|
||
"under memory pressure",
|
||
request_id,
|
||
)
|
||
|
||
if self._store_cache_executor is not None:
|
||
# Hand host memcpy and disk write to the
|
||
# background executor after the owner thread
|
||
# has materialized KV arrays. The gate counts
|
||
# cleanup slots that still own extracted cache
|
||
# references; backpressure is applied at
|
||
# admission in _schedule_waiting so cache
|
||
# persistence does not wait in the token loop
|
||
# after submission (#1496). note_submitted is
|
||
# called before submit, and note_done happens in
|
||
# _drain_pending_async_removes after the request
|
||
# cache references are released.
|
||
gate = self._store_cache_gate
|
||
if gate is not None:
|
||
gate.note_submitted()
|
||
try:
|
||
store_future = self._store_cache_executor.submit(
|
||
self._async_store_cache_worker,
|
||
request_id,
|
||
token_sequence_to_store,
|
||
cache_to_store,
|
||
model_cache_config,
|
||
intermediate_snapshots,
|
||
request.vlm_extra_keys_for_cache,
|
||
request.vlm_extra_key_token_start_for_cache,
|
||
request.vlm_extra_key_ranges_for_cache,
|
||
hot_cache_write_back,
|
||
)
|
||
except BaseException:
|
||
if gate is not None:
|
||
gate.note_done()
|
||
raise
|
||
self._inflight_store_futures[request_id] = store_future
|
||
self._inflight_store_info[request_id] = (
|
||
_InflightStoreInfo(
|
||
tokens=list(token_sequence_to_store),
|
||
extra_keys=request.vlm_extra_keys_for_cache,
|
||
extra_key_token_start=(
|
||
request.vlm_extra_key_token_start_for_cache
|
||
),
|
||
extra_key_ranges=(
|
||
request.vlm_extra_key_ranges_for_cache
|
||
),
|
||
)
|
||
)
|
||
else:
|
||
# Executor unavailable — synchronous fallback.
|
||
self._async_store_cache_worker(
|
||
request_id,
|
||
token_sequence_to_store,
|
||
cache_to_store,
|
||
model_cache_config,
|
||
intermediate_snapshots,
|
||
request.vlm_extra_keys_for_cache,
|
||
request.vlm_extra_key_token_start_for_cache,
|
||
request.vlm_extra_key_ranges_for_cache,
|
||
hot_cache_write_back,
|
||
)
|
||
logger.debug(
|
||
f"Submitted async store_cache for {request_id} "
|
||
f"({len(token_sequence_to_store)} tokens, "
|
||
f"{len(full_token_sequence)} total: "
|
||
f"{len(request.prompt_token_ids)} prompt + "
|
||
f"{len(request.output_token_ids)} output)"
|
||
)
|
||
except _BoundaryStoreUnavailable:
|
||
logger.debug(
|
||
"Skipping cache store for %s: no boundary-aligned "
|
||
"snapshot for non-sliceable cache state (all "
|
||
"captures skipped, e.g. by the speculative-decode "
|
||
"skew guard); storing live state would corrupt "
|
||
"later prefix hits",
|
||
request_id,
|
||
)
|
||
block_table = None
|
||
if self.paged_cache_manager:
|
||
block_table = self.paged_cache_manager.get_block_table(
|
||
request_id
|
||
)
|
||
if block_table and self.paged_cache_manager:
|
||
self.paged_cache_manager.release_for_eviction(
|
||
block_table.block_ids
|
||
)
|
||
self.block_aware_cache.clear_request_entry(request_id)
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"Failed to submit async store for {request_id}: {e}"
|
||
)
|
||
else:
|
||
# No extracted_cache to store, but ensure block leak guard.
|
||
block_table = None
|
||
if self.paged_cache_manager:
|
||
block_table = self.paged_cache_manager.get_block_table(
|
||
request_id
|
||
)
|
||
if block_table is None and hasattr(request, "block_table"):
|
||
block_table = request.block_table
|
||
if block_table and self.paged_cache_manager:
|
||
self.paged_cache_manager.release_for_eviction(
|
||
block_table.block_ids
|
||
)
|
||
self.block_aware_cache.clear_request_entry(request_id)
|
||
|
||
# Remove from running
|
||
if request_id in self.running:
|
||
del self.running[request_id]
|
||
|
||
# batch_generator.remove(uid): defer until the async store_cache
|
||
# worker finishes so the BatchKVCache slot isn't reused while the
|
||
# worker is still reading buffer references via cache_to_store.
|
||
# _drain_pending_async_removes (next step) handles the actual
|
||
# mx.synchronize + remove + uid_maps cleanup. If we have no async
|
||
# store (no extracted_cache, executor missing, fallback fail),
|
||
# fall back to immediate remove for back-compat behavior.
|
||
if request_id in self.request_id_to_uid:
|
||
uid = self.request_id_to_uid[request_id]
|
||
if store_future is not None:
|
||
self._pending_async_removes.append((uid, request_id, store_future))
|
||
else:
|
||
# Synchronize in-flight GPU work before modifying batch state.
|
||
# batch_generator.remove() triggers lazy KV cache array slicing
|
||
# (BatchKVCache.filter) that replaces references to arrays still
|
||
# used by in-flight Metal command buffers from the previous
|
||
# batch_generator.next() call. Without this barrier the Metal
|
||
# driver can hit 'completeMemory() prepare count underflow'.
|
||
_safe_sync_stream(self._stream)
|
||
self._remove_uid_from_active_batch(uid)
|
||
if hasattr(self.model, "unregister_rope_delta"):
|
||
self.model.unregister_rope_delta(uid)
|
||
_unregister_uid_row(self.model, uid)
|
||
if uid in self.uid_to_request_id:
|
||
del self.uid_to_request_id[uid]
|
||
del self.request_id_to_uid[request_id]
|
||
|
||
# Clean up streaming detokenizer
|
||
self._cleanup_detokenizer(request_id)
|
||
|
||
# Clean up protocol-specific output parser session
|
||
self._cleanup_output_parser_session(request_id)
|
||
|
||
# Clean up VLM adapter state (position_ids, rope_deltas, pending embeddings)
|
||
if hasattr(self.model, "clear_vlm_position_state"):
|
||
self.model.clear_vlm_position_state()
|
||
if hasattr(self.model, "clear_pending_embeddings"):
|
||
self.model.clear_pending_embeddings()
|
||
|
||
# Drop any boundary snapshot for this request. The in-memory
|
||
# dict pop is safe — the async store worker holds its own
|
||
# reference to the snapshot dict via _BoundarySnapshotProvider.
|
||
self._boundary_cache_snapshots.pop(request_id, None)
|
||
# cleanup_request rmtree's the on-disk snapshot directory and
|
||
# races the worker's boundary_snapshot_store.load() calls. If
|
||
# an async store_future is in flight, defer cleanup until the
|
||
# worker finishes (handled in _drain_pending_async_removes).
|
||
if self._boundary_snapshot_store is not None and store_future is None:
|
||
self._boundary_snapshot_store.cleanup_request(request_id)
|
||
|
||
# Track as finished
|
||
self.finished_req_ids.add(request_id)
|
||
|
||
# Remove from requests dict to prevent memory leak.
|
||
# When async store_cache is in flight, keep _extracted_cache alive
|
||
# until the worker finishes — the worker holds a reference via
|
||
# cache_to_store argument, but request._extracted_cache pointing
|
||
# to the same data is the canonical owner. We pop here only when
|
||
# no future is pending; the future's done callback (or
|
||
# _drain_pending_async_removes) clears the request later.
|
||
if store_future is None:
|
||
req_to_remove = self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
if req_to_remove is not None:
|
||
req_to_remove._extracted_cache = None
|
||
req_to_remove.prompt_cache = None
|
||
else:
|
||
# Drop request from running but keep in self.requests so the
|
||
# async worker keeps the cache buffers alive via reachability.
|
||
# Cleanup happens in _drain_pending_async_removes.
|
||
pass
|
||
|
||
# Emit phase timing diagnostics when accumulated counts are meaningful.
|
||
# Helps diagnose cache-on overhead (boundary capture / store_cache /
|
||
# hot cache eviction). Logged at info level so operators can see it
|
||
# without enabling debug.
|
||
if finished_ids and self._phase_total_ms:
|
||
stats_parts = []
|
||
for phase, total_ms in sorted(self._phase_total_ms.items()):
|
||
count = self._phase_count.get(phase, 0)
|
||
if count == 0:
|
||
continue
|
||
stats_parts.append(f"{phase}={total_ms:.1f}ms/{count}")
|
||
if stats_parts:
|
||
logger.info("Cache phase timings: %s", ", ".join(stats_parts))
|
||
|
||
# Schedule deferred Metal cache cleanup after request completion.
|
||
if finished_ids:
|
||
# Schedule deferred Metal cache cleanup instead of clearing immediately.
|
||
# Immediate mx.clear_cache() after request completion races with IOKit's
|
||
# asynchronous completeMemory() callbacks — the kernel-level GPU memory
|
||
# reference counting can still be in-flight even after mx.synchronize()
|
||
# returns, causing 'prepare count underflow' kernel panics (#435).
|
||
# Deferring by _DEFERRED_CLEAR_DELAY generation steps (~10-40 ms) gives
|
||
# IOKit time to process callbacks while still reclaiming buffers fast
|
||
# enough to prevent TTFT spikes from pool bloat (#411).
|
||
#
|
||
# Use max() so that concurrent completions (max_num_seqs > 1) each get
|
||
# a full _DEFERRED_CLEAR_DELAY window counted from *their own* finish
|
||
# step. The old "only set if None" guard meant the second request's
|
||
# window was anchored to the first request's finish step, allowing the
|
||
# second request's KV cache blocks to be re-allocated before IOKit
|
||
# finished their completeMemory() callbacks (#557).
|
||
target = self._step_counter + self._DEFERRED_CLEAR_DELAY
|
||
if self._deferred_clear_at is None or target > self._deferred_clear_at:
|
||
self._deferred_clear_at = target
|
||
|
||
def _is_cache_corruption_error(self, error: Exception) -> bool:
|
||
"""Check if an error indicates cache corruption."""
|
||
return is_cache_corruption_error(error)
|
||
|
||
def _is_generation_overflow_error(self, error: Exception) -> bool:
|
||
"""Check for MLX/libc++ unordered-container overflow during decode."""
|
||
return isinstance(
|
||
error, OverflowError
|
||
) and self._GENERATION_OVERFLOW_PATTERN in str(error)
|
||
|
||
def _recover_from_cache_error(self) -> None:
|
||
"""Recover from cache corruption error."""
|
||
# Clear batch generator (this is the source of the corruption)
|
||
self.batch_generator = None
|
||
self._current_sampler_params = None
|
||
self._boundary_cache_snapshots.clear()
|
||
if self._boundary_snapshot_store is not None:
|
||
self._boundary_snapshot_store.cleanup_all()
|
||
self._boundary_snapshot_required = None
|
||
|
||
# Clear stale VLM position state to prevent re-corruption on retry
|
||
if hasattr(self.model, "clear_vlm_position_state"):
|
||
self.model.clear_vlm_position_state()
|
||
|
||
# Clear pending VLM embeddings
|
||
if hasattr(self.model, "clear_pending_embeddings"):
|
||
self.model.clear_pending_embeddings()
|
||
|
||
# Clear caches
|
||
if self.block_aware_cache is not None:
|
||
self.block_aware_cache.clear()
|
||
self._cache_rate_tracker.clear()
|
||
|
||
# Clear UID mappings
|
||
_unregister_uid_rows_for_model(self.model)
|
||
self.request_id_to_uid.clear()
|
||
self.uid_to_request_id.clear()
|
||
|
||
# Cancel any pending deferred Metal cache clear
|
||
self._deferred_clear_at = None
|
||
|
||
# Clear detokenizer state to prevent contamination after recovery
|
||
self._request_detokenizers.clear()
|
||
|
||
# Clear protocol-specific output parser sessions
|
||
self._output_parser_sessions.clear()
|
||
|
||
logger.info("Cache recovery completed")
|
||
|
||
def _recover_from_generation_overflow_error(self) -> None:
|
||
"""Reset decode state after MLX __next_prime overflow."""
|
||
self.batch_generator = None
|
||
self._current_sampler_params = None
|
||
self._boundary_snapshot_required = None
|
||
|
||
active_specprefill = self._specprefill_active_request_id
|
||
if active_specprefill is not None:
|
||
self._cleanup_specprefill(active_specprefill)
|
||
|
||
if hasattr(self.model, "clear_vlm_position_state"):
|
||
self.model.clear_vlm_position_state()
|
||
if hasattr(self.model, "clear_pending_embeddings"):
|
||
self.model.clear_pending_embeddings()
|
||
|
||
_unregister_uid_rows_for_model(self.model)
|
||
self.request_id_to_uid.clear()
|
||
self.uid_to_request_id.clear()
|
||
self._deferred_clear_at = None
|
||
self._request_detokenizers.clear()
|
||
self._output_parser_sessions.clear()
|
||
|
||
try:
|
||
_sync_and_clear_cache(self._stream)
|
||
except Exception as e:
|
||
logger.warning(
|
||
"Metal cache clear failed during generation overflow recovery: %s",
|
||
e,
|
||
)
|
||
|
||
logger.info("Generation overflow recovery completed")
|
||
|
||
def _reset_request_for_reprefill(self, request: Request) -> None:
|
||
"""Reset request-owned decode state so it can be prefilled again."""
|
||
request.status = RequestStatus.WAITING
|
||
request.batch_uid = None
|
||
request.prompt_cache = None
|
||
request.cached_tokens = 0
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
request.block_table = None
|
||
request.shared_prefix_blocks = 0
|
||
request.output_token_ids = []
|
||
request.output_text = ""
|
||
request.num_computed_tokens = 0
|
||
request._extracted_cache = None
|
||
request._model_cache_config = None
|
||
request.think_prefix_sent = False
|
||
|
||
def _reschedule_running_requests(
|
||
self, is_corruption: bool = False, max_corruption_retries: int = 3
|
||
) -> list[str]:
|
||
"""Move running requests back to waiting queue for retry.
|
||
|
||
Args:
|
||
is_corruption: If True, increment corruption retry counter and
|
||
fail requests that exceed max_corruption_retries.
|
||
max_corruption_retries: Max corruption retries before failing a request.
|
||
|
||
Returns:
|
||
List of request IDs that exceeded max retries (corruption only).
|
||
"""
|
||
failed_ids: list[str] = []
|
||
count = 0
|
||
for request_id, request in list(self.running.items()):
|
||
if is_corruption:
|
||
request.cache_corruption_retries += 1
|
||
if request.cache_corruption_retries > max_corruption_retries:
|
||
failed_ids.append(request_id)
|
||
del self.running[request_id]
|
||
# Clean up from requests dict (prevent memory leak)
|
||
req = self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
if req is not None:
|
||
req._extracted_cache = None
|
||
req.prompt_cache = None
|
||
continue
|
||
|
||
self._reset_request_for_reprefill(request)
|
||
|
||
# Move to waiting queue (at front for priority)
|
||
self.waiting.appendleft(request)
|
||
del self.running[request_id]
|
||
count += 1
|
||
|
||
if count > 0:
|
||
logger.info(f"Rescheduled {count} requests for re-prefill")
|
||
return failed_ids
|
||
|
||
def _reschedule_generation_overflow_requests(
|
||
self,
|
||
max_generation_overflow_retries: int = _MAX_GENERATION_OVERFLOW_RETRIES,
|
||
) -> list[str]:
|
||
"""Retry active requests serially after MLX generation overflow."""
|
||
retry_candidates: list[Request] = []
|
||
seen: set[str] = set()
|
||
prefilling_ids = {request.request_id for request in self.prefilling}
|
||
waiting_ids = {request.request_id for request in self.waiting}
|
||
|
||
def collect(request: Request | None) -> None:
|
||
if request is None:
|
||
return
|
||
request_id = request.request_id
|
||
if request_id in seen or request_id in self._inflight_store_futures:
|
||
return
|
||
if request.is_finished():
|
||
return
|
||
seen.add(request_id)
|
||
retry_candidates.append(request)
|
||
|
||
for request in self.running.values():
|
||
collect(request)
|
||
for request in self.prefilling:
|
||
collect(request)
|
||
for request_id, request in self.requests.items():
|
||
if request_id in self.running or request_id in prefilling_ids:
|
||
continue
|
||
if request_id in waiting_ids:
|
||
continue
|
||
collect(request)
|
||
|
||
collected_ids = {request.request_id for request in retry_candidates}
|
||
for request_id in collected_ids:
|
||
self.running.pop(request_id, None)
|
||
self._prefill_states.pop(request_id, None)
|
||
if collected_ids:
|
||
self.prefilling = deque(
|
||
request
|
||
for request in self.prefilling
|
||
if request.request_id not in collected_ids
|
||
)
|
||
|
||
failed_ids: list[str] = []
|
||
retryable: list[Request] = []
|
||
self._generation_overflow_recovery_ids.clear()
|
||
for request in retry_candidates:
|
||
request.generation_overflow_retries += 1
|
||
request_id = request.request_id
|
||
self._boundary_cache_snapshots.pop(request_id, None)
|
||
if self._boundary_snapshot_store is not None:
|
||
self._boundary_snapshot_store.cleanup_request(request_id)
|
||
get_prefill_tracker().remove(request_id)
|
||
if request.generation_overflow_retries > max_generation_overflow_retries:
|
||
failed_ids.append(request_id)
|
||
req = self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
if req is not None:
|
||
req._extracted_cache = None
|
||
req.prompt_cache = None
|
||
continue
|
||
|
||
self._reset_request_for_reprefill(request)
|
||
retryable.append(request)
|
||
self._generation_overflow_recovery_ids.add(request_id)
|
||
|
||
for request in reversed(retryable):
|
||
self.waiting.appendleft(request)
|
||
|
||
if retryable:
|
||
logger.info(
|
||
"Rescheduled %d request(s) for serial generation-overflow retry",
|
||
len(retryable),
|
||
)
|
||
return failed_ids
|
||
|
||
# Max times a single request is requeued after a prefill memory-pressure
|
||
# failure before we give up and emit a clean error to the client.
|
||
_MAX_PREFILL_OOM_RETRIES = 2
|
||
|
||
def _requeue_or_fail_prefill(self, request: "Request", error: Exception) -> bool:
|
||
"""Decide whether to requeue a prefill that hit the memory ceiling.
|
||
|
||
The three #1405 catch sites have already torn the request down
|
||
(released paged cache, popped ``self.requests``, removed the prefill
|
||
tracker entry, cleared Metal). This either resets the request and puts
|
||
it back on the waiting queue for a fresh attempt (returns ``True`` —
|
||
caller continues without emitting an error), or — when the retry
|
||
budget is exhausted or the failure is not a memory-pressure error —
|
||
returns ``False`` so the caller emits the clean
|
||
``finish_reason="error"``.
|
||
|
||
Only memory-limit failures are retried; any other RuntimeError fails
|
||
immediately so genuine model errors don't loop.
|
||
"""
|
||
if "Memory limit exceeded" not in str(error):
|
||
return False
|
||
if request.prefill_oom_retries >= self._MAX_PREFILL_OOM_RETRIES:
|
||
logger.warning(
|
||
"Prefill for %s exhausted %d memory-pressure retries; "
|
||
"failing with a clean error.",
|
||
request.request_id,
|
||
self._MAX_PREFILL_OOM_RETRIES,
|
||
)
|
||
return False
|
||
request.prefill_oom_retries += 1
|
||
|
||
# Reclaim before requeue so the retry starts from a lower baseline.
|
||
self._reclaim_prefill_headroom()
|
||
|
||
# Clear any SpecPrefill RoPE patch tied to this request so the retry
|
||
# re-scores cleanly.
|
||
if self._specprefill_active_request_id == request.request_id:
|
||
self._specprefill_active_request_id = None
|
||
|
||
# Restore mRoPE deltas if an external VLM prefill was interrupted before
|
||
# its own restore ran (value stashed on the request in
|
||
# _do_external_prefill). Benign for non-VLM requests (stash is None).
|
||
saved = getattr(request, "_prefill_saved_rope_deltas", None)
|
||
if saved is not None:
|
||
lm = getattr(self.model, "_language_model", None)
|
||
if lm is not None and hasattr(lm, "_rope_deltas"):
|
||
lm._rope_deltas = saved
|
||
request._prefill_saved_rope_deltas = None
|
||
|
||
# Reset scheduling + cache + output state to a clean pre-prefill state
|
||
# (mirrors _reschedule_running_requests). We deliberately drop
|
||
# cached_tokens / block_table so the retry does a cold full prefill and
|
||
# does not re-attach the large cached prefix that produced the same
|
||
# oversized SDPA span. VLM inputs/embeds are preserved.
|
||
request.status = RequestStatus.WAITING
|
||
request.batch_uid = None
|
||
request.prompt_cache = None
|
||
request.cached_tokens = 0
|
||
request.remaining_tokens = request.prompt_token_ids
|
||
request.block_table = None
|
||
request.shared_prefix_blocks = 0
|
||
request.output_token_ids = []
|
||
request.output_text = ""
|
||
request.num_computed_tokens = 0
|
||
request._extracted_cache = None
|
||
request._model_cache_config = None
|
||
request.think_prefix_sent = False
|
||
|
||
# Re-register (the catch site popped it) and requeue at the front. The
|
||
# retry is throttled from its first chunk by the now-populated transient
|
||
# EWMA, so it is strictly better-informed than this attempt.
|
||
self.requests[request.request_id] = request
|
||
self.waiting.appendleft(request)
|
||
logger.warning(
|
||
"Requeued %s for prefill retry %d/%d after memory pressure.",
|
||
request.request_id,
|
||
request.prefill_oom_retries,
|
||
self._MAX_PREFILL_OOM_RETRIES,
|
||
)
|
||
return True
|
||
|
||
def _pause_for_prefill_eviction(
|
||
self,
|
||
request: "Request",
|
||
eviction: PrefillEvictionRequest,
|
||
) -> None:
|
||
"""Hold a request until EngineCore can evict idle models asynchronously.
|
||
|
||
The request's prefix-cache state (prompt_cache, block_table,
|
||
cached_tokens, remaining_tokens) is deliberately left untouched so
|
||
a reconstructed prefix survives the pause and the retry prefills
|
||
only the uncached suffix instead of recomputing the prompt cold
|
||
(#2180).
|
||
"""
|
||
self._pending_prefill_eviction_request = eviction
|
||
request.status = RequestStatus.WAITING
|
||
request.batch_uid = None
|
||
self.waiting.appendleft(request)
|
||
logger.info(
|
||
"Paused request %s for prefill LRU eviction (reason=%s)",
|
||
request.request_id,
|
||
eviction.reason,
|
||
)
|
||
|
||
def step(self) -> 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. On cache corruption: clears all cache and reschedules requests
|
||
for re-prefill (no error raised to caller)
|
||
|
||
Returns:
|
||
SchedulerOutput with results of this step
|
||
"""
|
||
output = SchedulerOutput()
|
||
|
||
# Process pending aborts FIRST (thread-safe with hybrid executor)
|
||
self._process_pending_aborts()
|
||
|
||
# Drain a deferred between-turn reclaim requested by the memory
|
||
# enforcer (only acts when the scheduler is idle).
|
||
self._process_pending_reclaim()
|
||
|
||
# Drain async store_cache completions from prior steps. Each completed
|
||
# entry triggers the deferred batch_generator.remove(uid) on the
|
||
# inference thread. Inflight entries are left for a later step.
|
||
drained_async_removes = self._drain_pending_async_removes()
|
||
if drained_async_removes:
|
||
output.has_work = True
|
||
|
||
# Check memory pressure and evict if needed (tiered cache)
|
||
if self.memory_monitor is not None:
|
||
self._check_memory_pressure()
|
||
|
||
try:
|
||
# Advance in-flight chunked prefills (one chunk per request).
|
||
# Must run before _schedule_waiting() so that completing prefills
|
||
# are inserted into BatchGenerator before the decode step.
|
||
chunked_scheduled: list[Request] = []
|
||
chunked_rejected: list[RequestOutput] = []
|
||
if self.prefilling:
|
||
self._advance_chunked_prefills(chunked_scheduled, chunked_rejected)
|
||
|
||
# Schedule waiting requests
|
||
scheduled, rejected = self._schedule_waiting()
|
||
# Merge chunked-prefill completions into the scheduled list.
|
||
if chunked_scheduled:
|
||
scheduled = chunked_scheduled + scheduled
|
||
output.scheduled_request_ids = [r.request_id for r in scheduled]
|
||
output.num_scheduled_tokens = sum(r.num_prompt_tokens for r in scheduled)
|
||
if chunked_rejected:
|
||
output.outputs.extend(chunked_rejected)
|
||
output.has_work = True
|
||
if rejected:
|
||
output.outputs.extend(rejected)
|
||
output.has_work = True
|
||
if self._pending_prefill_eviction_request is not None:
|
||
output.prefill_eviction_request = self._pending_prefill_eviction_request
|
||
self._pending_prefill_eviction_request = None
|
||
output.has_work = True
|
||
|
||
# Run generation step if we have running requests.
|
||
# Use next_generated() which returns only GenerationBatch.Response
|
||
# objects (prefill is handled externally before insert).
|
||
if (
|
||
self.batch_generator is not None or self._vlm_mtp_active
|
||
) and self.running:
|
||
if self.batch_generator is not None:
|
||
responses = list(self.batch_generator.next_generated())
|
||
else:
|
||
responses = []
|
||
# Drive vlm_mtp generators alongside BatchGenerator. Order
|
||
# matters only for log determinism; _process_batch_responses
|
||
# is per-uid.
|
||
if self._vlm_mtp_active:
|
||
responses.extend(self._step_vlm_mtp())
|
||
output.has_work = True
|
||
|
||
if responses:
|
||
outputs, finished_ids = self._process_batch_responses(responses)
|
||
output.outputs.extend(outputs)
|
||
output.finished_request_ids.update(finished_ids)
|
||
|
||
# Periodic decode cache materialization for models whose
|
||
# KV cache update graph can otherwise grow for thousands of
|
||
# tokens. MiniMax-M3 has one lazy cache-update chain per
|
||
# layer; evaluating the cache state periodically cuts those
|
||
# references before Metal's resource-count limit is hit.
|
||
self._tokens_since_kv_cache_eval = getattr(
|
||
self, "_tokens_since_kv_cache_eval", 0
|
||
) + len(responses)
|
||
kv_eval_interval = self._decode_eval_kv_cache_interval
|
||
if (
|
||
kv_eval_interval > 0
|
||
and self._tokens_since_kv_cache_eval >= kv_eval_interval
|
||
):
|
||
with mx.stream(self._stream):
|
||
evaluated = _eval_generation_batch_cache(
|
||
self.batch_generator
|
||
)
|
||
logger.debug(
|
||
"Materialized decode KV cache state: %d arrays",
|
||
evaluated,
|
||
)
|
||
self._tokens_since_kv_cache_eval = 0
|
||
|
||
self._cleanup_finished(finished_ids)
|
||
|
||
# Periodic Metal allocator cleanup during long decodes.
|
||
# mx.random.categorical inside the sampler allocates a
|
||
# tiny scalar via gumbel → uniform on every call.
|
||
# omlx ships its own non-compiled sampler
|
||
# (omlx/utils/sampling.py) so that RNG state actually
|
||
# advances in the server, but the trade-off is that
|
||
# those scalars accumulate in the IOGPU residency set
|
||
# — macOS aborts at ~4096 entries. Long contexts
|
||
# (50k+) decoding thousands of tokens hit that limit
|
||
# mid-stream. Synchronise the generation stream first
|
||
# so any in-flight Metal command buffer that still
|
||
# references buffers we're about to drop has
|
||
# completed; the allocator only releases pool entries
|
||
# whose ref count is zero, but the sync guarantees
|
||
# there is no race window. Decode-only path —
|
||
# next_generated() returns nothing during prefill, so
|
||
# we never disrupt prefill activation buffers.
|
||
self._tokens_since_clear_cache = getattr(
|
||
self, "_tokens_since_clear_cache", 0
|
||
) + len(responses)
|
||
if self._tokens_since_clear_cache >= 1024:
|
||
_sync_and_clear_cache(self._stream)
|
||
self._tokens_since_clear_cache = 0
|
||
|
||
except _PrefillAbortedError:
|
||
# Prefill was interrupted by a pending abort.
|
||
# BatchGenerator is in an inconsistent state (partial
|
||
# prefill), so reset it entirely. Pending aborts will
|
||
# be processed at the start of the next step().
|
||
self.batch_generator = None
|
||
self._current_sampler_params = None
|
||
self._boundary_cache_snapshots.clear()
|
||
if self._boundary_snapshot_store is not None:
|
||
self._boundary_snapshot_store.cleanup_all()
|
||
self._boundary_snapshot_required = None
|
||
# Move any running requests back to waiting so they
|
||
# can be rescheduled with a fresh BatchGenerator.
|
||
self._reschedule_running_requests()
|
||
|
||
except (TypeError, AttributeError, ValueError) as e:
|
||
if self._is_cache_corruption_error(e):
|
||
import traceback
|
||
|
||
logger.warning(
|
||
f"Cache corruption detected: {e}, "
|
||
f"clearing cache and re-prefilling..."
|
||
)
|
||
logger.debug(f"Cache corruption traceback:\n{traceback.format_exc()}")
|
||
# Full reset: clear batch generator, all caches, VLM state
|
||
self._recover_from_cache_error()
|
||
# Reschedule requests for re-prefill from scratch.
|
||
# Requests exceeding max corruption retries are failed.
|
||
failed_ids = self._reschedule_running_requests(is_corruption=True)
|
||
for rid in failed_ids:
|
||
output.outputs.append(
|
||
RequestOutput(
|
||
request_id=rid,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=(
|
||
f"Cache corruption not recoverable "
|
||
f"after retries: {e}"
|
||
),
|
||
)
|
||
)
|
||
output.finished_request_ids.add(rid)
|
||
else:
|
||
raise
|
||
|
||
except OverflowError as e:
|
||
if self._is_generation_overflow_error(e):
|
||
import traceback
|
||
|
||
logger.warning(
|
||
"Generation overflow detected: %s; resetting decode state "
|
||
"and retrying affected requests serially",
|
||
e,
|
||
)
|
||
logger.debug(
|
||
"Generation overflow traceback:\n%s", traceback.format_exc()
|
||
)
|
||
self._recover_from_generation_overflow_error()
|
||
failed_ids = self._reschedule_generation_overflow_requests()
|
||
for rid in failed_ids:
|
||
output.outputs.append(
|
||
RequestOutput(
|
||
request_id=rid,
|
||
finished=True,
|
||
finish_reason="error",
|
||
error=(
|
||
"Generation overflow not recoverable after "
|
||
f"serial retry: {e}"
|
||
),
|
||
)
|
||
)
|
||
output.finished_request_ids.add(rid)
|
||
output.has_work = True
|
||
else:
|
||
raise
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
|
||
logger.error(
|
||
f"Error in batch generation step: {e}\n" f"{traceback.format_exc()}"
|
||
)
|
||
raise
|
||
|
||
# Clear finished tracking for next step
|
||
self.finished_req_ids = set()
|
||
self._refresh_generation_overflow_recovery_ids()
|
||
|
||
# Periodic Metal cache cleanup
|
||
self._step_counter += 1
|
||
should_clear = self._should_periodic_clear_cache()
|
||
# Deferred post-completion cleanup: fire once the step counter reaches
|
||
# the target set by _cleanup_finished() (#435, #557).
|
||
if (
|
||
self._deferred_clear_at is not None
|
||
and self._step_counter >= self._deferred_clear_at
|
||
):
|
||
should_clear = True
|
||
self._deferred_clear_at = None
|
||
if should_clear:
|
||
_sync_and_clear_cache(self._stream)
|
||
if (
|
||
self.config.gc_cleanup_interval > 0
|
||
and self._step_counter % self.config.gc_cleanup_interval == 0
|
||
):
|
||
gc.collect()
|
||
|
||
self._publish_admin_snapshot()
|
||
|
||
return output
|
||
|
||
def _publish_admin_snapshot(self) -> None:
|
||
"""Atomically publish a fresh admin-visible snapshot.
|
||
|
||
Called from step() on the engine thread, where running/waiting are
|
||
not concurrently mutated. The admin endpoint reads the reference via
|
||
snapshot_for_admin() and never iterates the live structures.
|
||
"""
|
||
self._admin_snapshot = {
|
||
"running_by_id": dict(self.running),
|
||
"waiting": list(self.waiting),
|
||
}
|
||
|
||
def snapshot_for_admin(self) -> dict[str, Any]:
|
||
"""Return the most recently published admin snapshot.
|
||
|
||
Reference read is GIL-atomic; the dict itself is no longer mutated
|
||
after publication. May be one step stale, which is fine for dashboard
|
||
polling.
|
||
"""
|
||
return self._admin_snapshot
|
||
|
||
def get_request(self, request_id: str) -> Request | None:
|
||
"""Get a request by ID."""
|
||
return self.requests.get(request_id)
|
||
|
||
def remove_finished_request(self, request_id: str) -> Request | None:
|
||
"""Remove a finished request from tracking."""
|
||
request = self.requests.pop(request_id, None)
|
||
self._clear_request_admission_bookkeeping(request_id)
|
||
return request
|
||
|
||
def get_stats(self) -> dict[str, Any]:
|
||
"""Get scheduler statistics."""
|
||
stats = {
|
||
"num_waiting": len(self.waiting),
|
||
"num_prefilling": len(self.prefilling),
|
||
"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,
|
||
}
|
||
# Include cache stats
|
||
if self.block_aware_cache is not None:
|
||
stats["ssd_cache"] = self.block_aware_cache.get_stats()
|
||
return stats
|
||
|
||
def get_cache_stats(self) -> dict[str, Any] | None:
|
||
"""Get cache statistics."""
|
||
if self.block_aware_cache is not None:
|
||
return self.block_aware_cache.get_stats()
|
||
return None
|
||
|
||
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.prefilling.clear()
|
||
self._prefill_states.clear()
|
||
self.running.clear()
|
||
self.requests.clear()
|
||
self.finished_req_ids.clear()
|
||
_unregister_uid_rows_for_model(self.model)
|
||
self.request_id_to_uid.clear()
|
||
self.uid_to_request_id.clear()
|
||
self._generation_overflow_recovery_ids.clear()
|
||
# Async store_cache bookkeeping. shutdown() drains these before us,
|
||
# but clear here too so reset() is safe to call standalone (e.g. tests
|
||
# or recovery paths) without leaking Request refs through stale futures.
|
||
self._pending_async_removes.clear()
|
||
self._inflight_store_futures.clear()
|
||
self._inflight_store_info.clear()
|
||
self._cache_freshness_waits.clear()
|
||
self._prefix_cache_prepared.clear()
|
||
self.batch_generator = None
|
||
self._current_sampler_params = None
|
||
self._boundary_cache_snapshots.clear()
|
||
if self._boundary_snapshot_store is not None:
|
||
self._boundary_snapshot_store.cleanup_all()
|
||
self._boundary_snapshot_required = None
|
||
|
||
# Clear caches
|
||
if self.block_aware_cache is not None:
|
||
self.block_aware_cache.clear()
|
||
self._cache_rate_tracker.clear()
|
||
|
||
# Clear detokenizers
|
||
self._request_detokenizers.clear()
|
||
|
||
# Clear protocol-specific output parser sessions
|
||
self._output_parser_sessions.clear()
|
||
|
||
# Cancel any pending deferred Metal cache clear
|
||
self._deferred_clear_at = None
|
||
|
||
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
|
||
|
||
# Release model and tokenizer references for GC
|
||
self.model = None
|
||
self.tokenizer = None
|
||
|
||
# Release all cache-related references for GC
|
||
if self._boundary_snapshot_store is not None:
|
||
try:
|
||
self._boundary_snapshot_store.shutdown()
|
||
except Exception as e:
|
||
logger.warning("Boundary snapshot store shutdown error: %s", e)
|
||
self._close_specprefill_draft_cache_manager()
|
||
self.paged_cache_manager = None
|
||
self._draft_prefix_cache = None
|
||
self._specprefill_draft_model = None
|
||
self.block_aware_cache = None
|
||
self.memory_monitor = None
|
||
self._boundary_snapshot_store = None
|
||
|
||
# Force garbage collection of any lingering cache objects
|
||
import gc
|
||
|
||
gc.collect()
|
||
|
||
logger.info("Deep reset completed - all caches cleared")
|
||
|
||
def shutdown(self) -> None:
|
||
"""
|
||
Graceful shutdown.
|
||
|
||
Flushes hot cache to SSD and closes the background writer.
|
||
paged SSD cache files are NOT cleared to allow reuse on reload.
|
||
"""
|
||
logger.info("Scheduler shutdown initiated...")
|
||
# The store-cache gate is a non-blocking counter (#1496), so there is
|
||
# no step-thread caller to wake here. Inflight futures are drained
|
||
# below before the executor is asked to shut down.
|
||
# Wait for any inflight async store_cache futures + drain pending
|
||
# batch_generator removes so the writer thread / underlying paged SSD
|
||
# cache see all blocks before close().
|
||
if self._store_cache_executor is not None:
|
||
try:
|
||
inflight = list(self._inflight_store_futures.values())
|
||
if inflight:
|
||
logger.info(
|
||
"Waiting for %d inflight async store_cache future(s)...",
|
||
len(inflight),
|
||
)
|
||
_done, not_done = concurrent.futures.wait(
|
||
inflight, timeout=FATAL_TEARDOWN_TIMEOUT_S
|
||
)
|
||
if not_done:
|
||
fatal_exit(
|
||
"Scheduler shutdown timed out after "
|
||
f"{FATAL_TEARDOWN_TIMEOUT_S:.0f}s waiting for "
|
||
f"{len(not_done)} async store_cache future(s)"
|
||
)
|
||
self._drain_pending_async_removes()
|
||
self._store_cache_executor.shutdown(wait=False)
|
||
# Final drain after the bounded wait. If all workers finished
|
||
# before the timeout, skipped entries are now drainable. If not,
|
||
# fatal_exit() above terminates the process instead of leaving
|
||
# a partially torn-down engine alive.
|
||
self._drain_pending_async_removes()
|
||
except Exception as e:
|
||
logger.warning(f"Async store_cache shutdown error: {e}")
|
||
self._store_cache_executor = None
|
||
self._store_cache_gate = None
|
||
self._inflight_store_futures.clear()
|
||
self._inflight_store_info.clear()
|
||
self._cache_freshness_waits.clear()
|
||
self._prefix_cache_prepared.clear()
|
||
if self._boundary_snapshot_store is not None:
|
||
try:
|
||
self._boundary_snapshot_store.cleanup_all()
|
||
except Exception as e:
|
||
logger.warning("Boundary snapshot store cleanup error: %s", e)
|
||
try:
|
||
self._boundary_snapshot_store.shutdown()
|
||
except Exception as e:
|
||
logger.warning("Boundary snapshot store shutdown error: %s", e)
|
||
self._boundary_snapshot_store = None
|
||
self._close_specprefill_draft_cache_manager()
|
||
self._draft_prefix_cache = None
|
||
self._specprefill_draft_model = None
|
||
if self.paged_ssd_cache_manager is not None:
|
||
self.paged_ssd_cache_manager.close()
|
||
self.paged_ssd_cache_manager = None
|
||
# Release whatever the per-path unregisters did not reach, so nothing
|
||
# survives this engine in the module-level row registry.
|
||
_unregister_uid_rows_for_model(self.model)
|
||
logger.info("Scheduler shutdown completed")
|
||
|
||
def adjust_store_cache_cap(self, pressure_level: str) -> None:
|
||
"""Resize the store-cache gate based on memory pressure (#1383).
|
||
|
||
Called from ProcessMemoryEnforcer on every poll. The cap walks one
|
||
step per poll toward its target so transient spikes don't oscillate
|
||
the cap. Bounded by [1, max_num_seqs]:
|
||
- ok pressure: grow cap back toward max_num_seqs.
|
||
- soft/hard pressure: shrink cap so KV cache backlog fits the system.
|
||
"""
|
||
gate = self._store_cache_gate
|
||
if gate is None:
|
||
return
|
||
current = gate.cap
|
||
if pressure_level == "ok":
|
||
new = min(self.config.max_num_seqs, current + 1)
|
||
else:
|
||
new = max(1, current - 1)
|
||
if new != current:
|
||
gate.set_cap(new)
|
||
logger.debug(
|
||
"store-cache queue cap: %d -> %d (pressure=%s)",
|
||
current,
|
||
new,
|
||
pressure_level,
|
||
)
|
||
|
||
# =========================================================================
|
||
# SSD Cache Methods
|
||
# =========================================================================
|
||
|
||
def _set_model_info_for_monitor(self) -> None:
|
||
"""Extract model info and set it on memory monitor for estimation."""
|
||
if self.memory_monitor is None:
|
||
return
|
||
|
||
try:
|
||
# Try to get model config
|
||
config = None
|
||
if hasattr(self.model, "config"):
|
||
config = self.model.config
|
||
elif hasattr(self.model, "args"):
|
||
config = self.model.args
|
||
|
||
if config is None:
|
||
logger.debug("Could not extract model config for memory estimation")
|
||
return
|
||
|
||
def _cfg_get(obj: Any, key: str, default: Any = None) -> Any:
|
||
if isinstance(obj, dict):
|
||
return obj.get(key, default)
|
||
return getattr(obj, key, default)
|
||
|
||
# VLM / multimodal configs (e.g. Qwen3.6-VL, Gemma-4) nest the
|
||
# language-model dimensions under a sub-config. Prefer
|
||
# ``text_config`` / ``language_config`` / ``llm_config`` when ANY
|
||
# of them exposes the LM layer count, even if the top-level config
|
||
# also has one — on some VLM packs (older Gemma-3, certain Llava /
|
||
# HF auto-wrappers) the top-level field refers to the *vision
|
||
# encoder*, not the LM, and accepting it silently miscalibrates
|
||
# the SDPA-peak estimate by a constant factor (a 40-layer LM
|
||
# wrapped in a 33-layer vision tower under-estimates by ~20 %).
|
||
# Probe both ``num_hidden_layers`` and the legacy ``n_layer`` alias
|
||
# so a GPT-style nested config is also picked up. Falls back to the
|
||
# top-level config only when no sub-config has either field.
|
||
for sub_attr in ("text_config", "language_config", "llm_config"):
|
||
sub = _cfg_get(config, sub_attr)
|
||
if sub is not None and (
|
||
_cfg_get(sub, "num_hidden_layers") or _cfg_get(sub, "n_layer")
|
||
):
|
||
config = sub
|
||
break
|
||
|
||
# Extract KV cache dimensions
|
||
num_layers = _cfg_get(config, "num_hidden_layers") or _cfg_get(
|
||
config, "n_layer"
|
||
)
|
||
num_kv_heads = (
|
||
_cfg_get(config, "num_key_value_heads")
|
||
or _cfg_get(config, "num_attention_heads")
|
||
or _cfg_get(config, "n_head")
|
||
)
|
||
head_dim = _cfg_get(config, "head_dim")
|
||
hidden_size = _cfg_get(config, "hidden_size") or _cfg_get(config, "n_embd")
|
||
|
||
# Calculate head_dim if not directly available
|
||
if head_dim is None and hidden_size and num_kv_heads:
|
||
num_heads = _cfg_get(config, "num_attention_heads") or num_kv_heads
|
||
head_dim = hidden_size // num_heads
|
||
|
||
# Determine base dtype size for uncompressed KV cache elements.
|
||
base_dtype_size: float = 2 # Default float16/bfloat16
|
||
if hasattr(self.model, "dtype"):
|
||
if self.model.dtype == mx.float32:
|
||
base_dtype_size = 4
|
||
elif self.model.dtype == mx.bfloat16:
|
||
base_dtype_size = 2
|
||
dtype_size = base_dtype_size
|
||
|
||
# Extract num_attention_heads (query heads) for SDPA peak estimation
|
||
num_attention_heads = (
|
||
_cfg_get(config, "num_attention_heads")
|
||
or _cfg_get(config, "n_head")
|
||
or num_kv_heads
|
||
)
|
||
|
||
# Count KVCache layers for hybrid models
|
||
cache_list_for_tq = None
|
||
actual_kv_cache_layers = None
|
||
num_kv_cache_layers = num_layers
|
||
if not hasattr(self.model, "make_cache"):
|
||
actual_kv_cache_layers = num_layers
|
||
else:
|
||
try:
|
||
cache_list = self.model.make_cache()
|
||
cache_list_for_tq = cache_list
|
||
from mlx_lm.models.cache import CacheList, KVCache
|
||
|
||
def _count_kv(c: Any) -> int:
|
||
if type(c) is KVCache:
|
||
return 1
|
||
if isinstance(c, CacheList):
|
||
return sum(_count_kv(inner) for inner in c.caches)
|
||
return 0
|
||
|
||
actual_kv_cache_layers = sum(_count_kv(c) for c in cache_list)
|
||
num_kv_cache_layers = actual_kv_cache_layers
|
||
if num_kv_cache_layers == 0:
|
||
num_kv_cache_layers = num_layers # fallback
|
||
except Exception:
|
||
pass
|
||
|
||
if (
|
||
self._turboquant_kv_bits is not None
|
||
and isinstance(head_dim, int)
|
||
and not isinstance(head_dim, bool)
|
||
and head_dim > 0
|
||
and isinstance(actual_kv_cache_layers, int)
|
||
and actual_kv_cache_layers > 0
|
||
and (
|
||
self._turboquant_eligible(cache_list_for_tq)
|
||
if cache_list_for_tq is not None
|
||
else not (
|
||
self._model_uses_mla() or self._model_uses_attention_sinks()
|
||
)
|
||
)
|
||
):
|
||
tq_dtype_size = float(self._turboquant_kv_bits) / 8.0 + (2.0 / head_dim)
|
||
if (
|
||
self._turboquant_skip_last
|
||
and not isinstance(actual_kv_cache_layers, bool)
|
||
and actual_kv_cache_layers > 1
|
||
):
|
||
dtype_size = (
|
||
(actual_kv_cache_layers - 1) * tq_dtype_size + base_dtype_size
|
||
) / actual_kv_cache_layers
|
||
else:
|
||
dtype_size = tq_dtype_size
|
||
|
||
kv_bytes_per_token = (
|
||
estimate_mla_kv_bytes_per_token(
|
||
config,
|
||
cache_list_for_tq,
|
||
base_dtype_size,
|
||
)
|
||
if estimate_mla_kv_bytes_per_token is not None
|
||
else None
|
||
)
|
||
|
||
# Truthiness alone isn't enough — MagicMock proxies leaking
|
||
# through the descent (test scaffolds that don't fully spec
|
||
# ``model.config``) are truthy but fail any later numeric
|
||
# comparison (``> 128`` etc.) deep inside MemoryMonitor.
|
||
# Insist on real positive integers before calling.
|
||
def _pos_int(v: Any) -> bool:
|
||
return isinstance(v, int) and not isinstance(v, bool) and v > 0
|
||
|
||
if _pos_int(num_layers) and _pos_int(num_kv_heads) and _pos_int(head_dim):
|
||
self.memory_monitor.set_model_info(
|
||
num_layers=num_layers,
|
||
num_kv_heads=num_kv_heads,
|
||
head_dim=head_dim,
|
||
dtype_size=dtype_size,
|
||
num_attention_heads=num_attention_heads,
|
||
num_kv_cache_layers=num_kv_cache_layers,
|
||
# SDPA scores are materialized at the compute/activation
|
||
# dtype, not the (possibly fractional TurboQuant) KV width.
|
||
compute_dtype_size=base_dtype_size,
|
||
kv_bytes_per_token=kv_bytes_per_token,
|
||
)
|
||
logger.debug(
|
||
f"Model info for memory estimation: "
|
||
f"layers={num_layers} ({num_kv_cache_layers} KVCache), "
|
||
f"kv_heads={num_kv_heads}, q_heads={num_attention_heads}, "
|
||
f"head_dim={head_dim}, dtype_size={dtype_size}"
|
||
)
|
||
else:
|
||
logger.debug(
|
||
f"Incomplete model info: layers={num_layers}, "
|
||
f"kv_heads={num_kv_heads}, head_dim={head_dim}"
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.debug(f"Failed to extract model info: {e}")
|
||
|
||
def _infer_live_layer_cache_types(
|
||
self,
|
||
) -> tuple[list[str], float | None] | None:
|
||
"""Infer the layer-cache signature that future SSD saves will use.
|
||
|
||
Returns ``(layer_cache_types, turboquant_kv_bits)`` — the predicted
|
||
per-layer type names plus the depth requests will quantize at (None
|
||
when TurboQuant is inactive or ineligible) — or None when no
|
||
signature can be inferred.
|
||
"""
|
||
if not HAS_CACHE_TYPE_HANDLERS or ModelCacheConfig is None:
|
||
return None
|
||
|
||
# Build the cache list the same way the request path does
|
||
# (make_prompt_cache defers to model.make_cache when the model
|
||
# defines one, and falls back to plain per-layer KVCache otherwise).
|
||
# Requiring model.make_cache here made the refresh a silent no-op
|
||
# for every plain dense model, so the manager never learned the
|
||
# TurboQuant layout or bit depth (#2045).
|
||
try:
|
||
cache_list = make_prompt_cache(self.model)
|
||
except Exception as e:
|
||
logger.debug("Failed to build cache list for SSD signature: %s", e)
|
||
return None
|
||
|
||
if not isinstance(cache_list, (list, tuple)) or not cache_list:
|
||
return None
|
||
|
||
cache_list = list(cache_list)
|
||
try:
|
||
model_cache_config = ModelCacheConfig.from_cache_list(
|
||
cache_list,
|
||
model_name=self.config.model_name or "",
|
||
)
|
||
layer_cache_types = model_cache_config.get_type_names()
|
||
except Exception as e:
|
||
logger.debug("Failed to infer SSD layer cache signature: %s", e)
|
||
return None
|
||
|
||
if not layer_cache_types:
|
||
return None
|
||
|
||
if self._turboquant_kv_bits is None:
|
||
return layer_cache_types, None
|
||
|
||
try:
|
||
eligible = self._turboquant_eligible(cache_list)
|
||
except Exception as e:
|
||
# Fail safe: committing an un-rewritten (plain-KV) layout here
|
||
# would make the stale-signature sweep evict every valid
|
||
# TurboQuant block, so refuse to infer rather than guess.
|
||
logger.debug("Failed to evaluate TurboQuant SSD signature: %s", e)
|
||
return None
|
||
if not eligible:
|
||
return layer_cache_types, None
|
||
|
||
kv_indices = [
|
||
i for i, c in enumerate(cache_list) if _is_turboquant_kv_family_cache(c)
|
||
]
|
||
skip_last = self._turboquant_skip_last and len(kv_indices) > 1
|
||
last_kv_idx = kv_indices[-1] if skip_last else -1
|
||
for idx in kv_indices:
|
||
if idx != last_kv_idx and idx < len(layer_cache_types):
|
||
layer_cache_types[idx] = "TurboQuantKVCache"
|
||
|
||
# The depth is keyed off the same eligibility gate the request path
|
||
# uses, not the rewritten names: models whose convertible caches sit
|
||
# inside CacheList layers report bare "CacheList" names at every
|
||
# depth, so the bits field is the only signature discriminator for
|
||
# them (#2045).
|
||
return layer_cache_types, float(self._turboquant_kv_bits)
|
||
|
||
def refresh_ssd_layer_signature(self) -> list[str] | None:
|
||
"""Set the SSD manager's live layer signature before prefix lookup."""
|
||
manager = self.paged_ssd_cache_manager
|
||
if manager is None:
|
||
return None
|
||
|
||
inferred = self._infer_live_layer_cache_types()
|
||
if inferred is None:
|
||
if self._turboquant_kv_bits is not None:
|
||
logger.warning(
|
||
"Could not infer the SSD cache layer signature; "
|
||
"TurboQuant bit-depth compatibility checks stay disabled "
|
||
"for this session (stale-depth blocks are not swept)."
|
||
)
|
||
return None
|
||
layer_cache_types, turboquant_kv_bits = inferred
|
||
|
||
try:
|
||
set_signature = getattr(manager, "set_expected_layer_signature", None)
|
||
if callable(set_signature):
|
||
set_signature(
|
||
layer_cache_types,
|
||
turboquant_kv_bits=turboquant_kv_bits,
|
||
)
|
||
else:
|
||
manager.adopt_layer_signature_if_unset(layer_cache_types)
|
||
manager.invalidate_stale_layer_signature()
|
||
except Exception as e:
|
||
logger.warning("Failed to refresh SSD layer cache signature: %s", e)
|
||
return None
|
||
|
||
return layer_cache_types
|
||
|
||
def _init_tiered_cache(self) -> bool:
|
||
"""Initialize paged SSD cache components if configured.
|
||
|
||
In paged SSD-only mode:
|
||
- All KV cache data is stored on paged SSD via PagedSSDCacheManager
|
||
- PagedCacheManager only stores block metadata (no GPU memory for cache data)
|
||
- BatchGenerator handles GPU memory for active inference
|
||
"""
|
||
if not HAS_TIERED_CACHE:
|
||
if self.config.paged_ssd_cache_dir:
|
||
logger.warning(
|
||
"paged SSD cache requested but ssd_cache/memory_monitor modules "
|
||
"not available. Install required dependencies."
|
||
)
|
||
return False
|
||
|
||
# In paged SSD-only mode, paged_ssd_cache_dir is required
|
||
if not self.config.paged_ssd_cache_dir:
|
||
logger.debug(
|
||
"paged SSD cache not configured (no --ssd-cache-dir specified)"
|
||
)
|
||
return False
|
||
|
||
try:
|
||
cache_dir = (
|
||
Path(self.config.paged_ssd_cache_dir)
|
||
if self.config.paged_ssd_cache_dir
|
||
else None
|
||
)
|
||
|
||
# Pass current model identity so stale blocks from a prior model
|
||
# version (e.g., 30-layer cache after an upgrade to 40 layers via
|
||
# #1404) are unlinked at startup instead of triggering a layer
|
||
# mismatch reject on every prefix lookup. See #1413.
|
||
expected_num_layers = (
|
||
self.block_aware_cache.expected_num_layers
|
||
if self.block_aware_cache is not None
|
||
else 0
|
||
)
|
||
|
||
# Pending-writes queue sizing depends on per-block bytes
|
||
# (block_size × per-token KV). Pass the *final* scheduler
|
||
# block size — possibly adjusted from the config default by
|
||
# RotatingKVCache / ArraysCache logic earlier in
|
||
# ``__init__`` — and a model-derived per-token KV estimate
|
||
# from the memory monitor.
|
||
#
|
||
# Gate on ``has_model_info()`` rather than just non-None:
|
||
# ``estimate_block_memory(1)`` silently substitutes a
|
||
# 7B-class fiction (32 layers × 8 KV heads × 128 head_dim
|
||
# ≈ 128 KB/token) when dims were never set, and feeding
|
||
# that "default" value into the writer-queue formula gives
|
||
# the wrong cap on real workloads. When dims are missing
|
||
# (test fixtures with skeletal model.config, unusual VLM
|
||
# packs the nested-config walk doesn't recognise), pass
|
||
# the PagedSSDCacheManager's 200 KB default explicitly so
|
||
# the cap math degrades to a known constant instead of a
|
||
# model-class fiction. The auto-init in ``Scheduler.__init__``
|
||
# paired with ``_set_model_info_for_monitor()`` means the
|
||
# happy path here is ``has_model_info() is True``; this
|
||
# else branch only fires for skeletal test fixtures.
|
||
if self.memory_monitor is not None and self.memory_monitor.has_model_info():
|
||
# ``estimate_block_memory(1)`` returns all-layers K+V
|
||
# bytes for a single token at the dtype the monitor was
|
||
# configured with — exactly the per-token cost the
|
||
# queue cap needs to weigh.
|
||
expected_kv_bytes_per_token = self.memory_monitor.estimate_block_memory(
|
||
1
|
||
)
|
||
else:
|
||
expected_kv_bytes_per_token = 200_000 # PagedSSDCacheManager default
|
||
|
||
# Initialize paged SSD cache manager for SSD storage
|
||
self.paged_ssd_cache_manager = PagedSSDCacheManager(
|
||
cache_dir=cache_dir,
|
||
max_size_bytes=self.config.paged_ssd_cache_max_size,
|
||
hot_cache_max_bytes=self.config.hot_cache_max_size,
|
||
hot_cache_only=self.config.hot_cache_only,
|
||
hot_cache_budget=self.config.hot_cache_budget,
|
||
expected_model_name=self.config.model_name or "",
|
||
expected_num_layers=expected_num_layers,
|
||
expected_block_size=self.config.paged_cache_block_size,
|
||
expected_block_size_tokens=self.config.paged_cache_block_size,
|
||
expected_kv_bytes_per_token=expected_kv_bytes_per_token,
|
||
)
|
||
|
||
# Connect paged SSD cache manager to PagedCacheManager
|
||
if self.paged_cache_manager is not None:
|
||
self.paged_cache_manager.set_paged_ssd_cache_manager(
|
||
self.paged_ssd_cache_manager
|
||
)
|
||
|
||
# Connect paged SSD cache manager to BlockAwarePrefixCache for paged SSD-only mode
|
||
if self.block_aware_cache is not None:
|
||
self.block_aware_cache.set_paged_ssd_cache_manager(
|
||
self.paged_ssd_cache_manager
|
||
)
|
||
|
||
# Initialize boundary snapshot SSD store for offloading
|
||
# non-sliceable cache snapshots during prefill.
|
||
# Skip in hot_cache_only mode since snapshots would never be written.
|
||
if BoundarySnapshotSSDStore is not None and not self.config.hot_cache_only:
|
||
try:
|
||
self._boundary_snapshot_store = BoundarySnapshotSSDStore(
|
||
base_dir=Path(self.config.paged_ssd_cache_dir)
|
||
)
|
||
except Exception as e:
|
||
logger.debug(
|
||
"Failed to initialize boundary snapshot SSD store: %s", e
|
||
)
|
||
|
||
if self.config.hot_cache_only:
|
||
logger.info(
|
||
f"hot-cache-only mode enabled: "
|
||
f"hot_cache_max={self._format_bytes(self.config.hot_cache_max_size)}, "
|
||
f"block_size={self.config.paged_cache_block_size} tokens"
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"paged SSD cache enabled: "
|
||
f"cache_dir={self.config.paged_ssd_cache_dir}, "
|
||
f"max_size={self._format_bytes(self.config.paged_ssd_cache_max_size)}, "
|
||
f"block_size={self.config.paged_cache_block_size} tokens"
|
||
)
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize paged SSD cache: {e}")
|
||
self.paged_ssd_cache_manager = None
|
||
return False
|
||
|
||
def _disable_paged_cache_components(self) -> None:
|
||
"""Clear paged-cache runtime state after SSD cache setup fails."""
|
||
if self.paged_ssd_cache_manager is not None:
|
||
try:
|
||
self.paged_ssd_cache_manager.close()
|
||
except Exception as e:
|
||
logger.debug("Failed to close paged SSD cache manager: %s", e)
|
||
self.paged_ssd_cache_manager = None
|
||
self.paged_cache_manager = None
|
||
self.block_aware_cache = None
|
||
self._boundary_snapshot_store = None
|
||
|
||
def _check_memory_pressure(self) -> None:
|
||
"""Check memory and evict blocks if needed.
|
||
|
||
In paged SSD-only mode, memory pressure is not monitored since
|
||
KV cache data is stored on paged SSD, not GPU memory.
|
||
"""
|
||
# In paged SSD-only mode, memory_monitor is not used
|
||
# All KV cache data is on paged SSD, so no GPU memory pressure from PagedCache
|
||
pass
|
||
|
||
def _evict_blocks_permanently(self, bytes_to_free: int) -> int:
|
||
"""
|
||
Evict LRU blocks permanently (metadata cleanup).
|
||
|
||
In paged SSD-only mode, blocks don't store data in GPU memory.
|
||
This method just removes block metadata to free up slots.
|
||
|
||
Args:
|
||
bytes_to_free: Target bytes to free (used for estimation).
|
||
|
||
Returns:
|
||
Number of bytes freed (estimated).
|
||
"""
|
||
if self.paged_cache_manager is None or self.memory_monitor is None:
|
||
return 0
|
||
|
||
# Estimate how many blocks to evict
|
||
block_size = self.config.paged_cache_block_size
|
||
num_blocks_to_evict = self.memory_monitor.estimate_blocks_to_free(
|
||
bytes_to_free, block_size
|
||
)
|
||
|
||
# Get evictable blocks in LRU order
|
||
evictable = self.paged_cache_manager.get_evictable_blocks(num_blocks_to_evict)
|
||
|
||
if not evictable:
|
||
logger.debug("No evictable blocks found for permanent eviction")
|
||
return 0
|
||
|
||
freed = 0
|
||
evicted_count = 0
|
||
|
||
for block in evictable:
|
||
# In paged SSD-only mode, just clear metadata (data is on paged SSD)
|
||
if self.paged_cache_manager.evict_block_permanently(block.block_id):
|
||
freed += self.memory_monitor.estimate_block_memory(block_size)
|
||
evicted_count += 1
|
||
|
||
if freed >= bytes_to_free:
|
||
break
|
||
|
||
if evicted_count > 0:
|
||
logger.info(
|
||
f"Evicted {evicted_count} blocks permanently "
|
||
f"(~{self._format_bytes(freed)} estimated)"
|
||
)
|
||
|
||
return freed
|
||
|
||
def _evict_blocks_to_cold(self, bytes_to_free: int) -> int:
|
||
"""
|
||
Evict LRU blocks (with paged SSD cache configured).
|
||
|
||
In paged SSD-only mode, data is already on paged SSD, so this just evicts
|
||
block metadata from the index. The data remains on paged SSD and can
|
||
be re-discovered if the same token sequence is requested.
|
||
|
||
Args:
|
||
bytes_to_free: Target bytes to free (used for estimation).
|
||
|
||
Returns:
|
||
Number of bytes freed (estimated).
|
||
"""
|
||
if self.paged_cache_manager is None or self.paged_ssd_cache_manager is None:
|
||
return 0
|
||
|
||
if self.memory_monitor is None:
|
||
return 0
|
||
|
||
# Estimate how many blocks to evict
|
||
block_size = self.config.paged_cache_block_size
|
||
num_blocks_to_evict = self.memory_monitor.estimate_blocks_to_free(
|
||
bytes_to_free, block_size
|
||
)
|
||
|
||
# Get evictable blocks in LRU order
|
||
evictable = self.paged_cache_manager.get_evictable_blocks(num_blocks_to_evict)
|
||
|
||
if not evictable:
|
||
logger.debug("No evictable blocks found")
|
||
return 0
|
||
|
||
evicted_count = 0
|
||
|
||
for block in evictable:
|
||
# In paged SSD-only mode, data is already on paged SSD
|
||
# Just evict the block metadata
|
||
if self.paged_cache_manager.evict_block_permanently(block.block_id):
|
||
evicted_count += 1
|
||
|
||
# Estimate bytes freed based on block count
|
||
estimated_freed = evicted_count * self.memory_monitor.estimate_block_memory(
|
||
block_size
|
||
)
|
||
|
||
if evicted_count > 0:
|
||
logger.info(
|
||
f"Evicted {evicted_count} blocks from index "
|
||
f"(data preserved on paged SSD, ~{self._format_bytes(estimated_freed)} metadata freed)"
|
||
)
|
||
|
||
return estimated_freed
|
||
|
||
def _restore_block_from_cold(self, block_id: int, block_hash: bytes) -> bool:
|
||
"""
|
||
Restore a block from cold storage (deprecated in paged SSD-only mode).
|
||
|
||
In paged SSD-only mode, blocks don't store cache_data. Data is loaded
|
||
directly from SSD when needed via reconstruct_cache().
|
||
|
||
Kept for API compatibility.
|
||
|
||
Args:
|
||
block_id: Block ID to restore.
|
||
block_hash: Block's content hash.
|
||
|
||
Returns:
|
||
True if block exists in cold storage.
|
||
"""
|
||
if self.paged_ssd_cache_manager is None or self.paged_cache_manager is None:
|
||
return False
|
||
|
||
# In paged SSD-only mode, just verify block exists on paged SSD
|
||
if not self.paged_ssd_cache_manager.has_block(block_hash):
|
||
logger.warning(f"Block {block_id} not found in cold storage")
|
||
return False
|
||
|
||
# Touch the block to update LRU
|
||
block = (
|
||
self.paged_cache_manager.blocks[block_id]
|
||
if block_id < len(self.paged_cache_manager.blocks)
|
||
else None
|
||
)
|
||
if block:
|
||
block.touch()
|
||
|
||
logger.debug(
|
||
f"Block {block_id} verified on paged SSD (hash={block_hash.hex()[:16]}...)"
|
||
)
|
||
return True
|
||
|
||
def restore_cold_blocks_for_request(self, request_id: str) -> int:
|
||
"""
|
||
Verify all blocks needed for a request exist on paged SSD.
|
||
|
||
In paged SSD-only mode, blocks don't store cache_data. This method
|
||
just verifies that blocks exist on paged SSD.
|
||
|
||
Args:
|
||
request_id: Request ID.
|
||
|
||
Returns:
|
||
Number of blocks verified on paged SSD.
|
||
"""
|
||
if self.paged_cache_manager is None or self.paged_ssd_cache_manager is None:
|
||
return 0
|
||
|
||
if self.block_aware_cache is None:
|
||
return 0
|
||
|
||
# Get block table for request
|
||
block_table = self.paged_cache_manager.request_tables.get(request_id)
|
||
if block_table is None:
|
||
return 0
|
||
|
||
verified = 0
|
||
for block_id in block_table.block_ids:
|
||
block = self.paged_cache_manager.blocks[block_id]
|
||
if block.block_hash is not None:
|
||
if self._restore_block_from_cold(block_id, block.block_hash):
|
||
verified += 1
|
||
|
||
return verified
|
||
|
||
def _collect_cache_counters(self) -> dict[str, int] | None:
|
||
if self.block_aware_cache is None:
|
||
return None
|
||
|
||
prefix_stats = self.block_aware_cache.get_stats()
|
||
counters = {
|
||
"prefix_hits": prefix_stats.hits,
|
||
"prefix_misses": prefix_stats.misses,
|
||
"prefix_tokens_matched": prefix_stats.tokens_matched_total,
|
||
"prefix_tokens_requested": prefix_stats.tokens_requested_total,
|
||
"prefix_tokens_saved": prefix_stats.tokens_saved,
|
||
"evictions": prefix_stats.evictions,
|
||
}
|
||
|
||
if self.paged_ssd_cache_manager is not None:
|
||
ssd = self.paged_ssd_cache_manager.get_stats()
|
||
hot_hits = ssd.hot_cache_hits
|
||
total_loads = ssd.loads
|
||
counters.update(
|
||
{
|
||
"ssd_hot_hits": hot_hits,
|
||
"ssd_disk_loads": max(0, total_loads - hot_hits),
|
||
"ssd_saves": ssd.saves,
|
||
"ssd_errors": ssd.errors,
|
||
"hot_cache_evictions": ssd.hot_cache_evictions,
|
||
"hot_cache_promotions": ssd.hot_cache_promotions,
|
||
}
|
||
)
|
||
|
||
return counters
|
||
|
||
def get_ssd_cache_stats(self) -> dict[str, Any] | None:
|
||
"""Get paged SSD + prefix cache observability statistics."""
|
||
stats = {}
|
||
|
||
if self.paged_ssd_cache_manager is not None:
|
||
stats["ssd_cache"] = self.paged_ssd_cache_manager.get_stats()
|
||
|
||
if self.paged_cache_manager is not None:
|
||
stats["indexed_blocks"] = self.paged_cache_manager.cold_block_count
|
||
stats["block_size"] = self.config.paged_cache_block_size
|
||
|
||
if self.block_aware_cache is not None:
|
||
stats["prefix_cache"] = self.block_aware_cache.get_stats_dict()
|
||
|
||
counters = self._collect_cache_counters()
|
||
if counters:
|
||
stats["cache_rates"] = self._cache_rate_tracker.snapshot_and_get_rates(
|
||
counters
|
||
)
|
||
|
||
return stats if stats else None
|
||
|
||
# Alias for backwards compatibility
|
||
get_tiered_cache_stats = get_ssd_cache_stats
|
||
|
||
@staticmethod
|
||
def _format_bytes(bytes_value: int) -> str:
|
||
"""Format bytes as human-readable string."""
|
||
if bytes_value >= 1024**3:
|
||
return f"{bytes_value / 1024**3:.2f} GB"
|
||
elif bytes_value >= 1024**2:
|
||
return f"{bytes_value / 1024**2:.2f} MB"
|
||
elif bytes_value >= 1024:
|
||
return f"{bytes_value / 1024:.2f} KB"
|
||
else:
|
||
return f"{bytes_value} B"
|