wehub-resource-sync
306 lines
11 KiB
Python
306 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
|
|
"""
|
|
Request management for oMLX continuous batching.
|
|
|
|
This module provides Request and RequestStatus classes adapted from vLLM's
|
|
request management system, simplified for MLX backend.
|
|
"""
|
|
|
|
import enum
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
|
|
|
if TYPE_CHECKING:
|
|
from .cache.paged_cache import BlockTable
|
|
|
|
|
|
class RequestStatus(enum.IntEnum):
|
|
"""Status of a request in the scheduling system."""
|
|
|
|
# Request is waiting to be scheduled
|
|
WAITING = enum.auto()
|
|
# Request is currently being processed (generating tokens)
|
|
RUNNING = enum.auto()
|
|
# Request was preempted and needs to be resumed
|
|
PREEMPTED = enum.auto()
|
|
# Request finished successfully (hit stop token)
|
|
FINISHED_STOPPED = enum.auto()
|
|
# Request finished due to max_tokens limit
|
|
FINISHED_LENGTH_CAPPED = enum.auto()
|
|
# Request was aborted by user
|
|
FINISHED_ABORTED = enum.auto()
|
|
|
|
@staticmethod
|
|
def is_finished(status: "RequestStatus") -> bool:
|
|
"""Check if the status indicates a finished request."""
|
|
return status > RequestStatus.PREEMPTED
|
|
|
|
@staticmethod
|
|
def get_finish_reason(status: "RequestStatus") -> Optional[str]:
|
|
"""Get the finish reason string for a finished status."""
|
|
if status == RequestStatus.FINISHED_STOPPED:
|
|
return "stop"
|
|
elif status == RequestStatus.FINISHED_LENGTH_CAPPED:
|
|
return "length"
|
|
elif status == RequestStatus.FINISHED_ABORTED:
|
|
return "abort"
|
|
return None
|
|
|
|
|
|
@dataclass
|
|
class SamplingParams:
|
|
"""Sampling parameters for text generation."""
|
|
|
|
max_tokens: int = 256
|
|
temperature: float = 0.7
|
|
top_p: float = 0.9
|
|
top_k: int = 0 # 0 means disabled
|
|
min_p: float = 0.0
|
|
xtc_probability: float = 0.0
|
|
xtc_threshold: float = 0.1
|
|
repetition_penalty: float = 1.0
|
|
presence_penalty: float = 0.0
|
|
frequency_penalty: float = 0.0
|
|
stop: Optional[List[str]] = None
|
|
stop_token_ids: Optional[List[int]] = None
|
|
|
|
# Logprobs settings (memory optimization: disabled by default)
|
|
logprobs: bool = False # Whether to return logprobs
|
|
top_logprobs: Optional[int] = None # Number of top logprobs (1-20)
|
|
|
|
# Thinking budget (None = unlimited thinking)
|
|
thinking_budget: Optional[int] = None
|
|
|
|
# Compiled grammar for constrained decoding (xgrammar CompiledGrammar).
|
|
# Typed as Any to avoid a hard dependency on xgrammar at import time.
|
|
compiled_grammar: Any = None
|
|
|
|
# Seed for reproducible generation (best-effort, per OpenAI spec)
|
|
seed: Optional[int] = None
|
|
|
|
def __post_init__(self):
|
|
if self.stop is None:
|
|
self.stop = []
|
|
if self.stop_token_ids is None:
|
|
self.stop_token_ids = []
|
|
|
|
|
|
@dataclass
|
|
class Request:
|
|
"""
|
|
Represents a single inference request in the scheduling system.
|
|
|
|
Adapted from vLLM's Request class with simplifications for MLX backend.
|
|
|
|
Attributes:
|
|
request_id: Unique identifier for this request
|
|
prompt: The input prompt (string or token ids)
|
|
prompt_token_ids: Tokenized prompt
|
|
sampling_params: Parameters for generation
|
|
arrival_time: When the request was received
|
|
status: Current status of the request
|
|
num_prompt_tokens: Number of tokens in the prompt
|
|
num_computed_tokens: Number of tokens processed so far
|
|
output_token_ids: Generated token ids
|
|
output_text: Generated text (decoded)
|
|
"""
|
|
|
|
request_id: str
|
|
prompt: Union[str, List[int]]
|
|
sampling_params: SamplingParams
|
|
arrival_time: float = field(default_factory=time.monotonic)
|
|
priority: int = 0 # Lower is higher priority
|
|
|
|
# Set after tokenization
|
|
prompt_token_ids: Optional[List[int]] = None
|
|
num_prompt_tokens: int = 0
|
|
|
|
# Generation state
|
|
status: RequestStatus = RequestStatus.WAITING
|
|
num_computed_tokens: int = 0
|
|
output_token_ids: List[int] = field(default_factory=list)
|
|
output_text: str = ""
|
|
generation_started_at: Optional[float] = None
|
|
last_activity_at: Optional[float] = None
|
|
|
|
# For BatchGenerator integration
|
|
batch_uid: Optional[int] = None # UID assigned by BatchGenerator
|
|
|
|
# Prefix cache fields
|
|
prompt_cache: Optional[List[Any]] = None # Cached KV state from prefix cache
|
|
cached_tokens: int = 0 # Number of tokens retrieved from cache
|
|
remaining_tokens: Optional[List[int]] = None # Tokens still needing processing
|
|
|
|
# Paged cache fields (for BlockAwarePrefixCache)
|
|
block_table: Optional["BlockTable"] = None # Block table for paged cache
|
|
shared_prefix_blocks: int = 0 # Number of shared prefix blocks
|
|
|
|
# Multimodal content (images, video)
|
|
images: Optional[List[Any]] = None
|
|
videos: Optional[List[Any]] = None
|
|
|
|
# VLM (Vision-Language Model) fields
|
|
vlm_inputs_embeds: Optional[Any] = (
|
|
None # Pre-computed vision+text embeddings (mx.array)
|
|
)
|
|
vlm_extra_kwargs: Optional[Dict[str, Any]] = (
|
|
None # Model-specific kwargs (e.g., position_ids)
|
|
)
|
|
vlm_image_hash: Optional[str] = None # SHA256 hash of images for prefix cache
|
|
vlm_cache_key_start: int = 0 # Token index where image-specific cache keying starts
|
|
vlm_cache_key_ranges: Optional[List[Tuple[int, str]]] = (
|
|
None # [(token_start, cumulative_image_hash)]
|
|
)
|
|
rope_deltas: float = 0.0 # Per-request mRoPE position delta (set after VLM prefill)
|
|
|
|
@property
|
|
def vlm_extra_keys_for_cache(self) -> Optional[Tuple[str, ...]]:
|
|
"""Whole-request image hash wrapped as extra_keys tuple."""
|
|
if self.vlm_image_hash:
|
|
return (self.vlm_image_hash,)
|
|
return None
|
|
|
|
@property
|
|
def vlm_extra_key_token_start_for_cache(self) -> Optional[int]:
|
|
"""Token index where image-specific cache keying begins."""
|
|
if self.vlm_image_hash:
|
|
return self.vlm_cache_key_start
|
|
return None
|
|
|
|
@property
|
|
def vlm_extra_key_ranges_for_cache(
|
|
self,
|
|
) -> Optional[List[Tuple[int, Tuple[str, ...]]]]:
|
|
"""Segmented VLM cache key ranges for per-image-turn keying."""
|
|
if not self.vlm_cache_key_ranges:
|
|
return None
|
|
return [
|
|
(start, (image_hash,)) for start, image_hash in self.vlm_cache_key_ranges
|
|
]
|
|
|
|
# Metadata
|
|
finish_reason: Optional[str] = None
|
|
|
|
# Reasoning model support (for models with <think> tags)
|
|
needs_think_prefix: bool = False # True if prompt ends with <think> token
|
|
think_prefix_sent: bool = False # Track if prefix already sent
|
|
|
|
# Harmony model support (gpt-oss models)
|
|
is_harmony_model: bool = False # True if model uses Harmony format
|
|
|
|
# SpecPrefill (sparse prefill for MoE models)
|
|
specprefill_indices: Optional[Any] = None # mx.array of selected token indices
|
|
specprefill_total_tokens: int = 0 # Original total token count (M)
|
|
specprefill_position_offset: int = 0 # RoPE offset = M - N
|
|
specprefill_system_end: int = 0 # Token index where system prompt ends
|
|
|
|
# Cache corruption recovery
|
|
cache_corruption_retries: int = 0 # Per-request corruption retry counter
|
|
generation_overflow_retries: int = 0 # Per-request __next_prime retry counter
|
|
|
|
# Prefill memory-pressure recovery
|
|
prefill_oom_retries: int = 0 # Per-request prefill-OOM requeue counter
|
|
prefill_eviction_retries: int = (
|
|
0 # Per-request prefill-headroom eviction phase counter
|
|
)
|
|
|
|
@property
|
|
def num_output_tokens(self) -> int:
|
|
"""Number of output tokens generated so far."""
|
|
return len(self.output_token_ids)
|
|
|
|
@property
|
|
def num_tokens(self) -> int:
|
|
"""Total number of tokens (prompt + output)."""
|
|
return self.num_prompt_tokens + self.num_output_tokens
|
|
|
|
@property
|
|
def max_tokens(self) -> int:
|
|
"""Maximum output tokens for this request."""
|
|
return self.sampling_params.max_tokens
|
|
|
|
def is_finished(self) -> bool:
|
|
"""Check if request has finished."""
|
|
return RequestStatus.is_finished(self.status)
|
|
|
|
def get_finish_reason(self) -> Optional[str]:
|
|
"""Get the finish reason if finished."""
|
|
if self.finish_reason:
|
|
return self.finish_reason
|
|
return RequestStatus.get_finish_reason(self.status)
|
|
|
|
def append_output_token(self, token_id: int) -> None:
|
|
"""Append a generated token to the output."""
|
|
self.output_token_ids.append(token_id)
|
|
self.num_computed_tokens += 1
|
|
|
|
def set_finished(self, status: RequestStatus, reason: Optional[str] = None) -> None:
|
|
"""Mark the request as finished."""
|
|
self.status = status
|
|
self.finish_reason = reason or RequestStatus.get_finish_reason(status)
|
|
|
|
def __lt__(self, other: "Request") -> bool:
|
|
"""Compare requests for priority queue ordering."""
|
|
if self.priority != other.priority:
|
|
return self.priority < other.priority
|
|
return self.arrival_time < other.arrival_time
|
|
|
|
def __hash__(self) -> int:
|
|
return hash(self.request_id)
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
if not isinstance(other, Request):
|
|
return False
|
|
return self.request_id == other.request_id
|
|
|
|
|
|
@dataclass
|
|
class RequestOutput:
|
|
"""
|
|
Output for a single request after a generation step.
|
|
|
|
This is returned by the engine to communicate results back to the API layer.
|
|
"""
|
|
|
|
request_id: str
|
|
# New tokens generated in this step
|
|
new_token_ids: List[int] = field(default_factory=list)
|
|
new_text: str = ""
|
|
# Cumulative output
|
|
output_token_ids: List[int] = field(default_factory=list)
|
|
output_text: str = ""
|
|
# Status
|
|
finished: bool = False
|
|
finish_reason: Optional[str] = None
|
|
# Timing
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
# Internal producer-side timestamp for the first generated token included in
|
|
# this output. Consumers may receive aggregated chunks later than the token
|
|
# was actually produced, so benchmark code must not rely only on receive time.
|
|
generated_at: Optional[float] = None
|
|
# Internal producer-side timestamp for the latest generated token included
|
|
# in this output. This lets aggregated chunks preserve the decode interval.
|
|
generated_until: Optional[float] = None
|
|
|
|
# Tool calls (for Harmony and other models with tool calling support)
|
|
tool_calls: Optional[List[Dict[str, str]]] = None
|
|
# Prefix cache stats
|
|
cached_tokens: int = 0
|
|
# Error message (set when engine encounters an unrecoverable error)
|
|
error: Optional[str] = None
|
|
# Structured internal error classification for API-layer mapping.
|
|
error_code: Optional[str] = None
|
|
error_metadata: Optional[Dict[str, Any]] = None
|
|
|
|
@property
|
|
def usage(self) -> Dict[str, int]:
|
|
"""Return usage statistics compatible with OpenAI API."""
|
|
return {
|
|
"prompt_tokens": self.prompt_tokens,
|
|
"completion_tokens": self.completion_tokens,
|
|
"total_tokens": self.prompt_tokens + self.completion_tokens,
|
|
}
|