baseline6 (3fe05902) clean

This commit is contained in:
root
2026-08-26 01:58:32 +00:00
commit ecbbc80c75
80 changed files with 24763 additions and 0 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,26 @@
import os
def env_bool(name: str, default: bool = False) -> bool:
raw = os.getenv(name)
if raw is None:
return default
if raw in ("1", "true", "True", "yes", "YES", "on", "ON"):
return True
if raw in ("0", "false", "False", "no", "NO", "off", "OFF"):
return False
raise RuntimeError(f"{name} must be boolean, got {raw!r}")
def env_int(name: str, default: int, min_value: int, max_value: int) -> int:
raw = os.getenv(name)
if raw is None:
return default
try:
value = int(raw)
except ValueError as exc:
raise RuntimeError(f"{name} must be int, got {raw!r}") from exc
if not (min_value <= value <= max_value):
raise RuntimeError(
f"{name}={value} outside [{min_value}, {max_value}]")
return value

View File

@@ -0,0 +1,237 @@
import contextlib
import fnmatch
import functools
import json
import os
import re
import threading
import time
from vllm.logger import init_logger
logger = init_logger(__name__)
_EVENT_SCHEMA = "bi100-profile-event-v1"
_EVENT_VERSION = 1
_NAME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_.-]{0,63}$")
_FILTER_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_.*?-]{0,63}$")
def _strict_bool(name: str, default: str = "0") -> bool:
value = os.getenv(name, default).strip()
if value not in {"0", "1"}:
raise RuntimeError(f"{name} must be exactly 0 or 1, got {value!r}")
return value == "1"
_ENABLED = _strict_bool("BI100_PROFILE")
_INCLUDE_STARTUP = _strict_bool("BI100_PROFILE_INCLUDE_STARTUP")
_MODE = os.getenv("BI100_PROFILE_MODE", "sync").strip().lower()
_FILTERS = tuple(
item.strip()
for item in os.getenv("BI100_PROFILE_FILTER", "").split(",")
if item.strip()
)
if _ENABLED and _MODE not in {"sync", "event"}:
raise RuntimeError(f"unsupported BI100_PROFILE_MODE={_MODE!r}")
if _ENABLED and any(_FILTER_RE.fullmatch(pattern) is None
for pattern in _FILTERS):
raise RuntimeError("BI100_PROFILE_FILTER contains an invalid pattern")
_EVENT_RECORDS = []
_COUNTERS = {}
_LOCK = threading.Lock()
_FORWARD_INDEX = 0
_LAST_FLUSH_NS = None
_ACTIVE_FORWARD_TOKEN = None
_NEXT_FORWARD_TOKEN = 0
def _enabled_for(name: str) -> bool:
return (_ENABLED
and (not _FILTERS
or any(fnmatch.fnmatchcase(name, pattern)
for pattern in _FILTERS)))
def _skip_startup() -> bool:
return (not _INCLUDE_STARTUP
and os.getenv("BI100_IN_STARTUP_PROFILE") == "1")
def bi100_profile_event_enabled() -> bool:
return _ENABLED and _MODE == "event" and not _skip_startup()
def _begin_profile_forward():
global _ACTIVE_FORWARD_TOKEN, _NEXT_FORWARD_TOKEN
if not bi100_profile_event_enabled():
return None
with _LOCK:
_EVENT_RECORDS.clear()
_COUNTERS.clear()
token = _NEXT_FORWARD_TOKEN
_NEXT_FORWARD_TOKEN += 1
_ACTIVE_FORWARD_TOKEN = token
return token
def _abort_profile_forward(token) -> None:
global _ACTIVE_FORWARD_TOKEN
if token is None:
return
with _LOCK:
if _ACTIVE_FORWARD_TOKEN != token:
return
_EVENT_RECORDS.clear()
_COUNTERS.clear()
_ACTIVE_FORWARD_TOKEN = None
def bi100_profile_transaction(function):
"""Keep one top-level model forward isolated from failed forwards."""
@functools.wraps(function)
def wrapped(*args, **kwargs):
token = _begin_profile_forward()
if token is None:
return function(*args, **kwargs)
try:
result = function(*args, **kwargs)
except BaseException:
_abort_profile_forward(token)
raise
with _LOCK:
was_flushed = _ACTIVE_FORWARD_TOKEN != token
if not was_flushed:
_abort_profile_forward(token)
raise RuntimeError(
"BI100 profile transaction completed without a flush")
return result
return wrapped
def _normalize_metadata(metadata):
normalized = {}
for key, value in metadata.items():
if not isinstance(key, str) or _NAME_RE.fullmatch(key) is None:
raise TypeError("profile metadata keys must be bounded names")
if isinstance(value, bool):
normalized[key] = value
elif isinstance(value, int) and not isinstance(value, bool):
normalized[key] = value
elif isinstance(value, str) and len(value) <= 64:
normalized[key] = value
else:
raise TypeError(
"profile metadata values must be bool, int, or short strings")
return normalized
def bi100_profile_count(name: str, **metadata) -> None:
"""Record privacy-safe path metadata for the current model forward."""
if not bi100_profile_event_enabled() or not _enabled_for(name):
return
if not isinstance(name, str) or _NAME_RE.fullmatch(name) is None:
raise TypeError("profile counter name must be a bounded name")
normalized = _normalize_metadata(metadata)
encoded = json.dumps(
{"name": name, **normalized}, sort_keys=True, separators=(",", ":"))
with _LOCK:
_COUNTERS[encoded] = _COUNTERS.get(encoded, 0) + 1
@contextlib.contextmanager
def bi100_timer(name: str):
if not _enabled_for(name) or _skip_startup():
yield
return
import torch
if _MODE == "event":
started = torch.cuda.Event(enable_timing=True)
finished = torch.cuda.Event(enable_timing=True)
host_started_ns = time.monotonic_ns()
started.record()
try:
yield
finally:
finished.record()
with _LOCK:
_EVENT_RECORDS.append(
(name, started, finished, host_started_ns))
return
torch.cuda.synchronize()
t0 = time.perf_counter()
try:
yield
finally:
torch.cuda.synchronize()
logger.info("[BI100_PROFILE] %s %.3f ms", name,
(time.perf_counter() - t0) * 1000)
def bi100_profile_flush(*, tp_rank, **metadata):
"""Synchronize once and emit one aggregate event record per model forward."""
global _ACTIVE_FORWARD_TOKEN, _FORWARD_INDEX, _LAST_FLUSH_NS
if not bi100_profile_event_enabled():
return None
if (not isinstance(tp_rank, int) or isinstance(tp_rank, bool)
or not 0 <= tp_rank < 256):
raise TypeError("profile TP rank must be an integer in [0, 255]")
normalized_metadata = _normalize_metadata(metadata)
with _LOCK:
records = list(_EVENT_RECORDS)
counters = dict(_COUNTERS)
_EVENT_RECORDS.clear()
_COUNTERS.clear()
_ACTIVE_FORWARD_TOKEN = None
if not records:
return None
import torch
torch.cuda.synchronize()
flushed_ns = time.monotonic_ns()
regions = {}
model_started_ns = []
for name, started, finished, host_started_ns in records:
stats = regions.setdefault(name, {"count": 0, "total_ms": 0.0})
stats["count"] += 1
stats["total_ms"] += float(started.elapsed_time(finished))
if name == "model.forward":
model_started_ns.append(host_started_ns)
counter_rows = []
for encoded, count in sorted(counters.items()):
row = json.loads(encoded)
row["count"] = count
counter_rows.append(row)
first_model_started_ns = (
min(model_started_ns) if model_started_ns else None)
payload = {
"schema": _EVENT_SCHEMA,
"version": _EVENT_VERSION,
"tp_rank": tp_rank,
"forward_index": _FORWARD_INDEX,
"metadata": normalized_metadata,
"event_count": len(records),
"model_forward_event_count": len(model_started_ns),
"regions": regions,
"counters": counter_rows,
"host_model_start_to_flush_ms": (
(flushed_ns - first_model_started_ns) / 1_000_000
if first_model_started_ns is not None else None),
"host_gap_since_previous_flush_ms": (
(first_model_started_ns - _LAST_FLUSH_NS) / 1_000_000
if first_model_started_ns is not None
and _LAST_FLUSH_NS is not None
else None),
}
_FORWARD_INDEX += 1
_LAST_FLUSH_NS = flushed_ns
logger.info("[BI100_PROFILE_EVENT] %s",
json.dumps(payload, sort_keys=True, separators=(",", ":")))
return payload

View File

@@ -0,0 +1,398 @@
from __future__ import annotations
import os
import time
from collections.abc import Mapping
import torch
from vllm.logger import init_logger
logger = init_logger(__name__)
ENABLE_ENV = "BI100_BLOCK_MAJOR_CPU_KV"
TRACE_ENV = "BI100_BLOCK_MAJOR_CPU_KV_TRACE"
CPU_OFFLOAD_ENV = "BI100_CPU_KV_OFFLOAD"
HYBRID_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
NUM_ATTENTION_LAYERS = 10
KV_PLANES = 2
ELEMENTS_PER_PLANE_BLOCK = 4096
STAGING_BLOCKS = 512
STAGING_BUFFER_COUNT = 2
BYTES_PER_BLOCK = (
NUM_ATTENTION_LAYERS * KV_PLANES * ELEMENTS_PER_PLANE_BLOCK * 2
)
GPU_STAGING_BYTES = STAGING_BLOCKS * STAGING_BUFFER_COUNT * BYTES_PER_BLOCK
def _strict_binary_selector(
name: str,
environ: Mapping[str, str] | None = None,
) -> bool:
source = os.environ if environ is None else environ
raw = source.get(name, "0")
if raw == "0":
return False
if raw == "1":
return True
raise RuntimeError(f"{name} must be exactly '0' or '1', got {raw!r}")
def block_major_cpu_kv_enabled(
environ: Mapping[str, str] | None = None,
) -> bool:
return _strict_binary_selector(ENABLE_ENV, environ)
def block_major_cpu_kv_trace_enabled(
environ: Mapping[str, str] | None = None,
) -> bool:
return _strict_binary_selector(TRACE_ENV, environ)
def _require_block_major_runtime(
environ: Mapping[str, str] | None = None,
) -> None:
source = os.environ if environ is None else environ
if source.get(CPU_OFFLOAD_ENV, "0") != "1":
raise RuntimeError(
f"{ENABLE_ENV}=1 requires {CPU_OFFLOAD_ENV}=1")
if source.get(HYBRID_ACCOUNTING_ENV, "legacy40") != "full_attention":
raise RuntimeError(
f"{ENABLE_ENV}=1 requires "
f"{HYBRID_ACCOUNTING_ENV}=full_attention")
def reserve_block_major_gpu_blocks(
num_gpu_blocks: int,
cache_block_size: int,
environ: Mapping[str, str] | None = None,
) -> int:
if (not isinstance(num_gpu_blocks, int)
or isinstance(num_gpu_blocks, bool)
or num_gpu_blocks < 0):
raise ValueError("num_gpu_blocks must be a non-negative integer")
if not block_major_cpu_kv_enabled(environ):
return num_gpu_blocks
_require_block_major_runtime(environ)
if cache_block_size != BYTES_PER_BLOCK:
raise RuntimeError(
f"{ENABLE_ENV}=1 requires cache block size "
f"{BYTES_PER_BLOCK}, got {cache_block_size}")
reserved_blocks = (
GPU_STAGING_BYTES + cache_block_size - 1
) // cache_block_size
remaining_blocks = num_gpu_blocks - reserved_blocks
if remaining_blocks <= 0:
raise RuntimeError(
"block-major GPU staging leaves no usable GPU KV blocks")
logger.info(
"[BI100 BLOCK KV] capacity reserve blocks=%d bytes=%d "
"profiled_blocks=%d usable_blocks=%d",
reserved_blocks,
GPU_STAGING_BYTES,
num_gpu_blocks,
remaining_blocks,
)
return remaining_blocks
def validate_block_mapping(
mapping: torch.Tensor,
source_limit: int,
destination_limit: int,
) -> tuple[torch.Tensor, torch.Tensor]:
if not isinstance(mapping, torch.Tensor):
raise TypeError("block mapping must be a torch.Tensor")
if mapping.device.type != "cpu":
raise ValueError("block mapping must be on CPU")
if mapping.dtype != torch.int64:
raise ValueError("block mapping must use torch.int64")
if not mapping.is_contiguous():
raise ValueError("block mapping must be contiguous")
if mapping.dim() != 2 or mapping.shape[1] != 2:
raise ValueError("block mapping must have shape [N, 2]")
if source_limit <= 0 or destination_limit <= 0:
raise ValueError("block mapping limits must be positive")
sources: set[int] = set()
destinations: set[int] = set()
for row, pair in enumerate(mapping.tolist()):
source, destination = pair
if not 0 <= source < source_limit:
raise ValueError(
f"source block out of range at row {row}: {source}")
if not 0 <= destination < destination_limit:
raise ValueError(
f"destination block out of range at row {row}: "
f"{destination}")
if source in sources:
raise ValueError(f"duplicate source block: {source}")
if destination in destinations:
raise ValueError(f"duplicate destination block: {destination}")
sources.add(source)
destinations.add(destination)
return mapping[:, 0].contiguous(), mapping[:, 1].contiguous()
class BlockMajorCpuKVCache:
def __init__(
self,
gpu_cache: list[torch.Tensor],
num_cpu_blocks: int,
pin_memory: bool,
) -> None:
self._validate_gpu_cache(gpu_cache)
if block_major_cpu_kv_enabled():
_require_block_major_runtime()
if num_cpu_blocks <= 0:
raise RuntimeError(
f"{ENABLE_ENV}=1 requires a positive CPU block count")
if not pin_memory:
raise RuntimeError(
f"{ENABLE_ENV}=1 requires pinned CPU memory")
try:
from vllm import corex_block_major_kv_transfer as extension
except ImportError as exc:
raise RuntimeError(
"block-major CoreX extension is unavailable") from exc
self.extension = extension
self.gpu_cache = gpu_cache
self.device = gpu_cache[0].device
self.dtype = gpu_cache[0].dtype
self.num_gpu_blocks = gpu_cache[0].shape[1]
self.num_cpu_blocks = num_cpu_blocks
self.trace_enabled = block_major_cpu_kv_trace_enabled()
self.cpu_pool = torch.zeros(
(
num_cpu_blocks,
NUM_ATTENTION_LAYERS,
KV_PLANES,
ELEMENTS_PER_PLANE_BLOCK,
),
dtype=self.dtype,
device="cpu",
pin_memory=True,
)
if not self.cpu_pool.is_pinned():
raise RuntimeError("block-major CPU pool is not pinned")
# Preserve the public CacheEngine shape without allocating a second
# layer-major CPU cache. Transfer methods use cpu_pool directly.
self.layer_views = [
self.cpu_pool[:, layer, :, :].permute(1, 0, 2)
for layer in range(NUM_ATTENTION_LAYERS)
]
self.cpu_staging = [
torch.empty(
(
STAGING_BLOCKS,
NUM_ATTENTION_LAYERS,
KV_PLANES,
ELEMENTS_PER_PLANE_BLOCK,
),
dtype=self.dtype,
device="cpu",
pin_memory=True,
)
for _ in range(STAGING_BUFFER_COUNT)
]
if not all(staging.is_pinned() for staging in self.cpu_staging):
raise RuntimeError("block-major CPU staging is not pinned")
with torch.cuda.device(self.device):
self.gpu_staging = [
torch.empty_like(staging, device=self.device)
for staging in self.cpu_staging
]
self.events = [
torch.cuda.Event(enable_timing=False)
for _ in range(STAGING_BUFFER_COUNT)
]
self.error_flag = torch.zeros(
1, dtype=torch.int32, device=self.device)
logger.info(
"[BI100 BLOCK KV] enabled device=%s gpu_blocks=%d cpu_blocks=%d "
"layers=%d block_bytes=%d staging_blocks=%d staging_buffers=%d",
self.device,
self.num_gpu_blocks,
self.num_cpu_blocks,
NUM_ATTENTION_LAYERS,
BYTES_PER_BLOCK,
STAGING_BLOCKS,
STAGING_BUFFER_COUNT,
)
@staticmethod
def _validate_gpu_cache(gpu_cache: list[torch.Tensor]) -> None:
if len(gpu_cache) != NUM_ATTENTION_LAYERS:
raise RuntimeError(
f"{ENABLE_ENV}=1 requires exactly "
f"{NUM_ATTENTION_LAYERS} GPU attention caches, got "
f"{len(gpu_cache)}")
first = gpu_cache[0]
if first.device.type != "cuda":
raise RuntimeError("block-major GPU cache must be on CUDA")
if first.dtype != torch.float16:
raise RuntimeError("block-major GPU cache must use float16")
if (first.dim() != 3 or first.shape[0] != KV_PLANES
or first.shape[2] != ELEMENTS_PER_PLANE_BLOCK):
raise RuntimeError(
"block-major GPU cache must have shape [2, blocks, 4096]")
if not first.is_contiguous():
raise RuntimeError("block-major GPU cache must be contiguous")
for layer, tensor in enumerate(gpu_cache):
if tensor.device != first.device:
raise RuntimeError(
f"GPU cache layer {layer} is on a different device")
if tensor.dtype != first.dtype or tensor.shape != first.shape:
raise RuntimeError(
f"GPU cache layer {layer} has inconsistent geometry")
if not tensor.is_contiguous():
raise RuntimeError(
f"GPU cache layer {layer} is not contiguous")
def _to_gpu_ids(self, block_ids: torch.Tensor) -> torch.Tensor:
return block_ids.to(
device=self.device,
dtype=torch.int32,
non_blocking=False,
)
@staticmethod
def _chunks(
source: torch.Tensor,
destination: torch.Tensor,
gpu_ids: torch.Tensor,
):
for start in range(0, source.numel(), STAGING_BLOCKS):
end = min(start + STAGING_BLOCKS, source.numel())
yield (
source[start:end],
destination[start:end],
gpu_ids[start:end],
end - start,
)
def _begin(self) -> None:
self.error_flag.zero_()
def _finish(
self,
direction: str,
block_count: int,
started: float | None,
) -> None:
# check_error performs the final stream synchronization. This also
# makes every staging slot safe to reuse in the next CacheEngine call.
self.extension.check_error(self.error_flag)
if started is not None:
elapsed_ms = (time.perf_counter() - started) * 1000.0
logger.info(
"[BI100 BLOCK KV TRACE] direction=%s blocks=%d bytes=%d "
"elapsed_ms=%.3f",
direction,
block_count,
block_count * BYTES_PER_BLOCK,
elapsed_ms,
)
def swap_out(self, mapping: torch.Tensor) -> None:
started = time.perf_counter() if self.trace_enabled else None
source_gpu, destination_cpu = validate_block_mapping(
mapping,
source_limit=self.num_gpu_blocks,
destination_limit=self.num_cpu_blocks,
)
block_count = source_gpu.numel()
if block_count == 0:
return
source_gpu_ids = self._to_gpu_ids(source_gpu)
self._begin()
pending: tuple[int, torch.Tensor, int] | None = None
for index, (_, destination, gpu_ids, count) in enumerate(
self._chunks(
source_gpu, destination_cpu, source_gpu_ids)):
slot = index % STAGING_BUFFER_COUNT
self.extension.pack(
self.gpu_cache,
gpu_ids,
self.gpu_staging[slot],
self.error_flag,
count,
)
self.cpu_staging[slot][:count].copy_(
self.gpu_staging[slot][:count],
non_blocking=True,
)
self.events[slot].record()
if pending is not None:
pending_slot, pending_destination, pending_count = pending
self.events[pending_slot].synchronize()
self.extension.cpu_scatter(
self.cpu_staging[pending_slot],
self.cpu_pool,
pending_destination,
pending_count,
)
pending = (slot, destination, count)
if pending is not None:
pending_slot, pending_destination, pending_count = pending
self.events[pending_slot].synchronize()
self.extension.cpu_scatter(
self.cpu_staging[pending_slot],
self.cpu_pool,
pending_destination,
pending_count,
)
self._finish("d2h", block_count, started)
def swap_in(self, mapping: torch.Tensor) -> None:
started = time.perf_counter() if self.trace_enabled else None
source_cpu, destination_gpu = validate_block_mapping(
mapping,
source_limit=self.num_cpu_blocks,
destination_limit=self.num_gpu_blocks,
)
block_count = source_cpu.numel()
if block_count == 0:
return
destination_gpu_ids = self._to_gpu_ids(destination_gpu)
self._begin()
for index, (source, _, gpu_ids, count) in enumerate(
self._chunks(
source_cpu, destination_gpu, destination_gpu_ids)):
slot = index % STAGING_BUFFER_COUNT
if index >= STAGING_BUFFER_COUNT:
self.events[slot].synchronize()
self.extension.cpu_gather(
self.cpu_pool,
source,
self.cpu_staging[slot],
count,
)
self.gpu_staging[slot][:count].copy_(
self.cpu_staging[slot][:count],
non_blocking=True,
)
self.extension.scatter(
self.gpu_staging[slot],
gpu_ids,
self.gpu_cache,
self.error_flag,
count,
)
self.events[slot].record()
self._finish("h2d", block_count, started)

View File

@@ -0,0 +1,617 @@
import asyncio
import codecs
import json
from abc import ABC, abstractmethod
from collections import defaultdict
from functools import lru_cache, partial
from pathlib import Path
from typing import (Any, Awaitable, Dict, Generic, Iterable, List, Literal,
Mapping, Optional, Tuple, TypeVar, Union, cast)
# yapf conflicts with isort for this block
# yapf: disable
from openai.types.chat import (ChatCompletionAssistantMessageParam,
ChatCompletionContentPartImageParam)
from openai.types.chat import (
ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam)
from openai.types.chat import (ChatCompletionContentPartRefusalParam,
ChatCompletionContentPartTextParam)
from openai.types.chat import (
ChatCompletionMessageParam as OpenAIChatCompletionMessageParam)
from openai.types.chat import (ChatCompletionMessageToolCallParam,
ChatCompletionToolMessageParam)
# yapf: enable
# pydantic needs the TypedDict from typing_extensions
from pydantic import ConfigDict
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
from typing_extensions import Required, TypeAlias, TypedDict
from vllm.config import ModelConfig
from vllm.logger import init_logger
from vllm.multimodal import MultiModalDataDict
from vllm.multimodal.utils import (async_get_and_parse_audio,
async_get_and_parse_image,
get_and_parse_audio, get_and_parse_image)
from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
logger = init_logger(__name__)
class AudioURL(TypedDict, total=False):
url: Required[str]
"""
Either a URL of the audio or a data URL with base64 encoded audio data.
"""
class ChatCompletionContentPartAudioParam(TypedDict, total=False):
audio_url: Required[AudioURL]
type: Required[Literal["audio_url"]]
"""The type of the content part."""
class CustomChatCompletionContentPartParam(TypedDict, total=False):
__pydantic_config__ = ConfigDict(extra="allow") # type: ignore
type: Required[str]
"""The type of the content part."""
ChatCompletionContentPartParam: TypeAlias = Union[
OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam,
ChatCompletionContentPartRefusalParam,
CustomChatCompletionContentPartParam]
class CustomChatCompletionMessageParam(TypedDict, total=False):
"""Enables custom roles in the Chat Completion API."""
role: Required[str]
"""The role of the message's author."""
content: Union[str, List[ChatCompletionContentPartParam]]
"""The contents of the message."""
name: str
"""An optional name for the participant.
Provides the model information to differentiate between participants of the
same role.
"""
tool_call_id: Optional[str]
"""Tool call that this message is responding to."""
tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
"""The tool calls generated by the model, such as function calls."""
reasoning_content: Optional[str]
"""Reasoning / thinking content for assistant messages (vLLM extension).
When present in a previous assistant turn, it is rendered as
<think>...</think> before the main content so the model sees its own
chain-of-thought in subsequent turns."""
ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
CustomChatCompletionMessageParam]
# TODO: Make fields ReadOnly once mypy supports it
class ConversationMessage(TypedDict, total=False):
role: Required[str]
"""The role of the message's author."""
content: Optional[str]
"""The contents of the message"""
tool_call_id: Optional[str]
"""Tool call that this message is responding to."""
name: Optional[str]
"""The name of the function to call"""
tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
"""The tool calls generated by the model, such as function calls."""
reasoning_content: Optional[str]
"""Reasoning / thinking content for assistant messages.
Passed directly to the chat template (Qwen3 reads message.reasoning_content
natively) instead of being manually wrapped in <think>...</think>."""
ModalityStr = Literal["image", "audio", "video"]
_T = TypeVar("_T")
class BaseMultiModalItemTracker(ABC, Generic[_T]):
"""
Tracks multi-modal items in a given request and ensures that the number
of multi-modal items in a given request does not exceed the configured
maximum per prompt.
"""
def __init__(self, model_config: ModelConfig, tokenizer: AnyTokenizer):
super().__init__()
self._model_config = model_config
self._tokenizer = tokenizer
self._allowed_items = (model_config.multimodal_config.limit_per_prompt
if model_config.multimodal_config else {})
self._consumed_items = {k: 0 for k in self._allowed_items}
self._items: List[_T] = []
@staticmethod
@lru_cache(maxsize=None)
def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str:
return tokenizer.decode(token_index)
def _placeholder_str(self, modality: ModalityStr,
current_count: int) -> Optional[str]:
# TODO: Let user specify how to insert image tokens into prompt
# (similar to chat template)
hf_config = self._model_config.hf_config
model_type = hf_config.model_type
if modality == "image":
if model_type == "phi3_v":
# Workaround since this token is not defined in the tokenizer
return f"<|image_{current_count}|>"
if model_type == "minicpmv":
return "(<image>./</image>)"
if model_type in ("blip-2", "chatglm", "fuyu", "paligemma",
"pixtral"):
# These models do not use image tokens in the prompt
return None
if model_type == "qwen":
return f"Picture {current_count}: <img></img>"
if model_type.startswith("llava"):
return self._cached_token_str(self._tokenizer,
hf_config.image_token_index)
if model_type in ("chameleon", "internvl_chat", "NVLM_D"):
return "<image>"
if model_type == "mllama":
return "<|image|>"
if model_type in ("qwen2_vl", "qwen2_5_vl", "qwen3_5",
"qwen3_5_moe"):
return "<|vision_start|><|image_pad|><|vision_end|>"
if model_type == "molmo":
return ""
raise TypeError(f"Unknown model type: {model_type}")
elif modality == "audio":
if model_type == "ultravox":
return "<|reserved_special_token_0|>"
raise TypeError(f"Unknown model type: {model_type}")
elif modality == "video":
if model_type in ("qwen2_vl","qwen2_5_vl"):
return "<|vision_start|><|video_pad|><|vision_end|>"
raise TypeError(f"Unknown model type: {model_type}")
else:
raise TypeError(f"Unknown modality: {modality}")
@staticmethod
def _combine(items: List[MultiModalDataDict]) -> MultiModalDataDict:
mm_lists: Mapping[str, List[object]] = defaultdict(list)
# Merge all the multi-modal items
for single_mm_data in items:
for mm_key, mm_item in single_mm_data.items():
if isinstance(mm_item, list):
mm_lists[mm_key].extend(mm_item)
else:
mm_lists[mm_key].append(mm_item)
# Unpack any single item lists for models that don't expect multiple.
return {
mm_key: mm_list[0] if len(mm_list) == 1 else mm_list
for mm_key, mm_list in mm_lists.items()
}
def add(self, modality: ModalityStr, item: _T) -> Optional[str]:
"""
Add a multi-modal item to the current prompt and returns the
placeholder string to use, if any.
"""
allowed_count = self._allowed_items.get(modality, 1)
current_count = self._consumed_items.get(modality, 0) + 1
if current_count > allowed_count:
raise ValueError(
f"At most {allowed_count} {modality}(s) may be provided in "
"one request.")
self._consumed_items[modality] = current_count
self._items.append(item)
return self._placeholder_str(modality, current_count)
@abstractmethod
def create_parser(self) -> "BaseMultiModalContentParser":
raise NotImplementedError
class MultiModalItemTracker(BaseMultiModalItemTracker[MultiModalDataDict]):
def all_mm_data(self) -> Optional[MultiModalDataDict]:
return self._combine(self._items) if self._items else None
def create_parser(self) -> "BaseMultiModalContentParser":
return MultiModalContentParser(self)
class AsyncMultiModalItemTracker(
BaseMultiModalItemTracker[Awaitable[MultiModalDataDict]]):
async def all_mm_data(self) -> Optional[MultiModalDataDict]:
if self._items:
items = await asyncio.gather(*self._items)
return self._combine(items)
return None
def create_parser(self) -> "BaseMultiModalContentParser":
return AsyncMultiModalContentParser(self)
class BaseMultiModalContentParser(ABC):
def __init__(self) -> None:
super().__init__()
# multimodal placeholder_string : count
self._placeholder_counts: Dict[str, int] = defaultdict(lambda: 0)
def _add_placeholder(self, placeholder: Optional[str]):
if placeholder:
self._placeholder_counts[placeholder] += 1
def mm_placeholder_counts(self) -> Dict[str, int]:
return dict(self._placeholder_counts)
@abstractmethod
def parse_image(self, image_url: str) -> None:
raise NotImplementedError
@abstractmethod
def parse_audio(self, audio_url: str) -> None:
raise NotImplementedError
class MultiModalContentParser(BaseMultiModalContentParser):
def __init__(self, tracker: MultiModalItemTracker) -> None:
super().__init__()
self._tracker = tracker
def parse_image(self, image_url: str) -> None:
image = get_and_parse_image(image_url)
placeholder = self._tracker.add("image", image)
self._add_placeholder(placeholder)
def parse_audio(self, audio_url: str) -> None:
audio = get_and_parse_audio(audio_url)
placeholder = self._tracker.add("audio", audio)
self._add_placeholder(placeholder)
class AsyncMultiModalContentParser(BaseMultiModalContentParser):
def __init__(self, tracker: AsyncMultiModalItemTracker) -> None:
super().__init__()
self._tracker = tracker
def parse_image(self, image_url: str) -> None:
image_coro = async_get_and_parse_image(image_url)
placeholder = self._tracker.add("image", image_coro)
self._add_placeholder(placeholder)
def parse_audio(self, audio_url: str) -> None:
audio_coro = async_get_and_parse_audio(audio_url)
placeholder = self._tracker.add("audio", audio_coro)
self._add_placeholder(placeholder)
def validate_chat_template(chat_template: Optional[Union[Path, str]]):
"""Raises if the provided chat template appears invalid."""
if chat_template is None:
return
elif isinstance(chat_template, Path) and not chat_template.exists():
raise FileNotFoundError(
"the supplied chat template path doesn't exist")
elif isinstance(chat_template, str):
JINJA_CHARS = "{}\n"
if not any(c in chat_template
for c in JINJA_CHARS) and not Path(chat_template).exists():
raise ValueError(
f"The supplied chat template string ({chat_template}) "
f"appears path-like, but doesn't exist!")
else:
raise TypeError(
f"{type(chat_template)} is not a valid chat template type")
def load_chat_template(
chat_template: Optional[Union[Path, str]]) -> Optional[str]:
if chat_template is None:
return None
try:
with open(chat_template, "r") as f:
resolved_chat_template = f.read()
except OSError as e:
if isinstance(chat_template, Path):
raise
JINJA_CHARS = "{}\n"
if not any(c in chat_template for c in JINJA_CHARS):
msg = (f"The supplied chat template ({chat_template}) "
f"looks like a file path, but it failed to be "
f"opened. Reason: {e}")
raise ValueError(msg) from e
# If opening a file fails, set chat template to be args to
# ensure we decode so our escape are interpreted correctly
resolved_chat_template = codecs.decode(chat_template, "unicode_escape")
logger.info("Using supplied chat template:\n%s", resolved_chat_template)
return resolved_chat_template
# TODO: Let user specify how to insert multimodal tokens into prompt
# (similar to chat template)
def _get_full_multimodal_text_prompt(placeholder_counts: Dict[str, int],
text_prompt: str) -> str:
"""Combine multimodal prompts for a multimodal language model."""
# Look through the text prompt to check for missing placeholders
missing_placeholders: List[str] = []
for placeholder in placeholder_counts:
# For any existing placeholder in the text prompt, we leave it as is
placeholder_counts[placeholder] -= text_prompt.count(placeholder)
if placeholder_counts[placeholder] < 0:
raise ValueError(
f"Found more '{placeholder}' placeholders in input prompt than "
"actual multimodal data items.")
missing_placeholders.extend([placeholder] *
placeholder_counts[placeholder])
# NOTE: For now we always add missing placeholders at the front of
# the prompt. This may change to be customizable in the future.
return "\n".join(missing_placeholders + [text_prompt])
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
_ImageParser = partial(cast, ChatCompletionContentPartImageParam)
_AudioParser = partial(cast, ChatCompletionContentPartAudioParam)
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
MODEL_KEEP_MULTI_MODAL_CONTENT = {'mllama'}
def _parse_chat_message_content_parts(
role: str,
parts: Iterable[ChatCompletionContentPartParam],
mm_tracker: BaseMultiModalItemTracker,
) -> List[ConversationMessage]:
texts: List[str] = []
mm_parser = mm_tracker.create_parser()
keep_multimodal_content = \
mm_tracker._model_config.hf_config.model_type in \
MODEL_KEEP_MULTI_MODAL_CONTENT
has_image = False
for part in parts:
part_type = part["type"]
if part_type == "text":
text = _TextParser(part)["text"]
texts.append(text)
elif part_type == "image_url":
image_url = _ImageParser(part)["image_url"]
if image_url.get("detail", "auto") != "auto":
logger.warning(
"'image_url.detail' is currently not supported and "
"will be ignored.")
mm_parser.parse_image(image_url["url"])
has_image = True
elif part_type == "audio_url":
audio_url = _AudioParser(part)["audio_url"]
mm_parser.parse_audio(audio_url["url"])
elif part_type == "refusal":
text = _RefusalParser(part)["refusal"]
texts.append(text)
else:
raise NotImplementedError(f"Unknown part type: {part_type}")
text_prompt = "\n".join(texts)
if keep_multimodal_content:
text_prompt = "\n".join(texts)
role_content = [{'type': 'text', 'text': text_prompt}]
if has_image:
role_content = [{'type': 'image'}] + role_content
return [ConversationMessage(role=role,
content=role_content)] # type: ignore
else:
mm_placeholder_counts = mm_parser.mm_placeholder_counts()
if mm_placeholder_counts:
text_prompt = _get_full_multimodal_text_prompt(
mm_placeholder_counts, text_prompt)
return [ConversationMessage(role=role, content=text_prompt)]
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)
def _parse_chat_message_content(
message: ChatCompletionMessageParam,
mm_tracker: BaseMultiModalItemTracker,
) -> List[ConversationMessage]:
role = message["role"]
content = message.get("content")
if content is None:
content = []
elif isinstance(content, str):
content = [
ChatCompletionContentPartTextParam(type="text", text=content)
]
result = _parse_chat_message_content_parts(
role,
content, # type: ignore
mm_tracker,
)
for result_msg in result:
if role == 'assistant':
parsed_msg = _AssistantParser(message)
if "tool_calls" in parsed_msg:
result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
# Pass reasoning content as a dedicated field so the chat template
# can render it natively (Qwen3: message.reasoning_content branch).
# Accept both "reasoning" (new vllm) and "reasoning_content" (ours).
reasoning = (message.get("reasoning") # type: ignore[arg-type]
or message.get("reasoning_content")) # type: ignore[arg-type]
if reasoning and isinstance(reasoning, str):
result_msg["reasoning_content"] = reasoning
elif role == "tool":
parsed_msg = _ToolParser(message)
if "tool_call_id" in parsed_msg:
result_msg["tool_call_id"] = parsed_msg["tool_call_id"]
if "name" in message and isinstance(message["name"], str):
result_msg["name"] = message["name"]
return result
def _postprocess_messages(messages: List[ConversationMessage]) -> None:
# per the Transformers docs & maintainers, tool call arguments in
# assistant-role messages with tool_calls need to be dicts not JSON str -
# this is how tool-use chat templates will expect them moving forwards
# so, for messages that have tool_calls, parse the string (which we get
# from openAI format) to dict
for message in messages:
if (message["role"] == "assistant" and "tool_calls" in message
and message["tool_calls"] is not None):
if not isinstance(message["tool_calls"], list):
message["tool_calls"] = list(message["tool_calls"])
for item in message["tool_calls"]:
arguments = item["function"]["arguments"]
if isinstance(arguments, str):
try:
arguments = json.loads(arguments)
except json.JSONDecodeError as exc:
raise ValueError(
"Tool call arguments are not valid JSON.") from exc
elif not isinstance(arguments, dict):
raise TypeError(
"Tool call arguments must be a JSON object or a "
"JSON-encoded object string.")
if not isinstance(arguments, dict):
raise TypeError(
"Tool call arguments must decode to a JSON object.")
item["function"]["arguments"] = arguments
def parse_chat_messages(
messages: List[ChatCompletionMessageParam],
model_config: ModelConfig,
tokenizer: AnyTokenizer,
) -> Tuple[List[ConversationMessage], Optional[MultiModalDataDict]]:
conversation: List[ConversationMessage] = []
mm_tracker = MultiModalItemTracker(model_config, tokenizer)
for msg in messages:
sub_messages = _parse_chat_message_content(msg, mm_tracker)
conversation.extend(sub_messages)
_postprocess_messages(conversation)
return conversation, mm_tracker.all_mm_data()
def parse_chat_messages_futures(
messages: List[ChatCompletionMessageParam],
model_config: ModelConfig,
tokenizer: AnyTokenizer,
) -> Tuple[List[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]:
conversation: List[ConversationMessage] = []
mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)
for msg in messages:
sub_messages = _parse_chat_message_content(msg, mm_tracker)
conversation.extend(sub_messages)
_postprocess_messages(conversation)
return conversation, mm_tracker.all_mm_data()
def apply_hf_chat_template(
tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
conversation: List[ConversationMessage],
chat_template: Optional[str],
*,
tokenize: bool = False, # Different from HF's default
**kwargs: Any,
) -> str:
if chat_template is None and tokenizer.chat_template is None:
raise ValueError(
"As of transformers v4.44, default chat template is no longer "
"allowed, so you must provide a chat template if the tokenizer "
"does not define one.")
return tokenizer.apply_chat_template(
conversation=conversation, # type: ignore[arg-type]
chat_template=chat_template,
tokenize=tokenize,
**kwargs,
)
def apply_mistral_chat_template(
tokenizer: MistralTokenizer,
messages: List[ChatCompletionMessageParam],
chat_template: Optional[str] = None,
**kwargs: Any,
) -> List[int]:
if chat_template is not None:
logger.warning(
"'chat_template' cannot be overridden for mistral tokenizer.")
if "add_generation_prompt" in kwargs:
logger.warning(
"'add_generation_prompt' is not supported for mistral tokenizer, "
"so it will be ignored.")
if "continue_final_message" in kwargs:
logger.warning(
"'continue_final_message' is not supported for mistral tokenizer, "
"so it will be ignored.")
return tokenizer.apply_chat_template(
messages=messages,
**kwargs,
)

261
qwen3_6_scripts/cli_args.py Normal file
View File

@@ -0,0 +1,261 @@
"""
This file contains the command line arguments for the vLLM's
OpenAI-compatible server. It is kept in a separate file for documentation
purposes.
"""
import argparse
import json
import ssl
from typing import List, Optional, Sequence, Union
from vllm.engine.arg_utils import AsyncEngineArgs, nullable_str
from vllm.entrypoints.chat_utils import validate_chat_template
from vllm.entrypoints.openai.serving_engine import (LoRAModulePath,
PromptAdapterPath)
from vllm.entrypoints.openai.tool_parsers import ToolParserManager
from vllm.utils import FlexibleArgumentParser
class LoRAParserAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser,
namespace: argparse.Namespace,
values: Optional[Union[str, Sequence[str]]],
option_string: Optional[str] = None,
):
if values is None:
values = []
if isinstance(values, str):
raise TypeError("Expected values to be a list")
lora_list: List[LoRAModulePath] = []
for item in values:
if item in [None, '']: # Skip if item is None or empty string
continue
if '=' in item and ',' not in item: # Old format: name=path
name, path = item.split('=')
lora_list.append(LoRAModulePath(name, path))
else: # Assume JSON format
try:
lora_dict = json.loads(item)
lora = LoRAModulePath(**lora_dict)
lora_list.append(lora)
except json.JSONDecodeError:
parser.error(
f"Invalid JSON format for --lora-modules: {item}")
except TypeError as e:
parser.error(
f"Invalid fields for --lora-modules: {item} - {str(e)}"
)
setattr(namespace, self.dest, lora_list)
class PromptAdapterParserAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser,
namespace: argparse.Namespace,
values: Optional[Union[str, Sequence[str]]],
option_string: Optional[str] = None,
):
if values is None:
values = []
if isinstance(values, str):
raise TypeError("Expected values to be a list")
adapter_list: List[PromptAdapterPath] = []
for item in values:
name, path = item.split('=')
adapter_list.append(PromptAdapterPath(name, path))
setattr(namespace, self.dest, adapter_list)
def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
parser.add_argument("--host",
type=nullable_str,
default=None,
help="host name")
parser.add_argument("--port", type=int, default=8000, help="port number")
parser.add_argument(
"--uvicorn-log-level",
type=str,
default="info",
choices=['debug', 'info', 'warning', 'error', 'critical', 'trace'],
help="log level for uvicorn")
parser.add_argument("--allow-credentials",
action="store_true",
help="allow credentials")
parser.add_argument("--allowed-origins",
type=json.loads,
default=["*"],
help="allowed origins")
parser.add_argument("--allowed-methods",
type=json.loads,
default=["*"],
help="allowed methods")
parser.add_argument("--allowed-headers",
type=json.loads,
default=["*"],
help="allowed headers")
parser.add_argument("--api-key",
type=nullable_str,
default=None,
help="If provided, the server will require this key "
"to be presented in the header.")
parser.add_argument(
"--lora-modules",
type=nullable_str,
default=None,
nargs='+',
action=LoRAParserAction,
help="LoRA module configurations in either 'name=path' format"
"or JSON format. "
"Example (old format): 'name=path' "
"Example (new format): "
"'{\"name\": \"name\", \"local_path\": \"path\", "
"\"base_model_name\": \"id\"}'")
parser.add_argument(
"--prompt-adapters",
type=nullable_str,
default=None,
nargs='+',
action=PromptAdapterParserAction,
help="Prompt adapter configurations in the format name=path. "
"Multiple adapters can be specified.")
parser.add_argument("--chat-template",
type=nullable_str,
default=None,
help="The file path to the chat template, "
"or the template in single-line form "
"for the specified model")
parser.add_argument("--response-role",
type=nullable_str,
default="assistant",
help="The role name to return if "
"`request.add_generation_prompt=true`.")
parser.add_argument("--ssl-keyfile",
type=nullable_str,
default=None,
help="The file path to the SSL key file")
parser.add_argument("--ssl-certfile",
type=nullable_str,
default=None,
help="The file path to the SSL cert file")
parser.add_argument("--ssl-ca-certs",
type=nullable_str,
default=None,
help="The CA certificates file")
parser.add_argument(
"--ssl-cert-reqs",
type=int,
default=int(ssl.CERT_NONE),
help="Whether client certificate is required (see stdlib ssl module's)"
)
parser.add_argument(
"--root-path",
type=nullable_str,
default=None,
help="FastAPI root_path when app is behind a path based routing proxy")
parser.add_argument(
"--middleware",
type=nullable_str,
action="append",
default=[],
help="Additional ASGI middleware to apply to the app. "
"We accept multiple --middleware arguments. "
"The value should be an import path. "
"If a function is provided, vLLM will add it to the server "
"using @app.middleware('http'). "
"If a class is provided, vLLM will add it to the server "
"using app.add_middleware(). ")
parser.add_argument(
"--return-tokens-as-token-ids",
action="store_true",
help="When --max-logprobs is specified, represents single tokens as "
"strings of the form 'token_id:{token_id}' so that tokens that "
"are not JSON-encodable can be identified.")
parser.add_argument(
"--disable-frontend-multiprocessing",
action="store_true",
help="If specified, will run the OpenAI frontend server in the same "
"process as the model serving engine.")
parser.add_argument(
"--enable-auto-tool-choice",
action="store_true",
default=False,
help=
"Enable auto tool choice for supported models. Use --tool-call-parser"
"to specify which parser to use")
valid_tool_parsers = ToolParserManager.tool_parsers.keys()
parser.add_argument(
"--tool-call-parser",
type=str,
metavar="{" + ",".join(valid_tool_parsers) + "} or name registered in "
"--tool-parser-plugin",
default=None,
help=
"Select the tool call parser depending on the model that you're using."
" This is used to parse the model-generated tool call into OpenAI API "
"format. Required for --enable-auto-tool-choice.")
parser.add_argument(
"--tool-parser-plugin",
type=str,
default="",
help=
"Special the tool parser plugin write to parse the model-generated tool"
" into OpenAI API format, the name register in this plugin can be used "
"in --tool-call-parser.")
parser.add_argument(
"--reasoning-parser",
type=str,
default=None,
help=
"Select the reasoning parser to split <think>...</think> content into "
"reasoning_content vs content in the response. "
"Supported: qwen3")
parser = AsyncEngineArgs.add_cli_args(parser)
parser.add_argument('--max-log-len',
type=int,
default=None,
help='Max number of prompt characters or prompt '
'ID numbers being printed in log.'
'\n\nDefault: Unlimited')
parser.add_argument(
"--disable-fastapi-docs",
action='store_true',
default=False,
help="Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint"
)
return parser
def validate_parsed_serve_args(args: argparse.Namespace):
"""Quick checks for model serve args that raise prior to loading."""
if hasattr(args, "subparser") and args.subparser != "serve":
return
# Ensure that the chat template is valid; raises if it likely isn't
validate_chat_template(args.chat_template)
# Enable auto tool needs a tool call parser to be valid
if args.enable_auto_tool_choice and not args.tool_call_parser:
raise TypeError("Error: --enable-auto-tool-choice requires "
"--tool-call-parser")
def create_parser_for_docs() -> FlexibleArgumentParser:
parser_for_docs = FlexibleArgumentParser(
prog="-m vllm.entrypoints.openai.api_server")
return make_arg_parser(parser_for_docs)

View File

@@ -0,0 +1,291 @@
"""Shared GDN prefix-state cache contracts for the BI100 runtime."""
from __future__ import annotations
import os
from collections import OrderedDict
from dataclasses import dataclass
from typing import Iterable, List, Optional, Sequence, Tuple
GdnPrefixKey = Tuple[int, bytes]
GdnCapturePoint = Tuple[int, GdnPrefixKey]
_VALID_POLICIES = {"fine32", "admission64", "off"}
GDN_KERNEL_CHUNK_TOKENS = 64
GDN_DIRECT_MIN_REPLAY_TOKENS = 2
_VALID_RESTORE_MODES = {"direct", "hybrid64", "chunk64", "aligned"}
def _env_choice(name: str, default: str, choices: set[str]) -> str:
value = os.getenv(name, default).strip().lower()
if value not in choices:
allowed = ", ".join(sorted(choices))
raise RuntimeError(f"invalid {name}={value!r}; expected one of: {allowed}")
return value
def gdn_cache_policy_from_env() -> str:
return _env_choice("BI100_GDN_CACHE_POLICY", "fine32", _VALID_POLICIES)
def gdn_restore_mode_from_env() -> str:
return _env_choice(
"BI100_GDN_RESTORE_MODE", "direct", _VALID_RESTORE_MODES)
def gdn_restore_alignment(restore_mode: str, block_size: int,
scheduler_chunk_tokens: int) -> int:
"""Return the content boundary required by a restore mode."""
if block_size <= 0:
raise ValueError("block_size must be positive")
if restore_mode == "direct":
return block_size
if restore_mode in {"hybrid64", "chunk64"}:
alignment = GDN_KERNEL_CHUNK_TOKENS
elif restore_mode == "aligned":
alignment = scheduler_chunk_tokens
else:
raise ValueError(f"unknown GDN restore mode: {restore_mode}")
if alignment <= 0 or alignment % block_size != 0:
raise ValueError(
f"{restore_mode} GDN restore requires a positive alignment "
f"divisible by block_size={block_size}; got {alignment}")
return alignment
def make_prefix_key(block_count: int, digest: bytes) -> GdnPrefixKey:
if block_count <= 0:
raise ValueError("GDN prefix key requires at least one complete block")
if not isinstance(digest, bytes) or len(digest) != 32:
raise ValueError("GDN prefix digest must be exactly 32 bytes")
return block_count, digest
def keys_from_block_hashes(block_hashes: Sequence[bytes]) -> List[GdnPrefixKey]:
return [make_prefix_key(i + 1, digest)
for i, digest in enumerate(block_hashes)]
def strict_prefix_block_count(token_count: int, block_size: int) -> int:
if block_size <= 0:
raise ValueError("block_size must be positive")
if token_count <= 1:
return 0
return (token_count - 1) // block_size
def key_at_strict_boundary(block_hashes: Sequence[bytes], token_count: int,
block_size: int) -> Optional[GdnPrefixKey]:
block_count = min(
len(block_hashes), strict_prefix_block_count(token_count, block_size))
if block_count <= 0:
return None
return make_prefix_key(block_count, block_hashes[block_count - 1])
def final_capture_key(
block_hashes: Sequence[bytes], prompt_tokens: int, block_size: int,
restore_mode: str, replay_alignment: int) -> Optional[GdnPrefixKey]:
if restore_mode in {"direct", "hybrid64"}:
block_count = min(
len(block_hashes), strict_prefix_block_count(
prompt_tokens, block_size))
if (block_count > 0
and prompt_tokens - block_count * block_size
< GDN_DIRECT_MIN_REPLAY_TOKENS):
block_count -= 1
if block_count <= 0:
return None
return make_prefix_key(block_count, block_hashes[block_count - 1])
if restore_mode not in {"chunk64", "aligned"}:
raise ValueError(f"unknown GDN restore mode: {restore_mode}")
if (replay_alignment <= 0 or replay_alignment % block_size != 0
or prompt_tokens <= 1):
return None
boundary_tokens = ((prompt_tokens - 1) // replay_alignment
* replay_alignment)
block_count = min(len(block_hashes), boundary_tokens // block_size)
if block_count <= 0:
return None
return make_prefix_key(block_count, block_hashes[block_count - 1])
def restore_key_is_eligible(
key: GdnPrefixKey, prompt_tokens: int, block_size: int,
restore_mode: str, replay_alignment: int,
direct_final_key: Optional[GdnPrefixKey] = None) -> bool:
"""Return whether restoring ``key`` preserves the execution contract."""
make_prefix_key(*key)
if block_size <= 0:
raise ValueError("block_size must be positive")
boundary_tokens = key[0] * block_size
remaining_tokens = prompt_tokens - boundary_tokens
if remaining_tokens <= 0:
return False
if restore_mode == "direct":
return remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
if restore_mode == "hybrid64":
if direct_final_key is not None:
make_prefix_key(*direct_final_key)
return (remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
and replay_alignment > 0
and (boundary_tokens % replay_alignment == 0
or key == direct_final_key))
if restore_mode not in {"chunk64", "aligned"}:
raise ValueError(f"unknown GDN restore mode: {restore_mode}")
return (replay_alignment > 0
and boundary_tokens % replay_alignment == 0)
def capture_points_for_step(
targets: Iterable[GdnPrefixKey], physical_context_tokens: int,
logical_end_tokens: int, block_size: int) -> Tuple[GdnCapturePoint, ...]:
if physical_context_tokens < 0 or logical_end_tokens < 0:
raise ValueError("token positions must be non-negative")
if logical_end_tokens <= physical_context_tokens:
return ()
selected = {}
for key in targets:
make_prefix_key(*key)
boundary_tokens = key[0] * block_size
if physical_context_tokens < boundary_tokens <= logical_end_tokens:
selected[boundary_tokens - physical_context_tokens] = key
points = tuple(sorted(selected.items()))
if len(points) > 2:
raise ValueError("at most two GDN capture points are allowed per step")
return points
def cap_prefill_end_at_capture_boundary(
logical_start_tokens: int, logical_end_tokens: int,
targets: Iterable[GdnPrefixKey], block_size: int) -> int:
"""Stop a physical prefill step at its earliest pending capture boundary."""
if logical_start_tokens < 0 or logical_end_tokens < 0:
raise ValueError("token positions must be non-negative")
if logical_end_tokens < logical_start_tokens:
raise ValueError("logical end must not precede logical start")
if block_size <= 0:
raise ValueError("block_size must be positive")
capped_end = logical_end_tokens
for key in targets:
make_prefix_key(*key)
boundary_tokens = key[0] * block_size
if logical_start_tokens < boundary_tokens < capped_end:
capped_end = boundary_tokens
return capped_end
def canonical_direct_segment_offsets(
block_hashes: Sequence[bytes], physical_context_tokens: int,
logical_end_tokens: int, block_size: int,
scheduler_chunk_tokens: int) -> Tuple[int, ...]:
"""Reproduce cold fine32/direct segment boundaries after fast-forward."""
if physical_context_tokens < 0 or logical_end_tokens < 0:
raise ValueError("token positions must be non-negative")
if block_size <= 0 or scheduler_chunk_tokens <= 0:
raise ValueError("block and scheduler chunk sizes must be positive")
if scheduler_chunk_tokens % block_size != 0:
raise ValueError("scheduler chunk size must be divisible by block size")
if logical_end_tokens <= physical_context_tokens:
return ()
boundaries = set()
step_ends = list(range(scheduler_chunk_tokens, logical_end_tokens,
scheduler_chunk_tokens))
for step_end in (*step_ends, logical_end_tokens):
key = final_capture_key(block_hashes, step_end, block_size,
"direct", block_size)
if key is not None:
boundaries.add(key[0] * block_size)
boundaries.update(step_ends)
return tuple(
boundary - physical_context_tokens
for boundary in sorted(boundaries)
if physical_context_tokens < boundary < logical_end_tokens)
@dataclass(frozen=True)
class GdnCachePlan:
restore_key: Optional[GdnPrefixKey] = None
capture_points: Tuple[GdnCapturePoint, ...] = ()
evict_keys: Tuple[GdnPrefixKey, ...] = ()
class GdnPrefixStatePolicy:
"""Scheduler-owned state index with deterministic worker actions."""
def __init__(self, policy: str) -> None:
if policy not in _VALID_POLICIES:
raise ValueError(f"unknown GDN cache policy: {policy}")
self.policy = policy
self.capacity = {"fine32": 32, "admission64": 64, "off": 0}[policy]
self._resident: OrderedDict[GdnPrefixKey, None] = OrderedDict()
def __len__(self) -> int:
return len(self._resident)
def resident_keys(self) -> Tuple[GdnPrefixKey, ...]:
return tuple(self._resident)
def contains(self, key: GdnPrefixKey) -> bool:
return key in self._resident
def should_capture_final(self, key: GdnPrefixKey) -> bool:
"""Return whether a final state must be materialized on this request."""
make_prefix_key(*key)
if self.policy == "off":
return False
if self.policy == "admission64":
return key not in self._resident
return True
def select_restore(
self, live_prefix_keys: Sequence[GdnPrefixKey],
max_blocks: int) -> Optional[GdnPrefixKey]:
if self.capacity == 0 or max_blocks <= 0:
return None
best = None
for key in live_prefix_keys[:max_blocks]:
if key in self._resident:
best = key
if best is not None:
self._resident.move_to_end(best)
return best
def repeated_branch_candidate(
self, live_prefix_keys: Sequence[GdnPrefixKey],
max_blocks: int) -> Optional[GdnPrefixKey]:
"""Return a repeated raw-KV branch that lacks recurrent state.
A live KV hit proves that the content occurred in an earlier request;
the current request is therefore the second or later occurrence.
"""
if (self.policy != "admission64" or max_blocks <= 0
or not live_prefix_keys):
return None
candidate = live_prefix_keys[min(len(live_prefix_keys), max_blocks) - 1]
if candidate in self._resident:
return None
return candidate
def admit(self, keys: Iterable[GdnPrefixKey]) -> Tuple[GdnPrefixKey, ...]:
evicted: List[GdnPrefixKey] = []
if self.capacity == 0:
return ()
for key in keys:
make_prefix_key(*key)
if key in self._resident:
self._resident.move_to_end(key)
else:
self._resident[key] = None
while len(self._resident) > self.capacity:
evicted_key, _ = self._resident.popitem(last=False)
evicted.append(evicted_key)
return tuple(evicted)
def forget(self, keys: Iterable[GdnPrefixKey]) -> None:
for key in keys:
self._resident.pop(key, None)

View File

@@ -0,0 +1,62 @@
#!/usr/bin/env bash
set -euo pipefail
VLLM_ROOT=${1:?usage: install_prebuilt_corex.sh VLLM_ROOT}
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUNDLE_DIR=${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10
MANIFEST=${BUNDLE_DIR}/SHA256SUMS
[[ -d "$VLLM_ROOT" ]] || {
printf 'vLLM root does not exist: %s\n' "$VLLM_ROOT" >&2
exit 2
}
[[ -f "$MANIFEST" ]] || {
printf 'prebuilt CoreX manifest is missing: %s\n' "$MANIFEST" >&2
exit 2
}
mapfile -t artifacts < <(awk '{print $2}' "$MANIFEST")
[[ "${#artifacts[@]}" -eq 16 ]] || {
printf 'expected 16 prebuilt CoreX artifacts, found %s\n' \
"${#artifacts[@]}" >&2
exit 2
}
for artifact in "${artifacts[@]}"; do
[[ ("$artifact" == corex_*.so || "$artifact" == ix_full_bridge.so) && "$artifact" != */* ]] || {
printf 'invalid prebuilt artifact name: %s\n' "$artifact" >&2
exit 2
}
done
(
cd "$BUNDLE_DIR"
sha256sum --strict --check SHA256SUMS
)
for artifact in "${artifacts[@]}"; do
install -m 0755 "$BUNDLE_DIR/$artifact" "$VLLM_ROOT/$artifact"
done
python3 - "$VLLM_ROOT" "${artifacts[@]}" <<'PY'
import pathlib
import struct
import sys
root = pathlib.Path(sys.argv[1])
for name in sys.argv[2:]:
path = root / name
if not path.is_file() or path.stat().st_size == 0:
raise SystemExit(f"installed CoreX extension is empty: {path}")
header = path.read_bytes()[:20]
if len(header) < 20 or header[:4] != b"\x7fELF":
raise SystemExit(f"installed CoreX extension is not ELF: {path}")
if header[4:6] != b"\x02\x01":
raise SystemExit(
f"installed CoreX extension is not 64-bit little-endian ELF: {path}")
machine = struct.unpack_from("<H", header, 18)[0]
if machine != 62:
raise SystemExit(
f"installed CoreX extension is not x86-64 ELF: {path} machine={machine}")
print(f"[ok] installed prebuilt CoreX extension {path}")
PY

View File

@@ -0,0 +1,196 @@
"""
ix_fused_moe.py — Fused MoE pipeline via ixformer C++ API
Replaces the entire Python expert loop in qwen3_5.py with xllm's 7-step
fused pipeline:
topk_softmax → gen_idx → expand → group_gemm → silu → group_gemm → combine
Source: upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
Bridge: ex_engine/csrc/ix_moe_bridge.cpp → ixformer::infer namespace
Loading strategy:
1. Try prebuilt ix_moe_bridge.so from known locations
2. Try JIT compile from .cpp source
3. Return unavailable (caller falls back to Python loop)
"""
import os
import logging
import importlib
import torch
logger = logging.getLogger("ix_fused_moe")
_bridge = None
_loaded = False
def _try_load_prebuilt():
"""Load prebuilt ix_moe_bridge.so without JIT compilation."""
search_paths = [
# Deployed by patch_ops.sh into vllm package
os.path.join(os.path.dirname(__file__), "ix_moe_bridge.so"),
# Prebuilt directory
os.path.join(os.path.dirname(__file__), "prebuilt",
"corex-3.2.3-ivcore10", "ix_moe_bridge.so"),
# Workspace deployment
"/workspace/qwen3_6_scripts/ix_moe_bridge.so",
]
# Also check the vllm package directory
try:
import vllm
vllm_dir = os.path.dirname(vllm.__file__)
search_paths.append(os.path.join(vllm_dir, "ix_moe_bridge.so"))
except ImportError:
pass
for path in search_paths:
if os.path.isfile(path):
try:
spec = importlib.util.spec_from_file_location(
"ix_moe_bridge", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
fns = [x for x in dir(mod) if not x.startswith("_")]
logger.info("ix_moe_bridge loaded from %s: %s", path, fns)
return mod
except Exception as e:
logger.debug("Failed to load %s: %s", path, e)
return None
def _try_jit_compile():
"""JIT compile ix_moe_bridge.cpp using torch.utils.cpp_extension."""
import glob
cpp_paths = [
os.path.join(os.path.dirname(__file__), "..", "ex_engine",
"csrc", "ix_moe_bridge.cpp"),
os.path.join(os.path.dirname(__file__), "ix_moe_bridge.cpp"),
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
"/workspace/qwen3_6_scripts/ix_moe_bridge.cpp",
]
cpp_file = None
for p in cpp_paths:
p = os.path.normpath(p)
if os.path.isfile(p):
cpp_file = p
break
if cpp_file is None:
logger.debug("ix_moe_bridge.cpp not found in search paths")
return None
# Build link flags to find ixformer symbols
extra_ldflags = []
try:
import ixformer
ixf_dir = os.path.dirname(ixformer.__file__)
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
extra_ldflags.append(so)
extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}")
except ImportError:
pass
corex_lib = "/usr/local/corex/lib64"
if os.path.isdir(corex_lib):
for lib in ["libixattn.so", "libixformer.so", "libcublas.so"]:
p = os.path.join(corex_lib, lib)
if os.path.isfile(p):
extra_ldflags.append(p)
extra_ldflags.append(f"-Wl,-rpath,{corex_lib}")
try:
from torch.utils.cpp_extension import load
logger.info("JIT compiling ix_moe_bridge from %s", cpp_file)
mod = load(
name="ix_moe_bridge",
sources=[cpp_file],
extra_cflags=["-O2", "-std=c++17"],
extra_ldflags=extra_ldflags,
verbose=False,
)
fns = [x for x in dir(mod) if not x.startswith("_")]
logger.info("ix_moe_bridge JIT compiled: %s", fns)
return mod
except Exception as e:
logger.warning("JIT compile failed: %s", e)
return None
def _ensure_loaded():
global _bridge, _loaded
if _loaded:
return _bridge is not None
_loaded = True
_bridge = _try_load_prebuilt()
if _bridge is not None:
return True
_bridge = _try_jit_compile()
if _bridge is not None:
return True
logger.info("ix_moe_bridge unavailable — MoE will use Python loop fallback")
return False
def is_available():
"""Check if the fused MoE bridge is available."""
return _ensure_loaded()
# =========================================================================
# Public API — matches xllm's 7-step pipeline
# =========================================================================
def fused_moe_forward(
hidden_states: torch.Tensor, # (T, H)
router_logits: torch.Tensor, # (T, E)
w13: torch.Tensor, # (E, 2*I, H)
w2: torch.Tensor, # (E, H, I)
topk: int,
num_experts: int,
renormalize: bool = True,
) -> torch.Tensor:
"""Full fused MoE forward — replaces _pure_pytorch_experts().
Pipeline (matching xllm/core/layers/ilu/fused_moe.cpp):
1. topk_softmax — router_logits → (weights, expert_ids)
2. moe_gen_idx — expert_ids → permutation maps
3. moe_expand_input — gather tokens by expert
4. group_gemm 1 — w13 projection (gate+up)
5. silu_and_mul — fused activation
6. group_gemm 2 — w2 projection (down)
7. combine_result — weighted scatter back
"""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.fused_moe_forward(
hidden_states, router_logits, w13, w2,
topk, num_experts, renormalize)
def topk_softmax(gating_output, topk, renormalize=True):
"""Fused topk + softmax routing."""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.topk_softmax(gating_output, topk, renormalize)
def moe_gen_idx(expert_id, expert_num):
"""Build expert permutation maps."""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.moe_gen_idx(expert_id, expert_num)
def group_gemm(inputs, weights, token_count, output_n):
"""Batched expert GEMM via ixformer."""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.group_gemm(inputs, weights, token_count, output_n)

View File

@@ -0,0 +1,227 @@
from typing import Dict, List, Optional
import torch
from vllm.attention.backends.abstract import AttentionMetadata
class MambaCacheManager:
def __init__(self, dtype, num_mamba_layers, max_batch_size,
conv_state_shape, temporal_state_shape):
conv_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
conv_state_shape,
dtype=dtype,
device="cuda")
temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
temporal_state_shape,
dtype=dtype,
device="cuda")
self.mamba_cache = (conv_state, temporal_state)
# Maps between the request id and a dict that maps between the seq_id
# and its index inside the self.mamba_cache
self.mamba_cache_indices_mapping: Dict[str, Dict[int, int]] = {}
def current_run_tensors(self, input_ids: torch.Tensor,
attn_metadata: AttentionMetadata, **kwargs):
"""
Return the tensors for the current run's conv and ssm state.
"""
if "seqlen_agnostic_capture_inputs" not in kwargs:
# We get here only on Prefill/Eager mode runs
request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]
finished_requests_ids = kwargs["finished_requests_ids"]
self._release_finished_requests(finished_requests_ids)
mamba_cache_tensors = self._prepare_current_run_mamba_cache(
request_ids_to_seq_ids, finished_requests_ids)
else:
# CUDA graph capturing runs
mamba_cache_tensors = kwargs["seqlen_agnostic_capture_inputs"]
return mamba_cache_tensors
def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
"""
Copy the relevant Mamba cache into the CUDA graph input buffer
that was provided during the capture runs
(JambaForCausalLM.mamba_gc_cache_buffer).
"""
assert all(
key in kwargs
for key in ["request_ids_to_seq_ids", "finished_requests_ids"])
finished_requests_ids = kwargs["finished_requests_ids"]
request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]
self._release_finished_requests(finished_requests_ids)
self._prepare_current_run_mamba_cache(request_ids_to_seq_ids,
finished_requests_ids)
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
"""
Provide the CUDA graph capture runs with a buffer in adjusted size.
The buffer is used to maintain the Mamba Cache during the CUDA graph
replay runs.
"""
return tuple(buffer[:, :batch_size] for buffer in self.mamba_cache)
def _swap_mamba_cache(self, from_index: int, to_index: int):
assert len(self.mamba_cache) > 0
for cache_t in self.mamba_cache:
cache_t[:, [to_index,from_index]] = \
cache_t[:, [from_index,to_index]]
def _copy_mamba_cache(self, from_index: int, to_index: int):
assert len(self.mamba_cache) > 0
for cache_t in self.mamba_cache:
cache_t[:, to_index].copy_(cache_t[:, from_index],
non_blocking=True)
def _move_out_if_already_occupied(self, index: int,
all_occupied_indices: List[int]):
if index in all_occupied_indices:
first_free_index = self._first_free_index_in_mamba_cache()
# In case occupied, move the occupied to a new empty block
self._move_cache_index_and_mappings(from_index=index,
to_index=first_free_index)
def _assign_seq_id_to_mamba_cache_in_specific_dest(self, cur_rid: str,
seq_id: int,
destination_index: int):
"""
Assign (req_id,seq_id) pair to a `destination_index` index, if
already occupied, move the occupying index to a free index.
"""
all_occupied_indices = self._get_all_occupied_indices()
if cur_rid not in self.mamba_cache_indices_mapping:
self._move_out_if_already_occupied(
index=destination_index,
all_occupied_indices=all_occupied_indices)
for cache_t in self.mamba_cache:
cache_t[:, destination_index].zero_()
self.mamba_cache_indices_mapping[cur_rid] = {
seq_id: destination_index
}
elif seq_id not in (seq_ids2indices :=
self.mamba_cache_indices_mapping[cur_rid]):
# parallel sampling , where n > 1, assume prefill have
# already happened now we only need to copy the already
# existing cache into the siblings seq_ids caches
self._move_out_if_already_occupied(
index=destination_index,
all_occupied_indices=all_occupied_indices)
index_exists = list(seq_ids2indices.values())[0]
# case of decoding n>1, copy prefill cache to decoding indices
self._copy_mamba_cache(from_index=index_exists,
to_index=destination_index)
self.mamba_cache_indices_mapping[cur_rid][
seq_id] = destination_index
else:
# already exists
cache_index_already_exists = self.mamba_cache_indices_mapping[
cur_rid][seq_id]
if cache_index_already_exists != destination_index:
# In case the seq id already exists but not in
# the right destination, swap it with what's occupying it
self._swap_pair_indices_and_mappings(
from_index=cache_index_already_exists,
to_index=destination_index)
def _prepare_current_run_mamba_cache(
self, request_ids_to_seq_ids: Dict[str, list[int]],
finished_requests_ids: List[str]):
running_indices = []
request_ids_to_seq_ids_flatten = [
(req_id, seq_id)
for req_id, seq_ids in request_ids_to_seq_ids.items()
for seq_id in seq_ids
]
batch_size = len(request_ids_to_seq_ids_flatten)
for dest_index, (request_id,
seq_id) in enumerate(request_ids_to_seq_ids_flatten):
if request_id in finished_requests_ids:
# Do not allocate cache index for requests that run
# and finish right after
continue
self._assign_seq_id_to_mamba_cache_in_specific_dest(
request_id, seq_id, dest_index)
running_indices.append(dest_index)
self._clean_up_first_bs_blocks(batch_size, running_indices)
conv_state = self.mamba_cache[0][:, :batch_size]
temporal_state = self.mamba_cache[1][:, :batch_size]
return (conv_state, temporal_state)
def _get_all_occupied_indices(self):
return [
cache_idx
for seq_ids2indices in self.mamba_cache_indices_mapping.values()
for cache_idx in seq_ids2indices.values()
]
def _clean_up_first_bs_blocks(self, batch_size: int,
indices_for_current_run: List[int]):
# move out all of the occupied but currently not running blocks
# outside of the first n blocks
destination_indices = range(batch_size)
max_possible_batch_size = self.mamba_cache[0].shape[1]
for destination_index in destination_indices:
if destination_index in self._get_all_occupied_indices() and \
destination_index not in indices_for_current_run:
# move not running indices outside of the batch
all_other_indices = list(
range(batch_size, max_possible_batch_size))
first_avail_index = self._first_free_index_in_mamba_cache(
all_other_indices)
self._swap_indices(from_index=destination_index,
to_index=first_avail_index)
def _move_cache_index_and_mappings(self, from_index: int, to_index: int):
self._copy_mamba_cache(from_index=from_index, to_index=to_index)
self._update_mapping_index(from_index=from_index, to_index=to_index)
def _swap_pair_indices_and_mappings(self, from_index: int, to_index: int):
self._swap_mamba_cache(from_index=from_index, to_index=to_index)
self._swap_mapping_index(from_index=from_index, to_index=to_index)
def _swap_mapping_index(self, from_index: int, to_index: int):
for seq_ids2index in self.mamba_cache_indices_mapping.values():
for seq_id, index in seq_ids2index.items():
if from_index == index:
seq_ids2index.update({seq_id: to_index})
elif to_index == index:
seq_ids2index.update({seq_id: from_index})
def _update_mapping_index(self, from_index: int, to_index: int):
for seq_ids2index in self.mamba_cache_indices_mapping.values():
for seq_id, index in seq_ids2index.items():
if from_index == index:
seq_ids2index.update({seq_id: to_index})
return
def _release_finished_requests(self,
finished_seq_groups_req_ids: List[str]):
for req_id in finished_seq_groups_req_ids:
if req_id in self.mamba_cache_indices_mapping:
seq_mapping = self.mamba_cache_indices_mapping.pop(req_id)
for cache_idx in seq_mapping.values():
for cache_t in self.mamba_cache:
cache_t[:, cache_idx].zero_()
def _first_free_index_in_mamba_cache(
self, indices_range: Optional[List[int]] = None) -> int:
assert self.mamba_cache is not None
if indices_range is None:
max_possible_batch_size = self.mamba_cache[0].shape[1]
indices_range = list(range(max_possible_batch_size))
all_occupied_indices = self._get_all_occupied_indices()
for i in indices_range:
if i not in all_occupied_indices:
return i
raise Exception("Couldn't find a free spot in the mamba cache! This"
"should never happen")

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,91 @@
from patch_utils import package_root, replace_once
CACHE_ENGINE = package_root("vllm") / "worker" / "cache_engine.py"
IMPORT_ANCHOR = """\
from vllm.logger import init_logger
"""
IMPORT_REPLACEMENT = """\
from vllm.block_major_kv_cache import (
BlockMajorCpuKVCache,
block_major_cpu_kv_enabled,
)
from vllm.logger import init_logger
"""
ALLOCATION_ANCHOR = """\
self.gpu_cache = self._allocate_kv_cache(
self.num_gpu_blocks, self.device_config.device_type)
self.cpu_cache = self._allocate_kv_cache(self.num_cpu_blocks, "cpu")
"""
ALLOCATION_REPLACEMENT = """\
self.gpu_cache = self._allocate_kv_cache(
self.num_gpu_blocks, self.device_config.device_type)
self._bi100_block_major_cpu_kv = None
if block_major_cpu_kv_enabled():
self._bi100_block_major_cpu_kv = BlockMajorCpuKVCache(
self.gpu_cache,
self.num_cpu_blocks,
pin_memory=is_pin_memory_available(),
)
self.cpu_cache = self._bi100_block_major_cpu_kv.layer_views
else:
self.cpu_cache = self._allocate_kv_cache(
self.num_cpu_blocks, "cpu")
"""
SWAP_ANCHOR = """\
def swap_in(self, src_to_dst: torch.Tensor) -> None:
for i in range(self.num_attention_layers):
self.attn_backend.swap_blocks(self.cpu_cache[i], self.gpu_cache[i],
src_to_dst)
def swap_out(self, src_to_dst: torch.Tensor) -> None:
for i in range(self.num_attention_layers):
self.attn_backend.swap_blocks(self.gpu_cache[i], self.cpu_cache[i],
src_to_dst)
"""
SWAP_REPLACEMENT = """\
def swap_in(self, src_to_dst: torch.Tensor) -> None:
if self._bi100_block_major_cpu_kv is not None:
self._bi100_block_major_cpu_kv.swap_in(src_to_dst)
return
for i in range(self.num_attention_layers):
self.attn_backend.swap_blocks(self.cpu_cache[i], self.gpu_cache[i],
src_to_dst)
def swap_out(self, src_to_dst: torch.Tensor) -> None:
if self._bi100_block_major_cpu_kv is not None:
self._bi100_block_major_cpu_kv.swap_out(src_to_dst)
return
for i in range(self.num_attention_layers):
self.attn_backend.swap_blocks(self.gpu_cache[i], self.cpu_cache[i],
src_to_dst)
"""
replace_once(
CACHE_ENGINE,
IMPORT_ANCHOR,
IMPORT_REPLACEMENT,
required=True,
already_contains="from vllm.block_major_kv_cache import",
)
replace_once(
CACHE_ENGINE,
ALLOCATION_ANCHOR,
ALLOCATION_REPLACEMENT,
required=True,
already_contains="self._bi100_block_major_cpu_kv = None",
)
replace_once(
CACHE_ENGINE,
SWAP_ANCHOR,
SWAP_REPLACEMENT,
required=True,
already_contains="self._bi100_block_major_cpu_kv.swap_in",
)

View File

@@ -0,0 +1,46 @@
from patch_utils import package_root, replace_once
WORKER = package_root("vllm") / "worker" / "worker.py"
IMPORT_ANCHOR = """\
from vllm.logger import init_logger
"""
IMPORT_REPLACEMENT = """\
from vllm.block_major_kv_cache import reserve_block_major_gpu_blocks
from vllm.logger import init_logger
"""
CAPACITY_ANCHOR = """\
num_gpu_blocks = max(num_gpu_blocks, 0)
num_cpu_blocks = max(num_cpu_blocks, 0)
"""
CAPACITY_REPLACEMENT = """\
num_gpu_blocks = reserve_block_major_gpu_blocks(
num_gpu_blocks, cache_block_size)
num_gpu_blocks = max(num_gpu_blocks, 0)
num_cpu_blocks = max(num_cpu_blocks, 0)
"""
replace_once(
WORKER,
IMPORT_ANCHOR,
IMPORT_REPLACEMENT,
required=True,
already_contains=(
"from vllm.block_major_kv_cache import "
"reserve_block_major_gpu_blocks"
),
)
replace_once(
WORKER,
CAPACITY_ANCHOR,
CAPACITY_REPLACEMENT,
required=True,
already_contains=(
"num_gpu_blocks = reserve_block_major_gpu_blocks("
),
)

View File

@@ -0,0 +1,210 @@
"""Install the optional BI100 prefix-cache diagnostic trace."""
from patch_utils import package_root, replace_once, replace_one_of
VLLM_ROOT = package_root("vllm")
TARGET = VLLM_ROOT / "core" / "block_manager_v2.py"
OUTPUTS_TARGET = VLLM_ROOT / "outputs.py"
HELPER = '''
def _bi100_capture_cache_trace(self, seq_group, seq, block_table) -> None:
if os.getenv("BI100_CACHE_TRACE", "0") != "1":
return
session = getattr(self, "_bi100_trace_session", None)
if session is None:
session = hashlib.sha256(os.urandom(16)).hexdigest()[:16]
self._bi100_trace_session = session
self._bi100_trace_ordinal = getattr(self, "_bi100_trace_ordinal", 0) + 1
request_id_sha256 = hashlib.sha256(
str(seq_group.request_id).encode("utf-8")).hexdigest()[:16]
prompt_tokens = len(seq.get_token_ids())
requests = getattr(self, "_bi100_trace_requests", None)
if requests is None:
requests = {}
self._bi100_trace_requests = requests
requests[seq.seq_id] = {
"version": 4,
"trace_session_sha256": session,
"ordinal": self._bi100_trace_ordinal,
"request_id_sha256": request_id_sha256,
"prompt_tokens": prompt_tokens,
"prompt_allocated_blocks": (
(prompt_tokens + self.block_size - 1) // self.block_size
),
"block_size": self.block_size,
"capacity_blocks": self.num_total_gpu_blocks,
}
setattr(seq_group, "_bi100_cache_trace_seq_id", seq.seq_id)
setattr(seq_group, "_bi100_cache_trace_emit",
self._bi100_emit_cache_trace)
def _bi100_update_cache_trace(
self, seq, raw_kv_hit_blocks, restore_key, capture_actions,
evict_keys, policy) -> None:
if os.getenv("BI100_CACHE_TRACE", "0") != "1":
return
requests = getattr(self, "_bi100_trace_requests", None)
if not requests or seq.seq_id not in requests:
return
record = requests[seq.seq_id]
record["gdn_policy"] = policy
if "initial_raw_kv_contiguous_hit_blocks" not in record:
record["initial_raw_kv_contiguous_hit_blocks"] = max(
0, int(raw_kv_hit_blocks))
record["gdn_restore_digest_base64"] = (
base64.b64encode(restore_key[1]).decode("ascii")
if restore_key is not None else None)
record["raw_kv_contiguous_hit_blocks"] = max(
int(raw_kv_hit_blocks),
int(record.get("raw_kv_contiguous_hit_blocks", 0)))
effective_blocks = int(restore_key[0]) if restore_key is not None else 0
record["effective_gdn_hit_blocks"] = max(
effective_blocks, int(record.get("effective_gdn_hit_blocks", 0)))
admissions = record.setdefault("gdn_admissions", [])
for key, reason in capture_actions:
admissions.append({
"block_count": int(key[0]),
"digest_base64": base64.b64encode(key[1]).decode("ascii"),
"reason": str(reason),
})
evictions = record.setdefault("gdn_evictions", [])
for key in evict_keys:
evictions.append({
"block_count": int(key[0]),
"digest_base64": base64.b64encode(key[1]).decode("ascii"),
"reason": "capacity_lru",
})
def _bi100_finalize_cache_trace(self, seq, block_table) -> None:
if os.getenv("BI100_CACHE_TRACE", "0") != "1":
return
requests = getattr(self, "_bi100_trace_requests", None)
if not requests:
return
record = requests.get(seq.seq_id)
if record is None:
return
total_tokens = len(seq.get_token_ids())
block_hashes = block_table.get_content_hashes()
for block_hash in block_hashes:
if not isinstance(block_hash, bytes) or len(block_hash) != 32:
raise RuntimeError(
"BI100 cache trace requires 32-byte content hashes")
full_blocks = len(block_hashes)
record.update({
"total_tokens": total_tokens,
"allocated_blocks": (
(total_tokens + self.block_size - 1) // self.block_size
),
"full_blocks": full_blocks,
"hash_encoding": "sha256_base64",
"block_hashes": base64.b64encode(b"".join(block_hashes)).decode("ascii"),
"_finalized": True,
})
generated_tokens = max(0, total_tokens - record["prompt_tokens"])
record["generated_tokens"] = generated_tokens
def _bi100_emit_cache_trace(self, seq_group) -> None:
if os.getenv("BI100_CACHE_TRACE", "0") != "1":
return
seq_id = getattr(seq_group, "_bi100_cache_trace_seq_id", None)
requests = getattr(self, "_bi100_trace_requests", None)
if seq_id is None or not requests:
return
record = requests.pop(seq_id, None)
if record is None:
return
if record.pop("_finalized", False) is not True:
raise RuntimeError(
"BI100 cache trace emitted before block finalization")
metrics = getattr(seq_group, "metrics", None)
arrival = getattr(metrics, "arrival_time", None)
first_token = getattr(metrics, "first_token_time", None)
finished = getattr(metrics, "finished_time", None)
queue = getattr(metrics, "time_in_queue", None)
cached = getattr(metrics, "num_cached_tokens", None)
if any(value is None for value in (
arrival, first_token, finished, queue)):
raise RuntimeError(
"BI100 cache trace requires finalized request metrics")
record["ttft_s"] = max(0.0, float(first_token - arrival))
record["request_latency_s"] = max(
0.0, float(finished - arrival))
record["time_in_queue_s"] = max(0.0, float(queue))
record["observed_effective_cached_tokens"] = max(
0, int(cached or 0))
ttft_s = record["ttft_s"]
if ttft_s > 0:
record["observed_input_tps"] = record["prompt_tokens"] / ttft_s
generated_tokens = record["generated_tokens"]
if generated_tokens > 1:
decode_s = finished - first_token
if decode_s > 0:
record["observed_output_tps"] = (
(generated_tokens - 1) / decode_s)
print("[BI100_CACHE_TRACE] " + json.dumps(record, separators=(",", ":"),
sort_keys=True), flush=True)
'''
def main():
replace_once(TARGET, "from collections.abc import Mapping\n",
"from collections.abc import Mapping\nimport base64\nimport json\nimport os\n",
required=True, already_contains="import base64\n")
replace_once(TARGET, "class BlockSpaceManagerV2(BlockSpaceManager):\n",
"class BlockSpaceManagerV2(BlockSpaceManager):\n" + HELPER,
required=True, already_contains="def _bi100_capture_cache_trace(")
replace_once(TARGET,
" self.block_tables[seq.seq_id] = block_table\n\n # Track seq",
" self.block_tables[seq.seq_id] = block_table\n self._bi100_capture_cache_trace(\n seq_group, seq, block_table)\n\n # Track seq",
required=True,
already_contains="self.block_tables[seq.seq_id] = block_table\n"
" self._bi100_capture_cache_trace(")
replacements = []
for table_key in ("seq_id", "seq.seq_id"):
prefix = (
" self._last_access_blocks_tracker."
"update_seq_blocks_last_access(\n"
f" seq_id, self.block_tables[{table_key}]."
"physical_block_ids)\n")
replacements.append((
prefix + "\n # Untrack seq",
prefix + " self._bi100_finalize_cache_trace(\n"
f" seq, self.block_tables[{table_key}])\n\n"
" # Untrack seq",
))
replace_one_of(
TARGET,
replacements,
required=True,
already_contains=" self._bi100_finalize_cache_trace(\n"
" seq, self.block_tables[")
replace_once(
OUTPUTS_TARGET,
" seq_group.set_finished_time(finished_time)\n\n"
" init_args = (seq_group.request_id, prompt, prompt_token_ids,\n",
" seq_group.set_finished_time(finished_time)\n"
" if finished_time is not None:\n"
" cache_trace_emit = getattr(\n"
" seq_group, \"_bi100_cache_trace_emit\", None)\n"
" if callable(cache_trace_emit):\n"
" cache_trace_emit(seq_group)\n"
" delattr(seq_group, \"_bi100_cache_trace_emit\")\n"
" delattr(seq_group, \"_bi100_cache_trace_seq_id\")\n\n"
" init_args = (seq_group.request_id, prompt, prompt_token_ids,\n",
required=True,
already_contains="if finished_time is not None:\n"
" cache_trace_emit = getattr(\n",
)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,65 @@
#!/usr/bin/env python3
"""
patch_chat_template.py — Fix non-thinking mode '!!!!!' output
Root cause: Qwen3.5/3.6-MoE's chat_template adds '<think>\n\n</think>\n\n'
when enable_thinking=false. This empty think block causes the model to
degenerate into outputting nothing but '!'.
Fix: Remove the empty think block so the model generates directly.
Verified on real machine 2026-08-20:
- Model: /root/public-storage/models/Qwen/Qwen3.6-35B-A3B
- target repr: "{{- '<think>\\n\\n</think>\\n\\n' }}"
- match: True → patch applied successfully
"""
import json
import sys
import os
def patch_tokenizer_config(model_path: str) -> bool:
config_path = os.path.join(model_path, "tokenizer_config.json")
if not os.path.isfile(config_path):
print(f"[patch_chat_template] {config_path} not found")
return False
with open(config_path, "r") as f:
config = json.load(f)
template = config.get("chat_template", "")
if not template:
print("[patch_chat_template] No chat_template found")
return False
# The Jinja template contains literal backslash-n sequences: \n
# After json.load these remain as two-char sequences (backslash + n),
# NOT real newlines. Use raw string so Python doesn't interpret them.
target = r"{{- '<think>\n\n</think>\n\n' }}"
replacement = "{{- '' }}"
if target in template:
template = template.replace(target, replacement, 1)
config["chat_template"] = template
with open(config_path, "w") as f:
json.dump(config, f, ensure_ascii=False, indent=2)
print("[patch_chat_template] ✓ Removed empty <think></think> block for non-thinking mode")
return True
# Fallback: check if already patched
if "enable_thinking is false" in template and target not in template:
print("[patch_chat_template] Already patched or different format")
return True
print(f"[patch_chat_template] WARNING: Could not find target pattern")
# Debug: show what's actually there
idx = template.find("enable_thinking is false")
if idx >= 0:
print(f"[patch_chat_template] Context: {repr(template[idx:idx+150])}")
return False
if __name__ == "__main__":
model_path = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("MODEL_PATH", "/model")
success = patch_tokenizer_config(model_path)
sys.exit(0 if success else 1)

View File

@@ -0,0 +1,65 @@
from patch_utils import package_root, replace_once
CUSTOM_OPS = package_root("vllm") / "_custom_ops.py"
CLEAN_BLOCK = """\
def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
block_mapping: torch.Tensor) -> None:
ixf_F.swap_blocks(src, dst, block_mapping)
"""
COMPATIBLE_BLOCK = """\
def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
block_mapping: torch.Tensor) -> None:
# BI100 CoreX 3.2.3 exposes vllm_swap_blocks, while this vLLM build calls
# the newer swap_blocks name. Normalize the worker's CPU int64 [N, 2]
# tensor only for the legacy public API and fail fast on malformed maps.
native_swap_blocks = getattr(ixf_F, "swap_blocks", None)
if native_swap_blocks is not None:
native_swap_blocks(src, dst, block_mapping)
return
vendor_swap_blocks = getattr(ixf_F, "vllm_swap_blocks", None)
if vendor_swap_blocks is None:
raise RuntimeError(
"ixformer exposes neither swap_blocks nor vllm_swap_blocks")
if isinstance(block_mapping, torch.Tensor):
if block_mapping.device.type != "cpu":
raise ValueError("swap block mapping must be a CPU tensor")
if block_mapping.dtype != torch.int64:
raise ValueError("swap block mapping must use torch.int64")
if block_mapping.dim() != 2 or block_mapping.shape[1] != 2:
raise ValueError("swap block mapping must have shape [N, 2]")
pairs = block_mapping.tolist()
elif isinstance(block_mapping, dict):
pairs = list(block_mapping.items())
else:
raise TypeError("swap block mapping must be a tensor or dict")
normalized_mapping = {}
destinations = set()
for source, destination in pairs:
source = int(source)
destination = int(destination)
if source < 0 or destination < 0:
raise ValueError("swap block indices must be non-negative")
if source in normalized_mapping:
raise ValueError(f"duplicate swap source block: {source}")
if destination in destinations:
raise ValueError(
f"duplicate swap destination block: {destination}")
normalized_mapping[source] = destination
destinations.add(destination)
vendor_swap_blocks(src, dst, normalized_mapping)
"""
replace_once(
CUSTOM_OPS,
CLEAN_BLOCK,
COMPATIBLE_BLOCK,
required=True,
already_contains="BI100 CoreX 3.2.3 exposes vllm_swap_blocks",
)

View File

@@ -0,0 +1,61 @@
from patch_utils import package_root, replace_once
VLLM_ROOT = package_root("vllm")
MULTIPROC_GPU_EXECUTOR = VLLM_ROOT / "executor" / "multiproc_gpu_executor.py"
MULTIPROC_WORKER_UTILS = VLLM_ROOT / "executor" / "multiproc_worker_utils.py"
def ensure_import_os(path):
text = path.read_text()
if "import os\n" in text:
print(f"[skip] import os already present: {path}")
return
for anchor in ("import time\n", "import signal\n", "import sys\n"):
if anchor in text:
replace_once(
path,
anchor,
anchor + "import os\n",
required=True,
already_contains="import os\n",
)
return
raise RuntimeError(f"no import anchor found for os in {path}")
ensure_import_os(MULTIPROC_GPU_EXECUTOR)
ensure_import_os(MULTIPROC_WORKER_UTILS)
replace_once(
MULTIPROC_GPU_EXECUTOR,
"""logger = init_logger(__name__)\n""",
"""logger = init_logger(__name__)\n\n\ndef _bi100_startup_debug(message: str, *args) -> None:\n if os.getenv(\"BI100_EXECUTOR_STARTUP_DEBUG\") == \"1\":\n logger.info(\"[BI100 startup] \" + message, *args)\n""",
required=True,
already_contains="def _bi100_startup_debug(",
)
replace_once(
MULTIPROC_GPU_EXECUTOR,
""" self.driver_worker = self._create_worker(\n distributed_init_method=distributed_init_method)\n self._run_workers(\"init_device\")\n self._run_workers(\"load_model\",\n max_concurrent_workers=self.parallel_config.\n max_parallel_loading_workers)\n""",
""" _bi100_startup_debug(\"creating driver worker\")\n self.driver_worker = self._create_worker(\n distributed_init_method=distributed_init_method)\n _bi100_startup_debug(\"created driver worker\")\n _bi100_startup_debug(\"starting init_device\")\n self._run_workers(\"init_device\")\n _bi100_startup_debug(\"finished init_device\")\n _bi100_startup_debug(\"starting load_model\")\n self._run_workers(\"load_model\",\n max_concurrent_workers=self.parallel_config.\n max_parallel_loading_workers)\n _bi100_startup_debug(\"finished load_model\")\n""",
required=True,
already_contains='_bi100_startup_debug("starting init_device")',
)
replace_once(
MULTIPROC_GPU_EXECUTOR,
""" # Start all remote workers first.\n worker_outputs = [\n worker.execute_method(method, *args, **kwargs)\n for worker in self.workers\n ]\n\n driver_worker_method = getattr(self.driver_worker, method)\n driver_worker_output = driver_worker_method(*args, **kwargs)\n\n # Get the results of the workers.\n return [driver_worker_output\n ] + [output.get() for output in worker_outputs]\n""",
""" _bi100_startup_debug(\"enqueue remote method=%s workers=%d\", method,\n len(self.workers))\n # Start all remote workers first.\n worker_outputs = [\n worker.execute_method(method, *args, **kwargs)\n for worker in self.workers\n ]\n _bi100_startup_debug(\"remote enqueued method=%s\", method)\n\n driver_worker_method = getattr(self.driver_worker, method)\n _bi100_startup_debug(\"driver start method=%s\", method)\n driver_worker_output = driver_worker_method(*args, **kwargs)\n _bi100_startup_debug(\"driver done method=%s\", method)\n\n # Get the results of the workers.\n _bi100_startup_debug(\"waiting remote results method=%s\", method)\n remote_outputs = [output.get() for output in worker_outputs]\n _bi100_startup_debug(\"remote done method=%s\", method)\n return [driver_worker_output] + remote_outputs\n""",
required=True,
already_contains='_bi100_startup_debug("enqueue remote method=%s workers=%d"',
)
replace_once(
MULTIPROC_WORKER_UTILS,
""" task_id, method, args, kwargs = items\n try:\n executor = getattr(worker, method)\n output = executor(*args, **kwargs)\n except SystemExit:\n""",
""" task_id, method, args, kwargs = items\n if os.getenv(\"BI100_EXECUTOR_STARTUP_DEBUG\") == \"1\":\n logger.info(\"[BI100 worker] start method=%s\", method)\n try:\n executor = getattr(worker, method)\n output = executor(*args, **kwargs)\n if os.getenv(\"BI100_EXECUTOR_STARTUP_DEBUG\") == \"1\":\n logger.info(\"[BI100 worker] done method=%s\", method)\n except SystemExit:\n""",
required=True,
already_contains='logger.info("[BI100 worker] start method=%s", method)',
)

View File

@@ -0,0 +1,408 @@
"""Patch vLLM 0.6.3 prefix-cache and MRoPE chunk alignment bugs."""
from __future__ import annotations
import pathlib
from patch_utils import package_root, replace_once
HELPER_ANCHOR = """\
logger = init_logger(__name__)
LORA_WARMUP_RANK = 8"""
HELPER_REPLACEMENT = """\
logger = init_logger(__name__)
def _slice_mrope_positions(positions, start, stop, expected_len):
if positions is None or len(positions) != 3:
raise RuntimeError("MRoPE positions must contain three axes")
sliced = [axis[start:stop] for axis in positions]
lengths = [len(axis) for axis in sliced]
if lengths != [expected_len] * 3:
raise RuntimeError(
"MRoPE/input token length mismatch after chunk alignment: "
f"positions={lengths}, input_tokens={expected_len}, "
f"slice=({start}, {stop})")
return sliced
LORA_WARMUP_RANK = 8"""
PREFIX_PAST_ANCHOR = """\
if prefix_cache_len <= context_len:
# We already passed the cache hit region,
# so do normal computation.
pass"""
PREFIX_PAST_REPLACEMENT = """\
if prefix_cache_len <= context_len:
# We already passed the cache hit region,
# so do normal computation.
# Must clear prefix_cache_hit so _add_seq_group uses the full
# block_tables (prefix + previous-chunk blocks) instead of only
# computed_block_nums (prefix only). Without this, block_tables
# passed to _forward_prefix_pytorch is too narrow for context_len,
# causing an empty blk_ids slice and a zero-dim amax() crash.
inter_data.prefix_cache_hit = False"""
PARTIAL_HIT_ANCHOR = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][uncomputed_start:]
context_len = prefix_cache_len
inter_data.context_lens[seq_idx] = context_len
inter_data.query_lens[
seq_idx] = inter_data.seq_lens[seq_idx] - context_len"""
PARTIAL_HIT_REPLACEMENT = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][uncomputed_start:]
context_len = prefix_cache_len
inter_data.context_lens[seq_idx] = context_len
inter_data.query_lens[
seq_idx] = inter_data.seq_lens[seq_idx] - context_len
if inter_data.mrope_input_positions is not None:
positions = inter_data.mrope_input_positions[seq_idx]
if positions is not None:
inter_data.mrope_input_positions[seq_idx] = \\
_slice_mrope_positions(
positions, uncomputed_start, None,
inter_data.query_lens[seq_idx])"""
FULL_HIT_ANCHOR = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][-1:]
inter_data.query_lens[seq_idx] = 1
inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1"""
FULL_HIT_REPLACEMENT = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][-1:]
inter_data.query_lens[seq_idx] = 1
inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1
if inter_data.mrope_input_positions is not None:
positions = inter_data.mrope_input_positions[seq_idx]
if positions is not None:
inter_data.mrope_input_positions[seq_idx] = \\
_slice_mrope_positions(positions, -1, None, 1)"""
MULTIMODAL_MROPE_ANCHOR = """\
mrope_input_positions, mrope_position_delta = \\
MRotaryEmbedding.get_input_positions(
token_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
image_token_id=hf_config.image_token_id,
video_token_id=hf_config.video_token_id,
vision_start_token_id=hf_config.vision_start_token_id,
vision_end_token_id=hf_config.vision_end_token_id,
spatial_merge_size=hf_config.vision_config.
spatial_merge_size,
context_len=inter_data.context_lens[seq_idx],
)
seq_data.mrope_position_delta = mrope_position_delta
inter_data.mrope_input_positions[
seq_idx] = mrope_input_positions"""
MULTIMODAL_MROPE_REPLACEMENT = """\
# vLLM 0.6.3 returns positions through the end of token_ids,
# while chunked prefill sends only [context_len:seq_len].
# Compute the full MRoPE map once so the delta remains tied to
# the complete request, then select exactly the physical query.
mrope_input_positions, mrope_position_delta = \\
MRotaryEmbedding.get_input_positions(
token_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
image_token_id=hf_config.image_token_id,
video_token_id=hf_config.video_token_id,
vision_start_token_id=hf_config.vision_start_token_id,
vision_end_token_id=hf_config.vision_end_token_id,
spatial_merge_size=hf_config.vision_config.
spatial_merge_size,
context_len=0,
)
mrope_input_positions = _slice_mrope_positions(
mrope_input_positions,
inter_data.context_lens[seq_idx],
inter_data.seq_lens[seq_idx],
len(inter_data.input_tokens[seq_idx]))
seq_data.mrope_position_delta = mrope_position_delta
inter_data.mrope_input_positions[
seq_idx] = mrope_input_positions"""
MODEL_INPUT_FIELDS_ANCHOR = """\
multi_modal_kwargs: Optional[BatchedTensorInputs] = None
request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None"""
MODEL_INPUT_FIELDS_REPLACEMENT = """\
multi_modal_kwargs: Optional[BatchedTensorInputs] = None
# BI100 scheduler-owned GDN prefix-cache actions. These plain Python
# objects are included in the multiprocess model-input broadcast.
gdn_restore_key: Optional[Tuple[int, bytes]] = None
gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
gdn_segment_offsets: Optional[List[int]] = None
request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None"""
BASE_BROADCAST_ANCHOR = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
return tensor_dict
@classmethod"""
BASE_BROADCAST_REPLACEMENT = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"gdn_restore_key\": self.gdn_restore_key,
\"gdn_capture_points\": self.gdn_capture_points,
\"gdn_evict_keys\": self.gdn_evict_keys,
\"gdn_segment_offsets\": self.gdn_segment_offsets,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
return tensor_dict
@classmethod"""
SAMPLING_BROADCAST_ANCHOR = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
_add_sampling_metadata_broadcastable_dict(tensor_dict,
self.sampling_metadata)"""
SAMPLING_BROADCAST_REPLACEMENT = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"gdn_restore_key\": self.gdn_restore_key,
\"gdn_capture_points\": self.gdn_capture_points,
\"gdn_evict_keys\": self.gdn_evict_keys,
\"gdn_segment_offsets\": self.gdn_segment_offsets,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
_add_sampling_metadata_broadcastable_dict(tensor_dict,
self.sampling_metadata)"""
BUILDER_INIT_ANCHOR = """\
self.finished_requests_ids = finished_requests_ids
self.decode_only = True
# Intermediate data"""
BUILDER_INIT_REPLACEMENT = """\
self.finished_requests_ids = finished_requests_ids
self.decode_only = True
self.gdn_restore_key = None
self.gdn_capture_points = None
self.gdn_evict_keys = None
self.gdn_segment_offsets = None
# Intermediate data"""
ADD_SEQ_GROUP_ANCHOR = """\
def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
\"\"\"Add a sequence group to the builder.\"\"\"
seq_ids = seq_group_metadata.seq_data.keys()"""
ADD_SEQ_GROUP_REPLACEMENT = """\
def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
\"\"\"Add a sequence group to the builder.\"\"\"
gdn_actions = (
seq_group_metadata.gdn_restore_key,
seq_group_metadata.gdn_capture_points,
seq_group_metadata.gdn_evict_keys,
seq_group_metadata.gdn_segment_offsets,
)
if any(value is not None for value in gdn_actions):
if not seq_group_metadata.is_prompt:
raise RuntimeError(\"GDN prefix-cache actions require prefill\")
if any(value is not None for value in (
self.gdn_restore_key, self.gdn_capture_points,
self.gdn_evict_keys, self.gdn_segment_offsets)):
raise RuntimeError(
\"only one GDN prefix-cache action group is supported\")
(self.gdn_restore_key, self.gdn_capture_points,
self.gdn_evict_keys, self.gdn_segment_offsets) = gdn_actions
seq_ids = seq_group_metadata.seq_data.keys()"""
BUILD_RESULT_ANCHOR = """\
lora_mapping=lora_mapping,
lora_requests=lora_requests,
multi_modal_kwargs=multi_modal_kwargs,
request_ids_to_seq_ids=request_ids_to_seq_ids,"""
BUILD_RESULT_REPLACEMENT = """\
lora_mapping=lora_mapping,
lora_requests=lora_requests,
multi_modal_kwargs=multi_modal_kwargs,
gdn_restore_key=self.gdn_restore_key,
gdn_capture_points=self.gdn_capture_points,
gdn_evict_keys=self.gdn_evict_keys,
gdn_segment_offsets=self.gdn_segment_offsets,
request_ids_to_seq_ids=request_ids_to_seq_ids,"""
EXECUTE_KWARGS_ANCHOR = """\
seqlen_agnostic_kwargs = {
\"finished_requests_ids\": model_input.finished_requests_ids,
\"request_ids_to_seq_ids\": model_input.request_ids_to_seq_ids,
} if self.has_inner_state else {}
if (self.observability_config is not None"""
EXECUTE_KWARGS_REPLACEMENT = """\
seqlen_agnostic_kwargs = {
\"finished_requests_ids\": model_input.finished_requests_ids,
\"request_ids_to_seq_ids\": model_input.request_ids_to_seq_ids,
} if self.has_inner_state else {}
gdn_prefix_kwargs = {}
if model_input.gdn_restore_key is not None:
gdn_prefix_kwargs[\"gdn_restore_key\"] = model_input.gdn_restore_key
if model_input.gdn_capture_points is not None:
gdn_prefix_kwargs[\"gdn_capture_points\"] = (
model_input.gdn_capture_points)
if model_input.gdn_evict_keys is not None:
gdn_prefix_kwargs[\"gdn_evict_keys\"] = model_input.gdn_evict_keys
if model_input.gdn_segment_offsets is not None:
gdn_prefix_kwargs[\"gdn_segment_offsets\"] = (
model_input.gdn_segment_offsets)
if (self.observability_config is not None"""
MODEL_CALL_ANCHOR = """\
**MultiModalInputs.as_kwargs(multi_modal_kwargs,
device=self.device),
**seqlen_agnostic_kwargs)"""
MODEL_CALL_REPLACEMENT = """\
**MultiModalInputs.as_kwargs(multi_modal_kwargs,
device=self.device),
**seqlen_agnostic_kwargs,
**gdn_prefix_kwargs)"""
PROFILE_KV_LAYERS_ANCHOR = """\
num_layers = self.model_config.get_num_layers(self.parallel_config)"""
PROFILE_KV_LAYERS_REPLACEMENT = """\
num_layers = self.model_config.get_num_attention_layers(
self.parallel_config)"""
def patch_model_runner(model_runner: pathlib.Path) -> None:
replace_once(
model_runner,
HELPER_ANCHOR,
HELPER_REPLACEMENT,
required=True,
already_contains="def _slice_mrope_positions(",
)
replace_once(
model_runner,
PREFIX_PAST_ANCHOR,
PREFIX_PAST_REPLACEMENT,
required=True,
already_contains="Must clear prefix_cache_hit so _add_seq_group",
)
replace_once(
model_runner,
PARTIAL_HIT_ANCHOR,
PARTIAL_HIT_REPLACEMENT,
required=True,
already_contains="positions, uncomputed_start, None,",
)
replace_once(
model_runner,
FULL_HIT_ANCHOR,
FULL_HIT_REPLACEMENT,
required=True,
already_contains="_slice_mrope_positions(positions, -1, None, 1)",
)
replace_once(
model_runner,
MULTIMODAL_MROPE_ANCHOR,
MULTIMODAL_MROPE_REPLACEMENT,
required=True,
already_contains="Compute the full MRoPE map once",
)
replace_once(
model_runner,
MODEL_INPUT_FIELDS_ANCHOR,
MODEL_INPUT_FIELDS_REPLACEMENT,
already_contains="gdn_restore_key: Optional[Tuple[int, bytes]]",
)
replace_once(
model_runner,
BASE_BROADCAST_ANCHOR,
BASE_BROADCAST_REPLACEMENT,
already_contains=BASE_BROADCAST_REPLACEMENT,
)
replace_once(
model_runner,
SAMPLING_BROADCAST_ANCHOR,
SAMPLING_BROADCAST_REPLACEMENT,
already_contains=SAMPLING_BROADCAST_REPLACEMENT,
)
replace_once(
model_runner,
BUILDER_INIT_ANCHOR,
BUILDER_INIT_REPLACEMENT,
already_contains="self.gdn_restore_key = None",
)
replace_once(
model_runner,
ADD_SEQ_GROUP_ANCHOR,
ADD_SEQ_GROUP_REPLACEMENT,
already_contains="gdn_actions = (",
)
replace_once(
model_runner,
BUILD_RESULT_ANCHOR,
BUILD_RESULT_REPLACEMENT,
already_contains="gdn_restore_key=self.gdn_restore_key",
)
replace_once(
model_runner,
EXECUTE_KWARGS_ANCHOR,
EXECUTE_KWARGS_REPLACEMENT,
already_contains="gdn_prefix_kwargs = {}",
)
replace_once(
model_runner,
MODEL_CALL_ANCHOR,
MODEL_CALL_REPLACEMENT,
already_contains="**gdn_prefix_kwargs)",
)
replace_once(
model_runner,
PROFILE_KV_LAYERS_ANCHOR,
PROFILE_KV_LAYERS_REPLACEMENT,
required=True,
already_contains=PROFILE_KV_LAYERS_REPLACEMENT,
)
if __name__ == "__main__":
patch_model_runner(package_root("vllm") / "worker" / "model_runner.py")

525
qwen3_6_scripts/patch_ops.sh Executable file
View File

@@ -0,0 +1,525 @@
#!/usr/bin/env bash
# BI-V100 patch script for Qwen3.6-35B-A3B (Qwen3_5 MoE architecture)
#
# Triton situation on BI-V100:
# - Standard Triton 2.3.1 is already present in the image.
# - HAS_TRITON = False (hardcoded in vendor vllm), but Triton is still used
# for TP-mode cache management (custom_cache_manager / libentry).
# - The vendor's triton_utils/__init__.py, custom_cache_manager.py, libentry.py
# are already correct for standard Triton 2.3.1 — do NOT overwrite them.
# - DO NOT install BI-V150 corex Triton 2.1.0 (pkgs/triton): that causes
# GPU hang on BI-V100 because the Triton CUDA PTX kernels are incompatible.
# Recommended server start command for TP=4 support 256K, needs chunked prefill
# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
# --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
# --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
# --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
# --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
# --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
# --tool-call-parser qwen3_coder --reasoning-parser qwen3
#
# With prefix caching (GDN align-mode, requires chunked prefill):
# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
# --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
# --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
# --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
# --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
# --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
# --tool-call-parser qwen3_coder --reasoning-parser qwen3
set -eo pipefail
# cd into this script's directory so ./relative paths work
cd "$(dirname "${BASH_SOURCE[0]}")"
echo "[patch_ops] working directory: $(pwd)"
build_stage() { printf '[BI100 BUILD] %s\n' "$1" >&2; }
require_file() {
local path=$1
[[ -f "$path" ]] || {
printf 'required patch source is missing: %s\n' "$path" >&2
exit 2
}
}
install_patch_file() {
local source=$1
local target=$2
require_file "$source"
mkdir -p "$(dirname "$target")"
install -m 0644 "$source" "$target"
}
build_stage "patch script entered"
build_stage "checking offline transformers dependency"
# --- transformers: Qwen3_5 tokenizer / model files --------------------------
TRANSFORMERS_REQUIRED_VERSION="4.55.3"
if ! python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
import importlib.metadata
import sys
required = sys.argv[1]
try:
installed = importlib.metadata.version("transformers")
except importlib.metadata.PackageNotFoundError:
raise SystemExit(1)
raise SystemExit(0 if installed == required else 1)
PY
then
WHEEL_DIR="./wheels"
if ! ls "${WHEEL_DIR}/transformers-${TRANSFORMERS_REQUIRED_VERSION}"*.whl >/dev/null 2>&1; then
echo "transformers ${TRANSFORMERS_REQUIRED_VERSION} is required, but no offline wheel was found in ${WHEEL_DIR}" >&2
exit 2
fi
python3 -m pip install --no-index --no-deps --find-links="${WHEEL_DIR}" \
"transformers==${TRANSFORMERS_REQUIRED_VERSION}"
fi
python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
import importlib.metadata
import sys
required = sys.argv[1]
installed = importlib.metadata.version("transformers")
if installed != required:
raise SystemExit(
f"transformers version mismatch: expected {required}, got {installed}")
print(f"[ok] transformers {installed}")
PY
build_stage "discovering Python package roots"
python3 - <<'PY' > /tmp/qwen36_patch_paths.env
from patch_utils import package_root, shell_env_line
print(shell_env_line("VLLM_ROOT", package_root("vllm")))
print(shell_env_line("TRANSFORMERS_ROOT", package_root("transformers")))
PY
source /tmp/qwen36_patch_paths.env
echo "VLLM_ROOT=${VLLM_ROOT}"
echo "TRANSFORMERS_ROOT=${TRANSFORMERS_ROOT}"
[[ -d "$VLLM_ROOT" ]] || {
printf 'vLLM root does not exist: %s\n' "$VLLM_ROOT" >&2
exit 2
}
VLLM_OVERRIDE_ROOT="./vendor_overrides/vllm"
[[ -d "$VLLM_OVERRIDE_ROOT" ]] || {
printf 'vLLM override directory missing: %s\n' "$VLLM_OVERRIDE_ROOT" >&2
exit 2
}
build_stage "installing authoritative vLLM core block overrides"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/core/evictor_v2.py" \
"${VLLM_ROOT}/core/evictor_v2.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/core/block/cpu_kv_content_cache.py" \
"${VLLM_ROOT}/core/block/cpu_kv_content_cache.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/core/block/cpu_gpu_block_allocator.py" \
"${VLLM_ROOT}/core/block/cpu_gpu_block_allocator.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/core/block/prefix_caching_block.py" \
"${VLLM_ROOT}/core/block/prefix_caching_block.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/core/block/block_table.py" \
"${VLLM_ROOT}/core/block/block_table.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/core/block_manager_v2.py" \
"${VLLM_ROOT}/core/block_manager_v2.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/sampling_params.py" \
"${VLLM_ROOT}/sampling_params.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/model_executor/sampling_metadata.py" \
"${VLLM_ROOT}/model_executor/sampling_metadata.py"
install_patch_file \
"${VLLM_OVERRIDE_ROOT}/model_executor/layers/sampler.py" \
"${VLLM_ROOT}/model_executor/layers/sampler.py"
build_stage "installing hash-pinned CoreX 3.2.3 extensions (16 prebuilt .so)"
bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
build_stage "installing BI100 runtime modules"
cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py"
cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"
cp ./block_major_kv_cache.py "${VLLM_ROOT}/block_major_kv_cache.py"
cp ./gdn_prefix.py "${VLLM_ROOT}/gdn_prefix.py"
build_stage "installing CoreX paged-KV swap compatibility"
python3 ./patch_corex_swap_blocks.py
python3 ./patch_block_major_cache_engine.py
python3 ./patch_worker_cache_transfer_order.py
# --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
# The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently
# (standard Triton 2.3.1 PTX is not supported by the corex runtime either).
# Our paged_attn.py bypasses it entirely via _forward_prefix_pytorch, which
# utilizes K-tiling techniques, and also have _forward_decode_pytorch to bypass kernel
# when context length is high
cp ./paged_attn.py "${VLLM_ROOT}/attention/ops/paged_attn.py"
# --- model_runner.py: fix prefix_cache_hit stays True in chunked-prefill chunk 2+ ---
# Bug: _compute_for_prefix_cache_hit Case 1 (prefix_cache_len <= context_len)
# leaves prefix_cache_hit=True. Then _add_seq_group uses block_table=computed_block_nums
# (only the original prefix blocks), ignoring chunk-1 KV cache blocks.
# _forward_prefix_pytorch then gets an undersized block_tables and crashes with
# "amax(): Expected reduction dim -1 to have non-zero size" on the 2nd tile.
# Fix: set prefix_cache_hit=False for Case 1 so the full block_tables is used.
python3 ./patch_model_runner.py
build_stage "installing executor startup diagnostics"
python3 ./patch_executor_startup_debug.py
python3 ./patch_worker_startup_profile_guard.py
python3 ./patch_block_major_worker_capacity.py
build_stage "installing transformers Qwen3.5 model support"
cp -r ./qwen3_5 "${TRANSFORMERS_ROOT}/models/"
cp -r ./qwen3_5_moe "${TRANSFORMERS_ROOT}/models/"
python3 ./patch_transformers_qwen3_5.py
build_stage "installing vLLM Qwen3.6 model implementation"
# --- vllm model: Qwen3.6-35B-A3B (Qwen3_5 MoE arch) -------------------------
cp ./mamba_cache.py "${VLLM_ROOT}/model_executor/models/"
cp ./qwen3_5.py "${VLLM_ROOT}/model_executor/models/qwen3_5.py"
cp ./ix_fused_moe.py "${VLLM_ROOT}/model_executor/models/ix_fused_moe.py" || true
python3 ./patch_vllm_qwen3_5.py
# --- Deploy prebuilt .so into vllm package for import -----------------------
PREBUILT_DIR="./prebuilt/corex-3.2.3-ivcore10"
if [ -d "$PREBUILT_DIR" ]; then
for so_file in "$PREBUILT_DIR"/*.so; do
base=$(basename "$so_file" .so)
# Deploy corex_*.so as vllm submodules (import from vllm import corex_xxx)
cp "$so_file" "${VLLM_ROOT}/${base}.so" 2>/dev/null || true
echo "[patch_ops] deployed ${base}.so → ${VLLM_ROOT}/"
done
fi
# --- Rebuild corex_moe_direct_routed.so for BI-V100 warp_size=64 -----------
# The prebuilt .so was compiled with kWarpSize=32 which silently corrupts
# results on BI-V100 (64-wide warps). Rebuild from the fixed .cu source
# that uses kWarpSize=64 and 6-step shuffle reductions.
build_stage "rebuilding corex_moe_direct_routed.so (warp64)"
COREX_ROOT="${COREX_ROOT:-/usr/local/corex-3.2.3}"
if [ ! -d "$COREX_ROOT" ]; then
COREX_ROOT="/usr/local/corex"
fi
TORCH_ROOT="${TORCH_ROOT:-$(python3 -c 'import torch,os;print(os.path.dirname(torch.__file__))' 2>/dev/null || echo "${COREX_ROOT}/lib64/python3/dist-packages/torch")}"
DIRECT_ROUTED_SRC="./corex_moe_direct_routed.cu"
DIRECT_ROUTED_DST="${VLLM_ROOT}/corex_moe_direct_routed.so"
if [ -f "$DIRECT_ROUTED_SRC" ] && [ -x "${COREX_ROOT}/bin/clang++" ]; then
"${COREX_ROOT}/bin/clang++" \
-std=c++17 -O3 -shared -fPIC \
--cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
--no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
-DTORCH_EXTENSION_NAME=corex_moe_direct_routed \
-DTORCH_API_INCLUDE_EXTENSION_H \
-I"${TORCH_ROOT}/include" \
-I"${TORCH_ROOT}/include/torch/csrc/api/include" \
-I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
-I/usr/local/include/python3.10 \
"$DIRECT_ROUTED_SRC" \
-L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
-Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
-ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
-lc10_cuda -lc10 -lcudart \
-o "$DIRECT_ROUTED_DST" 2>&1 && \
echo "[patch_ops] REBUILT corex_moe_direct_routed.so (warp64) → ${DIRECT_ROUTED_DST}" || \
echo "[patch_ops] WARNING: corex_moe_direct_routed.so rebuild FAILED, using prebuilt"
elif [ ! -x "${COREX_ROOT}/bin/clang++" ]; then
echo "[patch_ops] WARNING: CoreX clang++ not found at ${COREX_ROOT}/bin/clang++, cannot rebuild direct_routed"
else
echo "[patch_ops] WARNING: ${DIRECT_ROUTED_SRC} not found, cannot rebuild direct_routed"
fi
# --- Deploy ix_bridge Python integration layer --------------------------------
build_stage "deploying ix_bridge operator replacements"
EX_ENGINE_DIR="$(cd "$(dirname "$0")/ex_engine" 2>/dev/null && pwd || echo "")"
if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR/python" ]; then
EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")"
fi
if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR/python" ]; then
EX_ENGINE_DIR="/workspace/ex_engine"
fi
if [ -d "$EX_ENGINE_DIR/python" ]; then
# Create ex_engine package inside vllm with correct Python package structure
mkdir -p "${VLLM_ROOT}/ex_engine/python"
mkdir -p "${VLLM_ROOT}/ex_engine/csrc"
# __init__.py with re-exports so both import styles work:
# from ex_engine.python import ix_ops_dispatch (direct)
# from vllm.ex_engine import ix_ops_dispatch (via re-export)
cat > "${VLLM_ROOT}/ex_engine/__init__.py" << 'INIT_EOF'
"""ex_engine — Algorithm factor replacement for BI-V100."""
# Re-export python subpackage members at top level for backward compat
# Allows: from vllm.ex_engine import ix_ops_dispatch
try:
from ex_engine.python.ix_ops_dispatch import *
from ex_engine.python import ix_ops_dispatch
from ex_engine.python import ix_ops
from ex_engine.python import patch_vllm_ops
except ImportError:
pass
INIT_EOF
echo '"""ex_engine.python — dispatch and bridge modules."""' > "${VLLM_ROOT}/ex_engine/python/__init__.py"
# Deploy ALL Python modules
cp "$EX_ENGINE_DIR/python/"*.py "${VLLM_ROOT}/ex_engine/python/"
echo "[patch_ops] deployed $(ls -1 "${VLLM_ROOT}/ex_engine/python/"*.py | wc -l) modules → ${VLLM_ROOT}/ex_engine/python/"
# Deploy bridge C++ source for JIT fallback
for cpp in "$EX_ENGINE_DIR"/csrc/ix_full_bridge*.cpp "$EX_ENGINE_DIR"/csrc/ix_moe_bridge.cpp; do
[ -f "$cpp" ] && cp "$cpp" "${VLLM_ROOT}/ex_engine/csrc/" && \
echo "[patch_ops] deployed $(basename $cpp) for JIT fallback"
done
# Create startup hook that patches vllm ops at import time
cat > "${VLLM_ROOT}/ix_startup_patch.py" << 'STARTUP_EOF'
"""Apply ix_ops patches at vllm startup."""
import logging
_logger = logging.getLogger("ix_startup_patch")
_applied = False
def apply():
global _applied
if _applied:
return 0
_applied = True
import sys, os
# Ensure ex_engine is importable
for p in ["/workspace/qwen3_6_scripts", "/workspace"]:
rp = os.path.realpath(p)
if os.path.isdir(rp) and rp not in sys.path:
sys.path.insert(0, rp)
n = 0
try:
from ex_engine.python.patch_vllm_ops import apply_all_patches
k = apply_all_patches()
n += k
if k > 0:
_logger.info("ix_startup_patch: %d bridge patches applied", k)
except Exception as e:
_logger.warning("ix_startup_patch: bridge patches failed: %s", e)
try:
from ex_engine.python.patch_vllm_hot_path import apply as apply_hot
k = apply_hot(strict=False)
n += k
if k > 0:
_logger.info("ix_startup_patch: %d hot-path patches applied", k)
except Exception as e:
_logger.warning("ix_startup_patch: hot-path patches failed: %s", e)
try:
from ex_engine.python.patch_fused_linear_allreduce import apply_patch as apply_fused_ar
apply_fused_ar()
n += 1
_logger.info("ix_startup_patch: fused linear_allreduce patch applied")
except Exception as e:
_logger.warning("ix_startup_patch: fused linear_allreduce patch failed: %s", e)
return n
# DO NOT call apply() at import time — registry subprocess would crash.
# apply() is called from qwen3_5.py model init instead.
STARTUP_EOF
echo "[patch_ops] deployed ix_startup_patch.py"
# Hook into vllm __init__.py to auto-apply patches on import
VLLM_INIT="${VLLM_ROOT}/__init__.py"
if [ -f "$VLLM_INIT" ]; then
if ! grep -q "ix_startup_patch" "$VLLM_INIT" 2>/dev/null; then
echo "" >> "$VLLM_INIT"
echo "# Auto-apply ix_bridge operator patches" >> "$VLLM_INIT"
echo "try:" >> "$VLLM_INIT"
echo " from vllm import ix_startup_patch" >> "$VLLM_INIT"
echo "except Exception:" >> "$VLLM_INIT"
echo " pass" >> "$VLLM_INIT"
echo "[patch_ops] hooked ix_startup_patch into vllm/__init__.py"
fi
fi
else
echo "[patch_ops] WARN: ex_engine/python not found, skip ix_bridge deployment"
fi
# --- sequence.py: fix completion_tokens inflation under chunked prefill ------
# Bug: get_output_token_ids_to_return(delta=True) with num_new_tokens=0
# returns _cached_all_token_ids[-0:] == [0:] (the ENTIRE prompt+output list).
# Each prefill chunk step adds prompt_len to previous_num_tokens, so a 10K
# prompt processed in 3 chunks inflates completion_tokens by ~30K.
# Also adds num_cached_tokens field to RequestMetrics for prefix-cache stats.
cp ./sequence.py "${VLLM_ROOT}/sequence.py"
# --- scheduler.py: record num_cached_tokens in RequestMetrics ----------------
# Reports only the longest prefix backed by both live KV blocks and an exact
# GDN restore state. Raw KV-only hits must not inflate cached_tokens.
# serving_chat.py exposes the value in the OpenAI-compatible usage details.
cp ./scheduler.py "${VLLM_ROOT}/core/scheduler.py"
build_stage "installing diagnostic initial allocation trace"
python3 ./patch_block_manager_cache_trace.py
build_stage "installing scheduler and attention patches"
# --- xformers: bypass cudnnFlashAttnForward (head_dim=256 > 128 limit) ------
# Injects _run_sdpa_fallback (pure matmul+softmax) into xformers.py.
# Required because head_dim=256 > 128 and ixformer flash attention either
# crashes (is_causal=True) or produces wrong output (attn_mask path).
# The fallback uses query_start_loc to derive actual query lengths, so it
# works correctly during profiling runs with chunked-prefill-style batches.
# also bypasses auto chunked prefill on
python3 ./patch_xformers_sdpa_seq.py
python3 ./patch_xformers_profile.py
build_stage "installing API parsers and serving modules"
# --- tool parser: Qwen3 XML tool call format ---------------------------------
# Registers "qwen3_coder" parser for Qwen3.6 XML-style tool calls:
# <tool_call><function=name><parameter=key>\nvalue\n</parameter></function></tool_call>
# Use at server start: --tool-call-parser qwen3_coder --enable-auto-tool-choice
cp ./qwen3coder_tool_parser.py "${VLLM_ROOT}/entrypoints/openai/tool_parsers/"
python3 ./patch_vllm_tool_parser.py
# --- reasoning parser: Qwen3 <think>...</think> split ------------------------
# Adds --reasoning-parser qwen3 support.
# Routes thinking tokens to reasoning_content, rest to content in the delta.
# Works together with --tool-call-parser qwen3_coder (think → tool call flow).
cp -r ./reasoning "${VLLM_ROOT}/"
cp ./protocol.py "${VLLM_ROOT}/entrypoints/openai/protocol.py"
cp ./cli_args.py "${VLLM_ROOT}/entrypoints/openai/cli_args.py"
cp ./serving_chat.py "${VLLM_ROOT}/entrypoints/openai/serving_chat.py"
cp ./serving_tokenization.py \
"${VLLM_ROOT}/entrypoints/openai/serving_tokenization.py"
cp ./api_server.py "${VLLM_ROOT}/entrypoints/openai/api_server.py"
cp ./chat_utils.py "${VLLM_ROOT}/entrypoints/chat_utils.py"
python3 - ./api_server.py \
"${VLLM_ROOT}/entrypoints/openai/api_server.py" <<'PY'
from pathlib import Path
import sys
source = Path(sys.argv[1]).read_bytes()
installed = Path(sys.argv[2]).read_bytes()
if source != installed:
raise SystemExit("runtime api_server overlay identity mismatch")
PY
# --- protocol.py identity check: ensure max_completion_tokens is accepted ---
python3 - ./protocol.py \
"${VLLM_ROOT}/entrypoints/openai/protocol.py" <<'PY'
from pathlib import Path
import sys
source = Path(sys.argv[1]).read_bytes()
installed = Path(sys.argv[2]).read_bytes()
if source != installed:
raise SystemExit("runtime protocol overlay identity mismatch")
# Verify max_completion_tokens field is declared (not just extra=allow)
if b"max_completion_tokens" not in installed:
raise SystemExit("protocol.py missing max_completion_tokens field")
PY
build_stage "building CUTLASS grouped GEMM (gemm_grouped.so)"
if [[ -f "${EX_ENGINE_DIR}/build_gemm_grouped.sh" ]]; then
bash "${EX_ENGINE_DIR}/build_gemm_grouped.sh" 2>&1 || {
echo "[WARN] gemm_grouped build failed — will use torch.mm fallback"
}
# Deploy compiled .so if it exists
for so in "${EX_ENGINE_DIR}"/gemm_grouped.so "${EX_ENGINE_DIR}"/csrc/gemm_grouped.so; do
if [[ -f "$so" ]]; then
cp "$so" "${VLLM_ROOT}/gemm_grouped.so"
echo "[patch_ops] deployed gemm_grouped.so → ${VLLM_ROOT}/"
break
fi
done
fi
build_stage "building CUTLASS batched GEMM (corex_batched_gemm.so)"
if [[ -f "${EX_ENGINE_DIR}/xllm_kernels/cuda/corex_batched_gemm_kernel.cu" ]]; then
python3 << PYEOF
import os, sys, shutil
try:
from torch.utils.cpp_extension import load
ex = "${EX_ENGINE_DIR}"
cutlass_inc = ""
for d in ["/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass/include",
"/usr/local/corex/include/cutlass", "/usr/include/cutlass"]:
if os.path.isdir(d):
cutlass_inc = d
break
if not cutlass_inc:
print("[batched_gemm] No cutlass headers — skip"); sys.exit(0)
mod = load(
name="corex_batched_gemm",
sources=[
os.path.join(ex, "xllm_kernels/cuda/corex_batched_gemm_kernel.cu"),
os.path.join(ex, "xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp"),
],
extra_include_paths=[cutlass_inc],
extra_cflags=["-O2", "-std=c++17"],
extra_cuda_cflags=["-O2", f"-I{cutlass_inc}"],
extra_ldflags=["/usr/local/corex/lib64/libcuinfer.so", "-Wl,-rpath,/usr/local/corex/lib64"],
verbose=False,
)
print("[batched_gemm] ✓ Compiled")
import importlib
spec = importlib.util.find_spec("corex_batched_gemm")
if spec and spec.origin:
shutil.copy2(spec.origin, "${VLLM_ROOT}/corex_batched_gemm.so")
print("[batched_gemm] ✓ Deployed to ${VLLM_ROOT}/")
except Exception as e:
print(f"[batched_gemm] WARN: {e}")
PYEOF
fi
build_stage "building MoE bridge (ix_moe_bridge.so)"
if [[ -f "${EX_ENGINE_DIR}/csrc/ix_moe_bridge.cpp" ]]; then
SCRIPT_DIR="${EX_ENGINE_DIR}" bash "${EX_ENGINE_DIR}/build_moe_bridge.sh" "${VLLM_ROOT}" 2>&1 || {
echo "[WARN] MoE bridge build failed — will use Python fallback"
}
# Deploy .so to all paths ix_fused_moe.py searches
for src in "${VLLM_ROOT}/ex_engine/ix_moe_bridge.so" \
"${EX_ENGINE_DIR}/prebuilt/ix_moe_bridge.so"; do
if [[ -f "$src" ]]; then
cp "$src" "${VLLM_ROOT}/ix_moe_bridge.so" 2>/dev/null || true
cp "$src" "${VLLM_ROOT}/model_executor/models/ix_moe_bridge.so" 2>/dev/null || true
echo "[patch_ops] deployed ix_moe_bridge.so to vllm search paths"
break
fi
done
fi
build_stage "deploying fused linear+allreduce bridge (ix_full_bridge_fused_ar.so)"
for src in "${EX_ENGINE_DIR}/prebuilt/ix_full_bridge_fused_ar.so" \
"${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10/ix_full_bridge_fused_ar.so"; do
if [[ -f "$src" ]]; then
cp "$src" "${VLLM_ROOT}/ex_engine/ix_full_bridge_fused_ar.so" 2>/dev/null || true
cp "$src" "${VLLM_ROOT}/model_executor/models/ix_full_bridge_fused_ar.so" 2>/dev/null || true
echo "[patch_ops] deployed ix_full_bridge_fused_ar.so from prebuilt"
break
fi
done
build_stage "deploying all ex_engine Python modules"
EX_PY_DIR="${VLLM_ROOT}/ex_engine/python"
mkdir -p "${EX_PY_DIR}"
if [[ -d "${EX_ENGINE_DIR}/python" ]]; then
cp "${EX_ENGINE_DIR}/python/"*.py "${EX_PY_DIR}/" 2>/dev/null
echo "[patch_ops] deployed $(ls -1 "${EX_PY_DIR}"/*.py 2>/dev/null | wc -l) Python modules → ${EX_PY_DIR}/"
fi
build_stage "patching chat template for non-thinking mode"
MODEL_DIR="${MODEL_DIR:-/model}"
if [ -f "${MODEL_DIR}/tokenizer_config.json" ]; then
python3 ./patch_chat_template.py "${MODEL_DIR}" || \
echo "[patch_ops] WARNING: chat template patch failed"
else
echo "[patch_ops] WARNING: ${MODEL_DIR}/tokenizer_config.json not found"
fi
build_stage "compiling submission Python sources"
find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
build_stage "verifying dlopen chain"
python3 ./verify_dlopen_chain.py --vllm-root "${VLLM_ROOT}" || {
echo "[WARN] dlopen chain verification found issues (non-fatal)"
}
build_stage "patch script completed"

View File

@@ -0,0 +1,100 @@
"""
Patches transformers 4.55.3 to register qwen3_5 and qwen3_5_moe model types.
Deploy steps on the remote machine:
1. patch_ops.sh locates transformers with importlib.util.find_spec.
2. cp -r modified_scripts/qwen3_5* into the detected transformers/models.
3. python3 modified_scripts/patch_transformers_qwen3_5.py
"""
import sys
from patch_utils import package_root, replace_once, replace_one_of
TRANSFORMERS_ROOT = package_root("transformers")
AUTO_CONFIG = TRANSFORMERS_ROOT / "models" / "auto" / "configuration_auto.py"
MODELS_INIT = TRANSFORMERS_ROOT / "models" / "__init__.py"
def main():
print(f"=== Patching {AUTO_CONFIG} ===")
replace_one_of(AUTO_CONFIG, [
# CONFIG_MAPPING_NAMES: insert qwen3_5 + qwen3_5_moe right after qwen3
(
'("qwen3", "Qwen3Config"),',
'("qwen3", "Qwen3Config"),\n ("qwen3_5", "Qwen3_5Config"),\n ("qwen3_5_moe", "Qwen3_5MoeConfig"),',
),
(
'("qwen3", "Qwen3Config")\n',
'("qwen3", "Qwen3Config"),\n ("qwen3_5", "Qwen3_5Config"),\n ("qwen3_5_moe", "Qwen3_5MoeConfig"),\n',
),
], required=True, already_contains='("qwen3_5_moe", "Qwen3_5MoeConfig")')
replace_one_of(AUTO_CONFIG, [
# MODEL_NAMES_MAPPING (model_type -> human readable name)
(
'("qwen3", "Qwen3"),',
'("qwen3", "Qwen3"),\n ("qwen3_5", "Qwen3_5"),\n ("qwen3_5_moe", "Qwen3_5_MoE"),',
),
(
'("qwen3", "Qwen3")\n',
'("qwen3", "Qwen3"),\n ("qwen3_5", "Qwen3_5"),\n ("qwen3_5_moe", "Qwen3_5_MoE"),\n',
),
], required=True, already_contains='("qwen3_5_moe", "Qwen3_5_MoE")')
print(f"\n=== Patching {MODELS_INIT} ===")
replace_once(
MODELS_INIT,
"from .qwen3 import *\n",
"from .qwen3 import *\n from .qwen3_5 import *\n from .qwen3_5_moe import *\n",
required=True,
already_contains="from .qwen3_5_moe import *")
# Verification
print("\n=== Verification ===")
try:
import importlib.util, types
def _load_config_mod(module_name, file_path):
spec = importlib.util.spec_from_file_location(module_name, file_path)
mod = importlib.util.module_from_spec(spec)
mod.__package__ = ".".join(module_name.split(".")[:-1])
pkg = sys.modules.setdefault("transformers", types.ModuleType("transformers"))
pkg.__path__ = [str(TRANSFORMERS_ROOT)]
cu = sys.modules.setdefault(
"transformers.configuration_utils", types.ModuleType("transformers.configuration_utils"))
class _PC:
def __init__(self, **kwargs):
return None
cu.PretrainedConfig = _PC
for sub in ("transformers.models", f"transformers.models.{module_name.split('.')[-2]}"):
m = sys.modules.setdefault(sub, types.ModuleType(sub))
m.__path__ = [str(TRANSFORMERS_ROOT)]
spec.loader.exec_module(mod)
return mod
mod27 = _load_config_mod(
"transformers.models.qwen3_5.configuration_qwen3_5",
str(TRANSFORMERS_ROOT / "models" / "qwen3_5" /
"configuration_qwen3_5.py"),
)
cfg = mod27.Qwen3_5Config()
print(f" Qwen3_5Config() smoke-test OK (model_type={cfg.model_type})")
mod35 = _load_config_mod(
"transformers.models.qwen3_5_moe.configuration_qwen3_5_moe",
str(TRANSFORMERS_ROOT / "models" / "qwen3_5_moe" /
"configuration_qwen3_5_moe.py"),
)
moe_cfg = mod35.Qwen3_5MoeConfig()
print(f" Qwen3_5MoeConfig() smoke-test OK (model_type={moe_cfg.model_type})")
t = moe_cfg.text_config
print(f" num_experts={t.num_experts}, top_k={t.num_experts_per_tok}, "
f"shared={t.shared_expert_intermediate_size}, layers={t.num_hidden_layers}")
except Exception as e:
print(f" [optional] smoke-test failed (may be fine at runtime): {e}")
print("\nDone.")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,81 @@
from __future__ import annotations
import importlib.util
import pathlib
import shlex
from typing import Iterable, Optional, Sequence, Tuple
def package_root(pkg: str) -> pathlib.Path:
spec = importlib.util.find_spec(pkg)
if spec is None:
raise RuntimeError(f"package not found: {pkg}")
if not spec.submodule_search_locations:
raise RuntimeError(f"package has no package root: {pkg}")
return pathlib.Path(next(iter(spec.submodule_search_locations))).resolve()
def ensure_file(path: pathlib.Path) -> pathlib.Path:
if not path.is_file():
raise FileNotFoundError(str(path))
return path
def ensure_dir(path: pathlib.Path) -> pathlib.Path:
if not path.is_dir():
raise FileNotFoundError(str(path))
return path
def replace_once(path: pathlib.Path,
old: str,
new: str,
*,
required: bool = True,
already_contains: Optional[str] = None) -> bool:
path = ensure_file(path)
text = path.read_text()
marker = already_contains if already_contains is not None else new
if marker in text:
print(f"[skip] already patched: {path}")
return False
if old not in text:
msg = f"anchor not found in {path}: {old[:120]!r}"
if required:
raise RuntimeError(msg)
print(f"[warn] {msg}")
return False
path.write_text(text.replace(old, new, 1))
print(f"[ok] patched: {path}")
return True
def replace_one_of(path: pathlib.Path,
replacements: Sequence[Tuple[str, str]],
*,
required: bool = True,
already_contains: Optional[str] = None) -> bool:
path = ensure_file(path)
text = path.read_text()
if already_contains is not None and already_contains in text:
print(f"[skip] already patched: {path}")
return False
for _, new in replacements:
if new in text:
print(f"[skip] already patched: {path}")
return False
for old, new in replacements:
if old in text:
path.write_text(text.replace(old, new, 1))
print(f"[ok] patched: {path}")
return True
anchors = ", ".join(repr(old[:80]) for old, _ in replacements)
msg = f"anchor not found in {path}; tried: {anchors}"
if required:
raise RuntimeError(msg)
print(f"[warn] {msg}")
return False
def shell_env_line(name: str, value: pathlib.Path) -> str:
return f"{name}={shlex.quote(str(value))}"

View File

@@ -0,0 +1,73 @@
"""
Patches the vLLM model registry and deploys the Qwen3_5 model file.
Deploy steps on the remote machine:
1. patch_ops.sh locates vLLM with importlib.util.find_spec.
2. cp modified_scripts/qwen3_5.py into the detected vllm model directory.
2. python3 modified_scripts/patch_vllm_qwen3_5.py
The registry patch installs Qwen3.6 aliases so /model/config.json does not
need to be edited by hand.
"""
import ast
from patch_utils import package_root, replace_once
VLLM_ROOT = package_root("vllm")
REGISTRY = VLLM_ROOT / "model_executor" / "models" / "registry.py"
MODEL = VLLM_ROOT / "model_executor" / "models" / "qwen3_5.py"
EXPECTED_REGISTRY_ENTRIES = (
'"Qwen3ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")',
'"Qwen3MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")',
'"Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")',
'"Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")',
'"Qwen3_6ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")',
'"Qwen3_6MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")',
)
def main():
print(f"=== Patching {REGISTRY} ===")
replace_once(
REGISTRY,
' "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),\n'
' "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),',
' "Qwen3ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n'
' "Qwen3MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),\n'
' "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n'
' "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),\n'
' "Qwen3_6ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n'
' "Qwen3_6MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),',
required=True,
already_contains='"Qwen3_6MoeForCausalLM"')
print("\n=== Static verification ===")
model_source = MODEL.read_text(encoding="utf-8")
tree = ast.parse(model_source, filename=str(MODEL))
class_names = {
node.name for node in tree.body if isinstance(node, ast.ClassDef)
}
required_classes = {"Qwen3_5ForCausalLM", "Qwen3_5MoeForCausalLM"}
missing_classes = required_classes - class_names
if missing_classes:
raise RuntimeError(
f"Qwen3.5 model classes missing: {sorted(missing_classes)}")
registry_source = REGISTRY.read_text(encoding="utf-8")
missing_entries = [
entry for entry in EXPECTED_REGISTRY_ENTRIES
if entry not in registry_source
]
if missing_entries:
raise RuntimeError(
f"Qwen3.5 registry entries missing: {missing_entries}")
print(" model syntax and class declarations verified without import")
print(f" registry aliases verified: {len(EXPECTED_REGISTRY_ENTRIES)}")
print("\nDone. Registry aliases installed; do not edit /model/config.json.")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,57 @@
"""
Patches vLLM 0.6.3 to register Qwen3CoderToolParser under the name "qwen3_coder".
Deploy steps on the remote machine (already called by patch_ops.sh):
1. patch_ops.sh locates vLLM with importlib.util.find_spec.
2. cp qwen3coder_tool_parser.py into the detected vllm tool_parsers.
2. python3 patch_vllm_tool_parser.py
Usage after patching:
--tool-call-parser qwen3_coder --enable-auto-tool-choice
"""
from patch_utils import ensure_dir, package_root, replace_once
VLLM_ROOT = package_root("vllm")
TOOL_PARSERS_DIR = VLLM_ROOT / "entrypoints" / "openai" / "tool_parsers"
INIT_FILE = TOOL_PARSERS_DIR / "__init__.py"
def main():
ensure_dir(TOOL_PARSERS_DIR)
print(f"=== Patching {INIT_FILE} ===")
replace_once(
INIT_FILE,
"from .mistral_tool_parser import MistralToolParser",
"from .mistral_tool_parser import MistralToolParser\n"
"from .qwen3coder_tool_parser import Qwen3CoderToolParser",
required=True,
already_contains="from .qwen3coder_tool_parser import Qwen3CoderToolParser")
replace_once(
INIT_FILE,
'"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser"\n]',
'"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser",\n'
' "Qwen3CoderToolParser"\n]',
required=True,
already_contains='"Qwen3CoderToolParser"')
print("\n=== Verification ===")
try:
import importlib.util
spec = importlib.util.spec_from_file_location(
"qwen3coder_tool_parser",
str(TOOL_PARSERS_DIR / "qwen3coder_tool_parser.py"),
)
mod = importlib.util.module_from_spec(spec)
print(f" Module spec loaded: {spec.name}")
print(" (full import requires torch/vllm runtime — skipping exec)")
except Exception as e:
print(f" [optional] spec check failed: {e}")
print("\nDone. Start vLLM server with:")
print(" --tool-call-parser qwen3_coder --enable-auto-tool-choice")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,37 @@
from patch_utils import package_root, replace_once
WORKER = package_root("vllm") / "worker" / "worker.py"
CLEAN_BLOCK = """\
if (worker_input.blocks_to_swap_in is not None
and worker_input.blocks_to_swap_in.numel() > 0):
self.cache_engine[virtual_engine].swap_in(
worker_input.blocks_to_swap_in)
if (worker_input.blocks_to_swap_out is not None
and worker_input.blocks_to_swap_out.numel() > 0):
self.cache_engine[virtual_engine].swap_out(
worker_input.blocks_to_swap_out)
"""
ORDERED_BLOCK = """\
# BI100 content-addressed CPU KV tier may preserve a victim and reuse
# that same GPU slot in one step. Complete every D2H before any H2D.
if (worker_input.blocks_to_swap_out is not None
and worker_input.blocks_to_swap_out.numel() > 0):
self.cache_engine[virtual_engine].swap_out(
worker_input.blocks_to_swap_out)
if (worker_input.blocks_to_swap_in is not None
and worker_input.blocks_to_swap_in.numel() > 0):
self.cache_engine[virtual_engine].swap_in(
worker_input.blocks_to_swap_in)
"""
replace_once(
WORKER,
CLEAN_BLOCK,
ORDERED_BLOCK,
required=True,
already_contains="Complete every D2H before any H2D",
)

View File

@@ -0,0 +1,34 @@
from patch_utils import package_root, replace_one_of
WORKER = package_root("vllm") / "worker" / "worker.py"
CLEAN_BLOCK = """\
# Execute a forward pass with dummy inputs to profile the memory usage
# of the model.
self.model_runner.profile_run()
"""
GUARDED_BLOCK = """\
# Execute a forward pass with dummy inputs to profile the memory usage
# of the model. Mark this synthetic pass so BI100_PROFILE can exclude
# it without changing vLLM's normal capacity calculation.
_bi100_prev_startup_profile = os.environ.get("BI100_IN_STARTUP_PROFILE")
os.environ["BI100_IN_STARTUP_PROFILE"] = "1"
try:
self.model_runner.profile_run()
finally:
if _bi100_prev_startup_profile is None:
os.environ.pop("BI100_IN_STARTUP_PROFILE", None)
else:
os.environ["BI100_IN_STARTUP_PROFILE"] = _bi100_prev_startup_profile
"""
replace_one_of(
WORKER,
[(CLEAN_BLOCK, GUARDED_BLOCK)],
required=True,
already_contains=(
"Mark this synthetic pass so BI100_PROFILE can exclude"),
)

View File

@@ -0,0 +1,121 @@
"""Install disabled-by-default M1-48 XFormers timing boundaries."""
from __future__ import annotations
from pathlib import Path
try:
from patch_utils import package_root, replace_once
except ModuleNotFoundError:
from .patch_utils import package_root, replace_once
IMPORT_OLD = "from vllm.logger import init_logger"
IMPORT_NEW = """\
from vllm.bi100_profile import bi100_timer
from vllm.logger import init_logger"""
KV_WRITE_OLD = """\
PagedAttention.write_to_paged_cache(key, value, key_cache,
value_cache,
updated_slot_mapping,
self.kv_cache_dtype,
k_scale, v_scale)"""
KV_WRITE_NEW = """\
with bi100_timer("xformers.kv_write"):
PagedAttention.write_to_paged_cache(
key, value, key_cache, value_cache,
updated_slot_mapping, self.kv_cache_dtype,
k_scale, v_scale)"""
DENSE_OLD = """\
out = self._run_memory_efficient_xformers_forward(
query, key, value, prefill_meta, attn_type=attn_type)"""
DENSE_NEW = """\
with bi100_timer("xformers.dense_prefill"):
out = self._run_memory_efficient_xformers_forward(
query, key, value, prefill_meta, attn_type=attn_type)"""
PAGED_OLD = """\
out = PagedAttention.forward_prefix(
query,
key,
value,
self.kv_cache_dtype,
key_cache,
value_cache,
prefill_meta.block_tables,
prefill_meta.query_start_loc,
prefill_meta.seq_lens_tensor,
prefill_meta.context_lens_tensor,
prefill_meta.max_query_len,
self.alibi_slopes,
self.sliding_window,
k_scale,
v_scale,
is_causal_decoder=(attn_type == AttentionType.DECODER),
)"""
PAGED_NEW = """\
with bi100_timer("xformers.paged_prefill"):
out = PagedAttention.forward_prefix(
query,
key,
value,
self.kv_cache_dtype,
key_cache,
value_cache,
prefill_meta.block_tables,
prefill_meta.query_start_loc,
prefill_meta.seq_lens_tensor,
prefill_meta.context_lens_tensor,
prefill_meta.max_query_len,
self.alibi_slopes,
self.sliding_window,
k_scale,
v_scale,
is_causal_decoder=(attn_type == AttentionType.DECODER),
)"""
def patch_file(path: Path) -> None:
replace_once(
path,
IMPORT_OLD,
IMPORT_NEW,
already_contains="from vllm.bi100_profile import bi100_timer",
)
replace_once(
path,
KV_WRITE_OLD,
KV_WRITE_NEW,
already_contains='bi100_timer("xformers.kv_write")',
)
replace_once(
path,
DENSE_OLD,
DENSE_NEW,
already_contains='bi100_timer("xformers.dense_prefill")',
)
replace_once(
path,
PAGED_OLD,
PAGED_NEW,
already_contains='bi100_timer("xformers.paged_prefill")',
)
text = path.read_text(encoding="utf-8")
canonical = "\n".join(line.rstrip(" \t") for line in text.split("\n"))
if not canonical.endswith("\n"):
canonical += "\n"
if canonical != text:
path.write_text(canonical, encoding="utf-8")
def main() -> None:
path = package_root("vllm") / "attention" / "backends" / "xformers.py"
print("=== patch_xformers_profile (M1-48 diagnostic timers) ===")
print(f"Target: {path}")
patch_file(path)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,427 @@
"""
策略顺序per-sequencefallback — 纯 PyTorch 数学实现
==========================================================
逐条序列用 matmul + softmax 手写 attention完全绕开所有硬件
flash attention kernelixformer / cudnnFlashAttnForward
背景:
Iluvatar cudnnFlashAttnForward 存在两个已知问题:
1. 不支持 is_causal=True报错
2. 使用 attn_mask 路径时数值结果不正确(静默错误,输出全为"!"
与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。
纯数学路径matmul + softmax在任何 PyTorch 后端上结果都正确。
优点:
数值正确,不依赖任何硬件特定 attention kernel。
峰值显存 = max(seq_len)² × H × dtype_size由 --max-model-len 控制。
缺点:
并发请求的 prefill attention 串行执行。
O(L²) 显存(无 flash attention 的 O(L) 优化)。
内存参考fp16H_local=6
max-model-len=4096 → 峰值 ~200 MB
max-model-len=8192 → 峰值 ~800 MB
max-model-len=16384 → 峰值 ~3.2 GB
额外 patcharg_utils.py
vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill无命令行
关闭选项),原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling
解决了该问题chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到
_forward_prefix_pytorch属于不必要的行为变更因此一并禁用该自动逻辑。
Deploy:
python3 modified_scripts/patch_xformers_sdpa_seq.py
"""
from patch_utils import package_root, replace_one_of, replace_once
VLLM_ROOT = package_root("vllm")
XFORMERS_PATH = VLLM_ROOT / "attention" / "backends" / "xformers.py"
ARG_UTILS_PATH = VLLM_ROOT / "engine" / "arg_utils.py"
LOGITS_PROC_PATH = (
VLLM_ROOT / "model_executor" / "layers" / "logits_processor.py")
OUTLINES_DECODING_PATH = (
VLLM_ROOT / "model_executor" / "guided_decoding" /
"outlines_decoding.py")
# _apply_logits_processors crashes when seq_groups is None (intermediate
# chunked-prefill chunks on the driver rank). Add an early-return guard.
_LP_OLD_BLOCK = """\
def _apply_logits_processors(
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> torch.Tensor:
found_logits_processors = False\
"""
_LP_NEW_BLOCK = """\
def _apply_logits_processors(
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> torch.Tensor:
if sampling_metadata.seq_groups is None: # intermediate chunked-prefill chunk
return logits
found_logits_processors = False\
"""
# Outlines' UNESCAPED_STRING accepts raw JSON control characters, including
# newlines and tabs. The generated text can therefore satisfy the CFG while
# still failing json.loads(). Use the RFC 8259 string character constraints.
_JSON_STRING_OLD_BLOCK = """\
| UNESCAPED_STRING
| SIGNED_NUMBER -> number
| "true" -> true
| "false" -> false
| "null" -> null
array : "[" [value ("," value)*] "]"
object : "{" [pair ("," pair)*] "}"
pair : UNESCAPED_STRING ":" value
%import common.UNESCAPED_STRING
%import common.SIGNED_NUMBER
%import common.WS
%ignore WS\
"""
_JSON_STRING_V1_BLOCK = r'''| JSON_STRING
| SIGNED_NUMBER -> number
| "true" -> true
| "false" -> false
| "null" -> null
array : "[" [value ("," value)*] "]"
object : "{" [pair ("," pair)*] "}"
pair : JSON_STRING ":" value
JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/
%import common.SIGNED_NUMBER
%import common.WS
%ignore WS'''
_JSON_STRING_NEW_BLOCK = r'''| JSON_STRING
| SIGNED_NUMBER -> number
| "true" -> true
| "false" -> false
| "null" -> null
array : "[" _ws [value (_ws "," _ws value)*] _ws "]"
object : "{" _ws [pair (_ws "," _ws pair)*] _ws "}"
pair : JSON_STRING _ws ":" _ws value
_ws : JSON_WS?
JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/
JSON_WS: /[ \t\r\n]{1,4}/
%import common.SIGNED_NUMBER'''
# vllm 0.6.3 自动开启 chunked prefill 的原始块
_ARG_OLD_BLOCK = """\
if (is_gpu and not use_sliding_window and not use_spec_decode
and not self.enable_lora
and not self.enable_prompt_adapter):
self.enable_chunked_prefill = True
logger.warning(
"Chunked prefill is enabled by default for models with "
"max_model_len > 32K. Currently, chunked prefill might "
"not work with some features or models. If you "
"encounter any issues, please disable chunked prefill "
"by setting --enable-chunked-prefill=False.")\
"""
_ARG_NEW_BLOCK = """\
if (is_gpu and not use_sliding_window and not use_spec_decode
and not self.enable_lora
and not self.enable_prompt_adapter):
pass # skip auto-enable: Q-tiling in _run_sdpa_fallback
# handles long-context memory without chunked prefill\
"""
_MM_PREFIX_OLD_BLOCK = """\
if model_config.is_multimodal_model:
if self.enable_prefix_caching:
logger.warning(
"--enable-prefix-caching is currently not "
"supported for multimodal models and has been disabled.")
self.enable_prefix_caching = False\
"""
_MM_PREFIX_NEW_BLOCK = """\
if model_config.is_multimodal_model:
architectures = getattr(model_config.hf_config,
"architectures", []) or []
qwen36_native_vision = "Qwen3_5MoeForCausalLM" in architectures
if self.enable_prefix_caching and qwen36_native_vision:
logger.info(
"Keeping prefix caching enabled for the Qwen3.6 native "
"vision path.")
elif self.enable_prefix_caching:
logger.warning(
"--enable-prefix-caching is currently not "
"supported for multimodal models and has been disabled.")
self.enable_prefix_caching = False\
"""
FALLBACK_METHOD = '''
def _run_sdpa_fallback(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: "XFormersMetadata",
) -> torch.Tensor:
"""纯数学 causal attention fallback带 Q-tiling 内存优化。
调用时机kv_cache.numel()==0profiling 阶段)。
此路径无 KV 缓存前缀KV 长度 == query 长度。
内存优化Q-tiling与 Flash Attention 同思路):
将 Q 分成 _Q_CHUNK 大小的子块逐块计算,每块峰值内存
O(_Q_CHUNK × q_len) 而非 O(q_len²)。
profiling 阶段序列可能达到 max_model_len如 20K tokens
不加 Q-tiling 会产生 9.6 GB 矩阵直接 OOM。
softmax 在 float32 下计算以防止 float16 溢出,结果转回原始 dtype。
Args:
query : [1, total_query_tokens, num_heads, head_dim]
key : [1, total_query_tokens, num_kv_heads, head_dim]
value : [1, total_query_tokens, num_kv_heads, head_dim]
Returns:
[1, total_query_tokens, num_heads, head_dim]
"""
_Q_CHUNK = 256 # 与 _forward_prefix_pytorch 的 _ATTN_Q_CHUNK 保持一致
assert attn_metadata.seq_lens is not None
orig_dtype = query.dtype
num_seqs = len(attn_metadata.seq_lens)
# 推导每条序列的实际 query 长度。
# 正常 prefill 时 q_len == seq_len如果将来遇到 chunked 场景,
# query_start_loc 记录的是真实 query token 数(非全序列长度)。
if (attn_metadata.query_start_loc is not None
and len(attn_metadata.query_start_loc) == num_seqs + 1):
q_lens = [
int(attn_metadata.query_start_loc[i + 1].item()) -
int(attn_metadata.query_start_loc[i].item())
for i in range(num_seqs)
]
else:
q_lens = list(attn_metadata.seq_lens)
q_flat = query.squeeze(0) # [T, H, D]
k_flat = key.squeeze(0) # [T, Hkv, D]
v_flat = value.squeeze(0)
output = torch.empty_like(q_flat)
seq_start = 0
for q_len in q_lens:
seq_end = seq_start + q_len
# 当前序列的完整 K/V此路径无前缀KV == Q
k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float() # [Hkv, q_len, D]
v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float() # [Hkv, q_len, D]
# GQA展开 KV heads 至与 query heads 一致
if k_s.shape[0] != self.num_heads:
n = self.num_heads // k_s.shape[0]
k_s = k_s.repeat_interleave(n, dim=0).contiguous()
v_s = v_s.repeat_interleave(n, dim=0).contiguous()
# k_pos 用于因果掩码
k_pos = torch.arange(q_len, device=query.device)
# Q-tiling分块处理 query峰值内存 O(_Q_CHUNK × q_len)
for qc_start in range(0, q_len, _Q_CHUNK):
qc_end = min(qc_start + _Q_CHUNK, q_len)
# [H, qc, D]
q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \
.permute(1, 0, 2).float()
# [H, qc, q_len]
attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale
# 因果掩码q_c 里位置 j 只能看 k_pos <= j相对位置
qc_q_pos = torch.arange(qc_start, qc_end, device=query.device)
mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1)
attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf"))
attn_w = torch.softmax(attn_w, dim=-1)
out_c = torch.matmul(attn_w, v_s).to(orig_dtype) # [H, qc, D]
output[seq_start + qc_start:seq_start + qc_end] = (
out_c.permute(1, 0, 2))
seq_start = seq_end
return output.unsqueeze(0) # [1, T, H, D]
'''
OLD_XFORMER_BLOCK = """\
self.attn_op = xops.fmha.flash.FwOp()
if self.alibi_slopes is None:
# Add the batch dimension.
query = query.unsqueeze(0)
key = key.unsqueeze(0)
value = value.unsqueeze(0)
out = xops.memory_efficient_attention_forward(
query,
key,
value,
attn_bias=attn_bias[0],
p=0.0,
scale=self.scale,
op = self.attn_op
)
return out.view_as(original_query)\
"""
NEW_XFORMER_BLOCK = """\
self.attn_op = xops.fmha.flash.FwOp()
if self.alibi_slopes is None:
# Add the batch dimension.
query = query.unsqueeze(0)
key = key.unsqueeze(0)
value = value.unsqueeze(0)
if self.head_size > 128:
out = self._run_sdpa_fallback(query, key, value, attn_metadata)
else:
out = xops.memory_efficient_attention_forward(
query,
key,
value,
attn_bias=attn_bias[0],
p=0.0,
scale=self.scale,
op=self.attn_op,
)
return out.view_as(original_query)\
"""
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
_PREFIX_CALL_OLD_BLOCK = """\
out = PagedAttention.forward_prefix(
query,
key,
value,
self.kv_cache_dtype,
key_cache,
value_cache,
prefill_meta.block_tables,
prefill_meta.query_start_loc,
prefill_meta.seq_lens_tensor,
prefill_meta.context_lens_tensor,
prefill_meta.max_query_len,
self.alibi_slopes,
self.sliding_window,
k_scale,
v_scale,
)\
"""
_PREFIX_CALL_NEW_BLOCK = """\
out = PagedAttention.forward_prefix(
query,
key,
value,
self.kv_cache_dtype,
key_cache,
value_cache,
prefill_meta.block_tables,
prefill_meta.query_start_loc,
prefill_meta.seq_lens_tensor,
prefill_meta.context_lens_tensor,
prefill_meta.max_query_len,
self.alibi_slopes,
self.sliding_window,
k_scale,
v_scale,
is_causal_decoder=(attn_type == AttentionType.DECODER),
)\
"""
def patch_file(path):
replace_once(
path,
INJECT_ANCHOR,
FALLBACK_METHOD + INJECT_ANCHOR,
required=True,
already_contains="def _run_sdpa_fallback(")
replace_once(
path,
OLD_XFORMER_BLOCK,
NEW_XFORMER_BLOCK,
required=True,
already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
replace_once(
path,
_PREFIX_CALL_OLD_BLOCK,
_PREFIX_CALL_NEW_BLOCK,
required=True,
already_contains=(
"is_causal_decoder=(attn_type == AttentionType.DECODER)"))
def patch_arg_utils(path):
replace_once(
path,
_ARG_OLD_BLOCK,
_ARG_NEW_BLOCK,
required=True,
already_contains="skip auto-enable: Q-tiling")
replace_once(
path,
_MM_PREFIX_OLD_BLOCK,
_MM_PREFIX_NEW_BLOCK,
required=True,
already_contains="Keeping prefix caching enabled for the Qwen3.6")
def patch_logits_processor(path):
replace_once(
path,
_LP_OLD_BLOCK,
_LP_NEW_BLOCK,
required=True,
already_contains="intermediate chunked-prefill chunk")
def patch_outlines_json_grammar(path):
replace_one_of(
path,
[
(_JSON_STRING_V1_BLOCK, _JSON_STRING_NEW_BLOCK),
(_JSON_STRING_OLD_BLOCK, _JSON_STRING_NEW_BLOCK),
],
required=True,
already_contains="JSON_WS:")
def main():
print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===")
print(f"Target: {XFORMERS_PATH}")
patch_file(XFORMERS_PATH)
print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===")
print(f"Target: {ARG_UTILS_PATH}")
patch_arg_utils(ARG_UTILS_PATH)
print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===")
print(f"Target: {LOGITS_PROC_PATH}")
patch_logits_processor(LOGITS_PROC_PATH)
print("\n=== patch_outlines_json_grammar (reject raw control chars) ===")
print(f"Target: {OUTLINES_DECODING_PATH}")
patch_outlines_json_grammar(OUTLINES_DECODING_PATH)
print("\nDone.")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,16 @@
534019b3c2ad2d2c65492b01a975874ee440026eda2e8666bc3c1dc8a0a0a6f6 corex_attn_head_rms_norm.so
7e2aafd8dc755b0ee16c3b9bb812b95548fc042bbaa840dd9db7d2c51a10474c corex_block_major_kv_transfer.so
ad4ea7707bb2f2bfe04e07a7ad5fd58a647232be70a3056937a0d738c8254bff corex_fused_paged_prefill.so
1856c86e3100415061aa698a48bdeff3fe785994b45b4e72a42cd9158552a7d8 corex_gdn_beta_decay.so
957c7518f5831299fc73f19a4ca2aa3c8231afe9ea7c979127b4f426cd9d6906 corex_gdn_causal_conv.so
ff1c1c67ec252ed993bf1840a33ef5c4b386b730ba6dc5bf995991013c226222 corex_gdn_chunk_recurrent.so
ec2d11fa82d9d0816a6da53e62605e962786fa20ecd5f62e50f9d43087fc4d67 corex_gdn_gated_norm.so
27b7ae2ce4fe173336355d72a2678d043df4bd1ed85e9231a99bfb81885a6ce3 corex_gdn_packed_decode.so
015b61046ad73d8f12d754f7a87d4f6cba33070af1c079879e15b71a94571670 corex_gdn_qk_map.so
0eb120e89608bb5b64ca4356a5d3d362121806d081ccc1ccf346dac472a819ec corex_moe_direct_routed.so
d26f2fa39c3921a95793786601e90cf6ebadd06f1d752af541bf82c21acbc1c9 corex_moe_exact_reduce.so
0d5b04639a62fb67de63c413b596fc617a3784dbe72e04a7a3337a9e556f5b90 corex_moe_index_combine.so
c3208c8e0c13f54dbe22a9cfc88bdc6ab040e920d6cae4bc0ecf7087880795f3 corex_moe_topk_softmax.so
50b0b44c1da779bb2c03419ed549aee9bb922d1f9bab8b7f11a3d91cca0d21c3 corex_moe_weight_gather.so
e944ec0528ed9b6cb74518de3c57e3730543a7bdebc872f993bfdc8424f13e6b corex_paged_kv_gather.so
7de1d41919eee8b1caf4ef16b812f65d4b4a9421827e65a58e69aa79ce293ae5 ix_full_bridge.so

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

1228
qwen3_6_scripts/protocol.py Normal file

File diff suppressed because it is too large Load Diff

2857
qwen3_6_scripts/qwen3_5.py Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,3 @@
from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig
__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"]

View File

@@ -0,0 +1,242 @@
# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0
# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3
import os
from typing import Optional, List
from ...configuration_utils import PretrainedConfig as PreTrainedConfig
# --- Local stubs for APIs not present in transformers 4.55.3 ---
# Always use these definitions; do NOT import from the older transformers
# as same-named functions there have incompatible signatures.
def layer_type_validation(layer_types, num_hidden_layers=None, attention=True):
allowed = {"full_attention", "linear_attention"}
if not all(lt in allowed for lt in layer_types):
raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}")
if num_hidden_layers is not None and num_hidden_layers != len(layer_types):
raise ValueError(
f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})"
)
HYBRID_KV_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
HYBRID_KV_ACCOUNTING_CONFIG = "bi100_hybrid_kv_accounting_mode"
LEGACY_KV_ACCOUNTING = "legacy40"
FULL_ATTENTION_KV_ACCOUNTING = "full_attention"
def _hybrid_kv_accounting_mode(environ=None, serialized_mode=None):
source = os.environ if environ is None else environ
environment_mode = source.get(HYBRID_KV_ACCOUNTING_ENV)
if (environment_mode is not None and serialized_mode is not None
and environment_mode != serialized_mode):
raise RuntimeError(
f"{HYBRID_KV_ACCOUNTING_ENV}={environment_mode!r} conflicts "
f"with serialized {HYBRID_KV_ACCOUNTING_CONFIG}="
f"{serialized_mode!r}")
mode = environment_mode or serialized_mode or LEGACY_KV_ACCOUNTING
if mode not in (LEGACY_KV_ACCOUNTING, FULL_ATTENTION_KV_ACCOUNTING):
raise RuntimeError(
f"{HYBRID_KV_ACCOUNTING_ENV} must be "
f"'{LEGACY_KV_ACCOUNTING}' or "
f"'{FULL_ATTENTION_KV_ACCOUNTING}', got {mode!r}")
return mode
def _vllm_layers_block_type(
layer_types,
environ=None,
serialized_mode=None,
):
"""Expose hybrid-layer ownership in the form vLLM 0.6.3 consumes."""
mode = _hybrid_kv_accounting_mode(environ, serialized_mode)
if mode == LEGACY_KV_ACCOUNTING:
return ["attention"] * len(layer_types)
return [
"attention" if layer_type == "full_attention" else layer_type
for layer_type in layer_types
]
try:
from typing import TypedDict
except ImportError:
RopeParameters = dict
else:
class RopeParameters(TypedDict, total=False):
rope_theta: float
rope_type: str
partial_rotary_factor: float
factor: float
# --- End stubs ---
class Qwen3_5TextConfig(PreTrainedConfig):
r"""
Configuration for the text backbone of Qwen3.5 / Qwen3.6-35B-A3B models.
model_type is "qwen3_5_text" (used internally by the nested config).
"""
model_type = "qwen3_5_text"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=248320,
hidden_size=4096,
intermediate_size=12288,
num_hidden_layers=32,
num_attention_heads=16,
num_key_value_heads=4,
hidden_act="silu",
max_position_embeddings=32768,
initializer_range=0.02,
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_parameters=None,
attention_bias=False,
attention_dropout=0.0,
head_dim=256,
linear_conv_kernel_dim=4,
linear_key_head_dim=128,
linear_value_head_dim=128,
linear_num_key_heads=16,
linear_num_value_heads=32,
layer_types=None,
pad_token_id=None,
bos_token_id=None,
eos_token_id=None,
**kwargs,
):
self.pad_token_id = pad_token_id
self.bos_token_id = bos_token_id
self.eos_token_id = eos_token_id
self.tie_word_embeddings = tie_word_embeddings
self.vocab_size = vocab_size
self.max_position_embeddings = max_position_embeddings
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
self.head_dim = head_dim
self.rope_parameters = rope_parameters
kwargs.setdefault("partial_rotary_factor", 0.25)
self.layer_types = layer_types
if self.layer_types is None:
interval_pattern = kwargs.get("full_attention_interval", 4)
self.layer_types = [
"linear_attention" if bool((i + 1) % interval_pattern) else "full_attention"
for i in range(self.num_hidden_layers)
]
layer_type_validation(self.layer_types, self.num_hidden_layers)
self.linear_conv_kernel_dim = linear_conv_kernel_dim
self.linear_key_head_dim = linear_key_head_dim
self.linear_value_head_dim = linear_value_head_dim
self.linear_num_key_heads = linear_num_key_heads
self.linear_num_value_heads = linear_num_value_heads
super().__init__(**kwargs)
class Qwen3_5VisionConfig(PreTrainedConfig):
model_type = "qwen3_5_vision"
def __init__(
self,
depth=27,
hidden_size=1152,
hidden_act="gelu_pytorch_tanh",
intermediate_size=4304,
num_heads=16,
in_channels=3,
patch_size=16,
spatial_merge_size=2,
temporal_patch_size=2,
out_hidden_size=3584,
num_position_embeddings=2304,
initializer_range=0.02,
**kwargs,
):
super().__init__(**kwargs)
self.depth = depth
self.hidden_size = hidden_size
self.hidden_act = hidden_act
self.intermediate_size = intermediate_size
self.num_heads = num_heads
self.in_channels = in_channels
self.patch_size = patch_size
self.spatial_merge_size = spatial_merge_size
self.temporal_patch_size = temporal_patch_size
self.out_hidden_size = out_hidden_size
self.num_position_embeddings = num_position_embeddings
self.initializer_range = initializer_range
class Qwen3_5Config(PreTrainedConfig):
r"""
Top-level configuration for Qwen3.5 / Qwen3.6-35B-A3B.
model_type = "qwen3_5" matches the model card / config.json.
Wraps Qwen3_5TextConfig (and optionally Qwen3_5VisionConfig for multimodal use).
For vLLM text-only inference only text_config is consumed.
"""
model_type = "qwen3_5"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
text_config=None,
vision_config=None,
image_token_id=248056,
video_token_id=248057,
vision_start_token_id=248053,
vision_end_token_id=248054,
tie_word_embeddings=False,
**kwargs,
):
serialized_mode = kwargs.pop(HYBRID_KV_ACCOUNTING_CONFIG, None)
serialized_layers = kwargs.pop("layers_block_type", None)
if isinstance(text_config, dict):
self.text_config = Qwen3_5TextConfig(**text_config)
elif text_config is None:
self.text_config = Qwen3_5TextConfig()
else:
self.text_config = text_config
if isinstance(vision_config, dict):
self.vision_config = Qwen3_5VisionConfig(**vision_config)
elif vision_config is None:
self.vision_config = Qwen3_5VisionConfig()
else:
self.vision_config = vision_config
self.image_token_id = image_token_id
self.video_token_id = video_token_id
self.vision_start_token_id = vision_start_token_id
self.vision_end_token_id = vision_end_token_id
self.tie_word_embeddings = tie_word_embeddings
super().__init__(**kwargs)
mode = _hybrid_kv_accounting_mode(
serialized_mode=serialized_mode)
layers_block_type = _vllm_layers_block_type(
self.text_config.layer_types, serialized_mode=mode)
if (serialized_layers is not None
and list(serialized_layers) != layers_block_type):
raise RuntimeError(
"serialized layers_block_type conflicts with "
f"{HYBRID_KV_ACCOUNTING_CONFIG}={mode!r}")
setattr(self, HYBRID_KV_ACCOUNTING_CONFIG, mode)
self.layers_block_type = layers_block_type
__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"]

View File

@@ -0,0 +1,3 @@
from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig
__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"]

View File

@@ -0,0 +1,252 @@
# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0
# Source: transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3
# Removes ignore_keys_at_rope_validation / base_model_tp_plan / base_model_pp_plan
# which are 5.x-only and irrelevant for vLLM inference.
import os
from typing import Optional
from ...configuration_utils import PretrainedConfig as PreTrainedConfig
# --- Local stubs for APIs not present in transformers 4.55.3 ---
def layer_type_validation(layer_types, num_hidden_layers=None, attention=True):
allowed = {"full_attention", "linear_attention"}
if not all(lt in allowed for lt in layer_types):
raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}")
if num_hidden_layers is not None and num_hidden_layers != len(layer_types):
raise ValueError(
f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})"
)
HYBRID_KV_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
HYBRID_KV_ACCOUNTING_CONFIG = "bi100_hybrid_kv_accounting_mode"
LEGACY_KV_ACCOUNTING = "legacy40"
FULL_ATTENTION_KV_ACCOUNTING = "full_attention"
def _hybrid_kv_accounting_mode(environ=None, serialized_mode=None):
source = os.environ if environ is None else environ
environment_mode = source.get(HYBRID_KV_ACCOUNTING_ENV)
if (environment_mode is not None and serialized_mode is not None
and environment_mode != serialized_mode):
raise RuntimeError(
f"{HYBRID_KV_ACCOUNTING_ENV}={environment_mode!r} conflicts "
f"with serialized {HYBRID_KV_ACCOUNTING_CONFIG}="
f"{serialized_mode!r}")
mode = environment_mode or serialized_mode or LEGACY_KV_ACCOUNTING
if mode not in (LEGACY_KV_ACCOUNTING, FULL_ATTENTION_KV_ACCOUNTING):
raise RuntimeError(
f"{HYBRID_KV_ACCOUNTING_ENV} must be "
f"'{LEGACY_KV_ACCOUNTING}' or "
f"'{FULL_ATTENTION_KV_ACCOUNTING}', got {mode!r}")
return mode
def _vllm_layers_block_type(
layer_types,
environ=None,
serialized_mode=None,
):
"""Expose hybrid-layer ownership in the form vLLM 0.6.3 consumes."""
mode = _hybrid_kv_accounting_mode(environ, serialized_mode)
if mode == LEGACY_KV_ACCOUNTING:
return ["attention"] * len(layer_types)
return [
"attention" if layer_type == "full_attention" else layer_type
for layer_type in layer_types
]
try:
from typing import TypedDict
except ImportError:
RopeParameters = dict
else:
class RopeParameters(TypedDict, total=False):
rope_theta: float
rope_type: str
partial_rotary_factor: float
factor: float
# --- End stubs ---
class Qwen3_5MoeTextConfig(PreTrainedConfig):
r"""
Configuration for the text backbone of Qwen3.5-MoE / Qwen3.6-35B-A3B models.
model_type is "qwen3_5_moe_text" (used internally by the nested config).
"""
model_type = "qwen3_5_moe_text"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=248320,
hidden_size=2048,
num_hidden_layers=40,
num_attention_heads=16,
num_key_value_heads=2,
hidden_act="silu",
max_position_embeddings=32768,
initializer_range=0.02,
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_parameters=None,
attention_bias=False,
attention_dropout=0.0,
head_dim=256,
linear_conv_kernel_dim=4,
linear_key_head_dim=128,
linear_value_head_dim=128,
linear_num_key_heads=16,
linear_num_value_heads=32,
moe_intermediate_size=512,
shared_expert_intermediate_size=512,
num_experts_per_tok=8,
num_experts=256,
output_router_logits=False,
router_aux_loss_coef=0.001,
layer_types=None,
pad_token_id=None,
bos_token_id=None,
eos_token_id=None,
**kwargs,
):
self.pad_token_id = pad_token_id
self.bos_token_id = bos_token_id
self.eos_token_id = eos_token_id
self.tie_word_embeddings = tie_word_embeddings
self.vocab_size = vocab_size
self.max_position_embeddings = max_position_embeddings
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
self.head_dim = head_dim
self.rope_parameters = rope_parameters
kwargs.setdefault("partial_rotary_factor", 0.25)
self.layer_types = layer_types
if self.layer_types is None:
interval_pattern = kwargs.get("full_attention_interval", 4)
self.layer_types = [
"linear_attention" if bool((i + 1) % interval_pattern) else "full_attention"
for i in range(self.num_hidden_layers)
]
layer_type_validation(self.layer_types, self.num_hidden_layers)
self.linear_conv_kernel_dim = linear_conv_kernel_dim
self.linear_key_head_dim = linear_key_head_dim
self.linear_value_head_dim = linear_value_head_dim
self.linear_num_key_heads = linear_num_key_heads
self.linear_num_value_heads = linear_num_value_heads
self.moe_intermediate_size = moe_intermediate_size
self.shared_expert_intermediate_size = shared_expert_intermediate_size
self.num_experts_per_tok = num_experts_per_tok
self.num_experts = num_experts
self.output_router_logits = output_router_logits
self.router_aux_loss_coef = router_aux_loss_coef
super().__init__(**kwargs)
class Qwen3_5MoeVisionConfig(PreTrainedConfig):
model_type = "qwen3_5_moe"
def __init__(
self,
depth=27,
hidden_size=1152,
hidden_act="gelu_pytorch_tanh",
intermediate_size=4304,
num_heads=16,
in_channels=3,
patch_size=16,
spatial_merge_size=2,
temporal_patch_size=2,
out_hidden_size=3584,
num_position_embeddings=2304,
initializer_range=0.02,
**kwargs,
):
super().__init__(**kwargs)
self.depth = depth
self.hidden_size = hidden_size
self.hidden_act = hidden_act
self.intermediate_size = intermediate_size
self.num_heads = num_heads
self.in_channels = in_channels
self.patch_size = patch_size
self.spatial_merge_size = spatial_merge_size
self.temporal_patch_size = temporal_patch_size
self.out_hidden_size = out_hidden_size
self.num_position_embeddings = num_position_embeddings
self.initializer_range = initializer_range
class Qwen3_5MoeConfig(PreTrainedConfig):
r"""
Top-level configuration for Qwen3.5-MoE / Qwen3.6-35B-A3B.
model_type = "qwen3_5_moe" matches the model card / config.json.
Wraps Qwen3_5MoeTextConfig (and optionally Qwen3_5MoeVisionConfig).
For vLLM text-only inference only text_config is consumed.
"""
model_type = "qwen3_5_moe"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
text_config=None,
vision_config=None,
image_token_id=248056,
video_token_id=248057,
vision_start_token_id=248053,
vision_end_token_id=248054,
tie_word_embeddings=False,
**kwargs,
):
serialized_mode = kwargs.pop(HYBRID_KV_ACCOUNTING_CONFIG, None)
serialized_layers = kwargs.pop("layers_block_type", None)
if isinstance(text_config, dict):
self.text_config = Qwen3_5MoeTextConfig(**text_config)
elif text_config is None:
self.text_config = Qwen3_5MoeTextConfig()
else:
self.text_config = text_config
if isinstance(vision_config, dict):
self.vision_config = Qwen3_5MoeVisionConfig(**vision_config)
elif vision_config is None:
self.vision_config = Qwen3_5MoeVisionConfig()
else:
self.vision_config = vision_config
self.image_token_id = image_token_id
self.video_token_id = video_token_id
self.vision_start_token_id = vision_start_token_id
self.vision_end_token_id = vision_end_token_id
self.tie_word_embeddings = tie_word_embeddings
super().__init__(**kwargs)
mode = _hybrid_kv_accounting_mode(
serialized_mode=serialized_mode)
layers_block_type = _vllm_layers_block_type(
self.text_config.layer_types, serialized_mode=mode)
if (serialized_layers is not None
and list(serialized_layers) != layers_block_type):
raise RuntimeError(
"serialized layers_block_type conflicts with "
f"{HYBRID_KV_ACCOUNTING_CONFIG}={mode!r}")
setattr(self, HYBRID_KV_ACCOUNTING_CONFIG, mode)
self.layers_block_type = layers_block_type
__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"]

View File

@@ -0,0 +1,519 @@
import ast
import json
import uuid
from typing import Any, Dict, List, Optional, Sequence, Union
import regex as re
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
ChatCompletionToolsParam,
DeltaFunctionCall, DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall, ToolCall)
from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
ToolParser, ToolParserManager)
from vllm.logger import init_logger
from vllm.transformers_utils.tokenizer import AnyTokenizer
logger = init_logger(__name__)
@ToolParserManager.register_module("qwen3_coder")
class Qwen3CoderToolParser(ToolParser):
"""
Tool parser for Qwen3 models using XML-style tool call format:
<tool_call><function=name><parameter=key>
value
</parameter></function></tool_call>
Port of vllm-original qwen3coder_tool_parser.py to vllm 0.6.3 API.
"""
def __init__(self, tokenizer: AnyTokenizer):
super().__init__(tokenizer)
self.current_tool_name_sent: bool = False
self.prev_tool_call_arr: List[Dict] = []
# Base class uses int; we override with string IDs
self.current_tool_id: Optional[str] = None # type: ignore[assignment]
self.streamed_args_for_tool: List[str] = []
self.tool_call_start_token: str = "<tool_call>"
self.tool_call_end_token: str = "</tool_call>"
self.tool_call_prefix: str = "<function="
self.function_end_token: str = "</function>"
self.parameter_prefix: str = "<parameter="
self.parameter_end_token: str = "</parameter>"
self.is_tool_call_started: bool = False
self._reset_streaming_state()
self.tool_call_complete_regex = re.compile(
r"<tool_call>(.*?)</tool_call>", re.DOTALL)
self.tool_call_regex = re.compile(
r"<tool_call>(.*?)</tool_call>|<tool_call>(.*?)$", re.DOTALL)
self.tool_call_function_regex = re.compile(
r"<function=(.*?)</function>|<function=(.*)$", re.DOTALL)
self.tool_call_parameter_regex = re.compile(
r"<parameter=(.*?)(?:</parameter>|(?=<parameter=)|(?=</function>)|$)",
re.DOTALL)
if not self.model_tokenizer:
raise ValueError(
"The model tokenizer must be passed to the ToolParser "
"constructor during construction.")
self.tool_call_start_token_id = self.vocab.get(
self.tool_call_start_token)
self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
if (self.tool_call_start_token_id is None
or self.tool_call_end_token_id is None):
raise RuntimeError(
"Qwen3 XML Tool parser could not locate tool call start/end "
"tokens in the tokenizer!")
logger.debug("vLLM Successfully imported tool parser %s !",
self.__class__.__name__)
def _generate_tool_call_id(self) -> str:
return f"call_{uuid.uuid4().hex[:24]}"
def _reset_streaming_state(self) -> None:
self.current_tool_index = 0
self.is_tool_call_started = False
self.header_sent = False
self.current_tool_id = None
self.current_function_name: Optional[str] = None
self.current_param_name: Optional[str] = None
self.current_param_value: str = ""
self.param_count = 0
self.in_param = False
self.in_function = False
self.accumulated_text: str = ""
self.json_started = False
self.json_closed = False
self.accumulated_params: Dict[str, Any] = {}
self.streaming_request: Optional[ChatCompletionRequest] = None
def _get_arguments_config(
self, func_name: str,
tools: Optional[List[ChatCompletionToolsParam]]) -> Dict:
if tools is None:
return {}
for config in tools:
if not hasattr(config, "type") or not (
hasattr(config, "function")
and hasattr(config.function, "name")):
continue
if config.type == "function" and config.function.name == func_name:
if not hasattr(config.function, "parameters"):
return {}
params = config.function.parameters
if isinstance(params, dict) and "properties" in params:
return params["properties"]
elif isinstance(params, dict):
return params
else:
return {}
logger.debug("Tool '%s' is not defined in the tools list.", func_name)
return {}
def _convert_param_value(self, param_value: str, param_name: str,
param_config: Dict, func_name: str) -> Any:
if param_value.lower() == "null":
return None
if param_name not in param_config:
if param_config != {}:
logger.debug(
"Parsed parameter '%s' is not defined in tool '%s', "
"returning string value.", param_name, func_name)
return param_value
if (isinstance(param_config[param_name], dict)
and "type" in param_config[param_name]):
param_type = str(
param_config[param_name]["type"]).strip().lower()
else:
param_type = "string"
if param_type in ["string", "str", "text", "varchar", "char", "enum"]:
return param_value
elif (param_type.startswith("int") or param_type.startswith("uint")
or param_type.startswith("long")
or param_type.startswith("short")
or param_type.startswith("unsigned")):
try:
return int(param_value)
except (ValueError, TypeError):
return param_value
elif param_type.startswith("num") or param_type.startswith("float"):
try:
v = float(param_value)
return int(v) if v - int(v) == 0 else v
except (ValueError, TypeError):
return param_value
elif param_type in ["boolean", "bool", "binary"]:
lower = param_value.lower()
if lower not in ["true", "false"]:
logger.debug(
"Parameter '%s' value '%s' is not boolean in tool '%s'.",
param_name, param_value, func_name)
return lower == "true"
else:
if (param_type in ["object", "array", "arr"]
or param_type.startswith("dict")
or param_type.startswith("list")):
try:
return json.loads(param_value)
except (json.JSONDecodeError, TypeError, ValueError):
logger.debug(
"Could not JSON-decode parameter '%s' for tool '%s'; "
"falling back to literal evaluation.",
param_name,
func_name,
exc_info=True)
try:
return ast.literal_eval(param_value)
except (ValueError, SyntaxError, TypeError):
logger.debug(
"Could not literal-eval parameter '%s' for tool '%s'; "
"returning string value.",
param_name,
func_name,
exc_info=True)
return param_value
def _parse_xml_function_call(
self, function_call_str: str,
tools: Optional[List[ChatCompletionToolsParam]]) -> ToolCall:
end_index = function_call_str.index(">")
function_name = function_call_str[:end_index]
param_config = self._get_arguments_config(function_name, tools)
parameters = function_call_str[end_index + 1:]
param_dict: Dict[str, Any] = {}
for match_text in self.tool_call_parameter_regex.findall(parameters):
idx = match_text.index(">")
param_name = match_text[:idx]
param_value = str(match_text[idx + 1:])
if param_value.startswith("\n"):
param_value = param_value[1:]
if param_value.endswith("\n"):
param_value = param_value[:-1]
param_dict[param_name] = self._convert_param_value(
param_value, param_name, param_config, function_name)
return ToolCall(
type="function",
function=FunctionCall(
name=function_name,
arguments=json.dumps(param_dict, ensure_ascii=False)))
def _get_function_calls(self, model_output: str) -> List[str]:
matched_ranges = self.tool_call_regex.findall(model_output)
raw_tool_calls = [
match[0] if match[0] else match[1] for match in matched_ranges
]
if not raw_tool_calls:
raw_tool_calls = [model_output]
raw_function_calls: List[tuple] = []
for tool_call in raw_tool_calls:
raw_function_calls.extend(
self.tool_call_function_regex.findall(tool_call))
return [match[0] if match[0] else match[1]
for match in raw_function_calls]
def extract_tool_calls(
self, model_output: str,
request: ChatCompletionRequest) -> ExtractedToolCallInformation:
if self.tool_call_prefix not in model_output:
return ExtractedToolCallInformation(tools_called=False,
tool_calls=[],
content=model_output)
try:
function_calls = self._get_function_calls(model_output)
if not function_calls:
return ExtractedToolCallInformation(tools_called=False,
tool_calls=[],
content=model_output)
tool_calls = [
self._parse_xml_function_call(fc, request.tools)
for fc in function_calls
]
self.prev_tool_call_arr.clear()
for tc in tool_calls:
self.prev_tool_call_arr.append({
"name": tc.function.name,
"arguments": tc.function.arguments,
})
content_index = model_output.find(self.tool_call_start_token)
idx = model_output.find(self.tool_call_prefix)
content_index = content_index if content_index >= 0 else idx
content = model_output[:content_index]
return ExtractedToolCallInformation(
tools_called=bool(tool_calls),
tool_calls=tool_calls,
content=content if content else None,
)
except Exception:
logger.exception("Error extracting tool call from response.")
return ExtractedToolCallInformation(tools_called=False,
tool_calls=[],
content=model_output)
def extract_tool_calls_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
request: ChatCompletionRequest,
) -> Union[DeltaMessage, None]:
if not previous_text:
self._reset_streaming_state()
self.streaming_request = request
if not delta_text:
if delta_token_ids and self.tool_call_end_token_id not in delta_token_ids:
complete_calls = len(
self.tool_call_complete_regex.findall(current_text))
if complete_calls > 0 and self.prev_tool_call_arr:
open_calls = (
current_text.count(self.tool_call_start_token) -
current_text.count(self.tool_call_end_token))
if open_calls == 0:
return DeltaMessage(content="")
elif not self.is_tool_call_started and current_text:
return DeltaMessage(content="")
return None
self.accumulated_text = current_text
if self.json_closed and not self.in_function:
tool_ends = current_text.count(self.tool_call_end_token)
if tool_ends > self.current_tool_index:
self.current_tool_index += 1
self.header_sent = False
self.param_count = 0
self.json_started = False
self.json_closed = False
self.accumulated_params = {}
tool_starts = current_text.count(self.tool_call_start_token)
if self.current_tool_index >= tool_starts:
self.is_tool_call_started = False
return None
if not self.is_tool_call_started:
if (self.tool_call_start_token_id in delta_token_ids
or self.tool_call_start_token in delta_text):
self.is_tool_call_started = True
if self.tool_call_start_token in delta_text:
content_before = delta_text[:delta_text.index(
self.tool_call_start_token)]
if content_before:
return DeltaMessage(content=content_before)
return None
else:
if (current_text.rstrip().endswith(self.tool_call_end_token)
and delta_text.strip() == ""):
return None
return DeltaMessage(content=delta_text)
tool_starts_count = current_text.count(self.tool_call_start_token)
if self.current_tool_index >= tool_starts_count:
return None
# Locate the current tool call's text slice
tool_start_positions: List[int] = []
search = 0
while True:
search = current_text.find(self.tool_call_start_token, search)
if search == -1:
break
tool_start_positions.append(search)
search += len(self.tool_call_start_token)
if self.current_tool_index >= len(tool_start_positions):
return None
tool_start_idx = tool_start_positions[self.current_tool_index]
tool_end_idx = current_text.find(self.tool_call_end_token,
tool_start_idx)
if tool_end_idx == -1:
tool_text = current_text[tool_start_idx:]
else:
tool_text = current_text[tool_start_idx:tool_end_idx +
len(self.tool_call_end_token)]
if not self.header_sent:
if self.tool_call_prefix in tool_text:
func_start = (tool_text.find(self.tool_call_prefix) +
len(self.tool_call_prefix))
func_end = tool_text.find(">", func_start)
if func_end != -1:
self.current_function_name = tool_text[func_start:func_end]
self.current_tool_id = self._generate_tool_call_id()
self.header_sent = True
self.in_function = True
self.prev_tool_call_arr.append({
"name": self.current_function_name,
"arguments": "{}",
})
self.streamed_args_for_tool.append("")
return DeltaMessage(tool_calls=[
DeltaToolCall(
index=self.current_tool_index,
id=self.current_tool_id,
function=DeltaFunctionCall(
name=self.current_function_name,
arguments=""),
type="function",
)
])
return None
if self.in_function:
if not self.json_started:
self.json_started = True
self.streamed_args_for_tool[self.current_tool_index] += "{"
return DeltaMessage(tool_calls=[
DeltaToolCall(
index=self.current_tool_index,
function=DeltaFunctionCall(arguments="{"),
)
])
# Collect all complete parameters in one pass (speculative-decode safe)
param_starts: List[int] = []
search = 0
while True:
search = tool_text.find(self.parameter_prefix, search)
if search == -1:
break
param_starts.append(search)
search += len(self.parameter_prefix)
json_fragments: List[str] = []
while not self.in_param and self.param_count < len(param_starts):
param_idx = param_starts[self.param_count]
param_start = param_idx + len(self.parameter_prefix)
remaining = tool_text[param_start:]
if ">" not in remaining:
break
name_end = remaining.find(">")
current_param_name = remaining[:name_end]
value_start = param_start + name_end + 1
value_text = tool_text[value_start:]
if value_text.startswith("\n"):
value_text = value_text[1:]
param_end_idx = value_text.find(self.parameter_end_token)
if param_end_idx == -1:
next_param = value_text.find(self.parameter_prefix)
func_end = value_text.find(self.function_end_token)
if next_param != -1 and (func_end == -1
or next_param < func_end):
param_end_idx = next_param
elif func_end != -1:
param_end_idx = func_end
else:
tool_end_in_value = value_text.find(
self.tool_call_end_token)
if tool_end_in_value != -1:
param_end_idx = tool_end_in_value
else:
break
if param_end_idx == -1:
break
param_value = value_text[:param_end_idx]
if param_value.endswith("\n"):
param_value = param_value[:-1]
self.accumulated_params[current_param_name] = param_value
param_config = self._get_arguments_config(
self.current_function_name or "",
self.streaming_request.tools
if self.streaming_request else None)
converted = self._convert_param_value(
param_value, current_param_name, param_config,
self.current_function_name or "")
serialized = json.dumps(converted, ensure_ascii=False)
sep = "" if self.param_count == 0 else ", "
key = json.dumps(current_param_name, ensure_ascii=False)
json_fragments.append(f"{sep}{key}: {serialized}")
self.param_count += 1
if json_fragments:
combined = "".join(json_fragments)
if self.current_tool_index < len(self.streamed_args_for_tool):
self.streamed_args_for_tool[
self.current_tool_index] += combined
else:
logger.warning(
"streamed_args_for_tool out of sync: index=%d len=%d",
self.current_tool_index,
len(self.streamed_args_for_tool))
return DeltaMessage(tool_calls=[
DeltaToolCall(
index=self.current_tool_index,
function=DeltaFunctionCall(arguments=combined),
)
])
# Emit closing brace when </function> is seen (after params are done)
if not self.json_closed and self.function_end_token in tool_text:
self.json_closed = True
func_start = (tool_text.find(self.tool_call_prefix) +
len(self.tool_call_prefix))
func_content_end = tool_text.find(self.function_end_token,
func_start)
if func_content_end != -1:
try:
parsed_tool = self._parse_xml_function_call(
tool_text[func_start:func_content_end],
self.streaming_request.tools
if self.streaming_request else None)
if self.current_tool_index < len(
self.prev_tool_call_arr):
self.prev_tool_call_arr[
self.current_tool_index]["arguments"] = (
parsed_tool.function.arguments)
except Exception:
logger.debug("Failed to parse tool call during "
"streaming: %s",
tool_text,
exc_info=True)
if self.current_tool_index < len(self.streamed_args_for_tool):
self.streamed_args_for_tool[
self.current_tool_index] += "}"
else:
logger.warning(
"streamed_args_for_tool out of sync: index=%d len=%d",
self.current_tool_index,
len(self.streamed_args_for_tool))
result = DeltaMessage(tool_calls=[
DeltaToolCall(
index=self.current_tool_index,
function=DeltaFunctionCall(arguments="}"),
)
])
self.in_function = False
self.accumulated_params = {}
return result
return None

View File

@@ -0,0 +1,16 @@
"""
Reasoning parser module for vLLM 0.6.3 (BI-V100 / Qwen3.6-35B-A3B adaptation).
Usage: --reasoning-parser qwen3
"""
from vllm.reasoning.abs_reasoning_parsers import ReasoningParser, ReasoningParserManager
__all__ = ["ReasoningParser", "ReasoningParserManager"]
# Lazy-register Qwen3 parser; imported on first get_reasoning_parser("qwen3").
ReasoningParserManager.register_lazy(
"qwen3",
"vllm.reasoning.qwen3_reasoning_parser",
"Qwen3ReasoningParser",
)

View File

@@ -0,0 +1,243 @@
"""
Abstract reasoning parser base classes for vLLM 0.6.3.
Adapted from vllm-original/vllm/reasoning/abs_reasoning_parsers.py:
- Removed vllm.entrypoints.mcp, vllm.utils.collection_utils, import_utils
- DeltaMessage from vllm 0.6.3 protocol path
- TokenizerLike -> AnyTokenizer
- ReasoningParserManager: simplified eager + lazy registration
"""
import importlib
from abc import abstractmethod
from collections.abc import Iterable, Sequence
from functools import cached_property
from typing import Any, Optional, TYPE_CHECKING
if TYPE_CHECKING:
from vllm.entrypoints.openai.protocol import DeltaMessage
from vllm.transformers_utils.tokenizer import AnyTokenizer
else:
DeltaMessage = Any
AnyTokenizer = Any
class ReasoningParser:
"""Abstract base for all reasoning parsers."""
def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs):
self.model_tokenizer = tokenizer
@cached_property
def vocab(self) -> dict:
return self.model_tokenizer.get_vocab()
@abstractmethod
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
"""Return True once the reasoning block has closed in input_ids."""
def is_reasoning_end_streaming(
self, input_ids: Sequence[int], delta_ids: Iterable[int]
) -> bool:
return self.is_reasoning_end(input_ids)
@abstractmethod
def extract_content_ids(self, input_ids: list) -> list:
"""Return token ids that belong to the content (post-reasoning) part."""
def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
return 0
@abstractmethod
def extract_reasoning(
self, model_output: str, request: Any
) -> "tuple[Optional[str], Optional[str]]":
"""
Split a complete model output into (reasoning_text, content_text).
Either part may be None.
"""
@abstractmethod
def extract_reasoning_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
) -> Optional["DeltaMessage"]:
"""
Extract reasoning from a streaming delta.
Returns a DeltaMessage with reasoning_content and/or content set,
or None if this delta should be suppressed (control token).
"""
class BaseThinkingReasoningParser(ReasoningParser):
"""
Base for parsers that use <start_token>...</end_token> delimiters.
Subclasses define start_token / end_token properties.
"""
@property
@abstractmethod
def start_token(self) -> str:
raise NotImplementedError
@property
@abstractmethod
def end_token(self) -> str:
raise NotImplementedError
def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs):
super().__init__(tokenizer, *args, **kwargs)
if not self.model_tokenizer:
raise ValueError("Tokenizer must be passed to ReasoningParser.")
if not self.start_token or not self.end_token:
raise ValueError("start_token and end_token must be defined.")
self.start_token_id: Optional[int] = self.vocab.get(self.start_token)
self.end_token_id: Optional[int] = self.vocab.get(self.end_token)
if self.start_token_id is None or self.end_token_id is None:
raise RuntimeError(
f"{self.__class__.__name__}: could not find think tokens "
f"'{self.start_token}'/'{self.end_token}' in tokenizer vocab."
)
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
for token_id in reversed(input_ids):
if token_id == self.start_token_id:
return False
if token_id == self.end_token_id:
return True
return False
def is_reasoning_end_streaming(
self, input_ids: Sequence[int], delta_ids: Iterable[int]
) -> bool:
return self.end_token_id in delta_ids
def extract_content_ids(self, input_ids: list) -> list:
if self.end_token_id not in input_ids[:-1]:
return []
return input_ids[input_ids.index(self.end_token_id) + 1:]
def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
count = 0
depth = 0
for tid in token_ids:
if tid == self.start_token_id:
depth += 1
elif tid == self.end_token_id:
if depth > 0:
depth -= 1
elif depth > 0:
count += 1
return count
def extract_reasoning(
self, model_output: str, request: Any
) -> "tuple[Optional[str], Optional[str]]":
# Strip <think> if the model generated it (old-style template).
parts = model_output.partition(self.start_token)
model_output = parts[2] if parts[1] else parts[0]
if self.end_token not in model_output:
return model_output, None
reasoning, _, content = model_output.partition(self.end_token)
return reasoning, content or None
def extract_reasoning_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
) -> Optional["DeltaMessage"]:
from vllm.entrypoints.openai.protocol import DeltaMessage as _DeltaMessage
# Suppress lone control tokens.
if len(delta_token_ids) == 1 and delta_token_ids[0] in (
self.start_token_id, self.end_token_id
):
return None
start_in_prev = self.start_token_id in previous_token_ids
start_in_delta = self.start_token_id in delta_token_ids
end_in_prev = self.end_token_id in previous_token_ids
end_in_delta = self.end_token_id in delta_token_ids
if start_in_prev:
if end_in_delta:
end_idx = delta_text.find(self.end_token)
reasoning = delta_text[:end_idx] if end_idx >= 0 else ""
content = delta_text[end_idx + len(self.end_token):] if end_idx >= 0 else None
return _DeltaMessage(
reasoning_content=reasoning or None,
content=content or None,
)
elif end_in_prev:
return _DeltaMessage(content=delta_text)
else:
return _DeltaMessage(reasoning_content=delta_text)
elif start_in_delta:
if end_in_delta:
start_idx = delta_text.find(self.start_token)
end_idx = delta_text.find(self.end_token)
reasoning = delta_text[start_idx + len(self.start_token):end_idx]
content = delta_text[end_idx + len(self.end_token):]
return _DeltaMessage(
reasoning_content=reasoning or None,
content=content or None,
)
else:
return _DeltaMessage(reasoning_content=delta_text)
else:
return _DeltaMessage(content=delta_text)
class ReasoningParserManager:
"""
Registry for ReasoningParser implementations.
Supports eager and lazy registration.
"""
_parsers: dict = {} # name -> class (eager)
_lazy: dict = {} # name -> (module_path, class_name)
@classmethod
def register_module(cls, name: str, parser_cls: type) -> None:
"""Eagerly register a ReasoningParser class."""
if not issubclass(parser_cls, ReasoningParser):
raise TypeError(f"{parser_cls} is not a ReasoningParser subclass.")
cls._parsers[name] = parser_cls
@classmethod
def register_lazy(cls, name: str, module_path: str, class_name: str) -> None:
"""Register a parser for deferred import."""
cls._lazy[name] = (module_path, class_name)
@classmethod
def get_reasoning_parser(cls, name: str) -> type:
if name in cls._parsers:
return cls._parsers[name]
if name in cls._lazy:
module_path, class_name = cls._lazy[name]
mod = importlib.import_module(module_path)
parser_cls = getattr(mod, class_name)
cls._parsers[name] = parser_cls
return parser_cls
registered = sorted(set(cls._parsers) | set(cls._lazy))
raise KeyError(
f"Reasoning parser '{name}' not found. "
f"Available: {registered}"
)
@classmethod
def list_registered(cls) -> list:
return sorted(set(cls._parsers) | set(cls._lazy))

View File

@@ -0,0 +1,112 @@
"""
Reasoning parser for Qwen3 / Qwen3.5 / Qwen3.6 model family.
Adapted from vllm-original/vllm/reasoning/qwen3_reasoning_parser.py.
The model uses <think>...</think> to wrap chain-of-thought output.
For Qwen3.5+ the chat template injects <think> into the prompt, so only
</think> appears in the generated tokens; older templates generate <think>
themselves. Both styles are handled.
"""
from typing import Optional, Sequence, Any
from vllm.reasoning.abs_reasoning_parsers import (
BaseThinkingReasoningParser,
ReasoningParserManager,
)
class Qwen3ReasoningParser(BaseThinkingReasoningParser):
def __init__(self, tokenizer: Any, *args, **kwargs):
super().__init__(tokenizer, *args, **kwargs)
chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
self.thinking_enabled = chat_kwargs.get("enable_thinking", True)
@property
def start_token(self) -> str:
return "<think>"
@property
def end_token(self) -> str:
return "</think>"
def extract_reasoning(
self, model_output: str, request: Any
) -> "tuple[Optional[str], Optional[str]]":
# Strip <think> if the model generated it (old template / edge case).
parts = model_output.partition(self.start_token)
model_output = parts[2] if parts[1] else parts[0]
if not self.thinking_enabled:
if self.end_token in model_output:
_, _, content = model_output.partition(self.end_token)
return None, content or ""
return None, model_output
if self.end_token not in model_output:
# Thinking enabled but output truncated before </think>.
return model_output, None
reasoning, _, content = model_output.partition(self.end_token)
return reasoning, content or None
def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
token_ids = list(token_ids)
if self.start_token_id in token_ids:
# Old-style template: model generates <think> itself.
# Use depth-counting from the base class.
return super().count_reasoning_tokens(token_ids)
elif self.end_token_id in token_ids:
# New-style template (Qwen3.5+): <think> is injected into the
# prompt, so output starts already inside the thinking block.
# Every token before </think> is a reasoning token.
return token_ids.index(self.end_token_id)
else:
# No </think> in output: either truncated (all reasoning)
# or thinking disabled (none).
return len(token_ids) if self.thinking_enabled else 0
def extract_reasoning_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
):
from vllm.entrypoints.openai.protocol import DeltaMessage
if not self.thinking_enabled:
return DeltaMessage(content=delta_text) if delta_text else None
# Strip <think> from delta if the model generates it itself.
if self.start_token_id in delta_token_ids:
start_idx = delta_text.find(self.start_token)
if start_idx >= 0:
delta_text = delta_text[start_idx + len(self.start_token):]
if self.end_token_id in delta_token_ids:
end_idx = delta_text.find(self.end_token)
if end_idx >= 0:
reasoning = delta_text[:end_idx]
content = delta_text[end_idx + len(self.end_token):]
if not reasoning and not content:
return None
return DeltaMessage(
reasoning_content=reasoning or None,
content=content or None,
)
return None
if not delta_text:
return None
elif self.end_token_id in previous_token_ids:
return DeltaMessage(content=delta_text)
else:
return DeltaMessage(reasoning_content=delta_text)
# Register immediately when this module is imported.
ReasoningParserManager.register_module("qwen3", Qwen3ReasoningParser)

1995
qwen3_6_scripts/scheduler.py Normal file

File diff suppressed because it is too large Load Diff

1406
qwen3_6_scripts/sequence.py Normal file

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,159 @@
from typing import List, Optional, Union
from vllm.config import ModelConfig
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import (apply_hf_chat_template,
apply_mistral_chat_template,
load_chat_template,
parse_chat_messages_futures)
from vllm.entrypoints.logger import RequestLogger
# yapf conflicts with isort for this block
# yapf: disable
from vllm.entrypoints.openai.protocol import (DetokenizeRequest,
DetokenizeResponse,
ErrorResponse,
TokenizeChatRequest,
TokenizeRequest,
TokenizeResponse)
# yapf: enable
from vllm.entrypoints.openai.serving_engine import (BaseModelPath,
LoRAModulePath,
OpenAIServing)
from vllm.logger import init_logger
from vllm.transformers_utils.tokenizer import MistralTokenizer
from vllm.utils import random_uuid
logger = init_logger(__name__)
class OpenAIServingTokenization(OpenAIServing):
def __init__(
self,
engine_client: EngineClient,
model_config: ModelConfig,
base_model_paths: List[BaseModelPath],
*,
lora_modules: Optional[List[LoRAModulePath]],
request_logger: Optional[RequestLogger],
chat_template: Optional[str],
):
super().__init__(engine_client=engine_client,
model_config=model_config,
base_model_paths=base_model_paths,
lora_modules=lora_modules,
prompt_adapters=None,
request_logger=request_logger)
# If this is None we use the tokenizer's default chat template
# the list of commonly-used chat template names for HF named templates
hf_chat_templates: List[str] = ['default', 'tool_use']
self.chat_template = chat_template \
if chat_template in hf_chat_templates \
else load_chat_template(chat_template)
async def create_tokenize(
self,
request: TokenizeRequest,
) -> Union[TokenizeResponse, ErrorResponse]:
error_check_ret = await self._check_model(request)
if error_check_ret is not None:
return error_check_ret
request_id = f"tokn-{random_uuid()}"
(
lora_request,
prompt_adapter_request,
) = self._maybe_get_adapters(request)
tokenizer = await self.engine_client.get_tokenizer(lora_request)
prompt: Union[str, List[int]]
if isinstance(request, TokenizeChatRequest):
model_config = self.model_config
conversation, mm_data_future = parse_chat_messages_futures(
request.messages, model_config, tokenizer)
mm_data = await mm_data_future
if mm_data:
logger.warning(
"Multi-modal inputs are ignored during tokenization")
if isinstance(tokenizer, MistralTokenizer):
prompt = apply_mistral_chat_template(
tokenizer,
messages=request.messages,
chat_template=self.chat_template,
add_generation_prompt=request.add_generation_prompt,
continue_final_message=request.continue_final_message,
**(request.chat_template_kwargs or {}),
)
else:
prompt = apply_hf_chat_template(
tokenizer,
conversation=conversation,
chat_template=self.chat_template,
add_generation_prompt=request.add_generation_prompt,
continue_final_message=request.continue_final_message,
**(request.chat_template_kwargs or {}),
)
else:
prompt = request.prompt
self._log_inputs(request_id,
prompt,
params=None,
lora_request=lora_request,
prompt_adapter_request=prompt_adapter_request)
# Silently ignore prompt adapter since it does not affect tokenization
prompt_input = self._tokenize_prompt_input(
request,
tokenizer,
prompt,
add_special_tokens=request.add_special_tokens,
)
input_ids = prompt_input["prompt_token_ids"]
return TokenizeResponse(tokens=input_ids,
count=len(input_ids),
max_model_len=self.max_model_len)
async def create_detokenize(
self,
request: DetokenizeRequest,
) -> Union[DetokenizeResponse, ErrorResponse]:
error_check_ret = await self._check_model(request)
if error_check_ret is not None:
return error_check_ret
request_id = f"tokn-{random_uuid()}"
(
lora_request,
prompt_adapter_request,
) = self._maybe_get_adapters(request)
tokenizer = await self.engine_client.get_tokenizer(lora_request)
self._log_inputs(request_id,
request.tokens,
params=None,
lora_request=lora_request,
prompt_adapter_request=prompt_adapter_request)
if prompt_adapter_request is not None:
raise NotImplementedError("Prompt adapter is not supported "
"for tokenization")
prompt_input = self._tokenize_prompt_input(
request,
tokenizer,
request.tokens,
)
input_text = prompt_input["prompt"]
return DetokenizeResponse(prompt=input_text)

View File

@@ -0,0 +1,456 @@
import math
from typing import List, Optional
from vllm.core.block.common import BlockList
from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator
from vllm.utils import Device, cdiv, chunk_list
class BlockTable:
"""A class to manage blocks for a specific sequence.
The BlockTable maps a sequence of tokens to a list of blocks, where each
block represents a contiguous memory allocation for a portion of the
sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is
responsible for allocating and freeing memory for the blocks.
Args:
block_size (int): The maximum number of tokens that can be stored in a
single block.
block_allocator (DeviceAwareBlockAllocator): The block allocator used to
manage memory for the blocks.
_blocks (Optional[List[Block]], optional): An optional list of existing
blocks to initialize the BlockTable with. If not provided, an empty
BlockTable is created.
max_block_sliding_window (Optional[int], optional): The number of
blocks to keep around for each sequance. If None, all blocks
are kept (eg., when sliding window is not used).
It should at least fit the sliding window size of the model.
Attributes:
_block_size (int): The maximum number of tokens that can be stored in a
single block.
_allocator (DeviceAwareBlockAllocator): The block allocator used to
manage memory for the blocks.
_blocks (Optional[List[Block]]): The list of blocks managed by this
BlockTable.
_num_full_slots (int): The number of tokens currently stored in the
blocks.
"""
def __init__(
self,
block_size: int,
block_allocator: DeviceAwareBlockAllocator,
_blocks: Optional[List[Block]] = None,
max_block_sliding_window: Optional[int] = None,
cache_namespace: Optional[bytes] = None,
):
self._block_size = block_size
self._allocator = block_allocator
self._cache_namespace = cache_namespace
if _blocks is None:
_blocks = []
self._blocks: BlockList = BlockList(_blocks)
self._max_block_sliding_window = max_block_sliding_window
self._num_full_slots = self._get_num_token_ids()
@staticmethod
def get_num_required_blocks(token_ids: List[int],
block_size: int,
num_lookahead_slots: int = 0) -> int:
"""Calculates the minimum number of blocks required to store a given
sequence of token IDs along with any look-ahead slots that may be
required (like in multi-step + chunked-prefill).
This assumes worst-case scenario, where every block requires a new
allocation (e.g. ignoring prefix caching).
Args:
token_ids (List[int]): The sequence of token IDs to be stored.
block_size (int): The maximum number of tokens that can be stored in
a single block.
num_lookahead_slots (int): look-ahead slots that the sequence may
require.
Returns:
int: The minimum number of blocks required to store the given
sequence of token IDs along with any required look-ahead slots.
"""
return cdiv(len(token_ids) + num_lookahead_slots, block_size)
def allocate(self,
token_ids: List[int],
device: Device = Device.GPU) -> None:
"""Allocates memory blocks for storing the given sequence of token IDs.
This method allocates the required number of blocks to store the given
sequence of token IDs.
Args:
token_ids (List[int]): The sequence of token IDs to be stored.
device (Device, optional): The device on which the blocks should be
allocated. Defaults to Device.GPU.
"""
assert not self._is_allocated
assert token_ids
blocks = self._allocate_blocks_for_token_ids(prev_block=None,
token_ids=token_ids,
device=device)
self.update(blocks)
self._num_full_slots = len(token_ids)
def update(self, blocks: List[Block]) -> None:
"""Resets the table to the newly provided blocks
(with their corresponding block ids)
"""
self._blocks.update(blocks)
def get_content_hashes(self) -> List[bytes]:
"""Returns block-level content hashes for full blocks in order."""
content_hashes: List[bytes] = []
for block in self._blocks:
block_hash = block.content_hash
if block_hash is not None:
content_hashes.append(block_hash)
return content_hashes
def append_token_ids(self,
token_ids: List[int],
num_lookahead_slots: int = 0,
num_computed_slots: Optional[int] = None) -> None:
"""Appends a sequence of token IDs to the existing blocks in the
BlockTable.
This method appends the given sequence of token IDs to the existing
blocks in the BlockTable. If there is not enough space in the existing
blocks, new blocks are allocated using the `ensure_num_empty_slots`
method to accommodate the additional tokens.
The token IDs are divided into chunks of size `block_size` (except for
the first chunk, which may be smaller), and each chunk is appended to a
separate block.
Args:
token_ids (List[int]): The sequence of token IDs to be appended.
num_computed_slots (Optional[int]): The number of KV cache slots
that are already filled (computed).
When sliding window is enabled, this is used to compute how many
blocks to drop at the front of the sequence.
Without sliding window, None can be passed.
Without chunked prefill, it should be the same as
_num_full_slots.
"""
assert self._is_allocated, "no blocks have been allocated"
assert len(self._blocks) > 0
# Drop blocks that are no longer needed due to sliding window
if self._max_block_sliding_window is not None:
null_block = self._allocator.allocate_or_get_null_block()
assert num_computed_slots is not None
end_block_idx = (num_computed_slots //
self._block_size) - self._max_block_sliding_window
for idx in range(0, end_block_idx):
b = self._blocks[idx]
if b is not null_block:
self._allocator.free(b)
self._blocks[idx] = null_block
# Ensure there are enough empty slots for the new tokens plus
# lookahead slots
self.ensure_num_empty_slots(num_empty_slots=len(token_ids) +
num_lookahead_slots)
# Update the blocks with the new tokens
first_block_idx = self._num_full_slots // self._block_size
token_blocks = self._chunk_token_blocks_for_append(token_ids)
for i, token_block in enumerate(token_blocks):
self._blocks.append_token_ids(first_block_idx + i, token_block)
self._num_full_slots += len(token_ids)
def ensure_num_empty_slots(self, num_empty_slots: int) -> None:
"""Ensures that the BlockTable has at least the specified number of
empty slots available.
This method checks if the BlockTable has enough empty slots (i.e.,
available space) to accommodate the requested number of tokens. If not,
it allocates additional blocks on the GPU to ensure that the required
number of empty slots is available.
Args:
num_empty_slots (int): The minimum number of empty slots required.
"""
# Currently the block table only supports
# appending tokens to GPU blocks.
device = Device.GPU
assert self._is_allocated
if self._num_empty_slots >= num_empty_slots:
return
slots_to_allocate = num_empty_slots - self._num_empty_slots
blocks_to_allocate = cdiv(slots_to_allocate, self._block_size)
for _ in range(blocks_to_allocate):
assert len(self._blocks) > 0
self._blocks.append(
self._allocator.allocate_mutable_block(
prev_block=self._blocks[-1], device=device))
def fork(self) -> "BlockTable":
"""Creates a new BlockTable instance with a copy of the blocks from the
current instance.
This method creates a new BlockTable instance with the same block size,
block allocator, and a copy of the blocks from the current instance. The
new BlockTable has its own independent set of blocks, but shares the
same underlying memory allocation with the original BlockTable.
Returns:
BlockTable: A new BlockTable instance with a copy of the blocks from
the current instance.
"""
assert self._is_allocated
assert len(self._blocks) > 0
forked_blocks = self._allocator.fork(self._blocks[-1])
return BlockTable(
block_size=self._block_size,
block_allocator=self._allocator,
_blocks=forked_blocks,
max_block_sliding_window=self._max_block_sliding_window,
cache_namespace=self._cache_namespace,
)
def free(self) -> None:
"""Frees the memory occupied by the blocks in the BlockTable.
This method iterates over all the blocks in the `_blocks` list and calls
the `free` method of the `_allocator` object to release the memory
occupied by each block. After freeing all the blocks, the `_blocks` list
is set to `None`.
"""
for block in self.blocks:
self._allocator.free(block)
self._blocks.reset()
@property
def physical_block_ids(self) -> List[int]:
"""Returns a list of physical block indices for the blocks in the
BlockTable.
This property returns a list of integers, where each integer represents
the physical block index of a corresponding block in the `_blocks` list.
The physical block index is a unique identifier for the memory location
occupied by the block.
Returns:
List[int]: A list of physical block indices for the blocks in the
BlockTable.
"""
return self._blocks.ids()
def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]:
"""Get the number of "unseen" tokens in the sequence.
Unseen tokens are tokens in the sequence corresponding to this block
table, but are not yet appended to this block table.
Args:
sequence_token_ids (List[int]): The list of token ids in the
sequence.
Returns:
List[int]: The postfix of sequence_token_ids that has not yet been
appended to the block table.
"""
# Since the block table is append-only, the unseen token ids are the
# ones after the appended ones.
return sequence_token_ids[self.num_full_slots:]
def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
token_ids: List[int],
device: Device) -> List[Block]:
blocks: List[Block] = []
block_token_ids = []
tail_token_ids = []
for cur_token_ids in chunk_list(token_ids, self._block_size):
if len(cur_token_ids) == self._block_size:
block_token_ids.append(cur_token_ids)
else:
tail_token_ids.append(cur_token_ids)
if block_token_ids:
blocks.extend(self._allocate_immutable_blocks(
prev_block=prev_block,
block_token_ids=block_token_ids,
device=device))
prev_block = blocks[-1]
if tail_token_ids:
assert len(tail_token_ids) == 1
cur_token_ids = tail_token_ids[0]
block = self._allocate_mutable_block(prev_block=prev_block,
device=device)
block.append_token_ids(cur_token_ids)
blocks.append(block)
return blocks
def _allocate_mutable_block(self, prev_block: Optional[Block],
device: Device) -> Block:
if self._cache_namespace is None:
return self._allocator.allocate_mutable_block(
prev_block=prev_block, device=device)
with_cache_namespace = getattr(
self._allocator, "allocate_mutable_block_with_cache_namespace",
None)
if callable(with_cache_namespace):
return with_cache_namespace(
prev_block=prev_block,
cache_namespace=self._cache_namespace,
device=device)
backend_allocators = getattr(self._allocator, "_allocators", None)
if isinstance(backend_allocators, dict):
device_allocator = backend_allocators.get(device)
if device_allocator is not None:
with_cache_namespace = getattr(
device_allocator,
"allocate_mutable_block_with_cache_namespace", None)
if callable(with_cache_namespace):
return with_cache_namespace(
prev_block=prev_block,
cache_namespace=self._cache_namespace)
return self._allocator.allocate_mutable_block(
prev_block=prev_block, device=device)
def _allocate_immutable_blocks(self,
prev_block: Optional[Block],
block_token_ids: List[List[int]],
device: Device) -> List[Block]:
if self._cache_namespace is None:
return self._allocator.allocate_immutable_blocks(
prev_block,
block_token_ids=block_token_ids,
device=device)
with_cache_namespace = getattr(
self._allocator, "allocate_immutable_blocks_with_cache_namespace", None)
if callable(with_cache_namespace):
return with_cache_namespace(
prev_block=prev_block,
block_token_ids=block_token_ids,
cache_namespace=self._cache_namespace,
device=device)
backend_allocator = getattr(self._allocator, "_allocators", None)
if isinstance(backend_allocator, dict):
device_allocator = backend_allocator.get(device)
if device_allocator is not None:
with_cache_namespace = getattr(
device_allocator,
"allocate_immutable_blocks_with_cache_namespace",
None)
if callable(with_cache_namespace):
return with_cache_namespace(
prev_block=prev_block,
block_token_ids=block_token_ids,
cache_namespace=self._cache_namespace)
# Fallback: keep behavior identical when no namespace-aware allocator
# is available.
return self._allocator.allocate_immutable_blocks(
prev_block,
block_token_ids=block_token_ids,
device=device)
def _get_all_token_ids(self) -> List[int]:
# NOTE: This function is O(seq_len); use sparingly.
token_ids: List[int] = []
if not self._is_allocated:
return token_ids
for block in self.blocks:
token_ids.extend(block.token_ids)
return token_ids
def _get_num_token_ids(self) -> int:
res = 0
for block in self.blocks:
res += len(block.token_ids)
return res
@property
def _is_allocated(self) -> bool:
return len(self._blocks) > 0
@property
def blocks(self) -> List[Block]:
return self._blocks.list()
@property
def _num_empty_slots(self) -> int:
assert self._is_allocated
return len(self._blocks) * self._block_size - self._num_full_slots
@property
def num_full_slots(self) -> int:
"""Returns the total number of tokens currently stored in the
BlockTable.
Returns:
int: The total number of tokens currently stored in the BlockTable.
"""
return self._num_full_slots
def get_num_blocks_touched_by_append_slots(
self, token_ids: List[int], num_lookahead_slots: int) -> int:
"""Determine how many blocks will be "touched" by appending the token
ids.
This is required for the scheduler to determine whether a sequence can
continue generation, or if it must be preempted.
"""
# Math below is equivalent to:
# all_token_ids = token_ids + [-1] * num_lookahead_slots
# token_blocks = self._chunk_token_blocks_for_append(all_token_ids)
# return len(token_blocks)
num_token_ids = len(token_ids) + num_lookahead_slots
first_chunk_size = self._block_size - (self._num_full_slots %
self._block_size)
num_token_blocks = (1 + math.ceil(
(num_token_ids - first_chunk_size) / self._block_size))
return num_token_blocks
def _chunk_token_blocks_for_append(
self, token_ids: List[int]) -> List[List[int]]:
"""Split the token ids into block-sized chunks so they can be easily
appended to blocks. The first such "token block" may have less token ids
than the block size, since the last allocated block may be partially
full.
If no token ids are provided, then no chunks are returned.
"""
if not token_ids:
return []
first_chunk_size = self._block_size - (self._num_full_slots %
self._block_size)
token_blocks = [token_ids[:first_chunk_size]]
token_blocks.extend(
chunk_list(token_ids[first_chunk_size:], self._block_size))
return token_blocks

View File

@@ -0,0 +1,475 @@
from typing import Dict, FrozenSet, List, Optional, Tuple
from vllm.core.block.cpu_kv_content_cache import (CpuKvContentCache,
cpu_kv_offload_enabled)
from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,
DeviceAwareBlockAllocator)
from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator
from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator
from vllm.utils import Device
class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):
"""A block allocator that can allocate blocks on both CPU and GPU memory.
This class implements the `DeviceAwareBlockAllocator` interface and provides
functionality for allocating and managing blocks of memory on both CPU and
GPU devices.
The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU
blocks, and allows for allocation, deallocation, forking, and swapping of
blocks across these memory pools.
"""
@staticmethod
def create(
allocator_type: str,
num_gpu_blocks: int,
num_cpu_blocks: int,
block_size: int,
) -> DeviceAwareBlockAllocator:
"""Creates a CpuGpuBlockAllocator instance with the specified
configuration.
This static method creates and returns a CpuGpuBlockAllocator instance
based on the provided parameters. It initializes the CPU and GPU block
allocators with the specified number of blocks, block size, and
allocator type.
Args:
allocator_type (str): The type of block allocator to use for CPU
and GPU blocks. Currently supported values are "naive" and
"prefix_caching".
num_gpu_blocks (int): The number of blocks to allocate for GPU
memory.
num_cpu_blocks (int): The number of blocks to allocate for CPU
memory.
block_size (int): The size of each block in number of tokens.
Returns:
DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the
specified configuration.
Notes:
- The block IDs are assigned contiguously, with GPU block IDs coming
before CPU block IDs.
"""
content_offload = cpu_kv_offload_enabled()
if content_offload and allocator_type != "prefix_caching":
raise RuntimeError(
"BI100_CPU_KV_OFFLOAD=1 requires prefix caching")
if content_offload and num_cpu_blocks <= 0:
raise RuntimeError(
"BI100_CPU_KV_OFFLOAD=1 requires at least one CPU KV block")
block_ids = list(range(num_gpu_blocks + num_cpu_blocks))
gpu_block_ids = block_ids[:num_gpu_blocks]
cpu_block_ids = block_ids[num_gpu_blocks:]
if allocator_type == "naive":
gpu_allocator: BlockAllocator = NaiveBlockAllocator(
create_block=NaiveBlock, # type: ignore
num_blocks=num_gpu_blocks,
block_size=block_size,
block_ids=gpu_block_ids,
)
cpu_allocator: BlockAllocator = NaiveBlockAllocator(
create_block=NaiveBlock, # type: ignore
num_blocks=num_cpu_blocks,
block_size=block_size,
block_ids=cpu_block_ids,
)
elif allocator_type == "prefix_caching":
gpu_allocator = PrefixCachingBlockAllocator(
num_blocks=num_gpu_blocks,
block_size=block_size,
block_ids=gpu_block_ids,
)
cpu_allocator = PrefixCachingBlockAllocator(
num_blocks=num_cpu_blocks,
block_size=block_size,
block_ids=cpu_block_ids,
)
else:
raise ValueError(f"Unknown allocator type {allocator_type=}")
return CpuGpuBlockAllocator(
cpu_block_allocator=cpu_allocator,
gpu_block_allocator=gpu_allocator,
cpu_content_cache=(CpuKvContentCache(num_cpu_blocks)
if content_offload else None),
)
def __init__(self, cpu_block_allocator: BlockAllocator,
gpu_block_allocator: BlockAllocator,
cpu_content_cache: Optional[CpuKvContentCache] = None):
assert not (
cpu_block_allocator.all_block_ids
& gpu_block_allocator.all_block_ids
), "cpu and gpu block allocators can't have intersection of block ids"
self._allocators = {
Device.CPU: cpu_block_allocator,
Device.GPU: gpu_block_allocator,
}
self._swap_mapping: Dict[int, int] = {}
self._null_block: Optional[Block] = None
self._cpu_content_cache = cpu_content_cache
self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}
for _, allocator in self._allocators.items():
for block_id in allocator.all_block_ids:
self._block_ids_to_allocator[block_id] = allocator
if self._cpu_content_cache is not None:
if not isinstance(gpu_block_allocator,
PrefixCachingBlockAllocator):
raise RuntimeError(
"CPU KV content tier requires PrefixCachingBlockAllocator")
if (self._cpu_content_cache.capacity !=
cpu_block_allocator.get_num_total_blocks()):
raise RuntimeError(
"CPU KV content capacity must cover the complete CPU cache")
gpu_block_allocator.set_external_cache_callbacks(
claim=self._claim_cpu_content,
load=self._stage_cpu_to_gpu,
cancel=self._cancel_cpu_claim,
store=self._stage_gpu_to_cpu,
)
@property
def content_offload_enabled(self) -> bool:
return self._cpu_content_cache is not None
def _claim_cpu_content(self, content_hash: bytes) -> Optional[int]:
assert self._cpu_content_cache is not None
return self._cpu_content_cache.claim_load(content_hash)
def _cancel_cpu_claim(self, content_hash: bytes, cpu_slot: int) -> None:
assert self._cpu_content_cache is not None
self._cpu_content_cache.cancel_load(content_hash, cpu_slot)
def _stage_cpu_to_gpu(self, content_hash: bytes, cpu_slot: int,
gpu_block_id: BlockId) -> None:
assert self._cpu_content_cache is not None
gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id)
self._cpu_content_cache.stage_load(
content_hash, cpu_slot, gpu_slot)
def _stage_gpu_to_cpu(self, content_hash: bytes,
gpu_block_id: BlockId) -> bool:
assert self._cpu_content_cache is not None
gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id)
return self._cpu_content_cache.stage_store(content_hash, gpu_slot)
def allocate_or_get_null_block(self) -> Block:
if self._null_block is None:
self._null_block = NullBlock(
self.allocate_mutable_block(None, Device.GPU))
return self._null_block
def allocate_mutable_block(self, prev_block: Optional[Block],
device: Device) -> Block:
"""Allocates a new mutable block on the specified device.
Args:
prev_block (Optional[Block]): The previous block to in the sequence.
Used for prefix hashing.
device (Device): The device on which to allocate the new block.
Returns:
Block: The newly allocated mutable block.
"""
return self._allocators[device].allocate_mutable_block(prev_block)
def allocate_immutable_blocks(self, prev_block: Optional[Block],
block_token_ids: List[List[int]],
device: Device) -> List[Block]:
"""Allocates a new group of immutable blocks with the provided block
token IDs on the specified device.
Args:
prev_block (Optional[Block]): The previous block in the sequence.
Used for prefix hashing.
block_token_ids (List[int]): The list of block token IDs to be
stored in the new blocks.
device (Device): The device on which to allocate the new block.
Returns:
List[Block]: The newly allocated list of immutable blocks
containing the provided block token IDs.
"""
return self._allocators[device].allocate_immutable_blocks(
prev_block, block_token_ids)
def allocate_immutable_block(self, prev_block: Optional[Block],
token_ids: List[int],
device: Device) -> Block:
"""Allocates a new immutable block with the provided token IDs on the
specified device.
Args:
prev_block (Optional[Block]): The previous block in the sequence.
Used for prefix hashing.
token_ids (List[int]): The list of token IDs to be stored in the new
block.
device (Device): The device on which to allocate the new block.
Returns:
Block: The newly allocated immutable block containing the provided
token IDs.
"""
return self._allocators[device].allocate_immutable_block(
prev_block, token_ids)
def free(self, block: Block) -> None:
"""Frees the memory occupied by the given block.
Args:
block (Block): The block to be freed.
"""
# Null block should never be freed
if isinstance(block, NullBlock):
return
block_id = block.block_id
assert block_id is not None
allocator = self._block_ids_to_allocator[block_id]
allocator.free(block)
def fork(self, last_block: Block) -> List[Block]:
"""Creates a new sequence of blocks that shares the same underlying
memory as the original sequence.
Args:
last_block (Block): The last block in the original sequence.
Returns:
List[Block]: A new list of blocks that shares the same memory as the
original sequence.
"""
# do not attempt to fork the null block
assert not isinstance(last_block, NullBlock)
block_id = last_block.block_id
assert block_id is not None
allocator = self._block_ids_to_allocator[block_id]
return allocator.fork(last_block)
def get_num_free_blocks(self, device: Device) -> int:
"""Returns the number of free blocks available on the specified device.
Args:
device (Device): The device for which to query the number of free
blocks. AssertionError is raised if None is passed.
Returns:
int: The number of free blocks available on the specified device.
"""
return self._allocators[device].get_num_free_blocks()
def get_num_total_blocks(self, device: Device) -> int:
return self._allocators[device].get_num_total_blocks()
def get_physical_block_id(self, device: Device, absolute_id: int) -> int:
"""Returns the zero-offset block id on certain device given the
absolute block id.
Args:
device (Device): The device for which to query relative block id.
absolute_id (int): The absolute block id for the block in
whole allocator.
Returns:
int: The zero-offset block id on certain device.
"""
return self._allocators[device].get_physical_block_id(absolute_id)
def swap(self, blocks: List[Block], src_device: Device,
dst_device: Device) -> Dict[int, int]:
"""Execute the swap for the given blocks from source_device
on to dest_device, save the current swap mapping and append
them to the accumulated `self._swap_mapping` for each
scheduling move.
Args:
blocks: List of blocks to be swapped.
src_device (Device): Device to swap the 'blocks' from.
dst_device (Device): Device to swap the 'blocks' to.
Returns:
Dict[int, int]: Swap mapping from source_device
on to dest_device.
"""
if self.content_offload_enabled:
raise RuntimeError(
"request-level preemption swap cannot share CPU slots with "
"BI100_CPU_KV_OFFLOAD")
src_block_ids = [block.block_id for block in blocks]
self._allocators[src_device].swap_out(blocks)
self._allocators[dst_device].swap_in(blocks)
dst_block_ids = [block.block_id for block in blocks]
current_swap_mapping: Dict[int, int] = {}
for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids):
if src_block_id is not None and dst_block_id is not None:
self._swap_mapping[src_block_id] = dst_block_id
current_swap_mapping[src_block_id] = dst_block_id
return current_swap_mapping
def get_num_full_blocks_touched(self, blocks: List[Block],
device: Device) -> int:
"""Returns the number of full blocks that will be touched by
swapping in/out the given blocks on to the 'device'.
Args:
blocks: List of blocks to be swapped.
device (Device): Device to swap the 'blocks' on.
Returns:
int: the number of full blocks that will be touched by
swapping in/out the given blocks on to the 'device'.
Non full blocks are ignored when deciding the number
of blocks to touch.
"""
return self._allocators[device].get_num_full_blocks_touched(blocks)
def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
"""Clears the copy-on-write (CoW) state and returns the mapping of
source to destination block IDs.
Returns:
List[Tuple[int, int]]: A list mapping source block IDs to
destination block IDs.
"""
# CoW only supported on GPU
device = Device.GPU
return self._allocators[device].clear_copy_on_writes()
def mark_blocks_as_accessed(self, block_ids: List[int],
now: float) -> None:
"""Mark blocks as accessed, only use for prefix caching."""
# Prefix caching only supported on GPU.
device = Device.GPU
return self._allocators[device].mark_blocks_as_accessed(block_ids, now)
def mark_blocks_as_computed(self, block_ids: List[int]) -> None:
"""Mark blocks as accessed, only use for prefix caching."""
# Prefix caching only supported on GPU.
device = Device.GPU
return self._allocators[device].mark_blocks_as_computed(block_ids)
def get_computed_block_ids(self, prev_computed_block_ids: List[int],
block_ids: List[int],
skip_last_block_id: bool) -> List[int]:
# Prefix caching only supported on GPU.
device = Device.GPU
return self._allocators[device].get_computed_block_ids(
prev_computed_block_ids, block_ids, skip_last_block_id)
def get_common_computed_block_ids(
self, computed_seq_block_ids: List[List[int]]) -> List[int]:
# Prefix caching only supported on GPU.
device = Device.GPU
return self._allocators[device].get_common_computed_block_ids(
computed_seq_block_ids)
@property
def all_block_ids(self) -> FrozenSet[int]:
return frozenset(self._block_ids_to_allocator.keys())
def get_prefix_cache_hit_rate(self, device: Device) -> float:
"""Prefix cache hit rate. -1 means not supported or disabled."""
assert device in self._allocators
return self._allocators[device].get_prefix_cache_hit_rate()
def get_and_reset_swaps(self) -> List[Tuple[int, int]]:
"""Returns and clears the mapping of source to destination block IDs.
Will be called after every swapping operations for now, and after every
schedule when BlockManagerV2 become default. Currently not useful.
Returns:
List[Tuple[int, int]]: A mapping of source to destination block IDs.
"""
mapping = self._swap_mapping.copy()
self._swap_mapping.clear()
return list(mapping.items())
def get_and_reset_prefix_swaps(
self) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
"""Return scheduler-owned (CPU->GPU, GPU->CPU) content maps."""
if self._cpu_content_cache is None:
return [], []
return self._cpu_content_cache.drain_step()
def begin_prefix_cache_step(self) -> None:
if self._cpu_content_cache is not None:
self._cpu_content_cache.begin_step()
class NullBlock(Block):
"""
Null blocks are used as a placeholders for KV cache blocks that have
been dropped due to sliding window.
This implementation just wraps an ordinary block and prevents it from
being modified. It also allows for testing if a block is NullBlock
via isinstance().
"""
def __init__(self, proxy: Block):
super().__init__()
self._proxy = proxy
def append_token_ids(self, token_ids: List[BlockId]):
raise ValueError("null block should not be modified")
@property
def block_id(self):
return self._proxy.block_id
@block_id.setter
def block_id(self, value: Optional[BlockId]):
raise ValueError("null block should not be modified")
@property
def token_ids(self) -> List[BlockId]:
return self._proxy.token_ids
@property
def num_tokens_total(self) -> int:
raise NotImplementedError(
"num_tokens_total is not used for null block")
@property
def num_empty_slots(self) -> BlockId:
return self._proxy.num_empty_slots
@property
def is_full(self):
return self._proxy.is_full
@property
def prev_block(self):
return self._proxy.prev_block
@property
def computed(self):
return self._proxy.computed
@computed.setter
def computed(self, value):
self._proxy.computed = value
@property
def last_accessed(self) -> float:
return self._proxy.last_accessed
@last_accessed.setter
def last_accessed(self, last_accessed_ts: float):
self._proxy.last_accessed = last_accessed_ts
@property
def content_hash(self):
return self._proxy.content_hash

View File

@@ -0,0 +1,255 @@
"""Scheduler-owned content index for an inclusive CPU KV cache tier."""
from __future__ import annotations
import heapq
import os
from collections import OrderedDict
from typing import Dict, List, Mapping, Optional, Set, Tuple
ContentHash = bytes
SwapMapping = List[Tuple[int, int]]
def cpu_kv_offload_enabled(
environ: Optional[Mapping[str, str]] = None,
) -> bool:
"""Read the experimental selector without accepting ambiguous values."""
source = os.environ if environ is None else environ
value = source.get("BI100_CPU_KV_OFFLOAD", "0")
if value == "0":
return False
if value == "1":
return True
raise RuntimeError(
"BI100_CPU_KV_OFFLOAD must be exactly '0' or '1', "
f"got {value!r}")
class CpuKvContentCache:
"""Track immutable KV blocks held in the worker's pinned CPU cache.
The scheduler owns this metadata and sends identical physical block maps
to every tensor-parallel worker. CPU copies are inclusive: loading a block
back to GPU does not remove its CPU entry. Slots touched by either transfer
direction are pinned for the whole scheduling step so a D2H destination
can never overwrite an H2D source before workers execute the maps.
"""
def __init__(self, capacity: int) -> None:
if capacity <= 0:
raise ValueError("CPU KV content cache capacity must be positive")
self.capacity = capacity
self._hash_to_slot: Dict[ContentHash, int] = {}
self._slot_to_hash: Dict[int, ContentHash] = {}
self._ready_slots: Set[int] = set()
self._lru: OrderedDict[int, None] = OrderedDict()
self._free_slots = list(range(capacity))
heapq.heapify(self._free_slots)
self._step_slots_in_use: Set[int] = set()
self._step_load_slots: Set[int] = set()
self._step_h2d: Dict[int, int] = {}
self._step_d2h: Dict[int, int] = {}
self._deferred_d2h: Dict[int, ContentHash] = {}
self._deferred_hashes: Set[ContentHash] = set()
self._pending_ready_slots: Set[int] = set()
self.hits = 0
self.misses = 0
self.stores = 0
self.deduplicated_stores = 0
self.evictions = 0
self.skipped_stores = 0
@staticmethod
def _validate_hash(content_hash: ContentHash) -> None:
if not isinstance(content_hash, bytes) or len(content_hash) != 32:
raise ValueError("CPU KV cache key must be a 32-byte content hash")
@staticmethod
def _validate_block_id(name: str, block_id: int) -> None:
if not isinstance(block_id, int) or isinstance(block_id, bool):
raise TypeError(f"{name} must be an integer")
if block_id < 0:
raise ValueError(f"{name} must be non-negative")
def _touch(self, slot: int) -> None:
self._lru.pop(slot, None)
self._lru[slot] = None
def _select_store_slot(self) -> Optional[int]:
if self._free_slots:
return heapq.heappop(self._free_slots)
for slot in self._lru:
if slot not in self._step_slots_in_use:
return slot
return None
def _commit_store(self, content_hash: ContentHash,
gpu_block: int, slot: int) -> None:
old_hash = self._slot_to_hash.get(slot)
if old_hash is not None:
if slot in self._step_slots_in_use:
raise RuntimeError("selected an in-use CPU KV slot for eviction")
del self._hash_to_slot[old_hash]
self._ready_slots.discard(slot)
self.evictions += 1
if slot in self._step_h2d:
raise RuntimeError(
"a CPU KV slot cannot be an H2D source and D2H destination "
"in one scheduler step")
if slot in self._step_d2h.values():
raise RuntimeError(f"duplicate D2H destination CPU slot {slot}")
self._hash_to_slot[content_hash] = slot
self._slot_to_hash[slot] = content_hash
self._ready_slots.discard(slot)
self._step_slots_in_use.add(slot)
self._step_d2h[gpu_block] = slot
self._touch(slot)
self.stores += 1
def begin_step(self) -> None:
"""Publish D2H stores returned by the preceding synchronous step."""
if (self._step_slots_in_use or self._step_h2d or self._step_d2h
or self._deferred_d2h or self._deferred_hashes):
raise RuntimeError("cannot begin a CPU KV step before draining it")
self._ready_slots.update(self._pending_ready_slots)
self._pending_ready_slots.clear()
def _require_step_started(self) -> None:
if self._pending_ready_slots:
raise RuntimeError(
"CPU KV step must begin before content lookup or eviction")
def claim_load(self, content_hash: ContentHash) -> Optional[int]:
"""Pin and return a ready CPU source for this scheduling step."""
self._validate_hash(content_hash)
self._require_step_started()
slot = self._hash_to_slot.get(content_hash)
if slot is None or slot not in self._ready_slots:
self.misses += 1
return None
if slot in self._step_slots_in_use:
raise RuntimeError(
f"CPU KV slot {slot} was claimed twice in one scheduler step")
self._step_slots_in_use.add(slot)
self._step_load_slots.add(slot)
self._touch(slot)
self.hits += 1
return slot
def cancel_load(self, content_hash: ContentHash, cpu_slot: int) -> None:
"""Release a claim when GPU allocation fails before H2D is staged."""
self._validate_hash(content_hash)
self._validate_block_id("cpu_slot", cpu_slot)
if self._hash_to_slot.get(content_hash) != cpu_slot:
raise RuntimeError("CPU KV load cancellation key/slot mismatch")
if cpu_slot in self._step_h2d:
raise RuntimeError("cannot cancel a CPU KV load after H2D staging")
if cpu_slot not in self._step_slots_in_use:
raise RuntimeError("cannot cancel an unclaimed CPU KV load")
self._step_slots_in_use.remove(cpu_slot)
self._step_load_slots.remove(cpu_slot)
def stage_load(self, content_hash: ContentHash, cpu_slot: int,
gpu_block: int) -> None:
"""Stage one CPU-to-GPU promotion after the GPU slot is reserved."""
self._validate_hash(content_hash)
self._validate_block_id("cpu_slot", cpu_slot)
self._validate_block_id("gpu_block", gpu_block)
if self._hash_to_slot.get(content_hash) != cpu_slot:
raise RuntimeError("CPU KV load key/slot mismatch")
if cpu_slot not in self._ready_slots:
raise RuntimeError("CPU KV load source is not ready")
if cpu_slot not in self._step_slots_in_use:
raise RuntimeError("CPU KV load source was not claimed")
if cpu_slot in self._step_h2d:
raise RuntimeError(f"duplicate H2D source CPU slot {cpu_slot}")
if gpu_block in self._step_h2d.values():
raise RuntimeError(f"duplicate H2D destination GPU block {gpu_block}")
if cpu_slot in self._step_d2h.values():
raise RuntimeError(
"a CPU KV slot cannot be an H2D source and D2H destination "
"in one scheduler step")
self._step_h2d[cpu_slot] = gpu_block
def stage_store(self, content_hash: ContentHash,
gpu_block: int) -> bool:
"""Stage a lazy GPU-to-CPU copy for an evicted immutable block."""
self._validate_hash(content_hash)
self._validate_block_id("gpu_block", gpu_block)
self._require_step_started()
if (content_hash in self._hash_to_slot
or content_hash in self._deferred_hashes):
slot = self._hash_to_slot.get(content_hash)
if slot is not None:
self._touch(slot)
self.deduplicated_stores += 1
return False
if gpu_block in self._step_d2h or gpu_block in self._deferred_d2h:
raise RuntimeError(f"duplicate D2H source GPU block {gpu_block}")
if self._free_slots:
self._commit_store(
content_hash, gpu_block, heapq.heappop(self._free_slots))
return True
# Do not replace resident content until every lookup in this scheduler
# step is known. A later H2D claim can refer to any current LRU entry.
self._deferred_d2h[gpu_block] = content_hash
self._deferred_hashes.add(content_hash)
return True
def _resolve_deferred_stores(self) -> None:
if self._step_load_slots:
self.skipped_stores += len(self._deferred_d2h)
else:
for gpu_block, content_hash in self._deferred_d2h.items():
slot = self._select_store_slot()
if slot is None:
self.skipped_stores += 1
continue
self._commit_store(content_hash, gpu_block, slot)
self._deferred_d2h.clear()
self._deferred_hashes.clear()
def drain_step(self) -> Tuple[SwapMapping, SwapMapping]:
"""Finalize this synchronous step and return (H2D, D2H) maps."""
self._resolve_deferred_stores()
transfer_slots = (
set(self._step_h2d) | set(self._step_d2h.values()))
if transfer_slots != self._step_slots_in_use:
raise RuntimeError(
"CPU KV scheduler step contains an uncommitted slot claim")
if set(self._step_h2d) & set(self._step_d2h.values()):
raise RuntimeError(
"CPU KV scheduler step reuses a CPU slot across directions")
if set(self._step_h2d) != self._step_load_slots:
raise RuntimeError(
"CPU KV scheduler step contains an unstaged load claim")
swap_in = sorted(self._step_h2d.items())
swap_out = sorted(self._step_d2h.items())
self._pending_ready_slots.update(self._step_d2h.values())
self._step_h2d.clear()
self._step_d2h.clear()
self._step_slots_in_use.clear()
self._step_load_slots.clear()
return swap_in, swap_out
def resident_slot(self, content_hash: ContentHash) -> Optional[int]:
self._validate_hash(content_hash)
return self._hash_to_slot.get(content_hash)
def is_ready(self, content_hash: ContentHash) -> bool:
self._validate_hash(content_hash)
slot = self._hash_to_slot.get(content_hash)
return slot is not None and slot in self._ready_slots
@property
def resident_count(self) -> int:
return len(self._hash_to_slot)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,769 @@
"""A block manager that manages token blocks."""
import hashlib
import os
import struct
from collections.abc import Mapping
from typing import Any, Dict, List, Optional, Sequence as GenericSequence, Tuple
try:
from PIL import Image
except Exception: # pragma: no cover - optional dependency in some envs
Image = None # type: ignore
try:
import torch
except Exception: # pragma: no cover - optional dependency in some envs
torch = None # type: ignore
from vllm.core.block.block_table import BlockTable
from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
from vllm.core.block.interfaces import Block
from vllm.core.block.prefix_caching_block import (ComputedBlocksTracker,
LastAccessBlocksTracker)
from vllm.core.block.utils import check_no_caching_or_swa_for_blockmgr_encdec
from vllm.core.interfaces import AllocStatus, BlockSpaceManager
from vllm.logger import init_logger
from vllm.sequence import Sequence, SequenceGroup, SequenceStatus
from vllm.utils import Device
SeqId = int
EncoderSeqId = str
logger = init_logger(__name__)
class BlockSpaceManagerV2(BlockSpaceManager):
"""BlockSpaceManager which manages the allocation of KV cache.
It owns responsibility for allocation, swapping, allocating memory for
autoregressively-generated tokens, and other advanced features such as
prefix caching, forking/copy-on-write, and sliding-window memory allocation.
This class implements the design described in
https://github.com/vllm-project/vllm/pull/3492.
Lookahead slots
The block manager has the notion of a "lookahead slot". These are slots
in the KV cache that are allocated for a sequence. Unlike the other
allocated slots, the content of these slots is undefined -- the worker
may use the memory allocations in any way.
In practice, a worker could use these lookahead slots to run multiple
forward passes for a single scheduler invocation. Each successive
forward pass would write KV activations to the corresponding lookahead
slot. This allows low inter-token latency use-cases, where the overhead
of continuous batching scheduling is amortized over >1 generated tokens.
Speculative decoding uses lookahead slots to store KV activations of
proposal tokens.
See https://github.com/vllm-project/vllm/pull/3250 for more information
on lookahead scheduling.
Args:
block_size (int): The size of each memory block.
num_gpu_blocks (int): The number of memory blocks allocated on GPU.
num_cpu_blocks (int): The number of memory blocks allocated on CPU.
watermark (float, optional): The threshold used for memory swapping.
Defaults to 0.01.
sliding_window (Optional[int], optional): The size of the sliding
window. Defaults to None.
enable_caching (bool, optional): Flag indicating whether caching is
enabled. Defaults to False.
"""
def __init__(
self,
block_size: int,
num_gpu_blocks: int,
num_cpu_blocks: int,
watermark: float = 0.01,
sliding_window: Optional[int] = None,
enable_caching: bool = False,
) -> None:
self.block_size = block_size
self.num_total_gpu_blocks = num_gpu_blocks
self.num_total_cpu_blocks = num_cpu_blocks
self.sliding_window = sliding_window
# max_block_sliding_window is the max number of blocks that need to be
# allocated
self.max_block_sliding_window = None
if sliding_window is not None:
# +1 here because // rounds down
num_blocks = sliding_window // block_size + 1
# +1 here because the last block may not be full,
# and so the sequence stretches one more block at the beginning
# For example, if sliding_window is 3 and block_size is 4,
# we may need 2 blocks when the second block only holds 1 token.
self.max_block_sliding_window = num_blocks + 1
self.watermark = watermark
assert watermark >= 0.0
self.enable_caching = enable_caching
self.watermark_blocks = int(watermark * num_gpu_blocks)
self.block_allocator = CpuGpuBlockAllocator.create(
allocator_type="prefix_caching" if enable_caching else "naive",
num_gpu_blocks=num_gpu_blocks,
num_cpu_blocks=num_cpu_blocks,
block_size=block_size,
)
self.block_tables: Dict[SeqId, BlockTable] = {}
self.cross_block_tables: Dict[EncoderSeqId, BlockTable] = {}
self._warned_mm_namespace_requests = set[str]()
self._request_local_namespace: Dict[str, bytes] = {}
self._runtime_cache_namespace = self._build_runtime_cache_namespace()
self._computed_blocks_tracker = ComputedBlocksTracker(
self.block_allocator)
self._last_access_blocks_tracker = LastAccessBlocksTracker(
self.block_allocator)
def can_allocate(self,
seq_group: SequenceGroup,
num_lookahead_slots: int = 0) -> AllocStatus:
# FIXME(woosuk): Here we assume that all sequences in the group share
# the same prompt. This may not be true for preempted sequences.
check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
num_required_blocks = BlockTable.get_num_required_blocks(
seq.get_token_ids(),
block_size=self.block_size,
num_lookahead_slots=num_lookahead_slots,
)
if seq_group.is_encoder_decoder():
encoder_seq = seq_group.get_encoder_seq()
assert encoder_seq is not None
num_required_blocks += BlockTable.get_num_required_blocks(
encoder_seq.get_token_ids(),
block_size=self.block_size,
)
if self.max_block_sliding_window is not None:
num_required_blocks = min(num_required_blocks,
self.max_block_sliding_window)
num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
device=Device.GPU)
# Use watermark to avoid frequent cache eviction.
if (self.num_total_gpu_blocks - num_required_blocks <
self.watermark_blocks):
return AllocStatus.NEVER
if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
return AllocStatus.OK
else:
return AllocStatus.LATER
def _allocate_sequence(
self,
seq: Sequence,
cache_namespace: Optional[bytes] = None,
) -> BlockTable:
block_table = BlockTable(
block_size=self.block_size,
block_allocator=self.block_allocator,
max_block_sliding_window=self.max_block_sliding_window,
cache_namespace=cache_namespace,
)
if seq.get_token_ids():
# Add blocks to the block table only if the sequence is non empty.
block_table.allocate(seq.get_token_ids())
return block_table
def allocate(self, seq_group: SequenceGroup) -> None:
# Allocate self-attention block tables for decoder sequences
waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)
assert not (set(seq.seq_id for seq in waiting_seqs)
& self.block_tables.keys()), "block table already exists"
# NOTE: Here we assume that all sequences in the group have the same
# prompt.
seq = waiting_seqs[0]
request_id = seq_group.request_id
cache_namespace = self._get_cache_namespace(
seq,
request_id=request_id,
seq_group=seq_group,
)
block_table: BlockTable = self._allocate_sequence(
seq,
cache_namespace=cache_namespace,
)
self.block_tables[seq.seq_id] = block_table
# Track seq
self._computed_blocks_tracker.add_seq(seq.seq_id)
self._last_access_blocks_tracker.add_seq(seq.seq_id)
# Assign the block table for each sequence.
for seq in waiting_seqs[1:]:
self.block_tables[seq.seq_id] = block_table.fork()
# Track seq
self._computed_blocks_tracker.add_seq(seq.seq_id)
self._last_access_blocks_tracker.add_seq(seq.seq_id)
# Allocate cross-attention block table for encoder sequence
#
# NOTE: Here we assume that all sequences in the group have the same
# encoder prompt.
request_id = seq_group.request_id
assert (request_id
not in self.cross_block_tables), \
"block table already exists"
check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
if seq_group.is_encoder_decoder():
encoder_seq = seq_group.get_encoder_seq()
assert encoder_seq is not None
encoder_cache_namespace = self._get_cache_namespace(
encoder_seq,
request_id=request_id,
seq_group=seq_group)
block_table = self._allocate_sequence(
encoder_seq, cache_namespace=encoder_cache_namespace)
self.cross_block_tables[request_id] = block_table
@staticmethod
def _has_multi_modal_payload(multi_modal_data: Any) -> bool:
if multi_modal_data is None:
return False
if isinstance(multi_modal_data, Mapping):
try:
return len(multi_modal_data) > 0
except (TypeError, ValueError, RuntimeError, OSError,
OverflowError, AttributeError, LookupError, struct.error):
# Treat an unusual mapping as payload and let normalization
# either identify it or select request-local isolation.
return True
return True
def _get_cache_namespace(self, seq: Sequence, request_id: str,
seq_group: SequenceGroup) -> bytes:
digest = hashlib.sha256()
digest.update(b"bi100-request-prefix-namespace-v1|")
digest.update(self._runtime_cache_namespace)
digest.update(self._adapter_cache_namespace(seq_group))
multi_modal_data = seq.multi_modal_data
if self._has_multi_modal_payload(multi_modal_data):
try:
mm_namespace = self._hash_multi_modal_namespace(
multi_modal_data)
except (TypeError, ValueError, RuntimeError, OSError,
OverflowError, AttributeError, LookupError, struct.error):
if request_id not in self._warned_mm_namespace_requests:
logger.warning(
"Request %s has multimodal input that cannot be "
"normalized for cache namespace hashing. Falling "
"back to "
"request-local namespace isolation.",
request_id,
)
self._warned_mm_namespace_requests.add(request_id)
mm_namespace = self._request_local_fallback_cache_namespace(
request_id=request_id)
digest.update(b"mm|")
digest.update(mm_namespace)
else:
digest.update(b"text|")
return digest.digest()
def can_append_slots(self, seq_group: SequenceGroup,
num_lookahead_slots: int) -> bool:
"""Determine if there is enough space in the GPU KV cache to continue
generation of the specified sequence group.
We use a worst-case heuristic: assume each touched block will require a
new allocation (either via CoW or new block). We can append slots if the
number of touched blocks is less than the number of free blocks.
"Lookahead slots" are slots that are allocated in addition to the slots
for known tokens. The contents of the lookahead slots are not defined.
This is used by speculative decoding when speculating future tokens.
"""
num_touched_blocks = 0
for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
block_table = self.block_tables[seq.seq_id]
num_touched_blocks += (
block_table.get_num_blocks_touched_by_append_slots(
token_ids=block_table.get_unseen_token_ids(
seq.get_token_ids()),
num_lookahead_slots=num_lookahead_slots,
))
num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
Device.GPU)
return num_touched_blocks <= num_free_gpu_blocks
def append_slots(
self,
seq: Sequence,
num_lookahead_slots: int,
) -> List[Tuple[int, int]]:
block_table = self.block_tables[seq.seq_id]
block_table.append_token_ids(
token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()),
num_lookahead_slots=num_lookahead_slots,
num_computed_slots=seq.data.get_num_computed_tokens(),
)
# Return any new copy-on-writes.
new_cows = self.block_allocator.clear_copy_on_writes()
return new_cows
def free(self, seq: Sequence) -> None:
seq_id = seq.seq_id
if seq_id not in self.block_tables:
# Already freed or haven't been scheduled yet.
return
# Update seq block ids with the latest access time
self._last_access_blocks_tracker.update_seq_blocks_last_access(
seq_id, self.block_tables[seq.seq_id].physical_block_ids)
# Untrack seq
self._last_access_blocks_tracker.remove_seq(seq_id)
self._computed_blocks_tracker.remove_seq(seq_id)
# Free table/blocks
self.block_tables[seq_id].free()
del self.block_tables[seq_id]
def free_cross(self, seq_group: SequenceGroup) -> None:
request_id = seq_group.request_id
if request_id not in self.cross_block_tables:
# Already freed or hasn't been scheduled yet.
return
self.cross_block_tables[request_id].free()
del self.cross_block_tables[request_id]
def get_block_table(self, seq: Sequence) -> List[int]:
block_ids = self.block_tables[seq.seq_id].physical_block_ids
return block_ids # type: ignore
def get_cross_block_table(self, seq_group: SequenceGroup) -> List[int]:
request_id = seq_group.request_id
assert request_id in self.cross_block_tables
block_ids = self.cross_block_tables[request_id].physical_block_ids
assert all(b is not None for b in block_ids)
return block_ids # type: ignore
def access_all_blocks_in_seq(self, seq: Sequence, now: float):
if self.enable_caching:
# Record the latest access time for the sequence. The actual update
# of the block ids is deferred to the sequence free(..) call, since
# only during freeing of block ids, the blocks are actually added to
# the evictor (which is when the most updated time is required)
# (This avoids expensive calls to mark_blocks_as_accessed(..))
self._last_access_blocks_tracker.update_last_access(
seq.seq_id, now)
def mark_blocks_as_computed(self, seq_group: SequenceGroup,
token_chunk_size: int):
# If prefix caching is enabled, mark immutable blocks as computed
# right after they have been scheduled (for prefill). This assumes
# the scheduler is synchronous so blocks are actually computed when
# scheduling the next batch.
self.block_allocator.mark_blocks_as_computed([])
def get_common_computed_block_ids(
self, seqs: List[Sequence]) -> GenericSequence[int]:
"""Determine which blocks for which we skip prefill.
With prefix caching we can skip prefill for previously-generated blocks.
Currently, the attention implementation only supports skipping cached
blocks if they are a contiguous prefix of cached blocks.
This method determines which blocks can be safely skipped for all
sequences in the sequence group.
"""
computed_seq_block_ids = []
for seq in seqs:
computed_seq_block_ids.append(
self._computed_blocks_tracker.
get_cached_computed_blocks_and_update(
seq.seq_id,
self.block_tables[seq.seq_id].physical_block_ids))
# NOTE(sang): This assumes seq_block_ids doesn't contain any None.
return self.block_allocator.get_common_computed_block_ids(
computed_seq_block_ids) # type: ignore
def get_content_hashes(self, seq: Sequence) -> List[bytes]:
return self.block_tables[seq.seq_id].get_content_hashes()
def get_and_reset_prefix_swaps(
self) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
"""Return scheduler-owned (CPU->GPU, GPU->CPU) content transfers."""
return self.block_allocator.get_and_reset_prefix_swaps()
def begin_prefix_cache_step(self) -> None:
self.block_allocator.begin_prefix_cache_step()
def _build_runtime_cache_namespace(self) -> bytes:
"""Bind first-block hashes to the fixed model runtime identity."""
model = os.getenv("BI100_PREFIX_MODEL_FINGERPRINT",
"Qwen3.6-35B-A3B")
dtype = os.getenv("BI100_PREFIX_DTYPE", "float16")
tp_raw = os.getenv("BI100_PREFIX_TP_SIZE", "4")
try:
tp_size = int(tp_raw)
except ValueError as exc:
raise RuntimeError(
"BI100_PREFIX_TP_SIZE must be a positive integer") from exc
if tp_size <= 0:
raise RuntimeError(
"BI100_PREFIX_TP_SIZE must be a positive integer")
digest = hashlib.sha256()
digest.update(b"bi100-runtime-prefix-identity-v1|")
for label, value in (
(b"model", model),
(b"dtype", dtype),
(b"tp", str(tp_size)),
(b"block_size", str(self.block_size))):
encoded = value.encode("utf-8")
digest.update(label)
digest.update(struct.pack("!Q", len(encoded)))
digest.update(encoded)
return digest.digest()
@staticmethod
def _adapter_cache_namespace(seq_group: SequenceGroup) -> bytes:
digest = hashlib.sha256()
digest.update(b"bi100-adapter-prefix-identity-v1|")
lora = getattr(seq_group, "lora_request", None)
prompt_adapter = getattr(seq_group, "prompt_adapter_request", None)
identities = (
("lora", lora, ("lora_name", "lora_int_id", "lora_path",
"base_model_name")),
("prompt", prompt_adapter,
("prompt_adapter_name", "prompt_adapter_id",
"prompt_adapter_local_path",
"prompt_adapter_num_virtual_tokens")),
)
for kind, adapter, fields in identities:
digest.update(kind.encode("ascii"))
if adapter is None:
digest.update(b"none|")
continue
for field in fields:
value = str(getattr(adapter, field, ""))
encoded = value.encode("utf-8")
digest.update(field.encode("ascii"))
digest.update(struct.pack("!Q", len(encoded)))
digest.update(encoded)
return digest.digest()
def _request_local_fallback_cache_namespace(self,
request_id: str) -> bytes:
namespace = self._request_local_namespace.get(request_id)
if namespace is None:
digest = hashlib.sha256()
digest.update(b"multimodal-unsupported-request-local-v1|")
digest.update(self._runtime_cache_namespace)
digest.update(os.urandom(32))
digest.update(request_id.encode("utf-8"))
namespace = digest.digest()
self._request_local_namespace[request_id] = namespace
return namespace
def release_request_cache_namespace(self, request_id: str) -> None:
"""Release request-local isolation state after request completion."""
self._request_local_namespace.pop(request_id, None)
self._warned_mm_namespace_requests.discard(request_id)
def _hash_multi_modal_namespace(self, mm_data: Any) -> bytes:
digest = hashlib.sha256()
self._hash_multi_modal_obj(digest, mm_data)
return digest.digest()
@staticmethod
def _sort_map_keys(mm_map: Mapping[Any, Any]) -> List[Any]:
return sorted(mm_map.keys(), key=lambda key: repr(key))
@classmethod
def _hash_multi_modal_obj(cls, digest: Any, value: Any) -> None:
if value is None:
digest.update(b"none|")
return
if isinstance(value, Mapping):
digest.update(b"map|")
digest.update(struct.pack("!Q", len(value)))
for key in cls._sort_map_keys(value):
digest.update(b"k|")
cls._hash_multi_modal_obj(digest, key)
digest.update(b"v|")
cls._hash_multi_modal_obj(digest, value[key])
return
if isinstance(value, list):
digest.update(b"list|")
digest.update(struct.pack("!Q", len(value)))
for item in value:
cls._hash_multi_modal_obj(digest, item)
return
if isinstance(value, tuple):
digest.update(b"tuple|")
digest.update(struct.pack("!Q", len(value)))
for item in value:
cls._hash_multi_modal_obj(digest, item)
return
if isinstance(value, str):
encoded = value.encode()
digest.update(b"str|")
digest.update(struct.pack("!Q", len(encoded)))
digest.update(encoded)
return
if isinstance(value, bytes):
digest.update(b"bytes|")
digest.update(struct.pack("!Q", len(value)))
digest.update(value)
return
if isinstance(value, bytearray):
cls._hash_multi_modal_obj(digest, bytes(value))
return
if isinstance(value, bool):
digest.update(b"bool|")
digest.update(b"1" if value else b"0")
return
if isinstance(value, int):
digest.update(b"int|")
digest.update(str(value).encode())
return
if isinstance(value, float):
digest.update(b"float|")
digest.update(struct.pack("!d", value))
return
if torch is not None and isinstance(value, torch.Tensor):
digest.update(b"tensor|")
tensor = value.detach().cpu().contiguous()
digest.update(struct.pack("!Q", len(tensor.shape)))
for dim in tensor.shape:
digest.update(struct.pack("!Q", int(dim)))
digest.update(str(tensor.dtype).encode())
# Byte views work for bfloat16 and other dtypes that NumPy cannot
# materialize directly.
tensor_bytes = tensor.view(torch.uint8).numpy().tobytes()
digest.update(struct.pack("!Q", len(tensor_bytes)))
digest.update(tensor_bytes)
return
if Image is not None and isinstance(value, Image.Image):
digest.update(b"image|")
digest.update(value.mode.encode())
digest.update(struct.pack("!II", value.width, value.height))
image_bytes = value.tobytes()
digest.update(struct.pack("!Q", len(image_bytes)))
digest.update(image_bytes)
palette = value.getpalette()
digest.update(b"palette-mode|")
cls._hash_multi_modal_obj(
digest, getattr(getattr(value, "palette", None), "mode", None))
digest.update(b"palette|")
cls._hash_multi_modal_obj(digest, palette)
digest.update(b"transparency|")
cls._hash_multi_modal_obj(
digest, value.info.get("transparency"))
return
raise TypeError(f"Unsupported multimodal namespace value type {type(value)}")
def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
if parent_seq.seq_id not in self.block_tables:
# Parent sequence has either been freed or never existed.
return
src_block_table = self.block_tables[parent_seq.seq_id]
self.block_tables[child_seq.seq_id] = src_block_table.fork()
# Track child seq
self._computed_blocks_tracker.add_seq(child_seq.seq_id)
self._last_access_blocks_tracker.add_seq(child_seq.seq_id)
def can_swap_in(self, seq_group: SequenceGroup,
num_lookahead_slots: int) -> AllocStatus:
"""Returns the AllocStatus for the given sequence_group
with num_lookahead_slots.
Args:
sequence_group (SequenceGroup): The sequence group to swap in.
num_lookahead_slots (int): Number of lookahead slots used in
speculative decoding, default to 0.
Returns:
AllocStatus: The AllocStatus for the given sequence group.
"""
if self.block_allocator.content_offload_enabled:
return AllocStatus.NEVER
return self._can_swap(seq_group, Device.GPU, SequenceStatus.SWAPPED,
num_lookahead_slots)
def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:
"""Returns the block id mapping (from CPU to GPU) generated by
swapping in the given seq_group with num_lookahead_slots.
Args:
seq_group (SequenceGroup): The sequence group to swap in.
Returns:
List[Tuple[int, int]]: The mapping of swapping block from CPU
to GPU.
"""
physical_block_id_mapping = []
for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
blocks = self.block_tables[seq.seq_id].blocks
if len(blocks) == 0:
continue
seq_swap_mapping = self.block_allocator.swap(blocks=blocks,
src_device=Device.CPU,
dst_device=Device.GPU)
# Refresh the block ids of the table (post-swap)
self.block_tables[seq.seq_id].update(blocks)
seq_physical_block_id_mapping = {
self.block_allocator.get_physical_block_id(
Device.CPU, cpu_block_id):
self.block_allocator.get_physical_block_id(
Device.GPU, gpu_block_id)
for cpu_block_id, gpu_block_id in seq_swap_mapping.items()
}
physical_block_id_mapping.extend(
list(seq_physical_block_id_mapping.items()))
return physical_block_id_mapping
def can_swap_out(self, seq_group: SequenceGroup) -> bool:
"""Returns whether we can swap out the given sequence_group
with num_lookahead_slots.
Args:
seq_group (SequenceGroup): The sequence group to swap in.
num_lookahead_slots (int): Number of lookahead slots used in
speculative decoding, default to 0.
Returns:
bool: Whether it's possible to swap out current sequence group.
"""
if self.block_allocator.content_offload_enabled:
return False
alloc_status = self._can_swap(seq_group, Device.CPU,
SequenceStatus.RUNNING)
return alloc_status == AllocStatus.OK
def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:
"""Returns the block id mapping (from GPU to CPU) generated by
swapping out the given sequence_group with num_lookahead_slots.
Args:
sequence_group (SequenceGroup): The sequence group to swap in.
Returns:
List[Tuple[int, int]]: The mapping of swapping block from
GPU to CPU.
"""
physical_block_id_mapping = []
for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
blocks = self.block_tables[seq.seq_id].blocks
if len(blocks) == 0:
continue
seq_swap_mapping = self.block_allocator.swap(blocks=blocks,
src_device=Device.GPU,
dst_device=Device.CPU)
# Refresh the block ids of the table (post-swap)
self.block_tables[seq.seq_id].update(blocks)
seq_physical_block_id_mapping = {
self.block_allocator.get_physical_block_id(
Device.GPU, gpu_block_id):
self.block_allocator.get_physical_block_id(
Device.CPU, cpu_block_id)
for gpu_block_id, cpu_block_id in seq_swap_mapping.items()
}
physical_block_id_mapping.extend(
list(seq_physical_block_id_mapping.items()))
return physical_block_id_mapping
def get_num_free_gpu_blocks(self) -> int:
return self.block_allocator.get_num_free_blocks(Device.GPU)
def get_num_free_cpu_blocks(self) -> int:
return self.block_allocator.get_num_free_blocks(Device.CPU)
def get_prefix_cache_hit_rate(self, device: Device) -> float:
return self.block_allocator.get_prefix_cache_hit_rate(device)
def _can_swap(self,
seq_group: SequenceGroup,
device: Device,
status: SequenceStatus,
num_lookahead_slots: int = 0) -> AllocStatus:
"""Returns the AllocStatus for swapping in/out the given sequence_group
on to the 'device'.
Args:
sequence_group (SequenceGroup): The sequence group to swap in.
device (Device): device to swap the 'seq_group' on.
status (SequenceStatus): The status of sequence which is needed
for action. RUNNING for swap out and SWAPPED for swap in
num_lookahead_slots (int): Number of lookahead slots used in
speculative decoding, default to 0.
Returns:
AllocStatus: The AllocStatus for swapping in/out the given
sequence_group on to the 'device'.
"""
# First determine the number of blocks that will be touched by this
# swap. Then verify if there are available blocks in the device
# to perform the swap.
num_blocks_touched = 0
blocks: List[Block] = []
for seq in seq_group.get_seqs(status=status):
block_table = self.block_tables[seq.seq_id]
if block_table.blocks is not None:
# Compute the number blocks to touch for the tokens to be
# appended. This does NOT include the full blocks that need
# to be touched for the swap.
num_blocks_touched += \
block_table.get_num_blocks_touched_by_append_slots(
block_table.get_unseen_token_ids(seq.get_token_ids()),
num_lookahead_slots=num_lookahead_slots)
blocks.extend(block_table.blocks)
# Compute the number of full blocks to touch and add it to the
# existing count of blocks to touch.
num_blocks_touched += self.block_allocator.get_num_full_blocks_touched(
blocks, device=device)
watermark_blocks = 0
if device == Device.GPU:
watermark_blocks = self.watermark_blocks
if self.block_allocator.get_num_total_blocks(
device) < num_blocks_touched:
return AllocStatus.NEVER
elif self.block_allocator.get_num_free_blocks(
device) - num_blocks_touched >= watermark_blocks:
return AllocStatus.OK
else:
return AllocStatus.LATER

View File

@@ -0,0 +1,272 @@
import enum
import heapq
import os
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import Dict, List, OrderedDict, Tuple
ContentHash = bytes
class EvictionPolicy(enum.Enum):
"""Enum for eviction policy used by make_evictor to instantiate the correct
Evictor subclass.
"""
LRU = enum.auto()
FREQUENCY_AWARE = enum.auto()
class Evictor(ABC):
"""The Evictor subclasses should be used by the BlockAllocator class to
handle eviction of freed PhysicalTokenBlocks.
"""
@abstractmethod
def __init__(self):
pass
@abstractmethod
def __contains__(self, block_id: int) -> bool:
pass
@abstractmethod
def evict(self) -> Tuple[int, ContentHash]:
"""Runs the eviction algorithm and returns the evicted block's
content hash along with physical block id along with physical block id
"""
pass
@abstractmethod
def add(self, block_id: int, content_hash: ContentHash,
num_hashed_tokens: int,
last_accessed: float):
"""Adds block to the evictor, making it a candidate for eviction"""
pass
@abstractmethod
def update(self, block_id: int, last_accessed: float):
"""Update corresponding block's access time in metadata"""
pass
@abstractmethod
def remove(self, block_id: int):
"""Remove a given block id from the cache."""
pass
@property
@abstractmethod
def num_blocks(self) -> int:
pass
class BlockMetaData():
"""Data structure for storing key data describe cached block, so that
evitor could use to make its decision which one to choose for eviction
Here we use physical block id as the dict key, as there maybe several
blocks with the same content hash, but their physical id is unique.
"""
def __init__(self, content_hash: ContentHash, num_hashed_tokens: int,
last_accessed: float):
self.content_hash = content_hash
self.num_hashed_tokens = num_hashed_tokens
self.last_accessed = last_accessed
class LRUEvictor(Evictor):
"""Evicts in a least-recently-used order using the last_accessed timestamp
that's recorded in the PhysicalTokenBlock. If there are multiple blocks with
the same last_accessed time, then the one with the largest num_hashed_tokens
will be evicted. If two blocks each have the lowest last_accessed time and
highest num_hashed_tokens value, then one will be chose arbitrarily
"""
def __init__(self):
self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict()
def __contains__(self, block_id: int) -> bool:
return block_id in self.free_table
def evict(self) -> Tuple[int, ContentHash]:
if len(self.free_table) == 0:
raise ValueError("No usable cache memory left")
evicted_block, evicted_block_id = None, None
# The blocks with the lowest timestamps should be placed consecutively
# at the start of OrderedDict. Loop through all these blocks to
# find the one with maximum number of hashed tokens.
for _id, block in self.free_table.items():
if evicted_block is None:
evicted_block, evicted_block_id = block, _id
continue
if evicted_block.last_accessed < block.last_accessed:
break
if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
evicted_block, evicted_block_id = block, _id
assert evicted_block is not None
assert evicted_block_id is not None
self.free_table.pop(evicted_block_id)
return evicted_block_id, evicted_block.content_hash
def add(self, block_id: int, content_hash: ContentHash,
num_hashed_tokens: int,
last_accessed: float):
self.free_table[block_id] = BlockMetaData(content_hash,
num_hashed_tokens,
last_accessed)
def update(self, block_id: int, last_accessed: float):
self.free_table[block_id].last_accessed = last_accessed
def remove(self, block_id: int):
if block_id not in self.free_table:
raise ValueError(
"Attempting to remove block that's not in the evictor")
self.free_table.pop(block_id)
@property
def num_blocks(self) -> int:
return len(self.free_table)
class FrequencyAwareEvictor(Evictor):
"""Evict the least frequently reused logical prefix content first.
Content frequency survives physical block reuse. Heap entries carry a
generation and are lazily invalidated so eviction remains O(log N) without
allowing stale entries to grow without bound.
"""
_COMPACTION_FACTOR = 2
_COMPACTION_SLACK = 1
def __init__(self):
self.free_table: Dict[int, BlockMetaData] = {}
self.frequency_by_hash: Dict[ContentHash, int] = {}
self._heap: List[Tuple[int, float, int, int, int]] = []
self._generations: Dict[int, int] = {}
self._next_generation = 0
@staticmethod
def _validate_content_hash(content_hash: ContentHash) -> None:
if not isinstance(content_hash, bytes) or len(content_hash) != 32:
raise ValueError(
"frequency-aware eviction requires a 32-byte content hash")
def __contains__(self, block_id: int) -> bool:
return block_id in self.free_table
def _heap_key(self, block_id: int, block: BlockMetaData,
generation: int) -> Tuple[int, float, int, int, int]:
return (
self.frequency_by_hash[block.content_hash],
block.last_accessed,
-block.num_hashed_tokens,
block_id,
generation,
)
def _push(self, block_id: int) -> None:
self._next_generation += 1
generation = self._next_generation
self._generations[block_id] = generation
heapq.heappush(
self._heap,
self._heap_key(
block_id, self.free_table[block_id], generation),
)
def _compact_if_needed(self) -> None:
limit = (
self._COMPACTION_FACTOR * len(self.free_table)
+ self._COMPACTION_SLACK
)
if len(self._heap) <= limit:
return
self._heap = [
self._heap_key(block_id, block, self._generations[block_id])
for block_id, block in self.free_table.items()
]
heapq.heapify(self._heap)
def evict(self) -> Tuple[int, ContentHash]:
if not self.free_table:
raise ValueError("No usable cache memory left")
while self._heap:
entry = heapq.heappop(self._heap)
frequency, _, _, block_id, generation = entry
block = self.free_table.get(block_id)
if (
block is None
or self._generations.get(block_id) != generation
):
continue
if self.frequency_by_hash[block.content_hash] != frequency:
heapq.heappush(
self._heap,
self._heap_key(block_id, block, generation),
)
continue
block = self.free_table.pop(block_id)
self._generations.pop(block_id)
self._compact_if_needed()
return block_id, block.content_hash
raise RuntimeError("Evictor heap has no usable entry")
def add(self, block_id: int, content_hash: ContentHash,
num_hashed_tokens: int, last_accessed: float):
self._validate_content_hash(content_hash)
self.frequency_by_hash[content_hash] = (
self.frequency_by_hash.get(content_hash, 0) + 1)
self.free_table[block_id] = BlockMetaData(
content_hash, num_hashed_tokens, last_accessed)
self._push(block_id)
self._compact_if_needed()
def update(self, block_id: int, last_accessed: float):
self.free_table[block_id].last_accessed = last_accessed
self._push(block_id)
self._compact_if_needed()
def remove(self, block_id: int):
if block_id not in self.free_table:
raise ValueError(
"Attempting to remove block that's not in the evictor")
self.free_table.pop(block_id)
self._generations.pop(block_id)
self._compact_if_needed()
@property
def num_blocks(self) -> int:
return len(self.free_table)
def eviction_policy_from_env(
environ: Mapping[str, str] | None = None,
) -> EvictionPolicy:
source = os.environ if environ is None else environ
value = source.get("BI100_KV_EVICTION_POLICY", "lru").strip().lower()
policies = {
"lru": EvictionPolicy.LRU,
"frequency": EvictionPolicy.FREQUENCY_AWARE,
}
if value not in policies:
raise ValueError(
"BI100_KV_EVICTION_POLICY must be one of: frequency, lru")
return policies[value]
def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:
if eviction_policy == EvictionPolicy.LRU:
return LRUEvictor()
elif eviction_policy == EvictionPolicy.FREQUENCY_AWARE:
return FrequencyAwareEvictor()
else:
raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,644 @@
from array import array
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import torch
from vllm.sampling_params import SamplingParams, SamplingType
from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData,
SequenceGroupMetadata)
from vllm.utils import (PyObjectCache, async_tensor_h2d,
is_pin_memory_available, make_tensor_with_pad)
_SAMPLING_EPS = 1e-5
@dataclass
class SequenceGroupToSample:
# |---------- N-1 iteration --------|
# |---------------- N iteration ---------------------|
# |- tokenA -|......................|-- newTokens ---|
# |---------- context_len ----------|
# |-------------------- seq_len ----------------------|
# |-- query_len ---|
# Sequence ids for the sequence group in a previous step.
seq_ids: List[int]
sampling_params: SamplingParams
# seq_id -> sequence data.
seq_data: Dict[int, SequenceData]
# The length of the sequence (all tokens seen in the past + new token to
# compute attention) of the sequence group. None if it is in a decode
# stage.
seq_len: Optional[int]
# The length of new query tokens to compute in the current step. None if it
# is in a decode stage. The length of query_len <= seq_len if chunked
# prefill is enabled.
query_len: Optional[int]
# A random number generator for sampling.
generator: Optional[torch.Generator]
# True if the sequence group is in prefill stage. False if it is in a
# decode stage.
is_prompt: bool
# Query token indices from logits. to compute prompt logprob. Empty if
# prompt logprob is not required.
prompt_logprob_indices: List[int]
# Output offsets within this prefill chunk. Sparse diagnostic requests use
# this to retain the standard full-length prompt-logprob response shape.
prompt_logprob_output_indices: List[int]
# Sample token indices from logits. Empty if sampling is not required.
sample_indices: List[int]
@property
def do_sample(self):
return len(self.sample_indices) > 0
def __post_init__(self):
if len(self.prompt_logprob_indices) > 0:
assert self.sampling_params.prompt_logprobs is not None
assert (len(self.prompt_logprob_indices)
== len(self.prompt_logprob_output_indices))
assert self.prompt_logprob_output_indices == sorted(
set(self.prompt_logprob_output_indices))
if self.is_prompt:
assert self.seq_len is not None
assert self.query_len is not None
assert all(
0 <= index < self.query_len
for index in self.prompt_logprob_output_indices)
def gen_seq_group_to_sample_builder(num_seqs: int):
return lambda: SequenceGroupToSample(
seq_ids=[0] * num_seqs,
sampling_params=None,
seq_data=None, # type: ignore
seq_len=0,
query_len=0,
generator=None,
is_prompt=True,
prompt_logprob_indices=[],
prompt_logprob_output_indices=[],
sample_indices=[],
)
class SamplingMetadataCache:
"""Used to cache SamplingMetadata objects between scheduler iterations"""
def __init__(self):
self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {}
def get_cached_seq_group_to_sample(self, num_seqs):
if num_seqs not in self._seq_group_to_sample_cache:
self._seq_group_to_sample_cache[num_seqs] = PyObjectCache(
gen_seq_group_to_sample_builder(num_seqs))
obj = self._seq_group_to_sample_cache[num_seqs].get_object()
return obj
def reset(self):
for cache in self._seq_group_to_sample_cache.values():
cache.reset()
class SamplingMetadata:
"""Metadata for input sequences. Used in sampler.
The usage is as follow;
```
hidden_states = execute_model(...)
logits = hidden_states[sampling_metadata.selected_token_indices]
sample(logits)
def sample(logits):
# Use categorized_sample_indices for sampling....
```
Args:
seq_groups: List of batched sequence groups.
selected_token_indices: (num_query_tokens_to_logprob). Indices to find
logits from the initial model output hidden states.
categorized_sample_indices: SamplingType -> token indices to sample.
Each token indices is 2D tensor of (num_indices, num_indices) where
the first item means the sample index within the returned logit
(before pruning padding), and the second item means the sample
index after pruning using selected_token_indices.
For example, if the returned logit is [1, 2, 3], and we select
[1, 2] for sampling, the pruned logit will be [2, 3]. In this case,
The first tuple is [1, 2] (sampled index within original logit),
and the second tuple is [0, 1] (sampled index within pruned logit).
num_prompts: Number of prompt sequence groups in seq_groups.
skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU
serialization of token outputs.
reuse_sampling_tensors: Indicates if we want to reuse sampling
tensors that are part of the sampler forward pass. Currently,
it is mainly used for multi-step decode.
"""
def __init__(
self,
seq_groups: List[SequenceGroupToSample],
selected_token_indices: torch.Tensor,
categorized_sample_indices: Dict[SamplingType, torch.Tensor],
num_prompts: int,
skip_sampler_cpu_output: bool = False,
reuse_sampling_tensors: bool = False,
) -> None:
self.seq_groups = seq_groups
self.selected_token_indices = selected_token_indices
self.categorized_sample_indices = categorized_sample_indices
self.num_prompts = num_prompts
self.skip_sampler_cpu_output = skip_sampler_cpu_output
self.reuse_sampling_tensors = reuse_sampling_tensors
@staticmethod
def prepare(
seq_group_metadata_list: List[SequenceGroupMetadata],
seq_lens: List[int],
query_lens: List[int],
device: str,
pin_memory: bool,
generators: Optional[Dict[str, torch.Generator]] = None,
cache: Optional[SamplingMetadataCache] = None,
) -> "SamplingMetadata":
(
seq_groups,
selected_token_indices,
categorized_sample_indices,
num_prompts,
) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens,
device, generators, cache)
selected_token_indices = async_tensor_h2d(
selected_token_indices,
dtype=torch.long,
target_device=device,
pin_memory=pin_memory,
)
categorized_sample_indices = {
t: async_tensor_h2d(
seq_ids,
dtype=torch.int,
target_device=device,
pin_memory=pin_memory,
)
for t, seq_ids in categorized_sample_indices.items()
}
sampling_metadata = SamplingMetadata(
seq_groups=seq_groups,
selected_token_indices=selected_token_indices,
categorized_sample_indices=categorized_sample_indices,
num_prompts=num_prompts,
)
return sampling_metadata
def __repr__(self) -> str:
return (
"SamplingMetadata("
f"seq_groups={self.seq_groups}, "
f"selected_token_indices={self.selected_token_indices}, "
f"categorized_sample_indices={self.categorized_sample_indices}), ")
def _get_prompt_logprob_output_indices(
sampling_params: SamplingParams,
seq_data: SequenceData,
prompt_logprob_len: int,
) -> List[int]:
if sampling_params.prompt_logprobs is None or prompt_logprob_len <= 0:
return []
positions = sampling_params.prompt_logprob_positions
computed_len = seq_data.get_num_computed_tokens()
available_next_tokens = max(
0,
len(seq_data.prompt_token_ids) - computed_len - 1,
)
materialized_len = min(prompt_logprob_len, available_next_tokens)
if positions is None:
return list(range(materialized_len))
output_indices = [
position - computed_len - 1
for position in positions
if computed_len < position
<= computed_len + materialized_len
]
assert output_indices == sorted(set(output_indices))
assert all(0 <= index < materialized_len for index in output_indices)
return output_indices
def _prepare_seq_groups(
seq_group_metadata_list: List[SequenceGroupMetadata],
seq_lens: List[int],
query_lens: List[int],
device: str,
generators: Optional[Dict[str, torch.Generator]] = None,
cache: Optional[SamplingMetadataCache] = None,
) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType,
List[int]], int, ]:
"""Prepare sequence groups and indices for sampling.
Args:
seq_group_metadata_list: A list of sequence group to batch.
seq_lens: A list of sequence lens per sequence group.
Index of prompt len should match with seq_group_metadata_list.
query_lens: A list of query lengths. Prompt lens include the length
of entire prompt tokens, and it could be shorter.
device: A device to use for random number generators,
`SequenceGroupToSample.generator`.
generators: A store of per-request random number generators used
for seeded requests.
Returns:
seq_groups: A list of sequence group to sample.
selected_token_indices: See the definition from `SamplingMetadata`.
categorized_sample_indices: See the definition from `SamplingMetadata`.
num_prompts: Total number of prompts from `seq_group_metadata_list`.
"""
# Batched sequence groups for the current model forward stsep.
seq_groups: List[SequenceGroupToSample] = []
# A list of token indices to sample/compute logprob. It is used to
# prune the outcome logits from the model for the performance.
selected_token_indices: List[int] = []
# Used for selected_token_indices.
model_output_idx = 0
# Sampling type -> (
# indices to sample/prompt logprob within pruned output logits,
# indices to sample within pruned logits)
categorized_sample_indices: Dict[SamplingType, List[int]] = {
t: []
for t in SamplingType
}
# Index of logits to compute logprob. Logits include both prompt logprob
# and sample logprob indices.
logit_idx = 0
# Total number of prompts from given sequence groups.
num_prompts = 0
for i, seq_group_metadata in enumerate(seq_group_metadata_list):
seq_ids = seq_group_metadata.seq_data.keys()
if cache is not None:
sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids))
for j, seq_id in enumerate(seq_ids):
sample_obj.seq_ids[j] = seq_id
sample_obj.prompt_logprob_indices.clear()
sample_obj.prompt_logprob_output_indices.clear()
sample_obj.sample_indices.clear()
sampling_params = seq_group_metadata.sampling_params
is_prompt = seq_group_metadata.is_prompt
generator: Optional[torch.Generator] = None
# If the current seq group is in decode stage, it is None.
seq_len: Optional[int] = None
query_len: Optional[int] = None
prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices
if cache is not None else [])
prompt_logprob_output_indices: List[int] = (
sample_obj.prompt_logprob_output_indices
if cache is not None else [])
sample_indices: List[int] = (sample_obj.sample_indices
if cache is not None else [])
do_sample = seq_group_metadata.do_sample
if seq_group_metadata.is_prompt:
if sampling_params.seed is not None:
generator = torch.Generator(device=device).manual_seed(
sampling_params.seed)
if generators is not None:
generators[seq_group_metadata.request_id] = generator
num_prompts += 1
num_prefill_sample = len(seq_ids)
assert num_prefill_sample == 1
assert query_lens is not None and seq_lens is not None
query_len, seq_len = query_lens[i], seq_lens[i]
# If we need sampling, exclude num_prefill_sample tokens from
# prompt logprob.
prompt_logprob_len = (query_len - num_prefill_sample
if do_sample else query_len)
sample_len = num_prefill_sample if do_sample else 0
else:
# Decode
prompt_logprob_len = 0
query_len = query_lens[i] if query_lens is not None else 1
sample_len = len(seq_ids) * query_len if do_sample else 0
if sampling_params.seed is not None and generators is not None:
generator = generators.get(seq_group_metadata.request_id)
seq_data = next(iter(seq_group_metadata.seq_data.values()))
prompt_logprob_output_indices.extend(
_get_prompt_logprob_output_indices(
sampling_params,
seq_data,
prompt_logprob_len,
))
# Update indices to select from the model output.
"""
This blocks computes selected_token_indices which is used in the
following way.
hidden_states = model(...)
logits = hidden_states[selected_token_indices]
"""
if sampling_params.prompt_logprobs is not None:
selected_token_indices.extend(
model_output_idx + output_index
for output_index in prompt_logprob_output_indices)
model_output_idx += prompt_logprob_len
if do_sample:
selected_token_indices.extend(
range(model_output_idx, model_output_idx + sample_len))
model_output_idx += sample_len
# We now find indices for logprob computation and sampling.
"""
This block computes categorized_sample_indices which is used in the
following way.
hidden_states = model(...)
logits = hidden_states[selected_token_indices]
def sample(logits):
# Use categorized_sample_indices for sampling.
# prompt_logprob_indices to find prompt logprob indices.
# sample_indices to find sample indices.
"""
if sampling_params.prompt_logprobs is not None:
prompt_logprob_indices.extend(
range(logit_idx,
logit_idx + len(prompt_logprob_output_indices)))
logit_idx += len(prompt_logprob_output_indices)
if do_sample:
sample_indices.extend(range(logit_idx, logit_idx + sample_len))
categorized_sample_indices[sampling_params.sampling_type].extend(
list(range(logit_idx, logit_idx + sample_len)))
logit_idx += sample_len
if cache is not None:
sample_obj.sampling_params = sampling_params
sample_obj.seq_data = seq_group_metadata.seq_data
sample_obj.seq_len = seq_len
sample_obj.query_len = query_len
sample_obj.generator = generator
sample_obj.is_prompt = is_prompt
else:
sample_obj = SequenceGroupToSample(
seq_ids=list(seq_ids),
sampling_params=sampling_params,
seq_data=seq_group_metadata.seq_data,
seq_len=seq_len,
query_len=query_len,
generator=generator,
is_prompt=is_prompt,
prompt_logprob_indices=list(prompt_logprob_indices),
prompt_logprob_output_indices=list(
prompt_logprob_output_indices),
sample_indices=list(sample_indices),
)
assert (len(sample_obj.prompt_logprob_indices)
== len(sample_obj.prompt_logprob_output_indices))
seq_groups.append(sample_obj)
if cache is not None:
cache.reset()
return (seq_groups, selected_token_indices, categorized_sample_indices,
num_prompts)
@dataclass
class SamplingTensors:
"""Tensors for sampling."""
temperatures: torch.Tensor
top_ps: torch.Tensor
top_ks: torch.Tensor
min_ps: torch.Tensor
presence_penalties: torch.Tensor
frequency_penalties: torch.Tensor
repetition_penalties: torch.Tensor
prompt_tokens: torch.Tensor
output_tokens: torch.Tensor
@classmethod
def from_sampling_metadata(
cls,
sampling_metadata: "SamplingMetadata",
vocab_size: int,
device: torch.device,
dtype: torch.dtype,
) -> Tuple["SamplingTensors", bool, bool, bool]:
prompt_tokens: List[array] = []
output_tokens: List[array] = []
top_ks: List[int] = []
temperatures: List[float] = []
top_ps: List[float] = []
min_ps: List[float] = []
presence_penalties: List[float] = []
frequency_penalties: List[float] = []
repetition_penalties: List[float] = []
do_penalties = False
do_top_p_top_k = False
do_min_p = False
assert sampling_metadata.seq_groups is not None
for seq_group in sampling_metadata.seq_groups:
seq_ids = seq_group.seq_ids
sampling_params = seq_group.sampling_params
temperature = sampling_params.temperature
p = sampling_params.presence_penalty
f = sampling_params.frequency_penalty
r = sampling_params.repetition_penalty
top_p = sampling_params.top_p
min_p = sampling_params.min_p
# k should not be greater than the vocab size.
top_k = min(sampling_params.top_k, vocab_size)
top_k = vocab_size if top_k == -1 else top_k
if temperature < _SAMPLING_EPS:
# NOTE: Zero temperature means deterministic sampling
# (i.e., greedy sampling or beam search).
# Set the temperature to 1 to avoid division by zero.
temperature = 1.0
if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS
or top_k != vocab_size):
do_top_p_top_k = True
if not do_min_p and min_p > _SAMPLING_EPS:
do_min_p = True
if not do_penalties and (abs(p) >= _SAMPLING_EPS
or abs(f) >= _SAMPLING_EPS
or abs(r - 1.0) >= _SAMPLING_EPS):
do_penalties = True
is_prompt = seq_group.is_prompt
if is_prompt and sampling_params.prompt_logprobs is not None:
# For tokens in the prompt that we only need to get
# their logprobs
query_len = seq_group.query_len
assert query_len is not None
prefill_len = len(seq_group.prompt_logprob_indices)
temperatures += [temperature] * prefill_len
top_ps += [top_p] * prefill_len
top_ks += [top_k] * prefill_len
min_ps += [min_p] * prefill_len
presence_penalties += [0] * prefill_len
frequency_penalties += [0] * prefill_len
repetition_penalties += [1] * prefill_len
if seq_group.do_sample:
sample_lens = len(seq_group.sample_indices)
assert sample_lens >= len(seq_ids)
temperatures += [temperature] * sample_lens
top_ps += [top_p] * sample_lens
top_ks += [top_k] * sample_lens
min_ps += [min_p] * sample_lens
presence_penalties += [p] * sample_lens
frequency_penalties += [f] * sample_lens
repetition_penalties += [r] * sample_lens
if do_penalties:
for seq_group in sampling_metadata.seq_groups:
seq_ids = seq_group.seq_ids
if (seq_group.is_prompt
and sampling_params.prompt_logprobs is not None):
prefill_len = len(seq_group.prompt_logprob_indices)
prompt_tokens.extend(
array(VLLM_TOKEN_ID_ARRAY_TYPE)
for _ in range(prefill_len))
output_tokens.extend(
array(VLLM_TOKEN_ID_ARRAY_TYPE)
for _ in range(prefill_len))
if seq_group.do_sample:
for seq_id in seq_ids:
seq_data = seq_group.seq_data[seq_id]
prompt_tokens.append(seq_data.prompt_token_ids_array)
output_tokens.append(seq_data.output_token_ids_array)
sampling_tensors = SamplingTensors.from_lists(
temperatures,
top_ps,
top_ks,
min_ps,
presence_penalties,
frequency_penalties,
repetition_penalties,
prompt_tokens,
output_tokens,
vocab_size,
device,
dtype,
)
return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p)
@classmethod
def from_lists(
cls,
temperatures: List[float],
top_ps: List[float],
top_ks: List[int],
min_ps: List[float],
presence_penalties: List[float],
frequency_penalties: List[float],
repetition_penalties: List[float],
prompt_tokens: List[array],
output_tokens: List[array],
vocab_size: int,
device: torch.device,
dtype: torch.dtype,
) -> "SamplingTensors":
# Note that the performance will be very bad without
# pinned memory.
pin_memory = is_pin_memory_available()
do_penalties = prompt_tokens or output_tokens
if do_penalties:
prompt_t = make_tensor_with_pad(
prompt_tokens,
vocab_size,
device="cpu",
dtype=torch.int64,
pin_memory=pin_memory,
)
output_t = make_tensor_with_pad(
output_tokens,
vocab_size,
device="cpu",
dtype=torch.int64,
pin_memory=pin_memory,
)
else:
empty_tensor = torch.empty(0, device=device, dtype=torch.long)
prompt_t = empty_tensor
output_t = empty_tensor
temperatures_t = torch.tensor(
temperatures,
device="cpu",
dtype=dtype,
pin_memory=pin_memory,
)
top_ps_t = torch.tensor(
top_ps,
device="cpu",
dtype=dtype,
pin_memory=pin_memory,
)
min_ps_t = torch.tensor(
min_ps,
device="cpu",
dtype=dtype,
pin_memory=pin_memory,
)
presence_penalties_t = torch.tensor(
presence_penalties,
device="cpu",
dtype=dtype,
pin_memory=pin_memory,
)
frequency_penalties_t = torch.tensor(
frequency_penalties,
device="cpu",
dtype=dtype,
pin_memory=pin_memory,
)
repetition_penalties_t = torch.tensor(
repetition_penalties,
device="cpu",
dtype=dtype,
pin_memory=pin_memory,
)
top_ks_t = torch.tensor(
top_ks,
device="cpu",
dtype=torch.int,
pin_memory=pin_memory,
)
# Because the memory is pinned, we can do non-blocking
# transfer to device.
return cls(
temperatures=temperatures_t.to(device=device, non_blocking=True),
top_ps=top_ps_t.to(device=device, non_blocking=True),
top_ks=top_ks_t.to(device=device, non_blocking=True),
min_ps=min_ps_t.to(device=device, non_blocking=True),
presence_penalties=presence_penalties_t.to(device=device,
non_blocking=True),
frequency_penalties=frequency_penalties_t.to(device=device,
non_blocking=True),
repetition_penalties=repetition_penalties_t.to(device=device,
non_blocking=True),
prompt_tokens=prompt_t.to(device=device, non_blocking=True),
output_tokens=output_t.to(device=device, non_blocking=True),
)

View File

@@ -0,0 +1,520 @@
"""Sampling parameters for text generation."""
import copy
from dataclasses import dataclass
from enum import Enum, IntEnum
from functools import cached_property
from typing import Any, Callable, Dict, List, Optional, Set, Union
import msgspec
import torch
from pydantic import BaseModel
from typing_extensions import Annotated
from vllm.logger import init_logger
logger = init_logger(__name__)
_SAMPLING_EPS = 1e-5
_MAX_TEMP = 1e-2
class SamplingType(IntEnum):
GREEDY = 0
RANDOM = 1
RANDOM_SEED = 2
LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor],
Callable[[List[int], List[int], torch.Tensor],
torch.Tensor]]
"""LogitsProcessor is a function that takes a list
of previously generated tokens, the logits tensor
for the next token and, optionally, prompt tokens as a
first argument, and returns a modified tensor of logits
to sample from."""
# maybe make msgspec?
@dataclass
class GuidedDecodingParams:
"""One of these fields will be used to build a logit processor."""
json: Optional[Union[str, Dict]] = None
regex: Optional[str] = None
choice: Optional[List[str]] = None
grammar: Optional[str] = None
json_object: Optional[bool] = None
"""These are other options that can be set"""
backend: Optional[str] = None
whitespace_pattern: Optional[str] = None
@staticmethod
def from_optional(
json: Optional[Union[Dict, BaseModel, str]],
regex: Optional[str] = None,
choice: Optional[List[str]] = None,
grammar: Optional[str] = None,
json_object: Optional[bool] = None,
backend: Optional[str] = None,
whitespace_pattern: Optional[str] = None,
) -> "GuidedDecodingParams":
# Extract json schemas from pydantic models
if isinstance(json, (BaseModel, type(BaseModel))):
json = json.model_json_schema()
return GuidedDecodingParams(
json=json,
regex=regex,
choice=choice,
grammar=grammar,
json_object=json_object,
backend=backend,
whitespace_pattern=whitespace_pattern,
)
def __post_init__(self):
"""Validate that some fields are mutually exclusive."""
guide_count = sum([
self.json is not None, self.regex is not None, self.choice
is not None, self.grammar is not None, self.json_object is not None
])
if guide_count > 1:
raise ValueError(
"You can only use one kind of guided decoding but multiple are "
f"specified: {self.__dict__}")
class RequestOutputKind(Enum):
# Return entire output so far in every RequestOutput
CUMULATIVE = 0
# Return only deltas in each RequestOutput
DELTA = 1
# Do not return intermediate RequestOuputs
FINAL_ONLY = 2
class SamplingParams(
msgspec.Struct,
omit_defaults=True, # type: ignore[call-arg]
# required for @cached_property.
dict=True): # type: ignore[call-arg]
"""Sampling parameters for text generation.
Overall, we follow the sampling parameters from the OpenAI text completion
API (https://platform.openai.com/docs/api-reference/completions/create).
In addition, we support beam search, which is not supported by OpenAI.
Args:
n: Number of output sequences to return for the given prompt.
best_of: Number of output sequences that are generated from the prompt.
From these `best_of` sequences, the top `n` sequences are returned.
`best_of` must be greater than or equal to `n`. By default,
`best_of` is set to `n`.
presence_penalty: Float that penalizes new tokens based on whether they
appear in the generated text so far. Values > 0 encourage the model
to use new tokens, while values < 0 encourage the model to repeat
tokens.
frequency_penalty: Float that penalizes new tokens based on their
frequency in the generated text so far. Values > 0 encourage the
model to use new tokens, while values < 0 encourage the model to
repeat tokens.
repetition_penalty: Float that penalizes new tokens based on whether
they appear in the prompt and the generated text so far. Values > 1
encourage the model to use new tokens, while values < 1 encourage
the model to repeat tokens.
temperature: Float that controls the randomness of the sampling. Lower
values make the model more deterministic, while higher values make
the model more random. Zero means greedy sampling.
top_p: Float that controls the cumulative probability of the top tokens
to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
top_k: Integer that controls the number of top tokens to consider. Set
to -1 to consider all tokens.
min_p: Float that represents the minimum probability for a token to be
considered, relative to the probability of the most likely token.
Must be in [0, 1]. Set to 0 to disable this.
seed: Random seed to use for the generation.
stop: List of strings that stop the generation when they are generated.
The returned output will not contain the stop strings.
stop_token_ids: List of tokens that stop the generation when they are
generated. The returned output will contain the stop tokens unless
the stop tokens are special tokens.
include_stop_str_in_output: Whether to include the stop strings in
output text. Defaults to False.
ignore_eos: Whether to ignore the EOS token and continue generating
tokens after the EOS token is generated.
max_tokens: Maximum number of tokens to generate per output sequence.
min_tokens: Minimum number of tokens to generate per output sequence
before EOS or stop_token_ids can be generated
logprobs: Number of log probabilities to return per output token.
When set to None, no probability is returned. If set to a non-None
value, the result includes the log probabilities of the specified
number of most likely tokens, as well as the chosen tokens.
Note that the implementation follows the OpenAI API: The API will
always return the log probability of the sampled token, so there
may be up to `logprobs+1` elements in the response.
prompt_logprobs: Number of log probabilities to return per prompt token.
detokenize: Whether to detokenize the output. Defaults to True.
skip_special_tokens: Whether to skip special tokens in the output.
spaces_between_special_tokens: Whether to add spaces between special
tokens in the output. Defaults to True.
logits_processors: List of functions that modify logits based on
previously generated tokens, and optionally prompt tokens as
a first argument.
truncate_prompt_tokens: If set to an integer k, will use only the last k
tokens from the prompt (i.e., left truncation). Defaults to None
(i.e., no truncation).
guided_decoding: If provided, the engine will construct a guided
decoding logits processor from these parameters. Defaults to None.
logit_bias: If provided, the engine will construct a logits processor
that applies these logit biases. Defaults to None.
allowed_token_ids: If provided, the engine will construct a logits
processor which only retains scores for the given token ids.
Defaults to None.
prompt_logprob_positions: Optional prompt-token positions whose logits
should be materialized. None preserves the standard all-position
prompt-logprob behavior.
"""
n: int = 1
best_of: Optional[int] = None
_real_n: Optional[int] = None
presence_penalty: float = 0.0
frequency_penalty: float = 0.0
repetition_penalty: float = 1.0
temperature: float = 1.0
top_p: float = 1.0
top_k: int = -1
min_p: float = 0.0
seed: Optional[int] = None
stop: Optional[Union[str, List[str]]] = None
stop_token_ids: Optional[List[int]] = None
ignore_eos: bool = False
max_tokens: Optional[int] = 16
min_tokens: int = 0
logprobs: Optional[int] = None
prompt_logprobs: Optional[int] = None
# NOTE: This parameter is only exposed at the engine level for now.
# It is not exposed in the OpenAI API server, as the OpenAI API does
# not support returning only a list of token IDs.
detokenize: bool = True
skip_special_tokens: bool = True
spaces_between_special_tokens: bool = True
# Optional[List[LogitsProcessor]] type. We use Any here because
# Optional[List[LogitsProcessor]] type is not supported by msgspec.
logits_processors: Optional[Any] = None
include_stop_str_in_output: bool = False
truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None
output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE
# The below fields are not supposed to be used as an input.
# They are set in post_init.
output_text_buffer_length: int = 0
_all_stop_token_ids: Set[int] = msgspec.field(default_factory=set)
# Fields used to construct logits processors
guided_decoding: Optional[GuidedDecodingParams] = None
logit_bias: Optional[Dict[int, float]] = None
allowed_token_ids: Optional[List[int]] = None
prompt_logprob_positions: Optional[List[int]] = None
@staticmethod
def from_optional(
n: Optional[int] = 1,
best_of: Optional[int] = None,
presence_penalty: Optional[float] = 0.0,
frequency_penalty: Optional[float] = 0.0,
repetition_penalty: Optional[float] = 1.0,
temperature: Optional[float] = 1.0,
top_p: Optional[float] = 1.0,
top_k: int = -1,
min_p: float = 0.0,
seed: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
stop_token_ids: Optional[List[int]] = None,
include_stop_str_in_output: bool = False,
ignore_eos: bool = False,
max_tokens: Optional[int] = 16,
min_tokens: int = 0,
logprobs: Optional[int] = None,
prompt_logprobs: Optional[int] = None,
detokenize: bool = True,
skip_special_tokens: bool = True,
spaces_between_special_tokens: bool = True,
logits_processors: Optional[List[LogitsProcessor]] = None,
truncate_prompt_tokens: Optional[Annotated[int,
msgspec.Meta(ge=1)]] = None,
output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE,
guided_decoding: Optional[GuidedDecodingParams] = None,
logit_bias: Optional[Union[Dict[int, float], Dict[str, float]]] = None,
allowed_token_ids: Optional[List[int]] = None,
prompt_logprob_positions: Optional[List[int]] = None,
) -> "SamplingParams":
if logit_bias is not None:
logit_bias = {
int(token): bias
for token, bias in logit_bias.items()
}
return SamplingParams(
n=1 if n is None else n,
best_of=best_of,
presence_penalty=0.0
if presence_penalty is None else presence_penalty,
frequency_penalty=0.0
if frequency_penalty is None else frequency_penalty,
repetition_penalty=1.0
if repetition_penalty is None else repetition_penalty,
temperature=1.0 if temperature is None else temperature,
top_p=1.0 if top_p is None else top_p,
top_k=top_k,
min_p=min_p,
seed=seed,
stop=stop,
stop_token_ids=stop_token_ids,
include_stop_str_in_output=include_stop_str_in_output,
ignore_eos=ignore_eos,
max_tokens=max_tokens,
min_tokens=min_tokens,
logprobs=logprobs,
prompt_logprobs=prompt_logprobs,
detokenize=detokenize,
skip_special_tokens=skip_special_tokens,
spaces_between_special_tokens=spaces_between_special_tokens,
logits_processors=logits_processors,
truncate_prompt_tokens=truncate_prompt_tokens,
output_kind=output_kind,
guided_decoding=guided_decoding,
logit_bias=logit_bias,
allowed_token_ids=allowed_token_ids,
prompt_logprob_positions=prompt_logprob_positions,
)
def __post_init__(self) -> None:
# how we deal with `best_of``:
# if `best_of`` is not set, we default to `n`;
# if `best_of`` is set, we set `n`` to `best_of`,
# and set `_real_n`` to the original `n`.
# when we return the result, we will check
# if we need to return `n` or `_real_n` results
if self.best_of:
if self.best_of < self.n:
raise ValueError(
f"best_of must be greater than or equal to n, "
f"got n={self.n} and best_of={self.best_of}.")
self._real_n = self.n
self.n = self.best_of
if 0 < self.temperature < _MAX_TEMP:
logger.warning(
"temperature %s is less than %s, which may cause numerical "
"errors nan or inf in tensors. We have maxed it out to %s.",
self.temperature, _MAX_TEMP, _MAX_TEMP)
self.temperature = max(self.temperature, _MAX_TEMP)
if self.seed == -1:
self.seed = None
else:
self.seed = self.seed
if self.stop is None:
self.stop = []
elif isinstance(self.stop, str):
self.stop = [self.stop]
else:
self.stop = list(self.stop)
if self.stop_token_ids is None:
self.stop_token_ids = []
else:
self.stop_token_ids = list(self.stop_token_ids)
self.logprobs = 1 if self.logprobs is True else self.logprobs
self.prompt_logprobs = (1 if self.prompt_logprobs is True else
self.prompt_logprobs)
if self.prompt_logprob_positions is not None:
self.prompt_logprob_positions = list(
self.prompt_logprob_positions)
# Number of characters to hold back for stop string evaluation
# until sequence is finished.
if self.stop and not self.include_stop_str_in_output:
self.output_text_buffer_length = max(len(s) for s in self.stop) - 1
self._verify_args()
if self.temperature < _SAMPLING_EPS:
# Zero temperature means greedy sampling.
self.top_p = 1.0
self.top_k = -1
self.min_p = 0.0
self._verify_greedy_sampling()
# eos_token_id is added to this by the engine
self._all_stop_token_ids = set(self.stop_token_ids)
def _verify_args(self) -> None:
if not isinstance(self.n, int):
raise ValueError(f"n must be an int, but is of "
f"type {type(self.n)}")
if self.n < 1:
raise ValueError(f"n must be at least 1, got {self.n}.")
if not -2.0 <= self.presence_penalty <= 2.0:
raise ValueError("presence_penalty must be in [-2, 2], got "
f"{self.presence_penalty}.")
if not -2.0 <= self.frequency_penalty <= 2.0:
raise ValueError("frequency_penalty must be in [-2, 2], got "
f"{self.frequency_penalty}.")
if not 0.0 < self.repetition_penalty <= 2.0:
raise ValueError("repetition_penalty must be in (0, 2], got "
f"{self.repetition_penalty}.")
if self.temperature < 0.0:
raise ValueError(
f"temperature must be non-negative, got {self.temperature}.")
if not 0.0 < self.top_p <= 1.0:
raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
if self.top_k < -1 or self.top_k == 0:
raise ValueError(f"top_k must be -1 (disable), or at least 1, "
f"got {self.top_k}.")
if not isinstance(self.top_k, int):
raise TypeError(
f"top_k must be an integer, got {type(self.top_k).__name__}")
if not 0.0 <= self.min_p <= 1.0:
raise ValueError("min_p must be in [0, 1], got "
f"{self.min_p}.")
if self.max_tokens is not None and self.max_tokens < 1:
raise ValueError(
f"max_tokens must be at least 1, got {self.max_tokens}.")
if self.min_tokens < 0:
raise ValueError(f"min_tokens must be greater than or equal to 0, "
f"got {self.min_tokens}.")
if self.max_tokens is not None and self.min_tokens > self.max_tokens:
raise ValueError(
f"min_tokens must be less than or equal to "
f"max_tokens={self.max_tokens}, got {self.min_tokens}.")
if self.logprobs is not None and self.logprobs < 0:
raise ValueError(
f"logprobs must be non-negative, got {self.logprobs}.")
if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
raise ValueError(f"prompt_logprobs must be non-negative, got "
f"{self.prompt_logprobs}.")
if self.prompt_logprob_positions is not None:
if self.prompt_logprobs is None:
raise ValueError(
"prompt_logprob_positions requires prompt_logprobs.")
if (
not self.prompt_logprob_positions
or any(
not isinstance(position, int)
or isinstance(position, bool)
or position <= 0
for position in self.prompt_logprob_positions
)
or self.prompt_logprob_positions
!= sorted(set(self.prompt_logprob_positions))
):
raise ValueError(
"prompt_logprob_positions must be a sorted unique list "
"of positive integers.")
if (self.truncate_prompt_tokens is not None
and self.truncate_prompt_tokens < 1):
raise ValueError(f"truncate_prompt_tokens must be >= 1, "
f"got {self.truncate_prompt_tokens}")
assert isinstance(self.stop, list)
if any(not stop_str for stop_str in self.stop):
raise ValueError("stop cannot contain an empty string.")
if self.stop and not self.detokenize:
raise ValueError(
"stop strings are only supported when detokenize is True. "
"Set detokenize=True to use stop.")
if self.best_of != self._real_n and self.output_kind == (
RequestOutputKind.DELTA):
raise ValueError("best_of must equal n to use output_kind=DELTA")
def _verify_greedy_sampling(self) -> None:
if self.n > 1:
raise ValueError("n must be 1 when using greedy sampling, "
f"got {self.n}.")
def update_from_generation_config(
self,
generation_config: Dict[str, Any],
model_eos_token_id: Optional[int] = None) -> None:
"""Update if there are non-default values from generation_config"""
if model_eos_token_id is not None:
# Add the eos token id into the sampling_params to support
# min_tokens processing.
self._all_stop_token_ids.add(model_eos_token_id)
# Update eos_token_id for generation
if (eos_ids := generation_config.get("eos_token_id")) is not None:
# it can be either int or list of int
eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
if model_eos_token_id is not None:
# We don't need to include the primary eos_token_id in
# stop_token_ids since it's handled separately for stopping
# purposes.
eos_ids.discard(model_eos_token_id)
if eos_ids:
self._all_stop_token_ids.update(eos_ids)
if not self.ignore_eos:
eos_ids.update(self.stop_token_ids)
self.stop_token_ids = list(eos_ids)
@cached_property
def sampling_type(self) -> SamplingType:
if self.temperature < _SAMPLING_EPS:
return SamplingType.GREEDY
if self.seed is not None:
return SamplingType.RANDOM_SEED
return SamplingType.RANDOM
@property
def all_stop_token_ids(self) -> Set[int]:
return self._all_stop_token_ids
def clone(self) -> "SamplingParams":
"""Deep copy excluding LogitsProcessor objects.
LogitsProcessor objects are excluded because they may contain an
arbitrary, nontrivial amount of data.
See https://github.com/vllm-project/vllm/issues/3087
"""
logit_processor_refs = None if self.logits_processors is None else {
id(lp): lp
for lp in self.logits_processors
}
return copy.deepcopy(self, memo=logit_processor_refs)
def __repr__(self) -> str:
return (
f"SamplingParams(n={self.n}, "
f"presence_penalty={self.presence_penalty}, "
f"frequency_penalty={self.frequency_penalty}, "
f"repetition_penalty={self.repetition_penalty}, "
f"temperature={self.temperature}, "
f"top_p={self.top_p}, "
f"top_k={self.top_k}, "
f"min_p={self.min_p}, "
f"seed={self.seed}, "
f"stop={self.stop}, "
f"stop_token_ids={self.stop_token_ids}, "
f"include_stop_str_in_output={self.include_stop_str_in_output}, "
f"ignore_eos={self.ignore_eos}, "
f"max_tokens={self.max_tokens}, "
f"min_tokens={self.min_tokens}, "
f"logprobs={self.logprobs}, "
f"prompt_logprobs={self.prompt_logprobs}, "
"prompt_logprob_positions="
f"{self.prompt_logprob_positions}, "
f"skip_special_tokens={self.skip_special_tokens}, "
"spaces_between_special_tokens="
f"{self.spaces_between_special_tokens}, "
f"truncate_prompt_tokens={self.truncate_prompt_tokens}), "
f"guided_decoding={self.guided_decoding}")
class BeamSearchParams(
msgspec.Struct,
omit_defaults=True, # type: ignore[call-arg]
# required for @cached_property.
dict=True): # type: ignore[call-arg]
"""Beam search parameters for text generation."""
beam_width: int
max_tokens: int
ignore_eos: bool = False
temperature: float = 0.0
length_penalty: float = 1.0

View File

@@ -0,0 +1,276 @@
#!/usr/bin/env python3
"""verify_dlopen_chain.py — Verify complete dlopen SO → Python → Model call chain.
Run on BI-V100 to identify gaps before submitting.
Usage: python3 verify_dlopen_chain.py [--vllm-root /path/to/vllm]
Checks:
1. All prebuilt .so files are installed and loadable as Python extension modules
2. All env-gated kernel dispatch paths have matching .so
3. protocol.py correctly accepts max_completion_tokens
4. qwen3_5.py import chain is complete (no silent None fallbacks for enabled kernels)
"""
import argparse
import ctypes
import importlib
import importlib.util
import json
import os
import pathlib
import struct
import sys
import traceback
PASS = "\033[92m✓ PASS\033[0m"
FAIL = "\033[91m✗ FAIL\033[0m"
SKIP = "\033[93m- SKIP\033[0m"
INFO = "\033[94m INFO\033[0m"
results = {"pass": 0, "fail": 0, "skip": 0}
def check(name, condition, detail=""):
if condition:
results["pass"] += 1
print(f" {PASS} {name}" + (f"{detail}" if detail else ""))
else:
results["fail"] += 1
print(f" {FAIL} {name}" + (f"{detail}" if detail else ""))
def skip(name, reason=""):
results["skip"] += 1
print(f" {SKIP} {name}" + (f"{reason}" if reason else ""))
def section(title):
print(f"\n{'='*60}")
print(f" {title}")
print(f"{'='*60}")
# ---------- Section 1: prebuilt SO ELF validation ----------
EXPECTED_SO = [
"corex_attn_head_rms_norm",
"corex_block_major_kv_transfer",
"corex_fused_paged_prefill",
"corex_gdn_beta_decay",
"corex_gdn_causal_conv",
"corex_gdn_chunk_recurrent",
"corex_gdn_gated_norm",
"corex_gdn_packed_decode",
"corex_gdn_qk_map",
"corex_moe_direct_routed",
"corex_moe_exact_reduce",
"corex_moe_index_combine",
"corex_moe_topk_softmax",
"corex_moe_weight_gather",
"corex_paged_kv_gather",
]
# Env var → module name mapping (from qwen3_5.py)
ENV_KERNEL_MAP = {
"BI100_GDN_COREX_CAUSAL_CONV": "corex_gdn_causal_conv",
"BI100_GDN_COREX_GATED_NORM": "corex_gdn_gated_norm",
"BI100_GDN_COREX_BETA_DECAY": "corex_gdn_beta_decay",
"BI100_GDN_COREX_QK_MAP": "corex_gdn_qk_map",
"BI100_GDN_COREX_PACKED_DECODE": "corex_gdn_packed_decode",
"BI100_ATTN_COREX_HEAD_RMS_NORM": "corex_attn_head_rms_norm",
"BI100_MOE_COREX_EXACT_REDUCE": "corex_moe_exact_reduce",
"BI100_MOE_COREX_WEIGHT_GATHER": "corex_moe_weight_gather",
"BI100_MOE_COREX_DIRECT_ROUTED": "corex_moe_direct_routed",
"BI100_MOE_COREX_TOPK_SOFTMAX": "corex_moe_topk_softmax",
"BI100_MOE_COREX_INDEX_COMBINE": "corex_moe_index_combine",
}
def find_vllm_root():
spec = importlib.util.find_spec("vllm")
if spec and spec.submodule_search_locations:
return pathlib.Path(next(iter(spec.submodule_search_locations)))
return None
def check_elf(path):
"""Validate file is a 64-bit x86-64 ELF."""
if not path.exists():
return False, "file not found"
if path.stat().st_size == 0:
return False, "empty file"
header = path.read_bytes()[:20]
if len(header) < 20 or header[:4] != b"\x7fELF":
return False, "not ELF"
if header[4:6] != b"\x02\x01":
return False, "not 64-bit LE"
machine = struct.unpack_from("<H", header, 18)[0]
if machine != 62:
return False, f"not x86-64 (machine={machine})"
return True, f"ok ({path.stat().st_size} bytes)"
def verify_so_installations(vllm_root):
section("1. Prebuilt SO Installation & ELF Validation")
for name in EXPECTED_SO:
so_path = vllm_root / f"{name}.so"
ok, detail = check_elf(so_path)
check(f"{name}.so", ok, detail)
def verify_so_importable(vllm_root):
section("2. SO Python Import Chain (torch extension)")
for name in EXPECTED_SO:
so_path = vllm_root / f"{name}.so"
if not so_path.exists():
skip(f"import vllm.{name}", "SO not installed")
continue
try:
mod = importlib.import_module(f"vllm.{name}")
funcs = [f for f in dir(mod) if not f.startswith("_")]
check(f"import vllm.{name}", True,
f"exports: {', '.join(funcs[:5])}")
except Exception as e:
check(f"import vllm.{name}", False, str(e)[:120])
def verify_env_kernel_dispatch():
section("3. Env-Gated Kernel Dispatch Consistency")
for env_var, module_name in ENV_KERNEL_MAP.items():
env_val = os.environ.get(env_var, "<unset>")
enabled = env_val in ("1", "true", "True")
try:
mod = importlib.import_module(f"vllm.{module_name}")
available = mod is not None
except Exception:
available = False
if enabled and not available:
check(f"{env_var}={env_val}{module_name}",
False, "ENABLED but SO not loadable — will crash!")
elif enabled and available:
check(f"{env_var}={env_val}{module_name}",
True, "enabled + available")
elif not enabled:
check(f"{env_var}={env_val}{module_name}",
True, f"disabled (available={available})")
def verify_protocol():
section("4. Protocol max_completion_tokens Acceptance")
try:
from vllm.entrypoints.openai.protocol import ChatCompletionRequest
# Simulate a request with max_completion_tokens
test_data = {
"model": "llm",
"messages": [{"role": "user", "content": "test"}],
"max_completion_tokens": 8192,
}
req = ChatCompletionRequest(**test_data)
# After fold_max_completion_tokens, max_tokens should be 8192
check("max_completion_tokens accepted",
req.max_tokens == 8192,
f"max_tokens={req.max_tokens}")
# Also test with thinking param
test_data2 = {
"model": "llm",
"messages": [{"role": "user", "content": "test"}],
"max_completion_tokens": 32768,
"thinking": {"type": "enabled", "budget_tokens": 10000},
}
req2 = ChatCompletionRequest(**test_data2)
check("max_completion_tokens+thinking accepted",
req2.max_tokens == 32768,
f"max_tokens={req2.max_tokens}, thinking={req2.thinking}")
except Exception as e:
check("protocol import/validation", False, str(e)[:200])
def verify_model_imports():
section("5. qwen3_5.py Model Import Chain")
try:
# Don't actually import the full model (needs CUDA), just check the file
vllm_root = find_vllm_root()
if vllm_root is None:
skip("qwen3_5.py", "vllm not installed")
return
model_path = vllm_root / "model_executor" / "models" / "qwen3_5.py"
check("qwen3_5.py installed",
model_path.exists(),
str(model_path))
if model_path.exists():
source = model_path.read_text()
# Check all corex imports are present
for name in EXPECTED_SO:
if f"from vllm import {name}" in source or \
f"import {name}" in source:
check(f"qwen3_5 imports {name}", True)
else:
# Not all SOs are imported directly in qwen3_5.py
# Some are used via other modules
if name in ("corex_block_major_kv_transfer",
"corex_fused_paged_prefill",
"corex_paged_kv_gather"):
skip(f"qwen3_5 imports {name}",
"used via paged_attn/block_major modules")
else:
check(f"qwen3_5 imports {name}", False,
"import not found in source")
except Exception as e:
check("model import chain", False, str(e)[:200])
def verify_paged_attn_chain(vllm_root):
section("6. Paged Attention dlopen Chain")
if vllm_root is None:
skip("paged_attn chain", "vllm not installed")
return
paged = vllm_root / "attention" / "ops" / "paged_attn.py"
check("paged_attn.py installed", paged.exists(), str(paged))
if paged.exists():
source = paged.read_text()
for name in ("corex_paged_kv_gather",
"corex_fused_paged_prefill",
"corex_block_major_kv_transfer"):
found = name in source
if found:
check(f"paged_attn uses {name}", True)
else:
skip(f"paged_attn uses {name}", "not referenced")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--vllm-root", type=pathlib.Path, default=None)
args = parser.parse_args()
vllm_root = args.vllm_root or find_vllm_root()
if vllm_root is None:
print("ERROR: Cannot find vllm installation. Use --vllm-root.")
sys.exit(1)
print(f"vllm root: {vllm_root}")
verify_so_installations(vllm_root)
verify_so_importable(vllm_root)
verify_env_kernel_dispatch()
verify_protocol()
verify_model_imports()
verify_paged_attn_chain(vllm_root)
section("SUMMARY")
total = results["pass"] + results["fail"] + results["skip"]
print(f" Pass: {results['pass']}/{total}")
print(f" Fail: {results['fail']}/{total}")
print(f" Skip: {results['skip']}/{total}")
if results["fail"] > 0:
print(f"\n ⚠️ {results['fail']} checks FAILED — fix before submitting!")
sys.exit(1)
else:
print(f"\n ✅ All checks passed!")
sys.exit(0)
if __name__ == "__main__":
main()