forked from Karylab-cklius/vllm
Merge branch 'main' into wentao-fix-dcp-IMA-for-v2
This commit is contained in:
+57
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euxo pipefail
|
||||
|
||||
# Nightly e2e test for prefetch offloading with a MoE model.
|
||||
# Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights
|
||||
# and validates GSM8K accuracy matches baseline (no offloading).
|
||||
#
|
||||
# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
|
||||
THRESHOLD=${1:-0.25}
|
||||
NUM_Q=${2:-1319}
|
||||
PORT=${3:-8030}
|
||||
OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
|
||||
mkdir -p "${OUT_DIR}"
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
timeout 600 bash -c '
|
||||
until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
|
||||
sleep 1
|
||||
done'
|
||||
}
|
||||
|
||||
MODEL="deepseek-ai/DeepSeek-V2-Lite"
|
||||
|
||||
cleanup() {
|
||||
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
|
||||
kill "${SERVER_PID}" 2>/dev/null || true
|
||||
for _ in {1..20}; do
|
||||
kill -0 "${SERVER_PID}" 2>/dev/null || break
|
||||
sleep 0.5
|
||||
done
|
||||
kill -9 "${SERVER_PID}" 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
vllm serve "$MODEL" \
|
||||
--max-model-len 2048 \
|
||||
--offload-group-size 8 \
|
||||
--offload-num-in-group 2 \
|
||||
--offload-prefetch-step 1 \
|
||||
--offload-params w13_weight w2_weight \
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_prefetch_offload.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} prefetch_offload: accuracy {acc:.3f}")
|
||||
assert acc >= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}"
|
||||
PY
|
||||
|
||||
cleanup
|
||||
SERVER_PID=
|
||||
@@ -28,3 +28,12 @@ steps:
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
|
||||
|
||||
- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030
|
||||
|
||||
@@ -30,7 +30,7 @@ steps:
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_8
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
|
||||
@@ -24,6 +24,11 @@ steps:
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server 1)
|
||||
timeout_in_minutes: 130
|
||||
|
||||
@@ -73,3 +73,29 @@ steps:
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
|
||||
|
||||
- label: GPQA Eval (GPT-OSS) (H100)
|
||||
timeout_in_minutes: 120
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
- tests/evals/gpt_oss/
|
||||
commands:
|
||||
- uv pip install --system 'gpt-oss[eval]==0.0.5'
|
||||
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
|
||||
|
||||
- label: GPQA Eval (GPT-OSS) (B200)
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
- tests/evals/gpt_oss/
|
||||
commands:
|
||||
- uv pip install --system 'gpt-oss[eval]==0.0.5'
|
||||
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
|
||||
|
||||
@@ -153,33 +153,6 @@ steps:
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
- label: GPT-OSS Eval (H100)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- tests/evals/gpt_oss
|
||||
- vllm/model_executor/models/gpt_oss.py
|
||||
- vllm/model_executor/layers/quantization/mxfp4.py
|
||||
commands:
|
||||
- uv pip install --system 'gpt-oss[eval]==0.0.5'
|
||||
- pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
|
||||
|
||||
- label: GPT-OSS Eval (B200)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- tests/evals/gpt_oss
|
||||
- vllm/model_executor/models/gpt_oss.py
|
||||
- vllm/model_executor/layers/quantization/mxfp4.py
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
commands:
|
||||
- uv pip install --system 'gpt-oss[eval]==0.0.5'
|
||||
- pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
|
||||
|
||||
- label: Batch Invariance (H100)
|
||||
timeout_in_minutes: 25
|
||||
device: h100
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
# doc: https://github.com/pytorch/test-infra/blob/main/tools/stronghold/docs/bc_linter_config.md
|
||||
version: 1
|
||||
paths:
|
||||
# We temporarily disable globally, and will only enable with `annotations.include`
|
||||
# include:
|
||||
# - "vllm/v1/attetion/*.py"
|
||||
# - "vllm/v1/core/*.py"
|
||||
exclude:
|
||||
- "**/*.py"
|
||||
|
||||
scan:
|
||||
functions: true # check free functions and methods
|
||||
classes: true # check classes/dataclasses
|
||||
public_only: true # ignore names starting with "_" at any level
|
||||
|
||||
annotations:
|
||||
include: # decorators that force‑include a symbol
|
||||
- name: "bc_linter_include" # matched by simple name or dotted suffix
|
||||
propagate_to_members: false # for classes, include methods/inner classes
|
||||
exclude: # decorators that force‑exclude a symbol
|
||||
- name: "bc_linter_skip" # matched by simple name or dotted suffix
|
||||
propagate_to_members: true # for classes, exclude methods/inner classes
|
||||
|
||||
excluded_violations: [] # e.g. ["ParameterRenamed", "FieldTypeChanged"]
|
||||
@@ -1,29 +0,0 @@
|
||||
name: BC Lint
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types:
|
||||
- opened
|
||||
- synchronize
|
||||
- reopened
|
||||
- labeled
|
||||
- unlabeled
|
||||
|
||||
jobs:
|
||||
bc_lint:
|
||||
if: github.repository_owner == 'vllm-project'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Run BC Lint Action
|
||||
uses: pytorch/test-infra/.github/actions/bc-lint@main
|
||||
with:
|
||||
repo: ${{ github.event.pull_request.head.repo.full_name }}
|
||||
base_sha: ${{ github.event.pull_request.base.sha }}
|
||||
head_sha: ${{ github.event.pull_request.head.sha }}
|
||||
suppression: ${{ contains(github.event.pull_request.labels.*.name, 'suppress-bc-linter') }}
|
||||
docs_link: 'https://github.com/pytorch/test-infra/wiki/BC-Linter'
|
||||
config_dir: .github
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: true
|
||||
@@ -649,9 +649,3 @@ ASYNC_REQUEST_FUNCS = {
|
||||
"sglang": async_request_openai_completions,
|
||||
"llama.cpp": async_request_openai_completions,
|
||||
}
|
||||
|
||||
OPENAI_COMPATIBLE_BACKENDS = [
|
||||
k
|
||||
for k, v in ASYNC_REQUEST_FUNCS.items()
|
||||
if v in (async_request_openai_completions, async_request_openai_chat_completions)
|
||||
]
|
||||
|
||||
@@ -1,78 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from types import TracebackType
|
||||
from typing import Any
|
||||
|
||||
|
||||
def convert_to_pytorch_benchmark_format(
|
||||
args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any]
|
||||
) -> list:
|
||||
"""
|
||||
Save the benchmark results in the format used by PyTorch OSS benchmark with
|
||||
on metric per record
|
||||
https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database
|
||||
"""
|
||||
records = []
|
||||
if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False):
|
||||
return records
|
||||
|
||||
for name, benchmark_values in metrics.items():
|
||||
record = {
|
||||
"benchmark": {
|
||||
"name": "vLLM benchmark",
|
||||
"extra_info": {
|
||||
"args": vars(args),
|
||||
},
|
||||
},
|
||||
"model": {
|
||||
"name": args.model,
|
||||
},
|
||||
"metric": {
|
||||
"name": name,
|
||||
"benchmark_values": benchmark_values,
|
||||
"extra_info": extra_info,
|
||||
},
|
||||
}
|
||||
|
||||
tp = record["benchmark"]["extra_info"]["args"].get("tensor_parallel_size")
|
||||
# Save tensor_parallel_size parameter if it's part of the metadata
|
||||
if not tp and "tensor_parallel_size" in extra_info:
|
||||
record["benchmark"]["extra_info"]["args"]["tensor_parallel_size"] = (
|
||||
extra_info["tensor_parallel_size"]
|
||||
)
|
||||
|
||||
records.append(record)
|
||||
|
||||
return records
|
||||
|
||||
|
||||
class InfEncoder(json.JSONEncoder):
|
||||
def clear_inf(self, o: Any):
|
||||
if isinstance(o, dict):
|
||||
return {k: self.clear_inf(v) for k, v in o.items()}
|
||||
elif isinstance(o, list):
|
||||
return [self.clear_inf(v) for v in o]
|
||||
elif isinstance(o, float) and math.isinf(o):
|
||||
return "inf"
|
||||
return o
|
||||
|
||||
def iterencode(self, o: Any, *args, **kwargs) -> Any:
|
||||
return super().iterencode(self.clear_inf(o), *args, **kwargs)
|
||||
|
||||
|
||||
def write_to_json(filename: str, records: list) -> None:
|
||||
with open(filename, "w") as f:
|
||||
json.dump(
|
||||
records,
|
||||
f,
|
||||
cls=InfEncoder,
|
||||
default=lambda o: f"<{type(o).__name__} object is not JSON serializable>",
|
||||
)
|
||||
|
||||
|
||||
# Collect time and generate time metrics
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Cutlass bench utils
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
|
||||
@@ -86,15 +85,3 @@ def make_rand_sparse_tensors(
|
||||
|
||||
# Compressed B, Metadata, Original A, B
|
||||
return b_compressed, e, a, b
|
||||
|
||||
|
||||
def make_n_rand_sparse_tensors(
|
||||
num_tensors: int, dtype: torch.dtype, m: int, n: int, k: int
|
||||
) -> tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]:
|
||||
ABs = []
|
||||
for _ in range(num_tensors):
|
||||
b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)
|
||||
if b_comp is not None:
|
||||
ABs.append(make_rand_sparse_tensors(dtype, m, n, k))
|
||||
BComps, Es, As, Bs = zip(*ABs)
|
||||
return list(BComps), list(Es), list(As), list(Bs)
|
||||
|
||||
@@ -1,45 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""Token bucket rate limiter implementation"""
|
||||
|
||||
def __init__(self, rate_limit):
|
||||
self.rate_limit = rate_limit # Requests per second
|
||||
self.num_available_tokens = rate_limit # Available tokens
|
||||
self.last_refill = time.monotonic() # Last token refill time
|
||||
self.lock = asyncio.Lock() # Synchronization lock
|
||||
|
||||
async def acquire(self):
|
||||
"""Acquire a token from the rate limiter"""
|
||||
while True:
|
||||
async with self.lock:
|
||||
current_time = time.monotonic()
|
||||
elapsed = current_time - self.last_refill
|
||||
|
||||
# Refill num_available_tokens if more than 1 second has passed
|
||||
if elapsed > 1.0:
|
||||
self.num_available_tokens = self.rate_limit
|
||||
self.last_refill = current_time
|
||||
|
||||
# Check if num_available_tokens are available
|
||||
if self.num_available_tokens > 0:
|
||||
self.num_available_tokens -= 1
|
||||
return True
|
||||
|
||||
# Calculate wait time if no num_available_tokens available
|
||||
wait_time = 1.0 - elapsed
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
async def __aenter__(self):
|
||||
"""Enter async context manager - acquire token"""
|
||||
await self.acquire()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_value, traceback):
|
||||
"""Exit async context manager - no cleanup needed"""
|
||||
pass
|
||||
@@ -1,39 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
from collections import deque
|
||||
|
||||
|
||||
class RequestQueue:
|
||||
"""Request queue manager with concurrency control"""
|
||||
|
||||
def __init__(self, max_concurrent, max_queue_size):
|
||||
# Maximum concurrent requests
|
||||
self.max_concurrent = max_concurrent
|
||||
self.max_queue_size = max_queue_size # Maximum queue size
|
||||
# Concurrency control
|
||||
self.semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self.queue = deque() # Request queue
|
||||
self.queue_size = 0 # Current queue size
|
||||
self.lock = asyncio.Lock() # Sync queue Lock
|
||||
|
||||
async def enqueue(self, task):
|
||||
"""Add a request task to the queue"""
|
||||
async with self.lock:
|
||||
if self.queue_size >= self.max_queue_size:
|
||||
return False
|
||||
|
||||
self.queue.append(task)
|
||||
self.queue_size += 1
|
||||
return True
|
||||
|
||||
async def process(self):
|
||||
"""Process queued requests using semaphore for concurrency control"""
|
||||
while True:
|
||||
if self.queue:
|
||||
async with self.semaphore, self.lock:
|
||||
task = self.queue.popleft()
|
||||
self.queue_size -= 1
|
||||
await task
|
||||
await asyncio.sleep(0.01) # Yield control to event loop
|
||||
@@ -30,6 +30,9 @@ import torch.distributed as dist
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
|
||||
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
|
||||
FlashInferAllReduce,
|
||||
)
|
||||
from vllm.distributed.device_communicators.pynccl import (
|
||||
PyNcclCommunicator,
|
||||
register_nccl_symmetric_ops,
|
||||
@@ -44,7 +47,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Default sequence lengths to benchmark
|
||||
DEFAULT_SEQUENCE_LENGTHS = [128, 512, 1024, 2048, 4096, 8192]
|
||||
DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192]
|
||||
|
||||
# Fixed hidden size and dtype for all benchmarks
|
||||
HIDDEN_SIZE = 8192
|
||||
@@ -81,6 +84,7 @@ class CommunicatorBenchmark:
|
||||
self.symm_mem_comm = None
|
||||
self.symm_mem_comm_multimem = None
|
||||
self.symm_mem_comm_two_shot = None
|
||||
self.fi_ar_comm = None
|
||||
|
||||
self._init_communicators()
|
||||
|
||||
@@ -161,6 +165,22 @@ class CommunicatorBenchmark:
|
||||
)
|
||||
self.symm_mem_comm_two_shot = None
|
||||
|
||||
try:
|
||||
self.fi_ar_comm = FlashInferAllReduce(
|
||||
group=self.cpu_group,
|
||||
device=self.device,
|
||||
)
|
||||
if not self.fi_ar_comm.disabled:
|
||||
logger.info("Rank %s: FlashInferAllReduce initialized", self.rank)
|
||||
else:
|
||||
logger.info("Rank %s: FlashInferAllReduce disabled", self.rank)
|
||||
self.fi_ar_comm = None
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Rank %s: Failed to initialize FlashInferAllReduce: %s", self.rank, e
|
||||
)
|
||||
self.fi_ar_comm = None
|
||||
|
||||
def benchmark_allreduce(
|
||||
self, sequence_length: int, num_warmup: int, num_trials: int
|
||||
) -> dict[str, float]:
|
||||
@@ -180,7 +200,8 @@ class CommunicatorBenchmark:
|
||||
lambda t, c=comm: c.custom_all_reduce(t),
|
||||
lambda t, c=comm: c.should_custom_ar(t),
|
||||
comm.capture(),
|
||||
"1stage", # env variable value
|
||||
{"VLLM_CUSTOM_ALLREDUCE_ALGO": "1stage"},
|
||||
None, # no destroy function
|
||||
)
|
||||
)
|
||||
# CustomAllreduce two-shot
|
||||
@@ -190,7 +211,8 @@ class CommunicatorBenchmark:
|
||||
lambda t, c=comm: c.custom_all_reduce(t),
|
||||
lambda t, c=comm: c.should_custom_ar(t),
|
||||
comm.capture(),
|
||||
"2stage", # env variable value
|
||||
{"VLLM_CUSTOM_ALLREDUCE_ALGO": "2stage"},
|
||||
None, # no destroy function
|
||||
)
|
||||
)
|
||||
|
||||
@@ -202,7 +224,8 @@ class CommunicatorBenchmark:
|
||||
lambda t, c=comm: c.all_reduce(t),
|
||||
lambda t: True, # Always available if initialized
|
||||
nullcontext(),
|
||||
None, # no env variable needed
|
||||
{}, # no env variable needed
|
||||
None, # no destroy function
|
||||
)
|
||||
)
|
||||
communicators.append(
|
||||
@@ -211,7 +234,8 @@ class CommunicatorBenchmark:
|
||||
lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t),
|
||||
lambda t: True, # Always available if initialized
|
||||
nullcontext(),
|
||||
None, # no env variable needed
|
||||
{}, # no env variable needed
|
||||
None, # no destroy function
|
||||
)
|
||||
)
|
||||
|
||||
@@ -223,7 +247,8 @@ class CommunicatorBenchmark:
|
||||
lambda t, c=comm: c.all_reduce(t),
|
||||
lambda t, c=comm: c.should_use_symm_mem(t),
|
||||
nullcontext(),
|
||||
None, # no env variable needed
|
||||
{}, # no env variable needed
|
||||
None, # no destroy function
|
||||
)
|
||||
)
|
||||
|
||||
@@ -235,29 +260,67 @@ class CommunicatorBenchmark:
|
||||
lambda t, c=comm: c.all_reduce(t),
|
||||
lambda t, c=comm: c.should_use_symm_mem(t),
|
||||
nullcontext(),
|
||||
None, # no env variable needed
|
||||
{}, # no env variable needed
|
||||
None, # no destroy function needed
|
||||
)
|
||||
)
|
||||
|
||||
if self.fi_ar_comm is not None:
|
||||
comm = self.fi_ar_comm
|
||||
communicators.append(
|
||||
(
|
||||
"flashinfer_trtllm",
|
||||
lambda t, c=comm: c.all_reduce(t),
|
||||
lambda t, c=comm: c.should_use_fi_ar(t),
|
||||
nullcontext(),
|
||||
{"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm"},
|
||||
lambda c=comm: c.destroy(),
|
||||
)
|
||||
)
|
||||
communicators.append(
|
||||
(
|
||||
"flashinfer_mnnvl",
|
||||
lambda t, c=comm: c.all_reduce(t),
|
||||
lambda t, c=comm: c.should_use_fi_ar(t),
|
||||
nullcontext(),
|
||||
{"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "mnnvl"},
|
||||
lambda c=comm: c.destroy(),
|
||||
)
|
||||
)
|
||||
|
||||
# Benchmark each communicator
|
||||
for name, allreduce_fn, should_use_fn, context, env_var in communicators:
|
||||
# Set environment variable if needed
|
||||
if env_var is not None:
|
||||
os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = env_var
|
||||
else:
|
||||
# Clear the environment variable to avoid interference
|
||||
os.environ.pop("VLLM_CUSTOM_ALLREDUCE_ALGO", None)
|
||||
|
||||
latency = self.benchmark_allreduce_single(
|
||||
sequence_length,
|
||||
allreduce_fn,
|
||||
should_use_fn,
|
||||
context,
|
||||
num_warmup,
|
||||
num_trials,
|
||||
)
|
||||
if latency is not None:
|
||||
results[name] = latency
|
||||
for (
|
||||
name,
|
||||
allreduce_fn,
|
||||
should_use_fn,
|
||||
context,
|
||||
env_dict,
|
||||
destroy_fn,
|
||||
) in communicators:
|
||||
# Save original values and apply new environment variables
|
||||
saved_env = {key: os.environ.get(key) for key in env_dict}
|
||||
for key, value in env_dict.items():
|
||||
os.environ[key] = value
|
||||
try:
|
||||
latency = self.benchmark_allreduce_single(
|
||||
sequence_length,
|
||||
allreduce_fn,
|
||||
should_use_fn,
|
||||
context,
|
||||
num_warmup,
|
||||
num_trials,
|
||||
)
|
||||
if latency is not None:
|
||||
results[name] = latency
|
||||
finally:
|
||||
if destroy_fn is not None:
|
||||
destroy_fn()
|
||||
# Restore environment variables to their original state
|
||||
for key, original_value in saved_env.items():
|
||||
if original_value is None:
|
||||
os.environ.pop(key, None)
|
||||
else:
|
||||
os.environ[key] = original_value
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@@ -5,8 +5,11 @@
|
||||
Benchmark for FlashInfer fused collective operations vs standard operations.
|
||||
|
||||
This benchmark compares:
|
||||
1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
|
||||
2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
|
||||
1. FlashInfer's allreduce_fusion with trtllm backend
|
||||
(fused allreduce + rmsnorm + optional FP8/FP4 quant)
|
||||
2. FlashInfer's allreduce_fusion with mnnvl backend
|
||||
(fused allreduce + rmsnorm only, no quantization support)
|
||||
3. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
|
||||
|
||||
Usage with torchrun:
|
||||
torchrun --nproc_per_node=2 benchmark_fused_collective.py
|
||||
@@ -48,8 +51,12 @@ SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Try to import FlashInfer
|
||||
TorchDistBackend = None
|
||||
try:
|
||||
import flashinfer.comm as flashinfer_comm # type: ignore
|
||||
from flashinfer.comm.mnnvl import ( # type: ignore
|
||||
TorchDistBackend,
|
||||
)
|
||||
|
||||
if not (
|
||||
hasattr(flashinfer_comm, "allreduce_fusion")
|
||||
@@ -74,11 +81,15 @@ _FI_MAX_SIZES = {
|
||||
8: 64 * MiB, # 64MB
|
||||
}
|
||||
|
||||
# Global workspace tensor for FlashInfer
|
||||
_FI_WORKSPACE = None
|
||||
# Global workspace tensors for FlashInfer (keyed by backend name)
|
||||
_FI_WORKSPACES: dict = {}
|
||||
|
||||
# Backends to benchmark
|
||||
FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
|
||||
|
||||
|
||||
def setup_flashinfer_workspace(
|
||||
backend: str,
|
||||
world_size: int,
|
||||
rank: int,
|
||||
hidden_dim: int,
|
||||
@@ -86,41 +97,54 @@ def setup_flashinfer_workspace(
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
"""Setup FlashInfer workspace for fused allreduce operations."""
|
||||
global _FI_WORKSPACE
|
||||
global FI_WORKSPACES
|
||||
|
||||
if flashinfer_comm is None:
|
||||
return None, None
|
||||
return None
|
||||
|
||||
if world_size not in _FI_MAX_SIZES:
|
||||
logger.warning("FlashInfer not supported for world size %s", world_size)
|
||||
return None, None
|
||||
return None
|
||||
|
||||
try:
|
||||
kwargs = {}
|
||||
if TorchDistBackend is not None:
|
||||
kwargs["comm_backend"] = TorchDistBackend(group=dist.group.WORLD)
|
||||
|
||||
workspace = flashinfer_comm.create_allreduce_fusion_workspace(
|
||||
backend="trtllm",
|
||||
backend=backend,
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
_FI_WORKSPACE = workspace
|
||||
_FI_WORKSPACES[backend] = workspace
|
||||
return workspace
|
||||
except Exception as e:
|
||||
logger.error("Failed to setup FlashInfer workspace: %s", e)
|
||||
logger.error(
|
||||
"Failed to setup FlashInfer workspace (backend=%s): %s", backend, e
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def cleanup_flashinfer_workspace(workspace):
|
||||
"""Cleanup FlashInfer workspace."""
|
||||
if flashinfer_comm is None or workspace is None:
|
||||
def cleanup_flashinfer_workspaces():
|
||||
"""Cleanup all FlashInfer workspaces."""
|
||||
if flashinfer_comm is None:
|
||||
return
|
||||
|
||||
try:
|
||||
workspace.destroy()
|
||||
except Exception as e:
|
||||
logger.error("Failed to cleanup FlashInfer workspace: %s", e)
|
||||
for backend, workspace in _FI_WORKSPACES.items():
|
||||
try:
|
||||
workspace.destroy()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to cleanup FlashInfer workspace (backend=%s): %s",
|
||||
backend,
|
||||
e,
|
||||
)
|
||||
_FI_WORKSPACES.clear()
|
||||
|
||||
|
||||
class FlashInferFusedAllReduceParams:
|
||||
@@ -134,7 +158,7 @@ class FlashInferFusedAllReduceParams:
|
||||
self.fp32_acc = True
|
||||
self.max_token_num = max_token_num
|
||||
|
||||
def get_trtllm_fused_allreduce_kwargs(self):
|
||||
def get_flashinfer_fused_allreduce_kwargs(self):
|
||||
return {
|
||||
"launch_with_pdl": self.launch_with_pdl,
|
||||
"fp32_acc": self.fp32_acc,
|
||||
@@ -147,11 +171,12 @@ def flashinfer_fused_allreduce_rmsnorm(
|
||||
rms_gamma: torch.Tensor,
|
||||
rms_eps: float,
|
||||
allreduce_params: "FlashInferFusedAllReduceParams",
|
||||
workspace: object,
|
||||
use_oneshot: bool,
|
||||
norm_out: torch.Tensor | None = None,
|
||||
):
|
||||
"""FlashInfer fused allreduce + rmsnorm operation."""
|
||||
if flashinfer_comm is None or _FI_WORKSPACE is None:
|
||||
if flashinfer_comm is None or workspace is None:
|
||||
raise RuntimeError("FlashInfer not available or workspace not initialized")
|
||||
|
||||
if norm_out is None:
|
||||
@@ -160,9 +185,13 @@ def flashinfer_fused_allreduce_rmsnorm(
|
||||
else:
|
||||
residual_out = input_tensor
|
||||
|
||||
layout_code = None
|
||||
if workspace.backend == "trtllm":
|
||||
layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
|
||||
|
||||
flashinfer_comm.allreduce_fusion(
|
||||
input=input_tensor,
|
||||
workspace=_FI_WORKSPACE,
|
||||
workspace=workspace,
|
||||
pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
|
||||
residual_in=residual,
|
||||
residual_out=residual_out,
|
||||
@@ -171,10 +200,10 @@ def flashinfer_fused_allreduce_rmsnorm(
|
||||
rms_eps=rms_eps,
|
||||
quant_out=None,
|
||||
scale_out=None,
|
||||
layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
|
||||
layout_code=layout_code,
|
||||
scale_factor=None,
|
||||
use_oneshot=use_oneshot,
|
||||
**allreduce_params.get_trtllm_fused_allreduce_kwargs(),
|
||||
**allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
@@ -185,12 +214,16 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
|
||||
rms_eps: float,
|
||||
scale_factor: torch.Tensor,
|
||||
allreduce_params: FlashInferFusedAllReduceParams,
|
||||
workspace: object,
|
||||
use_oneshot: bool = True,
|
||||
norm_out: torch.Tensor | None = None,
|
||||
quant_out: torch.Tensor | None = None,
|
||||
):
|
||||
"""FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
|
||||
if flashinfer_comm is None or _FI_WORKSPACE is None:
|
||||
"""FlashInfer fused allreduce + rmsnorm + FP8 quantization.
|
||||
|
||||
Note: Only supported by the trtllm backend.
|
||||
"""
|
||||
if flashinfer_comm is None or workspace is None:
|
||||
raise RuntimeError("FlashInfer not available or workspace not initialized")
|
||||
|
||||
if norm_out is None:
|
||||
@@ -201,7 +234,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
|
||||
|
||||
flashinfer_comm.allreduce_fusion(
|
||||
input=input_tensor,
|
||||
workspace=_FI_WORKSPACE,
|
||||
workspace=workspace,
|
||||
pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
|
||||
residual_in=residual,
|
||||
residual_out=residual_out,
|
||||
@@ -213,7 +246,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
|
||||
layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
|
||||
scale_factor=scale_factor,
|
||||
use_oneshot=use_oneshot,
|
||||
**allreduce_params.get_trtllm_fused_allreduce_kwargs(),
|
||||
**allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
@@ -224,13 +257,17 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
|
||||
rms_eps: float,
|
||||
input_global_scale: torch.Tensor,
|
||||
allreduce_params: FlashInferFusedAllReduceParams,
|
||||
workspace: object,
|
||||
quant_out: torch.Tensor,
|
||||
use_oneshot: bool,
|
||||
output_scale: torch.Tensor,
|
||||
norm_out: torch.Tensor | None = None,
|
||||
):
|
||||
"""FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
|
||||
if flashinfer_comm is None or _FI_WORKSPACE is None:
|
||||
"""FlashInfer fused allreduce + rmsnorm + FP4 quantization.
|
||||
|
||||
Note: Only supported by the trtllm backend.
|
||||
"""
|
||||
if flashinfer_comm is None or workspace is None:
|
||||
raise RuntimeError("FlashInfer not available or workspace not initialized")
|
||||
|
||||
if norm_out is None:
|
||||
@@ -241,7 +278,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
|
||||
|
||||
flashinfer_comm.allreduce_fusion(
|
||||
input=input_tensor,
|
||||
workspace=_FI_WORKSPACE,
|
||||
workspace=workspace,
|
||||
pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
|
||||
residual_in=residual,
|
||||
residual_out=residual_out,
|
||||
@@ -253,7 +290,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
|
||||
layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
|
||||
scale_factor=input_global_scale,
|
||||
use_oneshot=use_oneshot,
|
||||
**allreduce_params.get_trtllm_fused_allreduce_kwargs(),
|
||||
**allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
@@ -386,13 +423,16 @@ def run_benchmarks(
|
||||
dtype: torch.dtype,
|
||||
use_residual: bool,
|
||||
allreduce_params: FlashInferFusedAllReduceParams | None,
|
||||
workspaces: dict,
|
||||
quant_modes: set[str],
|
||||
no_oneshot: bool,
|
||||
):
|
||||
"""Run all benchmarks for given configuration.
|
||||
|
||||
Args:
|
||||
quant_mode: "none", "fp8_only", "fp4_only", or "all"
|
||||
allreduce_params: Shared parameters for FlashInfer fused allreduce.
|
||||
workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace.
|
||||
quant_modes: Set of quantization modes: "none", "fp8", "fp4".
|
||||
"""
|
||||
(
|
||||
input_tensor,
|
||||
@@ -454,10 +494,11 @@ def run_benchmarks(
|
||||
logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e)
|
||||
results["standard_allreduce_rmsnorm_native_compiled"] = float("inf")
|
||||
|
||||
# FlashInfer Fused AllReduce + RMSNorm Oneshot/Twoshot
|
||||
if flashinfer_comm is not None and allreduce_params is not None:
|
||||
# FlashInfer Fused AllReduce + RMSNorm (all backends)
|
||||
for backend, workspace in workspaces.items():
|
||||
for use_oneshot in use_oneshot_options:
|
||||
suffix = "_oneshot" if use_oneshot else "_twoshot"
|
||||
key = f"flashinfer_{backend}_fused_allreduce_rmsnorm{suffix}"
|
||||
try:
|
||||
time_ms = benchmark_operation(
|
||||
flashinfer_fused_allreduce_rmsnorm,
|
||||
@@ -467,14 +508,17 @@ def run_benchmarks(
|
||||
rms_gamma=rms_gamma,
|
||||
rms_eps=rms_eps,
|
||||
allreduce_params=allreduce_params,
|
||||
workspace=workspace,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = time_ms
|
||||
results[key] = time_ms
|
||||
except Exception as e:
|
||||
logger.error("FlashInfer Fused AllReduce+RMSNorm failed: %s", e)
|
||||
results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = float(
|
||||
"inf"
|
||||
logger.error(
|
||||
"FlashInfer (%s) Fused AllReduce+RMSNorm failed: %s",
|
||||
backend,
|
||||
e,
|
||||
)
|
||||
results[key] = float("inf")
|
||||
|
||||
if "fp8" in quant_modes:
|
||||
# Standard AllReduce + RMSNorm + FP8 Quant
|
||||
@@ -540,10 +584,12 @@ def run_benchmarks(
|
||||
"inf"
|
||||
)
|
||||
|
||||
# FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot
|
||||
if flashinfer_comm is not None and allreduce_params is not None:
|
||||
# FlashInfer Fused AllReduce + RMSNorm + FP8 Quant (trtllm only)
|
||||
if "trtllm" in workspaces:
|
||||
trtllm_ws = workspaces["trtllm"]
|
||||
for use_oneshot in use_oneshot_options:
|
||||
suffix = "_oneshot" if use_oneshot else "_twoshot"
|
||||
key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp8_quant{suffix}"
|
||||
try:
|
||||
time_ms = benchmark_operation(
|
||||
flashinfer_fused_allreduce_rmsnorm_fp8_quant,
|
||||
@@ -555,19 +601,16 @@ def run_benchmarks(
|
||||
scale_factor=scale_fp8,
|
||||
quant_out=quant_out_fp8,
|
||||
allreduce_params=allreduce_params,
|
||||
workspace=trtllm_ws,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
|
||||
time_ms
|
||||
)
|
||||
results[key] = time_ms
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s",
|
||||
"FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP8 failed: %s",
|
||||
e,
|
||||
)
|
||||
results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
|
||||
float("inf")
|
||||
)
|
||||
results[key] = float("inf")
|
||||
|
||||
if "fp4" in quant_modes and current_platform.has_device_capability(100):
|
||||
# Standard AllReduce + RMSNorm + FP4 Quant
|
||||
@@ -627,10 +670,12 @@ def run_benchmarks(
|
||||
"inf"
|
||||
)
|
||||
|
||||
# FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot
|
||||
if flashinfer_comm is not None and allreduce_params is not None:
|
||||
# FlashInfer Fused AllReduce + RMSNorm + FP4 Quant (trtllm only)
|
||||
if "trtllm" in workspaces:
|
||||
trtllm_ws = workspaces["trtllm"]
|
||||
for use_oneshot in use_oneshot_options:
|
||||
suffix = "_oneshot" if use_oneshot else "_twoshot"
|
||||
key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp4_quant{suffix}"
|
||||
try:
|
||||
time_ms = benchmark_operation(
|
||||
flashinfer_fused_allreduce_rmsnorm_fp4_quant,
|
||||
@@ -641,49 +686,18 @@ def run_benchmarks(
|
||||
rms_eps=rms_eps,
|
||||
input_global_scale=scale_fp4,
|
||||
allreduce_params=allreduce_params,
|
||||
workspace=trtllm_ws,
|
||||
quant_out=fp4_quant_out,
|
||||
output_scale=fp4_output_scale,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
|
||||
time_ms
|
||||
)
|
||||
results[key] = time_ms
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s",
|
||||
"FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP4 failed: %s",
|
||||
e,
|
||||
)
|
||||
results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
|
||||
float("inf")
|
||||
)
|
||||
|
||||
# FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot
|
||||
if flashinfer_comm is not None and allreduce_params is not None:
|
||||
try:
|
||||
time_ms = benchmark_operation(
|
||||
flashinfer_fused_allreduce_rmsnorm_fp4_quant,
|
||||
input_tensor,
|
||||
residual=residual,
|
||||
norm_out=norm_out,
|
||||
rms_gamma=rms_gamma,
|
||||
rms_eps=rms_eps,
|
||||
input_global_scale=scale_fp4,
|
||||
allreduce_params=allreduce_params,
|
||||
quant_out=fp4_quant_out,
|
||||
output_scale=fp4_output_scale,
|
||||
use_oneshot=False,
|
||||
)
|
||||
results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = (
|
||||
time_ms
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s",
|
||||
e,
|
||||
)
|
||||
results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float(
|
||||
"inf"
|
||||
)
|
||||
results[key] = float("inf")
|
||||
|
||||
return results
|
||||
|
||||
@@ -1021,8 +1035,7 @@ def main():
|
||||
|
||||
configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
|
||||
|
||||
# Setup FlashInfer workspace if available
|
||||
workspace = None
|
||||
# Setup FlashInfer workspaces for all backends
|
||||
allreduce_params = None
|
||||
|
||||
if flashinfer_comm is not None:
|
||||
@@ -1037,15 +1050,17 @@ def main():
|
||||
args.hidden_dim * max_element_size
|
||||
)
|
||||
|
||||
workspace = setup_flashinfer_workspace(
|
||||
world_size,
|
||||
rank,
|
||||
args.hidden_dim,
|
||||
max_num_token,
|
||||
dtype=workspace_dtype,
|
||||
)
|
||||
for backend in FLASHINFER_BACKENDS:
|
||||
setup_flashinfer_workspace(
|
||||
backend=backend,
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
hidden_dim=args.hidden_dim,
|
||||
max_token_num=max_num_token,
|
||||
dtype=workspace_dtype,
|
||||
)
|
||||
|
||||
if workspace is not None:
|
||||
if _FI_WORKSPACES:
|
||||
allreduce_params = FlashInferFusedAllReduceParams(
|
||||
max_token_num=max_num_token,
|
||||
)
|
||||
@@ -1071,6 +1086,7 @@ def main():
|
||||
dtype,
|
||||
use_residual,
|
||||
allreduce_params,
|
||||
workspaces=_FI_WORKSPACES,
|
||||
quant_modes=quant_modes,
|
||||
no_oneshot=args.no_oneshot,
|
||||
)
|
||||
@@ -1109,11 +1125,13 @@ def main():
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
if workspace is not None:
|
||||
cleanup_flashinfer_workspace(workspace)
|
||||
cleanup_flashinfer_workspaces()
|
||||
|
||||
dist.barrier()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
main()
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct SSMParamsBase {
|
||||
using index_t = uint32_t;
|
||||
using index_t = size_t;
|
||||
|
||||
int batch, dim, seqlen, dstate, n_groups, n_chunks;
|
||||
int dim_ngroups_ratio;
|
||||
|
||||
@@ -107,7 +107,9 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
int64_t outOffset =
|
||||
rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
out[outOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -140,7 +142,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
|
||||
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
|
||||
int sf_n_unpadded = int(n / CVT_FP4_ELTS_PER_THREAD);
|
||||
|
||||
int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
|
||||
int grid_x = std::min(
|
||||
|
||||
@@ -109,7 +109,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols,
|
||||
int32_t sf_n_unpadded, Type const* __restrict__ in,
|
||||
int32_t sf_n_unpadded, int32_t num_packed_cols,
|
||||
Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out,
|
||||
uint32_t* __restrict__ SFout) {
|
||||
@@ -131,7 +132,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
// Iterate over all rows and cols including padded ones -
|
||||
// ensures we visit every single scale factor address to initialize it.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
|
||||
if (colIdx < sf_n_unpadded) {
|
||||
if (colIdx < num_packed_cols) {
|
||||
PackedVec in_vec;
|
||||
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
|
||||
@@ -222,7 +223,8 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
} else {
|
||||
int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
|
||||
int num_packed_cols = n / CVT_FP4_ELTS_PER_THREAD;
|
||||
int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));
|
||||
int grid_x = std::min(
|
||||
m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
@@ -232,8 +234,8 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
|
||||
<<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, input_ptr,
|
||||
input_sf_ptr,
|
||||
<<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, num_packed_cols,
|
||||
input_ptr, input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
|
||||
+8
-4
@@ -132,8 +132,10 @@ ENV UV_LINK_MODE=copy
|
||||
# Verify GCC version
|
||||
RUN gcc --version
|
||||
|
||||
# Ensure CUDA compatibility library is loaded
|
||||
RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig
|
||||
# Enable CUDA forward compatibility by setting '-e VLLM_ENABLE_CUDA_COMPATIBILITY=1'
|
||||
# Only needed for datacenter/professional GPUs with older drivers.
|
||||
# See: https://docs.nvidia.com/deploy/cuda-compatibility/
|
||||
ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
|
||||
|
||||
# ============================================================
|
||||
# SLOW-CHANGING DEPENDENCIES BELOW
|
||||
@@ -560,8 +562,10 @@ ENV UV_HTTP_TIMEOUT=500
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE=copy
|
||||
|
||||
# Ensure CUDA compatibility library is loaded
|
||||
RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig
|
||||
# Enable CUDA forward compatibility by setting '-e VLLM_ENABLE_CUDA_COMPATIBILITY=1'
|
||||
# Only needed for datacenter/professional GPUs with older drivers.
|
||||
# See: https://docs.nvidia.com/deploy/cuda-compatibility/
|
||||
ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
|
||||
|
||||
# ============================================================
|
||||
# SLOW-CHANGING DEPENDENCIES BELOW
|
||||
|
||||
+15
-3
@@ -6,8 +6,7 @@ ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
|
||||
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
add-apt-repository -y ppa:kobuk-team/intel-graphics
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
|
||||
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
@@ -28,9 +27,22 @@ RUN apt clean && apt-get update -y && \
|
||||
python3-pip
|
||||
|
||||
RUN apt update && apt upgrade -y && \
|
||||
apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
|
||||
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
|
||||
|
||||
# Install UMD
|
||||
RUN mkdir neo && \
|
||||
cd neo && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
|
||||
dpkg -i *.deb && \
|
||||
cd .. && \
|
||||
rm -rf neo
|
||||
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
|
||||
@@ -4,6 +4,11 @@ This section guides you through running benchmark tests with the extensive datas
|
||||
|
||||
It's a living document, updated as new features and datasets become available.
|
||||
|
||||
!!! tip
|
||||
The benchmarks described on this page are mainly for evaluating specific vLLM features as well as regression testing.
|
||||
|
||||
For benchmarking production vLLM servers, we recommend [GuideLLM](https://github.com/vllm-project/guidellm), an established performance benchmarking framework with live progress updates and automatic report generation. It is also more flexible than `vllm bench serve` in terms of dataset loading, request formatting, and workload patterns.
|
||||
|
||||
## Dataset Overview
|
||||
|
||||
<style>
|
||||
|
||||
+47
-41
@@ -1,10 +1,15 @@
|
||||
# Parameter Sweeps
|
||||
|
||||
`vllm bench sweep` is a suite of commands designed to run benchmarks across multiple configurations and compare them by visualizing the results.
|
||||
|
||||
## Online Benchmark
|
||||
|
||||
### Basic
|
||||
|
||||
`vllm bench sweep serve` automatically starts `vllm serve` and runs `vllm bench serve` to evaluate vLLM over multiple configurations.
|
||||
`vllm bench sweep serve` starts `vllm serve` and iteratively runs `vllm bench serve` for each server configuration.
|
||||
|
||||
!!! tip
|
||||
If you only need to run benchmarks for a single server configuration, consider using [GuideLLM](https://github.com/vllm-project/guidellm), an established performance benchmarking framework with live progress updates and automatic report generation. It is also more flexible than `vllm bench serve` in terms of dataset loading, request formatting, and workload patterns.
|
||||
|
||||
Follow these steps to run the script:
|
||||
|
||||
@@ -50,14 +55,17 @@ Follow these steps to run the script:
|
||||
```json
|
||||
[
|
||||
{
|
||||
"_benchmark_name": "scenario_A",
|
||||
"random_input_len": 128,
|
||||
"random_output_len": 32
|
||||
},
|
||||
{
|
||||
"_benchmark_name": "scenario_B",
|
||||
"random_input_len": 256,
|
||||
"random_output_len": 64
|
||||
},
|
||||
{
|
||||
"_benchmark_name": "scenario_C",
|
||||
"random_input_len": 512,
|
||||
"random_output_len": 128
|
||||
}
|
||||
@@ -77,6 +85,8 @@ vllm bench sweep serve \
|
||||
-o benchmarks/results
|
||||
```
|
||||
|
||||
By default, each parameter combination is benchmarked 3 times to make the results more reliable. You can adjust the number of runs by setting `--num-runs`.
|
||||
|
||||
!!! important
|
||||
If both `--serve-params` and `--bench-params` are passed, the script will iterate over the Cartesian product between them.
|
||||
You can use `--dry-run` to preview the commands to be run.
|
||||
@@ -86,60 +96,40 @@ vllm bench sweep serve \
|
||||
In case you are using a custom `--serve-cmd`, you can override the commands used for resetting the state by setting `--after-bench-cmd`.
|
||||
|
||||
!!! note
|
||||
By default, each parameter combination is run 3 times to make the results more reliable. You can adjust the number of runs by setting `--num-runs`.
|
||||
You should set `_benchmark_name` to provide a human-readable name for parameter combinations involving many variables.
|
||||
This becomes mandatory if the file name would otherwise exceed the maximum path length allowed by the filesystem.
|
||||
|
||||
!!! tip
|
||||
You can use the `--resume` option to continue the parameter sweep if one of the runs failed.
|
||||
|
||||
### SLA auto-tuner
|
||||
You can use the `--resume` option to continue the parameter sweep if an unexpected error occurs, e.g., timeout when connecting to HF Hub.
|
||||
|
||||
`vllm bench sweep serve_sla` is a wrapper over `vllm bench sweep serve` that tunes either the request rate or concurrency (choose using `--sla-variable`) in order to satisfy the SLA constraints given by `--sla-params`.
|
||||
### SLA Scanner
|
||||
|
||||
For example, to ensure E2E latency within different target values for 99% of requests:
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"p99_e2el_ms": "<=200"
|
||||
},
|
||||
{
|
||||
"p99_e2el_ms": "<=500"
|
||||
},
|
||||
{
|
||||
"p99_e2el_ms": "<=1000"
|
||||
},
|
||||
{
|
||||
"p99_e2el_ms": "<=2000"
|
||||
}
|
||||
]
|
||||
```
|
||||
`vllm bench sweep serve_sla` is a variant of `vllm bench sweep serve` that scans through values of request rate or concurrency (choose using `--sla-variable`) in order to find the tradeoff between latency and throughput. The results can then be [visualized](#visualization) to determine the feasible SLAs.
|
||||
|
||||
Example command:
|
||||
|
||||
```bash
|
||||
vllm bench sweep serve_sla \
|
||||
--serve-cmd 'vllm serve meta-llama/Llama-2-7b-chat-hf' \
|
||||
--bench-cmd 'vllm bench serve --model meta-llama/Llama-2-7b-chat-hf --backend vllm --endpoint /v1/completions --dataset-name sharegpt --dataset-path benchmarks/ShareGPT_V3_unfiltered_cleaned_split.json' \
|
||||
--bench-cmd 'vllm bench serve --model meta-llama/Llama-2-7b-chat-hf --backend vllm --endpoint /v1/completions --dataset-name sharegpt --dataset-path benchmarks/ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 100' \
|
||||
--serve-params benchmarks/serve_hparams.json \
|
||||
--bench-params benchmarks/bench_hparams.json \
|
||||
--sla-params benchmarks/sla_hparams.json \
|
||||
--sla-variable max_concurrency \
|
||||
--bench-params benchmarks/bench_hparams.json
|
||||
-o benchmarks/results
|
||||
```
|
||||
|
||||
The algorithm for adjusting the SLA variable is as follows:
|
||||
The algorithm for scanning through different values of `sla_variable` can be summarized as follows:
|
||||
|
||||
1. Run the benchmark once with maximum possible QPS, and once with minimum possible QPS. For each run, calculate the distance of the SLA metrics from their targets, resulting in data points of QPS vs SLA distance.
|
||||
2. Perform spline interpolation between the data points to estimate the QPS that results in zero SLA distance.
|
||||
3. Run the benchmark with the estimated QPS and add the resulting data point to the history.
|
||||
4. Repeat Steps 2 and 3 until the maximum QPS that passes SLA and the minimum QPS that fails SLA in the history are close enough to each other.
|
||||
1. Run the benchmark once with `sla_variable = 1` to simulate serial inference. This results in the lowest possible latency and throughput.
|
||||
2. Run the benchmark once with `sla_variable = num_prompts` to simulate batch inference over the whole dataset. This results in the highest possible latency and throughput.
|
||||
3. Estimate the maximum value of `sla_variable` that can be supported by the server without oversaturating it.
|
||||
4. Run the benchmark over intermediate values of `sla_variable` uniformly using the remaining iterations.
|
||||
|
||||
!!! important
|
||||
SLA tuning is applied over each combination of `--serve-params`, `--bench-params`, and `--sla-params`.
|
||||
You can override the number of iterations in the algorithm by setting `--sla-iters`.
|
||||
|
||||
For a given combination of `--serve-params` and `--bench-params`, we share the benchmark results across `--sla-params` to avoid rerunning benchmarks with the same SLA variable value.
|
||||
!!! tip
|
||||
This is our equivalent of [GuideLLM's `--profile sweep`](https://github.com/vllm-project/guidellm/blob/v0.5.3/src/guidellm/benchmark/profiles.py#L575).
|
||||
|
||||
### Startup
|
||||
## Startup Benchmark
|
||||
|
||||
`vllm bench sweep startup` runs `vllm bench startup` across parameter combinations to compare cold/warm startup time for different engine settings.
|
||||
|
||||
@@ -202,15 +192,28 @@ vllm bench sweep startup \
|
||||
|
||||
`vllm bench sweep plot` can be used to plot performance curves from parameter sweep results.
|
||||
|
||||
Example command:
|
||||
Control the variables to plot via `--var-x` and `--var-y`, optionally applying `--filter-by` and `--bin-by` to the values. The plot is organized according to `--fig-by`, `--row-by`, `--col-by`, and `--curve-by`.
|
||||
|
||||
Example commands for visualizing [SLA Scanner](#sla-scanner) results:
|
||||
|
||||
```bash
|
||||
# Latency increases as the request rate increases
|
||||
vllm bench sweep plot benchmarks/results/<timestamp> \
|
||||
--var-x max_concurrency \
|
||||
--var-x request_rate \
|
||||
--var-y p99_ttft_ms \
|
||||
--row-by random_input_len \
|
||||
--col-by random_output_len \
|
||||
--curve-by api_server_count,max_num_batched_tokens \
|
||||
--filter-by 'max_concurrency<=1024'
|
||||
--curve-by max_num_seqs,max_num_batched_tokens \
|
||||
--filter-by 'request_rate<=128'
|
||||
|
||||
# Tradeoff between latency and throughput
|
||||
vllm bench sweep plot benchmarks/results/<timestamp> \
|
||||
--var-x request_throughput \
|
||||
--var-y median_ttft_ms \
|
||||
--row-by random_input_len \
|
||||
--col-by random_output_len \
|
||||
--curve-by max_num_seqs,max_num_batched_tokens \
|
||||
--filter-by 'request_rate<=128'
|
||||
```
|
||||
|
||||
!!! tip
|
||||
@@ -233,3 +236,6 @@ Example:
|
||||
vllm bench sweep plot_pareto benchmarks/results/<timestamp> \
|
||||
--label-by max_concurrency,tensor_parallel_size,pipeline_parallel_size
|
||||
```
|
||||
|
||||
!!! tip
|
||||
You can use `--dry-run` to preview the figures to be plotted.
|
||||
|
||||
@@ -297,6 +297,23 @@ You can add any other [engine-args](https://docs.vllm.ai/en/latest/configuration
|
||||
RUN uv pip install --system git+https://github.com/huggingface/transformers.git
|
||||
```
|
||||
|
||||
#### Running on Systems with Older CUDA Drivers
|
||||
|
||||
vLLM's Docker image comes with [CUDA compatibility libraries](https://docs.nvidia.com/deploy/cuda-compatibility/index.html) pre-installed. This allows you to run vLLM on systems with NVIDIA drivers that are older than the CUDA Toolkit version used in the image, but only supports select professional and datacenter NVIDIA GPUs.
|
||||
|
||||
To enable this feature, set the `VLLM_ENABLE_CUDA_COMPATIBILITY` environment variable to `1` or `true` when running the container:
|
||||
|
||||
```bash
|
||||
docker run --runtime nvidia --gpus all \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
-p 8000:8000 \
|
||||
--env "HF_TOKEN=<secret>" \
|
||||
--env "VLLM_ENABLE_CUDA_COMPATIBILITY=1" \
|
||||
vllm/vllm-openai <args...>
|
||||
```
|
||||
|
||||
This will automatically configure `LD_LIBRARY_PATH` to point to the compatibility libraries before loading PyTorch and other dependencies.
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
# --8<-- [start:build-image-from-source]
|
||||
|
||||
|
||||
@@ -498,6 +498,67 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
|
||||
- Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
|
||||
- Reranking (text + multi-modal): [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
|
||||
|
||||
### Llama Nemotron Multimodal Embedding Models
|
||||
|
||||
Llama Nemotron VL Embedding models combine the bidirectional Llama embedding backbone
|
||||
(from `nvidia/llama-nemotron-embed-1b-v2`) with SigLIP as the vision encoder to produce
|
||||
single-vector embeddings from text and/or images.
|
||||
|
||||
| Architecture | Backbone | Example HF Models |
|
||||
|---|---|---|
|
||||
| `LlamaNemotronVLModel` | Bidirectional Llama + SigLIP | `nvidia/llama-nemotron-embed-vl-1b-v2` |
|
||||
|
||||
Start the server:
|
||||
|
||||
```shell
|
||||
vllm serve nvidia/llama-nemotron-embed-vl-1b-v2 \
|
||||
--trust-remote-code \
|
||||
--chat-template examples/pooling/embed/template/nemotron_embed_vl.jinja
|
||||
```
|
||||
|
||||
!!! note
|
||||
The chat template bundled with this model's tokenizer is not suitable for
|
||||
the embeddings API. Use the provided override template above when serving
|
||||
with the `messages`-based (chat-style) embeddings endpoint.
|
||||
|
||||
The override template uses the message `role` to automatically prepend the
|
||||
appropriate prefix: set `role` to `"query"` for queries (prepends `query: `)
|
||||
or `"document"` for passages (prepends `passage: `). Any other role omits
|
||||
the prefix.
|
||||
|
||||
Embed text queries:
|
||||
|
||||
```shell
|
||||
curl -s http://localhost:8000/v1/embeddings -H "Content-Type: application/json" -d '{
|
||||
"model": "nvidia/llama-nemotron-embed-vl-1b-v2",
|
||||
"messages": [
|
||||
{
|
||||
"role": "query",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is machine learning?"}
|
||||
]
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Embed images via the chat-style `messages` field:
|
||||
|
||||
```shell
|
||||
curl -s http://localhost:8000/v1/embeddings -H "Content-Type: application/json" -d '{
|
||||
"model": "nvidia/llama-nemotron-embed-vl-1b-v2",
|
||||
"messages": [
|
||||
{
|
||||
"role": "document",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
|
||||
{"type": "text", "text": "Describe the image."}
|
||||
]
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
### BAAI/bge-m3
|
||||
|
||||
The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
|
||||
|
||||
@@ -372,6 +372,7 @@ th {
|
||||
| `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ |
|
||||
| `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |
|
||||
| `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
|
||||
| `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ |
|
||||
| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ |
|
||||
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
|
||||
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -820,6 +821,7 @@ The following table lists those that are tested in vLLM.
|
||||
|--------------|--------|--------|-------------------|----------------------|---------------------------|
|
||||
| `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | |
|
||||
| `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
|
||||
| `LlamaNemotronVLModel` | Llama Nemotron Embedding + SigLIP | T + I | `nvidia/llama-nemotron-embed-vl-1b-v2` | | |
|
||||
| `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ |
|
||||
| `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ |
|
||||
| `Qwen3VLForConditionalGeneration`<sup>C</sup> | Qwen3-VL | T + I + V | `Qwen/Qwen3-VL-Embedding-2B`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -318,7 +318,32 @@ This indicates vLLM failed to initialize the NCCL communicator, possibly due to
|
||||
|
||||
## CUDA error: the provided PTX was compiled with an unsupported toolchain
|
||||
|
||||
If you see an error like `RuntimeError: CUDA error: the provided PTX was compiled with an unsupported toolchain.`, it means that the CUDA PTX in vLLM's wheels was compiled with a toolchain unsupported by your system. The released vLLM wheels have to be compiled with a specific version of CUDA toolkit, and the compiled code might fail to run on lower versions of CUDA drivers. Read [cuda compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/) for more details. The solution is to install `cuda-compat` package from your package manager. For example, on Ubuntu, you can run `sudo apt-get install cuda-compat-12-9`, and then add `export LD_LIBRARY_PATH=/usr/local/cuda-12.9/compat:$LD_LIBRARY_PATH` to your `.bashrc` file. When successfully installed, you should see that the output of `nvidia-smi` will show `CUDA Version: 12.9`. Note that we use CUDA 12.9 as an example here, you may want to install a higher version of cuda-compat package in case vLLM's default CUDA version goes higher.
|
||||
If you see an error like `RuntimeError: CUDA error: the provided PTX was compiled with an unsupported toolchain`, it means that the CUDA PTX in vLLM's wheels was compiled with a toolchain unsupported by your system. This section also applies if you get the error `RuntimeError: The NVIDIA driver on your system is too old`.
|
||||
|
||||
The released vLLM wheels are compiled with a specific version of CUDA toolkit, and the compiled code might fail to run on lower versions of CUDA drivers. Read [CUDA compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/) for more details. **This is only supported on select professional and datacenter NVIDIA GPUs.**
|
||||
|
||||
If you are using the vLLM official Docker image, you can solve this by adding `-e VLLM_ENABLE_CUDA_COMPATIBILITY=1` to your `docker run` command. This will enable the pre-installed CUDA forward compatibility libraries.
|
||||
|
||||
If you are running vLLM outside of Docker, the solution is to install the `cuda-compat` package from your package manager with the [CUDA repository](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/) enabled. For example, on Ubuntu, you can run `sudo apt-get install cuda-compat-12-9`, and then set `export VLLM_ENABLE_CUDA_COMPATIBILITY=1` and `export VLLM_CUDA_COMPATIBILITY_PATH="/usr/local/cuda-12.9/compat"`.
|
||||
|
||||
On Conda, you can install the `conda-forge::cuda-compat` package (e.g., `conda install -c conda-forge cuda-compat=12.9`), then after activating the environment, set `export VLLM_ENABLE_CUDA_COMPATIBILITY=1` and `export VLLM_CUDA_COMPATIBILITY_PATH="${CONDA_PREFIX}/cuda-compat"`.
|
||||
|
||||
You can verify the configuration works by running a minimal Python script that initializes CUDA via vLLM:
|
||||
|
||||
```bash
|
||||
export VLLM_ENABLE_CUDA_COMPATIBILITY=1
|
||||
export VLLM_CUDA_COMPATIBILITY_PATH="/usr/local/cuda-12.9/compat"
|
||||
|
||||
python3 - << 'EOF'
|
||||
import vllm
|
||||
import torch
|
||||
|
||||
print(f"CUDA available: {torch.cuda.is_available()}")
|
||||
print(f"CUDA device count: {torch.cuda.device_count()}")
|
||||
EOF
|
||||
```
|
||||
|
||||
Note that we use CUDA 12.9 as an example here, and you may want to install a higher version of cuda-compat package in case vLLM's default CUDA version goes higher.
|
||||
|
||||
## ptxas fatal: Value 'sm_110a' is not defined for option 'gpu-name'
|
||||
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
{%- if messages | length > 1 -%}
|
||||
{{ raise_exception('Embedding models should only embed one message at a time') }}
|
||||
{%- endif -%}
|
||||
|
||||
{% set vars = namespace(prefix='', images=[], texts=[]) %}
|
||||
{%- for message in messages -%}
|
||||
{%- if message['role'] == 'query' -%}
|
||||
{%- set vars.prefix = 'query: ' %}
|
||||
{%- elif message['role'] == 'document' -%}
|
||||
{%- set vars.prefix = 'passage: ' %}
|
||||
{%- endif -%}
|
||||
{%- for content in message['content'] -%}
|
||||
{%- if content['type'] == 'text' -%}
|
||||
{%- set vars.texts = vars.texts + [content['text']] %}
|
||||
{%- elif content['type'] == 'image' -%}
|
||||
{%- set vars.images = vars.images + ['<image> '] %}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- endfor -%}
|
||||
{{- bos_token }}{{ vars.prefix }}{{ (vars.images + vars.texts) | join('') }}
|
||||
@@ -1033,7 +1033,7 @@ setup(
|
||||
ext_modules=ext_modules,
|
||||
install_requires=get_requirements(),
|
||||
extras_require={
|
||||
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy"],
|
||||
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
|
||||
"tensorizer": ["tensorizer==2.10.1"],
|
||||
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
|
||||
"runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test prefetch offloading correctness with Llama model."""
|
||||
|
||||
from ..utils import compare_two_settings
|
||||
|
||||
|
||||
def test_prefetch_offload_llama():
|
||||
"""Test prefetch CPU offloading with Llama-3.2-1B-Instruct.
|
||||
|
||||
Compares outputs between:
|
||||
1. Baseline (no offloading)
|
||||
2. Prefetch offloading (group_size=8, num_in_group=2, prefetch_step=1)
|
||||
|
||||
This tests prefetching-based offloading on a dense model.
|
||||
"""
|
||||
compare_two_settings(
|
||||
"meta-llama/Llama-3.2-1B-Instruct",
|
||||
[
|
||||
# Prefetch offloading configuration
|
||||
"--offload-group-size",
|
||||
"8",
|
||||
"--offload-num-in-group",
|
||||
"2",
|
||||
"--offload-prefetch-step",
|
||||
"1",
|
||||
# Selective offloading: only MLP weights
|
||||
"--offload-params",
|
||||
"gate_up_proj",
|
||||
"down_proj",
|
||||
],
|
||||
[], # Baseline: no offloading
|
||||
)
|
||||
@@ -1,298 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.benchmarks.sweep.param_sweep import ParameterSweepItem
|
||||
from vllm.benchmarks.sweep.serve_sla import _get_sla_run_path, solve_sla
|
||||
from vllm.benchmarks.sweep.server import ServerProcess
|
||||
from vllm.benchmarks.sweep.sla_sweep import (
|
||||
SLACriterionBase,
|
||||
SLALessThan,
|
||||
SLALessThanOrEqualTo,
|
||||
SLASweepItem,
|
||||
)
|
||||
|
||||
|
||||
def _set_return_value(
|
||||
var2metric: Callable[[ParameterSweepItem], list[dict[str, float]]],
|
||||
):
|
||||
"""
|
||||
Create a patch for run_sla with a specific function
|
||||
indicating the relationship between the benchmark combination
|
||||
(which includes the SLA variable) and the SLA criterion.
|
||||
"""
|
||||
|
||||
def mock_run_sla(
|
||||
server: ServerProcess | None,
|
||||
bench_cmd: list[str],
|
||||
*,
|
||||
serve_comb: ParameterSweepItem,
|
||||
bench_comb: ParameterSweepItem,
|
||||
iter_path: Path,
|
||||
num_runs: int,
|
||||
dry_run: bool,
|
||||
):
|
||||
iter_data = var2metric(bench_comb)
|
||||
|
||||
summary_path = _get_sla_run_path(iter_path, run_number=None)
|
||||
summary_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with summary_path.open("w") as f:
|
||||
json.dump(iter_data, f, indent=4)
|
||||
|
||||
return iter_data
|
||||
|
||||
return patch("vllm.benchmarks.sweep.serve_sla.run_sla", side_effect=mock_run_sla)
|
||||
|
||||
|
||||
def _var2metric_linear():
|
||||
def wrapped(bench_comb):
|
||||
x = float(bench_comb["request_rate"])
|
||||
y = x
|
||||
|
||||
return [{"request_throughput": y}]
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _var2metric_concave(elbow_point: float):
|
||||
def wrapped(bench_comb):
|
||||
x = float(bench_comb["request_rate"])
|
||||
if x < elbow_point:
|
||||
y = 0.5 * (x - elbow_point) + elbow_point
|
||||
else:
|
||||
y = 1.5 * (x - elbow_point) + elbow_point
|
||||
|
||||
return [{"request_throughput": y}]
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _var2metric_convex(elbow_point: float):
|
||||
def wrapped(bench_comb):
|
||||
x = float(bench_comb["request_rate"])
|
||||
if x < elbow_point:
|
||||
y = 1.5 * (x - elbow_point) + elbow_point
|
||||
else:
|
||||
y = 0.5 * (x - elbow_point) + elbow_point
|
||||
|
||||
return [{"request_throughput": y}]
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _var2metric_quadratic(y_intercept: float):
|
||||
def wrapped(bench_comb):
|
||||
x = float(bench_comb["request_rate"])
|
||||
y = y_intercept + 0.1 * x**2
|
||||
|
||||
return [{"request_throughput": y}]
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _var2metric_sqrt(y_intercept: float):
|
||||
def wrapped(bench_comb):
|
||||
x = float(bench_comb["request_rate"])
|
||||
y = y_intercept + 10 * x**0.5
|
||||
|
||||
return [{"request_throughput": y}]
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _run_solve_sla(
|
||||
var2metric: Callable[[ParameterSweepItem], list[dict[str, float]]],
|
||||
criterion: SLACriterionBase,
|
||||
base_path: Path,
|
||||
min_value: int = 1,
|
||||
max_value: int = 100,
|
||||
):
|
||||
with _set_return_value(var2metric):
|
||||
result = solve_sla(
|
||||
server=None,
|
||||
bench_cmd=[],
|
||||
serve_comb=ParameterSweepItem(),
|
||||
bench_comb=ParameterSweepItem(),
|
||||
sla_comb=SLASweepItem({"request_throughput": criterion}),
|
||||
base_path=base_path,
|
||||
num_runs=1,
|
||||
dry_run=False,
|
||||
sla_variable="request_rate",
|
||||
sla_min_value=min_value,
|
||||
sla_max_value=max_value,
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def test_solve_linear_sla_le(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_linear(),
|
||||
SLALessThanOrEqualTo(target=32),
|
||||
tmp_path,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 32
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
1: True,
|
||||
32: True,
|
||||
33: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_linear_sla_lt(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_linear(),
|
||||
SLALessThan(target=32),
|
||||
tmp_path,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 31
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
1: True,
|
||||
31: True,
|
||||
32: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_linear_sla_oob(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_linear(),
|
||||
SLALessThanOrEqualTo(target=32),
|
||||
tmp_path,
|
||||
min_value=64,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 64
|
||||
assert history.get_min_failing() == 64
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
64: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_concave_sla_le(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_concave(elbow_point=32),
|
||||
SLALessThanOrEqualTo(target=24),
|
||||
tmp_path,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 16
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
1: True,
|
||||
7: True,
|
||||
13: True,
|
||||
15: True,
|
||||
16: True,
|
||||
17: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_convex_sla_le(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_convex(elbow_point=32),
|
||||
SLALessThanOrEqualTo(target=24),
|
||||
tmp_path,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 26
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
1: True,
|
||||
48: False,
|
||||
30: False,
|
||||
24: True,
|
||||
26: True,
|
||||
27: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_quadratic_sla_le(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_quadratic(y_intercept=10),
|
||||
SLALessThanOrEqualTo(target=50),
|
||||
tmp_path,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 20
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
1: True,
|
||||
4: True,
|
||||
20: True,
|
||||
21: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_sqrt_sla_le(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_sqrt(y_intercept=10),
|
||||
SLALessThanOrEqualTo(target=100),
|
||||
tmp_path,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 81
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
100: False,
|
||||
1: True,
|
||||
89: False,
|
||||
81: True,
|
||||
82: False,
|
||||
}
|
||||
|
||||
|
||||
def test_solve_reuse_history(tmp_path):
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_linear(),
|
||||
SLALessThanOrEqualTo(target=10),
|
||||
tmp_path,
|
||||
min_value=1,
|
||||
max_value=20,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 10
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
20: False,
|
||||
1: True,
|
||||
10: True,
|
||||
11: False,
|
||||
}
|
||||
|
||||
sla_data, history = _run_solve_sla(
|
||||
_var2metric_linear(),
|
||||
SLALessThanOrEqualTo(target=30),
|
||||
tmp_path,
|
||||
min_value=21,
|
||||
max_value=40,
|
||||
)
|
||||
|
||||
assert history.get_max_passing() == 30
|
||||
|
||||
assert {val: margin <= 0 for val, margin in history.items()} == {
|
||||
# Items from the past run
|
||||
# (the margins are different because the target changed)
|
||||
20: True,
|
||||
1: True,
|
||||
10: True,
|
||||
11: True,
|
||||
# Items from this run
|
||||
40: False,
|
||||
30: True,
|
||||
31: False,
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_cuda_platform():
|
||||
"""
|
||||
Fixture that returns a factory for creating mocked CUDA platforms.
|
||||
|
||||
Usage:
|
||||
def test_something(mock_cuda_platform):
|
||||
with mock_cuda_platform(is_cuda=True, capability=(9, 0)):
|
||||
# test code
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None):
|
||||
mock_platform = MagicMock()
|
||||
mock_platform.is_cuda.return_value = is_cuda
|
||||
if capability is not None:
|
||||
mock_platform.get_device_capability.return_value = DeviceCapability(
|
||||
*capability
|
||||
)
|
||||
with patch("vllm.platforms.current_platform", mock_platform):
|
||||
yield mock_platform
|
||||
|
||||
return _mock_platform
|
||||
@@ -94,7 +94,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
|
||||
run_model(full_compilation_config, model_name, **model_kwargs)
|
||||
|
||||
num_compile_ranges = len(full_compilation_config.get_compile_ranges())
|
||||
assert num_compile_ranges in [1, 2]
|
||||
assert num_compile_ranges in [1, 2, 3]
|
||||
|
||||
print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}")
|
||||
print("Fusion results:")
|
||||
@@ -107,12 +107,33 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
|
||||
|
||||
# Now check the matches
|
||||
for match_name in matches_check:
|
||||
num_ranges_activated = (
|
||||
1 if match_name == "ar_rms_fusion" else num_compile_ranges
|
||||
)
|
||||
n_expected = tp_size * num_ranges_activated
|
||||
|
||||
log_matches = list(int(ms) for ms in log_matches_dict[match_name])
|
||||
|
||||
# AR+RMS skips the largest range; SP skips the smallest.
|
||||
# When both are enabled, AR+RMS activation count is
|
||||
# model-dependent (hidden_size affects threshold), so derive
|
||||
# from log data.
|
||||
if (
|
||||
match_name == "ar_rms_fusion"
|
||||
and "sequence_parallel" in matches_check
|
||||
and num_compile_ranges >= 2
|
||||
):
|
||||
assert (
|
||||
len(log_matches) >= tp_size and len(log_matches) % tp_size == 0
|
||||
), (
|
||||
f"Expected multiple of {tp_size} ar_rms log entries, "
|
||||
f"found {len(log_matches)}"
|
||||
)
|
||||
num_ranges_activated = len(log_matches) // tp_size
|
||||
elif (
|
||||
match_name in ("ar_rms_fusion", "sequence_parallel")
|
||||
and num_compile_ranges >= 2
|
||||
):
|
||||
num_ranges_activated = num_compile_ranges - 1
|
||||
else:
|
||||
num_ranges_activated = num_compile_ranges
|
||||
|
||||
n_expected = tp_size * num_ranges_activated
|
||||
assert len(log_matches) == n_expected, (
|
||||
f"Could not find {n_expected} {match_name} "
|
||||
f"(found {len(log_matches)}) in:\n {log_holder.text}"
|
||||
@@ -122,8 +143,8 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
|
||||
|
||||
if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
|
||||
# AR+rms+quant takes precedence over rms+quant if activated.
|
||||
# That means we get full matching where ar+rms+quant was not activated,
|
||||
# and less where it was
|
||||
# That means we get full matching where ar+rms+quant was not
|
||||
# activated, and less where it was (only the smallest range).
|
||||
assert sum(m == expected_matches for m in log_matches) == tp_size * (
|
||||
num_ranges_activated - 1
|
||||
), "Expecting full rms+quant fusion where ar+rms+quant not activated"
|
||||
@@ -135,6 +156,43 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
|
||||
f"Expecting at least {expected_matches - matches.ar_rms_fusion} "
|
||||
f"where ar+rms+quant was activated"
|
||||
)
|
||||
elif (
|
||||
match_name == "async_tp"
|
||||
and "sequence_parallel" in matches_check
|
||||
and num_compile_ranges >= 2
|
||||
):
|
||||
# AsyncTP only finds patterns on ranges where SP ran.
|
||||
n_sp_ranges = num_compile_ranges - 1
|
||||
assert (
|
||||
sum(m == expected_matches for m in log_matches)
|
||||
== tp_size * n_sp_ranges
|
||||
), (
|
||||
f"Expecting {expected_matches} async_tp on "
|
||||
f"{tp_size * n_sp_ranges} SP-range entries, "
|
||||
f"found: {log_matches}"
|
||||
)
|
||||
assert sum(m == 0 for m in log_matches) == tp_size, (
|
||||
f"Expecting 0 async_tp on {tp_size} small-range entries "
|
||||
f"(no SP), found: {log_matches}"
|
||||
)
|
||||
elif (
|
||||
match_name == "ar_rms_fusion"
|
||||
and "sequence_parallel" in matches_check
|
||||
and num_compile_ranges >= 2
|
||||
):
|
||||
# SP consumes allreduce patterns first, so AR+RMS finds
|
||||
# full matches only on the smallest range (no SP).
|
||||
assert sum(m == expected_matches for m in log_matches) == tp_size, (
|
||||
f"Expecting {expected_matches} ar_rms on "
|
||||
f"{tp_size} small-range entries, found: {log_matches}"
|
||||
)
|
||||
assert sum(m == 0 for m in log_matches) == tp_size * (
|
||||
num_ranges_activated - 1
|
||||
), (
|
||||
f"Expecting 0 ar_rms on "
|
||||
f"{tp_size * (num_ranges_activated - 1)} large-range "
|
||||
f"entries (SP took precedence), found: {log_matches}"
|
||||
)
|
||||
else:
|
||||
expected_matches_list = [expected_matches] * n_expected
|
||||
assert sorted(log_matches) == expected_matches_list, (
|
||||
@@ -142,7 +200,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
|
||||
f"found: {sorted(log_matches)}"
|
||||
)
|
||||
|
||||
if match_name == "ar_rms_fusion":
|
||||
if match_name == "ar_rms_fusion" and num_compile_ranges >= 2:
|
||||
log_matches = re.findall(
|
||||
r"pass_manager.py:\d+] Skipping "
|
||||
r".*AllReduceFusionPass.* with compile range",
|
||||
@@ -155,4 +213,17 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
|
||||
f"(found {len(log_matches)}) in:\n {log_holder.text}"
|
||||
)
|
||||
|
||||
if match_name == "sequence_parallel" and num_compile_ranges >= 2:
|
||||
log_matches = re.findall(
|
||||
r"pass_manager.py:\d+] Skipping "
|
||||
r".*SequenceParallelismPass.* with compile range",
|
||||
log_holder.text,
|
||||
)
|
||||
|
||||
n_expected = tp_size * (num_compile_ranges - num_ranges_activated)
|
||||
assert len(log_matches) == n_expected, (
|
||||
f'Could not find {n_expected} "Skipping SequenceParallelismPass" '
|
||||
f"(found {len(log_matches)}) in:\n {log_holder.text}"
|
||||
)
|
||||
|
||||
return run
|
||||
|
||||
@@ -66,6 +66,9 @@ def test_tp2_async_tp_fp8_fusions(
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
enable_sp=True,
|
||||
fuse_gemm_comms=True,
|
||||
fuse_allreduce_rms=False,
|
||||
# Override threshold for testing (models have small hidden_size)
|
||||
sp_min_token_num=512,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -123,6 +126,9 @@ def test_tp2_async_tp_fusions(
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
enable_sp=True,
|
||||
fuse_gemm_comms=True,
|
||||
fuse_allreduce_rms=False,
|
||||
# Override threshold for testing (models have small hidden_size)
|
||||
sp_min_token_num=512,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -141,3 +147,130 @@ def test_tp2_async_tp_fusions(
|
||||
matches_check,
|
||||
tp_size=2,
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize(
|
||||
"model_name, matches_fn, model_kwargs, hf_overrides",
|
||||
[llama3_8b_fp8, llama4_scout_fp8],
|
||||
)
|
||||
@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
|
||||
@pytest.mark.parametrize("n_layers", [4])
|
||||
@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
|
||||
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
|
||||
def test_tp2_sp_ar_rms_fp8_fusions(
|
||||
model_name: str,
|
||||
matches_fn: Callable[[int], Matches],
|
||||
model_kwargs: dict,
|
||||
hf_overrides: Callable[[int], dict],
|
||||
attn_backend: AttentionBackendCase,
|
||||
n_layers: int,
|
||||
custom_ops: str,
|
||||
inductor_graph_partition: bool,
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
matches = matches_fn(n_layers)
|
||||
|
||||
if is_blackwell():
|
||||
# Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns
|
||||
monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel")
|
||||
|
||||
# Reduce size of model and skip weight loading time
|
||||
model_kwargs["hf_overrides"] = hf_overrides(n_layers)
|
||||
model_kwargs["load_format"] = "dummy"
|
||||
model_kwargs["max_model_len"] = 1024
|
||||
|
||||
compilation_config = dict(
|
||||
use_inductor_graph_partition=inductor_graph_partition,
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
fuse_attn_quant=True,
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
enable_sp=True,
|
||||
fuse_gemm_comms=True,
|
||||
fuse_allreduce_rms=True,
|
||||
# Override threshold for testing (models have small hidden_size)
|
||||
sp_min_token_num=512,
|
||||
),
|
||||
)
|
||||
|
||||
matches_check = [
|
||||
"rms_quant_fusion",
|
||||
"act_quant_fusion",
|
||||
"norm_rope_fusion",
|
||||
"attn_quant_fusion",
|
||||
"ar_rms_fusion",
|
||||
"sequence_parallel",
|
||||
"async_tp",
|
||||
]
|
||||
|
||||
run_e2e_fusion_test(
|
||||
model_name,
|
||||
matches,
|
||||
model_kwargs,
|
||||
attn_backend,
|
||||
compilation_config,
|
||||
matches_check,
|
||||
tp_size=2,
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize(
|
||||
"model_name, matches_fn, model_kwargs, hf_overrides",
|
||||
[llama3_8b, qwen3_a3b],
|
||||
)
|
||||
@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
|
||||
@pytest.mark.parametrize("n_layers", [4])
|
||||
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
|
||||
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
|
||||
def test_tp2_sp_ar_rms_fusions(
|
||||
model_name: str,
|
||||
matches_fn: Callable[[int], Matches],
|
||||
model_kwargs: dict,
|
||||
hf_overrides: Callable[[int], dict],
|
||||
attn_backend: AttentionBackendCase,
|
||||
n_layers: int,
|
||||
custom_ops: str,
|
||||
inductor_graph_partition: bool,
|
||||
run_e2e_fusion_test,
|
||||
):
|
||||
matches = matches_fn(n_layers)
|
||||
|
||||
# Reduce size of model and skip weight loading time
|
||||
model_kwargs["hf_overrides"] = hf_overrides(n_layers)
|
||||
model_kwargs["load_format"] = "dummy"
|
||||
model_kwargs["max_model_len"] = 1024
|
||||
|
||||
compilation_config = dict(
|
||||
use_inductor_graph_partition=inductor_graph_partition,
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
enable_sp=True,
|
||||
fuse_gemm_comms=True,
|
||||
fuse_allreduce_rms=True,
|
||||
# Override threshold for testing (models have small hidden_size)
|
||||
sp_min_token_num=512,
|
||||
),
|
||||
)
|
||||
|
||||
matches_check = [
|
||||
"norm_rope_fusion",
|
||||
"ar_rms_fusion",
|
||||
"sequence_parallel",
|
||||
"async_tp",
|
||||
]
|
||||
|
||||
run_e2e_fusion_test(
|
||||
model_name,
|
||||
matches,
|
||||
model_kwargs,
|
||||
attn_backend,
|
||||
compilation_config,
|
||||
matches_check,
|
||||
tp_size=2,
|
||||
)
|
||||
|
||||
@@ -142,7 +142,6 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
|
||||
*(scaled_fp4_quant(w, wg) for w, wg in zip(self.w, wgscale))
|
||||
)
|
||||
self.wq, self.wscale = list(wq_gen), list(wscale_gen)
|
||||
print(f"{self.wq=}, {self.wscale=}")
|
||||
|
||||
def forward(self, hidden_states):
|
||||
# avoid having graph input be an arg to a pattern directly
|
||||
@@ -199,6 +198,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
|
||||
@pytest.mark.parametrize("hidden_size", [64])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
|
||||
@pytest.mark.parametrize("flashinfer_allreduce_backend", ["trtllm", "mnnvl"])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
@pytest.mark.skipif(
|
||||
not find_spec("flashinfer")
|
||||
@@ -215,6 +215,7 @@ def test_all_reduce_fusion_pass_replace(
|
||||
dtype: torch.dtype,
|
||||
enable_rms_norm_custom_op,
|
||||
enable_quant_fp8_custom_op,
|
||||
flashinfer_allreduce_backend,
|
||||
):
|
||||
num_processes = 2
|
||||
if (
|
||||
@@ -238,6 +239,7 @@ def test_all_reduce_fusion_pass_replace(
|
||||
dtype,
|
||||
enable_rms_norm_custom_op,
|
||||
enable_quant_fp8_custom_op,
|
||||
flashinfer_allreduce_backend,
|
||||
),
|
||||
nprocs=nprocs,
|
||||
)
|
||||
@@ -255,6 +257,7 @@ def all_reduce_fusion_pass_on_test_model(
|
||||
dtype: torch.dtype,
|
||||
enable_rms_norm_custom_op,
|
||||
enable_quant_fp8_custom_op,
|
||||
flashinfer_allreduce_backend,
|
||||
):
|
||||
set_random_seed(0)
|
||||
|
||||
@@ -270,6 +273,7 @@ def all_reduce_fusion_pass_on_test_model(
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -317,6 +321,10 @@ def all_reduce_fusion_pass_on_test_model(
|
||||
compiled_model = torch.compile(model, backend=backend)
|
||||
compiled_model(hidden_states)
|
||||
|
||||
results_unfused = model(hidden_states)
|
||||
results_fused = compiled_model(hidden_states)
|
||||
torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2)
|
||||
|
||||
assert all_reduce_fusion_pass.matched_count == 4, (
|
||||
f"{all_reduce_fusion_pass.matched_count=}"
|
||||
)
|
||||
|
||||
@@ -421,6 +421,7 @@ def test_cudagraph_sizes_post_init(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
eliminate_noops=True,
|
||||
sp_min_token_num=512 if enable_sp else None,
|
||||
),
|
||||
cudagraph_mode=cudagraph_mode,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.compilation.passes.fusion.sequence_parallelism import (
|
||||
SP_MIN_HIDDEN_SIZE,
|
||||
SP_MIN_PER_GPU_SIZE_MB,
|
||||
get_sequence_parallelism_threshold,
|
||||
)
|
||||
|
||||
|
||||
class TestGetSequenceParallelismThreshold:
|
||||
"""Tests for get_sequence_parallelism_threshold function."""
|
||||
|
||||
def test_non_cuda_returns_none(self, mock_cuda_platform):
|
||||
"""Non-CUDA platforms should return None."""
|
||||
with mock_cuda_platform(is_cuda=False):
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=8192, tp_size=2, element_size=2
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_unsupported_device_capability_returns_none(self, mock_cuda_platform):
|
||||
"""Unsupported device capabilities (e.g., sm80) should return None."""
|
||||
with mock_cuda_platform(capability=(8, 0)):
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=8192, tp_size=2, element_size=2
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_small_hidden_size_returns_none(self, mock_cuda_platform):
|
||||
"""H100 with hidden_size below threshold should return None."""
|
||||
with mock_cuda_platform(capability=(9, 0)):
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=4096,
|
||||
tp_size=2,
|
||||
element_size=2, # 4096 < 8192
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_h100_large_model_returns_threshold(self, mock_cuda_platform):
|
||||
"""H100 with large enough hidden_size should return calculated threshold."""
|
||||
with mock_cuda_platform(capability=(9, 0)):
|
||||
hidden_size = 8192
|
||||
tp_size = 2
|
||||
element_size = 2 # float16/bfloat16
|
||||
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=hidden_size,
|
||||
tp_size=tp_size,
|
||||
element_size=element_size,
|
||||
)
|
||||
|
||||
# Verify calculation: (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024
|
||||
MiB = 1024 * 1024
|
||||
expected = int(
|
||||
(SP_MIN_PER_GPU_SIZE_MB[90] * tp_size * MiB)
|
||||
// (hidden_size * element_size)
|
||||
)
|
||||
assert result == expected
|
||||
assert result == 1024
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"hidden_size,tp_size,element_size,expected",
|
||||
[
|
||||
# Boundary: exactly at min hidden size threshold, tp_size=1
|
||||
# (8 * 1 * 1024 * 1024) // (8192 * 2) = 512
|
||||
(8192, 1, 2, 512),
|
||||
# Larger hidden size reduces token threshold
|
||||
# (8 * 1 * 1024 * 1024) // (16384 * 2) = 256
|
||||
(16384, 1, 2, 256),
|
||||
# Larger tp_size increases token threshold
|
||||
# (8 * 4 * 1024 * 1024) // (8192 * 2) = 2048
|
||||
(8192, 4, 2, 2048),
|
||||
# Larger element_size (fp32) reduces token threshold
|
||||
# (8 * 2 * 1024 * 1024) // (8192 * 4) = 512
|
||||
(8192, 2, 4, 512),
|
||||
],
|
||||
)
|
||||
def test_threshold_calculation_variations(
|
||||
self, mock_cuda_platform, hidden_size, tp_size, element_size, expected
|
||||
):
|
||||
"""Test threshold calculation with various parameter combinations."""
|
||||
with mock_cuda_platform(capability=(9, 0)):
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=hidden_size,
|
||||
tp_size=tp_size,
|
||||
element_size=element_size,
|
||||
)
|
||||
assert result == expected
|
||||
|
||||
def test_hidden_size_boundary(self, mock_cuda_platform):
|
||||
"""Test behavior at the exact hidden_size boundary."""
|
||||
with mock_cuda_platform(capability=(9, 0)):
|
||||
# Just below threshold
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=SP_MIN_HIDDEN_SIZE[90] - 1,
|
||||
tp_size=2,
|
||||
element_size=2,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Exactly at threshold
|
||||
result = get_sequence_parallelism_threshold(
|
||||
hidden_size=SP_MIN_HIDDEN_SIZE[90],
|
||||
tp_size=2,
|
||||
element_size=2,
|
||||
)
|
||||
assert result is not None
|
||||
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for CUDA forward compatibility path logic in env_override.py.
|
||||
|
||||
Verifies the opt-in LD_LIBRARY_PATH manipulation for CUDA compat libs,
|
||||
including env var parsing, path detection, and deduplication.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Import the functions directly (they're module-level in env_override)
|
||||
# We must import them without triggering the module-level side effects,
|
||||
# so we import the functions by name after the module is already loaded.
|
||||
from vllm.env_override import (
|
||||
_get_torch_cuda_version,
|
||||
_maybe_set_cuda_compatibility_path,
|
||||
)
|
||||
|
||||
|
||||
class TestCudaCompatibilityEnvParsing:
|
||||
"""Test VLLM_ENABLE_CUDA_COMPATIBILITY env var parsing."""
|
||||
|
||||
def test_disabled_by_default(self, monkeypatch):
|
||||
"""Compat path is NOT set when env var is absent."""
|
||||
monkeypatch.delenv("VLLM_ENABLE_CUDA_COMPATIBILITY", raising=False)
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
assert (
|
||||
"LD_LIBRARY_PATH" not in os.environ
|
||||
or os.environ.get("LD_LIBRARY_PATH", "") == ""
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "False", "no", ""])
|
||||
def test_disabled_values(self, monkeypatch, value):
|
||||
"""Various falsy values should not activate compat path."""
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", value)
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
# LD_LIBRARY_PATH should not be set (or remain empty)
|
||||
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
assert "compat" not in ld_path
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "True", " 1 ", " TRUE "])
|
||||
def test_enabled_values_with_valid_path(self, monkeypatch, tmp_path, value):
|
||||
"""Truthy values activate compat path when a valid path exists."""
|
||||
compat_dir = tmp_path / "compat"
|
||||
compat_dir.mkdir()
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", value)
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir))
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
assert str(compat_dir) in ld_path
|
||||
|
||||
|
||||
class TestCudaCompatibilityPathDetection:
|
||||
"""Test path detection: custom override, conda, default."""
|
||||
|
||||
def test_custom_path_override(self, monkeypatch, tmp_path):
|
||||
"""VLLM_CUDA_COMPATIBILITY_PATH takes highest priority."""
|
||||
custom_dir = tmp_path / "my-compat"
|
||||
custom_dir.mkdir()
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(custom_dir))
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
assert ld_path.startswith(str(custom_dir))
|
||||
|
||||
def test_conda_prefix_fallback(self, monkeypatch, tmp_path):
|
||||
"""Falls back to $CONDA_PREFIX/cuda-compat if custom not set."""
|
||||
conda_dir = tmp_path / "conda-env"
|
||||
compat_dir = conda_dir / "cuda-compat"
|
||||
compat_dir.mkdir(parents=True)
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.delenv("VLLM_CUDA_COMPATIBILITY_PATH", raising=False)
|
||||
monkeypatch.setenv("CONDA_PREFIX", str(conda_dir))
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
assert str(compat_dir) in ld_path
|
||||
|
||||
def test_no_valid_path_does_nothing(self, monkeypatch):
|
||||
"""When enabled but no valid path exists, LD_LIBRARY_PATH unchanged."""
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", "/nonexistent/path")
|
||||
monkeypatch.delenv("CONDA_PREFIX", raising=False)
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
with patch("vllm.env_override._get_torch_cuda_version", return_value=None):
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
assert os.environ.get("LD_LIBRARY_PATH", "") == ""
|
||||
|
||||
def test_default_cuda_path_fallback(self, monkeypatch, tmp_path):
|
||||
"""Falls back to /usr/local/cuda-{ver}/compat via torch version."""
|
||||
fake_cuda = tmp_path / "cuda-12.8" / "compat"
|
||||
fake_cuda.mkdir(parents=True)
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.delenv("VLLM_CUDA_COMPATIBILITY_PATH", raising=False)
|
||||
monkeypatch.delenv("CONDA_PREFIX", raising=False)
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
with (
|
||||
patch("vllm.env_override._get_torch_cuda_version", return_value="12.8"),
|
||||
patch(
|
||||
"vllm.env_override.os.path.isdir",
|
||||
side_effect=lambda p: p == "/usr/local/cuda-12.8/compat"
|
||||
or os.path.isdir(p),
|
||||
),
|
||||
):
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
assert "/usr/local/cuda-12.8/compat" in ld_path
|
||||
|
||||
|
||||
class TestCudaCompatibilityLdPathManipulation:
|
||||
"""Test LD_LIBRARY_PATH prepend and deduplication logic."""
|
||||
|
||||
def test_prepends_to_empty_ld_path(self, monkeypatch, tmp_path):
|
||||
"""Compat path is set when LD_LIBRARY_PATH is empty."""
|
||||
compat_dir = tmp_path / "compat"
|
||||
compat_dir.mkdir()
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir))
|
||||
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
assert os.environ["LD_LIBRARY_PATH"] == str(compat_dir)
|
||||
|
||||
def test_prepends_to_existing_ld_path(self, monkeypatch, tmp_path):
|
||||
"""Compat path is prepended before existing entries."""
|
||||
compat_dir = tmp_path / "compat"
|
||||
compat_dir.mkdir()
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir))
|
||||
monkeypatch.setenv("LD_LIBRARY_PATH", "/usr/lib:/other/lib")
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
ld_path = os.environ["LD_LIBRARY_PATH"]
|
||||
parts = ld_path.split(os.pathsep)
|
||||
assert parts[0] == str(compat_dir)
|
||||
assert "/usr/lib" in parts
|
||||
assert "/other/lib" in parts
|
||||
|
||||
def test_deduplicates_existing_compat_path(self, monkeypatch, tmp_path):
|
||||
"""If compat path already in LD_LIBRARY_PATH, move to front."""
|
||||
compat_dir = tmp_path / "compat"
|
||||
compat_dir.mkdir()
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir))
|
||||
monkeypatch.setenv(
|
||||
"LD_LIBRARY_PATH",
|
||||
f"/usr/lib:{compat_dir}:/other/lib",
|
||||
)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
ld_path = os.environ["LD_LIBRARY_PATH"]
|
||||
parts = ld_path.split(os.pathsep)
|
||||
assert parts[0] == str(compat_dir)
|
||||
assert parts.count(str(compat_dir)) == 1
|
||||
|
||||
def test_already_at_front_is_noop(self, monkeypatch, tmp_path):
|
||||
"""If compat path is already first, don't modify LD_LIBRARY_PATH."""
|
||||
compat_dir = tmp_path / "compat"
|
||||
compat_dir.mkdir()
|
||||
original = f"{compat_dir}:/usr/lib"
|
||||
monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1")
|
||||
monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir))
|
||||
monkeypatch.setenv("LD_LIBRARY_PATH", original)
|
||||
_maybe_set_cuda_compatibility_path()
|
||||
assert os.environ["LD_LIBRARY_PATH"] == original
|
||||
|
||||
|
||||
class TestGetTorchCudaVersion:
|
||||
"""Test _get_torch_cuda_version() helper."""
|
||||
|
||||
def test_returns_string_when_torch_available(self):
|
||||
"""Should return a CUDA version string like '12.8'."""
|
||||
version = _get_torch_cuda_version()
|
||||
# torch is installed in vllm's environment
|
||||
assert version is None or isinstance(version, str)
|
||||
|
||||
def test_returns_none_when_torch_missing(self):
|
||||
"""Should return None when torch is not importable."""
|
||||
with patch(
|
||||
"vllm.env_override.importlib.util.find_spec",
|
||||
return_value=None,
|
||||
):
|
||||
assert _get_torch_cuda_version() is None
|
||||
@@ -2,31 +2,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import ResponseFunctionToolCall, ResponseReasoningItem
|
||||
from openai.types.responses.response_output_item import McpCall
|
||||
from openai_harmony import Author, Message, Role, TextContent
|
||||
from openai_harmony import Message, Role
|
||||
|
||||
from tests.entrypoints.openai.utils import verify_harmony_messages
|
||||
from vllm.entrypoints.openai.parser.harmony_utils import (
|
||||
auto_drop_analysis_messages,
|
||||
get_encoding,
|
||||
get_system_message,
|
||||
has_custom_tools,
|
||||
parse_chat_input_to_harmony_message,
|
||||
parse_chat_output,
|
||||
parse_input_to_harmony_message,
|
||||
parse_output_message,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.harmony import (
|
||||
response_previous_input_to_harmony,
|
||||
)
|
||||
|
||||
|
||||
class TestCommonParseInputToHarmonyMessage:
|
||||
"""
|
||||
Tests for scenarios that are common to both Chat Completion
|
||||
parse_chat_input_to_harmony_message and Responsees API
|
||||
parse_input_to_harmony_message functions.
|
||||
parse_chat_input_to_harmony_message and Responses API
|
||||
response_previous_input_to_harmony functions.
|
||||
"""
|
||||
|
||||
@pytest.fixture(
|
||||
params=[parse_chat_input_to_harmony_message, parse_input_to_harmony_message]
|
||||
params=[parse_chat_input_to_harmony_message, response_previous_input_to_harmony]
|
||||
)
|
||||
def parse_function(self, request):
|
||||
return request.param
|
||||
@@ -211,81 +211,6 @@ class TestCommonParseInputToHarmonyMessage:
|
||||
assert messages[0].content[1].text == "actual text"
|
||||
|
||||
|
||||
class TestParseInputToHarmonyMessage:
|
||||
"""
|
||||
Tests for scenarios that are specific to the Responses API
|
||||
parse_input_to_harmony_message function.
|
||||
"""
|
||||
|
||||
def test_message_with_empty_content(self):
|
||||
"""Test parsing message with empty string content."""
|
||||
chat_msg = {
|
||||
"role": "user",
|
||||
"content": "",
|
||||
}
|
||||
|
||||
messages = parse_input_to_harmony_message(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].content[0].text == ""
|
||||
|
||||
def test_tool_message_with_string_content(self):
|
||||
"""Test parsing tool message with string content."""
|
||||
chat_msg = {
|
||||
"role": "tool",
|
||||
"name": "get_weather",
|
||||
"content": "The weather in San Francisco is sunny, 72°F",
|
||||
}
|
||||
|
||||
messages = parse_input_to_harmony_message(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].author.role == Role.TOOL
|
||||
assert messages[0].author.name == "functions.get_weather"
|
||||
assert (
|
||||
messages[0].content[0].text == "The weather in San Francisco is sunny, 72°F"
|
||||
)
|
||||
assert messages[0].channel == "commentary"
|
||||
|
||||
def test_tool_message_with_array_content(self):
|
||||
"""Test parsing tool message with array content."""
|
||||
chat_msg = {
|
||||
"role": "tool",
|
||||
"name": "search_results",
|
||||
"content": [
|
||||
{"type": "text", "text": "Result 1: "},
|
||||
{"type": "text", "text": "Result 2: "},
|
||||
{
|
||||
"type": "image",
|
||||
"url": "http://example.com/img.png",
|
||||
}, # Should be ignored
|
||||
{"type": "text", "text": "Result 3"},
|
||||
],
|
||||
}
|
||||
|
||||
messages = parse_input_to_harmony_message(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].author.role == Role.TOOL
|
||||
assert messages[0].author.name == "functions.search_results"
|
||||
assert messages[0].content[0].text == "Result 1: Result 2: Result 3"
|
||||
|
||||
def test_tool_message_with_empty_content(self):
|
||||
"""Test parsing tool message with None content."""
|
||||
chat_msg = {
|
||||
"role": "tool",
|
||||
"name": "empty_tool",
|
||||
"content": None,
|
||||
}
|
||||
|
||||
messages = parse_input_to_harmony_message(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].author.role == Role.TOOL
|
||||
assert messages[0].author.name == "functions.empty_tool"
|
||||
assert messages[0].content[0].text == ""
|
||||
|
||||
|
||||
class TestParseChatInputToHarmonyMessage:
|
||||
"""
|
||||
Tests for scenarios that are specific to the Chat Completion API
|
||||
@@ -840,192 +765,47 @@ class TestParseChatOutput:
|
||||
assert reasoning == "I've thought hard about this."
|
||||
assert final_content == "The answer is 4."
|
||||
|
||||
def test_parse_chat_output_commentary_with_recipient_excluded(self) -> None:
|
||||
"""Commentary with a recipient (tool call) should not appear in
|
||||
final_content — those are handled separately by the tool parser.
|
||||
|
||||
class TestParseOutputMessage:
|
||||
"""Tests for parse_output_message function."""
|
||||
|
||||
def test_commentary_with_no_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient=None (preambles) creates reasoning items.
|
||||
|
||||
Per Harmony format, commentary channel can contain preambles to calling
|
||||
multiple functions - explanatory text with no recipient.
|
||||
The first message is a preamble (visible), the second is a tool
|
||||
call (excluded). Only the preamble should appear in final_content.
|
||||
"""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "I will now search for the weather information."
|
||||
harmony_str = (
|
||||
"<|channel|>commentary"
|
||||
"<|message|>Let me check the weather.<|end|>"
|
||||
"<|start|>assistant to=functions.get_weather"
|
||||
"<|channel|>commentary"
|
||||
'<|message|>{"location": "SF"}<|end|>'
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
# recipient is None by default, representing a preamble
|
||||
token_ids = get_encoding().encode(harmony_str, allowed_special="all")
|
||||
reasoning, final_content, _ = parse_chat_output(token_ids)
|
||||
assert reasoning is None
|
||||
assert final_content == "Let me check the weather."
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
def test_parse_chat_output_interrupted_preamble(self) -> None:
|
||||
"""Partial/interrupted preamble (commentary without recipient) should
|
||||
appear in final_content, not reasoning."""
|
||||
harmony_str = "<|channel|>commentary<|message|>I'll search for that"
|
||||
token_ids = get_encoding().encode(harmony_str, allowed_special="all")
|
||||
reasoning, final_content, _ = parse_chat_output(token_ids)
|
||||
assert reasoning is None
|
||||
assert final_content == "I'll search for that"
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert (
|
||||
output_items[0].content[0].text
|
||||
== "I will now search for the weather information."
|
||||
def test_parse_chat_output_preamble_then_final(self) -> None:
|
||||
"""Preamble followed by a final message should both appear in
|
||||
final_content, joined by newline."""
|
||||
harmony_str = (
|
||||
"<|channel|>commentary"
|
||||
"<|message|>Let me look that up.<|end|>"
|
||||
"<|start|>assistant<|channel|>final"
|
||||
"<|message|>The answer is 42.<|end|>"
|
||||
)
|
||||
assert output_items[0].content[0].type == "reasoning_text"
|
||||
|
||||
def test_commentary_with_function_recipient_creates_function_call(self):
|
||||
"""Test commentary with recipient='functions.X' creates function calls."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, '{"location": "San Francisco", "units": "celsius"}'
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].type == "function_call"
|
||||
assert output_items[0].name == "get_weather"
|
||||
assert (
|
||||
output_items[0].arguments
|
||||
== '{"location": "San Francisco", "units": "celsius"}'
|
||||
)
|
||||
assert output_items[0].call_id.startswith("call_")
|
||||
assert output_items[0].id.startswith("fc_")
|
||||
|
||||
def test_commentary_with_python_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='python' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "import numpy as np\nprint(np.array([1, 2, 3]))"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("python")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert (
|
||||
output_items[0].content[0].text
|
||||
== "import numpy as np\nprint(np.array([1, 2, 3]))"
|
||||
)
|
||||
|
||||
def test_commentary_with_browser_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='browser' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Navigating to the specified URL"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("browser")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == "Navigating to the specified URL"
|
||||
|
||||
def test_commentary_with_container_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='container' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Running command in container"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("container")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == "Running command in container"
|
||||
|
||||
def test_commentary_with_empty_content_and_no_recipient(self):
|
||||
"""Test edge case: empty commentary with recipient=None."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, "")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].content[0].text == ""
|
||||
|
||||
def test_commentary_with_multiple_contents_and_no_recipient(self):
|
||||
"""Test multiple content items in commentary with no recipient."""
|
||||
contents = [
|
||||
TextContent(text="Step 1: Analyze the request"),
|
||||
TextContent(text="Step 2: Prepare to call functions"),
|
||||
]
|
||||
message = Message.from_role_and_contents(Role.ASSISTANT, contents)
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 2
|
||||
assert all(isinstance(item, ResponseReasoningItem) for item in output_items)
|
||||
assert output_items[0].content[0].text == "Step 1: Analyze the request"
|
||||
assert output_items[1].content[0].text == "Step 2: Prepare to call functions"
|
||||
|
||||
def test_commentary_with_multiple_function_calls(self):
|
||||
"""Test multiple function calls in commentary channel."""
|
||||
contents = [
|
||||
TextContent(text='{"location": "San Francisco"}'),
|
||||
TextContent(text='{"location": "New York"}'),
|
||||
]
|
||||
message = Message.from_role_and_contents(Role.ASSISTANT, contents)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 2
|
||||
assert all(isinstance(item, ResponseFunctionToolCall) for item in output_items)
|
||||
assert output_items[0].name == "get_weather"
|
||||
assert output_items[1].name == "get_weather"
|
||||
assert output_items[0].arguments == '{"location": "San Francisco"}'
|
||||
assert output_items[1].arguments == '{"location": "New York"}'
|
||||
|
||||
def test_commentary_with_unknown_recipient_creates_mcp_call(self):
|
||||
"""Test that commentary with unknown recipient creates MCP call."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("custom_tool")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == "custom_tool"
|
||||
assert output_items[0].server_label == "custom_tool"
|
||||
|
||||
def test_analysis_channel_creates_reasoning(self):
|
||||
"""Test that analysis channel creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Analyzing the problem step by step..."
|
||||
)
|
||||
message = message.with_channel("analysis")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert (
|
||||
output_items[0].content[0].text == "Analyzing the problem step by step..."
|
||||
)
|
||||
|
||||
def test_non_assistant_message_returns_empty(self):
|
||||
"""Test that non-assistant messages return empty list.
|
||||
|
||||
Per the implementation, tool messages to assistant (e.g., search results)
|
||||
are not included in final output to align with OpenAI behavior.
|
||||
"""
|
||||
message = Message.from_author_and_content(
|
||||
Author.new(Role.TOOL, "functions.get_weather"),
|
||||
"The weather is sunny, 72°F",
|
||||
)
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 0
|
||||
token_ids = get_encoding().encode(harmony_str, allowed_special="all")
|
||||
reasoning, final_content, _ = parse_chat_output(token_ids)
|
||||
assert reasoning is None
|
||||
assert final_content == "Let me look that up.\nThe answer is 42."
|
||||
|
||||
|
||||
def test_has_custom_tools() -> None:
|
||||
@@ -1037,165 +817,27 @@ def test_has_custom_tools() -> None:
|
||||
)
|
||||
|
||||
|
||||
def test_parse_mcp_call_basic() -> None:
|
||||
"""Test that MCP calls are parsed with correct type and server_label."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"path": "/tmp"}')
|
||||
message = message.with_recipient("filesystem")
|
||||
message = message.with_channel("commentary")
|
||||
class TestGetSystemMessage:
|
||||
"""Tests for get_system_message channel configuration."""
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
def test_commentary_channel_present_without_custom_tools(self) -> None:
|
||||
"""Commentary channel must be valid even without custom tools."""
|
||||
sys_msg = get_system_message(with_custom_tools=False)
|
||||
valid_channels = sys_msg.content[0].channel_config.valid_channels
|
||||
assert "commentary" in valid_channels
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == "filesystem"
|
||||
assert output_items[0].server_label == "filesystem"
|
||||
assert output_items[0].arguments == '{"path": "/tmp"}'
|
||||
assert output_items[0].status == "completed"
|
||||
def test_commentary_channel_present_with_custom_tools(self) -> None:
|
||||
"""Commentary channel present when custom tools are enabled."""
|
||||
sys_msg = get_system_message(with_custom_tools=True)
|
||||
valid_channels = sys_msg.content[0].channel_config.valid_channels
|
||||
assert "commentary" in valid_channels
|
||||
|
||||
|
||||
def test_parse_mcp_call_dotted_recipient() -> None:
|
||||
"""Test that dotted recipients extract the tool name correctly."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"cmd": "ls"}')
|
||||
message = message.with_recipient("repo_browser.list")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = parse_output_message(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].name == "list"
|
||||
assert output_items[0].server_label == "repo_browser"
|
||||
|
||||
|
||||
def test_mcp_vs_function_call() -> None:
|
||||
"""Test that function calls are not parsed as MCP calls."""
|
||||
func_message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
func_message = func_message.with_recipient("functions.my_tool")
|
||||
func_message = func_message.with_channel("commentary")
|
||||
|
||||
func_items = parse_output_message(func_message)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
|
||||
|
||||
def test_mcp_vs_builtin_tools() -> None:
|
||||
"""Test that built-in tools (python, container) are not parsed as MCP calls."""
|
||||
# Test python (built-in tool) - should be reasoning, not MCP
|
||||
python_message = Message.from_role_and_content(Role.ASSISTANT, "print('hello')")
|
||||
python_message = python_message.with_recipient("python")
|
||||
python_message = python_message.with_channel("commentary")
|
||||
|
||||
python_items = parse_output_message(python_message)
|
||||
|
||||
assert len(python_items) == 1
|
||||
assert not isinstance(python_items[0], McpCall)
|
||||
assert python_items[0].type == "reasoning"
|
||||
|
||||
|
||||
def test_parse_remaining_state_commentary_channel() -> None:
|
||||
"""Test parse_remaining_state with commentary channel and various recipients."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
from vllm.entrypoints.openai.parser.harmony_utils import parse_remaining_state
|
||||
|
||||
# Test 1: functions.* recipient → should return function tool call
|
||||
parser_func = Mock()
|
||||
parser_func.current_content = '{"arg": "value"}'
|
||||
parser_func.current_role = Role.ASSISTANT
|
||||
parser_func.current_channel = "commentary"
|
||||
parser_func.current_recipient = "functions.my_tool"
|
||||
|
||||
func_items = parse_remaining_state(parser_func)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
assert func_items[0].name == "my_tool"
|
||||
assert func_items[0].status == "in_progress"
|
||||
|
||||
# Test 2: MCP tool (not builtin) → should return MCP call
|
||||
parser_mcp = Mock()
|
||||
parser_mcp.current_content = '{"path": "/tmp"}'
|
||||
parser_mcp.current_role = Role.ASSISTANT
|
||||
parser_mcp.current_channel = "commentary"
|
||||
parser_mcp.current_recipient = "filesystem"
|
||||
|
||||
mcp_items = parse_remaining_state(parser_mcp)
|
||||
|
||||
assert len(mcp_items) == 1
|
||||
assert isinstance(mcp_items[0], McpCall)
|
||||
assert mcp_items[0].type == "mcp_call"
|
||||
assert mcp_items[0].name == "filesystem"
|
||||
assert mcp_items[0].server_label == "filesystem"
|
||||
assert mcp_items[0].status == "in_progress"
|
||||
|
||||
# Test 3: Built-in tool (python)
|
||||
# should NOT return MCP call, falls through to reasoning
|
||||
parser_builtin = Mock()
|
||||
parser_builtin.current_content = "print('hello')"
|
||||
parser_builtin.current_role = Role.ASSISTANT
|
||||
parser_builtin.current_channel = "commentary"
|
||||
parser_builtin.current_recipient = "python"
|
||||
|
||||
builtin_items = parse_remaining_state(parser_builtin)
|
||||
|
||||
# Should fall through to reasoning logic
|
||||
assert len(builtin_items) == 1
|
||||
assert not isinstance(builtin_items[0], McpCall)
|
||||
assert builtin_items[0].type == "reasoning"
|
||||
|
||||
|
||||
def test_parse_remaining_state_analysis_channel() -> None:
|
||||
"""Test parse_remaining_state with analysis channel and various recipients."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
from vllm.entrypoints.openai.parser.harmony_utils import parse_remaining_state
|
||||
|
||||
# Test 1: functions.* recipient → should return function tool call
|
||||
parser_func = Mock()
|
||||
parser_func.current_content = '{"arg": "value"}'
|
||||
parser_func.current_role = Role.ASSISTANT
|
||||
parser_func.current_channel = "analysis"
|
||||
parser_func.current_recipient = "functions.my_tool"
|
||||
|
||||
func_items = parse_remaining_state(parser_func)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
assert func_items[0].name == "my_tool"
|
||||
assert func_items[0].status == "in_progress"
|
||||
|
||||
# Test 2: MCP tool (not builtin) → should return MCP call
|
||||
parser_mcp = Mock()
|
||||
parser_mcp.current_content = '{"query": "test"}'
|
||||
parser_mcp.current_role = Role.ASSISTANT
|
||||
parser_mcp.current_channel = "analysis"
|
||||
parser_mcp.current_recipient = "database"
|
||||
|
||||
mcp_items = parse_remaining_state(parser_mcp)
|
||||
|
||||
assert len(mcp_items) == 1
|
||||
assert isinstance(mcp_items[0], McpCall)
|
||||
assert mcp_items[0].type == "mcp_call"
|
||||
assert mcp_items[0].name == "database"
|
||||
assert mcp_items[0].server_label == "database"
|
||||
assert mcp_items[0].status == "in_progress"
|
||||
|
||||
# Test 3: Built-in tool (container)
|
||||
# should NOT return MCP call, falls through to reasoning
|
||||
parser_builtin = Mock()
|
||||
parser_builtin.current_content = "docker run"
|
||||
parser_builtin.current_role = Role.ASSISTANT
|
||||
parser_builtin.current_channel = "analysis"
|
||||
parser_builtin.current_recipient = "container"
|
||||
|
||||
builtin_items = parse_remaining_state(parser_builtin)
|
||||
|
||||
# Should fall through to reasoning logic
|
||||
assert len(builtin_items) == 1
|
||||
assert not isinstance(builtin_items[0], McpCall)
|
||||
assert builtin_items[0].type == "reasoning"
|
||||
def test_all_standard_channels_present(self) -> None:
|
||||
"""All three standard Harmony channels should always be valid."""
|
||||
for with_tools in (True, False):
|
||||
sys_msg = get_system_message(with_custom_tools=with_tools)
|
||||
valid_channels = sys_msg.content[0].channel_config.valid_channels
|
||||
for channel in ("analysis", "commentary", "final"):
|
||||
assert channel in valid_channels, (
|
||||
f"{channel} missing when with_custom_tools={with_tools}"
|
||||
)
|
||||
|
||||
@@ -712,15 +712,14 @@ async def test_function_calling_required(client: OpenAI, model_name: str):
|
||||
async def test_system_message_with_tools(client: OpenAI, model_name: str):
|
||||
from vllm.entrypoints.openai.parser.harmony_utils import get_system_message
|
||||
|
||||
# Test with custom tools enabled - commentary channel should be available
|
||||
sys_msg = get_system_message(with_custom_tools=True)
|
||||
valid_channels = sys_msg.content[0].channel_config.valid_channels
|
||||
assert "commentary" in valid_channels
|
||||
|
||||
# Test with custom tools disabled - commentary channel should be removed
|
||||
sys_msg = get_system_message(with_custom_tools=False)
|
||||
valid_channels = sys_msg.content[0].channel_config.valid_channels
|
||||
assert "commentary" not in valid_channels
|
||||
# Commentary channel should always be present (needed for preambles)
|
||||
# regardless of whether custom tools are enabled
|
||||
for with_tools in (True, False):
|
||||
sys_msg = get_system_message(with_custom_tools=with_tools)
|
||||
valid_channels = sys_msg.content[0].channel_config.valid_channels
|
||||
assert "commentary" in valid_channels, (
|
||||
f"commentary channel missing when with_custom_tools={with_tools}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -0,0 +1,463 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for vllm.entrypoints.openai.responses.harmony."""
|
||||
|
||||
from openai.types.responses import (
|
||||
ResponseFunctionToolCall,
|
||||
ResponseOutputMessage,
|
||||
ResponseReasoningItem,
|
||||
)
|
||||
from openai.types.responses.response_output_item import McpCall
|
||||
from openai_harmony import Author, Message, Role, TextContent
|
||||
|
||||
from vllm.entrypoints.openai.responses.harmony import (
|
||||
harmony_to_response_output,
|
||||
parser_state_to_response_output,
|
||||
response_previous_input_to_harmony,
|
||||
)
|
||||
|
||||
|
||||
class TestResponsePreviousInputToHarmony:
|
||||
"""
|
||||
Tests for scenarios that are specific to the Responses API
|
||||
response_previous_input_to_harmony function.
|
||||
"""
|
||||
|
||||
def test_message_with_empty_content(self):
|
||||
"""Test parsing message with empty string content."""
|
||||
chat_msg = {
|
||||
"role": "user",
|
||||
"content": "",
|
||||
}
|
||||
|
||||
messages = response_previous_input_to_harmony(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].content[0].text == ""
|
||||
|
||||
def test_tool_message_with_string_content(self):
|
||||
"""Test parsing tool message with string content."""
|
||||
chat_msg = {
|
||||
"role": "tool",
|
||||
"name": "get_weather",
|
||||
"content": "The weather in San Francisco is sunny, 72°F",
|
||||
}
|
||||
|
||||
messages = response_previous_input_to_harmony(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].author.role == Role.TOOL
|
||||
assert messages[0].author.name == "functions.get_weather"
|
||||
assert (
|
||||
messages[0].content[0].text == "The weather in San Francisco is sunny, 72°F"
|
||||
)
|
||||
assert messages[0].channel == "commentary"
|
||||
|
||||
def test_tool_message_with_array_content(self):
|
||||
"""Test parsing tool message with array content."""
|
||||
chat_msg = {
|
||||
"role": "tool",
|
||||
"name": "search_results",
|
||||
"content": [
|
||||
{"type": "text", "text": "Result 1: "},
|
||||
{"type": "text", "text": "Result 2: "},
|
||||
{
|
||||
"type": "image",
|
||||
"url": "http://example.com/img.png",
|
||||
}, # Should be ignored
|
||||
{"type": "text", "text": "Result 3"},
|
||||
],
|
||||
}
|
||||
|
||||
messages = response_previous_input_to_harmony(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].author.role == Role.TOOL
|
||||
assert messages[0].author.name == "functions.search_results"
|
||||
assert messages[0].content[0].text == "Result 1: Result 2: Result 3"
|
||||
|
||||
def test_tool_message_with_empty_content(self):
|
||||
"""Test parsing tool message with None content."""
|
||||
chat_msg = {
|
||||
"role": "tool",
|
||||
"name": "empty_tool",
|
||||
"content": None,
|
||||
}
|
||||
|
||||
messages = response_previous_input_to_harmony(chat_msg)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].author.role == Role.TOOL
|
||||
assert messages[0].author.name == "functions.empty_tool"
|
||||
assert messages[0].content[0].text == ""
|
||||
|
||||
|
||||
class TestHarmonyToResponseOutput:
|
||||
"""Tests for harmony_to_response_output function."""
|
||||
|
||||
def test_commentary_with_no_recipient_creates_message(self):
|
||||
"""Test that commentary with recipient=None (preambles) creates message items.
|
||||
|
||||
Per Harmony format, preambles are intended to be shown to end-users,
|
||||
unlike analysis channel content which is hidden reasoning.
|
||||
See: https://cookbook.openai.com/articles/openai-harmony
|
||||
"""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "I will now search for the weather information."
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
# recipient is None by default, representing a preamble
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseOutputMessage)
|
||||
assert output_items[0].type == "message"
|
||||
assert output_items[0].role == "assistant"
|
||||
assert output_items[0].status == "completed"
|
||||
assert len(output_items[0].content) == 1
|
||||
assert output_items[0].content[0].type == "output_text"
|
||||
assert (
|
||||
output_items[0].content[0].text
|
||||
== "I will now search for the weather information."
|
||||
)
|
||||
|
||||
def test_commentary_with_function_recipient_creates_function_call(self):
|
||||
"""Test commentary with recipient='functions.X' creates function calls."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, '{"location": "San Francisco", "units": "celsius"}'
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].type == "function_call"
|
||||
assert output_items[0].name == "get_weather"
|
||||
assert (
|
||||
output_items[0].arguments
|
||||
== '{"location": "San Francisco", "units": "celsius"}'
|
||||
)
|
||||
assert output_items[0].call_id.startswith("call_")
|
||||
assert output_items[0].id.startswith("fc_")
|
||||
|
||||
def test_commentary_with_python_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='python' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "import numpy as np\nprint(np.array([1, 2, 3]))"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("python")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert (
|
||||
output_items[0].content[0].text
|
||||
== "import numpy as np\nprint(np.array([1, 2, 3]))"
|
||||
)
|
||||
|
||||
def test_commentary_with_browser_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='browser' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Navigating to the specified URL"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("browser")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == "Navigating to the specified URL"
|
||||
|
||||
def test_commentary_with_container_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='container' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Running command in container"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("container")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == "Running command in container"
|
||||
|
||||
def test_commentary_with_empty_content_and_no_recipient(self):
|
||||
"""Test edge case: empty commentary with recipient=None."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, "")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseOutputMessage)
|
||||
assert output_items[0].content[0].text == ""
|
||||
|
||||
def test_commentary_with_multiple_contents_and_no_recipient(self):
|
||||
"""Test multiple content items in commentary with no recipient."""
|
||||
contents = [
|
||||
TextContent(text="Step 1: Analyze the request"),
|
||||
TextContent(text="Step 2: Prepare to call functions"),
|
||||
]
|
||||
message = Message.from_role_and_contents(Role.ASSISTANT, contents)
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
# _parse_final_message returns single ResponseOutputMessage with
|
||||
# multiple contents
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseOutputMessage)
|
||||
assert len(output_items[0].content) == 2
|
||||
assert output_items[0].content[0].text == "Step 1: Analyze the request"
|
||||
assert output_items[0].content[1].text == "Step 2: Prepare to call functions"
|
||||
|
||||
def test_commentary_with_multiple_function_calls(self):
|
||||
"""Test multiple function calls in commentary channel."""
|
||||
contents = [
|
||||
TextContent(text='{"location": "San Francisco"}'),
|
||||
TextContent(text='{"location": "New York"}'),
|
||||
]
|
||||
message = Message.from_role_and_contents(Role.ASSISTANT, contents)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 2
|
||||
assert all(isinstance(item, ResponseFunctionToolCall) for item in output_items)
|
||||
assert output_items[0].name == "get_weather"
|
||||
assert output_items[1].name == "get_weather"
|
||||
assert output_items[0].arguments == '{"location": "San Francisco"}'
|
||||
assert output_items[1].arguments == '{"location": "New York"}'
|
||||
|
||||
def test_commentary_with_unknown_recipient_creates_mcp_call(self):
|
||||
"""Test that commentary with unknown recipient creates MCP call."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("custom_tool")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == "custom_tool"
|
||||
assert output_items[0].server_label == "custom_tool"
|
||||
|
||||
def test_analysis_channel_creates_reasoning(self):
|
||||
"""Test that analysis channel creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Analyzing the problem step by step..."
|
||||
)
|
||||
message = message.with_channel("analysis")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert (
|
||||
output_items[0].content[0].text == "Analyzing the problem step by step..."
|
||||
)
|
||||
|
||||
def test_non_assistant_message_returns_empty(self):
|
||||
"""Test that non-assistant messages return empty list.
|
||||
|
||||
Per the implementation, tool messages to assistant (e.g., search results)
|
||||
are not included in final output to align with OpenAI behavior.
|
||||
"""
|
||||
message = Message.from_author_and_content(
|
||||
Author.new(Role.TOOL, "functions.get_weather"),
|
||||
"The weather is sunny, 72°F",
|
||||
)
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 0
|
||||
|
||||
|
||||
def test_parse_mcp_call_basic() -> None:
|
||||
"""Test that MCP calls are parsed with correct type and server_label."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"path": "/tmp"}')
|
||||
message = message.with_recipient("filesystem")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == "filesystem"
|
||||
assert output_items[0].server_label == "filesystem"
|
||||
assert output_items[0].arguments == '{"path": "/tmp"}'
|
||||
assert output_items[0].status == "completed"
|
||||
|
||||
|
||||
def test_parse_mcp_call_dotted_recipient() -> None:
|
||||
"""Test that dotted recipients extract the tool name correctly."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"cmd": "ls"}')
|
||||
message = message.with_recipient("repo_browser.list")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].name == "list"
|
||||
assert output_items[0].server_label == "repo_browser"
|
||||
|
||||
|
||||
def test_mcp_vs_function_call() -> None:
|
||||
"""Test that function calls are not parsed as MCP calls."""
|
||||
func_message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
func_message = func_message.with_recipient("functions.my_tool")
|
||||
func_message = func_message.with_channel("commentary")
|
||||
|
||||
func_items = harmony_to_response_output(func_message)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
|
||||
|
||||
def test_mcp_vs_builtin_tools() -> None:
|
||||
"""Test that built-in tools (python, container) are not parsed as MCP calls."""
|
||||
# Test python (built-in tool) - should be reasoning, not MCP
|
||||
python_message = Message.from_role_and_content(Role.ASSISTANT, "print('hello')")
|
||||
python_message = python_message.with_recipient("python")
|
||||
python_message = python_message.with_channel("commentary")
|
||||
|
||||
python_items = harmony_to_response_output(python_message)
|
||||
|
||||
assert len(python_items) == 1
|
||||
assert not isinstance(python_items[0], McpCall)
|
||||
assert python_items[0].type == "reasoning"
|
||||
|
||||
|
||||
def test_parser_state_to_response_output_commentary_channel() -> None:
|
||||
"""Test parser_state_to_response_output with commentary
|
||||
channel and various recipients."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
# Test 1: functions.* recipient -> should return function tool call
|
||||
parser_func = Mock()
|
||||
parser_func.current_content = '{"arg": "value"}'
|
||||
parser_func.current_role = Role.ASSISTANT
|
||||
parser_func.current_channel = "commentary"
|
||||
parser_func.current_recipient = "functions.my_tool"
|
||||
|
||||
func_items = parser_state_to_response_output(parser_func)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
assert func_items[0].name == "my_tool"
|
||||
assert func_items[0].status == "in_progress"
|
||||
|
||||
# Test 2: MCP tool (not builtin) -> should return MCP call
|
||||
parser_mcp = Mock()
|
||||
parser_mcp.current_content = '{"path": "/tmp"}'
|
||||
parser_mcp.current_role = Role.ASSISTANT
|
||||
parser_mcp.current_channel = "commentary"
|
||||
parser_mcp.current_recipient = "filesystem"
|
||||
|
||||
mcp_items = parser_state_to_response_output(parser_mcp)
|
||||
|
||||
assert len(mcp_items) == 1
|
||||
assert isinstance(mcp_items[0], McpCall)
|
||||
assert mcp_items[0].type == "mcp_call"
|
||||
assert mcp_items[0].name == "filesystem"
|
||||
assert mcp_items[0].server_label == "filesystem"
|
||||
assert mcp_items[0].status == "in_progress"
|
||||
|
||||
# Test 3: Built-in tool (python)
|
||||
# should NOT return MCP call, returns reasoning (internal tool interaction)
|
||||
parser_builtin = Mock()
|
||||
parser_builtin.current_content = "print('hello')"
|
||||
parser_builtin.current_role = Role.ASSISTANT
|
||||
parser_builtin.current_channel = "commentary"
|
||||
parser_builtin.current_recipient = "python"
|
||||
|
||||
builtin_items = parser_state_to_response_output(parser_builtin)
|
||||
|
||||
# Built-in tools explicitly return reasoning
|
||||
assert len(builtin_items) == 1
|
||||
assert not isinstance(builtin_items[0], McpCall)
|
||||
assert builtin_items[0].type == "reasoning"
|
||||
|
||||
# Test 4: No recipient (preamble) → should return message, not reasoning
|
||||
parser_preamble = Mock()
|
||||
parser_preamble.current_content = "I'll search for that information now."
|
||||
parser_preamble.current_role = Role.ASSISTANT
|
||||
parser_preamble.current_channel = "commentary"
|
||||
parser_preamble.current_recipient = None
|
||||
|
||||
preamble_items = parser_state_to_response_output(parser_preamble)
|
||||
|
||||
assert len(preamble_items) == 1
|
||||
assert isinstance(preamble_items[0], ResponseOutputMessage)
|
||||
assert preamble_items[0].type == "message"
|
||||
assert preamble_items[0].content[0].text == "I'll search for that information now."
|
||||
assert preamble_items[0].status == "incomplete" # streaming
|
||||
|
||||
|
||||
def test_parser_state_to_response_output_analysis_channel() -> None:
|
||||
"""Test parser_state_to_response_output with analysis
|
||||
channel and various recipients."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
# Test 1: functions.* recipient -> should return function tool call
|
||||
parser_func = Mock()
|
||||
parser_func.current_content = '{"arg": "value"}'
|
||||
parser_func.current_role = Role.ASSISTANT
|
||||
parser_func.current_channel = "analysis"
|
||||
parser_func.current_recipient = "functions.my_tool"
|
||||
|
||||
func_items = parser_state_to_response_output(parser_func)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
assert func_items[0].name == "my_tool"
|
||||
assert func_items[0].status == "in_progress"
|
||||
|
||||
# Test 2: MCP tool (not builtin) -> should return MCP call
|
||||
parser_mcp = Mock()
|
||||
parser_mcp.current_content = '{"query": "test"}'
|
||||
parser_mcp.current_role = Role.ASSISTANT
|
||||
parser_mcp.current_channel = "analysis"
|
||||
parser_mcp.current_recipient = "database"
|
||||
|
||||
mcp_items = parser_state_to_response_output(parser_mcp)
|
||||
|
||||
assert len(mcp_items) == 1
|
||||
assert isinstance(mcp_items[0], McpCall)
|
||||
assert mcp_items[0].type == "mcp_call"
|
||||
assert mcp_items[0].name == "database"
|
||||
assert mcp_items[0].server_label == "database"
|
||||
assert mcp_items[0].status == "in_progress"
|
||||
|
||||
# Test 3: Built-in tool (container)
|
||||
# should NOT return MCP call, falls through to reasoning
|
||||
parser_builtin = Mock()
|
||||
parser_builtin.current_content = "docker run"
|
||||
parser_builtin.current_role = Role.ASSISTANT
|
||||
parser_builtin.current_channel = "analysis"
|
||||
parser_builtin.current_recipient = "container"
|
||||
|
||||
builtin_items = parser_state_to_response_output(parser_builtin)
|
||||
|
||||
# Should fall through to reasoning logic
|
||||
assert len(builtin_items) == 1
|
||||
assert not isinstance(builtin_items[0], McpCall)
|
||||
assert builtin_items[0].type == "reasoning"
|
||||
@@ -97,16 +97,16 @@ class TestMCPToolServerUnit:
|
||||
assert server.get_tool_description("test_server", allowed_tools=[]) is None
|
||||
|
||||
def test_builtin_tools_consistency(self):
|
||||
"""MCP_BUILTIN_TOOLS must match _BUILTIN_TOOL_TO_MCP_SERVER_LABEL values."""
|
||||
"""MCP_BUILTIN_TOOLS must match BUILTIN_TOOL_TO_MCP_SERVER_LABEL values."""
|
||||
from vllm.entrypoints.openai.parser.harmony_utils import (
|
||||
_BUILTIN_TOOL_TO_MCP_SERVER_LABEL,
|
||||
BUILTIN_TOOL_TO_MCP_SERVER_LABEL,
|
||||
MCP_BUILTIN_TOOLS,
|
||||
)
|
||||
|
||||
assert set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) == MCP_BUILTIN_TOOLS, (
|
||||
assert set(BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) == MCP_BUILTIN_TOOLS, (
|
||||
f"MCP_BUILTIN_TOOLS {MCP_BUILTIN_TOOLS} does not match "
|
||||
f"_BUILTIN_TOOL_TO_MCP_SERVER_LABEL values "
|
||||
f"{set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())}"
|
||||
f"BUILTIN_TOOL_TO_MCP_SERVER_LABEL values "
|
||||
f"{set(BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())}"
|
||||
)
|
||||
|
||||
|
||||
@@ -172,13 +172,13 @@ class TestMCPEnabled:
|
||||
recipient = message.get("recipient")
|
||||
if recipient and recipient.startswith("python"):
|
||||
tool_call_found = True
|
||||
assert message.get("channel") == "analysis"
|
||||
assert message.get("channel") == "commentary"
|
||||
author = message.get("author", {})
|
||||
if author.get("role") == "tool" and (author.get("name") or "").startswith(
|
||||
"python"
|
||||
):
|
||||
tool_response_found = True
|
||||
assert message.get("channel") == "analysis"
|
||||
assert message.get("channel") == "commentary"
|
||||
|
||||
assert tool_call_found, (
|
||||
f"No Python tool call found. "
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
# imports for structured outputs tests
|
||||
import json
|
||||
from collections import defaultdict
|
||||
|
||||
import jsonschema
|
||||
import openai # use the official client for correctness check
|
||||
@@ -13,6 +14,11 @@ import requests
|
||||
import torch
|
||||
from openai import BadRequestError
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
# any model with a chat template should work here
|
||||
@@ -815,3 +821,203 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA
|
||||
|
||||
assert chat_output.keys() == invocation_output.keys()
|
||||
assert chat_output["choices"] == invocation_output["choices"]
|
||||
|
||||
|
||||
# Test n parameter for chat completions
|
||||
# Tests that the n parameter works correctly for regular sampling
|
||||
# (non-beam search) in chat completions, addressing issue #34305.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_chat_completion_n_parameter_non_streaming(
|
||||
client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test that n parameter returns multiple choices for non-streaming requests."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the opposite of big?"},
|
||||
]
|
||||
|
||||
# Test with n=3
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_completion_tokens=20,
|
||||
temperature=0.7,
|
||||
n=3,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 3
|
||||
|
||||
# Verify each choice has content and correct index
|
||||
for i, choice in enumerate(chat_completion.choices):
|
||||
assert choice.index == i
|
||||
assert choice.message.content is not None
|
||||
assert len(choice.message.content) > 0
|
||||
|
||||
# Verify all responses are different (highly likely with temperature > 0)
|
||||
contents = [choice.message.content for choice in chat_completion.choices]
|
||||
assert len(set(contents)) > 1, "Expected different responses with n=3"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_chat_completion_n_parameter_streaming(
|
||||
client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test that n parameter returns multiple choices for streaming requests."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
]
|
||||
|
||||
stream = await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_completion_tokens=15,
|
||||
temperature=0.7,
|
||||
n=2,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect all chunks using defaultdict for dynamic handling
|
||||
chunks_by_index = defaultdict(list)
|
||||
async for chunk in stream:
|
||||
for choice in chunk.choices:
|
||||
if choice.delta.content:
|
||||
chunks_by_index[choice.index].append(choice.delta.content)
|
||||
|
||||
# Verify both choices received content
|
||||
assert len(chunks_by_index[0]) > 0, "Choice 0 received no content chunks"
|
||||
assert len(chunks_by_index[1]) > 0, "Choice 1 received no content chunks"
|
||||
|
||||
# Reconstruct full responses
|
||||
response_0 = "".join(chunks_by_index[0])
|
||||
response_1 = "".join(chunks_by_index[1])
|
||||
|
||||
assert len(response_0) > 0, "Choice 0 has empty response"
|
||||
assert len(response_1) > 0, "Choice 1 has empty response"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_chat_completion_n_with_seed(client: openai.AsyncOpenAI, model_name: str):
|
||||
"""Test that n parameter works correctly with seed parameter."""
|
||||
messages = [
|
||||
{"role": "user", "content": "Say hello."},
|
||||
]
|
||||
|
||||
# Test that seed parameter is accepted and works with n > 1
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.8,
|
||||
n=2,
|
||||
seed=42,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Verify we get n=2 choices
|
||||
assert len(chat_completion.choices) == 2
|
||||
|
||||
# Verify both choices have valid content
|
||||
for i, choice in enumerate(chat_completion.choices):
|
||||
assert choice.index == i
|
||||
assert choice.message.content is not None
|
||||
assert len(choice.message.content) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_chat_completion_n_equals_1(client: openai.AsyncOpenAI, model_name: str):
|
||||
"""Test that n=1 (default) still works correctly."""
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello!"},
|
||||
]
|
||||
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.7,
|
||||
n=1,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 1
|
||||
assert chat_completion.choices[0].index == 0
|
||||
assert chat_completion.choices[0].message.content is not None
|
||||
|
||||
|
||||
# Unit tests for n parameter in ChatCompletionRequest.to_sampling_params()
|
||||
def test_chat_completion_request_n_parameter_to_sampling_params():
|
||||
"""Test that n parameter is correctly passed to SamplingParams."""
|
||||
# Test with n=3
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
n=3,
|
||||
max_tokens=10,
|
||||
)
|
||||
|
||||
sampling_params = request.to_sampling_params(
|
||||
max_tokens=10,
|
||||
default_sampling_params={},
|
||||
)
|
||||
|
||||
assert isinstance(sampling_params, SamplingParams)
|
||||
assert sampling_params.n == 3, f"Expected n=3, got n={sampling_params.n}"
|
||||
|
||||
|
||||
def test_chat_completion_request_n_parameter_default():
|
||||
"""Test that n parameter defaults to 1."""
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
# n not specified, should default to 1
|
||||
max_tokens=10,
|
||||
)
|
||||
|
||||
assert request.n == 1, "n should default to 1"
|
||||
sampling_params = request.to_sampling_params(
|
||||
max_tokens=10,
|
||||
default_sampling_params={},
|
||||
)
|
||||
|
||||
# SamplingParams.from_optional converts None to 1
|
||||
assert sampling_params.n == 1, f"Expected n=1 (default), got n={sampling_params.n}"
|
||||
|
||||
|
||||
def test_chat_completion_request_n_parameter_various_values():
|
||||
"""Test n parameter with various values."""
|
||||
for n_value in [1, 2, 5, 10]:
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Test"}],
|
||||
n=n_value,
|
||||
max_tokens=10,
|
||||
)
|
||||
|
||||
sampling_params = request.to_sampling_params(
|
||||
max_tokens=10,
|
||||
default_sampling_params={},
|
||||
)
|
||||
|
||||
assert sampling_params.n == n_value, (
|
||||
f"Expected n={n_value}, got n={sampling_params.n}"
|
||||
)
|
||||
|
||||
@@ -180,20 +180,13 @@ class TestExtractHarmonyStreamingDelta:
|
||||
|
||||
assert delta_message.tool_calls[0].index == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"channel,recipient",
|
||||
[
|
||||
("commentary", None),
|
||||
("commentary", "browser.search"),
|
||||
],
|
||||
)
|
||||
def test_returns_tool_call_preambles(self, channel, recipient):
|
||||
"""Test that invalid tool recipient on commentary is treated as content."""
|
||||
def test_returns_preambles_as_content(self):
|
||||
"""Test that commentary with no recipient (preamble) is user content."""
|
||||
parser = MockStreamableParser()
|
||||
delta_text = "some text"
|
||||
|
||||
token_states = [
|
||||
TokenState(channel=channel, recipient=recipient, text=delta_text)
|
||||
TokenState(channel="commentary", recipient=None, text=delta_text)
|
||||
]
|
||||
|
||||
delta_message, tools_streamed = extract_harmony_streaming_delta(
|
||||
@@ -211,6 +204,7 @@ class TestExtractHarmonyStreamingDelta:
|
||||
[
|
||||
(None, None),
|
||||
("unknown_channel", None),
|
||||
("commentary", "browser.search"),
|
||||
],
|
||||
)
|
||||
def test_returns_none_for_invalid_inputs(self, channel, recipient):
|
||||
|
||||
@@ -26,6 +26,9 @@ from vllm.entrypoints.openai.responses.serving import (
|
||||
_extract_allowed_tools_from_mcp_requests,
|
||||
extract_tool_types,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
StreamingState,
|
||||
)
|
||||
from vllm.inputs.data import TokensPrompt
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.sampling_params import SamplingParams
|
||||
@@ -439,3 +442,115 @@ class TestExtractAllowedToolsFromMcpRequests:
|
||||
"server1": ["tool1"],
|
||||
"server2": ["tool2"],
|
||||
}
|
||||
|
||||
|
||||
class TestHarmonyPreambleStreaming:
|
||||
"""Tests for preamble (commentary with no recipient) streaming events."""
|
||||
|
||||
@staticmethod
|
||||
def _make_ctx(*, channel, recipient, delta="hello"):
|
||||
"""Build a lightweight mock StreamingHarmonyContext."""
|
||||
ctx = MagicMock()
|
||||
ctx.last_content_delta = delta
|
||||
ctx.parser.current_channel = channel
|
||||
ctx.parser.current_recipient = recipient
|
||||
return ctx
|
||||
|
||||
@staticmethod
|
||||
def _make_previous_item(*, channel, recipient, text="preamble text"):
|
||||
"""Build a lightweight mock previous_item (openai_harmony Message)."""
|
||||
content_part = MagicMock()
|
||||
content_part.text = text
|
||||
item = MagicMock()
|
||||
item.channel = channel
|
||||
item.recipient = recipient
|
||||
item.content = [content_part]
|
||||
return item
|
||||
|
||||
def test_preamble_delta_emits_text_events(self) -> None:
|
||||
"""commentary + recipient=None should emit output_text.delta events."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_content_delta_events,
|
||||
)
|
||||
|
||||
ctx = self._make_ctx(channel="commentary", recipient=None)
|
||||
state = StreamingState()
|
||||
|
||||
events = emit_content_delta_events(ctx, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.delta" in type_names
|
||||
assert "response.output_item.added" in type_names
|
||||
|
||||
def test_preamble_delta_second_token_no_added(self) -> None:
|
||||
"""Second preamble token should emit delta only, not added again."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_content_delta_events,
|
||||
)
|
||||
|
||||
ctx = self._make_ctx(channel="commentary", recipient=None, delta="w")
|
||||
state = StreamingState()
|
||||
state.sent_output_item_added = True
|
||||
state.current_item_id = "msg_test"
|
||||
state.current_content_index = 0
|
||||
|
||||
events = emit_content_delta_events(ctx, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.delta" in type_names
|
||||
assert "response.output_item.added" not in type_names
|
||||
|
||||
def test_commentary_with_function_recipient_not_preamble(self) -> None:
|
||||
"""commentary + recipient='functions.X' must NOT use preamble path."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_content_delta_events,
|
||||
)
|
||||
|
||||
ctx = self._make_ctx(
|
||||
channel="commentary",
|
||||
recipient="functions.get_weather",
|
||||
)
|
||||
state = StreamingState()
|
||||
|
||||
events = emit_content_delta_events(ctx, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.delta" not in type_names
|
||||
|
||||
def test_preamble_done_emits_text_done_events(self) -> None:
|
||||
"""Completed preamble should emit text done + content_part done +
|
||||
output_item done, same shape as final channel."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_previous_item_done_events,
|
||||
)
|
||||
|
||||
previous = self._make_previous_item(channel="commentary", recipient=None)
|
||||
state = StreamingState()
|
||||
state.current_item_id = "msg_test"
|
||||
state.current_output_index = 0
|
||||
state.current_content_index = 0
|
||||
|
||||
events = emit_previous_item_done_events(previous, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.done" in type_names
|
||||
assert "response.content_part.done" in type_names
|
||||
assert "response.output_item.done" in type_names
|
||||
|
||||
def test_commentary_with_recipient_no_preamble_done(self) -> None:
|
||||
"""commentary + recipient='functions.X' should route to function call
|
||||
done, not preamble done."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_previous_item_done_events,
|
||||
)
|
||||
|
||||
previous = self._make_previous_item(
|
||||
channel="commentary", recipient="functions.get_weather"
|
||||
)
|
||||
state = StreamingState()
|
||||
state.current_item_id = "fc_test"
|
||||
|
||||
events = emit_previous_item_done_events(previous, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.done" not in type_names
|
||||
|
||||
@@ -236,6 +236,44 @@ def test_reasoning_tokens_counting(mock_parser):
|
||||
assert context.num_output_tokens == 4
|
||||
|
||||
|
||||
def test_preamble_tokens_not_counted_as_reasoning(mock_parser):
|
||||
"""Preambles (commentary with no recipient) are visible user text,
|
||||
not hidden reasoning. They must NOT inflate num_reasoning_tokens."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
|
||||
mock_parser.current_channel = "commentary"
|
||||
mock_parser.current_recipient = None # preamble
|
||||
|
||||
mock_output = create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
output_token_ids=[4, 5, 6],
|
||||
num_cached_tokens=0,
|
||||
)
|
||||
context.append_output(mock_output)
|
||||
|
||||
assert context.num_reasoning_tokens == 0
|
||||
assert context.num_output_tokens == 3
|
||||
|
||||
|
||||
def test_commentary_with_recipient_counted_as_reasoning(mock_parser):
|
||||
"""Commentary directed at a tool (recipient != None) is hidden from
|
||||
the user, so it should still count as reasoning tokens."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
|
||||
mock_parser.current_channel = "commentary"
|
||||
mock_parser.current_recipient = "python"
|
||||
|
||||
mock_output = create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
output_token_ids=[4, 5, 6],
|
||||
num_cached_tokens=0,
|
||||
)
|
||||
context.append_output(mock_output)
|
||||
|
||||
assert context.num_reasoning_tokens == 3
|
||||
assert context.num_output_tokens == 3
|
||||
|
||||
|
||||
def test_zero_tokens_edge_case():
|
||||
"""Test behavior with all zero token counts."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
# GPQA Evaluation using GPT-OSS
|
||||
|
||||
This directory contains GPQA evaluation tests using the GPT-OSS evaluation package and vLLM server.
|
||||
|
||||
## Usage
|
||||
|
||||
### Run tests with pytest (like buildkite)
|
||||
|
||||
```bash
|
||||
# H200
|
||||
pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
|
||||
--config-list-file=configs/models-h200.txt
|
||||
|
||||
# B200
|
||||
pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
|
||||
--config-list-file=configs/models-b200.txt
|
||||
```
|
||||
|
||||
## Configuration Format
|
||||
|
||||
Model configs in `configs/` directory use this YAML format:
|
||||
|
||||
```yaml
|
||||
model_name: "openai/gpt-oss-20b"
|
||||
metric_threshold: 0.568 # Minimum expected accuracy
|
||||
reasoning_effort: "low" # Reasoning effort level (default: "low")
|
||||
server_args: "--tensor-parallel-size 2" # Server arguments
|
||||
startup_max_wait_seconds: 1800 # Max wait for server startup (default: 1800)
|
||||
env: # Environment variables (optional)
|
||||
SOME_VAR: "value"
|
||||
```
|
||||
|
||||
The `server_args` field accepts any arguments that can be passed to `vllm serve`.
|
||||
|
||||
The `env` field accepts a dictionary of environment variables to set for the server process.
|
||||
|
||||
## Adding New Models
|
||||
|
||||
1. Create a new YAML config file in the `configs/` directory
|
||||
2. Add the filename to the appropriate `models-*.txt` file
|
||||
|
||||
## Tiktoken Encoding Files
|
||||
|
||||
The tiktoken encoding files required by the vLLM server are automatically downloaded from OpenAI's public blob storage on first run:
|
||||
|
||||
- `cl100k_base.tiktoken`
|
||||
- `o200k_base.tiktoken`
|
||||
|
||||
Files are cached in the `data/` directory. The `TIKTOKEN_ENCODINGS_BASE` environment variable is automatically set to point to this directory when running evaluations.
|
||||
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
model_name: "openai/gpt-oss-20b"
|
||||
metric_threshold: 0.568
|
||||
reasoning_effort: "low"
|
||||
server_args: "--tensor-parallel-size 2"
|
||||
@@ -0,0 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
model_name: "openai/gpt-oss-20b"
|
||||
metric_threshold: 0.568
|
||||
reasoning_effort: "low"
|
||||
server_args: "--tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: "1"
|
||||
@@ -0,0 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
model_name: "openai/gpt-oss-20b"
|
||||
metric_threshold: 0.568
|
||||
reasoning_effort: "low"
|
||||
server_args: "--tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: "1"
|
||||
@@ -0,0 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
model_name: "openai/gpt-oss-20b"
|
||||
metric_threshold: 0.568
|
||||
reasoning_effort: "low"
|
||||
server_args: "--tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_MXFP4_USE_MARLIN: "1"
|
||||
@@ -0,0 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
model_name: "openai/gpt-oss-20b"
|
||||
metric_threshold: 0.568
|
||||
reasoning_effort: "low"
|
||||
server_args: "--tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: "1"
|
||||
@@ -0,0 +1,5 @@
|
||||
# B200 model configurations for GPQA evaluation
|
||||
# Tests different environment variable combinations
|
||||
gpt-oss-20b-flashinfer-mxfp4-bf16.yaml
|
||||
gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml
|
||||
gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml
|
||||
@@ -0,0 +1,5 @@
|
||||
# H100 model configurations for GPQA evaluation
|
||||
# Tests different environment variable combinations
|
||||
gpt-oss-20b-baseline.yaml
|
||||
gpt-oss-20b-flashinfer-mxfp4-bf16.yaml
|
||||
gpt-oss-20b-marlin.yaml
|
||||
@@ -4,13 +4,61 @@
|
||||
Pytest configuration for GPT-OSS evaluation tests.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def pytest_addoption(parser):
|
||||
"""Add command line options for pytest."""
|
||||
parser.addoption("--model", action="store", help="Model name to evaluate")
|
||||
"""Add custom command line options."""
|
||||
parser.addoption(
|
||||
"--metric", action="store", type=float, help="Expected metric threshold"
|
||||
)
|
||||
parser.addoption(
|
||||
"--server-args", action="store", default="", help="Additional server arguments"
|
||||
"--config-list-file",
|
||||
required=True,
|
||||
help="File containing list of config files to test",
|
||||
)
|
||||
|
||||
|
||||
def pytest_generate_tests(metafunc):
|
||||
"""Generate test parameters from config files."""
|
||||
if "config_filename" in metafunc.fixturenames:
|
||||
config_list_file = metafunc.config.getoption("--config-list-file")
|
||||
|
||||
# Handle both relative and absolute paths
|
||||
config_list_path = Path(config_list_file)
|
||||
if not config_list_path.is_absolute():
|
||||
# If relative, try relative to test directory first
|
||||
test_dir_path = Path(__file__).parent / config_list_file
|
||||
if test_dir_path.exists():
|
||||
config_list_path = test_dir_path
|
||||
else:
|
||||
# Try relative to current working directory
|
||||
config_list_path = Path.cwd() / config_list_file
|
||||
|
||||
print(f"Looking for config list at: {config_list_path}")
|
||||
|
||||
config_files = []
|
||||
if config_list_path.exists():
|
||||
# Determine config directory (same directory as the list file)
|
||||
config_dir = config_list_path.parent
|
||||
|
||||
with open(config_list_path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#"):
|
||||
config_path = config_dir / line
|
||||
print(f"Checking config file: {config_path}")
|
||||
if config_path.exists():
|
||||
config_files.append(config_path)
|
||||
print(f" Found: {config_path}")
|
||||
else:
|
||||
print(f" Missing: {config_path}")
|
||||
else:
|
||||
print(f"Config list file not found: {config_list_path}")
|
||||
|
||||
# Generate test parameters
|
||||
if config_files:
|
||||
metafunc.parametrize(
|
||||
"config_filename",
|
||||
config_files,
|
||||
ids=[config_file.stem for config_file in config_files],
|
||||
)
|
||||
else:
|
||||
print("No config files found, test will be skipped")
|
||||
|
||||
@@ -5,22 +5,48 @@ GPQA evaluation using vLLM server and GPT-OSS evaluation package.
|
||||
|
||||
Usage:
|
||||
pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
|
||||
--model openai/gpt-oss-20b \
|
||||
--metric 0.58 \
|
||||
--server-args "--tensor-parallel-size 2"
|
||||
--config-list-file=configs/models-h200.txt
|
||||
"""
|
||||
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
import regex as re
|
||||
import yaml
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
TOL = 0.05 # Absolute tolerance for accuracy comparison
|
||||
|
||||
# Path to tiktoken encoding files
|
||||
TIKTOKEN_DATA_DIR = Path(__file__).parent / "data"
|
||||
|
||||
def run_gpqa_eval(model_name: str, base_url: str) -> float:
|
||||
# Tiktoken encoding files to download
|
||||
TIKTOKEN_FILES = {
|
||||
"cl100k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken",
|
||||
"o200k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken",
|
||||
}
|
||||
|
||||
|
||||
def ensure_tiktoken_files():
|
||||
"""Download tiktoken encoding files if they don't exist."""
|
||||
TIKTOKEN_DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for filename, url in TIKTOKEN_FILES.items():
|
||||
filepath = TIKTOKEN_DATA_DIR / filename
|
||||
if not filepath.exists():
|
||||
print(f"Downloading {filename} from {url}...")
|
||||
urllib.request.urlretrieve(url, filepath)
|
||||
print(f" Downloaded to {filepath}")
|
||||
else:
|
||||
print(f" {filename} already exists.")
|
||||
|
||||
|
||||
def run_gpqa_eval(model_name: str, base_url: str, reasoning_effort: str) -> float:
|
||||
"""Run GPQA evaluation using the gpt-oss evaluation package."""
|
||||
|
||||
# Build the command to run the evaluation
|
||||
@@ -33,7 +59,7 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float:
|
||||
"--model",
|
||||
model_name,
|
||||
"--reasoning-effort",
|
||||
"low",
|
||||
reasoning_effort,
|
||||
"--base-url",
|
||||
base_url,
|
||||
"--n-threads",
|
||||
@@ -41,16 +67,29 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float:
|
||||
]
|
||||
|
||||
try:
|
||||
# Set up environment for the evaluation subprocess
|
||||
# Inherit current environment and add required variables
|
||||
eval_env = os.environ.copy()
|
||||
eval_env["OPENAI_API_KEY"] = "dummy"
|
||||
|
||||
# Run the evaluation
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
text=True,
|
||||
capture_output=True,
|
||||
timeout=1800, # 30 minute timeout
|
||||
env={"OPENAI_API_KEY": "dummy"},
|
||||
env=eval_env,
|
||||
)
|
||||
|
||||
print("Evaluation process output:\n", result.stdout)
|
||||
print("Evaluation process stdout:\n", result.stdout)
|
||||
print("Evaluation process stderr:\n", result.stderr)
|
||||
print(f"Evaluation process return code: {result.returncode}")
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Evaluation failed with exit code {result.returncode}:\n"
|
||||
f"stdout: {result.stdout}\nstderr: {result.stderr}"
|
||||
)
|
||||
|
||||
# Parse the output to extract the score
|
||||
match = re.search(r"'metric':\s*([\d.]+)", result.stdout)
|
||||
@@ -64,47 +103,62 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float:
|
||||
|
||||
except subprocess.TimeoutExpired as e:
|
||||
raise RuntimeError("Evaluation timed out") from e
|
||||
except subprocess.CalledProcessError as e:
|
||||
raise RuntimeError(
|
||||
f"Evaluation failed with exit code {e.returncode}:\n"
|
||||
f"stdout: {e.stdout}\nstderr: {e.stderr}"
|
||||
) from e
|
||||
|
||||
|
||||
def test_gpqa_correctness(request):
|
||||
"""Test GPQA correctness for GPT-OSS model."""
|
||||
def test_gpqa_correctness(config_filename):
|
||||
"""Test GPQA correctness for a given model configuration."""
|
||||
# Ensure tiktoken files are downloaded
|
||||
ensure_tiktoken_files()
|
||||
|
||||
# Get command line arguments
|
||||
model_name = request.config.getoption("--model")
|
||||
expected_metric = request.config.getoption("--metric")
|
||||
server_args_str = request.config.getoption("--server-args")
|
||||
# Verify tiktoken files exist
|
||||
for filename in TIKTOKEN_FILES:
|
||||
filepath = TIKTOKEN_DATA_DIR / filename
|
||||
assert filepath.exists(), f"Tiktoken file not found: {filepath}"
|
||||
|
||||
# Parse server arguments
|
||||
server_args = []
|
||||
if server_args_str:
|
||||
server_args = server_args_str.split()
|
||||
eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
|
||||
|
||||
# Parse server arguments from config (use shlex to handle quoted strings)
|
||||
server_args_str = eval_config.get("server_args", "")
|
||||
server_args = shlex.split(server_args_str) if server_args_str else []
|
||||
|
||||
# Add standard server arguments
|
||||
server_args.extend(
|
||||
[
|
||||
"--trust-remote-code",
|
||||
"--enforce-eager",
|
||||
"--disable-uvicorn-access-log",
|
||||
]
|
||||
)
|
||||
|
||||
print(f"Starting GPQA evaluation for model: {model_name}")
|
||||
print(f"Expected metric threshold: {expected_metric}")
|
||||
# Build server environment with tiktoken path and any config-specified vars
|
||||
server_env = {"TIKTOKEN_ENCODINGS_BASE": str(TIKTOKEN_DATA_DIR)}
|
||||
if eval_config.get("env"):
|
||||
server_env.update(eval_config["env"])
|
||||
|
||||
reasoning_effort = eval_config.get("reasoning_effort", "low")
|
||||
|
||||
print(f"Starting GPQA evaluation for model: {eval_config['model_name']}")
|
||||
print(f"Expected metric threshold: {eval_config['metric_threshold']}")
|
||||
print(f"Reasoning effort: {reasoning_effort}")
|
||||
print(f"Server args: {' '.join(server_args)}")
|
||||
print(f"Server environment variables: {server_env}")
|
||||
|
||||
# Launch server and run evaluation
|
||||
with RemoteOpenAIServer(
|
||||
model_name, server_args, max_wait_seconds=1800
|
||||
eval_config["model_name"],
|
||||
server_args,
|
||||
env_dict=server_env,
|
||||
max_wait_seconds=eval_config.get("startup_max_wait_seconds", 1800),
|
||||
) as remote_server:
|
||||
base_url = remote_server.url_for("v1")
|
||||
print(f"Server started at: {base_url}")
|
||||
|
||||
measured_metric = run_gpqa_eval(model_name, base_url)
|
||||
measured_metric = run_gpqa_eval(
|
||||
eval_config["model_name"], base_url, reasoning_effort
|
||||
)
|
||||
expected_metric = eval_config["metric_threshold"]
|
||||
|
||||
print(f"GPQA Results for {model_name}:")
|
||||
print(f"GPQA Results for {eval_config['model_name']}:")
|
||||
print(f" Measured metric: {measured_metric:.4f}")
|
||||
print(f" Expected metric: {expected_metric:.4f}")
|
||||
print(f" Tolerance: {TOL:.4f}")
|
||||
@@ -115,4 +169,4 @@ def test_gpqa_correctness(request):
|
||||
f"{expected_metric:.4f} - {TOL:.4f} = {expected_metric - TOL:.4f}"
|
||||
)
|
||||
|
||||
print(f"✅ GPQA test passed for {model_name}")
|
||||
print(f"GPQA test passed for {eval_config['model_name']}")
|
||||
|
||||
@@ -8,5 +8,4 @@ server_args: >-
|
||||
--tensor-parallel-size 2
|
||||
--enable-expert-parallel
|
||||
--speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":1}'
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
--moe-backend=flashinfer_trtllm
|
||||
|
||||
@@ -7,5 +7,4 @@ server_args: >-
|
||||
--tensor-parallel-size 2
|
||||
--enable-expert-parallel
|
||||
--async-scheduling
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
--moe-backend=flashinfer_trtllm
|
||||
|
||||
@@ -2,7 +2,6 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"
|
||||
accuracy_threshold: 0.92
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=triton"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "0"
|
||||
VLLM_USE_DEEP_GEMM: "0"
|
||||
|
||||
+1
-4
@@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "masked_gemm"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"
|
||||
|
||||
+1
-4
@@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "masked_gemm"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl"
|
||||
|
||||
+1
-4
@@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"
|
||||
|
||||
+1
-4
@@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_trtllm"
|
||||
|
||||
@@ -2,8 +2,4 @@ model_name: "meta-llama/Llama-4-Scout-17B-16E-Instruct"
|
||||
accuracy_threshold: 0.92
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP16: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"
|
||||
accuracy_threshold: 0.92
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"
|
||||
accuracy_threshold: 0.92
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
|
||||
|
||||
@@ -2,6 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"
|
||||
accuracy_threshold: 0.92
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "0"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "mistralai/Mixtral-8x7B-v0.1"
|
||||
accuracy_threshold: 0.58
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP16: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -3,7 +3,4 @@
|
||||
# accuracy_threshold: 0.62
|
||||
# num_questions: 1319
|
||||
# num_fewshot: 5
|
||||
# server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
# env:
|
||||
# VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
# VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
# server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
+1
-4
@@ -2,7 +2,4 @@ model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
|
||||
accuracy_threshold: 0.29
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
|
||||
|
||||
+1
-4
@@ -2,7 +2,4 @@ model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.29
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,6 +2,4 @@ model_name: "Qwen/Qwen3-30B-A3B"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP16: "1"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
|
||||
|
||||
@@ -2,7 +2,6 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "0"
|
||||
VLLM_USE_DEEP_GEMM: "0"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block"
|
||||
accuracy_threshold: 0.85
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,6 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block"
|
||||
accuracy_threshold: 0.85
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "0"
|
||||
VLLM_USE_DEEP_GEMM: "0"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
|
||||
|
||||
@@ -2,6 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "0"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass"
|
||||
|
||||
@@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
|
||||
|
||||
+1
-3
@@ -2,6 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "0"
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=cutlass"
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests that triton_kernel_moe_forward correctly applies expert_map
|
||||
remapping when expert parallelism (EP) is enabled.
|
||||
|
||||
Previously, legacy_routing was always used and it produced routing data
|
||||
with global expert IDs that didn't correspond to local weight indices,
|
||||
causing illegal memory access with EP. The fix splits routing: when
|
||||
expert_map is provided, topk selection is performed first, expert_map is
|
||||
applied to remap global→local IDs, and make_routing_data builds routing
|
||||
structures from the local IDs.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.mxfp4 import (
|
||||
Mxfp4Backend,
|
||||
Mxfp4MoEMethod,
|
||||
)
|
||||
|
||||
|
||||
def _make_mock_moe_config(ep_size: int = 1) -> MagicMock:
|
||||
"""Create a mock FusedMoEConfig with the given EP size."""
|
||||
parallel_config = MagicMock()
|
||||
parallel_config.ep_size = ep_size
|
||||
|
||||
moe_config = MagicMock()
|
||||
moe_config.ep_size = ep_size
|
||||
moe_config.is_lora_enabled = False
|
||||
moe_config.moe_parallel_config = parallel_config
|
||||
return moe_config
|
||||
|
||||
|
||||
class TestMxfp4TritonIsMonolithic:
|
||||
"""Verify that is_monolithic is always True for the TRITON backend,
|
||||
regardless of EP size, since triton_kernel_moe_forward now handles
|
||||
expert_map remapping internally."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend,ep_size,expected_monolithic",
|
||||
[
|
||||
# TRITON is always monolithic (handles EP via expert_map remapping)
|
||||
(Mxfp4Backend.TRITON, 1, True),
|
||||
(Mxfp4Backend.TRITON, 2, True),
|
||||
(Mxfp4Backend.TRITON, 4, True),
|
||||
# SM100 backends are always monolithic
|
||||
(Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM, 1, True),
|
||||
(Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM, 2, True),
|
||||
(Mxfp4Backend.SM100_FI_MXFP4_BF16, 1, True),
|
||||
(Mxfp4Backend.SM100_FI_MXFP4_BF16, 2, True),
|
||||
# MARLIN is never monolithic
|
||||
(Mxfp4Backend.MARLIN, 1, False),
|
||||
(Mxfp4Backend.MARLIN, 2, False),
|
||||
],
|
||||
ids=[
|
||||
"triton-no-ep",
|
||||
"triton-ep2",
|
||||
"triton-ep4",
|
||||
"sm100-trtllm-no-ep",
|
||||
"sm100-trtllm-ep2",
|
||||
"sm100-bf16-no-ep",
|
||||
"sm100-bf16-ep2",
|
||||
"marlin-no-ep",
|
||||
"marlin-ep2",
|
||||
],
|
||||
)
|
||||
@patch(
|
||||
"vllm.model_executor.layers.quantization.mxfp4.get_mxfp4_backend",
|
||||
)
|
||||
@patch(
|
||||
"vllm.model_executor.layers.quantization.mxfp4.get_current_vllm_config",
|
||||
)
|
||||
def test_is_monolithic(
|
||||
self,
|
||||
mock_get_config,
|
||||
mock_get_backend,
|
||||
backend,
|
||||
ep_size,
|
||||
expected_monolithic,
|
||||
):
|
||||
"""is_monolithic should be True for TRITON regardless of EP size."""
|
||||
mock_get_backend.return_value = backend
|
||||
|
||||
mock_compilation_config = MagicMock()
|
||||
mock_compilation_config.max_cudagraph_capture_size = 1024
|
||||
mock_vllm_config = MagicMock()
|
||||
mock_vllm_config.compilation_config = mock_compilation_config
|
||||
mock_get_config.return_value = mock_vllm_config
|
||||
|
||||
moe_config = _make_mock_moe_config(ep_size=ep_size)
|
||||
method = Mxfp4MoEMethod(moe_config)
|
||||
|
||||
assert method.is_monolithic == expected_monolithic, (
|
||||
f"Expected is_monolithic={expected_monolithic} for "
|
||||
f"backend={backend.name}, ep_size={ep_size}, "
|
||||
f"but got {method.is_monolithic}."
|
||||
)
|
||||
|
||||
|
||||
class TestTritonMoeForwardExpertMap:
|
||||
"""Test that triton_kernel_moe_forward applies expert_map remapping
|
||||
when expert_map is provided (EP active)."""
|
||||
|
||||
@pytest.mark.parametrize("expert_map_present", [False, True])
|
||||
def test_routing_path_selection(self, expert_map_present):
|
||||
"""Verify that the EP-aware routing path is taken when expert_map
|
||||
is present, and the legacy_routing path is taken otherwise."""
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
# This is a structural test: we mock the routing functions to
|
||||
# verify the correct path is exercised.
|
||||
mock_expert_map = (
|
||||
torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.layers.fused_moe."
|
||||
"gpt_oss_triton_kernels_moe.legacy_routing"
|
||||
) as mock_legacy,
|
||||
patch("triton_kernels.topk.topk") as mock_topk,
|
||||
patch(
|
||||
"vllm.model_executor.layers.fused_moe."
|
||||
"gpt_oss_triton_kernels_moe.make_routing_data"
|
||||
) as mock_make_routing,
|
||||
patch(
|
||||
"vllm.model_executor.layers.fused_moe."
|
||||
"gpt_oss_triton_kernels_moe.triton_kernel_fused_experts"
|
||||
) as mock_fused_experts,
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
|
||||
triton_kernel_moe_forward,
|
||||
)
|
||||
|
||||
# Set up return values
|
||||
mock_routing_data = MagicMock()
|
||||
mock_gather = MagicMock()
|
||||
mock_scatter = MagicMock()
|
||||
|
||||
if expert_map_present:
|
||||
sparse_result = MagicMock()
|
||||
sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32)
|
||||
sparse_result.vals = torch.tensor([[0.6, 0.4]])
|
||||
mock_topk.return_value = sparse_result
|
||||
mock_make_routing.return_value = (
|
||||
mock_routing_data,
|
||||
mock_gather,
|
||||
mock_scatter,
|
||||
)
|
||||
else:
|
||||
mock_legacy.return_value = (
|
||||
mock_routing_data,
|
||||
mock_gather,
|
||||
mock_scatter,
|
||||
)
|
||||
|
||||
mock_fused_experts.return_value = torch.zeros((1, 8), device=device)
|
||||
|
||||
hidden = torch.randn((1, 8), device=device)
|
||||
w1 = torch.randn((2, 8, 16), device=device)
|
||||
w2 = torch.randn((2, 8, 8), device=device)
|
||||
logits = torch.randn((1, 4), device=device)
|
||||
|
||||
triton_kernel_moe_forward(
|
||||
hidden_states=hidden,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
gating_output=logits,
|
||||
topk=2,
|
||||
renormalize=True,
|
||||
expert_map=mock_expert_map,
|
||||
)
|
||||
|
||||
if expert_map_present:
|
||||
# EP path: should use topk + make_routing_data, NOT
|
||||
# legacy_routing
|
||||
mock_topk.assert_called_once()
|
||||
mock_make_routing.assert_called_once()
|
||||
mock_legacy.assert_not_called()
|
||||
# expert_map should be None in the fused_experts call
|
||||
# (already applied)
|
||||
call_kwargs = mock_fused_experts.call_args
|
||||
assert call_kwargs[1].get("expert_map") is None or (
|
||||
len(call_kwargs[0]) > 0
|
||||
)
|
||||
else:
|
||||
# Non-EP path: should use legacy_routing
|
||||
mock_legacy.assert_called_once()
|
||||
mock_topk.assert_not_called()
|
||||
mock_make_routing.assert_not_called()
|
||||
@@ -103,14 +103,14 @@ def dummy_model(default_vllm_config) -> nn.Module:
|
||||
("output", ColumnParallelLinear(50, 10)),
|
||||
("outact", nn.Sigmoid()),
|
||||
# Special handling for lm_head & sampler
|
||||
("lm_head", ParallelLMHead(512, 10)),
|
||||
("logits_processor", LogitsProcessor(512)),
|
||||
("lm_head", ParallelLMHead(32064, 10)),
|
||||
("logits_processor", LogitsProcessor(32064)),
|
||||
]
|
||||
)
|
||||
)
|
||||
model.config = MagicMock()
|
||||
model.embedding_modules = {"lm_head": "lm_head"}
|
||||
model.unpadded_vocab_size = 32000
|
||||
model.unpadded_vocab_size = 32064
|
||||
return model
|
||||
|
||||
|
||||
@@ -136,8 +136,8 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module:
|
||||
("gate_up_proj", MergedColumnParallelLinear(50, [5, 5])),
|
||||
("outact", nn.Sigmoid()),
|
||||
# Special handling for lm_head & sampler
|
||||
("lm_head", ParallelLMHead(512, 10)),
|
||||
("logits_processor", LogitsProcessor(512)),
|
||||
("lm_head", ParallelLMHead(32064, 10)),
|
||||
("logits_processor", LogitsProcessor(32064)),
|
||||
]
|
||||
)
|
||||
)
|
||||
@@ -149,7 +149,7 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module:
|
||||
],
|
||||
}
|
||||
model.embedding_modules = {"lm_head": "lm_head"}
|
||||
model.unpadded_vocab_size = 32000
|
||||
model.unpadded_vocab_size = 32064
|
||||
|
||||
return model
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ from vllm.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.lora.ops.triton_ops import fused_moe_lora
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
@@ -244,8 +245,9 @@ def use_torch(
|
||||
return torch.stack(outputs, dim=0)
|
||||
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
DEVICES = [f"cuda:{0}"]
|
||||
DEVICES = [f"{DEVICE_TYPE}:{0}"]
|
||||
SEED = [42]
|
||||
|
||||
|
||||
|
||||
@@ -353,7 +353,7 @@ def test_embeddings(
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("num_loras", [1, 2, 4])
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 256512])
|
||||
@pytest.mark.parametrize("vocab_size", [64000, 256512, 258048])
|
||||
@pytest.mark.parametrize("stage", STAGES)
|
||||
def test_lm_head_logits_processor(
|
||||
default_vllm_config, dist_init, num_loras, device, vocab_size, stage
|
||||
@@ -468,6 +468,31 @@ def test_lm_head_logits_processor(
|
||||
torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("vocab_size", [512, 32000, 258049, 300000])
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
def test_lm_head_logits_processor_invalid_vocab_size(
|
||||
default_vllm_config, dist_init, vocab_size, device
|
||||
) -> None:
|
||||
"""Test that LogitsProcessorWithLoRA raises ValueError for invalid vocab sizes."""
|
||||
if current_platform.is_cuda_alike():
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
torch.set_default_device(device)
|
||||
max_loras = 8
|
||||
lora_config = LoRAConfig(
|
||||
max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
|
||||
)
|
||||
|
||||
logits_processor = LogitsProcessor(vocab_size)
|
||||
lora_logits_processor = LogitsProcessorWithLoRA(
|
||||
logits_processor, 1024, torch.float16, device, None
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="vocab size must be > 32000 and <= 258048"):
|
||||
lora_logits_processor.create_lora_weights(max_loras, lora_config)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("num_loras", [1, 2, 4])
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
|
||||
@@ -395,6 +395,7 @@ def test_kernels(
|
||||
Tests LoRA kernels.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
torch.cuda.set_device(device)
|
||||
set_random_seed(seed)
|
||||
|
||||
if op_type == "shrink":
|
||||
@@ -447,6 +448,7 @@ def test_kernels_hidden_size(
|
||||
Tests SGMV and LoRA kernels.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
torch.cuda.set_device(device)
|
||||
set_random_seed(seed)
|
||||
|
||||
if op_type == "shrink":
|
||||
|
||||
@@ -0,0 +1,298 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.lora.utils import (
|
||||
PunicaTensors,
|
||||
assert_close,
|
||||
generate_data,
|
||||
generate_data_for_expand_nslices,
|
||||
)
|
||||
from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def torch_bgmv_expand(
|
||||
inputs: torch.Tensor,
|
||||
lora_b_weights: torch.Tensor,
|
||||
output_tensor: torch.Tensor,
|
||||
lora_indices_tensor: torch.Tensor,
|
||||
add_inputs: bool = True,
|
||||
):
|
||||
selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
|
||||
if len(selected_loras.shape) == 4:
|
||||
selected_loras = selected_loras.squeeze(dim=1)
|
||||
inputs = inputs.to(dtype=output_tensor.dtype)
|
||||
outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
|
||||
|
||||
limit = output_tensor.shape[0]
|
||||
if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
|
||||
limit = 1
|
||||
|
||||
# LoRA adapter and model may add different amounts of padding to output
|
||||
common_len = min(outputs.shape[1], output_tensor.shape[1])
|
||||
|
||||
if add_inputs:
|
||||
output_tensor[:, :common_len] += outputs[:limit, :common_len]
|
||||
else:
|
||||
output_tensor[:, :common_len] = outputs[:limit, :common_len]
|
||||
|
||||
|
||||
def torch_bgmv_shrink(
|
||||
inputs: torch.Tensor,
|
||||
lora_b_weights: torch.Tensor,
|
||||
output_tensor: torch.Tensor,
|
||||
lora_indices_tensor: torch.Tensor,
|
||||
scaling: float = 1.0,
|
||||
):
|
||||
selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
|
||||
if len(selected_loras.shape) == 4:
|
||||
selected_loras = selected_loras.squeeze(dim=1)
|
||||
inputs = inputs.to(dtype=output_tensor.dtype)
|
||||
outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
|
||||
|
||||
output_tensor[:, : outputs.shape[1]] = scaling * outputs[:]
|
||||
|
||||
|
||||
def torch_bgmv_expand_slice(
|
||||
inputs: torch.Tensor,
|
||||
lora_b_weights: torch.Tensor,
|
||||
output_tensor: torch.Tensor,
|
||||
lora_indices_tensor: torch.Tensor,
|
||||
slice_offset: int,
|
||||
slice_size: int,
|
||||
add_inputs: bool = True,
|
||||
):
|
||||
selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
|
||||
inputs = inputs.to(dtype=output_tensor.dtype)
|
||||
if len(selected_loras.shape) == 4:
|
||||
selected_loras = selected_loras.squeeze(dim=1)
|
||||
outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
|
||||
|
||||
if add_inputs:
|
||||
output_tensor[:, slice_offset : slice_offset + slice_size] += outputs[:]
|
||||
else:
|
||||
output_tensor[:, slice_offset : slice_offset + slice_size] = outputs[:]
|
||||
|
||||
|
||||
def check_bgmv_shrink(
|
||||
batches: int,
|
||||
num_loras: int,
|
||||
rank: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
scaling: float,
|
||||
):
|
||||
"""
|
||||
Compare vllm.bgmv_shrink against a reference implementation.
|
||||
"""
|
||||
seq_length = 1
|
||||
data: PunicaTensors = generate_data(
|
||||
batches,
|
||||
hidden_size,
|
||||
num_loras,
|
||||
rank,
|
||||
seq_length,
|
||||
dtype,
|
||||
"shrink",
|
||||
device,
|
||||
)
|
||||
|
||||
bgmv_shrink(
|
||||
data.inputs_tensor,
|
||||
data.lora_weights,
|
||||
data.our_out_tensor,
|
||||
data.token_lora_mapping,
|
||||
scaling,
|
||||
)
|
||||
|
||||
torch_bgmv_shrink(
|
||||
data.inputs_tensor,
|
||||
data.lora_weights,
|
||||
data.ref_out_tensor,
|
||||
data.token_lora_mapping,
|
||||
scaling,
|
||||
)
|
||||
|
||||
data.ref_out_tensor = data.ref_out_tensor.to(torch.float32)
|
||||
assert_close(data.our_out_tensor, data.ref_out_tensor)
|
||||
|
||||
|
||||
def check_bgmv_expand(
|
||||
batches: int,
|
||||
num_loras: int,
|
||||
rank: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
add_inputs: bool,
|
||||
):
|
||||
"""
|
||||
Compare vllm.bgmv_expand against a reference implementation.
|
||||
"""
|
||||
seq_length = 1
|
||||
data: PunicaTensors = generate_data(
|
||||
batches,
|
||||
hidden_size,
|
||||
num_loras,
|
||||
rank,
|
||||
seq_length,
|
||||
dtype,
|
||||
"expand",
|
||||
device,
|
||||
)
|
||||
|
||||
bgmv_expand(
|
||||
data.inputs_tensor,
|
||||
data.lora_weights,
|
||||
data.our_out_tensor,
|
||||
data.token_lora_mapping,
|
||||
add_inputs=add_inputs,
|
||||
)
|
||||
torch_bgmv_expand(
|
||||
data.inputs_tensor,
|
||||
data.lora_weights,
|
||||
data.ref_out_tensor,
|
||||
data.token_lora_mapping,
|
||||
add_inputs=add_inputs,
|
||||
)
|
||||
assert_close(data.ref_out_tensor, data.our_out_tensor)
|
||||
|
||||
|
||||
def check_bgmv_expand_slice(
|
||||
batches: int,
|
||||
num_loras: int,
|
||||
rank: int,
|
||||
hidden_size: int,
|
||||
nslices: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
add_inputs: bool,
|
||||
):
|
||||
"""
|
||||
Compare vllm.bgmv_expand_slice against a reference implementation.
|
||||
"""
|
||||
seq_length = 1
|
||||
data: PunicaTensors = generate_data_for_expand_nslices(
|
||||
batches,
|
||||
hidden_size,
|
||||
num_loras,
|
||||
rank,
|
||||
seq_length,
|
||||
dtype,
|
||||
nslices,
|
||||
device,
|
||||
)
|
||||
|
||||
slice_offset = 0
|
||||
for index in range(nslices):
|
||||
bgmv_expand_slice(
|
||||
data.inputs_tensor,
|
||||
data.lora_weights[index],
|
||||
data.our_out_tensor,
|
||||
data.token_lora_mapping,
|
||||
slice_offset,
|
||||
slice_size=hidden_size,
|
||||
add_inputs=add_inputs,
|
||||
)
|
||||
torch_bgmv_expand_slice(
|
||||
data.inputs_tensor,
|
||||
data.lora_weights[index],
|
||||
data.ref_out_tensor,
|
||||
data.token_lora_mapping,
|
||||
slice_offset,
|
||||
slice_size=hidden_size,
|
||||
add_inputs=add_inputs,
|
||||
)
|
||||
|
||||
slice_offset += hidden_size
|
||||
assert_close(data.ref_out_tensor, data.our_out_tensor)
|
||||
|
||||
|
||||
# General tests params that tests for variations in all dimensions
|
||||
# except hidden_size.
|
||||
test_params = {
|
||||
"hidden_sizes": [2049],
|
||||
"batches": [4],
|
||||
"num_loras": [4],
|
||||
"max_ranks": [32],
|
||||
}
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
DEVICES = [f"xpu:{0}"]
|
||||
SEED = [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batches", test_params["batches"])
|
||||
@pytest.mark.parametrize("num_loras", test_params["num_loras"])
|
||||
@pytest.mark.parametrize("rank", test_params["max_ranks"])
|
||||
@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("seed", SEED)
|
||||
@pytest.mark.parametrize("op_type", ["shrink", "expand"])
|
||||
@pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform")
|
||||
def test_bgmv(
|
||||
batches: int,
|
||||
num_loras: int,
|
||||
rank: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
seed: int,
|
||||
op_type: str,
|
||||
):
|
||||
if op_type == "shrink":
|
||||
check_bgmv_shrink(
|
||||
batches=batches,
|
||||
num_loras=num_loras,
|
||||
rank=rank,
|
||||
hidden_size=hidden_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
scaling=0.5,
|
||||
)
|
||||
else:
|
||||
check_bgmv_expand(
|
||||
batches=batches,
|
||||
num_loras=num_loras,
|
||||
rank=rank,
|
||||
hidden_size=hidden_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
add_inputs=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batches", test_params["batches"])
|
||||
@pytest.mark.parametrize("num_loras", test_params["num_loras"])
|
||||
@pytest.mark.parametrize("rank", test_params["max_ranks"])
|
||||
@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
|
||||
@pytest.mark.parametrize("nslices", [2, 3])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("seed", SEED)
|
||||
@pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform")
|
||||
def test_bgmv_expand_nslices(
|
||||
batches: int,
|
||||
num_loras: int,
|
||||
rank: int,
|
||||
hidden_size: int,
|
||||
nslices: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
seed: int,
|
||||
):
|
||||
check_bgmv_expand_slice(
|
||||
batches=batches,
|
||||
num_loras=num_loras,
|
||||
rank=rank,
|
||||
hidden_size=hidden_size,
|
||||
nslices=nslices,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
add_inputs=True,
|
||||
)
|
||||
@@ -2,6 +2,9 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass
|
||||
|
||||
from packaging.version import Version
|
||||
from transformers import __version__ as TRANSFORMERS_VERSION
|
||||
|
||||
import vllm
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.lora.request import LoRARequest
|
||||
@@ -18,15 +21,25 @@ class TestConfig:
|
||||
enable_tower_connector_lora: bool = False
|
||||
max_model_len: int = 8192
|
||||
gpu_memory_utilization: float = 0.85
|
||||
mm_processor_kwargs: dict[str, int] | None = None
|
||||
mm_processor_kwargs: dict[str, object] | None = None
|
||||
mm_processor_cache_gb: float = 4
|
||||
|
||||
def __post_init__(self):
|
||||
if self.mm_processor_kwargs is None:
|
||||
self.mm_processor_kwargs = {
|
||||
"min_pixels": 28 * 28,
|
||||
"max_pixels": 1280 * 28 * 28,
|
||||
}
|
||||
# There is a bug in transformers v4 where size is ignored by
|
||||
# `Qwen2VLProcessor.__call__`
|
||||
if Version(TRANSFORMERS_VERSION) < Version("5.2.0"):
|
||||
self.mm_processor_kwargs = {
|
||||
"min_pixels": 28 * 28,
|
||||
"max_pixels": 1280 * 28 * 28,
|
||||
}
|
||||
else:
|
||||
self.mm_processor_kwargs = {
|
||||
"size": {
|
||||
"shortest_edge": 28 * 28,
|
||||
"longest_edge": 1280 * 28 * 28,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class Qwen2VLTester:
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Pytest configuration for vLLM multimodal tests."""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
@@ -9,6 +10,23 @@ import torch
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
"""Early ROCm configuration that must happen before test collection."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
# Disable skinny GEMM on ROCm to avoid non-deterministic results
|
||||
# from atomic reductions in wvSplitKrc kernel.
|
||||
# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
|
||||
os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
|
||||
warnings.warn(
|
||||
"ROCm: Set VLLM_ROCM_USE_SKINNY_GEMM=0 to avoid non-deterministic "
|
||||
"results from skinny GEMM atomic reductions",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
"""Configure ROCm-specific settings based on collected tests."""
|
||||
if not current_platform.is_rocm():
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for LlamaNemotronVL embedding model (nvidia/llama-nemotron-embed-vl-1b-v2).
|
||||
|
||||
This model uses SigLIP vision encoder with bidirectional LLaMA for embeddings.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
|
||||
from ...utils import check_embeddings_close
|
||||
|
||||
# Prefixes used by the model API
|
||||
QUERY_PREFIX = "query: "
|
||||
PASSAGE_PREFIX = "passage: "
|
||||
|
||||
# Text prompts for text-only embedding
|
||||
HF_TEXT_PROMPTS = [
|
||||
# T -> X (text embedding queries)
|
||||
f"{QUERY_PREFIX}The label of the object is stop sign",
|
||||
f"{QUERY_PREFIX}cherry blossom",
|
||||
]
|
||||
|
||||
# Image prompts using the model's expected format
|
||||
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
# I -> X (image embedding as passage/document)
|
||||
"stop_sign": f"{PASSAGE_PREFIX}<image>",
|
||||
"cherry_blossom": f"{PASSAGE_PREFIX}<image>",
|
||||
}
|
||||
)
|
||||
|
||||
MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
|
||||
|
||||
|
||||
def _run_test(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
input_texts: list[str],
|
||||
input_images: PromptImageInput,
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Run embedding comparison test between HF and vLLM.
|
||||
|
||||
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
|
||||
"""
|
||||
# Run vLLM inference first
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=2048,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.embed(input_texts, images=input_images)
|
||||
|
||||
# Run HF inference using the model's encode_queries/encode_documents API
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
|
||||
hf_outputs = []
|
||||
for text, image in zip(input_texts, input_images):
|
||||
with torch.inference_mode():
|
||||
if text.startswith(QUERY_PREFIX):
|
||||
# Strip prefix and use encode_queries for query texts
|
||||
query_text = text[len(QUERY_PREFIX) :]
|
||||
embedding = hf_model.model.encode_queries([query_text])
|
||||
elif text.startswith(PASSAGE_PREFIX):
|
||||
# Strip prefix and use encode_documents for passages/images
|
||||
passage_text = text[len(PASSAGE_PREFIX) :]
|
||||
if image is not None:
|
||||
# Image document - pass image to encode_documents
|
||||
embedding = hf_model.model.encode_documents(
|
||||
images=[image],
|
||||
texts=[passage_text],
|
||||
)
|
||||
else:
|
||||
# Text-only document
|
||||
embedding = hf_model.model.encode_documents(
|
||||
texts=[passage_text]
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
|
||||
)
|
||||
|
||||
hf_outputs.append(embedding[0].tolist())
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_outputs,
|
||||
embeddings_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
def test_models_text(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Test text-only embedding."""
|
||||
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
|
||||
input_texts = [text for text, _ in input_texts_images]
|
||||
input_images = [image for _, image in input_texts_images]
|
||||
|
||||
_run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
input_texts,
|
||||
input_images, # type: ignore
|
||||
model,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
def test_models_image(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Test image embedding."""
|
||||
input_texts_images = [
|
||||
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
|
||||
]
|
||||
input_texts = [text for text, _ in input_texts_images]
|
||||
input_images = [image for _, image in input_texts_images]
|
||||
|
||||
_run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
input_texts,
|
||||
input_images,
|
||||
model,
|
||||
dtype=dtype,
|
||||
)
|
||||
@@ -150,8 +150,11 @@ class TestGemma3nAudioTensorLogic:
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", [GEMMA3_MODEL_ID])
|
||||
@pytest.mark.parametrize("mm_processor_kwargs", [{}])
|
||||
def test_get_image_size_with_most_features(
|
||||
image_assets: ImageTestAssets, model_id: str
|
||||
image_assets: ImageTestAssets,
|
||||
model_id: str,
|
||||
mm_processor_kwargs: dict[str, object],
|
||||
):
|
||||
ctx = build_model_context(
|
||||
model_id,
|
||||
@@ -160,15 +163,14 @@ def test_get_image_size_with_most_features(
|
||||
)
|
||||
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
|
||||
|
||||
hf_processor_mm_kwargs: dict[str, object] = {}
|
||||
hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
|
||||
hf_processor = processor.info.get_hf_processor(**mm_processor_kwargs)
|
||||
|
||||
max_image_size = processor.info.get_image_size_with_most_features()
|
||||
max_tokens = processor.info.get_num_image_tokens(
|
||||
image_width=max_image_size.width,
|
||||
image_height=max_image_size.height,
|
||||
processor=hf_processor,
|
||||
mm_kwargs=hf_processor_mm_kwargs,
|
||||
mm_kwargs=mm_processor_kwargs,
|
||||
)
|
||||
|
||||
prompt = "<start_of_image>"
|
||||
@@ -179,7 +181,7 @@ def test_get_image_size_with_most_features(
|
||||
processed_inputs = processor(
|
||||
prompt,
|
||||
mm_items=processor.info.parse_mm_data(mm_data),
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
hf_processor_mm_kwargs=mm_processor_kwargs,
|
||||
)
|
||||
mm_kwargs_data = processed_inputs["mm_kwargs"].get_data()
|
||||
num_patches_tensor = mm_kwargs_data["num_patches"]
|
||||
|
||||
@@ -0,0 +1,358 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are
|
||||
correctly assigned to audio/image/video token positions.
|
||||
|
||||
Regression test for: https://github.com/vllm-project/vllm/issues/34506
|
||||
- Non-interleaved mixed modalities (audio + image + video) should correctly
|
||||
assign audio embeddings to audio positions, image to image, video to video.
|
||||
- Interleaved (use_audio_in_video) should also work correctly.
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
check_interleaved_audio_video,
|
||||
merge_interleaved_embeddings,
|
||||
)
|
||||
|
||||
# Fake token IDs
|
||||
AUDIO_TOKEN_ID = 1001
|
||||
IMAGE_TOKEN_ID = 1002
|
||||
VIDEO_TOKEN_ID = 1003
|
||||
TEXT_TOKEN_ID = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_token_seq(
|
||||
audio_n: int, image_n: int, video_n: int, text_prefix: int = 3, text_sep: int = 2
|
||||
):
|
||||
"""
|
||||
Build a flat token sequence:
|
||||
[text_prefix] [AUDIO * audio_n] [text_sep] [IMAGE * image_n]
|
||||
[text_sep] [VIDEO * video_n] [text_sep]
|
||||
Returns (input_ids tensor, is_multimodal mask, positions dict).
|
||||
"""
|
||||
tokens = (
|
||||
[TEXT_TOKEN_ID] * text_prefix
|
||||
+ [AUDIO_TOKEN_ID] * audio_n
|
||||
+ [TEXT_TOKEN_ID] * text_sep
|
||||
+ [IMAGE_TOKEN_ID] * image_n
|
||||
+ [TEXT_TOKEN_ID] * text_sep
|
||||
+ [VIDEO_TOKEN_ID] * video_n
|
||||
+ [TEXT_TOKEN_ID] * text_sep
|
||||
)
|
||||
input_ids = torch.tensor(tokens)
|
||||
is_multimodal = (
|
||||
(input_ids == AUDIO_TOKEN_ID)
|
||||
| (input_ids == IMAGE_TOKEN_ID)
|
||||
| (input_ids == VIDEO_TOKEN_ID)
|
||||
)
|
||||
return input_ids, is_multimodal
|
||||
|
||||
|
||||
def make_interleaved_seq(
|
||||
video_chunks: list[int], audio_chunks: list[int], text_prefix: int = 2
|
||||
):
|
||||
"""
|
||||
Build an interleaved sequence like use_audio_in_video:
|
||||
[text] [V*v0] [A*a0] [V*v1] [A*a1] ...
|
||||
"""
|
||||
tokens = [TEXT_TOKEN_ID] * text_prefix
|
||||
for v, a in zip(video_chunks, audio_chunks):
|
||||
tokens += [VIDEO_TOKEN_ID] * v + [AUDIO_TOKEN_ID] * a
|
||||
input_ids = torch.tensor(tokens)
|
||||
is_multimodal = (input_ids == VIDEO_TOKEN_ID) | (input_ids == AUDIO_TOKEN_ID)
|
||||
return input_ids, is_multimodal
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for check_interleaved_audio_video
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCheckInterleavedAudioVideo:
|
||||
def test_non_interleaved_audio_then_video(self):
|
||||
"""Audio entirely before video → not interleaved."""
|
||||
input_ids, is_multimodal = make_token_seq(5, 0, 4)
|
||||
is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
|
||||
is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
|
||||
assert not check_interleaved_audio_video(
|
||||
is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
|
||||
)
|
||||
|
||||
def test_non_interleaved_with_image(self):
|
||||
"""Audio + image + video (the mixed_modalities case) → not interleaved."""
|
||||
input_ids, is_multimodal = make_token_seq(5, 4, 6)
|
||||
is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
|
||||
is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
|
||||
assert not check_interleaved_audio_video(
|
||||
is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
|
||||
)
|
||||
|
||||
def test_no_audio(self):
|
||||
"""Video only → not interleaved."""
|
||||
input_ids, is_multimodal = make_token_seq(0, 0, 6)
|
||||
is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
|
||||
is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
|
||||
assert not check_interleaved_audio_video(
|
||||
is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
|
||||
)
|
||||
|
||||
def test_interleaved(self):
|
||||
"""V A V A interleaved → True."""
|
||||
input_ids, is_multimodal = make_interleaved_seq([4, 4], [3, 3])
|
||||
is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
|
||||
is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
|
||||
assert check_interleaved_audio_video(
|
||||
is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for embed_input_ids via a minimal mock
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_mock_model(hidden: int = 8):
|
||||
"""
|
||||
Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration
|
||||
that has enough structure to run embed_input_ids.
|
||||
"""
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
model = Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration)
|
||||
|
||||
# Config with token IDs
|
||||
cfg = Mock()
|
||||
cfg.video_token_index = VIDEO_TOKEN_ID
|
||||
cfg.audio_token_index = AUDIO_TOKEN_ID
|
||||
model.config = cfg
|
||||
|
||||
# embed_input_ids: simply embed each token as a one-hot-like vector
|
||||
# token_id * ones so we can verify which embedding ends up where.
|
||||
def fake_lm_embed(ids: torch.Tensor) -> torch.Tensor:
|
||||
# Use .clone() so the tensor is contiguous (expand() creates a strided
|
||||
# view with shared memory, which masked_scatter_ cannot handle).
|
||||
return ids.float().unsqueeze(-1).expand(-1, hidden).clone()
|
||||
|
||||
lang_model = Mock()
|
||||
lang_model.embed_input_ids = fake_lm_embed
|
||||
model.get_language_model = Mock(return_value=lang_model)
|
||||
|
||||
# _embed_text_input_ids: delegate to SupportsMultiModal's implementation
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal
|
||||
|
||||
model._embed_text_input_ids = (
|
||||
lambda *a, **kw: SupportsMultiModal._embed_text_input_ids(model, *a, **kw)
|
||||
)
|
||||
|
||||
# super().embed_input_ids → use SupportsMultiModal.embed_input_ids
|
||||
def fake_super_embed(
|
||||
ids, mm_embs=None, *, is_multimodal=None, handle_oov_mm_token=False
|
||||
):
|
||||
return SupportsMultiModal.embed_input_ids(
|
||||
model,
|
||||
ids,
|
||||
mm_embs,
|
||||
is_multimodal=is_multimodal,
|
||||
handle_oov_mm_token=handle_oov_mm_token,
|
||||
)
|
||||
|
||||
# Bind embed_input_ids as the real method
|
||||
model.embed_input_ids = (
|
||||
lambda *a, **kw: Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(
|
||||
model, *a, **kw
|
||||
)
|
||||
)
|
||||
|
||||
# Store super-embed for use inside the method
|
||||
model._super_embed_input_ids = fake_super_embed
|
||||
|
||||
return model, hidden
|
||||
|
||||
|
||||
def build_mm_embeds(
|
||||
audio_n, image_n, video_n, hidden, audio_val=10.0, image_val=20.0, video_val=30.0
|
||||
):
|
||||
"""
|
||||
Build multimodal_embeddings list in position order (audio, image, video).
|
||||
Each embedding is filled with a distinct constant so we can verify placement.
|
||||
"""
|
||||
embs = []
|
||||
if audio_n:
|
||||
embs.append(torch.full((audio_n, hidden), audio_val))
|
||||
if image_n:
|
||||
embs.append(torch.full((image_n, hidden), image_val))
|
||||
if video_n:
|
||||
embs.append(torch.full((video_n, hidden), video_val))
|
||||
return embs
|
||||
|
||||
|
||||
class TestEmbedInputIds:
|
||||
def _run(self, audio_n, image_n, video_n, hidden=8):
|
||||
"""
|
||||
Run embed_input_ids for a non-interleaved mixed-modality sequence.
|
||||
Returns (result_embeds, input_ids, is_multimodal).
|
||||
"""
|
||||
input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n)
|
||||
mm_embeds = build_mm_embeds(audio_n, image_n, video_n, hidden)
|
||||
|
||||
model, _ = make_mock_model(hidden)
|
||||
result = model.embed_input_ids(
|
||||
input_ids, mm_embeds, is_multimodal=is_multimodal
|
||||
)
|
||||
return result, input_ids, is_multimodal
|
||||
|
||||
def test_audio_only(self):
|
||||
"""Audio-only: audio positions get audio embeddings."""
|
||||
audio_n, hidden = 5, 8
|
||||
audio_val = 10.0
|
||||
result, input_ids, is_multimodal = self._run(audio_n, 0, 0, hidden)
|
||||
|
||||
audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
|
||||
"Audio positions should get audio embeddings"
|
||||
)
|
||||
|
||||
def test_video_only(self):
|
||||
"""Video-only: video positions get video embeddings."""
|
||||
video_n, hidden = 6, 8
|
||||
video_val = 30.0
|
||||
result, input_ids, is_multimodal = self._run(0, 0, video_n, hidden)
|
||||
|
||||
video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
|
||||
"Video positions should get video embeddings"
|
||||
)
|
||||
|
||||
def test_mixed_modalities_audio_goes_to_audio_pos(self):
|
||||
"""
|
||||
Regression test for GitHub issue #34506:
|
||||
With audio + image + video (non-interleaved), audio positions must
|
||||
receive audio embeddings (not image or video embeddings).
|
||||
"""
|
||||
audio_n, image_n, video_n, hidden = 5, 4, 6, 8
|
||||
audio_val, image_val, video_val = 10.0, 20.0, 30.0
|
||||
|
||||
result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden)
|
||||
|
||||
audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
image_pos = (input_ids == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
|
||||
mean_a = result[audio_pos].mean().item()
|
||||
assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
|
||||
f"Audio emb wrong: expected {audio_val}, got mean={mean_a:.1f}"
|
||||
)
|
||||
|
||||
mean_i = result[image_pos].mean().item()
|
||||
assert result[image_pos].allclose(torch.full((image_n, hidden), image_val)), (
|
||||
f"Image emb wrong: expected {image_val}, got mean={mean_i:.1f}"
|
||||
)
|
||||
|
||||
mean_v = result[video_pos].mean().item()
|
||||
assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
|
||||
f"Video emb wrong: expected {video_val}, got mean={mean_v:.1f}"
|
||||
)
|
||||
|
||||
def test_text_positions_unchanged(self):
|
||||
"""Text positions should keep their text embeddings."""
|
||||
audio_n, image_n, video_n, hidden = 3, 2, 4, 8
|
||||
result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden)
|
||||
|
||||
text_pos = (~is_multimodal).nonzero(as_tuple=True)[0]
|
||||
# Text tokens have value TEXT_TOKEN_ID=0, so embed → 0.0
|
||||
assert result[text_pos].allclose(torch.zeros(len(text_pos), hidden)), (
|
||||
"Text positions should keep text embeddings"
|
||||
)
|
||||
|
||||
def test_interleaved_use_audio_in_video(self):
|
||||
"""
|
||||
Interleaved (use_audio_in_video): video chunks interleaved with audio.
|
||||
Video embeddings must go to video positions, audio to audio positions.
|
||||
"""
|
||||
hidden = 8
|
||||
audio_val, video_val = 10.0, 30.0
|
||||
# Two video chunks of 4, two audio chunks of 3
|
||||
video_chunks = [4, 4]
|
||||
audio_chunks = [3, 3]
|
||||
input_ids, is_multimodal = make_interleaved_seq(video_chunks, audio_chunks)
|
||||
|
||||
video_n = sum(video_chunks) # 8
|
||||
audio_n = sum(audio_chunks) # 6
|
||||
|
||||
# mm_embeds come in [video, audio] order (video feature first in
|
||||
# mm_features when positions are the same for use_audio_in_video)
|
||||
mm_embeds = [
|
||||
torch.full((video_n, hidden), video_val),
|
||||
torch.full((audio_n, hidden), audio_val),
|
||||
]
|
||||
|
||||
model, _ = make_mock_model(hidden)
|
||||
result = model.embed_input_ids(
|
||||
input_ids, mm_embeds, is_multimodal=is_multimodal
|
||||
)
|
||||
|
||||
video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
|
||||
|
||||
assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
|
||||
"Interleaved: video positions should get video embeddings"
|
||||
)
|
||||
|
||||
assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
|
||||
"Interleaved: audio positions should get audio embeddings"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for merge_interleaved_embeddings helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMergeInterleavedEmbeddings:
|
||||
def test_basic_interleaved(self):
|
||||
"""Video chunks + audio chunks scattered to correct positions."""
|
||||
hidden = 4
|
||||
input_ids, is_multimodal = make_interleaved_seq([3, 3], [2, 2])
|
||||
|
||||
is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
|
||||
is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
|
||||
num_video = is_video.sum().item() # 6
|
||||
num_audio = is_audio.sum().item() # 4
|
||||
|
||||
inputs_embeds = torch.zeros(len(input_ids), hidden)
|
||||
mm_embeds = [
|
||||
torch.full((num_video, hidden), 30.0),
|
||||
torch.full((num_audio, hidden), 10.0),
|
||||
]
|
||||
|
||||
result = merge_interleaved_embeddings(
|
||||
inputs_embeds,
|
||||
mm_embeds,
|
||||
is_video,
|
||||
is_audio,
|
||||
is_multimodal,
|
||||
num_video,
|
||||
num_audio,
|
||||
)
|
||||
|
||||
video_pos = is_video.nonzero(as_tuple=True)[0]
|
||||
audio_pos = is_audio.nonzero(as_tuple=True)[0]
|
||||
assert result[video_pos].allclose(torch.full((num_video, hidden), 30.0))
|
||||
assert result[audio_pos].allclose(torch.full((num_audio, hidden), 10.0))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -2,6 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
from packaging.version import Version
|
||||
from transformers import __version__ as TRANSFORMERS_VERSION
|
||||
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
|
||||
@@ -15,6 +17,16 @@ from ...utils import build_model_context
|
||||
[
|
||||
({}, 1426, (5704, 1176)),
|
||||
({"min_pixels": 64**2, "max_pixels": 512**2}, 330, (1320, 1176)),
|
||||
(
|
||||
{
|
||||
"size": {
|
||||
"shortest_edge": 64**2,
|
||||
"longest_edge": 512**2,
|
||||
},
|
||||
},
|
||||
330,
|
||||
(1320, 1176),
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_imgs", [1, 2])
|
||||
@@ -29,6 +41,12 @@ def test_processor_override(
|
||||
kwargs_on_init: bool,
|
||||
):
|
||||
"""Ensure Qwen2VLMultiModalProcessor handles min/max pixels properly."""
|
||||
if (
|
||||
Version(TRANSFORMERS_VERSION) < Version("5.2.0")
|
||||
and "size" in mm_processor_kwargs
|
||||
):
|
||||
pytest.skip("`size` ignored by `Qwen2VLProcessor.__call__`")
|
||||
|
||||
ctx = build_model_context(
|
||||
model_id,
|
||||
mm_processor_kwargs=mm_processor_kwargs if kwargs_on_init else None,
|
||||
@@ -60,21 +78,34 @@ def test_processor_override(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", ["Qwen/Qwen2-VL-2B-Instruct"])
|
||||
@pytest.mark.parametrize("max_pixels", [1280 * 28 * 28, 1283 * 28 * 28])
|
||||
@pytest.mark.parametrize(
|
||||
"mm_processor_kwargs",
|
||||
[
|
||||
{"min_pixels": 28 * 28, "max_pixels": 1280 * 28 * 28},
|
||||
{"min_pixels": 28 * 28, "max_pixels": 1283 * 28 * 28},
|
||||
{"size": {"shortest_edge": 28 * 28, "longest_edge": 1280 * 28 * 28}},
|
||||
{"size": {"shortest_edge": 28 * 28, "longest_edge": 1283 * 28 * 28}},
|
||||
],
|
||||
)
|
||||
def test_get_image_size_with_most_features(
|
||||
image_assets: ImageTestAssets,
|
||||
model_id: str,
|
||||
max_pixels: int,
|
||||
mm_processor_kwargs: dict[str, object],
|
||||
):
|
||||
if (
|
||||
Version(TRANSFORMERS_VERSION) < Version("5.2.0")
|
||||
and "size" in mm_processor_kwargs
|
||||
):
|
||||
pytest.skip("`size` ignored by `Qwen2VLProcessor.__call__`")
|
||||
|
||||
ctx = build_model_context(
|
||||
model_id,
|
||||
mm_processor_kwargs={"max_pixels": max_pixels},
|
||||
mm_processor_kwargs=mm_processor_kwargs,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
|
||||
|
||||
hf_processor_mm_kwargs: dict[str, object] = {}
|
||||
hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
|
||||
hf_processor = processor.info.get_hf_processor(**mm_processor_kwargs)
|
||||
merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
|
||||
|
||||
max_image_size = processor.info.get_image_size_with_most_features()
|
||||
@@ -82,7 +113,7 @@ def test_get_image_size_with_most_features(
|
||||
image_width=max_image_size.width,
|
||||
image_height=max_image_size.height,
|
||||
image_processor=hf_processor.image_processor,
|
||||
mm_kwargs=hf_processor_mm_kwargs,
|
||||
mm_kwargs=mm_processor_kwargs,
|
||||
)
|
||||
|
||||
prompt = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
@@ -91,7 +122,7 @@ def test_get_image_size_with_most_features(
|
||||
processed_inputs = processor(
|
||||
prompt,
|
||||
mm_items=processor.info.parse_mm_data(mm_data),
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
hf_processor_mm_kwargs=mm_processor_kwargs,
|
||||
)
|
||||
grid_thw = processed_inputs["mm_kwargs"].get_data()["image_grid_thw"].tolist()
|
||||
t, h, w = grid_thw[0]
|
||||
|
||||
@@ -206,6 +206,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"BailingMoeV2ForCausalLM": _HfExamplesInfo(
|
||||
"inclusionAI/Ling-mini-2.0", trust_remote_code=True
|
||||
),
|
||||
"BailingMoeV2_5ForCausalLM": _HfExamplesInfo(
|
||||
"inclusionAI/Ring-2.5-1T", trust_remote_code=True
|
||||
),
|
||||
"BambaForCausalLM": _HfExamplesInfo(
|
||||
"ibm-ai-platform/Bamba-9B-v1",
|
||||
extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"},
|
||||
@@ -595,6 +598,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
|
||||
"ColModernVBertForRetrieval": _HfExamplesInfo(
|
||||
"ModernVBERT/colmodernvbert-merged",
|
||||
),
|
||||
"LlamaNemotronVLModel": _HfExamplesInfo(
|
||||
"nvidia/llama-nemotron-embed-vl-1b-v2", trust_remote_code=True
|
||||
),
|
||||
"LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
|
||||
"Phi3VForCausalLM": _HfExamplesInfo(
|
||||
"TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
|
||||
|
||||
@@ -85,34 +85,34 @@ def can_initialize(
|
||||
)
|
||||
)
|
||||
def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
|
||||
can_initialize(
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
|
||||
hf_overrides=HF_OVERRIDE_MM,
|
||||
extra_args=["--moe-backend=flashinfer_cutlass"],
|
||||
)
|
||||
|
||||
|
||||
def test_llama4_fp8_tensor_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
|
||||
can_initialize(
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
|
||||
hf_overrides=HF_OVERRIDE_MM,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
|
||||
def test_llama4_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
|
||||
can_initialize(
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", hf_overrides=HF_OVERRIDE_MM
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
|
||||
hf_overrides=HF_OVERRIDE_MM,
|
||||
extra_args=["--moe-backend=flashinfer_cutlass"],
|
||||
)
|
||||
|
||||
|
||||
def test_llama4_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
|
||||
can_initialize(
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", hf_overrides=HF_OVERRIDE_MM
|
||||
"nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
|
||||
hf_overrides=HF_OVERRIDE_MM,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
|
||||
@@ -120,8 +120,11 @@ def test_llama4_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
|
||||
def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1")
|
||||
can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT)
|
||||
can_initialize(
|
||||
"deepseek-ai/DeepSeek-V3.1",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=deep_gemm"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
@@ -131,27 +134,35 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):
|
||||
)
|
||||
)
|
||||
def test_deepseek_fp8_block_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
|
||||
can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT)
|
||||
can_initialize(
|
||||
"deepseek-ai/DeepSeek-V3.1",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_cutlass"],
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
|
||||
can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT)
|
||||
can_initialize(
|
||||
"deepseek-ai/DeepSeek-V3.1",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
|
||||
can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT)
|
||||
can_initialize(
|
||||
"nvidia/DeepSeek-R1-0528-FP4-v2",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_cutlass"],
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
|
||||
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
|
||||
can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT)
|
||||
can_initialize(
|
||||
"nvidia/DeepSeek-R1-0528-FP4-v2",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
|
||||
## GPT-OSS ##
|
||||
@@ -184,5 +195,8 @@ def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
|
||||
def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
|
||||
can_initialize("Qwen/Qwen3-Next-80B-A3B-Instruct", hf_overrides=HF_OVERRIDE_TEXT)
|
||||
can_initialize(
|
||||
"Qwen/Qwen3-Next-80B-A3B-Instruct",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
@@ -1327,6 +1327,57 @@ def multi_gpu_test(*, num_gpus: int):
|
||||
return wrapper
|
||||
|
||||
|
||||
def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None):
|
||||
"""
|
||||
Mark a test to only run when the GPU count falls within [min_gpus, max_gpus].
|
||||
|
||||
Examples:
|
||||
@gpu_tier_mark(min_gpus=2) # only on multi-GPU
|
||||
@gpu_tier_mark(max_gpus=1) # only on single-GPU
|
||||
@gpu_tier_mark(min_gpus=2, max_gpus=4) # 2-4 GPUs only
|
||||
"""
|
||||
gpu_count = cuda_device_count_stateless()
|
||||
marks = []
|
||||
|
||||
if min_gpus > 1:
|
||||
marks.append(pytest.mark.distributed(num_gpus=min_gpus))
|
||||
|
||||
reasons = []
|
||||
if gpu_count < min_gpus:
|
||||
reasons.append(f"Need at least {min_gpus} GPUs (have {gpu_count})")
|
||||
if max_gpus is not None and gpu_count > max_gpus:
|
||||
reasons.append(f"Need at most {max_gpus} GPUs (have {gpu_count})")
|
||||
|
||||
if reasons:
|
||||
marks.append(pytest.mark.skipif(True, reason="; ".join(reasons)))
|
||||
|
||||
return marks
|
||||
|
||||
|
||||
def single_gpu_only(f=None):
|
||||
"""Skip this test when running in a multi-GPU environment."""
|
||||
marks = gpu_tier_mark(max_gpus=1)
|
||||
|
||||
def wrapper(func):
|
||||
for mark in reversed(marks):
|
||||
func = mark(func)
|
||||
return func
|
||||
|
||||
return wrapper(f) if f is not None else wrapper
|
||||
|
||||
|
||||
def multi_gpu_only(*, num_gpus: int = 2):
|
||||
"""Skip this test when running on fewer than num_gpus GPUs."""
|
||||
marks = gpu_tier_mark(min_gpus=num_gpus)
|
||||
|
||||
def wrapper(f):
|
||||
for mark in reversed(marks):
|
||||
f = mark(f)
|
||||
return f
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
async def completions_with_server_args(
|
||||
prompts: list[str],
|
||||
model_name: str,
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import Any
|
||||
import pytest
|
||||
import torch._dynamo.config as dynamo_config
|
||||
|
||||
from tests.utils import large_gpu_mark, single_gpu_only
|
||||
from vllm import SamplingParams
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.platforms import current_platform
|
||||
@@ -36,6 +37,7 @@ default_params = dict(
|
||||
)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
def test_without_spec_decoding(
|
||||
sample_json_schema,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
@@ -95,6 +97,8 @@ def test_without_spec_decoding(
|
||||
run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=16)
|
||||
def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test consistency and acceptance rates with some different combos of
|
||||
preemption, executor, async scheduling, prefill chunking,
|
||||
|
||||
+258
-149
@@ -9,7 +9,13 @@ import pytest
|
||||
import torch
|
||||
|
||||
from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline
|
||||
from tests.utils import get_attn_backend_list_based_on_platform, large_gpu_mark
|
||||
from tests.utils import (
|
||||
get_attn_backend_list_based_on_platform,
|
||||
large_gpu_mark,
|
||||
multi_gpu_marks,
|
||||
multi_gpu_only,
|
||||
single_gpu_only,
|
||||
)
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.assets.base import VLLM_S3_BUCKET_URL
|
||||
from vllm.assets.image import VLM_IMAGES_DIR
|
||||
@@ -160,6 +166,8 @@ def reset_torch_dynamo():
|
||||
},
|
||||
],
|
||||
)
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=20)
|
||||
def test_ngram_and_suffix_correctness(
|
||||
speculative_config: dict,
|
||||
model_name: str,
|
||||
@@ -175,6 +183,8 @@ def test_ngram_and_suffix_correctness(
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=20)
|
||||
def test_suffix_decoding_acceptance(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
@@ -242,6 +252,8 @@ def test_suffix_decoding_acceptance(
|
||||
],
|
||||
ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
|
||||
)
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=24)
|
||||
def test_speculators_model_integration(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
@@ -319,137 +331,7 @@ def test_speculators_model_integration(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
[
|
||||
"model_setup",
|
||||
"mm_enabled",
|
||||
"enable_chunked_prefill",
|
||||
"model_impl",
|
||||
"expected_accuracy_threshold",
|
||||
],
|
||||
[
|
||||
(
|
||||
("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.8, # ref: 90%
|
||||
),
|
||||
(
|
||||
("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
|
||||
False,
|
||||
False,
|
||||
"transformers",
|
||||
0.8, # ref: 90%
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle3",
|
||||
"Qwen/Qwen3-VL-8B-Instruct",
|
||||
"taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.8, # ref: 90%
|
||||
marks=pytest.mark.skip(
|
||||
reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
|
||||
),
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle3",
|
||||
"Qwen/Qwen2.5-VL-7B-Instruct",
|
||||
"Rayzl/qwen2.5-vl-7b-eagle3-sgl",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.7, # TODO, update this with a reference value when re-enabling this case
|
||||
marks=pytest.mark.skip(
|
||||
reason="Skipping due to its head_dim not being a a multiple of 32"
|
||||
),
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-3.1-8B-Instruct",
|
||||
"yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
True,
|
||||
"auto",
|
||||
0.7, # ref: 75%-80%
|
||||
marks=large_gpu_mark(min_gb=40),
|
||||
), # works on 4x H100
|
||||
(
|
||||
(
|
||||
"eagle3",
|
||||
"meta-llama/Llama-3.1-8B-Instruct",
|
||||
"yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.7, # ref: 75%-80%
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
"morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
|
||||
4,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.8, # ref: 90%
|
||||
# marks=large_gpu_mark(min_gb=80),
|
||||
), # works on 4x H100
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
"morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
|
||||
4,
|
||||
),
|
||||
True,
|
||||
True,
|
||||
"auto",
|
||||
0.8, # ref: 90%
|
||||
marks=large_gpu_mark(min_gb=80),
|
||||
), # works on 4x H100
|
||||
(
|
||||
(
|
||||
"eagle",
|
||||
"eagle618/deepseek-v3-random",
|
||||
"eagle618/eagle-deepseek-v3-random",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.0, # dummy model, skip gsm8k check
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"qwen3_eagle3",
|
||||
"qwen3_eagle3-transformers",
|
||||
"qwen3_vl_eagle3",
|
||||
"qwen2_5_vl_eagle3",
|
||||
"llama3_eagle",
|
||||
"llama3_eagle3",
|
||||
"llama4_eagle",
|
||||
"llama4_eagle_mm",
|
||||
"deepseek_eagle",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
|
||||
def test_eagle_correctness(
|
||||
def _run_eagle_correctness(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
model_setup: tuple[str, str, str, int],
|
||||
@@ -460,14 +342,10 @@ def test_eagle_correctness(
|
||||
attn_backend: str,
|
||||
):
|
||||
"""
|
||||
Compare the outputs of a original LLM and a speculative LLM
|
||||
which should be the same when using eagle speculative decoding. Due to some variance
|
||||
in the engine, it is possible for some outputs to differ, so we expect that at least
|
||||
6/10 output tokens match exactly, and that the GSM8k accuracy is above
|
||||
a precomputed reference threshold for each model.
|
||||
Compare the outputs of an original LLM and a speculative LLM
|
||||
which should be the same when using eagle speculative decoding.
|
||||
"""
|
||||
if attn_backend == "TREE_ATTN":
|
||||
# TODO: Fix this flaky test
|
||||
pytest.skip(
|
||||
"TREE_ATTN is flaky in the test disable for now until it can be "
|
||||
"resolved (see https://github.com/vllm-project/vllm/issues/22922)"
|
||||
@@ -484,17 +362,17 @@ def test_eagle_correctness(
|
||||
f"transformers>={required}, but got {installed}"
|
||||
)
|
||||
|
||||
# Generate test prompts inside the function instead of using fixture
|
||||
test_prompts = get_test_prompts(mm_enabled)
|
||||
# Determine attention config
|
||||
# Scout requires default backend selection because vision encoder has
|
||||
# head_dim 88 being incompatible with FLASH_ATTN and needs to fall back
|
||||
# to Flex Attn
|
||||
|
||||
if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
|
||||
if current_platform.is_rocm():
|
||||
# TODO: Enable Flex Attn for spec_decode on ROCm
|
||||
pytest.skip("Flex Attn for spec_decode not supported on ROCm currently")
|
||||
attention_config = None # Let it fall back to default
|
||||
print(
|
||||
"FLASH_ATTN for spec_decode not supported on "
|
||||
"ROCm currently. Changing to FLEX_ATTENTION backend."
|
||||
)
|
||||
attention_config = {"backend": "FLEX_ATTENTION"}
|
||||
else:
|
||||
attention_config = None
|
||||
else:
|
||||
attention_config = {"backend": attn_backend}
|
||||
|
||||
@@ -509,7 +387,9 @@ def test_eagle_correctness(
|
||||
|
||||
if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
|
||||
if "deepseek" in model_setup[1].lower():
|
||||
pytest.skip("ROCM_AITER_FA for deepseek not supported on ROCm platform")
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
m.delenv("VLLM_MLA_DISABLE", raising=False)
|
||||
attention_config = {"backend": "TRITON_MLA"}
|
||||
else:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
|
||||
@@ -563,14 +443,235 @@ def test_eagle_correctness(
|
||||
print(f"ref_output: {ref_output.outputs[0].text}")
|
||||
print(f"spec_output: {spec_output.outputs[0].text}")
|
||||
|
||||
# Heuristic: expect at least 60% of the prompts to match exactly
|
||||
# Upon failure, inspect the outputs to check for inaccuracy.
|
||||
assert matches > int(0.6 * len(ref_outputs))
|
||||
del spec_llm
|
||||
torch.cuda.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
@pytest.mark.parametrize(
|
||||
[
|
||||
"model_setup",
|
||||
"mm_enabled",
|
||||
"enable_chunked_prefill",
|
||||
"model_impl",
|
||||
"expected_accuracy_threshold",
|
||||
],
|
||||
[
|
||||
(
|
||||
(
|
||||
"eagle",
|
||||
"eagle618/deepseek-v3-random",
|
||||
"eagle618/eagle-deepseek-v3-random",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.0,
|
||||
),
|
||||
],
|
||||
ids=["deepseek_eagle"],
|
||||
)
|
||||
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
|
||||
def test_eagle_correctness_light(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
model_setup: tuple[str, str, str, int],
|
||||
mm_enabled: bool,
|
||||
expected_accuracy_threshold: float,
|
||||
enable_chunked_prefill: bool,
|
||||
model_impl: str,
|
||||
attn_backend: str,
|
||||
):
|
||||
_run_eagle_correctness(
|
||||
monkeypatch,
|
||||
sampling_config,
|
||||
model_setup,
|
||||
mm_enabled,
|
||||
expected_accuracy_threshold,
|
||||
enable_chunked_prefill,
|
||||
model_impl,
|
||||
attn_backend,
|
||||
)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=24)
|
||||
@pytest.mark.parametrize(
|
||||
[
|
||||
"model_setup",
|
||||
"mm_enabled",
|
||||
"enable_chunked_prefill",
|
||||
"model_impl",
|
||||
"expected_accuracy_threshold",
|
||||
],
|
||||
[
|
||||
(
|
||||
("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.8,
|
||||
),
|
||||
(
|
||||
("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
|
||||
False,
|
||||
False,
|
||||
"transformers",
|
||||
0.8,
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle3",
|
||||
"Qwen/Qwen3-VL-8B-Instruct",
|
||||
"taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.8,
|
||||
marks=pytest.mark.skip(
|
||||
reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
|
||||
),
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle3",
|
||||
"Qwen/Qwen2.5-VL-7B-Instruct",
|
||||
"Rayzl/qwen2.5-vl-7b-eagle3-sgl",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.7,
|
||||
marks=pytest.mark.skip(
|
||||
reason="Skipping due to its head_dim not being a multiple of 32"
|
||||
),
|
||||
),
|
||||
(
|
||||
(
|
||||
"eagle3",
|
||||
"meta-llama/Llama-3.1-8B-Instruct",
|
||||
"yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.7,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"qwen3_eagle3",
|
||||
"qwen3_eagle3-transformers",
|
||||
"qwen3_vl_eagle3",
|
||||
"qwen2_5_vl_eagle3",
|
||||
"llama3_eagle3",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
|
||||
def test_eagle_correctness_medium(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
model_setup: tuple[str, str, str, int],
|
||||
mm_enabled: bool,
|
||||
expected_accuracy_threshold: float,
|
||||
enable_chunked_prefill: bool,
|
||||
model_impl: str,
|
||||
attn_backend: str,
|
||||
):
|
||||
_run_eagle_correctness(
|
||||
monkeypatch,
|
||||
sampling_config,
|
||||
model_setup,
|
||||
mm_enabled,
|
||||
expected_accuracy_threshold,
|
||||
enable_chunked_prefill,
|
||||
model_impl,
|
||||
attn_backend,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
[
|
||||
"model_setup",
|
||||
"mm_enabled",
|
||||
"enable_chunked_prefill",
|
||||
"model_impl",
|
||||
"expected_accuracy_threshold",
|
||||
],
|
||||
[
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-3.1-8B-Instruct",
|
||||
"yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
|
||||
1,
|
||||
),
|
||||
False,
|
||||
True,
|
||||
"auto",
|
||||
0.7,
|
||||
marks=large_gpu_mark(min_gb=40),
|
||||
id="llama3_eagle",
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
"morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
|
||||
4,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"auto",
|
||||
0.8,
|
||||
marks=multi_gpu_marks(num_gpus=4),
|
||||
id="llama4_eagle",
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
"morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
|
||||
4,
|
||||
),
|
||||
True,
|
||||
True,
|
||||
"auto",
|
||||
0.8,
|
||||
marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=80)],
|
||||
id="llama4_eagle_mm",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
|
||||
def test_eagle_correctness_heavy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
model_setup: tuple[str, str, str, int],
|
||||
mm_enabled: bool,
|
||||
expected_accuracy_threshold: float,
|
||||
enable_chunked_prefill: bool,
|
||||
model_impl: str,
|
||||
attn_backend: str,
|
||||
):
|
||||
_run_eagle_correctness(
|
||||
monkeypatch,
|
||||
sampling_config,
|
||||
model_setup,
|
||||
mm_enabled,
|
||||
expected_accuracy_threshold,
|
||||
enable_chunked_prefill,
|
||||
model_impl,
|
||||
attn_backend,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
["model_setup", "mm_enabled", "expected_accuracy_threshold"],
|
||||
[
|
||||
@@ -579,6 +680,8 @@ def test_eagle_correctness(
|
||||
],
|
||||
ids=["mimo", "deepseek"],
|
||||
)
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=20)
|
||||
def test_mtp_correctness(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sampling_config: SamplingParams,
|
||||
@@ -694,11 +797,13 @@ cases = [
|
||||
|
||||
@pytest.mark.parametrize("args", cases)
|
||||
@pytest.mark.parametrize("enforce_eager", [True, False])
|
||||
@single_gpu_only
|
||||
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
|
||||
args.enforce_eager = enforce_eager
|
||||
assert_draft_model_correctness(args)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
def test_draft_model_realistic_example():
|
||||
args = ArgsTest(
|
||||
target_model="Qwen/Qwen3-1.7B",
|
||||
@@ -713,6 +818,7 @@ def test_draft_model_realistic_example():
|
||||
assert_draft_model_correctness(args)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
def test_draft_model_parallel_drafting():
|
||||
args = ArgsTest(
|
||||
target_model="Qwen/Qwen3-1.7B",
|
||||
@@ -738,6 +844,7 @@ def test_draft_model_parallel_drafting():
|
||||
ids=["target_quantized", "draft_quantized"],
|
||||
)
|
||||
@pytest.mark.parametrize("enforce_eager", [True, False])
|
||||
@single_gpu_only
|
||||
def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
|
||||
tgt_model, draft_model = models
|
||||
sd_case = ArgsTest(
|
||||
@@ -749,6 +856,7 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
|
||||
assert_draft_model_correctness(sd_case)
|
||||
|
||||
|
||||
@multi_gpu_only(num_gpus=2)
|
||||
def test_draft_model_tensor_parallelism():
|
||||
"""Ensure spec decode works when running with TP > 1."""
|
||||
_skip_if_insufficient_gpus_for_tp(2)
|
||||
@@ -764,6 +872,7 @@ def test_draft_model_tensor_parallelism():
|
||||
assert_draft_model_correctness(sd_case)
|
||||
|
||||
|
||||
@multi_gpu_only(num_gpus=2)
|
||||
def test_draft_model_engine_args_tensor_parallelism():
|
||||
"""Ensure the vllm_config for the draft model is created correctly,
|
||||
and independently of the target model (quantization, TP, etc.)"""
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user