wehub-resource-sync
1602 lines
64 KiB
Python
1602 lines
64 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
DFlash engine for block diffusion speculative decoding.
|
|
|
|
This engine wraps dflash-mlx (>= 0.1.5) to provide 3-4x faster decoding on
|
|
Apple Silicon for Qwen and Gemma4 model families. By default it serves all
|
|
requests through dflash; setting ``model_settings.dflash_max_ctx`` opts into
|
|
evicting the dflash models and delegating long-context requests to omlx's
|
|
BatchedEngine/VLMBatchedEngine (paged cache, SSD cache, continuous batching).
|
|
"""
|
|
|
|
import asyncio
|
|
import copy
|
|
import gc
|
|
import json
|
|
import logging
|
|
import threading
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import mlx.core as mx
|
|
|
|
from ..adapter.output_parser import detect_output_parser
|
|
from ..api.tool_calling import convert_tools_for_template
|
|
from ..api.utils import clean_special_tokens, detect_and_strip_partial
|
|
from ..memory_monitor import (
|
|
MemoryMonitor,
|
|
raise_if_prefill_exceeds,
|
|
set_model_info_from_model,
|
|
)
|
|
from ..utils.generation_config import load_generation_config_token_ids
|
|
from ..utils.model_loading import maybe_apply_pre_load_patches
|
|
from ..utils.proc_memory import get_phys_footprint
|
|
from ..utils.tokenizer import create_streaming_detokenizer
|
|
from .base import BaseEngine, GenerationOutput, _warn_scheduler_unreachable_once
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def is_dflash_compatible(model_path: str | Path) -> tuple[bool, str]:
|
|
"""Decide whether ``model_path`` can run on the current dflash backend.
|
|
|
|
DFlash 0.1.5 registers QwenGdnTargetOps and Gemma4TargetOps. The
|
|
top-level ``model_type`` is the canonical discriminator: Gemma4 multimodal
|
|
configs use ``gemma4`` at the top, while MTP-only variants (e.g. the
|
|
Gemma4 ``-assistant`` checkpoint) declare ``gemma4_assistant`` even
|
|
though their nested ``text_config.model_type`` is still ``gemma4_text``.
|
|
Reading top-level only keeps the gate aligned with what dflash will
|
|
actually load.
|
|
|
|
Returns:
|
|
(is_compatible, reason). ``reason`` is empty when compatible.
|
|
"""
|
|
config_path = Path(model_path) / "config.json"
|
|
if not config_path.exists():
|
|
return False, f"config.json not found at {config_path}"
|
|
try:
|
|
with open(config_path, encoding="utf-8") as f:
|
|
cfg = json.load(f)
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
return False, f"failed to read config.json: {e}"
|
|
|
|
model_type = str(cfg.get("model_type") or "").lower()
|
|
|
|
is_qwen = "qwen" in model_type
|
|
is_gemma4 = model_type in ("gemma4", "gemma4_text")
|
|
if not (is_qwen or is_gemma4):
|
|
return False, (
|
|
f"DFlash supports only Qwen and Gemma4 models "
|
|
f"(model_type='{cfg.get('model_type', '')}')"
|
|
)
|
|
return True, ""
|
|
|
|
|
|
class _DFlashPrefillGuard:
|
|
"""Prefill-memory guard target for DFlash's primary (speculative) path,
|
|
which bypasses the Scheduler entirely.
|
|
|
|
Holds a ``MemoryMonitor`` (built from the target model's dims) plus the two
|
|
watermark attrs the ``ProcessMemoryEnforcer`` writes on schedulers each tick
|
|
(``_prefill_memory_guard``, ``_memory_hard_limit_bytes``). The enforcer
|
|
resolves this object via ``_resolve_scheduler`` so DFlash receives the same
|
|
ceiling as scheduler-driven engines, and ``preflight_or_raise`` delegates to
|
|
the shared ``raise_if_prefill_exceeds`` so the estimate + HTTP-400 mapping
|
|
match ``Scheduler.preflight_or_raise`` exactly.
|
|
"""
|
|
|
|
def __init__(self, memory_monitor: MemoryMonitor, prefill_step_size: int):
|
|
self.memory_monitor = memory_monitor
|
|
self._prefill_step_size = prefill_step_size
|
|
self._last_mlx_active_memory_bytes: int = 0
|
|
# Written by ProcessMemoryEnforcer._propagate_memory_limit each tick.
|
|
self._prefill_memory_guard: bool = False
|
|
self._memory_hard_limit_bytes: int = 0
|
|
|
|
def record_mlx_active_memory(self, active_bytes: int) -> None:
|
|
self._last_mlx_active_memory_bytes = max(0, int(active_bytes))
|
|
|
|
def _current_usage_bytes(self) -> int:
|
|
return max(self._last_mlx_active_memory_bytes, get_phys_footprint())
|
|
|
|
def preflight_or_raise(
|
|
self,
|
|
*,
|
|
num_prompt_tokens: int,
|
|
request_id: str | None = None,
|
|
) -> None:
|
|
# Deliberately no cached_tokens parameter: a DFlash prefix-cache hit
|
|
# reconstructs the matched KV into active memory, so the full prompt
|
|
# must always be charged (see DFlashEngine.preflight_chat).
|
|
raise_if_prefill_exceeds(
|
|
self.memory_monitor,
|
|
prefill_memory_guard=self._prefill_memory_guard,
|
|
hard_limit_bytes=self._memory_hard_limit_bytes,
|
|
current_usage_bytes=self._current_usage_bytes(),
|
|
prefill_step_size=self._prefill_step_size,
|
|
num_prompt_tokens=num_prompt_tokens,
|
|
request_id=request_id,
|
|
)
|
|
|
|
|
|
class DFlashEngine(BaseEngine):
|
|
"""
|
|
DFlash speculative decoding engine with optional batched fallback.
|
|
|
|
For prompts within ``model_settings.dflash_max_ctx`` (or always, when the
|
|
threshold is None), uses block diffusion speculative decoding for 3-4x
|
|
faster generation. When the threshold is exceeded, evicts dflash models
|
|
from memory and delegates to a fallback engine (BatchedEngine or
|
|
VLMBatchedEngine) that provides paged cache, SSD cache, and continuous
|
|
batching.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_name: str,
|
|
draft_model_path: str,
|
|
draft_quant_enabled: bool | None = None,
|
|
draft_quant_weight_bits: int | None = None,
|
|
draft_quant_activation_bits: int | None = None,
|
|
draft_quant_group_size: int | None = None,
|
|
model_settings: Any | None = None,
|
|
fallback_engine_type: str = "batched",
|
|
scheduler_config: Any | None = None,
|
|
omlx_ssd_cache_dir: str | Path | None = None,
|
|
):
|
|
self._model_name = model_name
|
|
self._draft_model_path = draft_model_path
|
|
self._draft_quant_enabled = draft_quant_enabled
|
|
self._draft_quant_weight_bits = draft_quant_weight_bits
|
|
self._draft_quant_activation_bits = draft_quant_activation_bits
|
|
self._draft_quant_group_size = draft_quant_group_size
|
|
self._model_settings = model_settings
|
|
self._fallback_engine_type = fallback_engine_type
|
|
# Snapshot the scheduler config now: the engine pool mutates the
|
|
# shared instance (model_name/model_path) on every model load, and
|
|
# the fallback engine may start long after this engine was built.
|
|
self._scheduler_config = (
|
|
copy.copy(scheduler_config) if scheduler_config else scheduler_config
|
|
)
|
|
self._omlx_ssd_cache_dir = (
|
|
Path(omlx_ssd_cache_dir) if omlx_ssd_cache_dir else None
|
|
)
|
|
|
|
self._target_model = None
|
|
self._target_ops = None
|
|
self._draft_model = None
|
|
self._draft_backend = None
|
|
self._tokenizer_obj = None
|
|
self._executor_tokenizer = None
|
|
self._loaded = False
|
|
self._active_request = False
|
|
self._model_type_str = None
|
|
self._fallback_engine: BaseEngine | None = None
|
|
self._in_fallback_mode = False
|
|
self._fallback_lock = asyncio.Lock()
|
|
# Primary-mode prefill memory guard. DFlash bypasses the scheduler, so
|
|
# it can't receive the enforcer's watermarks through one; this holder
|
|
# stands in (built in start(), resolved by the enforcer).
|
|
self._prefill_guard: _DFlashPrefillGuard | None = None
|
|
self._runtime_context: Any | None = None
|
|
self._dflash_prefix_cache: Any | None = None
|
|
self._suppress_token_ids: set[int] = set()
|
|
# Protocol-specific output parser factory (gemma4 / harmony).
|
|
# Detected once in start() after the target model is loaded; None means
|
|
# the streaming detokenizer is used as-is (qwen, llama, etc.).
|
|
self._output_parser_factory: Any | None = None
|
|
|
|
self._max_dflash_ctx = (
|
|
getattr(model_settings, "dflash_max_ctx", None) if model_settings else None
|
|
)
|
|
self._in_memory_cache_enabled = (
|
|
bool(getattr(model_settings, "dflash_in_memory_cache", True))
|
|
if model_settings
|
|
else True
|
|
)
|
|
self._in_memory_cache_max_entries = int(
|
|
getattr(model_settings, "dflash_in_memory_cache_max_entries", 4)
|
|
if model_settings
|
|
else 4
|
|
)
|
|
self._in_memory_cache_max_bytes = int(
|
|
getattr(model_settings, "dflash_in_memory_cache_max_bytes", 8 * 1024**3)
|
|
if model_settings
|
|
else 8 * 1024**3
|
|
)
|
|
self._ssd_cache_requested = (
|
|
bool(getattr(model_settings, "dflash_ssd_cache", False))
|
|
if model_settings
|
|
else False
|
|
)
|
|
self._ssd_cache_max_bytes = int(
|
|
getattr(model_settings, "dflash_ssd_cache_max_bytes", 20 * 1024**3)
|
|
if model_settings
|
|
else 20 * 1024**3
|
|
)
|
|
# None → let dflash-mlx pick its own default (window=1024, sink=64, verify="adaptive").
|
|
# `getattr` returns None for missing attrs so older settings files keep working.
|
|
self._draft_window_size = (
|
|
getattr(model_settings, "dflash_draft_window_size", None)
|
|
if model_settings
|
|
else None
|
|
)
|
|
self._draft_sink_size = (
|
|
getattr(model_settings, "dflash_draft_sink_size", None)
|
|
if model_settings
|
|
else None
|
|
)
|
|
self._verify_mode = (
|
|
getattr(model_settings, "dflash_verify_mode", None)
|
|
if model_settings
|
|
else None
|
|
)
|
|
|
|
@property
|
|
def model_name(self) -> str:
|
|
return self._model_name
|
|
|
|
@property
|
|
def tokenizer(self) -> Any:
|
|
return self._tokenizer_obj
|
|
|
|
@property
|
|
def model_type(self) -> str | None:
|
|
return self._model_type_str
|
|
|
|
@staticmethod
|
|
def _build_quant_spec(
|
|
weight_bits: int | None,
|
|
activation_bits: int | None,
|
|
group_size: int | None,
|
|
) -> str:
|
|
"""Convert draft quantization config into dflash 0.1.5's spec string format.
|
|
|
|
None values fall back to dflash defaults (w4a16:gs64), so the spec stays
|
|
valid when a profile or external API sets `enabled=True` without filling
|
|
in every bit value.
|
|
"""
|
|
wb = weight_bits if weight_bits is not None else 4
|
|
ab = activation_bits if activation_bits is not None else 16
|
|
gs = group_size if group_size is not None else 64
|
|
return f"w{wb}a{ab}:gs{gs}"
|
|
|
|
def _resolve_dflash_l2_dir(self) -> Path | None:
|
|
"""Compute the dflash L2 cache directory under the omlx SSD cache root."""
|
|
if not self._ssd_cache_requested:
|
|
return None
|
|
if self._omlx_ssd_cache_dir is None:
|
|
logger.warning(
|
|
"DFlash SSD cache requested but omlx paged SSD cache directory is "
|
|
"not configured; disabling L2."
|
|
)
|
|
return None
|
|
if not self._in_memory_cache_enabled:
|
|
logger.warning("DFlash SSD cache requires in-memory cache; disabling L2.")
|
|
return None
|
|
return self._omlx_ssd_cache_dir / "dflash_l2"
|
|
|
|
def _build_runtime_context(self) -> Any:
|
|
from dflash_mlx.runtime.config import runtime_config_from_defaults
|
|
from dflash_mlx.runtime.context import build_runtime_context
|
|
|
|
l2_dir = self._resolve_dflash_l2_dir()
|
|
l2_enabled = l2_dir is not None
|
|
cfg = runtime_config_from_defaults(
|
|
prefix_cache=self._in_memory_cache_enabled,
|
|
prefix_cache_max_entries=self._in_memory_cache_max_entries,
|
|
prefix_cache_max_bytes=self._in_memory_cache_max_bytes,
|
|
prefix_cache_l2=l2_enabled,
|
|
prefix_cache_l2_dir=str(l2_dir) if l2_dir else "",
|
|
# Per-model L2 disk budget. dflash-mlx's _evict_to_budget drops the
|
|
# oldest snapshots once dflash_l2/ exceeds this, so the directory
|
|
# stays bounded instead of filling the disk (issue #1326).
|
|
prefix_cache_l2_max_bytes=self._ssd_cache_max_bytes if l2_enabled else 0,
|
|
# None → dflash-mlx fills in DEFAULT_RUNTIME_CONFIG values.
|
|
draft_window_size=self._draft_window_size,
|
|
draft_sink_size=self._draft_sink_size,
|
|
verify_mode=self._verify_mode,
|
|
)
|
|
return build_runtime_context(cfg)
|
|
|
|
async def start(self) -> None:
|
|
if self._loaded:
|
|
return
|
|
|
|
from ..engine_core import get_mlx_executor
|
|
|
|
loop = asyncio.get_running_loop()
|
|
runtime_context = self._build_runtime_context()
|
|
|
|
def _load_models():
|
|
from dflash_mlx.draft_backend import EagerDraftBackend
|
|
from dflash_mlx.runtime.loading import (
|
|
load_draft_bundle,
|
|
load_target_bundle,
|
|
)
|
|
|
|
# Apply the same pre-load patches BatchedEngine uses before
|
|
# mlx_lm.load() runs. MTP-bearing targets (e.g. Qwen3.6 *-mtp)
|
|
# need the MTP-compat sanitize patch or stock mlx-lm double-shifts
|
|
# the already-converted norm and emits garbage. dflash and mtp are
|
|
# mutually exclusive per model_settings, so this never attaches an
|
|
# MTP head; it only fixes sanitize. See issue #1318.
|
|
maybe_apply_pre_load_patches(
|
|
self._model_name, model_settings=self._model_settings
|
|
)
|
|
|
|
# Wrap dflash's hook installers so we can revert the class-level
|
|
# __call__ patches when this engine stops. Without this, a later
|
|
# Native MTP load on the same process would see leftover dflash
|
|
# hooks and crash with TypeError on n_confirmed (issue #1388).
|
|
# Idempotent — only wraps once per process.
|
|
from ..patches.dflash_lifecycle import install_dflash_lifecycle_wrap
|
|
|
|
install_dflash_lifecycle_wrap()
|
|
|
|
target_bundle = load_target_bundle(
|
|
self._model_name,
|
|
quantize_kv_cache=bool(
|
|
getattr(runtime_context.runtime, "quantize_kv_cache", False)
|
|
),
|
|
verify_config=getattr(runtime_context, "verify", None),
|
|
)
|
|
draft, draft_meta = load_draft_bundle(
|
|
self._draft_model_path,
|
|
draft_quant=(
|
|
self._build_quant_spec(
|
|
self._draft_quant_weight_bits,
|
|
self._draft_quant_activation_bits,
|
|
self._draft_quant_group_size,
|
|
)
|
|
if self._draft_quant_enabled
|
|
else None
|
|
),
|
|
)
|
|
draft_backend = EagerDraftBackend()
|
|
return target_bundle, draft, draft_backend
|
|
|
|
result = await loop.run_in_executor(get_mlx_executor(), _load_models)
|
|
target_bundle, self._draft_model, self._draft_backend = result
|
|
self._runtime_context = runtime_context
|
|
self._target_model = target_bundle.model
|
|
self._tokenizer_obj = target_bundle.tokenizer
|
|
self._target_ops = target_bundle.target_ops
|
|
target_meta = target_bundle.meta
|
|
|
|
# Deep-copy tokenizer for executor-thread usage (dflash generation).
|
|
# The original self._tokenizer_obj stays for event-loop operations
|
|
# (encode, apply_chat_template, count_chat_tokens).
|
|
# See: https://github.com/huggingface/tokenizers/issues/537
|
|
self._executor_tokenizer = copy.deepcopy(self._tokenizer_obj)
|
|
|
|
# Extract model_type from config
|
|
config = target_meta.get("config", {})
|
|
if isinstance(config, dict):
|
|
self._model_type_str = config.get("model_type")
|
|
elif hasattr(config, "model_type"):
|
|
self._model_type_str = config.model_type
|
|
|
|
suppress_ref = (
|
|
getattr(self._executor_tokenizer, "name_or_path", None) or self._model_name
|
|
)
|
|
suppress_ids = load_generation_config_token_ids(suppress_ref, "suppress_tokens")
|
|
self._suppress_token_ids = suppress_ids or set()
|
|
if self._suppress_token_ids:
|
|
logger.info(
|
|
"DFlash loaded %d suppress token(s) from generation_config.json: %s",
|
|
len(self._suppress_token_ids),
|
|
self._suppress_token_ids,
|
|
)
|
|
|
|
# Detect protocol-specific output parser (gemma4 channel markers,
|
|
# harmony channels). Scheduler-driven engines apply this via
|
|
# OutputParserSession.process_token per request; dflash bypasses the
|
|
# scheduler so we do the same wiring inline in our two generate paths.
|
|
parser_config = config if isinstance(config, dict) else None
|
|
try:
|
|
self._output_parser_factory = detect_output_parser(
|
|
self._model_name, self._executor_tokenizer, parser_config
|
|
)
|
|
except Exception as exc:
|
|
logger.debug(f"output parser detect failed: {exc}")
|
|
self._output_parser_factory = None
|
|
|
|
# Build the primary-mode prefill memory guard from the target model's
|
|
# dims (same estimation Scheduler uses). Best-effort: on failure the
|
|
# guard stays None and preflight degrades to a no-op — the request runs
|
|
# unguarded rather than being falsely rejected.
|
|
self._prefill_guard = None
|
|
if self._target_model is not None:
|
|
try:
|
|
monitor = MemoryMonitor(
|
|
max_kv_cache_memory=None, eviction_enabled=False
|
|
)
|
|
set_model_info_from_model(monitor, self._target_model)
|
|
step = (
|
|
getattr(self._scheduler_config, "prefill_step_size", 2048) or 2048
|
|
)
|
|
self._prefill_guard = _DFlashPrefillGuard(monitor, step)
|
|
except Exception as exc:
|
|
# Warn (not debug): a missing guard silently disables the
|
|
# prefill OOM protection this engine relies on.
|
|
logger.warning(f"DFlash prefill guard init failed: {exc}")
|
|
self._prefill_guard = None
|
|
|
|
self._loaded = True
|
|
self._in_fallback_mode = False
|
|
max_ctx_display = (
|
|
"unlimited" if self._max_dflash_ctx is None else self._max_dflash_ctx
|
|
)
|
|
# Resolved values dflash-mlx actually ended up using (None settings → dflash default).
|
|
runtime_cfg = getattr(self._runtime_context, "runtime", None)
|
|
window_used = getattr(runtime_cfg, "draft_window_size", "?")
|
|
sink_used = getattr(runtime_cfg, "draft_sink_size", "?")
|
|
verify_used = getattr(runtime_cfg, "verify_mode", "?")
|
|
logger.info(
|
|
f"DFlashEngine loaded: target={self._model_name}, "
|
|
f"draft={self._draft_model_path}, "
|
|
f"max_ctx={max_ctx_display}, "
|
|
f"fallback={self._fallback_engine_type}, "
|
|
f"l1_cache={self._in_memory_cache_enabled}, "
|
|
f"l2_cache={self._resolve_dflash_l2_dir() is not None}, "
|
|
f"draft_window={window_used}, draft_sink={sink_used}, verify={verify_used}"
|
|
)
|
|
|
|
def _record_prefill_guard_active_memory(self) -> None:
|
|
guard = self._prefill_guard
|
|
if guard is None:
|
|
return
|
|
try:
|
|
guard.record_mlx_active_memory(mx.get_active_memory())
|
|
except Exception as exc:
|
|
logger.debug(f"DFlash active-memory sample failed: {exc}")
|
|
|
|
@staticmethod
|
|
def _begin_runtime_cache_request() -> Any | None:
|
|
"""Mark a DFlash cache request boundary when supported by dflash-mlx."""
|
|
try:
|
|
from dflash_mlx.cache.manager import current_runtime_cache_manager
|
|
except ImportError:
|
|
return None
|
|
try:
|
|
manager = current_runtime_cache_manager()
|
|
except Exception as exc:
|
|
logger.debug(f"current_runtime_cache_manager failed: {exc}")
|
|
return None
|
|
if manager is None:
|
|
return None
|
|
begin = getattr(manager, "begin_request", None)
|
|
if not callable(begin):
|
|
return None
|
|
try:
|
|
begin()
|
|
except Exception as exc:
|
|
logger.debug(f"dflash cache begin_request failed: {exc}")
|
|
return None
|
|
return manager
|
|
|
|
@staticmethod
|
|
def _end_runtime_cache_request(manager: Any | None) -> None:
|
|
if manager is None:
|
|
return
|
|
end = getattr(manager, "end_request", None)
|
|
if not callable(end):
|
|
return
|
|
try:
|
|
end()
|
|
except Exception as exc:
|
|
logger.debug(f"dflash cache end_request failed: {exc}")
|
|
|
|
async def _evict_dflash_and_start_fallback(self) -> None:
|
|
"""Evict dflash models from memory, verify release, then start fallback engine."""
|
|
from dflash_mlx.cache.manager import shutdown_runtime_cache_manager
|
|
|
|
from ..engine_core import get_mlx_executor
|
|
|
|
loop = asyncio.get_running_loop()
|
|
pre_active = mx.get_active_memory()
|
|
|
|
# Release dflash model and cache references
|
|
shutdown_runtime_cache_manager()
|
|
self._dflash_prefix_cache = None
|
|
self._runtime_context = None
|
|
self._target_model = None
|
|
self._target_ops = None
|
|
self._draft_model = None
|
|
self._draft_backend = None
|
|
self._executor_tokenizer = None
|
|
self._output_parser_factory = None
|
|
# Deliberately keep self._prefill_guard alive across the transition.
|
|
# Its MemoryMonitor holds only dims (no model ref), so it stays valid
|
|
# and cheap. Nulling it here would open a window — guard gone, fallback
|
|
# scheduler not yet started — where _resolve_scheduler returns None
|
|
# (spurious "scheduler unreachable" warning + an unguarded admission
|
|
# gap). The enforcer prefers the fallback scheduler once it is up, and
|
|
# fallback-mode preflight delegates to the fallback engine, so the
|
|
# stale guard is simply never consulted. Cleared in stop().
|
|
# The fallback engine (BatchedEngine / VLMBatchedEngine) starts next.
|
|
# Revert dflash's class patches now so the fallback's model loads
|
|
# onto clean linear_attn / self_attn classes (issue #1388).
|
|
try:
|
|
from ..patches.dflash_lifecycle import restore_dflash_class_patches
|
|
|
|
restore_dflash_class_patches()
|
|
except Exception as exc:
|
|
logger.debug(f"restore_dflash_class_patches (evict): {exc}")
|
|
|
|
# Force memory reclaim with settle barrier
|
|
gc.collect()
|
|
await loop.run_in_executor(
|
|
get_mlx_executor(),
|
|
lambda: (mx.synchronize(), mx.clear_cache()),
|
|
)
|
|
|
|
# Poll for actual memory release (same pattern as engine_pool._unload_engine)
|
|
for settle_round in range(10):
|
|
active_now = mx.get_active_memory()
|
|
freed = pre_active - active_now
|
|
if freed > 0:
|
|
logger.info(
|
|
f"DFlash models evicted: freed={freed / 1024**3:.2f}GB "
|
|
f"(round {settle_round + 1})"
|
|
)
|
|
break
|
|
await asyncio.sleep(0.5)
|
|
gc.collect()
|
|
await loop.run_in_executor(
|
|
get_mlx_executor(),
|
|
lambda: (mx.synchronize(), mx.clear_cache()),
|
|
)
|
|
else:
|
|
logger.warning("DFlash model eviction: memory settle timed out")
|
|
|
|
# Start fallback engine
|
|
if self._fallback_engine_type == "vlm":
|
|
from .vlm import VLMBatchedEngine
|
|
|
|
self._fallback_engine = VLMBatchedEngine(
|
|
model_name=self._model_name,
|
|
scheduler_config=self._scheduler_config,
|
|
model_settings=self._model_settings,
|
|
)
|
|
else:
|
|
from .batched import BatchedEngine
|
|
|
|
self._fallback_engine = BatchedEngine(
|
|
model_name=self._model_name,
|
|
scheduler_config=self._scheduler_config,
|
|
model_settings=self._model_settings,
|
|
)
|
|
await self._fallback_engine.start()
|
|
self._in_fallback_mode = True
|
|
logger.info(f"DFlash fallback engine started: {self._fallback_engine_type}")
|
|
|
|
async def stop(self) -> None:
|
|
from dflash_mlx.cache.manager import shutdown_runtime_cache_manager
|
|
|
|
if self._fallback_engine is not None:
|
|
await self._fallback_engine.stop()
|
|
self._fallback_engine = None
|
|
try:
|
|
shutdown_runtime_cache_manager()
|
|
except Exception as exc:
|
|
logger.debug(f"shutdown_runtime_cache_manager: {exc}")
|
|
self._dflash_prefix_cache = None
|
|
self._runtime_context = None
|
|
self._target_model = None
|
|
self._target_ops = None
|
|
self._draft_model = None
|
|
self._draft_backend = None
|
|
self._tokenizer_obj = None
|
|
self._executor_tokenizer = None
|
|
self._output_parser_factory = None
|
|
self._prefill_guard = None
|
|
self._in_fallback_mode = False
|
|
self._loaded = False
|
|
# Revert class-level __call__ patches dflash installed during start().
|
|
# Required so a subsequent Native MTP load on the same process sees
|
|
# clean classes instead of leftover dflash hooks (issue #1388).
|
|
try:
|
|
from ..patches.dflash_lifecycle import restore_dflash_class_patches
|
|
|
|
restore_dflash_class_patches()
|
|
except Exception as exc:
|
|
logger.debug(f"restore_dflash_class_patches: {exc}")
|
|
logger.info("DFlashEngine stopped")
|
|
|
|
def _apply_chat_template(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict] | None = None,
|
|
chat_template_kwargs: dict[str, Any] | None = None,
|
|
is_partial: bool | None = None,
|
|
) -> str:
|
|
"""Apply chat template to messages.
|
|
|
|
Args:
|
|
messages: List of chat messages
|
|
tools: Optional tool definitions
|
|
chat_template_kwargs: Optional kwargs for the chat template
|
|
(e.g. enable_thinking, reasoning_effort).
|
|
is_partial: Explicit partial-mode signal from the API server.
|
|
``True``/``False`` — server has already decided; the ``partial``
|
|
key is cleaned from message dicts but no detection is performed.
|
|
``None`` (default) — auto-detect from messages for backward
|
|
compatibility with direct engine callers.
|
|
"""
|
|
if hasattr(self._tokenizer_obj, "apply_chat_template"):
|
|
if is_partial is None:
|
|
is_partial = detect_and_strip_partial(messages)
|
|
else:
|
|
# Server already resolved partial; just clean residual keys
|
|
# so the chat template never sees the non-standard field.
|
|
for msg in messages:
|
|
msg.pop("partial", None)
|
|
template_kwargs = {
|
|
"tokenize": False,
|
|
"add_generation_prompt": not is_partial,
|
|
}
|
|
if is_partial:
|
|
template_kwargs["continue_final_message"] = True
|
|
if tools:
|
|
template_kwargs["tools"] = tools
|
|
if chat_template_kwargs:
|
|
template_kwargs.update(chat_template_kwargs)
|
|
try:
|
|
return self._tokenizer_obj.apply_chat_template(
|
|
messages, **template_kwargs
|
|
)
|
|
except TypeError:
|
|
if chat_template_kwargs:
|
|
for key in chat_template_kwargs:
|
|
template_kwargs.pop(key, None)
|
|
template_kwargs.pop("tools", None)
|
|
template_kwargs.pop("enable_thinking", None)
|
|
return self._tokenizer_obj.apply_chat_template(
|
|
messages, **template_kwargs
|
|
)
|
|
else:
|
|
prompt = "\n".join(f"{m['role']}: {m['content']}" for m in messages)
|
|
return prompt + "\nassistant:"
|
|
|
|
def count_chat_tokens(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict] | None = None,
|
|
chat_template_kwargs: dict[str, Any] | None = None,
|
|
is_partial: bool | None = None,
|
|
) -> int:
|
|
"""Count prompt tokens for chat messages after applying chat template.
|
|
|
|
Args:
|
|
messages: List of chat messages
|
|
tools: Optional tool definitions
|
|
chat_template_kwargs: Optional kwargs for chat template
|
|
is_partial: Explicit partial-mode signal (see _apply_chat_template).
|
|
|
|
Returns:
|
|
Number of prompt tokens
|
|
"""
|
|
template_tools = convert_tools_for_template(tools) if tools else None
|
|
prompt = self._apply_chat_template(
|
|
messages,
|
|
template_tools,
|
|
chat_template_kwargs=chat_template_kwargs,
|
|
is_partial=is_partial,
|
|
)
|
|
return len(self._tokenizer_obj.encode(prompt))
|
|
|
|
async def preflight_chat(
|
|
self,
|
|
messages: list,
|
|
tools: list | None = None,
|
|
request_id: str | None = None,
|
|
**kwargs,
|
|
) -> None:
|
|
"""Prefill-memory preflight for chat requests.
|
|
|
|
DFlash bypasses the scheduler, so it implements the front-door guard
|
|
itself (BaseEngine's no-op would leave primary-mode prefills
|
|
unprotected). In fallback mode it delegates to the fallback engine,
|
|
whose scheduler runs the full guard. Raises ``PrefillMemoryExceededError``
|
|
(→ HTTP 400) when the prompt's prefill peak would exceed the ceiling.
|
|
Mirrors ``BatchedEngine.preflight_chat``.
|
|
"""
|
|
if not self._loaded:
|
|
await self.start()
|
|
if self._in_fallback_mode and self._fallback_engine is not None:
|
|
await self._fallback_engine.preflight_chat(
|
|
messages, tools=tools, request_id=request_id, **kwargs
|
|
)
|
|
return
|
|
if self._prefill_guard is None:
|
|
_warn_scheduler_unreachable_once(
|
|
self, "preflight_chat", "primary-mode prefill guard unavailable"
|
|
)
|
|
return
|
|
try:
|
|
num_tokens = self.count_chat_tokens(
|
|
messages,
|
|
tools,
|
|
chat_template_kwargs=kwargs.get("chat_template_kwargs"),
|
|
is_partial=kwargs.get("is_partial"),
|
|
)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"DFlashEngine.preflight_chat: token count raised %s; skipping "
|
|
"prefill memory check, real chat path will surface the error",
|
|
type(e).__name__,
|
|
)
|
|
return
|
|
# Deliberately no cached_tokens: a DFlash prefix-cache hit
|
|
# *reconstructs* the matched KV into active memory (dflash_mlx
|
|
# ``hydrate_target_cache`` clones every array), so the full prompt's
|
|
# KV is allocated this request — unlike the scheduler's resident
|
|
# paged cache. Subtracting hit tokens here would under-count and
|
|
# defeat the OOM guard.
|
|
self._prefill_guard.preflight_or_raise(
|
|
num_prompt_tokens=num_tokens, request_id=request_id
|
|
)
|
|
|
|
async def preflight_completion(
|
|
self,
|
|
prompt: str,
|
|
request_id: str | None = None,
|
|
**kwargs,
|
|
) -> None:
|
|
"""Prefill-memory preflight for plain completions. See ``preflight_chat``."""
|
|
if not self._loaded:
|
|
await self.start()
|
|
if self._in_fallback_mode and self._fallback_engine is not None:
|
|
await self._fallback_engine.preflight_completion(
|
|
prompt, request_id=request_id, **kwargs
|
|
)
|
|
return
|
|
if self._prefill_guard is None:
|
|
_warn_scheduler_unreachable_once(
|
|
self,
|
|
"preflight_completion",
|
|
"primary-mode prefill guard unavailable",
|
|
)
|
|
return
|
|
try:
|
|
num_tokens = len(self._tokenizer_obj.encode(prompt))
|
|
except Exception as e:
|
|
logger.warning(
|
|
"DFlashEngine.preflight_completion: tokenizer.encode raised %s; "
|
|
"skipping prefill memory check, real completion path will "
|
|
"surface the error",
|
|
type(e).__name__,
|
|
)
|
|
return
|
|
# Deliberately no cached_tokens — see preflight_chat: a prefix-cache
|
|
# hit reconstructs KV into active memory, so the full prompt is
|
|
# charged.
|
|
self._prefill_guard.preflight_or_raise(
|
|
num_prompt_tokens=num_tokens, request_id=request_id
|
|
)
|
|
|
|
@property
|
|
def supports_multimodal_fallback(self) -> bool:
|
|
return self._fallback_engine_type == "vlm"
|
|
|
|
_MULTIMODAL_TYPES = frozenset({"image", "image_url", "input_image"})
|
|
|
|
@staticmethod
|
|
def _has_multimodal_content(messages: list[dict]) -> bool:
|
|
for msg in messages:
|
|
content = msg.get("content")
|
|
if isinstance(content, list):
|
|
for part in content:
|
|
if (
|
|
isinstance(part, dict)
|
|
and part.get("type") in DFlashEngine._MULTIMODAL_TYPES
|
|
):
|
|
return True
|
|
return False
|
|
|
|
def _should_fallback(self, prompt_tokens: list[int]) -> bool:
|
|
if self._max_dflash_ctx is None:
|
|
return False
|
|
return len(prompt_tokens) >= self._max_dflash_ctx
|
|
|
|
def _get_think_token_id(self, attr: str) -> int | None:
|
|
"""Safely read think_start_id / think_end_id from the tokenizer."""
|
|
try:
|
|
return getattr(self._tokenizer_obj, attr, None)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
|
|
def _detect_needs_think_prefix(self, prompt_tokens: list[int]) -> bool:
|
|
"""Detect if prompt ends with an open <think> tag (thinking enabled).
|
|
|
|
DFlash bypasses the scheduler, so the ``<think>\\n`` prefix that the
|
|
scheduler normally prepends to the first chunk for reasoning models
|
|
must be reproduced here. Mirrors ``Scheduler._detect_needs_think_prefix``.
|
|
|
|
Returns False for disabled-thinking patterns like <think></think>
|
|
where </think> immediately follows <think> in the prompt tail.
|
|
"""
|
|
if not prompt_tokens:
|
|
return False
|
|
|
|
think_start_id = self._get_think_token_id("think_start_id")
|
|
if think_start_id is None and self._tokenizer_obj is not None:
|
|
try:
|
|
tid = self._tokenizer_obj.convert_tokens_to_ids("<think>")
|
|
if tid == getattr(self._tokenizer_obj, "unk_token_id", None):
|
|
return False
|
|
think_start_id = tid
|
|
except (AttributeError, KeyError, TypeError):
|
|
return False
|
|
|
|
if not think_start_id:
|
|
return False
|
|
|
|
last_tokens = list(prompt_tokens[-3:])
|
|
if think_start_id not in last_tokens:
|
|
return False
|
|
|
|
last_idx = len(last_tokens) - 1 - last_tokens[::-1].index(think_start_id)
|
|
after_start = last_tokens[last_idx + 1 :]
|
|
|
|
if after_start:
|
|
think_end_id = self._get_think_token_id("think_end_id")
|
|
if think_end_id is not None and think_end_id in after_start:
|
|
return False
|
|
if self._tokenizer_obj is not None:
|
|
try:
|
|
tid = self._tokenizer_obj.convert_tokens_to_ids("</think>")
|
|
unk = getattr(self._tokenizer_obj, "unk_token_id", None)
|
|
if tid != unk and tid in after_start:
|
|
return False
|
|
except (AttributeError, KeyError, TypeError):
|
|
pass
|
|
|
|
return True
|
|
|
|
def _think_prefix_text(self) -> str:
|
|
"""Return the opening think tag string to prepend (e.g. '<think>\\n')."""
|
|
tag = getattr(self._tokenizer_obj, "think_start", "<think>")
|
|
return f"{tag}\n"
|
|
|
|
def _stream_dflash_events(
|
|
self,
|
|
prompt_tokens: list[int],
|
|
max_tokens: int,
|
|
):
|
|
"""Build the dflash event iterator with prefix cache plumbed in."""
|
|
from dflash_mlx.runtime import get_stop_token_ids, stream_dflash_generate
|
|
from dflash_mlx.server.prefix_cache_flow import PrefixCacheFlow
|
|
|
|
stop_ids = get_stop_token_ids(self._executor_tokenizer)
|
|
|
|
# Build a minimal model_provider shim for the prefix cache flow.
|
|
# ``model_key`` is consumed as a tuple where index 0 = target id and
|
|
# index 2 = draft id; the middle slot is unused on the dflash side.
|
|
# ``tokenizer`` and ``cli_args`` are required since dflash-mlx 1ba6713 —
|
|
# build_prefix_key hashes the chat template / policy. cli_args=None
|
|
# makes chat_template_args fall back to {}.
|
|
class _ModelProviderShim:
|
|
model_key = (self._model_name, None, self._draft_model_path)
|
|
model = self._target_model
|
|
target_ops = self._target_ops
|
|
tokenizer = self._executor_tokenizer
|
|
cli_args = None
|
|
|
|
prefix_flow = PrefixCacheFlow.for_request(
|
|
model_provider=_ModelProviderShim(),
|
|
draft_model=self._draft_model,
|
|
tokenizer=self._executor_tokenizer,
|
|
prompt=prompt_tokens,
|
|
max_new_tokens=max_tokens,
|
|
runtime_context=self._runtime_context,
|
|
)
|
|
|
|
event_iter = stream_dflash_generate(
|
|
target_model=self._target_model,
|
|
target_ops=self._target_ops,
|
|
tokenizer=self._executor_tokenizer,
|
|
draft_model=self._draft_model,
|
|
draft_backend=self._draft_backend,
|
|
prompt="",
|
|
max_new_tokens=max_tokens,
|
|
stop_token_ids=stop_ids,
|
|
suppress_token_ids=(
|
|
sorted(self._suppress_token_ids) if self._suppress_token_ids else None
|
|
),
|
|
prompt_tokens_override=prompt_tokens,
|
|
prefix_snapshot=prefix_flow.snapshot,
|
|
snapshot_service=prefix_flow.snapshot_service,
|
|
stable_prefix_len=prefix_flow.stable_prefix_len,
|
|
prefix_cache_active=prefix_flow.cache_active,
|
|
publish_generation_snapshot=prefix_flow.publish_generation_snapshot,
|
|
prefix_hit_kind=str(getattr(prefix_flow, "hit_kind", "miss") or "miss"),
|
|
runtime_context=self._runtime_context,
|
|
)
|
|
if hasattr(prefix_flow, "snapshot"):
|
|
prefix_flow.snapshot = None
|
|
return event_iter, prefix_flow, stop_ids
|
|
|
|
@staticmethod
|
|
def _cached_tokens_from_flow(prefix_flow) -> int:
|
|
"""Prompt tokens served from the DFlash prefix snapshot (prefill skipped).
|
|
|
|
On a hit ``PrefixCacheFlow.hit_tokens`` is the number of matched prompt
|
|
tokens; surfacing it as ``cached_tokens`` makes DFlash report prefix-cache
|
|
hits like the batched engine does instead of always reporting 0 (#1441).
|
|
"""
|
|
if prefix_flow is None:
|
|
return 0
|
|
return max(0, int(getattr(prefix_flow, "hit_tokens", 0) or 0))
|
|
|
|
def _run_generate_streaming(
|
|
self,
|
|
prompt_tokens: list[int],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
queue: asyncio.Queue,
|
|
loop: asyncio.AbstractEventLoop,
|
|
stop_event: threading.Event,
|
|
) -> None:
|
|
"""Run dflash generation with streaming on MLX executor thread.
|
|
|
|
``stop_event`` is set by the async consumer when it stops reading
|
|
(client disconnect / abort). Polling it between events lets the loop
|
|
return promptly so the single MLX executor thread is freed for the
|
|
next request.
|
|
"""
|
|
from dflash_mlx.engine.events import SummaryEvent, TokenEvent
|
|
|
|
event_iter = None
|
|
cache_manager = None
|
|
try:
|
|
self._record_prefill_guard_active_memory()
|
|
event_iter, prefix_flow, stop_ids = self._stream_dflash_events(
|
|
prompt_tokens=prompt_tokens,
|
|
max_tokens=max_tokens,
|
|
)
|
|
cache_manager = self._begin_runtime_cache_request()
|
|
self._record_prefill_guard_active_memory()
|
|
|
|
# Protocol-specific parser (gemma4 channel markers → <think> tags,
|
|
# harmony channels → <think>/visible split). When active it owns
|
|
# detokenization too, so the standard streaming detokenizer is
|
|
# only created when no parser is available.
|
|
parser_session = (
|
|
self._output_parser_factory.create_session(self._executor_tokenizer)
|
|
if self._output_parser_factory is not None
|
|
else None
|
|
)
|
|
detokenizer = None
|
|
if parser_session is None:
|
|
detokenizer = create_streaming_detokenizer(
|
|
self._executor_tokenizer,
|
|
model_path=self._model_name,
|
|
)
|
|
if detokenizer is not None:
|
|
detokenizer.reset()
|
|
|
|
for event in event_iter:
|
|
if stop_event.is_set():
|
|
logger.info("DFlash generation aborted by client")
|
|
break
|
|
|
|
if isinstance(event, TokenEvent):
|
|
token_id = int(event.token_id)
|
|
# Skip EOS/stop tokens from output
|
|
if token_id in stop_ids:
|
|
continue
|
|
if parser_session is not None:
|
|
result = parser_session.process_token(token_id)
|
|
text = result.stream_text
|
|
elif detokenizer is not None:
|
|
detokenizer.add_token(token_id)
|
|
text = detokenizer.last_segment
|
|
else:
|
|
text = self._executor_tokenizer.decode([token_id])
|
|
# Parser sessions can emit empty stream_text on protocol
|
|
# marker tokens — skip the chunk so clients don't see a
|
|
# flood of empty deltas.
|
|
if not text:
|
|
continue
|
|
asyncio.run_coroutine_threadsafe(
|
|
queue.put((text, [token_id], False, None)), loop
|
|
)
|
|
|
|
elif isinstance(event, SummaryEvent):
|
|
# Flush any buffered tail from the parser (e.g. close an
|
|
# unterminated <think> block) before the metrics chunk so
|
|
# the client sees a well-formed final delta.
|
|
if parser_session is not None:
|
|
final = parser_session.finalize()
|
|
tail = final.stream_text
|
|
if tail:
|
|
asyncio.run_coroutine_threadsafe(
|
|
queue.put((tail, [], False, None)), loop
|
|
)
|
|
|
|
gen_tokens = int(event.generation_tokens)
|
|
accept_ratio = float(event.acceptance_ratio)
|
|
cycles = int(event.cycles_completed)
|
|
elapsed_us = int(event.elapsed_us)
|
|
elapsed_s = elapsed_us / 1e6 if elapsed_us else 0
|
|
gen_tps = gen_tokens / elapsed_s if elapsed_s > 0 else 0
|
|
fallback = bool(event.fallback_ar)
|
|
logger.info(
|
|
f"DFlash generation complete: "
|
|
f"{gen_tokens} tokens, "
|
|
f"{gen_tps:.1f} tok/s, "
|
|
f"acceptance={accept_ratio:.1%}, "
|
|
f"cycles={cycles}"
|
|
f"{', fallback=AR' if fallback else ''}"
|
|
)
|
|
metrics = {
|
|
"prompt_tokens": int(event.prompt_token_count),
|
|
"completion_tokens": gen_tokens,
|
|
"acceptance_ratio": accept_ratio,
|
|
"cycles_completed": cycles,
|
|
# Prefix-snapshot hit count, surfaced on the final
|
|
# (usage) chunk so the API reports cached_tokens (#1441).
|
|
"cached_tokens": self._cached_tokens_from_flow(prefix_flow),
|
|
}
|
|
asyncio.run_coroutine_threadsafe(
|
|
queue.put(("", [], True, metrics)), loop
|
|
)
|
|
|
|
# Cycle, memory, prefill, and snapshot events are consumed by the
|
|
# runtime cache manager and metrics layers — omlx does not surface
|
|
# them so all other event types are intentionally ignored.
|
|
|
|
except Exception as e:
|
|
logger.error(f"DFlash streaming generation error: {e}")
|
|
asyncio.run_coroutine_threadsafe(
|
|
queue.put(("", [], True, {"error": str(e)})), loop
|
|
)
|
|
finally:
|
|
# Closing the dflash generator throws GeneratorExit on its next
|
|
# yield, releasing kernel state and any draft cache it holds.
|
|
self._record_prefill_guard_active_memory()
|
|
if event_iter is not None:
|
|
close = getattr(event_iter, "close", None)
|
|
if close is not None:
|
|
try:
|
|
close()
|
|
except Exception as exc:
|
|
logger.debug(f"event_iter.close() raised: {exc}")
|
|
self._end_runtime_cache_request(cache_manager)
|
|
# Always send a sentinel so the async consumer doesn't deadlock
|
|
# when an abort happened before the dflash summary was emitted.
|
|
asyncio.run_coroutine_threadsafe(
|
|
queue.put(("", [], True, {"aborted": stop_event.is_set()})),
|
|
loop,
|
|
)
|
|
self._active_request = False
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 0,
|
|
min_p: float = 0.0,
|
|
repetition_penalty: float = 1.0,
|
|
presence_penalty: float = 0.0,
|
|
stop: list[str] | None = None,
|
|
**kwargs,
|
|
) -> GenerationOutput:
|
|
if not self._loaded:
|
|
await self.start()
|
|
|
|
prompt_tokens = self._tokenizer_obj.encode(prompt)
|
|
|
|
# Fallback: evict dflash models, start LLM/VLM engine
|
|
if self._should_fallback(prompt_tokens):
|
|
async with self._fallback_lock:
|
|
if not self._in_fallback_mode:
|
|
logger.info(
|
|
f"DFlash context fallback: {len(prompt_tokens)} >= {self._max_dflash_ctx}, "
|
|
f"evicting dflash models and switching to {self._fallback_engine_type} engine"
|
|
)
|
|
await self._evict_dflash_and_start_fallback()
|
|
return await self._fallback_engine.generate(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
stop=stop,
|
|
**kwargs,
|
|
)
|
|
|
|
# Already in fallback mode but short context came in.
|
|
# Stay in fallback mode (reloading dflash models is expensive).
|
|
if self._in_fallback_mode:
|
|
return await self._fallback_engine.generate(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
stop=stop,
|
|
**kwargs,
|
|
)
|
|
|
|
from ..engine_core import get_mlx_executor
|
|
|
|
loop = asyncio.get_running_loop()
|
|
stop_event = threading.Event()
|
|
|
|
def _run():
|
|
from dflash_mlx.engine.events import SummaryEvent, TokenEvent
|
|
|
|
event_iter = None
|
|
cache_manager = None
|
|
# Per-request parser session (gemma4 channel markers, harmony
|
|
# channels). Lives only inside the executor thread so the parser
|
|
# state cannot leak across requests.
|
|
parser_session = (
|
|
self._output_parser_factory.create_session(self._executor_tokenizer)
|
|
if self._output_parser_factory is not None
|
|
else None
|
|
)
|
|
try:
|
|
self._record_prefill_guard_active_memory()
|
|
event_iter, prefix_flow, stop_ids = self._stream_dflash_events(
|
|
prompt_tokens=prompt_tokens,
|
|
max_tokens=max_tokens,
|
|
)
|
|
cache_manager = self._begin_runtime_cache_request()
|
|
self._record_prefill_guard_active_memory()
|
|
tokens: list[int] = []
|
|
parsed_visible_parts: list[str] = []
|
|
summary: SummaryEvent | None = None
|
|
for event in event_iter:
|
|
if stop_event.is_set():
|
|
logger.info("DFlash generation aborted by client")
|
|
break
|
|
if isinstance(event, TokenEvent):
|
|
token_id = int(event.token_id)
|
|
if token_id in stop_ids:
|
|
continue
|
|
tokens.append(token_id)
|
|
if parser_session is not None:
|
|
result = parser_session.process_token(token_id)
|
|
if result.visible_text:
|
|
parsed_visible_parts.append(result.visible_text)
|
|
elif isinstance(event, SummaryEvent):
|
|
summary = event
|
|
if parser_session is not None:
|
|
final = parser_session.finalize()
|
|
if final.visible_text:
|
|
parsed_visible_parts.append(final.visible_text)
|
|
return (
|
|
summary,
|
|
tokens,
|
|
parser_session,
|
|
parsed_visible_parts,
|
|
prefix_flow,
|
|
)
|
|
finally:
|
|
self._record_prefill_guard_active_memory()
|
|
if event_iter is not None:
|
|
close = getattr(event_iter, "close", None)
|
|
if close is not None:
|
|
try:
|
|
close()
|
|
except Exception as exc:
|
|
logger.debug(f"event_iter.close() raised: {exc}")
|
|
self._end_runtime_cache_request(cache_manager)
|
|
self._active_request = False
|
|
|
|
self._active_request = True
|
|
future = loop.run_in_executor(get_mlx_executor(), _run)
|
|
try:
|
|
summary, generated, parser_session, parsed_visible_parts, prefix_flow = (
|
|
await asyncio.shield(asyncio.wrap_future(future))
|
|
)
|
|
except asyncio.CancelledError:
|
|
stop_event.set()
|
|
logger.info("DFlash generate cancelled, waiting for executor to drain")
|
|
try:
|
|
await asyncio.wait_for(asyncio.wrap_future(future), timeout=10.0)
|
|
except TimeoutError:
|
|
logger.warning("DFlash executor did not exit within 10s after abort")
|
|
except Exception:
|
|
pass
|
|
raise
|
|
|
|
if parser_session is not None:
|
|
# Parser already converted protocol markers to <think>...</think>
|
|
# and stripped channel marker tokens, so just join the visible
|
|
# segments. Don't re-decode the raw token list — that would
|
|
# reintroduce the raw markers and double-buffer detokenization.
|
|
text = "".join(parsed_visible_parts)
|
|
else:
|
|
text = self._tokenizer_obj.decode(generated, skip_special_tokens=True)
|
|
text = clean_special_tokens(text)
|
|
|
|
# Reasoning models (Qwen3.x with enable_thinking, DeepSeek,
|
|
# MiniMax, ...) have <think>\n at the END of the prompt, so the
|
|
# model's first generated token is already INSIDE the thinking
|
|
# block. The opening tag never appears in the output, which would
|
|
# prevent extract_thinking / ThinkingParser from separating
|
|
# reasoning from content. Prepend the tag here so the API layer
|
|
# can split them correctly. Skipped when a parser session is
|
|
# active because gemma4/harmony parsers already emit <think> tags
|
|
# themselves and prepending would double the marker.
|
|
if self._detect_needs_think_prefix(prompt_tokens):
|
|
text = self._think_prefix_text() + text
|
|
|
|
prompt_token_count = (
|
|
int(summary.prompt_token_count)
|
|
if summary is not None
|
|
else len(prompt_tokens)
|
|
)
|
|
completion_token_count = (
|
|
int(summary.generation_tokens) if summary is not None else len(generated)
|
|
)
|
|
return GenerationOutput(
|
|
text=text,
|
|
tokens=generated,
|
|
prompt_tokens=prompt_token_count,
|
|
completion_tokens=completion_token_count,
|
|
cached_tokens=self._cached_tokens_from_flow(prefix_flow),
|
|
finish_reason="stop",
|
|
)
|
|
|
|
async def stream_generate(
|
|
self,
|
|
prompt: str,
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 0,
|
|
min_p: float = 0.0,
|
|
repetition_penalty: float = 1.0,
|
|
presence_penalty: float = 0.0,
|
|
stop: list[str] | None = None,
|
|
**kwargs,
|
|
) -> AsyncIterator[GenerationOutput]:
|
|
if not self._loaded:
|
|
await self.start()
|
|
|
|
prompt_tokens = self._tokenizer_obj.encode(prompt)
|
|
|
|
# Fallback: evict dflash models, start LLM/VLM engine
|
|
if self._should_fallback(prompt_tokens):
|
|
async with self._fallback_lock:
|
|
if not self._in_fallback_mode:
|
|
logger.info(
|
|
f"DFlash context fallback: {len(prompt_tokens)} >= {self._max_dflash_ctx}, "
|
|
f"evicting dflash models and switching to {self._fallback_engine_type} engine"
|
|
)
|
|
await self._evict_dflash_and_start_fallback()
|
|
async for output in self._fallback_engine.stream_generate(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
stop=stop,
|
|
**kwargs,
|
|
):
|
|
yield output
|
|
return
|
|
|
|
# Already in fallback mode — stay there
|
|
if self._in_fallback_mode:
|
|
async for output in self._fallback_engine.stream_generate(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
stop=stop,
|
|
**kwargs,
|
|
):
|
|
yield output
|
|
return
|
|
|
|
prompt_len = len(prompt_tokens)
|
|
loop = asyncio.get_running_loop()
|
|
queue: asyncio.Queue = asyncio.Queue()
|
|
stop_event = threading.Event()
|
|
|
|
# Reasoning models put <think>\n at the end of the prompt, so dflash
|
|
# generates tokens already inside the thinking block. The streaming
|
|
# ThinkingParser starts in _in_thinking=False, so without prepending
|
|
# the opening tag on the first chunk the whole reasoning block leaks
|
|
# into content. Mirror Scheduler._detect_needs_think_prefix here.
|
|
# When a protocol-aware parser session is active (gemma4 / harmony),
|
|
# the parser emits <think> tags itself, so prepending here would
|
|
# double the opening marker — gate it on factory absence.
|
|
needs_think_prefix = (
|
|
self._output_parser_factory is None
|
|
and self._detect_needs_think_prefix(prompt_tokens)
|
|
)
|
|
think_prefix_pending = needs_think_prefix
|
|
|
|
from ..engine_core import get_mlx_executor
|
|
|
|
self._active_request = True
|
|
future = loop.run_in_executor(
|
|
get_mlx_executor(),
|
|
self._run_generate_streaming,
|
|
prompt_tokens,
|
|
max_tokens,
|
|
temperature,
|
|
queue,
|
|
loop,
|
|
stop_event,
|
|
)
|
|
|
|
total_text = ""
|
|
total_completion = 0
|
|
finished_normally = False
|
|
|
|
try:
|
|
while True:
|
|
new_text, new_tokens, finished, metrics = await queue.get()
|
|
|
|
if think_prefix_pending and new_text:
|
|
new_text = self._think_prefix_text() + new_text
|
|
think_prefix_pending = False
|
|
|
|
total_text += new_text
|
|
total_completion += len(new_tokens)
|
|
|
|
finish_reason = None
|
|
if finished:
|
|
finish_reason = "stop"
|
|
if metrics and metrics.get("error"):
|
|
finish_reason = "error"
|
|
finished_normally = True
|
|
|
|
yield GenerationOutput(
|
|
text=total_text,
|
|
new_text=new_text,
|
|
tokens=new_tokens,
|
|
prompt_tokens=prompt_len,
|
|
completion_tokens=total_completion,
|
|
# Carried only on the final (usage) chunk's metrics; 0 on
|
|
# token deltas so the server's per-chunk sum isn't inflated.
|
|
cached_tokens=(metrics or {}).get("cached_tokens", 0),
|
|
finished=finished,
|
|
finish_reason=finish_reason,
|
|
)
|
|
|
|
if finished:
|
|
break
|
|
finally:
|
|
# Signal the executor to stop so the next request isn't blocked
|
|
# behind a cancelled generation. Wait briefly for the dflash loop
|
|
# to break out at its next event boundary; the timeout caps how
|
|
# long the next request has to wait if the model is mid-cycle.
|
|
if not finished_normally:
|
|
stop_event.set()
|
|
logger.info("DFlash stream cancelled, waiting for executor to drain")
|
|
try:
|
|
await asyncio.wait_for(asyncio.wrap_future(future), timeout=10.0)
|
|
except TimeoutError:
|
|
logger.warning(
|
|
"DFlash executor did not exit within 10s after abort; "
|
|
"next request may still be queued"
|
|
)
|
|
except Exception as exc:
|
|
logger.debug(f"DFlash executor future raised: {exc}")
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 0,
|
|
min_p: float = 0.0,
|
|
repetition_penalty: float = 1.0,
|
|
presence_penalty: float = 0.0,
|
|
tools: list[dict] | None = None,
|
|
**kwargs,
|
|
) -> GenerationOutput:
|
|
if not self._loaded:
|
|
await self.start()
|
|
|
|
if self._in_fallback_mode:
|
|
return await self._fallback_engine.chat(
|
|
messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
tools=tools,
|
|
**kwargs,
|
|
)
|
|
|
|
if self._fallback_engine_type == "vlm" and self._has_multimodal_content(
|
|
messages
|
|
):
|
|
async with self._fallback_lock:
|
|
if not self._in_fallback_mode:
|
|
logger.info(
|
|
"DFlash multimodal fallback: image content detected, "
|
|
"switching to VLM engine"
|
|
)
|
|
await self._evict_dflash_and_start_fallback()
|
|
return await self._fallback_engine.chat(
|
|
messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
tools=tools,
|
|
**kwargs,
|
|
)
|
|
|
|
template_tools = convert_tools_for_template(tools) if tools else None
|
|
ct_kwargs = kwargs.pop("chat_template_kwargs", None)
|
|
is_partial = kwargs.pop("is_partial", None)
|
|
prompt = self._apply_chat_template(
|
|
messages,
|
|
template_tools,
|
|
chat_template_kwargs=ct_kwargs,
|
|
is_partial=is_partial,
|
|
)
|
|
|
|
return await self.generate(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
**kwargs,
|
|
)
|
|
|
|
async def stream_chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 0,
|
|
min_p: float = 0.0,
|
|
repetition_penalty: float = 1.0,
|
|
presence_penalty: float = 0.0,
|
|
tools: list[dict] | None = None,
|
|
**kwargs,
|
|
) -> AsyncIterator[GenerationOutput]:
|
|
if not self._loaded:
|
|
await self.start()
|
|
|
|
if self._in_fallback_mode:
|
|
async for output in self._fallback_engine.stream_chat(
|
|
messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
tools=tools,
|
|
**kwargs,
|
|
):
|
|
yield output
|
|
return
|
|
|
|
if self._fallback_engine_type == "vlm" and self._has_multimodal_content(
|
|
messages
|
|
):
|
|
async with self._fallback_lock:
|
|
if not self._in_fallback_mode:
|
|
logger.info(
|
|
"DFlash multimodal fallback: image content detected, "
|
|
"switching to VLM engine"
|
|
)
|
|
await self._evict_dflash_and_start_fallback()
|
|
async for output in self._fallback_engine.stream_chat(
|
|
messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
tools=tools,
|
|
**kwargs,
|
|
):
|
|
yield output
|
|
return
|
|
|
|
template_tools = convert_tools_for_template(tools) if tools else None
|
|
ct_kwargs = kwargs.pop("chat_template_kwargs", None)
|
|
is_partial = kwargs.pop("is_partial", None)
|
|
prompt = self._apply_chat_template(
|
|
messages,
|
|
template_tools,
|
|
chat_template_kwargs=ct_kwargs,
|
|
is_partial=is_partial,
|
|
)
|
|
|
|
async for output in self.stream_generate(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
min_p=min_p,
|
|
repetition_penalty=repetition_penalty,
|
|
presence_penalty=presence_penalty,
|
|
**kwargs,
|
|
):
|
|
yield output
|
|
|
|
@property
|
|
def scheduler(self) -> Any | None:
|
|
fallback = self._fallback_engine
|
|
if fallback is None:
|
|
return None
|
|
|
|
scheduler = getattr(fallback, "scheduler", None)
|
|
if scheduler is not None:
|
|
return scheduler
|
|
|
|
inner = getattr(fallback, "_engine", None)
|
|
if inner is None:
|
|
return None
|
|
inner_engine = getattr(inner, "engine", None)
|
|
if inner_engine is None:
|
|
return None
|
|
return getattr(inner_engine, "scheduler", None)
|
|
|
|
def has_active_requests(self) -> bool:
|
|
if (
|
|
self._fallback_engine is not None
|
|
and self._fallback_engine.has_active_requests()
|
|
):
|
|
return True
|
|
return self._active_request
|
|
|
|
def get_stats(self) -> dict[str, Any]:
|
|
return {
|
|
"engine_type": "dflash",
|
|
"model_name": self._model_name,
|
|
"draft_model": self._draft_model_path,
|
|
"max_dflash_ctx": self._max_dflash_ctx,
|
|
"fallback_engine_type": self._fallback_engine_type,
|
|
"in_fallback_mode": self._in_fallback_mode,
|
|
"loaded": self._loaded,
|
|
"in_memory_cache": self._in_memory_cache_enabled,
|
|
"ssd_cache": self._resolve_dflash_l2_dir() is not None,
|
|
}
|
|
|
|
def get_cache_stats(self) -> dict[str, Any] | None:
|
|
if self._fallback_engine is not None:
|
|
return self._fallback_engine.get_cache_stats()
|
|
return None
|