Compare commits

...
Author SHA1 Message Date
Bugen Zhao efa494e397 refactor setup.py to extract stuff into build_rust.py
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:39:17 +08:00
Bugen Zhao 74b77a593a update dockerfile image
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:38:24 +08:00
Bugen Zhao ec53889a3b simplify to only match .so, as .dylib will also be renamed to .so under macos
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:13:09 +08:00
Bugen Zhao ede3d4ddf6 skip tests when _rust_tool_parser is absent
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:12:48 +08:00
Bugen Zhao 7301a834fa fix glob issue
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:06:14 +08:00
Bugen Zhao 7104ac6b5d simplify setup.py
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 11:40:33 +00:00
Bugen Zhao 9a98eccb49 cover two calls case
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 2ce1e2ba71 do not infer finished but just leave it unhandled
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 2496f66f7f Run tool parser tests before tokenizer tests
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao d18b7f2723 Build Rust parser extension in Docker rust stage
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 80b633f076 ci: configure PyO3 Python for Rust tests
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 28f589ebb1 do not enable pyo3 on --all-features
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 98998279d4 fix fmt & try import
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen Zhao 71d208089e add tests with deepseek v4
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen Zhao 0ac6fc8352 Add Rust tool parser Python bridge
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen ZhaoandGitHub fdfb2566c0 [Rust Frontend] [CI] Unify Rust artifact builds with setuptools-rust (#44981)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 17:48:34 +08:00
Juan Pérez de AlgabaandGitHub 8a5cf1ccd6 [Security] Fix remote DoS via invalid recovered token reinjection (#44744)
Signed-off-by: jperezde <jperezde@redhat.com>
2026-06-10 02:31:43 -07:00
Kunshang JiandGitHub fe1d923afc [BUGFIX][XPU] fix xpu flash_attn_varlen_func interface (#45110)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-10 17:07:40 +08:00
JartXGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
32daf56b42 [Refactor] Rename rocm_moe.py to rocm_moe_rdna.py (#45011)
Signed-off-by: JartX <sagformas@epdcenter.es>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-10 17:02:09 +08:00
Andreas KaratzasandGitHub 82a42234be [ROCm][CI] Defer AITER sampler import and isolate server test PYTHONPATH (#44823)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-10 08:56:11 +00:00
Harry MellorandGitHub af9f583344 Revert "[Bugfix][CI] Gemma3 Transformers multimodal encoder profiling and build prompt-embedding fixtures" (#45029)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-10 01:37:03 -07:00
yihengandGitHub bd2d83ff31 [SpecDecode] Reduce TP communication for large-vocab draft models speculative decoding (#39419)
Signed-off-by: EanWang211123 <wangyiheng@sangfor.com.cn>
2026-06-10 07:59:24 +00:00
xiaohuguo2023andGitHub bb78168b21 [ROCm][gpt-oss] Hybrid CDNA4 swizzle gate for A8W4 MoE (#44804)
Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com>
2026-06-09 23:59:44 -07:00
89c6a41001 [Bench] Add BFCL dataset for vllm bench serve tool-calling workloads (#42457)
Signed-off-by: Li Zhang <lzhanga@amazon.com>
Co-authored-by: Li Zhang <lzhanga@amazon.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-06-09 23:59:18 -07:00
7fdfa6441d Model/colbert autoweightsloader (#44999)
Signed-off-by: Furkan Fidan <dev@yufufi.com>
Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-09 23:58:50 -07:00
Harry MellorandGitHub e9b728de8a Change from owning configs to owning config utils (#45058)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-10 06:40:25 +00:00
5828a205ef Fix Harmony tool descriptions for optional fields (#44686)
Signed-off-by: Varun Shenoy <varun.vinayak.shenoy@oracle.com>
Co-authored-by: Codex <codex@openai.com>
2026-06-09 23:29:22 -07:00
Yaoming ZhanandGitHub 7a74f31d2e [Rust Frontend] Add seed_oss and step3p5 reasoning parsers (#44552)
Signed-off-by: yzhan1 <zhanyaoming2014@gmail.com>
2026-06-09 23:01:33 -07:00
47930b59ca [Bugfix] Handle HWC images in ImageProcessorItems.get_image_size (#45057)
Signed-off-by: YellowFoxH4XOR <yellowfoxh4xor@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-10 05:35:50 +00:00
Flora FengandGitHub 6aec99f030 [Refactor] Remove dead states from chat completion serving (#45081)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-09 22:20:15 -07:00
bnellnmandGitHub f4966f8b3d [Bugfix] Fix weight loading issues caused by #41184 (#45054)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-06-10 01:20:13 -04:00
Mohammad Miadh AngkadandGitHub 2c9c07c85e [Bugfix][CI/Build] Fix Rust frontend build after chat conversion refactor (#45085)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
2026-06-09 20:04:41 -07:00
Change72andGitHub 320c52b134 [Bench] benchmark_serving_multi_turn: make non-standard conversation_id payload opt-in (#43756)
Signed-off-by: Change72 <cguo51@asu.edu>
2026-06-09 19:41:56 -07:00
6deb05e0e4 [Core][Model] Gemma4: Unified FA4 for all layers + FlashAttention mm_prefix support (#42175)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-06-09 17:45:39 -07:00
Flora FengandGitHub d82ac00923 [Refactor][Mistral] Extract parsing logic into MistralParser (#44596)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-10 00:12:23 +00:00
Bugen ZhaoandGitHub dac9e9a640 [Rust Frontend] Extract shared options in route helper params (#44884)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-09 17:02:35 -07:00
Wentao YeandGitHub d7607ad273 [Bug] Fix deepseek v4 OOM issue (#44914)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-09 15:47:06 -07:00
d955745d58 [ROCm][CI] fix test_rope_kvcache_fusion.py (#44678)
Signed-off-by: charlifu <charlifu@amd.com>
Co-authored-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
2026-06-09 21:53:46 +00:00
Micah WilliamsonandGitHub e1ed89dbee Revert "[Kernel] Speed up silu_and_mul_per_block_quant with warp-shuf… (#45066)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-09 14:12:06 -07:00
JimmyGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>simon-mo
1c2ffc6f88 feat(multi-turn-bench): add api_key and custom headers for multi turn benchmark (#44516)
Signed-off-by: Jimmy <jinmingyi1998@sina.cn>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: simon-mo <simon.mo@hey.com>
2026-06-09 14:00:07 -07:00
Jiangyun ZhuandGitHub ca4cfd8731 [Bugfix] fix qwen3.5 ep weight loading (#45002)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-06-09 13:55:30 -07:00
Micah WilliamsonandGitHub c9c1540e61 [ROCm][V2] Fix failed assertion in Llama models when using EAGLE with ROCM_AITER_FA (#44936)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-09 13:30:52 -05:00
c1d754d681 [Mooncake] Use all HCAs on multi-NIC hosts instead of GPU-indexed RNIC selection (#43799)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-06-09 11:05:36 -07:00
01d8cd92dd [ROCm][Perf] Use fused softplus-sqrt-topk router under AITER fused-MoE (#44945)
Co-authored-by: vLLM Contributor <contributor@vllm.ai>
2026-06-09 17:53:05 +00:00
a4b14b98c6 [Kernel] Speed up silu_and_mul_per_block_quant with warp-shuffle reduction + vectorized I/O (#44173)
Signed-off-by: SII-yangdian <yangdian@sii.edu.cn>
Co-authored-by: SII-yangdian <yangdian@sii.edu.cn>
2026-06-09 10:41:26 -07:00
Juan Pérez de AlgabaandGitHub cf1c906724 [Security] Fix image EXIF orientation and tRNS transparency handling (#44974)
Signed-off-by: jperezde <jperezde@redhat.com>
2026-06-09 09:34:44 -07:00
766ce2bb6b Fix MiDashengLM TP>1 crash in audio encoder attention (#44408)
Signed-off-by: Michał Ganczarenko <michal.ganczarenko@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-09 09:29:14 -07:00
Ronen SchafferGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
3d119f78f7 [Docs] Add KV offloading usage guide (single- and multi-tier) (#44415)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-09 19:20:23 +03:00
Juan Pérez de AlgabaandGitHub 1b1359c332 [Security] Fix DoS via audio decompression bomb in speech-to-text endpoint (#44970)
Signed-off-by: jperezde <jperezde@redhat.com>
2026-06-10 00:18:53 +08:00
Tyko NiemiandGitHub cad4ca12b8 [Bugfix] Add X-Session-ID from conversation_id in multi-turn benchmark (#44663)
Signed-off-by: Tyko Niemi <tyko.niemi@amd.com>
2026-06-09 08:57:00 -07:00
116 changed files with 4189 additions and 1809 deletions
@@ -110,6 +110,36 @@ install_uv() {
| env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
}
setup_pyo3_python() {
local python_version="${PYO3_PYTHON_VERSION:-3.12}"
log_section "Installing Python ${python_version} for PyO3 tests"
uv python install "$python_version"
PYO3_PYTHON="$(uv python find \
--managed-python \
--no-project \
--resolve-links \
"$python_version")"
export PYO3_PYTHON
local python_libdir
python_libdir="$("$PYO3_PYTHON" - <<'PY'
import pathlib
import sysconfig
libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
ldlibrary = sysconfig.get_config_var("LDLIBRARY")
assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
assert ldlibrary
assert (libdir / ldlibrary).exists(), libdir / ldlibrary
print(libdir)
PY
)"
export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
}
run_style_clippy() {
install_cargo_sort
@@ -132,6 +162,7 @@ run_style_clippy() {
run_tests() {
install_uv
setup_pyo3_python
install_cargo_nextest
log_section "Running cargo nextest"
+2 -2
View File
@@ -300,9 +300,9 @@ steps:
- tests/multimodal
- tests/renderers
- tests/standalone_tests/lazy_imports.py
- tests/tokenizers_
- tests/reasoning
- tests/tool_parsers
- tests/tokenizers_
- tests/parser
- tests/transformers_utils
- tests/config
@@ -315,9 +315,9 @@ steps:
- pytest -v -s test_ray_env.py
- pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s renderers
- pytest -v -s tokenizers_
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
- pytest -v -s tool_parsers
- pytest -v -s tokenizers_
- pytest -v -s parser
- pytest -v -s transformers_utils
- pytest -v -s config
+6 -1
View File
@@ -23,9 +23,14 @@
# Any change to the VllmConfig changes can have a large user-facing impact,
# so spam a lot of people
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
/vllm/config/cache.py @heheda12345
# Config utils
/vllm/config/utils.py @hmellor
/vllm/engine/arg_utils.py @hmellor
/vllm/utils/argparse_utils.py
# Entrypoints
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
+1
View File
@@ -4,6 +4,7 @@ include requirements/cuda.txt
include requirements/rocm.txt
include requirements/cpu.txt
include CMakeLists.txt
include tools/build_rust.py
recursive-include cmake *
recursive-include csrc *
@@ -65,6 +65,32 @@ class RequestArgs(NamedTuple):
limit_min_tokens: int # Use negative value for no limit
limit_max_tokens: int # Use negative value for no limit
timeout_sec: int
send_conversation_id: bool
headers: dict[str, str]
def parse_custom_header(header: str) -> tuple[str, str]:
separators = (":", "=")
for separator in separators:
if separator in header:
key, value = header.split(separator, 1)
key = key.strip()
value = value.strip()
if key:
return key, value
break
raise argparse.ArgumentTypeError(
"Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
)
def build_request_headers(
api_key: str | None, custom_headers: list[tuple[str, str]] | None
) -> dict[str, str]:
headers = dict(custom_headers or [])
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
class BenchmarkArgs(NamedTuple):
@@ -218,12 +244,11 @@ async def send_request(
max_tokens: int | None = None,
timeout_sec: int = 120,
conversation_id: str | None = None,
headers: dict[str, str] | None = None,
) -> ServerResponse:
payload = {
"model": model,
"messages": messages,
"seed": 0,
"temperature": 0.0,
}
if conversation_id is not None:
@@ -233,13 +258,17 @@ async def send_request(
payload["stream"] = True
payload["stream_options"] = {"include_usage": False}
if min_tokens is not None:
payload["min_tokens"] = min_tokens
# if min_tokens is not None:
# payload["min_tokens"] = min_tokens
if max_tokens is not None:
payload["max_tokens"] = max_tokens
headers = {"Content-Type": "application/json"}
request_headers = {"Content-Type": "application/json"}
if conversation_id is not None:
request_headers["X-Session-ID"] = str(conversation_id)
if headers is not None:
request_headers.update(headers)
# Calculate the timeout for the request
if max_tokens is not None:
@@ -265,7 +294,7 @@ async def send_request(
most_recent_timestamp: int = start_time
async with session.post(
url=chat_url, json=payload, headers=headers, timeout=timeout
url=chat_url, json=payload, headers=request_headers, timeout=timeout
) as response:
http_status = HTTPStatus(response.status)
if http_status == HTTPStatus.OK:
@@ -317,6 +346,8 @@ async def send_request(
latency = time.perf_counter_ns() - start_time
if ttft is None:
if stream:
valid_response = False
# The response was a single chunk
ttft = latency
@@ -423,7 +454,8 @@ async def send_turn(
min_tokens,
max_tokens,
req_args.timeout_sec,
conversation_id=conv_id,
conversation_id=conv_id if req_args.send_conversation_id else None,
headers=req_args.headers,
)
if response.valid is False:
@@ -872,6 +904,7 @@ def get_client_config(
# Arguments for API requests
chat_url = f"{args.url}/v1/chat/completions"
model_name = args.served_model_name if args.served_model_name else args.model
headers = build_request_headers(args.api_key, args.header)
req_args = RequestArgs(
chat_url=chat_url,
@@ -880,6 +913,8 @@ def get_client_config(
limit_min_tokens=args.limit_min_tokens,
limit_max_tokens=args.limit_max_tokens,
timeout_sec=args.request_timeout_sec,
send_conversation_id=args.send_conversation_id,
headers=headers,
)
return client_args, req_args
@@ -1245,19 +1280,19 @@ def process_statistics(
)
async def get_server_info(url: str) -> None:
async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
async with aiohttp.ClientSession() as session:
# Get server version (not mandatory, "version" endpoint may not exist)
url_version = f"{url}/version"
async with session.get(url_version) as response:
async with session.get(url_version, headers=headers) as response:
if HTTPStatus(response.status) == HTTPStatus.OK:
text = await response.text()
logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
# Get available models
url_models = f"{url}/v1/models"
async with session.get(url_models) as response:
async with session.get(url_models, headers=headers) as response:
if HTTPStatus(response.status) == HTTPStatus.OK:
text = await response.text()
logger.info(f"{Color.BLUE}Models:{Color.RESET}")
@@ -1323,6 +1358,22 @@ async def main() -> None:
help="Base URL for the LLM API server",
)
parser.add_argument(
"--api-key",
type=str,
default=None,
help="API key to send as an Authorization bearer token",
)
parser.add_argument(
"--header",
action="append",
type=parse_custom_header,
default=None,
metavar="KEY=VALUE",
help="Custom request header. Can be specified multiple times. "
"Accepts 'Header-Name: value' or 'Header-Name=value'.",
)
parser.add_argument(
"-p",
"--num-clients",
@@ -1437,6 +1488,22 @@ async def main() -> None:
help="Disable stream/streaming mode (set 'stream' to False in the API request)",
)
parser.add_argument(
"--send-conversation-id",
default=False,
action="store_true",
help=(
"Inject a `conversation_id` field into each Chat Completions "
"payload. This is a non-standard OpenAI extension consumed by "
"vLLM's disaggregated multi-turn proxy "
"(examples/disaggregated/disaggregated_serving/"
"disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. "
"Leave disabled (default) when targeting strict "
"OpenAI-compatible endpoints; enable when benchmarking the "
"disaggregated proxy."
),
)
parser.add_argument(
"-e",
"--excel-output",
@@ -1525,7 +1592,8 @@ async def main() -> None:
args.model, trust_remote_code=args.trust_remote_code
)
await get_server_info(args.url)
headers = build_request_headers(args.api_key, args.header)
await get_server_info(args.url, headers=headers)
# Load the input file (either conversations of configuration file)
logger.info(f"Reading input file: {args.input_file}")
+4 -15
View File
@@ -1,5 +1,5 @@
#!/bin/bash
# Build the vllm-rs Rust frontend binary and install it into the vllm package.
# Build vLLM Rust artifacts and install them into the vllm package.
# Usage: ./build_rust.sh [--debug]
#
# By default builds in release mode. Pass --debug for faster compile times
@@ -8,8 +8,6 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
RUST_DIR="$REPO_ROOT/rust"
TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -27,18 +25,9 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
fi
if [[ "${1:-}" == "--debug" ]]; then
PROFILE_ARGS=()
PROFILE_DIR="debug"
PROFILE_ARG="--debug"
else
PROFILE_ARGS=(--release)
PROFILE_DIR="release"
PROFILE_ARG="--release"
fi
cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
--manifest-path "$RUST_DIR/Cargo.toml" \
--bin vllm-rs \
--features native-tls-vendored
mkdir -p "$(dirname "$TARGET_PATH")"
cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
echo "Installed vllm-rs to $TARGET_PATH"
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
+22 -22
View File
@@ -255,51 +255,49 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
#################### RUST BUILD IMAGE ####################
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
# build stage doesn't need the rust toolchain, protoc, or the rust source.
# This stage runs in parallel with csrc-build/extensions-build.
FROM ${BUILD_BASE_IMAGE} AS rust-build
# This stage reuses the Python environment from base and runs in parallel with
# csrc-build/extensions-build.
FROM base AS rust-build
ARG BUILD_OS
ENV DEBIAN_FRONTEND=noninteractive
# Install a basic C toolchain (some rust crates compile C in their build.rs
# scripts) and unzip (used to extract the pinned protoc release below).
# Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
dnf install -y --setopt=install_weak_deps=False \
ca-certificates curl git gcc gcc-c++ make unzip \
make unzip \
&& dnf clean all && rm -rf /var/cache/dnf; \
else \
apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
make unzip \
&& rm -rf /var/lib/apt/lists/*; \
fi
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
# Copy only the rust workspace — the binary is the sole artifact we need.
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
# Build the release binary. Cache cargo registry/git and target/, but copy the
# binary out of the target/ cache mount so it persists into the image layer
# for later COPY --from=rust-build.
# Build the release artifacts. Cache cargo registry/git, but not target/,
# because stale target metadata can outlive source updates across BuildKit
# cache reuse.
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
#################### RUST BUILD IMAGE ####################
#################### CSRC BUILD IMAGE ####################
@@ -342,6 +340,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
WORKDIR /workspace
COPY pyproject.toml setup.py CMakeLists.txt ./
COPY tools/build_rust.py tools/build_rust.py
COPY cmake cmake/
COPY csrc csrc/
COPY vllm/envs.py vllm/envs.py
@@ -506,9 +505,10 @@ WORKDIR /workspace
COPY --from=csrc-build /workspace/dist /precompiled-wheels
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
+15 -15
View File
@@ -93,35 +93,34 @@ ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
ca-certificates curl git build-essential unzip python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
# Copy only the rust workspace — the binary is the sole artifact we need.
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
# Build the release binary. Cache cargo registry/git and target/, but copy the
# binary out of the target/ cache mount so it persists into the image layer
# for later COPY --from=rust-build.
# Build the release artifacts. Cache cargo registry/git, but not target/,
# because stale target metadata can outlive source updates across BuildKit
# cache reuse.
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
--mount=type=cache,target=/workspace/rust/target,sharing=locked \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
######################### BUILD IMAGE #########################
FROM base AS vllm-build
@@ -154,9 +153,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
+12 -11
View File
@@ -102,21 +102,22 @@ ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
ca-certificates curl git build-essential unzip python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
@@ -125,8 +126,7 @@ ENV CARGO_BUILD_JOBS=4
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
#################### RUST BUILD IMAGE ####################
#################### WHEEL BUILD IMAGE ####################
@@ -139,9 +139,10 @@ ENV UV_HTTP_TIMEOUT=500
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
RUN python3 use_existing_torch.py
+17 -16
View File
@@ -138,27 +138,25 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
ENV CARGO_NET_RETRY=10
ENV RUSTUP_MAX_RETRIES=10
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
cd ${COMMON_WORKDIR}/vllm \
&& uv pip install --system -r requirements/build/rust.txt
# Build the release binary. Cargo's registry/git caches can be written by
# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
# keeping the cache benefit. Copy the binary out so it persists into the image
# layer for later COPY --from=rust-build.
# keeping the cache benefit. Do not cache target/, because stale target metadata
# can outlive source updates across BuildKit cache reuse.
RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
--mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
--mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \
cd ${COMMON_WORKDIR}/vllm \
&& VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
&& test -x /tmp/vllm-rs
&& bash build_rust.sh \
&& test -x vllm/vllm-rs
# -----------------------
# vLLM native build stages
@@ -178,6 +176,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
# not invalidate the expensive native build layer.
COPY setup.py CMakeLists.txt ./
COPY tools/build_rust.py tools/build_rust.py
COPY cmake cmake/
COPY csrc csrc/
COPY vllm/envs.py vllm/envs.py
@@ -209,9 +208,10 @@ ENV VLLM_TARGET_DEVICE=rocm
COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
cd vllm \
@@ -418,9 +418,10 @@ FROM fetch_vllm AS build_vllm_wheel_release
ARG COMMON_WORKDIR
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
# Create /install directory for custom wheels
RUN mkdir -p /install
+12 -11
View File
@@ -6,21 +6,22 @@ ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
ca-certificates curl git build-essential unzip python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
@@ -29,8 +30,7 @@ ENV CARGO_BUILD_JOBS=4
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
@@ -213,9 +213,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
# don't invalidate heavy dependency and UCX/NIXL layers.
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
Binary file not shown.

Before

Width:  |  Height:  |  Size: 388 KiB

After

Width:  |  Height:  |  Size: 373 KiB

+41
View File
@@ -405,6 +405,47 @@ vllm bench serve \
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
#### BFCL (Tool-Calling) Benchmark
The Berkeley Function Calling Leaderboard (BFCL) dataset measures serving
latency and throughput on realistic tool-calling traffic. Each request
carries a per-sample `tools` schema and chat history, so the server must
expose `/v1/chat/completions` with an auto-tool-choice parser enabled.
The benchmark client always uses the `openai-chat` backend.
Start a tool-parser-enabled server, then run the bench. For example, with
`gpt-oss-20b`:
```bash
# Server
vllm serve openai/gpt-oss-20b \
--enable-auto-tool-choice \
--tool-call-parser openai \
--reasoning-parser openai_gptoss
# Client
vllm bench serve \
--backend openai-chat \
--endpoint /v1/chat/completions \
--model openai/gpt-oss-20b \
--dataset-name hf \
--dataset-path gorilla-llm/Berkeley-Function-Calling-Leaderboard \
--bfcl-categories simple,live_simple,multiple \
--num-prompts 200
```
`--bfcl-categories` is a comma-separated list of BFCL v3 category names
(without the `BFCL_v3_` prefix or `.json` suffix). Defaults to
`simple,live_simple,multiple`. Other supported non-multi-turn categories
include `parallel`, `live_parallel`, `parallel_multiple`,
`live_parallel_multiple`, `irrelevance`, `live_irrelevance`,
`live_relevance`, `java`, `javascript`, and `rest`. Multi-turn categories
are not yet supported.
The dataset class normalizes BFCL's loose schema dialect (`dict` →
`object`, `float` → `number`, `tuple` → `array`, `any` → `string`) so
modern grammar backends accept the translated tool definitions.
#### Other HuggingFaceDataset Examples
```bash
+2
View File
@@ -41,6 +41,8 @@ Now supports 9 types of connectors:
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
```
For multi-tier offloading (e.g., CPU + filesystem tier) and the full configuration reference, see the [KV Offloading Usage Guide](kv_offloading_usage.md).
- **FlexKVConnectorV1**: refer to [examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py](../../examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
```bash
+132
View File
@@ -0,0 +1,132 @@
# KV Offloading Usage Guide
This guide covers configuration of the [`OffloadingConnector`](disagg_prefill.md), which extends the prefix cache by offloading completed KV blocks to slower but larger tiers (CPU host memory, plus optional secondary tiers) as they are produced. Hits in the offload tiers are promoted back to GPU on demand. Transfers between GPU and CPU use DMA (`cudaMemcpyAsync`) and run asynchronously alongside model computation, so offloading adds minimal CPU- and GPU-core overhead.
!!! note
The `OffloadingConnector` currently supports CUDA, ROCm, and XPU only.
## Overview
Two specs are available, selected by the `spec_name` key in `kv_connector_extra_config`:
- `CPUOffloadingSpec` (default): single CPU tier. Completed GPU blocks are copied into pinned host memory.
- `TieringOffloadingSpec`: multi-tier. A CPU primary tier plus one or more secondary tiers.
Only the CPU primary tier has direct GPU access. Secondary tiers cannot read from or write to GPU memory; all GPU↔secondary transfers are staged through the CPU primary tier.
```mermaid
flowchart LR
GPU <--> CPU["CPU primary tier"]
CPU <--> S0["Secondary tier 0"]
CPU <--> S1["Secondary tier 1"]
CPU <--> SN["..."]
```
## Single-Tier Setup (CPU Only)
```bash
vllm serve <model> \
--kv-transfer-config '{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"block_size": 64,
"cpu_bytes_to_use": 1000000000
}
}'
```
## Multi-Tier Setup
Set `spec_name` to `"TieringOffloadingSpec"` and supply a `secondary_tiers` list. Each entry is a dict with a required `type` key plus tier-specific fields. The list is ordered: tier 0 is consulted before tier 1, and so on. See [Secondary Tiers](#secondary-tiers) for tier-specific keys.
```bash
vllm serve <model> \
--kv-transfer-config '{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"spec_name": "TieringOffloadingSpec",
"cpu_bytes_to_use": 10737418240,
"block_size": 16,
"eviction_policy": "lru",
"secondary_tiers": [
{
"type": "fs",
"root_dir": "/mnt/kv_cache",
"n_read_threads": 32,
"n_write_threads": 16
}
]
}
}'
```
## `kv_connector_extra_config` Reference
| Key | Required | Default | Scope | Notes |
| --- | --- | --- | --- | --- |
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Secondary Tiers
Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
### Filesystem (FS)
The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
| Key | Required | Default | Notes |
| --- | --- | --- | --- |
| `type` | yes | — | Must be `fs`. |
| `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). |
| `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
| `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
#### On-Disk Layout
Under `root_dir`, vLLM creates a subdirectory `<model>_<digest>`, where `<model>` is the model name with `/` replaced by `_` (so HuggingFace IDs like `meta-llama/Llama-3-8B` don't nest), and `<digest>` is a short SHA256 prefix derived from the run configuration (model, block size, parallelism, dtype, etc.). Runs with the same configuration share the same subdirectory; runs with different configurations live side-by-side under the same `root_dir` without colliding.
Inside that subdirectory, blocks are sharded across hash-prefix subdirectories to limit directory fan-out:
```text
<root_dir>/
<model>_<digest>/
config.json
<model>_<digest>_r<rank>/
<hhh>/ # first 3 hex chars of the block hash
<hh>_g<group_idx>/ # next 2 hex chars + KV cache group index
<hash_hex>.bin # full block hash (in hex)
```
`config.json` records the run (block size, number of KV groups, etc.) and is written on first start. Each rank writes blocks under its own `_r<rank>` sibling directory, so multiple ranks can safely share the same `root_dir`.
#### Cross-Process Sharing
To enable KV cache sharing between multiple vLLM instances using the same `root_dir` (e.g., via a shared PVC), the `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g., `"0"`) on every instance. Without this, each process initializes `NONE_HASH` (the chain-hash seed for block content hashes) with random bytes, producing different block filenames for identical token content.
```bash
PYTHONHASHSEED=0 vllm serve ...
```
## Tuning Tips
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
## Further Reading
- [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).
+15
View File
@@ -294,6 +294,21 @@ curl http://localhost:8000/v1/chat/completions \
!!! note
The `conversation_id` field is a non-standard extension to the OpenAI API. It is consumed by the proxy and not forwarded to the vLLM engine.
### Benchmarking the multi-turn proxy
[`benchmarks/multi_turn/benchmark_serving_multi_turn.py`](../../benchmarks/multi_turn/benchmark_serving_multi_turn.py) supports targeting the disaggregated multi-turn proxy with the `--send-conversation-id` flag, which injects a per-conversation `conversation_id` into every request payload so the proxy can key cross-turn KV cache reuse.
The flag is **off by default** so the benchmark is compatible with strict OpenAI-compatible frontends that reject unknown top-level fields. When benchmarking the multi-turn proxy you must pass it explicitly — otherwise every turn lands as a cache MISS and the bidirectional KV transfer path is never exercised.
```bash
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \
--model <MODEL> --served-model-name <NAME> \
--url http://<proxy_host>:8000 \
--input-file benchmarks/multi_turn/generate_multi_turn.json \
--num-clients 2 --max-active-conversations 6 \
--send-conversation-id
```
### Limitations
- Requires a stateful proxy (or equivalent router) to track and forward `kv_transfer_params` between turns.
+7
View File
@@ -504,6 +504,13 @@ Flags: `--tool-call-parser pythonic --chat-template {see_above}`
!!! warning
Llama's smaller models frequently fail to emit tool calls in the correct format. Results may vary depending on the model.
## Benchmarking Tool-Calling Performance
To measure serving latency and throughput on realistic tool-calling traffic,
use the BFCL (Berkeley Function Calling Leaderboard) dataset with
`vllm bench serve`. See the [BFCL benchmark example](../benchmarking/cli.md#bfcl-tool-calling-benchmark)
for the full server + client commands.
## How to Write a Tool Parser Plugin
A tool parser plugin is a Python file containing one or more ToolParser implementations. You can write a ToolParser similar to the `Hermes2ProToolParser` in [vllm/tool_parsers/hermes_tool_parser.py](../../vllm/tool_parsers/hermes_tool_parser.py).
@@ -35,12 +35,36 @@ Conversation isolation:
the JSON body) to scope the KV cache across turns. Without it, the
proxy cannot link turns and falls back to no-cache behavior.
``conversation_id`` is a non-standard extension to the OpenAI Chat
Completions schema, consumed by this proxy and not forwarded to the
vLLM engine. Strict OpenAI-compatible frontends reject unknown
fields, so clients must opt in only when targeting this proxy.
Usage:
python disagg_proxy_multiturn.py \\
--host 0.0.0.0 --port 8000 \\
--prefiller-host 10.0.0.1 --prefiller-port 8100 \\
--decoder-host 10.0.0.2 --decoder-port 8200
Benchmarking:
Use ``benchmarks/multi_turn/benchmark_serving_multi_turn.py`` with
the ``--send-conversation-id`` flag to inject a per-conversation
``conversation_id`` into every request so this proxy can key
cross-turn KV cache reuse. The flag is *off by default*: without
it the benchmark sends OpenAI-schema-compliant payloads and every
turn lands as a cache MISS in this proxy.
Example:
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\
--model <MODEL> --served-model-name <NAME> \\
--url http://<proxy_host>:8000 \\
--input-file generate_multi_turn.json \\
--num-clients 2 --max-active-conversations 6 \\
--send-conversation-id
See ``docs/features/nixl_connector_usage.md`` for the broader
bidirectional-KV-transfer setup these benchmarks exercise.
Dependencies:
pip install fastapi uvicorn httpx
"""
@@ -373,7 +397,9 @@ async def _handle_request(api_path: str, request: Request):
logger.warning(
"[%s] No conversation_id provided — KV cache reuse disabled "
"for this request. Add a 'conversation_id' field to enable "
"cross-turn KV sharing.",
"cross-turn KV sharing. When using "
"benchmarks/multi_turn/benchmark_serving_multi_turn.py, pass "
"--send-conversation-id (off by default).",
request_id,
)
+4
View File
@@ -0,0 +1,4 @@
# Dependencies for building Rust artifacts through setuptools-rust.
setuptools>=77.0.3,<81.0.0
setuptools-rust>=1.9.0
wheel
+86
View File
@@ -3458,6 +3458,75 @@ version = "0.1.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
[[package]]
name = "pyo3"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
dependencies = [
"libc",
"once_cell",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
]
[[package]]
name = "pyo3-build-config"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
dependencies = [
"libc",
"pyo3-build-config",
]
[[package]]
name = "pyo3-macros"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pyo3-macros-backend"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
dependencies = [
"heck",
"proc-macro2",
"pyo3-build-config",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pythonize"
version = "0.28.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
dependencies = [
"pyo3",
"serde",
"serde_json",
]
[[package]]
name = "qoi"
version = "0.4.1"
@@ -4669,6 +4738,12 @@ dependencies = [
"libc",
]
[[package]]
name = "target-lexicon"
version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "task-local"
version = "0.1.1"
@@ -5904,6 +5979,17 @@ dependencies = [
"winnow",
]
[[package]]
name = "vllm-tool-parser-py"
version = "0.1.0"
dependencies = [
"pyo3",
"pythonize",
"serde_json",
"thiserror-ext",
"vllm-tool-parser",
]
[[package]]
name = "walkdir"
version = "2.5.0"
+3
View File
@@ -12,6 +12,7 @@ members = [
"src/text",
"src/tokenizer",
"src/tool-parser",
"src/tool-parser/python",
]
resolver = "3"
@@ -60,6 +61,8 @@ prometheus-client = "0.24.0"
prometheus-client-derive-encode = "0.5.0"
prost = "0.14.3"
prost-types = "0.14.3"
pyo3 = "0.28.3"
pythonize = "0.28.0"
rand = "0.9.2"
reasoning-parser = "1.2.2"
reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
+1 -1
View File
@@ -71,7 +71,7 @@ To build the `vllm-rs` in isolation:
```bash
# from the local checkout
cargo install --path src/cmd --bin vllm-rs
./build_rust.sh
```
### Example Request
+1 -1
View File
@@ -272,6 +272,6 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, step3)"].assert_eq(&error.to_report_string());
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
}
}
+14 -2
View File
@@ -6,7 +6,8 @@ pub use vllm_reasoning_parser::{
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
Step3p5ReasoningParser,
};
use vllm_tokenizer::DynTokenizer;
@@ -25,7 +26,9 @@ pub mod names {
pub const MINIMAX_M2: &str = "minimax_m2";
pub const NEMOTRON_V3: &str = "nemotron_v3";
pub const QWEN3: &str = "qwen3";
pub const SEED_OSS: &str = "seed_oss";
pub const STEP3: &str = "step3";
pub const STEP3P5: &str = "step3p5";
}
/// Constructor signature for one registered reasoning parser implementation.
@@ -61,7 +64,9 @@ impl ReasoningParserFactory {
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
.register_parser::<Qwen3ReasoningParser>(names::QWEN3)
.register_parser::<Step3ReasoningParser>(names::STEP3);
.register_parser::<SeedOssReasoningParser>(names::SEED_OSS)
.register_parser::<Step3ReasoningParser>(names::STEP3)
.register_parser::<Step3p5ReasoningParser>(names::STEP3P5);
factory
.register_pattern("deepseek-r1", names::DEEPSEEK_R1)
@@ -77,7 +82,14 @@ impl ReasoningParserFactory {
.register_pattern("glm-4.5", names::GLM45)
.register_pattern("kimi-k2", names::KIMI_K2)
.register_pattern("kimi", names::KIMI)
// step3p5 patterns must precede `step3`: substring matching would
// otherwise route step3p5 IDs to step3.
.register_pattern("step-3p5", names::STEP3P5)
.register_pattern("step3p5", names::STEP3P5)
.register_pattern("step-3.5", names::STEP3P5)
.register_pattern("step3", names::STEP3)
.register_pattern("seed-oss", names::SEED_OSS)
.register_pattern("seedoss", names::SEED_OSS)
.register_pattern("minimax", names::MINIMAX_M2)
.register_pattern("mm-m2", names::MINIMAX_M2)
.register_pattern("cohere", names::COHERE_CMD)
@@ -32,8 +32,12 @@ fn factory_contains_and_lists_registered_parsers() {
let factory = ReasoningParserFactory::new();
assert!(factory.contains(names::QWEN3));
assert!(factory.contains(names::DEEPSEEK_V4));
assert!(factory.contains(names::SEED_OSS));
assert!(factory.contains(names::STEP3P5));
assert!(factory.list().contains(&names::QWEN3.to_string()));
assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
assert!(factory.list().contains(&names::SEED_OSS.to_string()));
assert!(factory.list().contains(&names::STEP3P5.to_string()));
}
#[test]
@@ -49,6 +53,41 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
);
}
#[test]
fn factory_routes_step3p5_models_to_dedicated_parser() {
let factory = ReasoningParserFactory::new();
// step3p5 patterns must beat the bare `step3` substring.
assert_eq!(
factory.resolve_name_for_model("step-3p5-instruct"),
Some(names::STEP3P5)
);
assert_eq!(
factory.resolve_name_for_model("step3p5"),
Some(names::STEP3P5)
);
assert_eq!(
factory.resolve_name_for_model("step-3.5-base"),
Some(names::STEP3P5)
);
assert_eq!(
factory.resolve_name_for_model("step3-base"),
Some(names::STEP3)
);
}
#[test]
fn factory_routes_seed_oss_models() {
let factory = ReasoningParserFactory::new();
assert_eq!(
factory.resolve_name_for_model("ByteDance-Seed/Seed-OSS-36B-Instruct"),
Some(names::SEED_OSS)
);
assert_eq!(
factory.resolve_name_for_model("seedoss-7b"),
Some(names::SEED_OSS)
);
}
#[test]
fn factory_rejects_unknown_parser_names() {
let tokenizer = Arc::new(FakeTokenizer);
+26
View File
@@ -144,6 +144,30 @@ impl RoundtripCase {
json_fmt: spaced_json_fmt(),
}
}
/// SeedOSS with `<seed:think>` / `</seed:think>` reasoning tags.
fn seed_oss() -> Self {
Self {
model_id: "ByteDance-Seed/Seed-OSS-36B-Instruct",
assistant_stop_suffix: "<seed:eos>",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
}
}
/// Step-3.5 with `<think>` / `</think>` reasoning tags and newline trimming.
fn step3p5() -> Self {
Self {
model_id: "stepfun-ai/Step-3.5-Flash",
assistant_stop_suffix: "<|im_end|>\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
}
}
}
macro_rules! roundtrip_tests {
@@ -168,6 +192,8 @@ roundtrip_tests! {
minimax_m25 => [reasoning_and_content, tool_call_mix],
deepseek_v4 => [reasoning_and_content, tool_call_mix],
glm47 => [reasoning_and_content, tool_call_mix],
seed_oss => [reasoning_and_content],
step3p5 => [reasoning_and_content],
// Note: Kimi K2.5 strips the reasoning content in history.
kimi_k25 => [tool_call_mix],
@@ -68,6 +68,11 @@ impl DelimitedReasoningParser {
.unwrap_or(self.default_in_reasoning);
}
/// Return whether the parser is currently inside a reasoning section.
pub(crate) fn in_reasoning(&self) -> bool {
self.current_in_reasoning
}
/// Parse one decoded text delta and return its reasoning/content split.
pub(crate) fn push(&mut self, delta: &str) -> ReasoningDelta {
self.buffer.push_str(delta);
+4
View File
@@ -20,6 +20,8 @@ mod delimited;
mod gemma4;
mod kimi;
mod qwen3;
mod seed_oss;
mod step3p5;
use thiserror::Error;
use vllm_tokenizer::DynTokenizer;
@@ -30,6 +32,8 @@ pub(crate) use self::delimited::DelimitedReasoningParser;
pub use self::gemma4::Gemma4ReasoningParser;
pub use self::kimi::KimiReasoningParser;
pub use self::qwen3::Qwen3ReasoningParser;
pub use self::seed_oss::SeedOssReasoningParser;
pub use self::step3p5::Step3p5ReasoningParser;
/// DeepSeek V3 currently shares the standard `<think>...</think>` parser.
pub type DeepSeekV3ReasoningParser = Qwen3ReasoningParser;
+147
View File
@@ -0,0 +1,147 @@
use vllm_tokenizer::DynTokenizer;
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
/// Reasoning parser for SeedOSS models using `<seed:think>`/`</seed:think>`
/// delimiters.
pub struct SeedOssReasoningParser {
inner: DelimitedReasoningParser,
}
impl SeedOssReasoningParser {
/// Create a SeedOSS parser backed by the shared delimited state machine.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(
tokenizer,
"<seed:think>",
"</seed:think>",
false,
)?,
})
}
}
impl ReasoningParser for SeedOssReasoningParser {
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tokenizer)?))
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids);
Ok(())
}
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
Ok(self.inner.push(delta))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
Ok(self.inner.finish())
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::SeedOssReasoningParser;
use crate::{ReasoningParser, tests::FakeTokenizer};
#[test]
fn without_prompt_markers_expects_start_token() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(
delta.content.as_deref(),
Some("implicit reasoning</seed:think>answer")
);
}
#[test]
fn picks_up_prompt_start_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
// Prompt prefills `<seed:think>` (id 10), opening reasoning before the stream.
parser.initialize(&[10]).unwrap();
let delta = parser.push("reason</seed:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn respects_prompt_end_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
// Prompt already closed reasoning with `</seed:think>` (id 11).
parser.initialize(&[11]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn handles_explicit_start_token() {
// An explicit start delimiter must not leak into reasoning text.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn streams_explicit_start_token_across_pushes() {
// Start token, reasoning body, end token, and content arrive in separate
// streaming deltas.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let mut reasoning = String::new();
let mut content = String::new();
for delta_str in [
"<seed:think>",
"Some ",
"reasoning ",
"content",
"</seed:think>",
"Final ",
"answer",
] {
let delta = parser.push(delta_str).unwrap();
if let Some(r) = delta.reasoning {
reasoning.push_str(&r);
}
if let Some(c) = delta.content {
content.push_str(&c);
}
}
assert_eq!(reasoning, "Some reasoning content");
assert_eq!(content, "Final answer");
}
#[test]
fn handles_partial_delimiters_across_pushes() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[10]).unwrap();
// Closing delimiter `</seed:think>` arrives in two halves.
let first = parser.push("reason</seed:").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
assert_eq!(first.content, None);
let second = parser.push("think>answer").unwrap();
assert_eq!(second.reasoning, None);
assert_eq!(second.content.as_deref(), Some("answer"));
}
}
+308
View File
@@ -0,0 +1,308 @@
use vllm_tokenizer::DynTokenizer;
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
/// Reasoning parser for Step3p5 outputs.
///
/// Step3p5 uses standard `<think>`/`</think>` delimiters but emits a `\n`
/// immediately before and/or after `</think>`. The parser drops these framing
/// newlines on both sides of the boundary, holding a trailing `\n` from
/// reasoning across pushes until either more reasoning text or `</think>`
/// arrives, and dropping a leading `\n` from the first content delta after
/// the boundary.
pub struct Step3p5ReasoningParser {
inner: DelimitedReasoningParser,
/// `\n` at end of last reasoning delta, held in case `</think>` follows.
pending_reasoning_newline: bool,
/// Last push ended on `</think>` without emitting content; the next
/// content delta's leading `\n` should be dropped.
just_ended_reasoning: bool,
}
impl Step3p5ReasoningParser {
/// Create a Step3p5 parser backed by the shared delimited state machine.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false)?,
pending_reasoning_newline: false,
just_ended_reasoning: false,
})
}
/// Drop framing newlines around `</think>` and track held-newline state.
fn process(
&mut self,
mut inner_delta: ReasoningDelta,
was_in_reasoning: bool,
now_in_reasoning: bool,
) -> ReasoningDelta {
// A `<think>...</think>` round-trip in one push still counts as a
// transition: the inner emits reasoning while ending in content mode.
let transitioned =
!now_in_reasoning && (was_in_reasoning || inner_delta.reasoning.is_some());
// Replay or drop a previously-held trailing reasoning newline.
if self.pending_reasoning_newline {
if let Some(reasoning) = inner_delta.reasoning.as_mut() {
reasoning.insert(0, '\n');
self.pending_reasoning_newline = false;
} else if transitioned {
// The held `\n` was the one right before `</think>`: drop it.
self.pending_reasoning_newline = false;
}
}
// Hold back a trailing reasoning `\n` until we know if `</think>` follows.
if let Some(reasoning) = inner_delta.reasoning.as_mut()
&& reasoning.ends_with('\n')
{
reasoning.pop();
if !transitioned {
self.pending_reasoning_newline = true;
}
}
// Drop a leading `\n` of content emitted right after `</think>`.
if let Some(content) = inner_delta.content.as_mut()
&& (transitioned || self.just_ended_reasoning)
&& content.starts_with('\n')
{
content.remove(0);
}
self.just_ended_reasoning = transitioned && inner_delta.content.is_none();
if inner_delta.reasoning.as_deref() == Some("") {
inner_delta.reasoning = None;
}
if inner_delta.content.as_deref() == Some("") {
inner_delta.content = None;
}
inner_delta
}
}
impl ReasoningParser for Step3p5ReasoningParser {
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tokenizer)?))
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids);
Ok(())
}
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
let was = self.inner.in_reasoning();
let inner_delta = self.inner.push(delta);
let now = self.inner.in_reasoning();
Ok(self.process(inner_delta, was, now))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
let was = self.inner.in_reasoning();
let inner_delta = self.inner.finish();
let now = self.inner.in_reasoning();
let mut delta = self.process(inner_delta, was, now);
// Emit a still-held newline rather than silently dropping it.
if self.pending_reasoning_newline {
match delta.reasoning.as_mut() {
Some(existing) => existing.push('\n'),
None => delta.reasoning = Some("\n".to_string()),
}
self.pending_reasoning_newline = false;
}
Ok(delta)
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::Step3p5ReasoningParser;
use crate::{ReasoningParser, tests::FakeTokenizer};
#[test]
fn picks_up_prompt_start_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
// Prompt prefills `<think>` (id 1), opening reasoning before the stream.
parser.initialize(&[1]).unwrap();
let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();
assert_eq!(
delta.reasoning.as_deref(),
Some("This is a reasoning section")
);
assert_eq!(delta.content.as_deref(), Some("This is the rest"));
}
#[test]
fn handles_unterminated_reasoning() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let pushed = parser.push("<think>reason without end").unwrap();
assert_eq!(pushed.reasoning.as_deref(), Some("reason without end"));
assert_eq!(pushed.content, None);
let flushed = parser.finish().unwrap();
assert!(flushed.is_empty());
}
#[test]
fn handles_empty_input() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let pushed = parser.push("").unwrap();
assert!(pushed.is_empty());
let flushed = parser.finish().unwrap();
assert!(flushed.is_empty());
}
#[test]
fn complex_newline_pattern_trims_only_single_framing_newline_each_side() {
// Only the immediately-adjacent framing `\n` is dropped on each side of
// `</think>`; surrounding newlines remain part of reasoning/content.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[1]).unwrap();
let delta = parser
.push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")
.unwrap();
assert_eq!(
delta.reasoning.as_deref(),
Some("\n This is a \n reasoning section\n\n")
);
assert_eq!(delta.content.as_deref(), Some("\nThis is the rest"));
}
#[test]
fn drops_framing_newlines_in_single_push() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn drops_framing_newlines_across_pushes() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
// The trailing `\n` from the first push is held until we know whether
// `</think>` follows.
let first = parser.push("<think>reason\n").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
assert_eq!(first.content, None);
// `</think>` arrives standalone; the held newline should be dropped.
let second = parser.push("</think>").unwrap();
assert!(second.is_empty());
// The leading newline of the first content delta is dropped.
let third = parser.push("\nanswer").unwrap();
assert_eq!(third.reasoning, None);
assert_eq!(third.content.as_deref(), Some("answer"));
}
#[test]
fn replays_held_newline_when_more_reasoning_follows() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason\n").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
let second = parser.push("more reason").unwrap();
assert_eq!(second.reasoning.as_deref(), Some("\nmore reason"));
assert_eq!(second.content, None);
}
#[test]
fn finish_flushes_held_newline_in_unterminated_stream() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason\n").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
let flushed = parser.finish().unwrap();
assert_eq!(flushed.reasoning.as_deref(), Some("\n"));
assert_eq!(flushed.content, None);
}
#[test]
fn preserves_inner_newlines_in_reasoning() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>line1\nline2</think>tail").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("line1\nline2"));
assert_eq!(delta.content.as_deref(), Some("tail"));
}
#[test]
fn trims_only_one_trailing_reasoning_newline() {
// Only the single framing newline immediately before `</think>` is
// dropped; earlier newlines in the reasoning body are preserved.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason\n\n</think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason\n"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn drops_only_first_content_newline_after_transition() {
// The leading-`\n` drop applies only to the first content delta after
// `</think>`; later deltas pass through untouched.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason</think>").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
assert_eq!(first.content, None);
let second = parser.push("\nfirst").unwrap();
assert_eq!(second.reasoning, None);
assert_eq!(second.content.as_deref(), Some("first"));
// A `\n` arriving in a later content delta must NOT be dropped.
let third = parser.push("\nsecond").unwrap();
assert_eq!(third.reasoning, None);
assert_eq!(third.content.as_deref(), Some("\nsecond"));
}
#[test]
fn passes_through_clean_boundary_without_framing_newlines() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason</think>tail").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("tail"));
}
#[test]
fn handles_empty_reasoning_section() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think></think>answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
}
+3 -1
View File
@@ -6,7 +6,7 @@ use super::{
DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
};
struct FakeTokenizer;
pub(crate) struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
@@ -32,6 +32,8 @@ impl Tokenizer for FakeTokenizer {
"<|END_THINKING|>" => Some(4),
"◁think▷" => Some(5),
"◁/think▷" => Some(6),
"<seed:think>" => Some(10),
"</seed:think>" => Some(11),
_ => None,
}
}
+2 -1
View File
@@ -14,7 +14,8 @@ mod utils;
use std::sync::{Arc, OnceLock};
use anyhow::{Context as _, Result};
use axum::{Router, serve::ListenerExt as _};
use axum::Router;
use axum::serve::ListenerExt as _;
pub use config::{Config, CoordinatorMode, HttpListenerMode};
use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until};
@@ -22,7 +22,7 @@ use vllm_llm::{
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
};
use self::convert::prepare_generate_request;
use self::convert::{ResponseOptions, prepare_generate_request};
use self::types::{
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
GenerateResponseStreamChoice, GenerateStreamResponse,
@@ -54,10 +54,7 @@ pub async fn generate(
);
let log_request = state.enable_log_requests;
let include_logprobs = prepared.include_logprobs;
let include_prompt_logprobs = prepared.include_prompt_logprobs;
let stream = prepared.stream;
let raw_stream = match state
.chat
.text()
@@ -80,9 +77,7 @@ pub async fn generate(
raw_stream,
prepared.request_id,
log_request,
prepared.include_usage,
prepared.include_continuous_usage,
include_logprobs,
prepared.options,
);
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
@@ -100,21 +95,11 @@ pub async fn generate(
}
};
if log_request {
info!(
parent: &request_span,
prompt_tokens = collected.prompt_token_ids.len(),
output_tokens = collected.token_ids.len(),
finish_reason = collected.finish_reason.as_str(),
"generate finished"
);
}
let response = match collect_generate(
collected,
prepared.request_id,
include_logprobs,
include_prompt_logprobs,
log_request,
prepared.options,
) {
Ok(response) => response,
Err(error) => return error.into_response(),
@@ -128,9 +113,13 @@ async fn generate_chunk_stream(
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
request_id: String,
log_request: bool,
include_usage: bool,
include_continuous_usage: bool,
include_logprobs: bool,
ResponseOptions {
include_usage,
include_continuous_usage,
include_logprobs,
// Ignored: raw generate streaming has no prompt-logprobs wire shape.
include_prompt_logprobs: _,
}: ResponseOptions,
mut y: TryYielder<GenerateStreamResponse, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
@@ -222,8 +211,15 @@ async fn generate_chunk_stream(
fn collect_generate(
collected: CollectedGenerateOutput,
request_id: String,
include_logprobs: bool,
include_prompt_logprobs: bool,
log_request: bool,
ResponseOptions {
// Ignored: non-streaming raw generate responses do not include usage.
include_usage: _,
// Ignored: continuous usage is a streaming-only option.
include_continuous_usage: _,
include_logprobs,
include_prompt_logprobs,
}: ResponseOptions,
) -> Result<GenerateResponse, ApiError> {
let logprobs = if include_logprobs {
let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
@@ -246,13 +242,23 @@ fn collect_generate(
} else {
None
};
let finish_reason = collected.finish_reason.as_str().to_string();
if log_request {
info!(
prompt_tokens = collected.prompt_token_ids.len(),
output_tokens = collected.token_ids.len(),
%finish_reason,
"generate finished"
);
}
Ok(GenerateResponse {
request_id,
choices: vec![GenerateResponseChoice {
index: 0,
logprobs,
finish_reason: Some(collected.finish_reason.as_str().to_string()),
finish_reason: Some(finish_reason),
token_ids: collected.token_ids,
}],
prompt_logprobs,
@@ -408,11 +414,19 @@ mod tests {
}),
]);
let chunks: Vec<_> =
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
.try_collect()
.await
.expect("collect chunks");
let chunks: Vec<_> = generate_chunk_stream(
stream,
"raw-stream".to_string(),
false,
ResponseOptions {
include_usage: true,
include_continuous_usage: true,
..Default::default()
},
)
.try_collect()
.await
.expect("collect chunks");
assert_eq!(chunks.len(), 2);
assert_eq!(
@@ -8,19 +8,29 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
/// Lowered generate request plus the response request ID.
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedRequest {
pub(super) struct PreparedRequest {
pub request_id: String,
pub text_request: TextRequest,
pub stream: bool,
/// Public response rendering options for route-layer helpers.
pub options: ResponseOptions,
}
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub(super) struct ResponseOptions {
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Whether the caller asked for usage on every streamed chunk.
pub include_continuous_usage: bool,
/// Whether the caller requested output logprobs on generate choices.
pub include_logprobs: bool,
/// Whether the caller requested top-level prompt logprobs.
pub include_prompt_logprobs: bool,
}
/// Validate and lower one raw generate request into the internal
/// text-generation format.
pub fn prepare_generate_request(
pub(super) fn prepare_generate_request(
request: GenerateRequest,
lora_resolution: &LoraModelResolution,
ctx: ResolvedRequestContext,
@@ -65,10 +75,12 @@ pub fn prepare_generate_request(
request_id: ctx.request_id,
text_request,
stream,
include_usage,
include_continuous_usage,
include_logprobs,
include_prompt_logprobs,
options: ResponseOptions {
include_usage,
include_continuous_usage,
include_logprobs,
include_prompt_logprobs,
},
})
}
@@ -158,7 +170,7 @@ mod tests {
)
.expect("prepare");
assert!(!prepared.include_usage);
assert!(!prepared.include_continuous_usage);
assert!(!prepared.options.include_usage);
assert!(!prepared.options.include_continuous_usage);
}
}
@@ -1,4 +1,4 @@
pub mod convert;
pub(crate) mod convert;
mod types;
mod validate;
@@ -23,8 +23,8 @@ use vllm_chat::{
};
use vllm_engine_core_client::protocol::StopReason;
use self::convert::{ResponseOptions, prepare_chat_request};
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::chat_completions::convert::prepare_chat_request;
use crate::routes::openai::chat_completions::types::{
AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
ChatCompletionResponse, ChatCompletionStreamChoice, ChatCompletionStreamResponse,
@@ -83,12 +83,7 @@ pub async fn chat_completions(
prepared.response_model,
created,
log_request,
prepared.include_usage,
prepared.requested_logprobs,
prepared.include_reasoning,
prepared.echo,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
prepared.options,
);
let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
@@ -99,12 +94,8 @@ pub async fn chat_completions(
prepared.request_id,
prepared.response_model,
created,
prepared.requested_logprobs,
prepared.include_prompt_logprobs,
prepared.include_reasoning,
prepared.echo,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
log_request,
prepared.options,
)
.instrument(request_span.clone())
.await
@@ -113,18 +104,6 @@ pub async fn chat_completions(
Err(error) => return error.into_response(),
};
if log_request {
let usage = response.usage.as_ref();
info!(
parent: &request_span,
model = %response.model,
prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
"chat completion finished"
);
}
Json(response).into_response()
}
}
@@ -134,12 +113,17 @@ async fn collect_chat_completion(
request_id: String,
response_model: String,
created: u64,
requested_logprobs: bool,
include_prompt_logprobs: bool,
include_reasoning: bool,
echo: Option<String>,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
log_request: bool,
ResponseOptions {
// Ignored: non-streaming responses always include usage.
include_usage: _,
requested_logprobs,
include_prompt_logprobs,
include_reasoning,
echo,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
) -> Result<ChatCompletionResponse, ApiError> {
let collected = stream.collect_message().await.map_err(|error| {
server_error!(
@@ -201,6 +185,16 @@ async fn collect_chat_completion(
};
let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32);
if log_request {
info!(
model = %response_model,
prompt_tokens = usage.prompt_tokens,
output_tokens = usage.completion_tokens.unwrap_or(0),
finish_reason = %finish_reason,
"chat completion finished"
);
}
Ok(ChatCompletionResponse {
id: request_id,
object: "chat.completion".to_string(),
@@ -238,12 +232,16 @@ async fn chat_completion_chunk_stream(
response_model: String,
created: u64,
log_request: bool,
include_usage: bool,
requested_logprobs: bool,
include_reasoning: bool,
echo: Option<String>,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
ResponseOptions {
include_usage,
requested_logprobs,
// Ignored: chat streaming prompt logprobs are rejected for Python parity.
include_prompt_logprobs: _,
include_reasoning,
echo,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
) -> Result<(), ApiError> {
let mut saw_tool_calls = false;
@@ -806,7 +804,7 @@ mod tests {
use vllm_engine_core_client::protocol::StopReason;
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
use super::{block_delta_chunk, chat_completion_chunk_stream, final_chunk};
use super::{ResponseOptions, block_delta_chunk, chat_completion_chunk_stream, final_chunk};
#[test]
fn text_chunk_uses_content_only_delta() {
@@ -932,12 +930,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
true,
None,
false,
false,
ResponseOptions {
requested_logprobs: true,
include_reasoning: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -996,12 +993,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
true,
None,
false,
false,
ResponseOptions {
requested_logprobs: true,
include_reasoning: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -1049,12 +1045,7 @@ mod tests {
"model".to_string(),
1,
false,
false,
false,
false,
None,
false,
false,
ResponseOptions::default(),
)
.collect::<Vec<_>>()
.await
@@ -1132,12 +1123,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
false,
None,
true,
false,
ResponseOptions {
requested_logprobs: true,
return_token_ids: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -1263,12 +1253,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
false,
None,
true,
false,
ResponseOptions {
requested_logprobs: true,
return_token_ids: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -1342,12 +1331,10 @@ mod tests {
"model".to_string(),
1,
false,
false,
false,
true,
None,
false,
false,
ResponseOptions {
include_reasoning: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -18,11 +18,19 @@ use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer
/// Lowered chat request plus the public response metadata carried by every SSE
/// chunk.
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedRequest {
pub(super) struct PreparedRequest {
/// Stable OpenAI-style request ID, reused as the external chat request ID.
pub request_id: String,
/// Public model ID echoed back to the client.
pub response_model: String,
/// Public response rendering options for route-layer helpers.
pub options: ResponseOptions,
/// Lowered chat request for `vllm-chat`.
pub chat_request: ChatRequest,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub(super) struct ResponseOptions {
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Whether the caller requested output logprobs on chat choices.
@@ -31,8 +39,6 @@ pub struct PreparedRequest {
pub include_prompt_logprobs: bool,
/// Whether to include reasoning content in OpenAI responses.
pub include_reasoning: bool,
/// Lowered chat request for `vllm-chat`.
pub chat_request: ChatRequest,
/// Last assistant-role message content to echo back when `echo=true`.
pub echo: Option<String>,
/// Whether to include token IDs alongside generated text.
@@ -46,7 +52,7 @@ pub struct PreparedRequest {
///
/// `lora_resolution.model_names` must be non-empty; the first entry is used as
/// the base `model` field in responses when no LoRA adapter is selected.
pub(crate) fn prepare_chat_request(
pub(super) fn prepare_chat_request(
request: ChatCompletionRequest,
lora_resolution: &LoraModelResolution,
ctx: ResolvedRequestContext,
@@ -146,14 +152,16 @@ pub(crate) fn prepare_chat_request(
Ok(PreparedRequest {
request_id,
response_model,
include_usage,
requested_logprobs,
include_prompt_logprobs,
include_reasoning,
options: ResponseOptions {
include_usage,
requested_logprobs,
include_prompt_logprobs,
include_reasoning,
echo,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
},
chat_request,
echo,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
})
}
pub(crate) fn normalize_generation_prompt_mode(
@@ -497,7 +505,7 @@ mod tests {
)
.expect("request is valid");
assert!(!prepared.include_reasoning);
assert!(!prepared.options.include_reasoning);
}
#[test]
@@ -867,8 +875,8 @@ mod tests {
)
.expect("request is valid");
assert!(prepared.requested_logprobs);
assert!(prepared.include_prompt_logprobs);
assert!(prepared.options.requested_logprobs);
assert!(prepared.options.include_prompt_logprobs);
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0));
assert_eq!(
prepared.chat_request.sampling_params.prompt_logprobs,
@@ -894,7 +902,7 @@ mod tests {
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3));
assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None);
assert!(!prepared.include_prompt_logprobs);
assert!(!prepared.options.include_prompt_logprobs);
}
#[test]
@@ -18,13 +18,13 @@ use tracing::{debug, error, info, trace};
use tracing_futures::Instrument as _;
use vllm_text::{DecodedTextEvent, FinishReason, TextOutputStream, TextOutputStreamExt as _};
use self::convert::{ResponseOptions, prepare_completion_request};
use super::utils::logprobs::{
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
text_len,
};
use super::utils::types::Usage;
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::completions::convert::prepare_completion_request;
use crate::routes::openai::completions::types::{
CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
CompletionStreamChoice, CompletionStreamResponse,
@@ -42,7 +42,6 @@ pub async fn completions(
ValidatedJson(body): ValidatedJson<CompletionRequest>,
) -> Response {
let stream = body.stream;
let logprobs = body.logprobs;
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
@@ -57,9 +56,7 @@ pub async fn completions(
);
let created = unix_timestamp();
let include_prompt_logprobs = prepared.text_request.sampling_params.prompt_logprobs.is_some();
let log_request = state.enable_log_requests;
let text_stream = match state
.chat
.text()
@@ -84,11 +81,7 @@ pub async fn completions(
prepared.response_model,
created,
log_request,
prepared.include_usage,
prepared.echo,
logprobs,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
prepared.options,
);
let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
@@ -99,11 +92,8 @@ pub async fn completions(
prepared.request_id,
prepared.response_model,
created,
prepared.echo,
logprobs,
include_prompt_logprobs,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
log_request,
prepared.options,
)
.instrument(request_span.clone())
.await
@@ -112,18 +102,6 @@ pub async fn completions(
Err(error) => return error.into_response(),
};
if log_request {
let usage = response.usage.as_ref();
info!(
parent: &request_span,
model = %response.model,
prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
"completion finished"
);
}
Json(response).into_response()
}
}
@@ -133,11 +111,16 @@ async fn collect_completion(
request_id: String,
response_model: String,
created: u64,
echo: Option<String>,
requested_logprobs: Option<u32>,
include_prompt_logprobs: bool,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
log_request: bool,
ResponseOptions {
// Ignored: non-streaming responses always include usage.
include_usage: _,
echo,
requested_logprobs,
include_prompt_logprobs,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
) -> Result<CompletionResponse, ApiError> {
let collected = stream
.collect_output()
@@ -175,6 +158,21 @@ async fn collect_completion(
None => collected.text,
Some(prompt) => format!("{prompt}{}", collected.text),
};
let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string();
let usage = Usage::from_counts(
collected.prompt_token_ids.len() as u32,
collected.token_ids.len() as u32,
);
if log_request {
info!(
model = %response_model,
prompt_tokens = usage.prompt_tokens,
output_tokens = usage.completion_tokens.unwrap_or(0),
%finish_reason,
"completion finished"
);
}
Ok(CompletionResponse {
id: request_id,
@@ -185,16 +183,13 @@ async fn collect_completion(
index: 0,
text,
logprobs,
finish_reason: Some(completion_finish_reason_to_openai(finish_reason)?.into()),
finish_reason: Some(finish_reason),
stop_reason,
prompt_logprobs,
token_ids: return_token_ids.then(|| collected.token_ids.clone()),
prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()),
}],
usage: Some(Usage::from_counts(
collected.prompt_token_ids.len() as u32,
collected.token_ids.len() as u32,
)),
usage: Some(usage),
system_fingerprint: None,
kv_transfer_params: collected.kv_transfer_params,
})
@@ -208,11 +203,15 @@ async fn completion_chunk_stream(
response_model: String,
created: u64,
log_request: bool,
include_usage: bool,
echo: Option<String>,
requested_logprobs: Option<u32>,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
ResponseOptions {
include_usage,
echo,
requested_logprobs,
// Ignored: streaming prompt logprobs are rejected for Python parity.
include_prompt_logprobs: _,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
mut y: TryYielder<CompletionSseChunk, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
@@ -432,7 +431,7 @@ mod tests {
FinishReason, Finished,
};
use super::{CompletionSseChunk, completion_chunk_stream, final_chunk};
use super::{CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk};
#[test]
fn final_chunk_maps_stop_finish_reason() {
@@ -527,11 +526,10 @@ mod tests {
"model".to_string(),
1,
false,
false,
None,
Some(1),
false,
false,
ResponseOptions {
requested_logprobs: Some(1),
..Default::default()
},
)
.collect::<Vec<_>>()
.await;
@@ -10,18 +10,28 @@ use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer
/// Lowered completion request plus the public response metadata carried by
/// every SSE chunk.
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedRequest {
pub(super) struct PreparedRequest {
/// Stable OpenAI-style request ID, reused as the external text request ID.
pub request_id: String,
/// Public model ID echoed back to the client.
pub response_model: String,
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Public response rendering options for route-layer helpers.
pub options: ResponseOptions,
/// Lowered text request for the shared `vllm-text` facade.
pub text_request: TextRequest,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub(super) struct ResponseOptions {
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Original text prompt that should be echoed back northbound when
/// `echo=true`.
pub echo: Option<String>,
/// Whether the caller requested output logprobs on completion choices.
pub requested_logprobs: Option<u32>,
/// Whether the caller requested choice-level prompt logprobs.
pub include_prompt_logprobs: bool,
/// Whether to include token IDs alongside generated text.
pub return_token_ids: bool,
/// Whether to format logprob tokens as `token_id:{id}`.
@@ -33,7 +43,7 @@ pub struct PreparedRequest {
///
/// `lora_resolution.model_names` must be non-empty; the first entry is used as
/// the base `model` field in responses when no LoRA adapter is selected.
pub(crate) fn prepare_completion_request(
pub(super) fn prepare_completion_request(
request: CompletionRequest,
lora_resolution: &LoraModelResolution,
ctx: ResolvedRequestContext,
@@ -64,6 +74,7 @@ pub(crate) fn prepare_completion_request(
let include_usage = (request.stream_options.as_ref())
.and_then(|options| options.include_usage)
.unwrap_or(false);
let include_prompt_logprobs = prompt_logprobs.is_some();
let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();
let structured_outputs =
@@ -116,11 +127,15 @@ pub(crate) fn prepare_completion_request(
Ok(PreparedRequest {
request_id,
response_model,
include_usage,
options: ResponseOptions {
include_usage,
echo,
requested_logprobs: request.logprobs,
include_prompt_logprobs,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
},
text_request,
echo,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
})
}
@@ -206,7 +221,7 @@ mod tests {
)
.expect("prepare");
assert!(prepared.include_usage);
assert!(prepared.options.include_usage);
assert_eq!(
prepared.text_request.prompt,
Prompt::TokenIds(vec![11, 22, 33])
@@ -250,7 +265,7 @@ mod tests {
)
.expect("prepare");
assert_eq!(prepared.echo, Some("hello".to_string()));
assert_eq!(prepared.options.echo, Some("hello".to_string()));
assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
}
+3 -92
View File
@@ -14,15 +14,14 @@ use std::{fmt, fs};
use axum::body::{Body, to_bytes};
use axum::http::{Request, StatusCode};
use bytes::Bytes;
use futures::StreamExt as _;
use rmpv::Value;
use serde_json::json;
use serial_test::serial;
use tower::{Service as _, ServiceExt as _};
use vllm_chat::{
ChatBackend, ChatContent, ChatContentPart, ChatEvent, ChatLlm, ChatMessage, ChatRenderer,
ChatRequest, ChatRole, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor,
DynChatRenderer, NewChatOutputProcessorOptions, SamplingParams,
ChatBackend, ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRenderer, ChatRequest,
ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,
NewChatOutputProcessorOptions,
};
use vllm_engine_core_client::protocol::logprobs::{
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
@@ -44,8 +43,6 @@ use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora};
use crate::lora::LoraModelResolution;
use crate::routes::openai::chat_completions::convert::prepare_chat_request;
use crate::state::AppState;
fn request_output(
@@ -3478,92 +3475,6 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() {
assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn chat_harness_streams_text_events() {
let (chat, engine_task) = test_chat_with_engine_handle().await;
let mut stream = chat
.chat(ChatRequest {
messages: vec![ChatMessage::text(ChatRole::User, "hello")],
sampling_params: SamplingParams {
max_tokens: Some(8),
..Default::default()
},
request_id: "chat-harness".to_string(),
..ChatRequest::for_test()
})
.await
.expect("submit chat request");
let mut saw_text = false;
let mut saw_done = false;
while let Some(event) = stream.next().await {
match event.expect("chat event") {
ChatEvent::BlockDelta { .. } => saw_text = true,
ChatEvent::Done { .. } => {
saw_done = true;
break;
}
ChatEvent::Start { .. }
| ChatEvent::LogprobsDelta { .. }
| ChatEvent::BlockStart { .. }
| ChatEvent::BlockEnd { .. }
| ChatEvent::ToolCallStart { .. }
| ChatEvent::ToolCallArgumentsDelta { .. }
| ChatEvent::ToolCallEnd { .. } => {}
}
}
engine_task.await.expect("mock engine task");
assert!(saw_text);
assert!(saw_done);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn prepared_openai_request_streams_text_events() {
let (chat, engine_task) = test_chat_with_engine_handle().await;
let prepared = prepare_chat_request(
serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
}))
.expect("decode request"),
&LoraModelResolution {
model_names: vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
lora_request: None,
},
crate::utils::ResolvedRequestContext::default(),
)
.expect("prepare request");
let mut stream = chat.chat(prepared.chat_request).await.expect("submit chat request");
let mut saw_text = false;
let mut saw_done = false;
while let Some(event) = stream.next().await {
match event.expect("chat event") {
ChatEvent::BlockDelta { .. } => saw_text = true,
ChatEvent::Done { .. } => {
saw_done = true;
break;
}
ChatEvent::Start { .. }
| ChatEvent::LogprobsDelta { .. }
| ChatEvent::BlockStart { .. }
| ChatEvent::BlockEnd { .. }
| ChatEvent::ToolCallStart { .. }
| ChatEvent::ToolCallArgumentsDelta { .. }
| ChatEvent::ToolCallEnd { .. } => {}
}
}
engine_task.await.expect("mock engine task");
assert!(saw_text);
assert!(saw_done);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn reasoning_blocks_are_mapped_to_reasoning_sse_chunks() {
-1
View File
@@ -10,7 +10,6 @@ use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::lora::LoraRequest;
use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
+19
View File
@@ -0,0 +1,19 @@
[package]
name = "vllm-tool-parser-py"
version.workspace = true
edition.workspace = true
license.workspace = true
[lib]
name = "_rust_tool_parser"
crate-type = ["cdylib", "rlib"]
[dependencies]
pyo3.workspace = true
pythonize = { workspace = true, features = ["serde_json"] }
serde_json.workspace = true
thiserror-ext.workspace = true
vllm-tool-parser.workspace = true
[lints]
workspace = true
+366
View File
@@ -0,0 +1,366 @@
//! Thin PyO3 bindings for `vllm_tool_parser`.
//!
//! This crate exposes the Rust tool parser trait and data shapes to Python
//! while keeping parser state, grammar, and schema-aware argument conversion in
//! Rust. Python callers should use this module as a typed bridge and keep any
//! vLLM protocol adaptation outside the binding.
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyModule};
use pythonize::{depythonize, pythonize};
use serde_json::Value;
use thiserror_ext::AsReport as _;
use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
macro_rules! tool_parser_factory {
($($parser:ident),+ $(,)?) => {
fn create_tool_parser(
name: &str,
tools: &[Tool],
) -> PyResult<Box<dyn ToolParser>> {
match name {
$(
stringify!($parser) => {
<vllm_tool_parser::$parser as ToolParser>::create(tools)
}
)+
_ => {
return Err(PyValueError::new_err(format!(
"unsupported tool parser `{name}`"
)));
}
}
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
};
}
// Export a tool parser to Python by registering it here.
tool_parser_factory! {
DeepSeekV4ToolParser, // for testing on Python side
}
#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
#[derive(Clone)]
struct PyTool(Tool);
#[pymethods]
impl PyTool {
#[new]
#[pyo3(signature = (name, description, parameters, strict=None))]
fn new(
name: String,
description: Option<String>,
parameters: &Bound<'_, PyAny>,
strict: Option<bool>,
) -> PyResult<Self> {
let parameters = depythonize::<Value>(parameters).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert tool parameters from Python to JSON: {error}"
))
})?;
Ok(Self(Tool {
name,
description,
parameters,
strict,
}))
}
#[getter]
fn name(&self) -> &str {
&self.0.name
}
#[getter]
fn description(&self) -> Option<&str> {
self.0.description.as_deref()
}
#[getter]
fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert tool parameters from JSON to Python: {error}"
))
})
}
#[getter]
fn strict(&self) -> Option<bool> {
self.0.strict
}
}
#[pyclass(
name = "ToolCallDelta",
module = "vllm._rust_tool_parser",
skip_from_py_object
)]
#[derive(Clone)]
struct PyToolCallDelta(ToolCallDelta);
#[pymethods]
impl PyToolCallDelta {
#[new]
#[pyo3(signature = (tool_index, name, arguments))]
fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
Self(ToolCallDelta {
tool_index,
name,
arguments,
})
}
#[getter]
fn tool_index(&self) -> usize {
self.0.tool_index
}
#[getter]
fn name(&self) -> Option<&str> {
self.0.name.as_deref()
}
#[getter]
fn arguments(&self) -> &str {
&self.0.arguments
}
}
#[pyclass(
name = "ToolParserOutput",
module = "vllm._rust_tool_parser",
skip_from_py_object
)]
#[derive(Clone)]
struct PyToolParserOutput(ToolParserOutput);
#[pymethods]
impl PyToolParserOutput {
#[new]
#[pyo3(signature = (normal_text="", calls=None))]
fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
let calls =
calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
Self(ToolParserOutput {
normal_text: normal_text.to_owned(),
calls,
})
}
#[getter]
fn normal_text(&self) -> &str {
&self.0.normal_text
}
#[getter]
fn calls(&self) -> Vec<PyToolCallDelta> {
self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
}
fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
self.0.append(other.0.clone());
}
fn coalesce_calls(&self) -> Self {
Self(self.0.clone().coalesce_calls())
}
}
#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
struct PyToolParser(Box<dyn ToolParser>);
impl PyToolParser {
fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
self.0
.parse_into(chunk, &mut output.0)
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
}
#[pymethods]
impl PyToolParser {
#[new]
fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
create_tool_parser(parser_name, &tools).map(Self)
}
fn parse_into(
&mut self,
chunk: &str,
mut output: PyRefMut<'_, PyToolParserOutput>,
) -> PyResult<()> {
self.parse_into_output(chunk, &mut output)
}
fn finish(&mut self) -> PyResult<PyToolParserOutput> {
self.0
.finish()
.map(PyToolParserOutput)
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
fn reset(&mut self) -> String {
self.0.reset()
}
fn preserve_special_tokens(&self) -> bool {
self.0.preserve_special_tokens()
}
}
#[pymodule]
fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyTool>()?;
m.add_class::<PyToolCallDelta>()?;
m.add_class::<PyToolParserOutput>()?;
m.add_class::<PyToolParser>()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
Python::initialize();
Python::attach(f)
}
fn tool_schema() -> Value {
json!({
"type": "object",
"properties": {
"user_id": {"type": "integer"},
"shipping": {
"type": "object",
"properties": {
"city": {"type": "string"},
"zip": {"type": "integer"}
}
}
}
})
}
fn build_call() -> String {
r#"<DSMLtool_calls>
<DSMLinvoke name="create_order">
<DSMLparameter name="user_id" string="false">42</DSMLparameter>
<DSMLparameter name="shipping" string="false">{"city":"Singapore","zip":18956}</DSMLparameter>
</DSMLinvoke>
</DSMLtool_calls>"#
.to_owned()
}
fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
let parameters = pythonize(py, &tool_schema()).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert test schema from JSON to Python: {error}"
))
})?;
Py::new(
py,
PyTool::new(
"create_order".to_owned(),
Some("Create an order".to_owned()),
&parameters,
None,
)?,
)
}
#[test]
fn tool_round_trips_typed_fields() {
with_python(|py| {
let tool = make_py_tool(py)?;
let borrowed = tool.borrow(py);
assert_eq!(borrowed.name(), "create_order");
assert_eq!(borrowed.description(), Some("Create an order"));
assert_eq!(borrowed.strict(), None);
let parameters = borrowed.parameters(py)?;
let parameters = depythonize::<Value>(parameters.bind(py))?;
assert_eq!(parameters, tool_schema());
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn output_append_and_coalesce_calls() {
with_python(|py| {
let first = Py::new(
py,
PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
)?;
let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
output.append(other.borrow(py));
let coalesced = output.coalesce_calls();
assert_eq!(coalesced.normal_text(), "text");
let calls = coalesced.calls();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].tool_index(), 0);
assert_eq!(calls[0].name(), Some("create_order"));
assert_eq!(calls[0].arguments(), "{\"a\":1}");
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn parser_parse_finish_and_preserve_special_tokens() {
with_python(|py| {
let tool = make_py_tool(py)?;
let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
assert!(parser.preserve_special_tokens());
let mut output = PyToolParserOutput::new(py, "", None);
parser.parse_into_output(&build_call(), &mut output)?;
let finish = Py::new(py, parser.finish()?)?;
output.append(finish.borrow(py));
let output = output.coalesce_calls();
assert_eq!(output.normal_text(), "");
let calls = output.calls();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name(), Some("create_order"));
assert_eq!(
serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
json!({
"user_id": 42,
"shipping": {
"city": "Singapore",
"zip": 18956
}
})
);
assert_eq!(parser.reset(), "");
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn parser_errors_for_unknown_name() {
with_python(|py| {
let tool = make_py_tool(py)?;
let error = match PyToolParser::new(py, "missing", vec![tool]) {
Ok(_) => panic!("missing parser name unexpectedly succeeded"),
Err(error) => error,
};
let message = format!("{error}");
assert!(message.contains("unsupported tool parser `missing`"));
PyResult::Ok(())
})
.unwrap();
}
}
+26 -52
View File
@@ -18,8 +18,6 @@ import torch
from packaging.version import Version, parse
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools_rust import Binding, RustExtension
from setuptools_rust.build import build_rust
from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
@@ -35,11 +33,12 @@ def load_module_from_path(module_name, path):
ROOT_DIR = Path(__file__).parent
logger = logging.getLogger(__name__)
PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"
# cannot import envs directly because it depends on vllm,
# which is not installed yet
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
rust_build = load_module_from_path(
"rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py")
)
VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
USE_PRECOMPILED_EXTENSIONS = envs.VLLM_USE_PRECOMPILED
@@ -49,11 +48,6 @@ USE_PRECOMPILED_RUST_FRONTEND = (
)
def should_require_rust_frontend() -> bool:
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
return value.lower() not in ("", "0", "false", "no")
if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
VLLM_TARGET_DEVICE = "cpu"
@@ -420,24 +414,6 @@ class precompiled_build_ext(build_ext):
return
class precompiled_build_rust(build_rust):
"""Skips local Rust builds when the precompiled wheel already ships vllm-rs."""
def run(self) -> None:
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
logger.info(
"Skipping local Rust build: using precompiled %s",
PRECOMPILED_RUST_FRONTEND_PATH,
)
return
logger.warning(
"Precompiled wheel did not provide %s; falling back to local Rust build.",
PRECOMPILED_RUST_FRONTEND_PATH,
)
super().run()
class precompiled_wheel_utils:
"""Extracts libraries and other files from an existing wheel."""
@@ -731,9 +707,6 @@ class precompiled_wheel_utils:
"vllm/_rocm_C.abi3.so",
}
)
if extract_rust_frontend:
exact_members.add("vllm/vllm-rs")
flash_attn_regex = re.compile(
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
@@ -756,6 +729,12 @@ class precompiled_wheel_utils:
if member.filename in exact_members:
file_members.append(member)
continue
if (
extract_rust_frontend
and rust_build.is_precompiled_artifact_member(member.filename)
):
file_members.append(member)
continue
if not extract_extensions:
continue
@@ -1109,6 +1088,12 @@ package_data = {
}
def add_vllm_package_data(filename: str) -> None:
vllm_files = package_data.setdefault("vllm", [])
if filename not in vllm_files:
vllm_files.append(filename)
# If using precompiled artifacts, extract and patch package_data in advance.
if USE_PRECOMPILED_RUST_FRONTEND:
wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url()
@@ -1121,12 +1106,10 @@ if USE_PRECOMPILED_RUST_FRONTEND:
for pkg, files in patch.items():
package_data.setdefault(pkg, []).extend(files)
# If the rust frontend binary is already present in the source tree (e.g.,
# pre-built in a separate Docker build stage), ship it as-is.
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
vllm_files = package_data.setdefault("vllm", [])
if "vllm-rs" not in vllm_files:
vllm_files.append("vllm-rs")
# Rust artifacts already present in the source tree (e.g., pre-built in a
# separate Docker build stage) are shipped as-is.
for rust_artifact in rust_build.find_precompiled_artifacts():
add_vllm_package_data(rust_artifact.name)
if _no_device():
ext_modules = []
@@ -1139,23 +1122,14 @@ else:
if USE_PRECOMPILED_EXTENSIONS
else cmake_build_ext,
}
if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists():
cmdclass["build_rust"] = precompiled_build_rust
if USE_PRECOMPILED_RUST_FRONTEND or rust_build.find_precompiled_artifacts():
cmdclass["build_rust"] = rust_build.precompiled_build_rust
# Rust frontend binary, built via setuptools-rust and installed into the
# package directory alongside the Python modules.
# TODO: we may use `RustBin` to directly install it into `bin` directory, but this
# requires extra work on using precompiled binaries.
rust_extensions = [
RustExtension(
target="vllm.vllm-rs",
path="rust/src/cmd/Cargo.toml",
args=["--bin", "vllm-rs"],
features=["native-tls-vendored"],
binding=Binding.Exec,
optional=not should_require_rust_frontend(),
),
]
# Rust artifacts, built via setuptools-rust and installed into the package
# directory alongside the Python modules.
rust_extensions = rust_build.rust_extensions(
optional=not rust_build.should_require_rust_frontend()
)
setup(
# static metadata should rather go in pyproject.toml
+341
View File
@@ -0,0 +1,341 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from transformers import AutoTokenizer, PreTrainedTokenizerBase
from vllm.benchmarks.datasets import BFCLDataset, get_samples
def _patch_hf_api(side_effect):
"""Return a patch context that swaps `hf_api()` to a stub whose
`.hf_hub_download` attribute uses `side_effect`."""
fake_api = MagicMock()
fake_api.hf_hub_download.side_effect = side_effect
return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api)
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
return AutoTokenizer.from_pretrained("gpt2")
_FAKE_ROWS = {
"simple": [
{
"id": "simple_0",
"question": [
[
{
"role": "user",
"content": "What is 2+2?",
}
]
],
"function": [
{
"name": "add",
"description": "Add two numbers.",
"parameters": {
"type": "dict",
"properties": {
"a": {"type": "integer", "description": "first"},
"b": {"type": "float", "description": "second"},
},
"required": ["a", "b"],
},
}
],
},
],
"live_simple": [
{
"id": "live_simple_0",
"question": [[{"role": "user", "content": "Tell me the weather."}]],
"function": [
{
"name": "get_weather",
"description": "Get weather.",
"parameters": {
"type": "dict",
"properties": {
"city": {"type": "any", "description": "city"},
"coords": {"type": "tuple", "description": "coords"},
},
"required": ["city"],
},
}
],
},
],
}
def _write_fake_files(tmp_path: Path) -> dict[str, Path]:
"""Write fake BFCL JSONL files mimicking the HF repo layout."""
paths = {}
for category, rows in _FAKE_ROWS.items():
p = tmp_path / f"BFCL_v3_{category}.json"
with p.open("w") as f:
for row in rows:
f.write(json.dumps(row) + "\n")
paths[category] = p
return paths
def _args_for_bfcl(categories: list[str] | None) -> argparse.Namespace:
return argparse.Namespace(
dataset_name="hf",
dataset_path="gorilla-llm/Berkeley-Function-Calling-Leaderboard",
hf_name=None,
hf_subset=None,
hf_split=None,
hf_output_len=64,
disable_shuffle=True,
num_prompts=2,
no_oversample=False,
no_stream=True,
seed=0,
request_id_prefix="",
trust_remote_code=False,
skip_chat_template=False,
enable_multimodal_chat=False,
backend="openai-chat",
bfcl_categories=categories,
)
@pytest.mark.benchmark
def test_bfcl_dataset_translates_schema_and_attaches_tools(
hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
) -> None:
"""BFCLDataset should translate schemas to OpenAI tool format, set
`messages` directly on SampleRequest, and attach tools/tool_choice via
request_overrides."""
paths = _write_fake_files(tmp_path)
def fake_download(_repo, filename, **_kwargs):
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
return str(paths[category])
args = _args_for_bfcl(categories=["simple", "live_simple"])
with _patch_hf_api(fake_download):
samples = get_samples(args, hf_tokenizer)
assert len(samples) == 2
for s in samples:
assert s.chat_messages is not None
assert isinstance(s.chat_messages, list)
assert s.chat_messages[0]["role"] == "user"
assert s.request_overrides is not None
assert "tools" in s.request_overrides
assert s.request_overrides["tool_choice"] == "auto"
# messages must NOT leak into request_overrides — it has its own
# typed field on SampleRequest.
assert "messages" not in s.request_overrides
tools = s.request_overrides["tools"]
assert len(tools) == 1
tool = tools[0]
assert tool["type"] == "function"
# Translated schema: dict -> object, float -> number,
# any -> string, tuple -> array.
params = tool["function"]["parameters"]
assert params["type"] == "object"
for prop in params["properties"].values():
assert prop["type"] in {"integer", "number", "string", "array"}
@pytest.mark.benchmark
def test_bfcl_dataset_requires_openai_chat_backend(
hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
args = _args_for_bfcl(categories=["simple"])
args.backend = "openai"
with pytest.raises(ValueError, match="openai-chat"):
get_samples(args, hf_tokenizer)
@pytest.mark.benchmark
def test_bfcl_dataset_missing_category_raises_clear_error(
hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
"""A typo'd category should produce an actionable ValueError, not an
opaque huggingface_hub exception."""
from huggingface_hub.errors import EntryNotFoundError
args = _args_for_bfcl(categories=["simpl"]) # typo
def raise_missing(_repo, filename, **_kwargs):
raise EntryNotFoundError(f"404 Not Found: {filename}")
with (
_patch_hf_api(raise_missing),
pytest.raises(ValueError, match=r"BFCL category 'simpl' not found"),
):
get_samples(args, hf_tokenizer)
@pytest.mark.benchmark
def test_chat_backend_uses_messages_field_when_set() -> None:
"""When RequestFuncInput.chat_messages is set, the chat backend must use
it verbatim and skip default content construction from `prompt`."""
import asyncio
from vllm.benchmarks.lib.endpoint_request_func import (
RequestFuncInput,
async_request_openai_chat_completions,
)
captured: dict = {}
class _FakeResp:
status = 500
reason = "stop-after-capture"
content = None
async def __aenter__(self):
return self
async def __aexit__(self, *exc):
return False
class _FakeSession:
def post(self, url, json, headers): # noqa: A002
captured["url"] = url
captured["payload"] = json
return _FakeResp()
messages = [
{"role": "system", "content": "you are helpful"},
{"role": "user", "content": "call add(3, 4)"},
]
req = RequestFuncInput(
prompt="IGNORED",
api_url="http://localhost:0/v1/chat/completions",
prompt_len=10,
output_len=16,
model="test-model",
chat_messages=messages,
extra_body={"tools": [{"type": "function", "function": {"name": "add"}}]},
)
asyncio.run(
async_request_openai_chat_completions(
request_func_input=req, session=_FakeSession()
)
)
payload = captured["payload"]
assert payload["messages"] is messages, (
"chat backend must forward RequestFuncInput.chat_messages verbatim "
"instead of constructing a default user message from `prompt`"
)
# extra_body still merges in as before (shallow, per-request wins).
assert payload["tools"][0]["function"]["name"] == "add"
@pytest.mark.benchmark
def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None:
"""prompt_len must reflect tokens from both messages *and* tool schemas,
so percentile buckets and input-distribution summaries aren't biased
low for tool-heavy traffic."""
paths = _write_fake_files(tmp_path)
def fake_download(_repo, filename, **_kwargs):
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
return str(paths[category])
captured: dict = {}
class _FakeTokenizer:
def apply_chat_template(
self, messages, tools=None, tokenize=False, add_generation_prompt=True
):
captured["tools"] = tools
base = " ".join(m.get("content", "") for m in messages)
tool_text = json.dumps(tools) if tools else ""
return base + " " + tool_text
def __call__(self, text):
# 1 "token" per whitespace-separated word.
return type("Enc", (), {"input_ids": text.split()})()
fake = _FakeTokenizer()
args = _args_for_bfcl(categories=["simple"])
args.num_prompts = 1
with _patch_hf_api(fake_download):
samples = get_samples(args, fake)
assert len(samples) == 1
assert captured["tools"] is not None, (
"apply_chat_template must be called with tools= so the schema "
"contributes to the prompt-length estimate"
)
assert len(captured["tools"]) == 1
assert captured["tools"][0]["function"]["name"] == "add"
# Sanity: prompt_len exceeds a messages-only estimate. The fake row's
# user message is "What is 2+2?" (3 whitespace-separated tokens).
assert samples[0].prompt_len > 3
@pytest.mark.benchmark
def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools(
tmp_path: Path,
) -> None:
"""Older tokenizers don't accept tools=; fallback must still produce a
non-zero prompt_len without crashing."""
paths = _write_fake_files(tmp_path)
def fake_download(_repo, filename, **_kwargs):
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
return str(paths[category])
class _LegacyTokenizer:
def apply_chat_template(self, messages, **kwargs):
if "tools" in kwargs:
raise TypeError("unexpected keyword argument 'tools'")
return " ".join(m.get("content", "") for m in messages)
def __call__(self, text):
return type("Enc", (), {"input_ids": text.split()})()
args = _args_for_bfcl(categories=["simple"])
args.num_prompts = 1
with _patch_hf_api(fake_download):
samples = get_samples(args, _LegacyTokenizer())
assert len(samples) == 1
assert samples[0].prompt_len > 0
@pytest.mark.benchmark
def test_bfcl_schema_translation_is_recursive() -> None:
"""_translate_schema must recurse into nested properties."""
input_schema = {
"type": "dict",
"properties": {
"nested": {
"type": "dict",
"properties": {
"value": {"type": "float"},
"tags": {"type": "tuple", "items": {"type": "any"}},
},
}
},
}
out = BFCLDataset._translate_schema(input_schema)
assert out["type"] == "object"
assert out["properties"]["nested"]["type"] == "object"
assert out["properties"]["nested"]["properties"]["value"]["type"] == "number"
assert out["properties"]["nested"]["properties"]["tags"]["type"] == "array"
nested_props = out["properties"]["nested"]["properties"]
assert nested_props["tags"]["items"]["type"] == "string"
@@ -382,11 +382,12 @@ def test_rope_kvcache_fusion(
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
# Cannot compare fp8_* directly here, cast to model dtype instead
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
torch.testing.assert_close(
kv_cache_unfused.view(dtype),
kv_cache_fused.view(dtype),
atol=ATOL,
rtol=RTOL,
kv_cache_unfused.to(dtype),
kv_cache_fused.to(dtype),
atol=1e-1,
rtol=1e-1,
)
@@ -569,17 +570,19 @@ def test_rope_static_qquant_kvcache_fusion(
else:
ATOL, RTOL = (1e-2, 1e-2)
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
torch.testing.assert_close(
q_unfused.to(torch.float32),
q_fused.to(torch.float32),
atol=ATOL,
rtol=RTOL,
atol=1e-1,
rtol=1e-1,
)
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
torch.testing.assert_close(
kv_cache_unfused.view(dtype),
kv_cache_fused.view(dtype),
atol=ATOL,
rtol=RTOL,
kv_cache_unfused.to(dtype),
kv_cache_fused.to(dtype),
atol=1e-1,
rtol=1e-1,
)
@@ -2,11 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from openai.types.responses import FunctionTool
from openai_harmony import DeveloperContent, Message, Role
from tests.entrypoints.openai.utils import verify_harmony_messages
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionToolsParam
from vllm.entrypoints.openai.parser.harmony_utils import (
auto_drop_analysis_messages,
create_tool_definition,
extract_function_from_recipient,
get_encoding,
get_system_message,
@@ -20,6 +23,58 @@ from vllm.entrypoints.openai.responses.harmony import (
response_previous_input_to_harmony,
)
_TOOL_PARAMETERS = {
"type": "object",
"properties": {"status": {"type": "string"}},
"required": ["status"],
"additionalProperties": False,
}
class TestCreateToolDefinition:
def test_chat_completion_omitted_description_defaults_to_empty_string(self):
tool = ChatCompletionToolsParam(
function={
"name": "report_status",
"parameters": _TOOL_PARAMETERS,
}
)
tool_definition = create_tool_definition(tool)
assert tool_definition.name == "report_status"
assert tool_definition.description == ""
assert tool_definition.parameters == _TOOL_PARAMETERS
def test_chat_completion_none_description_defaults_to_empty_string(self):
tool = ChatCompletionToolsParam(
function={
"name": "report_status",
"description": None,
"parameters": _TOOL_PARAMETERS,
}
)
tool_definition = create_tool_definition(tool)
assert tool_definition.name == "report_status"
assert tool_definition.description == ""
assert tool_definition.parameters == _TOOL_PARAMETERS
def test_response_tool_none_description_defaults_to_empty_string(self):
tool = FunctionTool(
name="report_status",
description=None,
parameters=_TOOL_PARAMETERS,
type="function",
)
tool_definition = create_tool_definition(tool)
assert tool_definition.name == "report_status"
assert tool_definition.description == ""
assert tool_definition.parameters == _TOOL_PARAMETERS
class TestIsFunctionRecipient:
@pytest.mark.parametrize(
@@ -7,3 +7,4 @@ server_args: >-
--max-model-len 4096
--data-parallel-size 2
--enable-expert-parallel
--no-enable-flashinfer-autotune
@@ -166,14 +166,14 @@ def test_op_absent_on_non_gfx1100(op_name):
@not_gfx1100
def test_rocm_moe_not_supported_on_non_gfx1100():
"""rocm_moe.is_supported() must return False on non-gfx1100 hardware."""
"""rocm_moe_rdna.is_supported() must return False on non-gfx1100 hardware."""
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
rocm_moe,
rocm_moe_rdna,
)
wq = type("WQ", (), {"num_bits": 4})()
assert rocm_moe.is_supported(wq) is False, (
"rocm_moe.is_supported() returned True on non-gfx1100 — "
assert rocm_moe_rdna.is_supported(wq) is False, (
"rocm_moe_rdna.is_supported() returned True on non-gfx1100 — "
"dispatch guard is broken"
)
@@ -371,43 +371,43 @@ class TestMoEDispatchMocked:
"""Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
def test_is_supported_false_when_mocked_cdna(self):
"""rocm_moe.is_supported() must return False when not on gfx1100."""
"""rocm_moe_rdna.is_supported() must return False when not on gfx1100."""
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
rocm_moe,
rocm_moe_rdna,
)
with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
@pytest.mark.parametrize("num_bits", [2, 3, 8, 16])
def test_is_supported_rejects_non_w4(self, num_bits):
"""is_supported() rejects non-4-bit even before checking arch."""
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
rocm_moe,
rocm_moe_rdna,
)
assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
def test_is_supported_false_when_op_missing(self):
"""is_supported() returns False when the C++ op doesn't exist."""
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
rocm_moe,
rocm_moe_rdna,
)
fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
with patch.object(torch, "ops", create=True) as mock_ops:
mock_ops._rocm_C = fake_rocm_c
assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
def test_is_supported_false_when_rocm_c_absent(self):
"""is_supported() returns False when _rocm_C doesn't exist at all."""
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
rocm_moe,
rocm_moe_rdna,
)
fake_ops = type("FakeOps", (), {})()
with patch.object(torch, "ops", fake_ops):
assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
class TestDenseKernelSelectionMocked:
@@ -475,10 +475,10 @@ class TestDenseKernelSelectionMocked:
class TestCompressedTensorsMoEDispatchGuard:
"""Verify compressed_tensors_moe.py only enters rocm_moe under is_rocm()."""
"""Verify compressed_tensors_moe.py only enters rocm_moe_rdna under is_rocm()."""
def test_rocm_guard_in_dispatch_source(self):
"""The rocm_moe import and call must be inside an is_rocm() check."""
"""The rocm_moe_rdna import and call must be inside an is_rocm() check."""
src = _read_pkg_source_or_skip(
"model_executor",
"layers",
@@ -498,6 +498,6 @@ class TestCompressedTensorsMoEDispatchGuard:
found_guard = True
break
assert found_guard, (
f"L{i}: rocm_moe reference not protected by "
f"L{i}: rocm_moe_rdna reference not protected by "
f"is_rocm() guard: {stripped}"
)
+17
View File
@@ -68,6 +68,23 @@ def test_audio_media_io_encode_base64(dummy_audio):
mock_write.assert_called_once()
def test_load_audio_max_duration_respected(dummy_audio_bytes):
"""Valid audio within the duration limit should load successfully."""
from io import BytesIO
y, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, max_duration_s=3600)
assert isinstance(y, np.ndarray)
assert len(y) > 0
def test_load_audio_max_duration_rejected(dummy_audio_bytes):
"""Audio exceeding the duration limit must be rejected during decode."""
from io import BytesIO
with pytest.raises(ValueError, match="exceeds maximum allowed duration"):
load_audio(BytesIO(dummy_audio_bytes), sr=None, max_duration_s=0.0001)
def test_audio_media_io_from_video(video_assets):
audio_io = AudioMediaIO()
video_path = video_assets[0].video_path
+86 -1
View File
@@ -1,12 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import struct
from io import BytesIO
from pathlib import Path
import numpy as np
import pytest
from PIL import Image, ImageChops
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.image import (
_has_transparency,
convert_image_mode,
normalize_image,
)
pytestmark = pytest.mark.cpu_test
@@ -38,3 +44,82 @@ def test_rgba_to_rgb():
assert converted_image_numpy[i][j][0] == 255
assert converted_image_numpy[i][j][1] == 255
assert converted_image_numpy[i][j][2] == 255
def test_palette_with_trns_to_rgb():
"""P-mode PNG with tRNS: transparent index should become white."""
# Built synthetically because the existing assets are RGBA PNGs;
# this vulnerability only affects P/L/RGB images carrying a tRNS
# chunk, which uses a different transparency mechanism.
img = Image.new("P", (4, 4))
palette = [0] * 768
palette[0:3] = [255, 0, 0]
palette[3:6] = [0, 0, 255]
img.putpalette(palette)
img.putpixel((0, 0), 0)
img.putpixel((1, 0), 1)
img.info["transparency"] = 1
assert _has_transparency(img)
converted = convert_image_mode(img, "RGB")
assert converted.mode == "RGB"
r, g, b = converted.getpixel((1, 0))
assert (r, g, b) == (255, 255, 255)
r, g, b = converted.getpixel((0, 0))
assert (r, g, b) == (255, 0, 0)
def test_l_mode_no_trns_to_rgb():
"""L-mode without transparency should convert directly."""
img = Image.new("L", (4, 4), 128)
assert not _has_transparency(img)
converted = convert_image_mode(img, "RGB")
assert converted.mode == "RGB"
assert converted.getpixel((0, 0)) == (128, 128, 128)
def test_exif_transpose_normalizes_orientation():
"""Image with EXIF orientation 3 (180-degree rotation) should be
normalized so pixel data matches visual display."""
# Built synthetically because the existing assets are PNGs which
# don't carry EXIF orientation metadata; we need a JPEG with an
# injected EXIF orientation=3 tag.
img = Image.new("RGB", (2, 1))
img.putpixel((0, 0), (255, 0, 0))
img.putpixel((1, 0), (0, 0, 255))
buf = BytesIO()
img.save(buf, format="JPEG")
jpeg_bytes = buf.getvalue()
exif_orientation_3 = (
b"\xff\xe1"
+ struct.pack(">H", 26)
+ b"Exif\x00\x00"
+ b"MM"
+ b"\x00\x2a"
+ b"\x00\x00\x00\x08"
+ b"\x00\x01"
+ b"\x01\x12"
+ b"\x00\x03"
+ b"\x00\x00\x00\x01"
+ b"\x00\x03"
+ b"\x00\x00"
+ b"\x00\x00\x00\x00"
)
soi = jpeg_bytes[:2]
rest = jpeg_bytes[2:]
patched = soi + exif_orientation_3 + rest
rotated = Image.open(BytesIO(patched))
normalized = normalize_image(rotated)
assert normalized.size == (2, 1)
def test_normalize_image_no_exif():
"""Images without EXIF should pass through unchanged."""
img = Image.new("RGB", (4, 4), (100, 100, 100))
result = normalize_image(img)
assert result.size == img.size
assert result.mode == img.mode
+51
View File
@@ -0,0 +1,51 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import numpy as np
import pytest
import torch
from PIL import Image
from vllm.multimodal.parse import ImageProcessorItems, VideoProcessorItems
H, W = 480, 640
@pytest.mark.parametrize(
"image",
[
Image.new("RGB", (W, H)),
# HWC, e.g. from np.array(PIL.Image)
np.zeros((H, W, 3), dtype=np.uint8),
torch.zeros((H, W, 3), dtype=torch.uint8),
# CHW, standard PyTorch / numpy convention
np.zeros((3, H, W), dtype=np.uint8),
torch.zeros((3, H, W), dtype=torch.uint8),
],
)
def test_image_size_hwc_chw(image):
"""Image sizes must be channel-layout agnostic.
`get_image_size` determines the multimodal placeholder count; reading an
HWC array (the layout `np.array(PIL.Image)` produces) as CHW yields a
bogus size and a placeholder/embedding count mismatch at inference time.
"""
items = ImageProcessorItems([image])
assert items.get_image_size(0) == (W, H)
@pytest.mark.parametrize(
"frame",
[
Image.new("RGB", (W, H)),
np.zeros((H, W, 3), dtype=np.uint8),
torch.zeros((H, W, 3), dtype=torch.uint8),
np.zeros((3, H, W), dtype=np.uint8),
torch.zeros((3, H, W), dtype=torch.uint8),
],
)
def test_frame_size_hwc_chw(frame):
"""`get_frame_size` must stay consistent with `get_image_size`."""
items = VideoProcessorItems([[frame]])
assert items.get_frame_size(0) == (W, H)
@@ -3,7 +3,6 @@
import json
from collections.abc import Generator
from typing import Any
from unittest.mock import MagicMock, patch
import partial_json_parser
@@ -29,22 +28,17 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
StructuralTagResponseFormat,
)
from vllm.entrypoints.openai.engine.protocol import FunctionCall as VllmFunctionCall
from vllm.reasoning.mistral_reasoning_parser import MistralReasoningParser
from vllm.sampling_params import StructuredOutputsParams
from vllm.tokenizers import TokenizerLike, get_tokenizer
from vllm.tokenizers.detokenizer_utils import detokenize_incrementally
from vllm.tokenizers.mistral import MistralTokenizer
from vllm.tool_parsers.mistral_tool_parser import (
_DEFAULT_JSON_SCHEMA,
MistralStreamingResult,
MistralToolCall,
MistralToolParser,
)
@@ -1578,382 +1572,3 @@ def test_grammar_from_tool_parser_set_by_adjust_request(
request = _make_request()
result = mistral_tool_parser.adjust_request(request)
assert result._grammar_from_tool_parser is True
@pytest.mark.parametrize(
"tool_calls, expected_len",
[
(None, 0),
([], 0),
([VllmFunctionCall(id="abc123xyz", name="f", arguments="{}")], 1),
([VllmFunctionCall(name="f", arguments="{}")], 1),
(
[
VllmFunctionCall(id="fixed1234", name="a", arguments='{"x": 1}'),
VllmFunctionCall(name="b", arguments='{"y": 2}'),
],
2,
),
],
ids=["none", "empty", "with_id", "without_id", "mixed"],
)
def test_build_non_streaming_tool_calls(
tool_calls: list[VllmFunctionCall] | None,
expected_len: int,
) -> None:
result = MistralToolParser.build_non_streaming_tool_calls(tool_calls)
assert len(result) == expected_len
if tool_calls is None:
return
for i, tc in enumerate(result):
assert isinstance(tc, MistralToolCall)
assert tc.type == "function"
input_tc = tool_calls[i]
if input_tc.id:
assert tc.id == input_tc.id
else:
assert len(tc.id) == 9
assert tc.id.isalnum()
assert tc.function.name == input_tc.name
assert tc.function.arguments == input_tc.arguments
class TestExtractMaybeReasoningAndToolStreaming:
r"""Tests for `MistralToolParser.extract_maybe_reasoning_and_tool_streaming`."""
@pytest.fixture
def parser(self) -> MistralToolParser:
mock_tokenizer = MagicMock()
mock_tokenizer.get_vocab.return_value = {"[TOOL_CALLS]": 1}
return MistralToolParser(mock_tokenizer)
@pytest.fixture
def request_obj(self) -> ChatCompletionRequest:
return _make_request()
@staticmethod
def _call(
parser: MistralToolParser,
request: ChatCompletionRequest,
*,
reasoning_parser: Any = None,
previous_text: str = "",
current_text: str = "hello",
delta_text: str = "hello",
previous_token_ids: list[int] | None = None,
current_token_ids: list[int] | None = None,
output_token_ids: list[int] | None = None,
reasoning_ended: bool = False,
prompt_is_reasoning_end: bool | None = None,
) -> MistralStreamingResult:
return parser.extract_maybe_reasoning_and_tool_streaming(
reasoning_parser=reasoning_parser,
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=previous_token_ids or [],
current_token_ids=current_token_ids or [1, 2, 3],
output_token_ids=output_token_ids or [1, 2, 3],
reasoning_ended=reasoning_ended,
prompt_is_reasoning_end=prompt_is_reasoning_end,
request=request,
)
def test_no_reasoning_tools_called(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
tool_delta = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=0,
function=DeltaFunctionCall(name="f", arguments="{}"),
)
]
)
with patch.object(
parser, "extract_tool_calls_streaming", return_value=tool_delta
):
result = self._call(parser, request_obj, reasoning_parser=None)
assert result == MistralStreamingResult(
delta_message=tool_delta,
reasoning_ended=False,
tools_called=True,
current_text="hello",
current_token_ids=[1, 2, 3],
)
def test_no_reasoning_no_tools(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
content_delta = DeltaMessage(content="hello")
with patch.object(
parser, "extract_tool_calls_streaming", return_value=content_delta
):
result = self._call(parser, request_obj, reasoning_parser=None)
assert result == MistralStreamingResult(
delta_message=content_delta,
reasoning_ended=False,
tools_called=False,
current_text="hello",
current_token_ids=[1, 2, 3],
)
def test_mistral_reasoning_parser_no_think_token(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
mock_rp = MagicMock(spec=MistralReasoningParser)
mock_rp.start_token_id = 999
content_delta = DeltaMessage(content="direct")
with patch.object(
parser, "extract_tool_calls_streaming", return_value=content_delta
):
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
current_token_ids=[1, 2, 3],
)
mock_rp.extract_reasoning_streaming.assert_not_called()
assert result == MistralStreamingResult(
delta_message=content_delta,
reasoning_ended=False,
tools_called=False,
current_text="hello",
current_token_ids=[1, 2, 3],
)
def test_mistral_reasoning_parser_with_think_token(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
mock_rp = MagicMock(spec=MistralReasoningParser)
mock_rp.start_token_id = 999
mock_rp.extract_reasoning_streaming.return_value = DeltaMessage(
reasoning="thinking..."
)
mock_rp.is_reasoning_end_streaming.return_value = False
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
current_token_ids=[1, 999, 3],
)
mock_rp.extract_reasoning_streaming.assert_called_once()
assert result == MistralStreamingResult(
delta_message=DeltaMessage(reasoning="thinking..."),
reasoning_ended=False,
tools_called=False,
current_text="hello",
current_token_ids=[1, 999, 3],
)
def test_non_mistral_reasoning_parser_always_expects_thinking(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
mock_rp = MagicMock()
mock_rp.start_token_id = 999
mock_rp.extract_reasoning_streaming.return_value = DeltaMessage(
reasoning="thinking..."
)
mock_rp.is_reasoning_end_streaming.return_value = False
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
current_token_ids=[1, 2, 3],
)
mock_rp.extract_reasoning_streaming.assert_called_once()
assert result == MistralStreamingResult(
delta_message=DeltaMessage(reasoning="thinking..."),
reasoning_ended=False,
tools_called=False,
current_text="hello",
current_token_ids=[1, 2, 3],
)
def test_reasoning_already_ended_no_reset(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
content_delta = DeltaMessage(content="content")
with patch.object(
parser, "extract_tool_calls_streaming", return_value=content_delta
) as mock_extract:
result = self._call(
parser,
request_obj,
reasoning_parser=MagicMock(),
reasoning_ended=True,
previous_text="prior_tool_text",
previous_token_ids=[10, 20],
current_text="prior_tool_texthello",
current_token_ids=[10, 20, 1, 2, 3],
)
_, call_kwargs = mock_extract.call_args
assert call_kwargs["previous_text"] == "prior_tool_text"
assert call_kwargs["previous_token_ids"] == [10, 20]
assert result == MistralStreamingResult(
delta_message=content_delta,
reasoning_ended=True,
tools_called=False,
current_text="prior_tool_texthello",
current_token_ids=[10, 20, 1, 2, 3],
)
def test_pre_v15_ignores_prompt_reasoning_end(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
mock_tokenizer = MagicMock(spec=MistralTokenizer)
mock_tokenizer.version = 13
parser.model_tokenizer = mock_tokenizer
mock_rp = MagicMock(spec=MistralReasoningParser)
mock_rp.start_token_id = 999
mock_rp.extract_reasoning_streaming.return_value = DeltaMessage(
reasoning="thinking..."
)
mock_rp.is_reasoning_end_streaming.return_value = False
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
prompt_is_reasoning_end=True,
current_token_ids=[999, 1, 2],
)
mock_rp.extract_reasoning_streaming.assert_called_once()
assert result == MistralStreamingResult(
delta_message=DeltaMessage(reasoning="thinking..."),
reasoning_ended=False,
tools_called=False,
current_text="hello",
current_token_ids=[999, 1, 2],
)
def test_non_pre_v15_prompt_reasoning_end(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
mock_tokenizer = MagicMock(spec=MistralTokenizer)
mock_tokenizer.version = 15
parser.model_tokenizer = mock_tokenizer
mock_rp = MagicMock(spec=MistralReasoningParser)
mock_rp.start_token_id = 999
content_delta = DeltaMessage(content="after reasoning")
with patch.object(
parser, "extract_tool_calls_streaming", return_value=content_delta
):
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
prompt_is_reasoning_end=True,
current_token_ids=[999, 1, 2],
output_token_ids=[10, 20, 30],
)
mock_rp.extract_reasoning_streaming.assert_not_called()
assert result == MistralStreamingResult(
delta_message=content_delta,
reasoning_ended=True,
tools_called=False,
current_text="hello",
current_token_ids=[10, 20, 30],
)
def test_reasoning_end_transition_with_content(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
"""When reasoning ends and the delta has content, that content is
cleared from delta_message and used as current_text for tool parsing."""
mock_rp = MagicMock()
mock_rp.start_token_id = 999
mock_rp.extract_reasoning_streaming.return_value = DeltaMessage(
reasoning="think", content="leftover"
)
mock_rp.is_reasoning_end_streaming.return_value = True
mock_rp.extract_content_ids.return_value = [50, 51]
content_delta = DeltaMessage(content="leftover")
with patch.object(
parser, "extract_tool_calls_streaming", return_value=content_delta
) as mock_extract:
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
current_token_ids=[999, 1, 2],
output_token_ids=[10, 20, 30],
)
mock_rp.extract_content_ids.assert_called_once_with([10, 20, 30])
_, call_kwargs = mock_extract.call_args
assert call_kwargs["previous_text"] == ""
assert call_kwargs["previous_token_ids"] == []
assert call_kwargs["delta_text"] == "leftover"
assert call_kwargs["current_token_ids"] == [50, 51]
assert result == MistralStreamingResult(
delta_message=content_delta,
reasoning_ended=True,
tools_called=False,
current_text="leftover",
current_token_ids=[50, 51],
)
def test_reasoning_end_transition_without_content(
self, parser: MistralToolParser, request_obj: ChatCompletionRequest
) -> None:
"""When reasoning ends but the delta has no content, current_text
is set to empty string."""
mock_rp = MagicMock()
mock_rp.start_token_id = 999
mock_rp.extract_reasoning_streaming.return_value = DeltaMessage(
reasoning="think"
)
mock_rp.is_reasoning_end_streaming.return_value = True
mock_rp.extract_content_ids.return_value = [50, 51]
empty_delta = DeltaMessage(content="")
with patch.object(
parser, "extract_tool_calls_streaming", return_value=empty_delta
) as mock_extract:
result = self._call(
parser,
request_obj,
reasoning_parser=mock_rp,
reasoning_ended=False,
current_token_ids=[999, 1, 2],
output_token_ids=[10, 20, 30],
)
_, call_kwargs = mock_extract.call_args
assert call_kwargs["delta_text"] == ""
assert call_kwargs["current_token_ids"] == [50, 51]
assert result == MistralStreamingResult(
delta_message=empty_delta,
reasoning_ended=True,
tools_called=False,
current_text="",
current_token_ids=[50, 51],
)
+260
View File
@@ -0,0 +1,260 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from collections.abc import Sequence
from unittest.mock import MagicMock
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
ChatCompletionToolsParam,
)
from vllm.tool_parsers.rust_tool_parser import RustToolParser
# The PyO3 extension is an optional build artifact; skip when absent.
_rust_tool_parser = pytest.importorskip("vllm._rust_tool_parser")
MOCK_TOKENIZER = MagicMock()
MOCK_TOKENIZER.get_vocab.return_value = {}
TC_START = "<DSMLtool_calls>"
TC_END = "</DSMLtool_calls>"
INV_START = '<DSMLinvoke name="'
INV_END = "</DSMLinvoke>"
PARAM_START = '<DSMLparameter name="'
PARAM_END = "</DSMLparameter>"
class DeepSeekV4RustToolParser(RustToolParser):
rust_parser_name = "DeepSeekV4ToolParser"
tool_call_start_token = TC_START
def sample_tools() -> list[ChatCompletionToolsParam]:
return [
ChatCompletionToolsParam(
type="function",
function={
"name": "get_weather",
"description": "Get weather for a location",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"date": {"type": "string"},
},
},
},
),
ChatCompletionToolsParam(
type="function",
function={
"name": "add",
"description": "Add two integers",
"parameters": {
"type": "object",
"properties": {
"x": {"type": "integer"},
"y": {"type": "integer"},
},
},
},
),
]
EXPECTED_CALLS = [
("get_weather", {"location": "SF", "date": "2024-01-16"}),
("add", {"x": 3, "y": 5}),
]
def build_invoke(
function_name: str,
params: Sequence[tuple[str, str, bool]],
) -> str:
param_text = "\n".join(
f'{PARAM_START}{name}" string="{str(is_string).lower()}">{value}{PARAM_END}'
for name, value, is_string in params
)
return f'{INV_START}{function_name}">\n{param_text}\n{INV_END}\n'
def build_tool_call() -> str:
weather = build_invoke(
"get_weather",
[
("location", "SF", True),
("date", "2024-01-16", True),
],
)
add = build_invoke(
"add",
[
("x", "3", False),
("y", "5", False),
],
)
return f"{TC_START}\n{weather}{add}{TC_END}"
def parse_streaming(
parser: DeepSeekV4RustToolParser,
text: str,
chunk_size: int,
) -> list:
deltas = []
previous_text = ""
for start in range(0, len(text), chunk_size):
delta_text = text[start : start + chunk_size]
current_text = previous_text + delta_text
delta = parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[1],
request=MagicMock(),
)
previous_text = current_text
if delta is not None:
deltas.append(delta)
delta = parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=previous_text,
delta_text="",
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[2],
request=MagicMock(),
)
if delta is not None:
deltas.append(delta)
return deltas
def collect_streamed_arguments(deltas: Sequence, tool_index: int = 0) -> str:
return "".join(
tool_call.function.arguments
for delta in deltas
for tool_call in delta.tool_calls or []
if (
tool_call.index == tool_index
and tool_call.function is not None
and tool_call.function.arguments is not None
)
)
def test_rust_tool_parser_extension_typed_api() -> None:
tools = [
_rust_tool_parser.Tool(
tool.function.name,
tool.function.description,
tool.function.parameters,
None,
)
for tool in sample_tools()
]
parser = _rust_tool_parser.ToolParser("DeepSeekV4ToolParser", tools)
output = _rust_tool_parser.ToolParserOutput()
parser.parse_into(build_tool_call(), output)
output.append(parser.finish())
output = output.coalesce_calls()
assert parser.preserve_special_tokens()
assert output.normal_text == ""
assert len(output.calls) == 2
for call, (name, arguments) in zip(output.calls, EXPECTED_CALLS):
assert call.name == name
assert json.loads(call.arguments) == arguments
def test_rust_tool_parser_adapter_extracts_complete_output() -> None:
tools = sample_tools()
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools)
result = parser.extract_tool_calls(
"Let me create it. " + build_tool_call(),
ChatCompletionRequest(messages=[], model="m", tools=tools),
)
assert result.tools_called
assert result.content == "Let me create it. "
assert len(result.tool_calls) == 2
for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS):
assert tool_call.function.name == name
assert json.loads(tool_call.function.arguments) == arguments
def test_rust_tool_parser_adapter_streaming_handles_multiple_calls() -> None:
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())
deltas = parse_streaming(parser, build_tool_call(), chunk_size=5)
names = [
tool_call.function.name
for delta in deltas
for tool_call in delta.tool_calls or []
if tool_call.function is not None and tool_call.function.name is not None
]
assert names == [name for name, _ in EXPECTED_CALLS]
for index, (_, arguments) in enumerate(EXPECTED_CALLS):
assert json.loads(collect_streamed_arguments(deltas, index)) == arguments
def test_rust_tool_parser_adapter_ignores_midstream_empty_delta() -> None:
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())
text = build_tool_call()
split_at = len(TC_START) + 8
deltas = []
previous_text = ""
for delta_text in (text[:split_at], "", text[split_at:], ""):
current_text = previous_text + delta_text
delta = parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[1],
request=MagicMock(),
)
previous_text = current_text
if delta is not None:
deltas.append(delta)
names = [
tool_call.function.name
for delta in deltas
for tool_call in delta.tool_calls or []
if tool_call.function is not None and tool_call.function.name is not None
]
assert names == [name for name, _ in EXPECTED_CALLS]
for index, (_, arguments) in enumerate(EXPECTED_CALLS):
assert json.loads(collect_streamed_arguments(deltas, index)) == arguments
def test_rust_tool_parser_adapter_adjust_request_is_opaque() -> None:
tools = sample_tools()
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools)
request = ChatCompletionRequest(
messages=[],
model="m",
tools=tools,
tool_choice="required",
skip_special_tokens=True,
)
adjusted = parser.adjust_request(request)
assert adjusted is request
assert adjusted.skip_special_tokens is False
assert adjusted.structured_outputs is None
+47 -3
View File
@@ -18,7 +18,7 @@ import tempfile
import threading
import time
import warnings
from collections.abc import Callable, Iterable, Sequence
from collections.abc import Callable, Iterable, MutableMapping, Sequence
from contextlib import ExitStack, contextmanager
from multiprocessing import Process, get_context
from pathlib import Path
@@ -149,6 +149,46 @@ ROCM_ENGINE_KWARGS: dict = (
if current_platform.is_rocm()
else {}
)
_TILELANG_TVM_PYTHONPATH_FRAGMENT = os.path.join(
"tilelang", "3rdparty", "tvm", "python"
)
def _sanitize_pythonpath_value(pythonpath: str | None) -> str:
if not pythonpath:
return ""
entries = []
for entry in pythonpath.split(os.pathsep):
normalized = entry.replace(os.sep, "/")
if _TILELANG_TVM_PYTHONPATH_FRAGMENT.replace(os.sep, "/") in normalized:
continue
entries.append(entry)
return os.pathsep.join(entries)
def _sanitize_pythonpath_env(env: MutableMapping[str, str]) -> None:
cleaned = _sanitize_pythonpath_value(env.get("PYTHONPATH"))
if cleaned:
env["PYTHONPATH"] = cleaned
else:
env.pop("PYTHONPATH", None)
def _sanitize_current_pythonpath_env() -> None:
_sanitize_pythonpath_env(os.environ)
@contextmanager
def _temporarily_sanitized_pythonpath_env():
original = os.environ.get("PYTHONPATH")
_sanitize_current_pythonpath_env()
try:
yield
finally:
if original is None:
os.environ.pop("PYTHONPATH", None)
else:
os.environ["PYTHONPATH"] = original
def requires_spawn_multiprocessing() -> bool:
@@ -253,7 +293,8 @@ class RemoteVLLMServer:
getattr(args, "show_hidden_metrics_for_version", None) is not None
)
self._pre_download_model(model, args)
with _temporarily_sanitized_pythonpath_env():
self._pre_download_model(model, args)
self._shutdown_complete = False
# Record GPU memory before server start so we know what
@@ -727,6 +768,7 @@ class RemoteOpenAIServer(RemoteVLLMServer):
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if env_dict is not None:
env.update(env_dict)
_sanitize_pythonpath_env(env)
serve_cmd = ["vllm", "serve", model, *vllm_serve_args]
print(f"Launching RemoteOpenAIServer with: {' '.join(serve_cmd)}")
print(f"Environment variables: {env}")
@@ -754,6 +796,7 @@ class RemoteLaunchRenderServer(RemoteVLLMServer):
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if env_dict is not None:
env.update(env_dict)
_sanitize_pythonpath_env(env)
serve_cmd = ["vllm", "launch", "render", model, *vllm_serve_args]
print(f"Launching RemoteLaunchRenderServer with: {' '.join(serve_cmd)}")
self.proc: subprocess.Popen = subprocess.Popen(
@@ -795,7 +838,8 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):
target=_run_in_new_process_group,
args=(self.child_process_fxn, env_dict, model, vllm_serve_args),
) # type: ignore[assignment]
self.proc.start()
with _temporarily_sanitized_pythonpath_env():
self.proc.start()
def __init__(
self,
@@ -283,94 +283,6 @@ def test_get_configured_preferred_segment_rejects_empty_override():
rdma_utils.get_configured_preferred_segment({"preferred_segment": " "})
def test_get_configured_worker_rnic_prefers_explicit_device_name(monkeypatch):
store_config = worker.MooncakeStoreConfig(
metadata_server="",
local_buffer_size=1,
protocol="rdma",
device_name="rocep139s0",
master_server_address="",
)
assert (
rdma_utils.get_configured_worker_rnic(
protocol=store_config.protocol,
configured_device=store_config.device_name,
)
== "rocep139s0"
)
def test_get_configured_worker_rnic_selects_device_from_explicit_csv(monkeypatch):
monkeypatch.setattr(
rdma_utils,
"get_current_physical_gpu_index",
lambda: 1,
)
store_config = worker.MooncakeStoreConfig(
metadata_server="",
local_buffer_size=1,
protocol="rdma",
device_name="rocep139s0,rocep140s0",
master_server_address="",
)
assert (
rdma_utils.get_configured_worker_rnic(
protocol=store_config.protocol,
configured_device=store_config.device_name,
)
== "rocep140s0"
)
def test_get_configured_worker_rnic_warns_and_returns_empty_for_rdma_with_no_device(
caplog, monkeypatch
):
"""No device configured + protocol=rdma → emit a clear warning and return ""
so the C++ side handles auto-selection. There is no Python-side fallback."""
monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
with caplog.at_level(logging.WARNING):
result = rdma_utils.get_configured_worker_rnic(
protocol="rdma",
configured_device="",
)
assert result == ""
warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
assert any("No RDMA devices specified" in r.message for r in warnings), (
f"expected fallback warning, got {[r.message for r in warnings]}"
)
def test_get_configured_worker_rnic_silent_for_tcp_with_no_device(caplog, monkeypatch):
"""protocol=tcp + no device → return "" silently (no RDMA, no warning)."""
monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
with caplog.at_level(logging.WARNING):
result = rdma_utils.get_configured_worker_rnic(
protocol="tcp",
configured_device="",
)
assert result == ""
warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
assert not any("RDMA" in r.message for r in warnings), (
"did not expect RDMA warning for tcp protocol, got "
f"{[r.message for r in warnings]}"
)
def test_get_configured_worker_rnic_rejects_short_explicit_csv(monkeypatch):
monkeypatch.setattr(
rdma_utils,
"get_current_physical_gpu_index",
lambda: 2,
)
with pytest.raises(ValueError, match="does not cover local GPU 2"):
rdma_utils.get_configured_worker_rnic(
protocol="rdma",
configured_device="rocep139s0,rocep140s0",
)
class _ReplicaDesc:
def __init__(self, tier: str):
self.tier = tier
+93
View File
@@ -998,6 +998,99 @@ def test_sample_recovered_tokens_uses_fp64_exponential_race_when_requested():
assert torch.equal(actual, expected)
@pytest.mark.parametrize("no_draft_probs", [True, False])
@pytest.mark.parametrize(
"vocab_size",
[
100, # below BLOCK_SIZE: single partial tile with many padding entries
8193, # BLOCK_SIZE + 1: only 1 valid entry in the last tile
10000, # non-aligned, moderate tail
151936, # real-world Qwen3 vocab size from the CVE report
],
)
def test_sample_recovered_tokens_vocab_boundary(vocab_size: int, no_draft_probs: bool):
"""Regression test for GHSA-8wr5-jm2h-8r4f.
When vocab_size is not a multiple of BLOCK_SIZE (8192), the last Triton
tile extends beyond the vocabulary. If all valid entries in that tail tile
have zero target probability, the out-of-range masked positions (score 0)
could win the tl.max tie-break, producing recovered_id >= vocab_size.
This test forces that scenario and asserts every recovered token is valid.
"""
BLOCK_SIZE = 8192
batch_size = 2
max_spec_len = 3
num_tokens = batch_size * max_spec_len
last_tile_start = (vocab_size // BLOCK_SIZE) * BLOCK_SIZE
target_probs = torch.rand(
num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE
)
if last_tile_start > 0:
# Zero out valid entries in the last partial tile so the only
# non-zero scores come from earlier, fully-covered tiles.
target_probs[:, last_tile_start:] = 0.0
else:
# vocab_size < BLOCK_SIZE: single tile. Concentrate all mass on
# entry 0 so the NO_DRAFT_PROBS path (which zeroes the draft
# token entry) can drive all valid scores to zero.
target_probs = torch.zeros_like(target_probs)
target_probs[:, 0] = 1.0
# Re-normalize so it's a valid distribution.
target_probs = target_probs / target_probs.sum(dim=-1, keepdim=True)
draft_probs = torch.rand(
num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE
)
draft_probs = torch.nn.functional.softmax(draft_probs, dim=-1)
if last_tile_start == 0:
# Force draft token to 0 so the NO_DRAFT_PROBS path zeroes the
# only non-zero entry, leaving all valid scores at zero.
draft_token_ids = torch.zeros(
num_tokens, 1, dtype=torch.int32, device=DEVICE_TYPE
)
else:
draft_token_ids = torch.randint(
0, vocab_size, (num_tokens, 1), dtype=torch.int32, device=DEVICE_TYPE
)
temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
generators = {
i: torch.Generator(device=DEVICE_TYPE).manual_seed(42 + i)
for i in range(batch_size)
}
sampling_metadata = create_sampling_metadata(
all_greedy=False, temperature=temperature, generators=generators
)
spec_decode_metadata = create_spec_decode_metadata(
draft_token_ids.reshape(batch_size, max_spec_len).tolist(),
torch.rand(num_tokens, vocab_size, device=DEVICE_TYPE),
)
recovered = sample_recovered_tokens(
max_spec_len,
spec_decode_metadata.num_draft_tokens,
spec_decode_metadata.cu_num_draft_tokens,
draft_token_ids.squeeze(-1),
None if no_draft_probs else draft_probs,
target_probs,
sampling_metadata,
device=DEVICE_TYPE,
)
assert (recovered >= 0).all(), (
f"Recovered token IDs contain negative values: "
f"{recovered[recovered < 0].tolist()}"
)
assert (recovered < vocab_size).all(), (
f"Recovered token IDs >= vocab_size ({vocab_size}): "
f"{recovered[recovered >= vocab_size].tolist()}"
)
########################### Tests for Synthetic Rejection Sampling #########
+45
View File
@@ -64,6 +64,51 @@ def test_sampler_threads_fp64_gumbel_to_topk_topp_sampler():
assert sampler.topk_topp_sampler.use_fp64_gumbel
def test_rocm_aiter_sampler_defers_import_when_generators_force_native(
monkeypatch: pytest.MonkeyPatch,
):
from vllm.v1.sample.ops import topk_topp_sampler
class MockPlatform:
@staticmethod
def is_cuda():
return False
@staticmethod
def is_cpu():
return False
@staticmethod
def is_xpu():
return False
class MockRocmAiterOps:
@staticmethod
def is_enabled():
return True
real_import = __import__
def guard_aiter_sampling_import(name, *args, **kwargs):
if name == "aiter.ops.sampling":
raise AssertionError("aiter sampling import should be deferred")
return real_import(name, *args, **kwargs)
monkeypatch.setattr(topk_topp_sampler, "current_platform", MockPlatform())
monkeypatch.setattr(topk_topp_sampler, "rocm_aiter_ops", MockRocmAiterOps())
monkeypatch.setattr("builtins.__import__", guard_aiter_sampling_import)
sampler = topk_topp_sampler.TopKTopPSampler()
logits = torch.randn(2, 8)
k = torch.full((2,), 2, dtype=torch.int32)
generators = {0: torch.Generator(device=logits.device).manual_seed(0)}
token_ids, logits_to_return = sampler(logits, generators, k, None)
assert token_ids.shape == (2,)
assert logits_to_return is None
def test_random_sample_uses_fp64_exponential_race_when_requested():
torch.set_default_device(DEVICE_TYPE)
probs = torch.tensor(
+137
View File
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Rust build support shared by `setup.py` and the standalone `build_rust.sh`.
This module is the single source of truth for the Rust artifacts shipped in
the vllm package: which crates are built, where their artifacts land, and how
precompiled artifacts are detected.
"""
from __future__ import annotations
import logging
import os
import sys
from pathlib import Path
from setuptools import setup
from setuptools_rust import Binding, RustExtension
from setuptools_rust.build import build_rust
logger = logging.getLogger(__name__)
ROOT_DIR = Path(__file__).resolve().parents[1]
PACKAGE_DIR = ROOT_DIR / "vllm"
def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
return [
RustExtension(
target="vllm.vllm-rs",
path="rust/src/cmd/Cargo.toml",
args=["--bin", "vllm-rs"],
features=["native-tls-vendored"],
binding=Binding.Exec,
optional=optional,
),
RustExtension(
target="vllm._rust_tool_parser",
path="rust/src/tool-parser/python/Cargo.toml",
features=["pyo3/abi3-py38"],
binding=Binding.PyO3,
optional=optional,
py_limited_api=True,
),
]
def should_require_rust_frontend() -> bool:
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
return value.lower() not in ("", "0", "false", "no")
def _expected_artifacts() -> list[tuple[str, Binding]]:
"""(basename, binding) of each artifact installed into the package."""
artifacts = []
for extension in rust_extensions():
for target in extension.target.values():
package, _, name = target.rpartition(".")
assert package == "vllm", f"unexpected Rust target: {target}"
artifacts.append((name, extension.binding))
return artifacts
def _is_artifact_file(filename: str, name: str, binding: Binding) -> bool:
# setuptools-rust installs Exec binaries under their bare name, and PyO3
# modules as `<module>.<ext-suffix>` where the suffix ends with `.so` on
# Linux and macOS alike (e.g. `_rust_foo.abi3.so`).
if binding == Binding.Exec:
return filename == name
return filename.endswith(".so") and filename.split(".", 1)[0] == name
def find_precompiled_artifacts() -> list[Path]:
"""Rust artifacts already present in the package directory."""
return sorted(
path
for path in PACKAGE_DIR.iterdir()
if any(_is_artifact_file(path.name, *spec) for spec in _expected_artifacts())
)
def missing_precompiled_artifacts() -> list[str]:
"""Expected-but-absent artifacts, as file patterns for diagnostics."""
present = [path.name for path in find_precompiled_artifacts()]
return [
str(PACKAGE_DIR / (name if binding == Binding.Exec else f"{name}.*.so"))
for name, binding in _expected_artifacts()
if not any(_is_artifact_file(filename, name, binding) for filename in present)
]
def is_precompiled_artifact_member(member_name: str) -> bool:
"""Whether a wheel member is a Rust artifact (e.g. `vllm/vllm-rs`)."""
package, _, filename = member_name.rpartition("/")
return package == "vllm" and any(
_is_artifact_file(filename, *spec) for spec in _expected_artifacts()
)
class precompiled_build_rust(build_rust):
"""Skips the local Rust build when all precompiled artifacts are present."""
def run(self) -> None:
missing = missing_precompiled_artifacts()
if not missing:
logger.info(
"Skipping local Rust build: using precompiled %s",
find_precompiled_artifacts(),
)
return
logger.warning(
"Precompiled Rust artifacts missing (%s); "
"falling back to local Rust build.",
", ".join(missing),
)
super().run()
def build_artifacts(build_rust_args: list[str]) -> None:
os.chdir(ROOT_DIR)
PACKAGE_DIR.mkdir(exist_ok=True)
setup(
name="vllm-rust-frontend-build",
packages=[],
rust_extensions=rust_extensions(optional=False),
script_args=["build_rust", "--quiet", "--inplace", *build_rust_args],
)
def main() -> None:
build_artifacts(sys.argv[1:])
if __name__ == "__main__":
main()
+3
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
from typing import TYPE_CHECKING
import torch
@@ -783,6 +784,8 @@ class xpu_ops:
return_softmax_lse: bool | None = False,
s_aux: torch.Tensor | None = None,
return_attn_probs: bool | None = False,
mask_mod: Callable | None = None,
aux_tensors: list | None = None,
):
assert cu_seqlens_k is not None or seqused_k is not None, (
"cu_seqlens_k or seqused_k must be provided"
+2
View File
@@ -6,6 +6,7 @@ from vllm.benchmarks.datasets.datasets import (
AIMODataset,
ASRDataset,
BenchmarkDataset,
BFCLDataset,
BlazeditDataset,
BurstGPTDataset,
ConversationDataset,
@@ -49,6 +50,7 @@ __all__ = [
"AIMODataset",
"ASRDataset",
"BenchmarkDataset",
"BFCLDataset",
"BlazeditDataset",
"BurstGPTDataset",
"ConversationDataset",
+250
View File
@@ -86,6 +86,14 @@ class SampleRequest:
lora_request: LoRARequest | None = None
request_id: str | None = None
timestamp: float | None = None
# Pre-built chat messages. When set, the chat backend uses this list
# directly and skips constructing messages from `prompt` + multimodal
# content. Mutually exclusive with the `prompt`-based path.
chat_messages: list[dict[str, Any]] | None = None
# Per-request fields merged into the request body (e.g. tools,
# tool_choice, response_format). Shallow-merged with --extra-body at
# dispatch time; per-request keys win.
request_overrides: dict | None = None
# -----------------------------------------------------------------------------
@@ -1822,6 +1830,19 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
"from the sampled HF dataset.",
)
bfcl_group = parser.add_argument_group(
"BFCL dataset options", description=BFCLDataset.__doc__
)
bfcl_group.add_argument(
"--bfcl-categories",
type=lambda s: [c.strip() for c in s.split(",") if c.strip()],
default=None,
help="Comma-separated list of BFCL v3 category names (without the "
"'BFCL_v3_' prefix or '.json' suffix) to sample from, e.g. "
"'simple,live_simple,multiple'. Defaults to "
f"'{','.join(BFCLDataset.DEFAULT_CATEGORIES)}'.",
)
prefix_repetition_group = parser.add_argument_group(
"prefix repetition dataset options"
)
@@ -2249,6 +2270,20 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
dataset_class = MMStarDataset
args.hf_split = args.hf_split if args.hf_split else "val"
args.hf_subset = None
elif (
args.dataset_path in BFCLDataset.SUPPORTED_DATASET_PATHS
or args.hf_name in BFCLDataset.SUPPORTED_DATASET_PATHS
):
if args.backend != "openai-chat":
raise ValueError(
"BFCL dataset requires the 'openai-chat' backend because "
"it sends per-request tool schemas via chat completions."
)
dataset_class = BFCLDataset
# BFCL does not use HF splits/subsets; stub values for base init.
args.hf_split = args.hf_split if args.hf_split else "train"
args.hf_subset = None
hf_kwargs = {"categories": args.bfcl_categories}
else:
supported_datasets = set(
[
@@ -4320,6 +4355,221 @@ class MMStarDataset(HuggingFaceDataset):
return sampled_requests
# -----------------------------------------------------------------------------
# BFCL (Berkeley Function Calling Leaderboard) Dataset Implementation
# -----------------------------------------------------------------------------
class BFCLDataset(HuggingFaceDataset):
"""Berkeley Function Calling Leaderboard dataset.
https://huggingface.co/datasets/gorilla-llm/Berkeley-Function-Calling-Leaderboard
BFCL ships one JSON-lines file per category at the repo root (e.g.
``BFCL_v3_simple.json``, ``BFCL_v3_live_simple.json``) rather than a
single HuggingFace split. Each record has ``{id, question, function}``
where ``function`` uses a non-OpenAI schema dialect (``"type": "dict"``).
This dataset loader:
- downloads the selected per-category files via ``hf_hub_download``
and interleaves rows round-robin so sampling is balanced
- translates BFCL function schemas to OpenAI tool format
- sets :attr:`SampleRequest.chat_messages` directly and attaches
``tools`` / ``tool_choice`` via :attr:`SampleRequest.request_overrides`,
producing production-alike tool calling traffic when used with an
``openai-chat`` backend
"""
DEFAULT_OUTPUT_LEN = 512
DEFAULT_CATEGORIES = ("simple", "live_simple", "multiple")
SUPPORTED_DATASET_PATHS = {
"gorilla-llm/Berkeley-Function-Calling-Leaderboard",
}
IS_MULTIMODAL = False
# BFCL primitive type names that are not valid JSON Schema types.
# Map them to the closest JSON Schema equivalent so that grammar
# backends (xgrammar, outlines) accept the translated tool schema.
_TYPE_REMAP = {
"dict": "object",
"float": "number",
"tuple": "array",
"any": "string",
}
def load_data(self) -> None:
"""Defer loading to :meth:`sample` where categories are known."""
self.data = None
def _resolve_categories(self, categories: list[str] | None) -> list[str]:
if not categories:
return list(self.DEFAULT_CATEGORIES)
resolved: list[str] = []
for c in categories:
c = c.strip()
if not c:
continue
resolved.append(c)
return resolved or list(self.DEFAULT_CATEGORIES)
def _load_category(self, category: str) -> list[dict]:
# Local import: huggingface_hub.errors is a small module and
# importing at call site keeps module import cheap for users who
# never touch BFCL.
from huggingface_hub.errors import EntryNotFoundError
filename = f"BFCL_v3_{category}.json"
try:
path = hf_api().hf_hub_download(
self.dataset_path, filename, repo_type="dataset"
)
except EntryNotFoundError as e:
defaults = ", ".join(self.DEFAULT_CATEGORIES)
raise ValueError(
f"BFCL category '{category}' not found: file '{filename}' "
f"does not exist in {self.dataset_path}. Check --bfcl-categories "
f"(defaults: {defaults})."
) from e
rows: list[dict] = []
with open(path) as f:
for line in f:
line = line.strip()
if not line:
continue
rows.append(json.loads(line))
return rows
@classmethod
def _translate_schema(cls, node: Any) -> Any:
"""Recursively translate BFCL-flavored JSON schema to strict JSON Schema."""
if isinstance(node, dict):
translated = {k: cls._translate_schema(v) for k, v in node.items()}
t = translated.get("type")
if isinstance(t, str) and t in cls._TYPE_REMAP:
translated["type"] = cls._TYPE_REMAP[t]
return translated
if isinstance(node, list):
return [cls._translate_schema(v) for v in node]
return node
@classmethod
def _to_openai_tools(cls, functions: list[dict]) -> list[dict]:
tools: list[dict] = []
for fn in functions:
translated = cls._translate_schema(fn)
tools.append({"type": "function", "function": translated})
return tools
def sample(
self,
tokenizer: TokenizerLike,
num_requests: int,
request_id_prefix: str = "",
no_oversample: bool = False,
output_len: int | None = None,
categories: list[str] | None = None,
**kwargs,
) -> list[SampleRequest]:
output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
categories = self._resolve_categories(categories)
per_category_rows: list[list[dict]] = [
self._load_category(c) for c in categories
]
# Round-robin interleave so that when --disable-shuffle is set,
# taking the first num_requests rows still yields balanced category
# coverage. When shuffle is on (the default) this ordering is
# randomized away, which is fine — the subsequent random sample is
# already balanced in expectation.
interleaved: list[dict] = []
max_len = max((len(rows) for rows in per_category_rows), default=0)
for i in range(max_len):
for rows in per_category_rows:
if i < len(rows):
interleaved.append(rows[i])
if not self.disable_shuffle:
rng = random.Random(self.random_seed)
rng.shuffle(interleaved)
sampled_requests: list[SampleRequest] = []
for row in interleaved:
if len(sampled_requests) >= num_requests:
break
question = row.get("question")
functions = row.get("function")
if not question or not functions:
continue
# BFCL question is list[list[dict]] — outer is turns. Use the
# first turn only; skip multi-turn categories in this loader.
if not isinstance(question, list) or not question:
continue
first_turn = question[0]
if not isinstance(first_turn, list) or not first_turn:
continue
messages = first_turn
if not isinstance(functions, list):
functions = [functions]
tools = self._to_openai_tools(functions)
# Best-effort prompt length for percentile bucketing. Pass tools=
# so modern chat templates (Llama 3.1+, Qwen, gpt-oss harmony,
# Hermes) render the tool schemas — without this, the estimate
# misses a significant chunk of the true input for BFCL traffic.
# Older tokenizers reject the kwarg; fall back to tools-free.
try:
rendered = tokenizer.apply_chat_template(
messages,
tools=tools,
tokenize=False,
add_generation_prompt=True,
)
except TypeError:
rendered = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
except Exception as e:
# Unexpected template failure — prompt_len will fall back to a
# plain-text concatenation. Log so the degraded estimate is
# visible instead of silently skewing latency buckets.
logger.warning(
"BFCL: apply_chat_template failed for a sample, falling "
"back to plain-text prompt length: %s",
e,
exc_info=True,
)
rendered = None
if rendered is not None:
prompt_len = len(tokenizer(rendered).input_ids)
else:
text = "\n".join(m.get("content", "") for m in messages)
prompt_len = len(tokenizer(text).input_ids)
# The chat backend uses `messages` directly; `prompt` is only
# kept as a fallback string for display/debug.
prompt_text = messages[-1].get("content", "") if messages else ""
sampled_requests.append(
SampleRequest(
prompt=prompt_text,
prompt_len=prompt_len,
expected_output_len=output_len,
request_id=request_id_prefix + str(len(sampled_requests)),
chat_messages=messages,
request_overrides={
"tools": tools,
"tool_choice": "auto",
},
)
)
self.maybe_oversample_requests(
sampled_requests, num_requests, request_id_prefix, no_oversample
)
return sampled_requests
# -----------------------------------------------------------------------------
# Speed Bench Dataset Implementation
# -----------------------------------------------------------------------------
+8 -1
View File
@@ -79,6 +79,10 @@ class RequestFuncInput:
ignore_eos: bool = False
language: str | None = None
request_id: str | None = None
# Pre-built chat messages. When set, `async_request_openai_chat_completions`
# uses this list directly and skips building messages from `prompt` and
# `multi_modal_content`.
chat_messages: list[dict[str, Any]] | None = None
@dataclass
@@ -343,7 +347,10 @@ async def async_request_openai_chat_completions(
api_url = request_func_input.api_url
_validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions")
messages = _get_chat_messages(request_func_input, mm_position=mm_position)
if request_func_input.chat_messages is not None:
messages = request_func_input.chat_messages
else:
messages = _get_chat_messages(request_func_input, mm_position=mm_position)
payload = {
"model": request_func_input.model_name
+17 -3
View File
@@ -57,6 +57,14 @@ from vllm.utils.network_utils import join_host_port
MILLISECONDS_TO_SECONDS_CONVERSION = 1000
def _merge_overrides(base: dict | None, override: dict | None) -> dict | None:
"""Shallow merge; per-request wins. Returns None if both are empty."""
if not base and not override:
return None
return {**(base or {}), **(override or {})}
TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and (
shutil.which("gnuplot") is not None
)
@@ -753,6 +761,8 @@ async def benchmark(
input_requests[0].expected_output_len,
input_requests[0].multi_modal_data,
)
test_extra_body = _merge_overrides(extra_body, input_requests[0].request_overrides)
test_chat_messages = input_requests[0].chat_messages
assert (
test_mm_content is None
@@ -773,7 +783,8 @@ async def benchmark(
multi_modal_content=test_mm_content,
ignore_eos=ignore_eos,
extra_headers=extra_headers,
extra_body=extra_body,
extra_body=test_extra_body,
chat_messages=test_chat_messages,
)
if ready_check_timeout_sec > 0:
@@ -850,7 +861,8 @@ async def benchmark(
multi_modal_content=test_mm_content,
ignore_eos=ignore_eos,
extra_headers=extra_headers,
extra_body=extra_body,
extra_body=test_extra_body,
chat_messages=test_chat_messages,
)
profile_output = await request_func(
request_func_input=profile_input, session=session
@@ -927,6 +939,7 @@ async def benchmark(
request.multi_modal_data,
request.request_id,
)
per_request_extra_body = _merge_overrides(extra_body, request.request_overrides)
req_model_id, req_model_name = model_id, model_name
if lora_modules:
req_lora_module = next(lora_modules)
@@ -943,8 +956,9 @@ async def benchmark(
multi_modal_content=mm_content,
ignore_eos=ignore_eos,
extra_headers=extra_headers,
extra_body=extra_body,
extra_body=per_request_extra_body,
request_id=request_id,
chat_messages=request.chat_messages,
)
tasks.append(
asyncio.create_task(
@@ -5,8 +5,6 @@
from collections.abc import Mapping
from typing import Any
import torch
import vllm.envs as envs
from vllm.logger import init_logger
@@ -20,22 +18,6 @@ def normalize_string_override(value: Any) -> str | None:
return normalized or None
def get_current_physical_gpu_index() -> int | None:
try:
from vllm.platforms import current_platform
except ImportError:
return None
try:
device_index = torch.accelerator.current_device_index()
physical_device_id = current_platform.device_id_to_physical_device_id(
device_index
)
return int(physical_device_id)
except Exception:
return None
def get_requester_local_hostname(local_ip: str) -> str:
override = normalize_string_override(envs.MOONCAKE_REQUESTER_LOCAL_HOSTNAME)
if override is not None:
@@ -62,54 +44,3 @@ def get_configured_preferred_segment(
)
return env_value
return None
def _get_explicit_worker_rnic(device_list: str) -> str:
entries = [entry.strip() for entry in device_list.split(",")]
if any(not entry for entry in entries):
raise ValueError(
"Mooncake worker device_name contains an empty RDMA device entry"
)
if len(entries) == 1:
return entries[0]
gpu_index = get_current_physical_gpu_index()
if gpu_index is None:
raise RuntimeError(
"Mooncake RDMA requester could not determine the local physical GPU index"
)
if gpu_index >= len(entries):
raise ValueError(
"Mooncake worker device list does not cover local GPU "
f"{gpu_index}: {device_list}"
)
device_name = entries[gpu_index]
logger.info(
"Mooncake selected worker RNIC %s from explicit device list for local GPU %s",
device_name,
gpu_index,
)
return device_name
def get_configured_worker_rnic(
*,
protocol: str,
configured_device: str,
) -> str:
normalized_device = normalize_string_override(configured_device)
if normalized_device is not None:
return _get_explicit_worker_rnic(normalized_device)
if protocol not in {"rdma", "efa"}:
return ""
logger.warning(
"No RDMA devices specified for Mooncake backend (protocol=%s). "
"Set 'device_name' in mooncake_config.json to a single RNIC name "
"or a comma-separated CSV indexed by physical GPU; falling back to "
"Mooncake's built-in auto-selection, which may converge on the same "
"NIC across all DP ranks and saturate bandwidth.",
protocol,
)
return ""
@@ -141,7 +141,7 @@ class MooncakeStoreConfig:
)
@staticmethod
def load_from_env() -> "MooncakeStoreConfig":
def load_from_config() -> "MooncakeStoreConfig":
config_path = os.getenv("MOONCAKE_CONFIG_PATH")
if not config_path:
raise ValueError(
@@ -983,16 +983,12 @@ class MooncakeStoreWorker:
)
# Initialize MooncakeDistributedStore with its own TransferEngine
store_config = MooncakeStoreConfig.load_from_env()
store_config = MooncakeStoreConfig.load_from_config()
extra_config = (
vllm_config.kv_transfer_config.kv_connector_extra_config
if vllm_config.kv_transfer_config
else {}
)
store_config.device_name = rdma_utils.get_configured_worker_rnic(
protocol=store_config.protocol,
configured_device=store_config.device_name,
)
self.store = MooncakeDistributedStore()
local_ip = get_ip()
local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
@@ -72,7 +72,7 @@ from vllm.renderers import ChatParams
from vllm.sampling_params import BeamSearchParams, SamplingParams
from vllm.tokenizers import TokenizerLike
from vllm.utils.collection_utils import as_list
from vllm.utils.mistral import is_mistral_tokenizer, is_mistral_tool_parser
from vllm.utils.mistral import is_mistral_tool_parser
if TYPE_CHECKING:
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
@@ -376,7 +376,6 @@ class OpenAIServingChat(OpenAIServing):
conversation,
tokenizer,
request_metadata,
reasoning_parser,
chat_template_kwargs=chat_template_kwargs,
)
@@ -405,7 +404,6 @@ class OpenAIServingChat(OpenAIServing):
conversation: list[ConversationMessage],
tokenizer: TokenizerLike,
request_metadata: RequestResponseMetadata,
reasoning_parser: ReasoningParser | None = None,
chat_template_kwargs: dict[str, Any] | None = None,
) -> AsyncGenerator[str, None]:
created_time = int(time.time())
@@ -425,43 +423,18 @@ class OpenAIServingChat(OpenAIServing):
harmony_tools_streamed = [False] * num_choices
tools_streamed = [False] * num_choices
is_mistral_grammar_path = request._grammar_from_tool_parser
if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
tool_choice_function_name = request.tool_choice.function.name
else:
tool_choice_function_name = None
# Determine whether tools are in use with "auto" tool choice
tool_choice_auto = (
not tool_choice_function_name
and self._should_stream_with_auto_tool_parsing(request)
)
all_previous_token_ids: list[list[int]] | None
if self.tool_call_id_type == "kimi_k2":
history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
else:
history_tool_call_cnt = 0
# Always track previous_texts for comprehensive output logging
previous_texts = [""] * num_choices
# Only one of these will be used, thus previous_texts and
# all_previous_token_ids will not be used twice in the same iteration.
if (
is_mistral_grammar_path
or tool_choice_auto
or tool_choice_function_name
or request.tool_choice == "required"
or reasoning_parser
):
all_previous_token_ids = [[] for _ in range(num_choices)]
reasoning_end_arr = [False] * num_choices
prompt_is_reasoning_end_arr: list[bool | None] = [None] * num_choices
else:
all_previous_token_ids = None
try:
if self.parser_cls is not None:
if tokenizer is None:
@@ -592,18 +565,6 @@ class OpenAIServingChat(OpenAIServing):
for output in res.outputs:
i = output.index
parser = parsers[i]
tool_parser = parser.tool_parser if parser is not None else None
if (
reasoning_parser
and res.prompt_token_ids
and prompt_is_reasoning_end_arr[i] is None
):
# only check once per choice, because prompt_token_ids
# are the same for all deltas in that choice
prompt_is_reasoning_end_arr[i] = (
reasoning_parser.is_reasoning_end(res.prompt_token_ids)
)
if finish_reason_sent[i]:
continue
@@ -656,27 +617,6 @@ class OpenAIServingChat(OpenAIServing):
delta_message: DeltaMessage | None
# just update previous_texts and previous_token_ids
if (
is_mistral_grammar_path
or tool_choice_auto
or tool_choice_function_name
or request.tool_choice == "required"
or reasoning_parser
):
assert previous_texts is not None
assert all_previous_token_ids is not None
previous_text = previous_texts[i]
previous_token_ids = all_previous_token_ids[i]
current_text = previous_text + delta_text
# avoid the None + list error.
if previous_token_ids:
current_token_ids = previous_token_ids + as_list(
output.token_ids
)
else:
current_token_ids = as_list(output.token_ids)
if self.use_harmony:
delta_message, tools_streamed_flag = (
extract_harmony_streaming_delta(
@@ -687,35 +627,6 @@ class OpenAIServingChat(OpenAIServing):
)
)
harmony_tools_streamed[i] |= tools_streamed_flag
# Mistral grammar path: combined reasoning + tool streaming
elif is_mistral_grammar_path:
from vllm.tool_parsers.mistral_tool_parser import (
MistralToolParser,
)
assert tool_parser is not None
assert isinstance(tool_parser, MistralToolParser)
assert reasoning_end_arr is not None
output_token_ids = as_list(output.token_ids)
result = tool_parser.extract_maybe_reasoning_and_tool_streaming(
reasoning_parser=reasoning_parser,
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=previous_token_ids,
current_token_ids=current_token_ids,
output_token_ids=output_token_ids,
reasoning_ended=reasoning_end_arr[i],
prompt_is_reasoning_end=(prompt_is_reasoning_end_arr[i]),
request=request,
)
delta_message = result.delta_message
reasoning_end_arr[i] = result.reasoning_ended
current_text = result.current_text
current_token_ids = result.current_token_ids
if result.tools_called:
tools_streamed[i] = True
elif parser is not None:
delta_message = parser.parse_delta(
delta_text=delta_text,
@@ -730,22 +641,7 @@ class OpenAIServingChat(OpenAIServing):
else:
delta_message = DeltaMessage(content=delta_text)
# update the previous values for the next iteration
if (
is_mistral_grammar_path
or tool_choice_auto
or tool_choice_function_name
or request.tool_choice == "required"
or reasoning_parser
) and not self.use_harmony:
assert previous_texts is not None
assert all_previous_token_ids is not None
previous_texts[i] = current_text
all_previous_token_ids[i] = current_token_ids
else:
# Update for comprehensive logging even in simple case
assert previous_texts is not None
previous_texts[i] += delta_text
previous_texts[i] += delta_text
# set the previous values for the next iteration
previous_num_tokens[i] += len(output.token_ids)
@@ -1067,32 +963,8 @@ class OpenAIServingChat(OpenAIServing):
tool_calls = []
auto_tools_called = False
if is_mistral_tokenizer(tokenizer):
from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
tool_call_class: type[ToolCall] = MistralToolCall
else:
tool_call_class = ToolCall
use_mistral_tool_parser = request._grammar_from_tool_parser
if use_mistral_tool_parser:
from vllm.tool_parsers.mistral_tool_parser import MistralToolParser
tool_call_items = MistralToolParser.build_non_streaming_tool_calls(
tool_calls
)
if tool_call_items:
auto_tools_called = (
request.tool_choice is None or request.tool_choice == "auto"
)
message = ChatMessage(
role=role,
reasoning=reasoning,
content=content,
tool_calls=tool_call_items,
)
elif (not self.enable_auto_tools or not self.tool_parser) and (
if (not self.enable_auto_tools or not self.tool_parser) and (
not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
and request.tool_choice != "required"
):
@@ -1102,70 +974,42 @@ class OpenAIServingChat(OpenAIServing):
request.tool_choice
and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam
):
tool_call_class_items = []
tool_call_items = []
tool_calls = tool_calls or []
for idx, tc in enumerate(tool_calls):
# Use native ID if available (e.g., Kimi K2),
# otherwise generate ID with correct id_type
if tc.id:
tool_call_class_items.append(
tool_call_class(id=tc.id, function=tc)
for tc in tool_calls:
if not tc.id:
tc.id = make_tool_call_id(
id_type=self.tool_call_id_type,
func_name=tc.name,
idx=history_tool_call_cnt,
)
else:
# Generate ID using the correct format (kimi_k2 or random),
# but leave it to the class if it's Mistral to preserve
# 9-char IDs
if is_mistral_tokenizer(tokenizer):
tool_call_class_items.append(tool_call_class(function=tc))
else:
generated_id = make_tool_call_id(
id_type=self.tool_call_id_type,
func_name=tc.name,
idx=history_tool_call_cnt,
)
tool_call_class_items.append(
tool_call_class(id=generated_id, function=tc)
)
tool_call_items.append(ToolCall(id=tc.id, function=tc))
history_tool_call_cnt += 1
message = ChatMessage(
role=role,
reasoning=reasoning,
content="",
tool_calls=tool_call_class_items,
content=content or "",
tool_calls=tool_call_items,
)
elif request.tool_choice and request.tool_choice == "required":
tool_call_class_items = []
tool_call_items = []
tool_calls = tool_calls or []
for idx, tool_call in enumerate(tool_calls):
# Use native ID if available,
# otherwise generate ID with correct id_type
if tool_call.id:
tool_call_class_items.append(
tool_call_class(id=tool_call.id, function=tool_call)
for tool_call in tool_calls:
if not tool_call.id:
tool_call.id = make_tool_call_id(
id_type=self.tool_call_id_type,
func_name=tool_call.name,
idx=history_tool_call_cnt,
)
else:
# Generate ID using the correct format (kimi_k2 or random),
# but leave it to the class if it's Mistral to preserve
# 9-char IDs
if is_mistral_tokenizer(tokenizer):
tool_call_class_items.append(
tool_call_class(function=tool_call)
)
else:
generated_id = make_tool_call_id(
id_type=self.tool_call_id_type,
func_name=tool_call.name,
idx=history_tool_call_cnt,
)
tool_call_class_items.append(
tool_call_class(id=generated_id, function=tool_call)
)
tool_call_items.append(
ToolCall(id=tool_call.id, function=tool_call)
)
history_tool_call_cnt += 1
message = ChatMessage(
role=role,
content="",
tool_calls=tool_call_class_items,
content=content or "",
tool_calls=tool_call_items,
reasoning=reasoning,
)
@@ -1181,34 +1025,17 @@ class OpenAIServingChat(OpenAIServing):
and self.enable_auto_tools
and self.tool_parser
):
# In the OpenAI API the finish_reason is "tools_called"
# if the tool choice is auto and the model produced a tool
# call. The same is not true for named function calls
auto_tools_called = tool_calls is not None and len(tool_calls) > 0
if tool_calls:
tool_call_items = []
for idx, tc in enumerate(tool_calls):
# Use native ID if available (e.g., Kimi K2),
# otherwise generate ID with correct id_type
if tc.id:
tool_call_items.append(
tool_call_class(id=tc.id, function=tc)
for tc in tool_calls:
if not tc.id:
tc.id = make_tool_call_id(
id_type=self.tool_call_id_type,
func_name=tc.name,
idx=history_tool_call_cnt,
)
else:
# Generate ID using the correct format (kimi_k2 or random),
# but leave it to the class if it's Mistral to preserve
# 9-char IDs
if is_mistral_tokenizer(tokenizer):
tool_call_items.append(tool_call_class(function=tc))
else:
generated_id = make_tool_call_id(
id_type=self.tool_call_id_type,
func_name=tc.name,
idx=history_tool_call_cnt,
)
tool_call_items.append(
tool_call_class(id=generated_id, function=tc)
)
tool_call_items.append(ToolCall(id=tc.id, function=tc))
history_tool_call_cnt += 1
message = ChatMessage(
role=role,
@@ -1218,18 +1045,10 @@ class OpenAIServingChat(OpenAIServing):
)
else:
# FOR NOW make it a chat message; we will have to detect
# the type to make it later.
ret_content = content
# try to use content return from tool parser first,
# tool parser may do some modify for the content.
if content and len(content) > 0:
ret_content = content
message = ChatMessage(
role=role,
reasoning=reasoning,
content=ret_content,
content=content,
)
# undetermined case that is still important to handle
@@ -1451,19 +1270,3 @@ class OpenAIServingChat(OpenAIServing):
)
return ChatCompletionLogProbs(content=logprobs_content)
def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest):
"""
Utility function to check if streamed tokens should go through the tool
call parser that was configured.
We only want to do this IF user-provided tools are set, a tool parser
is configured, "auto" tool choice is enabled, and the request's tool
choice field indicates that "auto" tool choice should be used.
"""
return (
request.tools
and self.tool_parser
and self.enable_auto_tools
and request.tool_choice in ["auto", None]
)
@@ -150,12 +150,12 @@ def create_tool_definition(tool: ChatCompletionToolsParam | Tool):
if isinstance(tool, ChatCompletionToolsParam):
return ToolDescription.new(
name=tool.function.name,
description=tool.function.description,
description=tool.function.description or "",
parameters=tool.function.parameters,
)
return ToolDescription.new(
name=tool.name,
description=tool.description,
description=tool.description or "",
parameters=tool.parameters,
)
@@ -115,6 +115,7 @@ class OpenAISpeechToText(OpenAIServing):
self.enable_force_include_usage = enable_force_include_usage
self.max_audio_filesize_mb = envs.VLLM_MAX_AUDIO_CLIP_FILESIZE_MB
self.max_audio_decode_duration_s: int = envs.VLLM_MAX_AUDIO_DECODE_DURATION_S
if self.model_cls.supports_segment_timestamp:
self.tokenizer = cast(
PreTrainedTokenizerBase,
@@ -216,7 +217,11 @@ class OpenAISpeechToText(OpenAIServing):
# pre-requisite for chunking, as it assumes Whisper SR.
try:
with io.BytesIO(audio_data) as buf:
y, sr = load_audio(buf, sr=self.asr_config.sample_rate)
y, sr = load_audio(
buf,
sr=self.asr_config.sample_rate,
max_duration_s=self.max_audio_decode_duration_s,
)
except Exception as exc:
raise ValueError("Invalid or unsupported audio file.") from exc
+9
View File
@@ -78,6 +78,7 @@ if TYPE_CHECKING:
VLLM_MEDIA_URL_ALLOW_REDIRECTS: bool = True
VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8
VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25
VLLM_MAX_AUDIO_DECODE_DURATION_S: int = 600
VLLM_VIDEO_LOADER_BACKEND: str = "opencv"
VLLM_MEDIA_CONNECTOR: str = "http"
VLLM_MM_HASHER_ALGORITHM: str = "blake3"
@@ -955,6 +956,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
"VLLM_MAX_AUDIO_CLIP_FILESIZE_MB": lambda: int(
os.getenv("VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", "25")
),
# Maximum decoded audio duration in seconds. Compressed audio files
# (e.g. OPUS at very low bitrate) can expand into gigabytes of float32
# PCM. This limit is enforced *during* decoding so the memory is never
# allocated. Default is 600s (10 minutes).
"VLLM_MAX_AUDIO_DECODE_DURATION_S": lambda: int(
os.getenv("VLLM_MAX_AUDIO_DECODE_DURATION_S", "600")
),
# Backend for Video IO — selects the frame-sampling algorithm.
# - "opencv": uniform sampling.
# - "opencv_dynamic": duration-aware dynamic sampling.
@@ -2156,6 +2164,7 @@ def compile_factors() -> dict[str, object]:
"VLLM_MEDIA_URL_ALLOW_REDIRECTS",
"VLLM_MEDIA_LOADING_THREAD_COUNT",
"VLLM_MAX_AUDIO_CLIP_FILESIZE_MB",
"VLLM_MAX_AUDIO_DECODE_DURATION_S",
"VLLM_VIDEO_LOADER_BACKEND",
"VLLM_MEDIA_CONNECTOR",
"VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME",
@@ -113,6 +113,12 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
from aiter.ops.triton.quant_moe import downcast_to_static_fp8
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
should_use_cdna4_mx_scale_swizzle,
)
_swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
assert quant_config.w1_precision is not None, (
"w1_precision in quant config can't be None"
)
@@ -135,7 +141,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
routing_data,
gather_indx=gather_indx,
gammas=gammas if apply_router_weight_on_input else None,
swizzle_mx_scale="CDNA4_SCALE",
swizzle_mx_scale=_swizzle_mx_scale,
out_dtype=torch.float8_e4m3fn,
apply_swiglu=True,
alpha=swiglu_alpha,
@@ -155,7 +161,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
routing_data,
scatter_indx=scatter_indx,
gammas=None if apply_router_weight_on_input else gammas,
swizzle_mx_scale="CDNA4_SCALE",
swizzle_mx_scale=_swizzle_mx_scale,
unpadded_N=unpadded_N_w2,
unpadded_K=unpadded_K_w2,
)
@@ -262,13 +262,37 @@ def fused_topk_bias(
topk_weights *= routed_scaling_factor
return topk_weights, topk_ids
if scoring_func == "sqrtsoftplus":
M = hidden_states.size(0)
topk_weights = torch.empty(
M, topk, dtype=torch.float32, device=hidden_states.device
)
topk_ids = torch.empty(
M,
topk,
dtype=torch.int32 if indices_type is None else indices_type,
device=hidden_states.device,
)
token_expert_indices = torch.empty(
M, topk, dtype=torch.int32, device=hidden_states.device
)
return vllm_topk_softplus_sqrt(
topk_weights,
topk_ids,
token_expert_indices,
gating_output,
renormalize,
e_score_correction_bias,
input_tokens,
hash_indices_table,
routed_scaling_factor,
)
n_routed_experts = gating_output.shape[-1]
if scoring_func == "softmax":
scores = gating_output.softmax(dim=-1)
elif scoring_func == "sigmoid":
scores = gating_output.sigmoid()
elif scoring_func == "sqrtsoftplus":
scores = F.softplus(gating_output).sqrt()
else:
raise ValueError(f"Unsupported scoring function: {scoring_func}")
if e_score_correction_bias is not None:
@@ -109,10 +109,10 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
# Native ROCm HIP kernels (RDNA3, etc.)
if current_platform.is_rocm():
from . import rocm_moe
from . import rocm_moe_rdna
if rocm_moe.is_supported(weight_quant):
return rocm_moe.make_method(
if rocm_moe_rdna.is_supported(weight_quant):
return rocm_moe_rdna.make_method(
weight_quant, input_quant, layer.moe_config
)
@@ -19,6 +19,20 @@ logger = init_logger(__name__)
CK_MXFP4_MOE_DIM_ALIGNMENT = 256
def should_use_cdna4_mx_scale_swizzle() -> bool:
"""Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950.
CDNA4 swizzle requires BLOCK_K%256==0; at TP>=4 the A8W4 dispatch
picks BK<256 tiles for the smaller per-rank shapes, so swizzle must
be off. Used by both the weight-load swizzle in `_swizzle_mxfp4` and
the kernel-argument gate in `aiter_mxfp4_w4a8_moe`; they must agree.
"""
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.platforms.rocm import on_gfx950
return on_gfx950() and get_tensor_model_parallel_world_size() <= 2
def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
"""weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
assert has_triton_kernels()
@@ -44,10 +58,8 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
value_layout = StridedLayout
scale_layout = StridedLayout
elif current_platform.is_rocm():
from vllm.platforms.rocm import on_gfx950
value_layout = StridedLayout
if on_gfx950():
if should_use_cdna4_mx_scale_swizzle():
try:
# triton < 3.6
from triton_kernels.tensor_details.layout import GFX950MXScaleLayout
+2 -2
View File
@@ -335,8 +335,8 @@ class AriaTextModel(LlamaModel, SupportsQuant):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
"experts.w13_weight": ["experts.fc1.weight"],
"experts.w2_weight": ["experts.fc2.weight"],
"experts.routed_experts.w13_weight": ["experts.fc1.weight"],
"experts.routed_experts.w2_weight": ["experts.fc2.weight"],
}
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+27 -79
View File
@@ -25,6 +25,7 @@ from torch import nn
from vllm.config import PoolerConfig, VllmConfig
from vllm.model_executor.layers.pooler import Pooler
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
from .bert import BertEmbeddingModel, BertModel
from .interfaces import HasInnerState, IsHybrid, SupportsLateInteraction
@@ -217,38 +218,12 @@ class ColBERTModel(ColBERTMixin, BertEmbeddingModel):
return self._build_colbert_pooler(pooler_config)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
def _strip(name: str) -> str:
for p in ("model.", "bert."):
if name.startswith(p):
name = name[len(p) :]
return name
weights_list = list(weights)
model_side: list[tuple[str, torch.Tensor]] = []
colbert_side: list[tuple[str, torch.Tensor]] = []
for name, weight in weights_list:
stripped = _strip(name)
# Handle different checkpoint naming conventions
if stripped in ("linear.weight", "colbert_linear.weight"):
colbert_side.append(("colbert_linear.weight", weight))
elif stripped.startswith("linear.") or stripped.startswith(
"colbert_linear."
):
new_name = stripped.replace("linear.", "colbert_linear.")
colbert_side.append((new_name, weight))
else:
model_side.append((stripped, weight))
loaded: set[str] = set()
loaded_model = self.model.load_weights(model_side)
loaded.update({"model." + n for n in loaded_model})
if colbert_side:
_, colbert_loaded = self._load_colbert_weights(colbert_side)
loaded.update(colbert_loaded)
return loaded
other_weights, colbert_loaded = self._load_colbert_weights(weights)
# Force "bert." to become "model."
mapper = WeightsMapper(orig_to_new_prefix={"bert.": "model."})
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(other_weights, mapper=mapper)
return loaded | colbert_loaded
# -----------------------------------------------------------------------
@@ -309,18 +284,14 @@ class ColBERTModernBertModel(ColBERTMixin, nn.Module):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
other_weights, colbert_loaded = self._load_colbert_weights(weights)
# Strip "model." prefix added by the embedding adapter
model_weights = [
(n[len("model.") :] if n.startswith("model.") else n, w)
for n, w in other_weights
]
loaded_model = self.model.load_weights(other_weights)
loaded = {f"model.{name}" for name in loaded_model} | colbert_loaded
loaded_model = self.model.load_weights(model_weights)
loaded = {"model." + n for n in loaded_model} | colbert_loaded
# When the ST projector was auto-loaded during init
# (not from the main checkpoint), mark its params as loaded
# so the weight validator doesn't complain.
# When the ST projector is loaded via `_build_colbert_pooler`, the weights
# might come from `colbert_loaded` or the pooler automatically falls back to
# load from `/1_Dense` etc.
# We need to mark its params as loaded so the weight validator doesn't complain
# when they are loaded via fallback.
if hasattr(self.pooler, "head"):
head = self.pooler.head
projector = getattr(head, "projector", None)
@@ -385,36 +356,15 @@ class ColBERTJinaRobertaModel(ColBERTMixin, nn.Module):
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
weights_list = list(weights)
model_side: list[tuple[str, torch.Tensor]] = []
colbert_side: list[tuple[str, torch.Tensor]] = []
other_weights, colbert_loaded = self._load_colbert_weights(weights)
for name, weight in weights_list:
stripped = name
# Strip "model." prefix added by the embedding adapter
if stripped.startswith("model."):
stripped = stripped[len("model.") :]
# Strip "roberta." prefix from checkpoint
if stripped.startswith("roberta."):
stripped = stripped[len("roberta.") :]
mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."})
if stripped in ("linear.weight", "colbert_linear.weight"):
colbert_side.append(("colbert_linear.weight", weight))
elif stripped.startswith("pooler."):
# Skip HF pooler weights (not used in ColBERT)
continue
else:
model_side.append((stripped, weight))
# Skip HF pooler weights (model.pooler.*) as they not used in ColBERT
loader = AutoWeightsLoader(self, skip_prefixes=["model.pooler."])
loaded: set[str] = set()
loaded_model = self.model.load_weights(model_side)
loaded.update({"model." + n for n in loaded_model})
if colbert_side:
_, colbert_loaded = self._load_colbert_weights(colbert_side)
loaded.update(colbert_loaded)
return loaded
loaded = loader.load_weights(other_weights, mapper=mapper)
return loaded | colbert_loaded
# -----------------------------------------------------------------------
@@ -491,17 +441,15 @@ class ColBERTLfm2Model(ColBERTMixin, nn.Module, HasInnerState, IsHybrid):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
other_weights, colbert_loaded = self._load_colbert_weights(weights)
# Strip "model." prefix added by the embedding adapter
model_weights = [
(n[len("model.") :] if n.startswith("model.") else n, w)
for n, w in other_weights
]
loaded_model = self.model.load_weights(model_weights)
loaded_model = self.model.load_weights(other_weights)
loaded = {f"model.{name}" for name in loaded_model} | colbert_loaded
# When the ST projector was auto-loaded during init
# (not from the main checkpoint), mark its params as loaded
# so the weight validator doesn't complain.
# When the ST projector is loaded via `_build_colbert_pooler`, the weights
# might come from `colbert_loaded` or the pooler automatically falls back to
# load from `/1_Dense` etc.
# We need to mark its params as loaded so the weight validator doesn't complain
# when they are loaded via fallback.
if hasattr(self.pooler, "head"):
head = self.pooler.head
projector = getattr(head, "projector", None)
+32 -33
View File
@@ -57,50 +57,49 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):
class Gemma4Config(VerifyAndUpdateConfig):
@staticmethod
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
"""Force unified attention backend for models with heterogeneous
head dimensions.
"""Configure attention for heterogeneous head dimensions.
Some Gemma4 variants use different head dimensions for
sliding window (head_dim) vs full attention (global_head_dim) layers.
When global_head_dim > 256, FlashAttention rejects those layers
(head_size <= 256 kernel limit), causing vLLM to select a different
backend for each layer type. This mixed-backend execution produces
numerical divergence and output corruption.
Gemma4 uses different head dimensions for sliding window
(head_dim) vs full attention (global_head_dim) layers. The
default FA3 on Hopper cannot handle head_dim > 256, which
causes mixed backend selection and numerical divergence.
The fix detects heterogeneous head dimensions from the model config
and forces TRITON_ATTN (which has no head_size ceiling) for all
layers when the user hasn't explicitly chosen a backend.
TODO: Heterogeneous head_sizes (head_dim != global_head_dim)
require NixlConnector changes to support per-layer KV transfer
with different head dimensions for prefill-decode disaggregation.
When FA4 is available we force it for ALL layers, giving a
uniform kernel path and avoiding the mixed FA3+FA4 penalty.
When FA4 is not available we fall back to Triton.
"""
hf_text_config = vllm_config.model_config.hf_text_config
head_dim = getattr(hf_text_config, "head_dim", None)
global_head_dim = getattr(hf_text_config, "global_head_dim", None)
# Only force Triton when head dimensions actually differ AND the
# larger one exceeds FlashAttention's kernel limit (head_size <= 256).
# This avoids unnecessary backend forcing on smaller models where
# the config carries global_head_dim but all layers can still use
# the same FA backend.
max_head_dim = max(head_dim or 0, global_head_dim or 0)
if (
head_dim is not None
and global_head_dim is not None
and head_dim != global_head_dim
and max_head_dim > 256
and vllm_config.attention_config.backend is None
):
from vllm.v1.attention.backends.registry import (
AttentionBackendEnum,
)
if head_dim is None or global_head_dim is None or head_dim == global_head_dim:
return
from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
from vllm.v1.attention.backends.registry import AttentionBackendEnum
max_head_dim = max(head_dim, global_head_dim)
if is_fa_version_supported(4) and max_head_dim <= 512:
if (
vllm_config.attention_config.flash_attn_version is None
and vllm_config.attention_config.backend
in (None, AttentionBackendEnum.FLASH_ATTN)
):
vllm_config.attention_config.flash_attn_version = 4
logger.info(
"Gemma4 model has heterogeneous head dimensions "
"(head_dim=%d, global_head_dim=%d). Using FA4 for "
"all layers to avoid mixed FA3/FA4 penalty.",
head_dim,
global_head_dim,
)
elif vllm_config.attention_config.backend is None:
vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
logger.info(
"Gemma4 model has heterogeneous head dimensions "
"(head_dim=%d, global_head_dim=%d). Forcing TRITON_ATTN "
"backend to prevent mixed-backend numerical divergence.",
"(head_dim=%d, global_head_dim=%d). FA4 not available, "
"forcing TRITON_ATTN backend.",
head_dim,
global_head_dim,
)
@@ -31,6 +31,7 @@ from vllm.model_executor.models.deepseek_v2 import (
)
from vllm.multimodal.inputs import NestedTensors
from .interfaces import LocalArgmaxMixin
from .utils import (
AutoWeightsLoader,
get_draft_quant_config,
@@ -309,7 +310,7 @@ class DeepseekV2Eagle3Model(nn.Module):
return loaded_params
class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM):
class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):
"""Eagle3 speculative decoding model for DeepseekV2/V3."""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+35
View File
@@ -1282,6 +1282,41 @@ def supports_any_eagle(
return supports_eagle(model) or supports_eagle3(model)
class LocalArgmaxMixin:
"""Mixin for draft model heads in speculative decoding.
Provides a D2T-aware ``get_top_tokens`` that preserves the
local-argmax communication reduction even when the draft vocabulary
is smaller than the target vocabulary.
When ``draft_id_to_target_id`` is present (shape ``(draft_vocab_size,)``,
containing per-token offset to target vocab id), the draft argmax index
``k`` is mapped to the target vocab id via::
target_id = k + draft_id_to_target_id[k]
This is mathematically equivalent to computing the full-vocab scatter
logits and taking the global argmax, but requires only
O(batch * 2 * tp_size) communication instead of O(batch * vocab_size).
Requires the subclass to expose:
``self.logits_processor``: LogitsProcessor
``self.lm_head``: ParallelLMHead
``self.draft_id_to_target_id`` (optional): nn.Parameter
"""
def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""Vocab-parallel argmax with optional D2T remapping."""
top = self.logits_processor.get_top_tokens(
self.lm_head,
hidden_states,
)
d2t = getattr(self, "draft_id_to_target_id", None)
if d2t is not None:
top = top + d2t[top]
return top
class EagleModelMixin:
aux_hidden_state_layers: tuple[int, ...] = ()
+2 -1
View File
@@ -62,6 +62,7 @@ from vllm.v1.attention.backend import AttentionType
from .adapters import as_embedding_model, as_seq_cls_model
from .interfaces import (
EagleModelMixin,
LocalArgmaxMixin,
SupportsEagle,
SupportsEagle3,
SupportsLoRA,
@@ -487,7 +488,7 @@ class LlamaModel(nn.Module, EagleModelMixin):
class LlamaForCausalLM(
nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
@@ -208,23 +208,6 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
) -> tuple[torch.Tensor, torch.Tensor]:
return self.model(input_ids, positions, hidden_states, inputs_embeds)
def get_top_tokens(
self,
hidden_states: torch.Tensor,
) -> torch.Tensor:
"""Vocab-parallel argmax without all-gathering full logits.
Falls back to full logits when draft_id_to_target_id remapping is
active, since the shared lm_head covers the full target vocab but
the draft model only predicts over a subset (draft_vocab_size).
"""
if (
hasattr(self, "draft_id_to_target_id")
and self.draft_id_to_target_id is not None
):
return self.compute_logits(hidden_states).argmax(dim=-1)
return self.logits_processor.get_top_tokens(self.lm_head, hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
def transform(inputs):
name, loaded_weight = inputs
+3 -3
View File
@@ -229,10 +229,10 @@ class DashengAttention(nn.Module):
)
def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
B, N, C = x.shape
B, N, _ = x.shape
qkv, _ = self.qkv(x)
qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads)
qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv.unbind(0)
@@ -243,7 +243,7 @@ class DashengAttention(nn.Module):
attn_mask=mask[:, None, None, :] if mask is not None else None,
)
x = x.transpose(1, 2).reshape(B, N, C)
x = x.transpose(1, 2).reshape(B, N, self.q_size)
x, _ = self.proj(x)
return x
+8 -2
View File
@@ -48,7 +48,13 @@ from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.config import set_default_rope_theta
from vllm.v1.attention.backend import AttentionType
from .interfaces import SupportsEagle, SupportsEagle3, SupportsLoRA, SupportsPP
from .interfaces import (
LocalArgmaxMixin,
SupportsEagle,
SupportsEagle3,
SupportsLoRA,
SupportsPP,
)
from .qwen2 import Qwen2MLP as Qwen3MLP
from .qwen2 import Qwen2Model
from .utils import AutoWeightsLoader, PPMissingLayer, extract_layer_index, maybe_prefix
@@ -259,7 +265,7 @@ class Qwen3Model(Qwen2Model):
class Qwen3ForCausalLM(
nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
):
packed_modules_mapping = {
"qkv_proj": [
+17 -7
View File
@@ -36,6 +36,9 @@ from vllm.distributed import (
get_pp_group,
)
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe import (
fused_moe_make_expert_params_mapping,
)
from vllm.model_executor.layers.layernorm import (
GemmaRMSNorm as Qwen3_5RMSNorm,
)
@@ -294,13 +297,20 @@ class Qwen3_5Model(Qwen3NextModel):
loaded_params: set[str] = set()
expert_params_mapping = self.get_expert_mapping()
is_fused_expert = False
base_layer = (
"base_layer." if any(".base_layer." in name for name in params_dict) else ""
)
fused_expert_params_mapping = [
(f"experts.{base_layer}w13_weight", "experts.gate_up_proj", 0, "w1"),
(f"experts.{base_layer}w2_weight", "experts.down_proj", 0, "w2"),
]
fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
self,
ckpt_gate_proj_name="gate_up_proj",
ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="gate_up_proj",
num_experts=1,
):
if shard_id == "w3":
continue
parts = ckpt_name.split(".")
fused_expert_params_mapping.append(
(f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
)
num_experts = (
self.config.num_experts if hasattr(self.config, "num_experts") else 0
)
+16 -8
View File
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.interfaces import LocalArgmaxMixin
from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5RMSNorm
from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts
from vllm.sequence import IntermediateTensors
@@ -209,13 +210,20 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
is_fused_expert = False
base_layer = (
"base_layer." if any(".base_layer." in name for name in params_dict) else ""
)
fused_expert_params_mapping = [
(f"experts.{base_layer}w13_weight", "experts.gate_up_proj", 0, "w1"),
(f"experts.{base_layer}w2_weight", "experts.down_proj", 0, "w2"),
]
fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
self,
ckpt_gate_proj_name="gate_up_proj",
ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="gate_up_proj",
num_experts=1,
):
if shard_id == "w3":
continue
parts = ckpt_name.split(".")
fused_expert_params_mapping.append(
(f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
)
num_experts = (
self.config.num_experts if hasattr(self.config, "num_experts") else 0
)
@@ -346,7 +354,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
"hidden_states": 0,
}
)
class Qwen3_5MTP(nn.Module, SupportsMultiModal):
class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
+12 -2
View File
@@ -187,8 +187,18 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
"base_layer." if any(".base_layer." in name for name in params_dict) else ""
)
fused_expert_params_mapping = [
(f"experts.{base_layer}w13_weight", "experts.gate_up_proj", 0, "w1"),
(f"experts.{base_layer}w2_weight", "experts.down_proj", 0, "w2"),
(
f"experts.routed_experts.{base_layer}w13_weight",
"experts.gate_up_proj",
0,
"w1",
),
(
f"experts.routed_experts.{base_layer}w2_weight",
"experts.down_proj",
0,
"w2",
),
]
num_experts = self.config.num_experts
for name, loaded_weight in weights:
+15 -3
View File
@@ -422,9 +422,21 @@ class Step3TextModel(nn.Module):
)
expert_params_mapping = [
(f".moe.experts.{base_layer}w13_weight", ".moe.gate_proj.weight", "w1"),
(f".moe.experts.{base_layer}w13_weight", ".moe.up_proj.weight", "w3"),
(f".moe.experts.{base_layer}w2_weight", ".moe.down_proj.weight", "w2"),
(
f".moe.experts.routed_experts.{base_layer}w13_weight",
".moe.gate_proj.weight",
"w1",
),
(
f".moe.experts.routed_experts.{base_layer}w13_weight",
".moe.up_proj.weight",
"w3",
),
(
f".moe.experts.routed_experts.{base_layer}w2_weight",
".moe.down_proj.weight",
"w2",
),
]
disable_moe_stacked_params = [data[1] for data in expert_params_mapping]
+24 -12
View File
@@ -635,36 +635,48 @@ class Step3p5Model(nn.Module):
# Old packed 3D format: .moe.gate_proj.weight [num_experts, out, in]
expert_params_mapping = [
(f".moe.experts.{base_layer}w13_weight", ".moe.gate_proj.weight", "w1"),
(f".moe.experts.{base_layer}w13_weight", ".moe.up_proj.weight", "w3"),
(f".moe.experts.{base_layer}w2_weight", ".moe.down_proj.weight", "w2"),
(
f".moe.experts.{base_layer}w13_weight_scale_2",
f".moe.experts.routed_experts.{base_layer}w13_weight",
".moe.gate_proj.weight",
"w1",
),
(
f".moe.experts.routed_experts.{base_layer}w13_weight",
".moe.up_proj.weight",
"w3",
),
(
f".moe.experts.routed_experts.{base_layer}w2_weight",
".moe.down_proj.weight",
"w2",
),
(
f".moe.experts.routed_experts.{base_layer}w13_weight_scale_2",
".moe.gate_proj.weight_scale_2",
"w1",
),
(
f".moe.experts.{base_layer}w13_weight_scale_2",
f".moe.experts.routed_experts.{base_layer}w13_weight_scale_2",
".moe.up_proj.weight_scale_2",
"w3",
),
(
f".moe.experts.{base_layer}w2_weight_scale_2",
f".moe.experts.routed_experts.{base_layer}w2_weight_scale_2",
".moe.down_proj.weight_scale_2",
"w2",
),
(
f".moe.experts.{base_layer}w13_weight_scale",
f".moe.experts.routed_experts.{base_layer}w13_weight_scale",
".moe.gate_proj.weight_scale",
"w1",
),
(
f".moe.experts.{base_layer}w13_weight_scale",
f".moe.experts.routed_experts.{base_layer}w13_weight_scale",
".moe.up_proj.weight_scale",
"w3",
),
(
f".moe.experts.{base_layer}w2_weight_scale",
f".moe.experts.routed_experts.{base_layer}w2_weight_scale",
".moe.down_proj.weight_scale",
"w2",
),
@@ -672,17 +684,17 @@ class Step3p5Model(nn.Module):
# input scales are stored as moe.{gate,up,down}_proj.input_scale
# rather than the standard per-expert format handled generically.
(
f".moe.experts.{base_layer}w13_input_scale",
f".moe.experts.routed_experts.{base_layer}w13_input_scale",
".moe.gate_proj.input_scale",
"w1",
),
(
f".moe.experts.{base_layer}w13_input_scale",
f".moe.experts.routed_experts.{base_layer}w13_input_scale",
".moe.up_proj.input_scale",
"w3",
),
(
f".moe.experts.{base_layer}w2_input_scale",
f".moe.experts.routed_experts.{base_layer}w2_input_scale",
".moe.down_proj.input_scale",
"w2",
),
@@ -72,35 +72,7 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
image_sizes=([height, width],), **mm_processor_kwargs
)
image_tokens = mm_tokens["num_image_tokens"][0]
return self._get_max_encoder_tokens(processor, mm_tokens) or image_tokens
@staticmethod
def _get_mm_values(mm_tokens: object, key: str) -> object:
if isinstance(mm_tokens, Mapping):
return mm_tokens.get(key)
return getattr(mm_tokens, key, None)
def _get_max_encoder_tokens(
self, processor: object, mm_tokens: object
) -> int | None:
if "gemma3" not in processor.__class__.__name__.lower():
return None
vision_config = getattr(self.get_hf_config(), "vision_config", None)
image_size = getattr(vision_config, "image_size", None)
patch_size = getattr(vision_config, "patch_size", None)
if not image_size or not patch_size:
return None
# Gemma3 pools each 64x64 SigLIP patch grid down to 256 image tokens.
# Profile the vision encoder against the pre-pooling patch-token count.
patches_per_image = (image_size // patch_size) ** 2
num_image_patches = self._get_mm_values(mm_tokens, "num_image_patches") or [1]
if isinstance(num_image_patches, int):
max_image_patches = num_image_patches
else:
max_image_patches = max(num_image_patches)
return patches_per_image * int(max_image_patches)
return image_tokens
def get_max_image_size(self):
return 10_000, 10_000 # hardcode for arbitrary very large size
+6 -3
View File
@@ -7,7 +7,10 @@ from __future__ import annotations
from typing import TYPE_CHECKING
from vllm.config import get_current_vllm_config
from vllm.model_executor.layers.fused_moe import MoERunner, UnquantizedFusedMoEMethod
from vllm.model_executor.layers.fused_moe import (
RoutedExperts,
UnquantizedFusedMoEMethod,
)
from vllm.model_executor.layers.quantization import QuantizationMethods
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
@@ -129,7 +132,7 @@ class DeepseekV4FP8Config(Fp8Config):
return None
def get_quant_method(self, layer, prefix):
if isinstance(layer, MoERunner):
if isinstance(layer, RoutedExperts):
if is_layer_skipped(
prefix=prefix,
ignored_layers=self.ignored_layers,
@@ -152,6 +155,6 @@ class DeepseekV4FP8Config(Fp8Config):
return super().get_quant_method(layer, prefix)
def is_mxfp4_quant(self, prefix, layer):
if not isinstance(layer, MoERunner) or self.expert_dtype != "fp4":
if not isinstance(layer, RoutedExperts) or self.expert_dtype != "fp4":
return False
return self.moe_quant_algo != "NVFP4"
+30 -6
View File
@@ -1,7 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from PIL import Image
import contextlib
from PIL import Image, ImageOps
def rescale_image_size(
@@ -16,6 +18,13 @@ def rescale_image_size(
return image
def normalize_image(image: Image.Image) -> Image.Image:
"""Normalize EXIF orientation so the pixel data matches visual display."""
with contextlib.suppress(Exception):
image = ImageOps.exif_transpose(image)
return image
def rgba_to_rgb(
image: Image.Image,
background_color: tuple[int, int, int] | list[int] = (255, 255, 255),
@@ -27,10 +36,25 @@ def rgba_to_rgb(
return converted
def convert_image_mode(image: Image.Image, to_mode: str):
def _has_transparency(image: Image.Image) -> bool:
"""Detect whether an image carries transparency data (RGBA, LA, PA,
or tRNS chunk in P/L/RGB PNGs)."""
if image.mode in ("RGBA", "LA", "PA"):
return True
return "transparency" in getattr(image, "info", {})
def convert_image_mode(
image: Image.Image,
to_mode: str,
background_color: tuple[int, int, int] | list[int] = (255, 255, 255),
) -> Image.Image:
if image.mode == to_mode:
return image
elif image.mode == "RGBA" and to_mode == "RGB":
return rgba_to_rgb(image)
else:
return image.convert(to_mode)
if to_mode == "RGB" and _has_transparency(image):
if image.mode != "RGBA":
image = image.convert("RGBA")
return rgba_to_rgb(image, background_color)
return image.convert(to_mode)
+58 -4
View File
@@ -47,6 +47,7 @@ def load_audio_pyav(
*,
sr: float | None = 22050,
mono: bool = True,
max_duration_s: float | None = None,
) -> tuple[npt.NDArray, float]:
"""Load an audio file using PyAV (FFmpeg), returning float32 mono waveform.
@@ -57,6 +58,10 @@ def load_audio_pyav(
Args:
path: A :class:`~io.BytesIO` buffer, a filesystem
:class:`~pathlib.Path`, or a string path.
max_duration_s: If set, abort decoding once the accumulated
sample count exceeds this many seconds of audio. Prevents
decompression-bomb attacks where a small compressed file
expands into gigabytes of PCM.
Returns:
``(waveform, sample_rate)`` where *waveform* is a 1-D float32
@@ -72,6 +77,30 @@ def load_audio_pyav(
native_sr = stream.rate
sr = sr or native_sr
# Early rejection from container/stream metadata to avoid
# wasting resources on decoding decompression bombs.
if max_duration_s is not None:
metadata_duration_s = None
if stream.duration and stream.time_base:
metadata_duration_s = float(stream.duration * stream.time_base)
elif container.duration:
metadata_duration_s = container.duration / 1_000_000
if (
metadata_duration_s is not None
and metadata_duration_s > max_duration_s
):
raise ValueError(
f"Audio exceeds maximum allowed duration of "
f"{max_duration_s}s (metadata reports "
f"{metadata_duration_s:.1f}s). This limit "
f"prevents decompression-bomb attacks."
)
max_samples = (
int(sr * max_duration_s) if max_duration_s is not None else None
)
total_samples = 0
chunks: list[npt.NDArray] = []
needs_resampling = not math.isclose(
float(sr),
@@ -88,9 +117,21 @@ def load_audio_pyav(
if needs_resampling:
assert resampler is not None
for out_frame in resampler.resample(frame):
chunks.append(out_frame.to_ndarray())
arr = out_frame.to_ndarray()
total_samples += arr.shape[-1]
chunks.append(arr)
else:
chunks.append(frame.to_ndarray())
arr = frame.to_ndarray()
total_samples += arr.shape[-1]
chunks.append(arr)
if max_samples is not None and total_samples > max_samples:
raise ValueError(
f"Audio exceeds maximum allowed duration of "
f"{max_duration_s}s (decoded {total_samples} "
f"samples at {sr}Hz). This limit prevents "
f"decompression-bomb attacks."
)
except (ValueError, ImportError):
raise
except Exception as e:
@@ -114,10 +155,20 @@ def load_audio_soundfile(
*,
sr: float | None = 22050,
mono: bool = True,
max_duration_s: float | None = None,
) -> tuple[np.ndarray, int]:
"""Load audio via soundfile"""
with soundfile.SoundFile(path) as f:
native_sr = f.samplerate
if max_duration_s is not None:
file_duration_s = f.frames / native_sr
if file_duration_s > max_duration_s:
raise ValueError(
f"Audio exceeds maximum allowed duration of "
f"{max_duration_s}s (file contains "
f"{file_duration_s:.1f}s at {native_sr}Hz). "
f"This limit prevents decompression-bomb attacks."
)
y = f.read(dtype="float32", always_2d=False).T
if mono and y.ndim > 1:
@@ -134,9 +185,12 @@ def load_audio(
*,
sr: float | None = 22050,
mono: bool = True,
max_duration_s: float | None = None,
):
try:
return load_audio_soundfile(path, sr=sr, mono=mono)
return load_audio_soundfile(
path, sr=sr, mono=mono, max_duration_s=max_duration_s
)
except ImportError as exc:
# soundfile (or resampy) is not installed — fall through to pyav.
# NOTE: this clause must stay BEFORE ``soundfile.LibsndfileError``
@@ -153,7 +207,7 @@ def load_audio(
if isinstance(path, BytesIO):
path.seek(0)
try:
return load_audio_pyav(path, sr=sr, mono=mono)
return load_audio_pyav(path, sr=sr, mono=mono, max_duration_s=max_duration_s)
except ImportError:
raise # Let PlaceholderModule's message ("install vllm[audio]") propagate.
except Exception as pyav_exc:
+5 -2
View File
@@ -11,7 +11,7 @@ from PIL import Image
from vllm.utils.serial_utils import tensor2base64
from ..image import convert_image_mode, rgba_to_rgb
from ..image import convert_image_mode, normalize_image, rgba_to_rgb
from .base import MediaIO, MediaWithBytes
MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0
@@ -65,11 +65,14 @@ class ImageMediaIO(MediaIO[Image.Image]):
elif image.mode == "RGBA" and self.image_mode == "RGB":
return rgba_to_rgb(image, self.rgba_background_color)
else:
return convert_image_mode(image, self.image_mode)
return convert_image_mode(
image, self.image_mode, self.rgba_background_color
)
def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]:
try:
image = Image.open(BytesIO(data))
image = normalize_image(image)
image.load()
image = self._convert_image_mode(image)
except (OSError, Image.UnidentifiedImageError) as e:
+7 -1
View File
@@ -334,7 +334,13 @@ class ImageProcessorItems(ProcessorBatchItems[HfImageItem | None]):
if isinstance(image, PILImage.Image):
return ImageSize(*image.size)
if isinstance(image, (np.ndarray, torch.Tensor)):
_, h, w = image.shape
if image.ndim == 3 and image.shape[-1] in (1, 3, 4):
# HWC format (e.g. from np.array(PIL.Image)).
# PIL images are always channels-last.
h, w = image.shape[0], image.shape[1]
else:
# CHW format (standard PyTorch / numpy convention).
_, h, w = image.shape
return ImageSize(w, h)
assert_never(image)
+3 -20
View File
@@ -43,7 +43,6 @@ from vllm.tool_parsers.streaming import (
extract_required_tool_call_streaming,
)
from vllm.utils import random_uuid
from vllm.utils.mistral import is_mistral_tool_parser
logger = init_logger(__name__)
@@ -546,14 +545,6 @@ class DelegatingParser(Parser):
if tool_parser is None:
return [], content
# When the Mistral grammar factory injected structured outputs,
# let the parser handle the output.
use_mistral_tool_parser = (
is_mistral_tool_parser(type(tool_parser))
and isinstance(request, ChatCompletionRequest)
and request._grammar_from_tool_parser
)
supports_required_and_named = tool_parser.supports_required_and_named
is_named_tool_choice = request.tool_choice and isinstance(
request.tool_choice,
@@ -570,11 +561,7 @@ class DelegatingParser(Parser):
)
tool_calls = list[FunctionCall]()
if (
is_named_tool_choice
and supports_required_and_named
and not use_mistral_tool_parser
):
if is_named_tool_choice and supports_required_and_named:
if content is None:
return [], None
tool_calls.append(
@@ -584,11 +571,7 @@ class DelegatingParser(Parser):
)
)
content = None
elif (
is_required_tool_choice
and supports_required_and_named
and not use_mistral_tool_parser
):
elif is_required_tool_choice and supports_required_and_named:
# "required" with standard JSON-based parsing
parsed_calls = []
with contextlib.suppress(ValidationError):
@@ -604,7 +587,7 @@ class DelegatingParser(Parser):
)
)
content = None
elif is_auto_tool_choice or use_mistral_tool_parser:
elif is_auto_tool_choice:
# Automatic Tool Call Parsing (also used as fallback for
# required/named when supports_required_and_named=False)
tool_call_info = tool_parser.extract_tool_calls(
+77
View File
@@ -0,0 +1,77 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
from typing import TYPE_CHECKING
from vllm.entrypoints.openai.engine.protocol import DeltaMessage, FunctionCall
from vllm.parser.abstract_parser import DelegatingParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
class MistralParser(DelegatingParser):
def __init__(self, tokenizer, tools=None, *args, **kwargs):
super().__init__(tokenizer, tools, *args, **kwargs)
from vllm.tool_parsers.mistral_tool_parser import MistralToolParser
if not isinstance(self._tool_parser, MistralToolParser):
raise ValueError(
"MistralParser requires --tool-call-parser mistral, "
f"got {self._tool_parser.__class__.__name__}."
)
def _maybe_force_auto_tool_parsing(
self, request: ChatCompletionRequest | ResponsesRequest
) -> None:
# When the Mistral grammar factory injected structured outputs,
# the model emits v11+ format ([TOOL_CALLS]name{args}) that the
# named/required parsers can't handle. Disable them so all
# tool_choice modes fall back to auto tool parsing via
# extract_tool_calls.
if getattr(request, "_grammar_from_tool_parser", False):
assert self._tool_parser is not None
self._tool_parser.supports_required_and_named = False
def parse(
self,
model_output: str,
request: ChatCompletionRequest | ResponsesRequest,
enable_auto_tools: bool = False,
) -> tuple[str | None, str | None, list[FunctionCall] | None]:
self._maybe_force_auto_tool_parsing(request)
reasoning, content, tool_calls = super().parse(
model_output, request, enable_auto_tools
)
if tool_calls:
from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
# Named/required tool_choice builds FunctionCalls without
# ID, backfill with Mistral-format IDs.
for tc in tool_calls:
if not tc.id:
tc.id = MistralToolCall.generate_random_id()
return reasoning, content, tool_calls
def parse_delta(
self,
delta_text: str,
delta_token_ids: list[int],
request: ChatCompletionRequest | ResponsesRequest,
prompt_token_ids: list[int] | None = None,
*,
finished: bool,
) -> DeltaMessage | None:
self._maybe_force_auto_tool_parsing(request)
return super().parse_delta(
delta_text,
delta_token_ids,
request,
prompt_token_ids,
finished=finished,
)
+9
View File
@@ -106,6 +106,15 @@ class ParserManager:
if reasoning_parser_cls is None and tool_parser_cls is None:
return None
from vllm.utils.mistral import is_mistral_tool_parser
if is_mistral_tool_parser(tool_parser_cls):
from vllm.parser.mistral import MistralParser
MistralParser.reasoning_parser_cls = reasoning_parser_cls
MistralParser.tool_parser_cls = tool_parser_cls
return MistralParser
from vllm.parser.abstract_parser import DelegatingParser
r_cls = reasoning_parser_cls
+10 -1
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
from collections.abc import Iterable, Sequence
from functools import cached_property
from typing import TYPE_CHECKING
@@ -76,6 +76,15 @@ class MistralReasoningParser(BaseThinkingReasoningParser):
has_eot_token = True
return False
def is_reasoning_end_streaming(
self, input_ids: Sequence[int], delta_ids: Iterable[int]
) -> bool:
if self.end_token_id in delta_ids:
return True
# Grammar's think? is optional — if [THINK] was never generated,
# reasoning was skipped entirely.
return self.start_token_id not in input_ids
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
"""
Extract the content
+2 -163
View File
@@ -5,11 +5,10 @@ from __future__ import annotations
import json
from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum, auto
from random import choices
from string import ascii_letters, digits
from typing import TYPE_CHECKING, Any
from typing import Any
import ijson
import regex as re
@@ -40,19 +39,14 @@ from vllm.entrypoints.openai.engine.protocol import (
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.reasoning.mistral_reasoning_parser import MistralReasoningParser
from vllm.sampling_params import StructuredOutputsParams
from vllm.tokenizers import TokenizerLike
from vllm.tokenizers.mistral import MistralTokenizer
from vllm.tool_parsers.abstract_tool_parser import (
Tool,
ToolParser,
)
from vllm.utils.mistral import is_mistral_tokenizer
if TYPE_CHECKING:
from vllm.reasoning import ReasoningParser
logger = init_logger(__name__)
ALPHANUMERIC = ascii_letters + digits
@@ -99,19 +93,6 @@ def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool:
return "[ARGS]" not in vocab
@dataclass
class MistralStreamingResult:
r"""Encapsulates the mutable state returned from
`MistralToolParser.extract_maybe_reasoning_and_tool_streaming`.
"""
delta_message: DeltaMessage | None
reasoning_ended: bool
tools_called: bool
current_text: str
current_token_ids: list[int]
class MistralToolParser(ToolParser):
r"""Tool call parser for Mistral models, intended for use with either:
@@ -281,148 +262,6 @@ class MistralToolParser(ToolParser):
request._grammar_from_tool_parser = True
return request
def extract_maybe_reasoning_and_tool_streaming(
self,
*,
reasoning_parser: ReasoningParser | None,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: list[int],
current_token_ids: list[int],
output_token_ids: Sequence[int],
reasoning_ended: bool,
prompt_is_reasoning_end: bool | None,
request: ChatCompletionRequest,
) -> MistralStreamingResult:
r"""Streaming extraction with reasoning followed by tool-call parsing.
This method encapsulates the combined reasoning extraction and
tool-call streaming logic so that the serving layer only needs a
thin routing branch.
The flow is:
1. If a *reasoning_parser* is present and reasoning has **not** ended,
extract reasoning tokens. Pre-v15 models may have pre-filled
`[THINK]...[/THINK]` in system prompts, so we skip the
prompt-level reasoning-end check for those.
2. Once reasoning ends (or if there is no reasoning parser), delegate
to `extract_tool_calls_streaming` and track whether tools were
called.
Args:
reasoning_parser: Optional reasoning parser instance.
previous_text: Accumulated text from prior chunks.
current_text: Full accumulated text including current chunk.
delta_text: New text in this chunk.
previous_token_ids: Token ids from prior chunks.
current_token_ids: Full token ids including current chunk.
output_token_ids: Raw output token ids from the engine.
reasoning_ended: Whether reasoning has already ended.
prompt_is_reasoning_end: Whether the prompt itself ends reasoning.
request: The originating chat completion request.
"""
delta_message: DeltaMessage | None = None
tools_called = False
reasoning_ended_at_entry = reasoning_ended
# For MistralReasoningParser, only enter the reasoning block when
# the model has actually emitted a [THINK] token. Other reasoning
# parsers always expect thinking to be present.
expect_thinking = (
not isinstance(reasoning_parser, MistralReasoningParser)
or reasoning_parser.start_token_id in current_token_ids
)
if reasoning_parser is not None and not reasoning_ended and expect_thinking:
# Pre-v15 models may have pre-filled [THINK]...[/THINK] in
# system prompts, so skip the prompt-level reasoning-end
# check and wait for the output's own end-of-think.
is_pre_v15 = (
isinstance(self.model_tokenizer, MistralTokenizer)
and self.model_tokenizer.version < 15
)
if not is_pre_v15 and prompt_is_reasoning_end:
reasoning_ended = True
current_token_ids = list(output_token_ids)
else:
delta_message = reasoning_parser.extract_reasoning_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
output_token_ids,
)
if reasoning_parser.is_reasoning_end_streaming(
current_token_ids, output_token_ids
):
reasoning_ended = True
current_token_ids = reasoning_parser.extract_content_ids(
list(output_token_ids)
)
if delta_message and delta_message.content:
current_text = delta_message.content
delta_message.content = None
else:
current_text = ""
if not reasoning_ended:
return MistralStreamingResult(
delta_message=delta_message,
reasoning_ended=False,
tools_called=False,
current_text=current_text,
current_token_ids=current_token_ids,
)
delta_token_ids = list(output_token_ids)
# On the iteration where reasoning just ended, reset the text/token
# state so the tool parser sees a clean history instead of the
# accumulated reasoning text.
if not reasoning_ended_at_entry and reasoning_ended:
previous_text = ""
previous_token_ids = []
delta_text = current_text
delta_token_ids = current_token_ids
delta_message = self.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=previous_token_ids,
current_token_ids=current_token_ids,
delta_token_ids=delta_token_ids,
request=request,
)
if delta_message and delta_message.tool_calls:
tools_called = True
return MistralStreamingResult(
delta_message=delta_message,
reasoning_ended=reasoning_ended,
tools_called=tools_called,
current_text=current_text,
current_token_ids=current_token_ids,
)
@staticmethod
def build_non_streaming_tool_calls(
tool_calls: list[FunctionCall] | None,
) -> list[ToolCall]:
r"""Build `MistralToolCall` items for non-streaming responses."""
if not tool_calls:
return []
return [
MistralToolCall(id=tc.id, function=tc)
if tc.id
else MistralToolCall(function=tc)
for tc in tool_calls
]
def extract_tool_calls(
self,
model_output: str,
@@ -536,7 +375,7 @@ class MistralToolParser(ToolParser):
return ExtractedToolCallInformation(
tools_called=True,
tool_calls=mistral_tool_calls,
content=content if len(content) > 0 else None,
content=content if content.strip() else None,
)
def extract_tool_calls_streaming(
+310
View File
@@ -0,0 +1,310 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib
from collections.abc import Sequence
from typing import Any
from openai.types.responses.function_tool import FunctionTool
from vllm.entrypoints.chat_utils import make_tool_call_id
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
ChatCompletionToolsParam,
)
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
logger = init_logger(__name__)
def _rust_tool_parser_module() -> Any:
try:
return importlib.import_module("vllm._rust_tool_parser")
except ImportError as exc:
raise RuntimeError(
"Rust tool parsing requires the vllm._rust_tool_parser PyO3 "
"extension. Rebuild vLLM with Rust frontend/extensions enabled."
) from exc
class RustToolParser(ToolParser):
"""Adapter from an opaque Rust parser to the vLLM ToolParser API.
Subclasses provide only model-specific configuration: the exact Rust parser
name and an optional tool-call start marker for fast complete-output
rejection.
This class keeps the vLLM-specific bridge work:
- convert vLLM tool definitions into the Rust ``Tool`` shape;
- translate typed Rust parser outputs into vLLM protocol objects; and
- maintain vLLM streaming bookkeeping used by finish-reason handling.
The parser grammar and incremental parser state stay in Rust.
"""
# Rust-backed parsers are opaque to Python by default. Do not use vLLM's
# standard JSON required/named handling; let the Rust parser consume the
# model's native tool-call syntax.
supports_required_and_named = False
rust_parser_name: str
tool_call_start_token: str | None = None
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
self._parser: Any | None = None
self._error: Exception | None = None
self._tool_call_ids: dict[int, str] = {}
if not self.model_tokenizer:
raise ValueError(
"The model tokenizer must be passed to the ToolParser "
"constructor during construction."
)
logger.debug(
"vLLM successfully imported tool parser %s", self.__class__.__name__
)
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
"""Adjust request options without installing Python-side constraints.
Rust-backed parsers are treated as source-of-truth opaque parsers. The
bridge intentionally avoids ``super().adjust_request()`` so Python does
not install JSON schema guidance or structural-tag constraints that may
conflict with the Rust parser's native grammar.
"""
if self._get_parser().preserve_special_tokens():
request.skip_special_tokens = False
return request
def _rust_tools(self) -> list[Any]:
"""Build Rust ``Tool`` objects from vLLM tool definitions."""
if not self.tools:
return []
tools: list[Any] = []
for tool in self.tools:
if isinstance(tool, FunctionTool):
name = tool.name
description = tool.description
parameters = tool.parameters or {}
strict = getattr(tool, "strict", None)
elif isinstance(tool, ChatCompletionToolsParam):
name = tool.function.name
description = tool.function.description
parameters = tool.function.parameters or {}
strict = getattr(tool.function, "strict", None)
else:
continue
tools.append(
_rust_tool_parser_module().Tool(name, description, parameters, strict)
)
return tools
def _new_parser(self) -> Any:
"""Create a fresh Rust parser with the current tool schemas."""
return _rust_tool_parser_module().ToolParser(
self.rust_parser_name, self._rust_tools()
)
def _get_parser(self) -> Any:
if self._parser is None:
self._parser = self._new_parser()
return self._parser
def _reset_streaming_state(self) -> None:
"""Reset parser state for a new request on a reused parser instance."""
self._parser = self._new_parser()
self._error = None
self._tool_call_ids.clear()
self.prev_tool_call_arr.clear()
self.streamed_args_for_tool.clear()
self.current_tool_id = -1
self.current_tool_name_sent = False
def _ensure_tool_state(self, index: int) -> None:
"""Grow vLLM streaming state arrays to contain ``index``."""
while len(self.prev_tool_call_arr) <= index:
self.prev_tool_call_arr.append({})
while len(self.streamed_args_for_tool) <= index:
self.streamed_args_for_tool.append("")
def _record_delta(
self, index: int, name: str | None, arguments: str | None
) -> str | None:
"""Mirror a Rust parser delta into vLLM streaming bookkeeping.
``prev_tool_call_arr`` and ``streamed_args_for_tool`` are read later by
the chat serving layer to decide the final ``tool_calls`` finish reason
and to flush any remaining argument bytes.
"""
tool_call_id = None
self._ensure_tool_state(index)
if name is not None:
tool_call_id = make_tool_call_id()
self._tool_call_ids[index] = tool_call_id
self.prev_tool_call_arr[index] = {"name": name, "arguments": {}}
self.current_tool_name_sent = True
if arguments is not None:
self.streamed_args_for_tool[index] += arguments
self.prev_tool_call_arr[index]["arguments"] = self.streamed_args_for_tool[
index
]
self.current_tool_id = index
return tool_call_id
def _delta_message_from_parser_output(
self, parser_output: Any | None
) -> DeltaMessage | None:
"""Translate one Rust parser output into a vLLM ``DeltaMessage``."""
if parser_output is None:
return None
normal_text = parser_output.normal_text or None
tool_calls: list[DeltaToolCall] = []
for tool_call in parser_output.calls:
index = tool_call.tool_index
name = tool_call.name
arguments: str | None = tool_call.arguments
if name is None and arguments is None:
continue
tool_call_id = self._record_delta(index, name, arguments)
tool_calls.append(
DeltaToolCall(
index=index,
id=tool_call_id,
type="function" if name is not None else None,
function=DeltaFunctionCall(
name=name,
arguments=arguments,
),
)
)
if normal_text is None and not tool_calls:
return None
return DeltaMessage(content=normal_text, tool_calls=tool_calls)
def _parse_complete(self, model_output: str) -> Any | None:
"""Parse complete model output with a throwaway Rust parser instance."""
parser = self._new_parser()
output = _rust_tool_parser_module().ToolParserOutput()
try:
parser.parse_into(model_output, output)
output.append(parser.finish())
except Exception:
logger.exception(
"Error parsing %s tool call output.", self.rust_parser_name
)
return None
return output.coalesce_calls()
def extract_tool_calls(
self,
model_output: str,
request: ChatCompletionRequest,
) -> ExtractedToolCallInformation:
"""Extract tool calls from complete model output (non-streaming)."""
if (
self.tool_call_start_token is not None
and self.tool_call_start_token not in model_output
):
return ExtractedToolCallInformation(
tools_called=False,
tool_calls=[],
content=model_output,
)
parsed = self._parse_complete(model_output)
if parsed is None:
return ExtractedToolCallInformation(
tools_called=False,
tool_calls=[],
content=model_output,
)
tool_calls: list[ToolCall] = []
self.prev_tool_call_arr.clear()
for parsed_tool_call in parsed.calls:
name = parsed_tool_call.name
arguments = parsed_tool_call.arguments or "{}"
if name is None:
continue
tool_calls.append(
ToolCall(
type="function",
function=FunctionCall(name=name, arguments=arguments),
)
)
self.prev_tool_call_arr.append({"name": name, "arguments": arguments})
if not tool_calls:
return ExtractedToolCallInformation(
tools_called=False,
tool_calls=[],
content=model_output,
)
content = parsed.normal_text or None
return ExtractedToolCallInformation(
tools_called=True,
tool_calls=tool_calls,
content=content,
)
def extract_tool_calls_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int], # pylint: disable=unused-argument
current_token_ids: Sequence[int], # pylint: disable=unused-argument
delta_token_ids: Sequence[int], # pylint: disable=unused-argument
request: ChatCompletionRequest, # pylint: disable=unused-argument
) -> DeltaMessage | None:
"""Extract tool calls from streaming model output.
The Rust parser owns the incremental buffer, so this adapter feeds only
the newest text delta and lets the serving layer handle final empty
chunks.
"""
# TODO: Add a final-chunk hook if streaming needs to call Rust finish().
if not previous_text:
self._reset_streaming_state()
if self._error is not None:
return None
parser_output = _rust_tool_parser_module().ToolParserOutput()
try:
self._get_parser().parse_into(delta_text, parser_output)
except Exception as error:
self._error = error
logger.exception(
"Error parsing %s streaming tool call output.",
self.rust_parser_name,
)
delta_message = self._delta_message_from_parser_output(parser_output)
if delta_message is not None:
return delta_message
return None

Some files were not shown because too many files have changed in this diff Show More