forked from Karylab-cklius/vllm
Compare commits
50
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
efa494e397 | ||
|
|
74b77a593a | ||
|
|
ec53889a3b | ||
|
|
ede3d4ddf6 | ||
|
|
7301a834fa | ||
|
|
7104ac6b5d | ||
|
|
9a98eccb49 | ||
|
|
2ce1e2ba71 | ||
|
|
2496f66f7f | ||
|
|
d18b7f2723 | ||
|
|
80b633f076 | ||
|
|
28f589ebb1 | ||
|
|
98998279d4 | ||
|
|
71d208089e | ||
|
|
0ac6fc8352 | ||
|
|
fdfb2566c0 | ||
|
|
8a5cf1ccd6 | ||
|
|
fe1d923afc | ||
|
|
32daf56b42 | ||
|
|
82a42234be | ||
|
|
af9f583344 | ||
|
|
bd2d83ff31 | ||
|
|
bb78168b21 | ||
|
|
89c6a41001 | ||
|
|
7fdfa6441d | ||
|
|
e9b728de8a | ||
|
|
5828a205ef | ||
|
|
7a74f31d2e | ||
|
|
47930b59ca | ||
|
|
6aec99f030 | ||
|
|
f4966f8b3d | ||
|
|
2c9c07c85e | ||
|
|
320c52b134 | ||
|
|
6deb05e0e4 | ||
|
|
d82ac00923 | ||
|
|
dac9e9a640 | ||
|
|
d7607ad273 | ||
|
|
d955745d58 | ||
|
|
e1ed89dbee | ||
|
|
1c2ffc6f88 | ||
|
|
ca4cfd8731 | ||
|
|
c9c1540e61 | ||
|
|
c1d754d681 | ||
|
|
01d8cd92dd | ||
|
|
a4b14b98c6 | ||
|
|
cf1c906724 | ||
|
|
766ce2bb6b | ||
|
|
3d119f78f7 | ||
|
|
1b1359c332 | ||
|
|
cad4ca12b8 |
@@ -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"
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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 |
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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).
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
Generated
+86
@@ -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"
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
@@ -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#"<|DSML|tool_calls>
|
||||
<|DSML|invoke name="create_order">
|
||||
<|DSML|parameter name="user_id" string="false">42</|DSML|parameter>
|
||||
<|DSML|parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</|DSML|parameter>
|
||||
</|DSML|invoke>
|
||||
</|DSML|tool_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()),
|
||||
¶meters,
|
||||
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();
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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],
|
||||
)
|
||||
|
||||
@@ -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 = "<|DSML|tool_calls>"
|
||||
TC_END = "</|DSML|tool_calls>"
|
||||
INV_START = '<|DSML|invoke name="'
|
||||
INV_END = "</|DSML|invoke>"
|
||||
PARAM_START = '<|DSML|parameter name="'
|
||||
PARAM_END = "</|DSML|parameter>"
|
||||
|
||||
|
||||
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
@@ -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
|
||||
|
||||
@@ -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 #########
|
||||
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
+3
-3
@@ -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
|
||||
|
||||
@@ -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 = ""):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = ""):
|
||||
|
||||
@@ -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, ...] = ()
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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": [
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user