forked from Karylab-cklius/vllm
Compare commits
93
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3827c8c55a | ||
|
|
ade81f17fe | ||
|
|
6042e66cd5 | ||
|
|
9f9a675b23 | ||
|
|
a07c4c5939 | ||
|
|
d3a51da92a | ||
|
|
186ea22efe | ||
|
|
4a9c07a0a2 | ||
|
|
9d37941017 | ||
|
|
4171ff6dd9 | ||
|
|
13025e71e8 | ||
|
|
71dfce6aa6 | ||
|
|
2aa4140402 | ||
|
|
86c3b5a808 | ||
|
|
160424a937 | ||
|
|
9511a3f8ee | ||
|
|
de527e1cec | ||
|
|
1976356ee6 | ||
|
|
cbf8f7028c | ||
|
|
6831650c40 | ||
|
|
ed42507f6d | ||
|
|
9571e99945 | ||
|
|
c97234c08b | ||
|
|
b188bab441 | ||
|
|
15d76f74e2 | ||
|
|
8fd6975479 | ||
|
|
5d18bf8b32 | ||
|
|
0788ff0a15 | ||
|
|
d72b0be33c | ||
|
|
42489e43c2 | ||
|
|
af5e6afa0a | ||
|
|
ee59a7c615 | ||
|
|
709eadbb0b | ||
|
|
90fc7f9109 | ||
|
|
675ec59aa9 | ||
|
|
80e60a6133 | ||
|
|
26e722f906 | ||
|
|
2c619e5e3f | ||
|
|
8a685be8d9 | ||
|
|
2465071510 | ||
|
|
cd43673668 | ||
|
|
35d44b4557 | ||
|
|
8ad54a991b | ||
|
|
92510edc32 | ||
|
|
a6c137521c | ||
|
|
4572a06afe | ||
|
|
5cc29cfb8b | ||
|
|
8fae54faff | ||
|
|
f7967577f5 | ||
|
|
af770b8e7b | ||
|
|
2ff3e436ad | ||
|
|
c2c4c4611a | ||
|
|
f38f8c9742 | ||
|
|
ec1d30c0f6 | ||
|
|
e3b2324ec4 | ||
|
|
dbf0da817a | ||
|
|
3bbb2046ff | ||
|
|
576fe50333 | ||
|
|
a0e50a4260 | ||
|
|
9fa5b25a23 | ||
|
|
ea97750414 | ||
|
|
067c5d9ad1 | ||
|
|
f5972a872f | ||
|
|
a9e15e040d | ||
|
|
542ca66357 | ||
|
|
fc8456c336 | ||
|
|
9ce8fad2a9 | ||
|
|
c38b8d5a31 | ||
|
|
60da0e1544 | ||
|
|
9609b1f18d | ||
|
|
a0c7081695 | ||
|
|
34ce0ffd1f | ||
|
|
0de5333989 | ||
|
|
a87cc50859 | ||
|
|
761e63e541 | ||
|
|
d12d201409 | ||
|
|
b3ad37c5db | ||
|
|
14561fabfd | ||
|
|
c77f3e1207 | ||
|
|
012dee9233 | ||
|
|
f1c664545b | ||
|
|
c870eb9e0f | ||
|
|
6af03f2394 | ||
|
|
1a6cf39dec | ||
|
|
f91808ae0d | ||
|
|
33a0d43c71 | ||
|
|
80d93fd6da | ||
|
|
ec85340531 | ||
|
|
2ff4e51152 | ||
|
|
95642441d0 | ||
|
|
a7c9f7b7ec | ||
|
|
a4bd661fb3 | ||
|
|
3ef9fd0f98 |
@@ -1,25 +1,37 @@
|
||||
#!/bin/bash
|
||||
|
||||
# This script runs test inside the corresponding ROCm docker container.
|
||||
# This script runs tests inside the corresponding ROCm docker container.
|
||||
# It handles both single-node and multi-node test configurations.
|
||||
#
|
||||
# Multi-node detection: Instead of matching on fragile group names, we detect
|
||||
# multi-node jobs structurally by looking for the bracket command syntax
|
||||
# "[node0_cmds] && [node1_cmds]" or via the NUM_NODES environment variable.
|
||||
set -o pipefail
|
||||
|
||||
# Export Python path
|
||||
export PYTHONPATH=".."
|
||||
|
||||
# Print ROCm version
|
||||
echo "--- Confirming Clean Initial State"
|
||||
while true; do
|
||||
sleep 3
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
break
|
||||
fi
|
||||
done
|
||||
###############################################################################
|
||||
# Helper Functions
|
||||
###############################################################################
|
||||
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
wait_for_clean_gpus() {
|
||||
local timeout=${1:-300}
|
||||
local start=$SECONDS
|
||||
echo "--- Waiting for clean GPU state (timeout: ${timeout}s)"
|
||||
while true; do
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
return
|
||||
fi
|
||||
if (( SECONDS - start >= timeout )); then
|
||||
echo "Error: GPUs did not reach clean state within ${timeout}s" >&2
|
||||
exit 1
|
||||
fi
|
||||
sleep 3
|
||||
done
|
||||
}
|
||||
|
||||
# cleanup older docker images
|
||||
cleanup_docker() {
|
||||
# Get Docker's root directory
|
||||
docker_root=$(docker info -f '{{.DockerRootDir}}')
|
||||
@@ -28,15 +40,12 @@ cleanup_docker() {
|
||||
exit 1
|
||||
fi
|
||||
echo "Docker root directory: $docker_root"
|
||||
# Check disk usage of the filesystem where Docker's root directory is located
|
||||
|
||||
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
|
||||
# Define the threshold
|
||||
threshold=70
|
||||
if [ "$disk_usage" -gt "$threshold" ]; then
|
||||
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
|
||||
# Remove dangling images (those that are not tagged and not used by any container)
|
||||
docker image prune -f
|
||||
# Remove unused volumes / force the system prune for old images as well.
|
||||
docker volume prune -f && docker system prune --force --filter "until=72h" --all
|
||||
echo "Docker images and volumes cleanup completed."
|
||||
else
|
||||
@@ -45,193 +54,258 @@ cleanup_docker() {
|
||||
}
|
||||
|
||||
cleanup_network() {
|
||||
for node in $(seq 0 $((NUM_NODES-1))); do
|
||||
if docker pr -a -q -f name="node${node}" | grep -q .; then
|
||||
docker stop "node${node}"
|
||||
local max_nodes=${NUM_NODES:-2}
|
||||
for node in $(seq 0 $((max_nodes - 1))); do
|
||||
if docker ps -a -q -f name="node${node}" | grep -q .; then
|
||||
docker stop "node${node}" || true
|
||||
fi
|
||||
done
|
||||
if docker network ls | grep docker-net; then
|
||||
docker network rm docker-net
|
||||
if docker network ls | grep -q docker-net; then
|
||||
docker network rm docker-net || true
|
||||
fi
|
||||
}
|
||||
|
||||
# Call the cleanup docker function
|
||||
is_multi_node() {
|
||||
local cmds="$1"
|
||||
# Primary signal: NUM_NODES environment variable set by the pipeline
|
||||
if [[ "${NUM_NODES:-1}" -gt 1 ]]; then
|
||||
return 0
|
||||
fi
|
||||
# Fallback: detect the bracket syntax structurally
|
||||
# Pattern: [...] && [...] (per-node command arrays)
|
||||
if [[ "$cmds" =~ \[.*\].*\&\&.*\[.*\] ]]; then
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
###############################################################################
|
||||
# Pytest marker re-quoting
|
||||
#
|
||||
# When commands are passed through Buildkite -> shell -> $* -> bash -c,
|
||||
# quotes around pytest -m marker expressions get stripped:
|
||||
# pytest -v -s -m 'not cpu_test' v1/core
|
||||
# becomes:
|
||||
# pytest -v -s -m not cpu_test v1/core
|
||||
#
|
||||
# pytest then interprets "cpu_test" as a file path, not part of the marker.
|
||||
# This function detects unquoted multi-word marker expressions and re-quotes
|
||||
# them so they survive the final bash -c expansion.
|
||||
###############################################################################
|
||||
|
||||
re_quote_pytest_markers() {
|
||||
local cmds="$1"
|
||||
# Pattern: -m not <identifier> -> -m 'not <identifier>'
|
||||
# Handles the common cases: 'not cpu_test', 'not slow_test', etc.
|
||||
cmds=$(echo "$cmds" | sed -E "s/-m not ([a-zA-Z_][a-zA-Z0-9_]*)/-m 'not \1'/g")
|
||||
echo "$cmds"
|
||||
}
|
||||
|
||||
###############################################################################
|
||||
# ROCm-specific pytest command rewrites
|
||||
#
|
||||
# These apply ignore flags and environment overrides for tests that are not
|
||||
# yet supported or behave differently on ROCm hardware. Kept as a single
|
||||
# function so new exclusions are easy to add in one place.
|
||||
###############################################################################
|
||||
|
||||
apply_rocm_test_overrides() {
|
||||
local cmds="$1"
|
||||
|
||||
# --- Model registry filter ---
|
||||
if [[ $cmds == *"pytest -v -s models/test_registry.py"* ]]; then
|
||||
cmds=${cmds//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
|
||||
fi
|
||||
|
||||
# --- LoRA: disable custom paged attention ---
|
||||
if [[ $cmds == *"pytest -v -s lora"* ]]; then
|
||||
cmds=${cmds//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
|
||||
fi
|
||||
|
||||
# --- Kernel ignores ---
|
||||
if [[ $cmds == *" kernels/core"* ]]; then
|
||||
cmds="${cmds} \
|
||||
--ignore=kernels/core/test_fused_quant_layernorm.py \
|
||||
--ignore=kernels/core/test_permute_cols.py"
|
||||
fi
|
||||
|
||||
if [[ $cmds == *" kernels/attention"* ]]; then
|
||||
cmds="${cmds} \
|
||||
--ignore=kernels/attention/test_attention_selector.py \
|
||||
--ignore=kernels/attention/test_encoder_decoder_attn.py \
|
||||
--ignore=kernels/attention/test_flash_attn.py \
|
||||
--ignore=kernels/attention/test_flashinfer.py \
|
||||
--ignore=kernels/attention/test_prefix_prefill.py \
|
||||
--ignore=kernels/attention/test_cascade_flash_attn.py \
|
||||
--ignore=kernels/attention/test_mha_attn.py \
|
||||
--ignore=kernels/attention/test_lightning_attn.py \
|
||||
--ignore=kernels/attention/test_attention.py"
|
||||
fi
|
||||
|
||||
if [[ $cmds == *" kernels/quantization"* ]]; then
|
||||
cmds="${cmds} \
|
||||
--ignore=kernels/quantization/test_int8_quant.py \
|
||||
--ignore=kernels/quantization/test_machete_mm.py \
|
||||
--ignore=kernels/quantization/test_block_fp8.py \
|
||||
--ignore=kernels/quantization/test_block_int8.py \
|
||||
--ignore=kernels/quantization/test_marlin_gemm.py \
|
||||
--ignore=kernels/quantization/test_cutlass_scaled_mm.py \
|
||||
--ignore=kernels/quantization/test_int8_kernel.py"
|
||||
fi
|
||||
|
||||
if [[ $cmds == *" kernels/mamba"* ]]; then
|
||||
cmds="${cmds} \
|
||||
--ignore=kernels/mamba/test_mamba_mixer2.py \
|
||||
--ignore=kernels/mamba/test_causal_conv1d.py \
|
||||
--ignore=kernels/mamba/test_mamba_ssm_ssd.py"
|
||||
fi
|
||||
|
||||
if [[ $cmds == *" kernels/moe"* ]]; then
|
||||
cmds="${cmds} \
|
||||
--ignore=kernels/moe/test_moe.py \
|
||||
--ignore=kernels/moe/test_cutlass_moe.py \
|
||||
--ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
|
||||
fi
|
||||
|
||||
# --- Entrypoint ignores ---
|
||||
if [[ $cmds == *" entrypoints/openai "* ]]; then
|
||||
cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
|
||||
--ignore=entrypoints/openai/test_audio.py \
|
||||
--ignore=entrypoints/openai/test_shutdown.py \
|
||||
--ignore=entrypoints/openai/test_completion.py \
|
||||
--ignore=entrypoints/openai/test_models.py \
|
||||
--ignore=entrypoints/openai/test_lora_adapters.py \
|
||||
--ignore=entrypoints/openai/test_return_tokens_as_ids.py \
|
||||
--ignore=entrypoints/openai/test_root_path.py \
|
||||
--ignore=entrypoints/openai/test_tokenization.py \
|
||||
--ignore=entrypoints/openai/test_prompt_validation.py "}
|
||||
fi
|
||||
|
||||
if [[ $cmds == *" entrypoints/llm "* ]]; then
|
||||
cmds=${cmds//" entrypoints/llm "/" entrypoints/llm \
|
||||
--ignore=entrypoints/llm/test_chat.py \
|
||||
--ignore=entrypoints/llm/test_accuracy.py \
|
||||
--ignore=entrypoints/llm/test_init.py \
|
||||
--ignore=entrypoints/llm/test_prompt_validation.py "}
|
||||
fi
|
||||
|
||||
# Clean up escaped newlines from --ignore appends
|
||||
cmds=$(echo "$cmds" | sed 's/ \\ / /g')
|
||||
|
||||
echo "$cmds"
|
||||
}
|
||||
|
||||
###############################################################################
|
||||
# Main
|
||||
###############################################################################
|
||||
|
||||
# --- GPU initialization ---
|
||||
echo "--- Confirming Clean Initial State"
|
||||
wait_for_clean_gpus
|
||||
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
|
||||
# --- Docker housekeeping ---
|
||||
cleanup_docker
|
||||
|
||||
echo "--- Resetting GPUs"
|
||||
|
||||
echo "reset" > /opt/amdgpu/etc/gpu_state
|
||||
wait_for_clean_gpus
|
||||
|
||||
while true; do
|
||||
sleep 3
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
# --- Pull test image ---
|
||||
echo "--- Pulling container"
|
||||
image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
|
||||
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
docker pull "${image_name}"
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
|
||||
docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
# --- Prepare commands ---
|
||||
echo "--- Running container"
|
||||
|
||||
HF_CACHE="$(realpath ~)/huggingface"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
commands=$@
|
||||
commands="$*"
|
||||
echo "Raw commands: $commands"
|
||||
|
||||
commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"pytest -v -s basic_correctness/test_basic_correctness.py"}
|
||||
|
||||
if [[ $commands == *"pytest -v -s models/test_registry.py"* ]]; then
|
||||
commands=${commands//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
|
||||
fi
|
||||
|
||||
commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"pytest -v -s compile/test_basic_correctness.py"}
|
||||
|
||||
if [[ $commands == *"pytest -v -s lora"* ]]; then
|
||||
commands=${commands//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
|
||||
fi
|
||||
|
||||
#ignore certain kernels tests
|
||||
if [[ $commands == *" kernels/core"* ]]; then
|
||||
commands="${commands} \
|
||||
--ignore=kernels/core/test_fused_quant_layernorm.py \
|
||||
--ignore=kernels/core/test_permute_cols.py"
|
||||
fi
|
||||
|
||||
if [[ $commands == *" kernels/attention"* ]]; then
|
||||
commands="${commands} \
|
||||
--ignore=kernels/attention/test_attention_selector.py \
|
||||
--ignore=kernels/attention/test_encoder_decoder_attn.py \
|
||||
--ignore=kernels/attention/test_flash_attn.py \
|
||||
--ignore=kernels/attention/test_flashinfer.py \
|
||||
--ignore=kernels/attention/test_prefix_prefill.py \
|
||||
--ignore=kernels/attention/test_cascade_flash_attn.py \
|
||||
--ignore=kernels/attention/test_mha_attn.py \
|
||||
--ignore=kernels/attention/test_lightning_attn.py \
|
||||
--ignore=kernels/attention/test_attention.py"
|
||||
fi
|
||||
|
||||
if [[ $commands == *" kernels/quantization"* ]]; then
|
||||
commands="${commands} \
|
||||
--ignore=kernels/quantization/test_int8_quant.py \
|
||||
--ignore=kernels/quantization/test_machete_mm.py \
|
||||
--ignore=kernels/quantization/test_block_fp8.py \
|
||||
--ignore=kernels/quantization/test_block_int8.py \
|
||||
--ignore=kernels/quantization/test_marlin_gemm.py \
|
||||
--ignore=kernels/quantization/test_cutlass_scaled_mm.py \
|
||||
--ignore=kernels/quantization/test_int8_kernel.py"
|
||||
fi
|
||||
|
||||
if [[ $commands == *" kernels/mamba"* ]]; then
|
||||
commands="${commands} \
|
||||
--ignore=kernels/mamba/test_mamba_mixer2.py \
|
||||
--ignore=kernels/mamba/test_causal_conv1d.py \
|
||||
--ignore=kernels/mamba/test_mamba_ssm_ssd.py"
|
||||
fi
|
||||
|
||||
if [[ $commands == *" kernels/moe"* ]]; then
|
||||
commands="${commands} \
|
||||
--ignore=kernels/moe/test_moe.py \
|
||||
--ignore=kernels/moe/test_cutlass_moe.py \
|
||||
--ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
|
||||
fi
|
||||
|
||||
#ignore certain Entrypoints/openai tests
|
||||
if [[ $commands == *" entrypoints/openai "* ]]; then
|
||||
commands=${commands//" entrypoints/openai "/" entrypoints/openai \
|
||||
--ignore=entrypoints/openai/test_audio.py \
|
||||
--ignore=entrypoints/openai/test_shutdown.py \
|
||||
--ignore=entrypoints/openai/test_completion.py \
|
||||
--ignore=entrypoints/openai/test_models.py \
|
||||
--ignore=entrypoints/openai/test_lora_adapters.py \
|
||||
--ignore=entrypoints/openai/test_return_tokens_as_ids.py \
|
||||
--ignore=entrypoints/openai/test_root_path.py \
|
||||
--ignore=entrypoints/openai/test_tokenization.py \
|
||||
--ignore=entrypoints/openai/test_prompt_validation.py "}
|
||||
fi
|
||||
|
||||
#ignore certain Entrypoints/llm tests
|
||||
if [[ $commands == *" entrypoints/llm "* ]]; then
|
||||
commands=${commands//" entrypoints/llm "/" entrypoints/llm \
|
||||
--ignore=entrypoints/llm/test_chat.py \
|
||||
--ignore=entrypoints/llm/test_accuracy.py \
|
||||
--ignore=entrypoints/llm/test_init.py \
|
||||
--ignore=entrypoints/llm/test_prompt_validation.py "}
|
||||
fi
|
||||
|
||||
commands=$(echo "$commands" | sed 's/ \\ / /g')
|
||||
# Fix quoting before ROCm overrides (so overrides see correct structure)
|
||||
commands=$(re_quote_pytest_markers "$commands")
|
||||
commands=$(apply_rocm_test_overrides "$commands")
|
||||
echo "Final commands: $commands"
|
||||
|
||||
# --ignore=entrypoints/openai/test_encoder_decoder.py \
|
||||
# --ignore=entrypoints/openai/test_embedding.py \
|
||||
# --ignore=entrypoints/openai/test_oot_registration.py
|
||||
# --ignore=entrypoints/openai/test_accuracy.py \
|
||||
# --ignore=entrypoints/openai/test_models.py <= Fails on MI250 but passes on MI300 as of 2025-03-13
|
||||
|
||||
|
||||
MYPYTHONPATH=".."
|
||||
|
||||
# Test that we're launching on the machine that has
|
||||
# proper access to GPUs
|
||||
# Verify GPU access
|
||||
render_gid=$(getent group render | cut -d: -f3)
|
||||
if [[ -z "$render_gid" ]]; then
|
||||
echo "Error: 'render' group not found. This is required for GPU access." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ $commands == *"VLLM_TEST_GROUP_NAME=mi325_4-2-node-tests-4-gpus-in-total"* ]]; then
|
||||
|
||||
# --- Route: multi-node vs single-node ---
|
||||
if is_multi_node "$commands"; then
|
||||
echo "--- Multi-node job detected"
|
||||
export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
|
||||
|
||||
if [[ "$commands" =~ ^(.*)"["(.*)"] && ["(.*)"]"$ ]]; then
|
||||
prefix=$( echo "${BASH_REMATCH[1]}" | sed 's/;//g')
|
||||
echo "PREFIX: ${prefix}"
|
||||
export composite_command="(command rocm-smi || true)"
|
||||
myIFS=$IFS
|
||||
IFS=','
|
||||
read -ra node0 <<< ${BASH_REMATCH[2]}
|
||||
read -ra node1 <<< ${BASH_REMATCH[3]}
|
||||
IFS=$myIFS
|
||||
for i in "${!node0[@]}";do
|
||||
command_node_0=$(echo ${node0[i]} | sed 's/\"//g')
|
||||
command_node_1=$(echo ${node1[i]} | sed 's/\"//g')
|
||||
|
||||
export commands="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
|
||||
echo "COMMANDS: ${commands}"
|
||||
composite_command=$(echo "${composite_command} && ${commands}")
|
||||
done
|
||||
/bin/bash -c "${composite_command}"
|
||||
cleanup_network
|
||||
# Parse the bracket syntax: prefix ; [node0_cmds] && [node1_cmds]
|
||||
# BASH_REMATCH[1] = prefix (everything before first bracket)
|
||||
# BASH_REMATCH[2] = comma-separated node0 commands
|
||||
# BASH_REMATCH[3] = comma-separated node1 commands
|
||||
if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then
|
||||
prefix=$(echo "${BASH_REMATCH[1]}" | sed 's/;//g')
|
||||
echo "PREFIX: ${prefix}"
|
||||
|
||||
export composite_command="(command rocm-smi || true)"
|
||||
saved_IFS=$IFS
|
||||
IFS=','
|
||||
read -ra node0 <<< "${BASH_REMATCH[2]}"
|
||||
read -ra node1 <<< "${BASH_REMATCH[3]}"
|
||||
IFS=$saved_IFS
|
||||
|
||||
if [[ ${#node0[@]} -ne ${#node1[@]} ]]; then
|
||||
echo "Warning: node0 has ${#node0[@]} commands, node1 has ${#node1[@]}. They will be paired by index."
|
||||
fi
|
||||
|
||||
for i in "${!node0[@]}"; do
|
||||
command_node_0=$(echo "${node0[i]}" | sed 's/\"//g')
|
||||
command_node_1=$(echo "${node1[i]}" | sed 's/\"//g')
|
||||
|
||||
step_cmd="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
|
||||
echo "COMMANDS: ${step_cmd}"
|
||||
composite_command="${composite_command} && ${step_cmd}"
|
||||
done
|
||||
|
||||
/bin/bash -c "${composite_command}"
|
||||
cleanup_network
|
||||
else
|
||||
echo "Failed to parse node commands! Exiting."
|
||||
cleanup_network
|
||||
exit 111
|
||||
echo "Multi-node job detected but failed to parse bracket command syntax."
|
||||
echo "Expected format: prefix ; [node0_cmd1, node0_cmd2] && [node1_cmd1, node1_cmd2]"
|
||||
echo "Got: $commands"
|
||||
cleanup_network
|
||||
exit 111
|
||||
fi
|
||||
else
|
||||
echo "--- Single-node job"
|
||||
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
|
||||
docker run \
|
||||
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
|
||||
--network=host \
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
-e HF_TOKEN \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${commands}"
|
||||
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
|
||||
--network=host \
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
-e HF_TOKEN \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${commands}"
|
||||
fi
|
||||
|
||||
+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=
|
||||
@@ -145,7 +145,7 @@ steps:
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- cd examples/offline_inference/new_weight_syncing && VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -28,3 +28,11 @@ steps:
|
||||
- pytest -v -s v1/engine/test_preprocess_error_handling.py
|
||||
# Run the rest of v1/engine tests
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- pytest -v -s v1/e2e
|
||||
- pytest -v -s v1/engine
|
||||
|
||||
@@ -65,6 +65,11 @@ steps:
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
timeout_in_minutes: 50
|
||||
|
||||
@@ -115,6 +115,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
|
||||
# e2e
|
||||
- pytest -v -s tests/models/quantization/test_nvfp4.py
|
||||
@@ -156,14 +157,3 @@ steps:
|
||||
- pytest -v -s kernels/moe/test_deepep_moe.py
|
||||
- pytest -v -s kernels/moe/test_pplx_cutlass_moe.py
|
||||
# - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main
|
||||
|
||||
- label: Kernels Fp4 MoE Test (B200)
|
||||
timeout_in_minutes: 60
|
||||
device: b200
|
||||
num_devices: 1
|
||||
optional: true
|
||||
commands:
|
||||
- pytest -v -s kernels/moe/test_cutedsl_moe.py
|
||||
- pytest -v -s kernels/moe/test_flashinfer_moe.py
|
||||
- pytest -v -s kernels/moe/test_nvfp4_moe.py
|
||||
- pytest -v -s kernels/moe/test_ocp_mx_moe.py
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -16,6 +16,7 @@ steps:
|
||||
- pytest -v -s v1/sample
|
||||
- pytest -v -s v1/logits_processors
|
||||
- pytest -v -s v1/worker
|
||||
# TODO: create another `optional` test group for slow tests
|
||||
- pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
@@ -25,6 +26,11 @@ steps:
|
||||
# Integration test for streaming correctness (requires special branch).
|
||||
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
|
||||
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Others (CPU)
|
||||
depends_on:
|
||||
@@ -147,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
|
||||
|
||||
@@ -55,6 +55,15 @@ steps:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
|
||||
- label: Language Models Test (PPL)
|
||||
timeout_in_minutes: 110
|
||||
@@ -73,6 +82,11 @@ steps:
|
||||
- tests/models/language/pooling
|
||||
commands:
|
||||
- pytest -v -s models/language/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
timeout_in_minutes: 110
|
||||
|
||||
+4
-4
@@ -5,14 +5,14 @@
|
||||
/vllm/compilation @zou3519 @youkaichao @ProExpertProg
|
||||
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
|
||||
/vllm/lora @jeejeelee
|
||||
/vllm/model_executor/layers/attention @LucasWilkinson
|
||||
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
|
||||
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
|
||||
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
|
||||
/vllm/model_executor/layers/mamba @tdoublep
|
||||
/vllm/model_executor/model_loader @22quinn
|
||||
/vllm/model_executor/layers/batch_invariant.py @yewentao256
|
||||
/vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
|
||||
/vllm/vllm_flash_attn @LucasWilkinson
|
||||
/vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
|
||||
CMakeLists.txt @tlrmchlsmth @LucasWilkinson
|
||||
|
||||
# Any change to the VllmConfig changes can have a large user-facing impact,
|
||||
@@ -43,14 +43,14 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
|
||||
/vllm/tool_parsers @aarnphm @chaunceyjiang
|
||||
|
||||
# vLLM V1
|
||||
/vllm/v1/attention @LucasWilkinson
|
||||
/vllm/v1/attention @LucasWilkinson @MatthewBonanni
|
||||
/vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
|
||||
/vllm/v1/attention/backends/mla @pavanimajety
|
||||
/vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety
|
||||
/vllm/v1/attention/backends/triton_attn.py @tdoublep
|
||||
/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
|
||||
/vllm/v1/sample @22quinn @houseroad @njhill
|
||||
/vllm/v1/spec_decode @benchislett @luccafong
|
||||
/vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
|
||||
/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett
|
||||
/vllm/v1/kv_cache_interface.py @heheda12345
|
||||
/vllm/v1/kv_offload @ApostaC @orozery
|
||||
|
||||
@@ -783,7 +783,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
SRCS "${DSV3_FUSED_A_GEMM_SRC}"
|
||||
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
|
||||
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -16,6 +16,8 @@ torch::Tensor get_scheduler_metadata(
|
||||
isa = cpu_attention::ISA::VEC16;
|
||||
} else if (isa_hint == "neon") {
|
||||
isa = cpu_attention::ISA::NEON;
|
||||
} else if (isa_hint == "vxe") {
|
||||
isa = cpu_attention::ISA::VXE;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
|
||||
}
|
||||
@@ -100,6 +102,8 @@ void cpu_attn_reshape_and_cache(
|
||||
return cpu_attention::ISA::VEC16;
|
||||
} else if (isa == "neon") {
|
||||
return cpu_attention::ISA::NEON;
|
||||
} else if (isa == "vxe") {
|
||||
return cpu_attention::ISA::VXE;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid ISA type: " + isa);
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
enum class ISA { AMX, VEC, VEC16, NEON };
|
||||
enum class ISA { AMX, VEC, VEC16, NEON, VXE };
|
||||
|
||||
template <ISA isa, typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl {};
|
||||
|
||||
@@ -0,0 +1,386 @@
|
||||
#ifndef CPU_ATTN_VXE_HPP
|
||||
#define CPU_ATTN_VXE_HPP
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
#include <vecintrin.h>
|
||||
#include <type_traits>
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// s390x Vector = 16 bytes (128 bits)
|
||||
#define BLOCK_SIZE_ALIGNMENT 32
|
||||
#define HEAD_SIZE_ALIGNMENT 32
|
||||
#define MAX_Q_HEAD_NUM_PER_ITER 16
|
||||
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
|
||||
__vector float& b1);
|
||||
|
||||
// [1] Float Specialization
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
|
||||
__vector float& b1) {
|
||||
// Explicitly cast to long long for offset, and float* for pointer
|
||||
b0 = vec_xl((long long)0, const_cast<float*>(p));
|
||||
b1 = vec_xl((long long)0, const_cast<float*>(p + 4));
|
||||
}
|
||||
|
||||
// [2] BFloat16 Specialization (Big Endian Fix)
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
// 1. Load 8 BF16s (16 bytes) into one vector
|
||||
// Explicit cast to unsigned short* for vec_xl to return vector unsigned short
|
||||
__vector unsigned short raw = vec_xl((long long)0, (unsigned short*)p);
|
||||
|
||||
// 2. Prepare Zero vector
|
||||
__vector unsigned short zeros = vec_splat_u16(0);
|
||||
|
||||
// 3. Merge High/Low to expand BF16 -> Float32
|
||||
// On Big Endian, a float is [BF16_bits | 16_zero_bits]
|
||||
b0 = (__vector float)vec_mergeh(raw, zeros);
|
||||
b1 = (__vector float)vec_mergel(raw, zeros);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
alignas(16) float tmp[8];
|
||||
|
||||
// Manual unroll / conversion
|
||||
tmp[0] = static_cast<float>(p[0]);
|
||||
tmp[1] = static_cast<float>(p[1]);
|
||||
tmp[2] = static_cast<float>(p[2]);
|
||||
tmp[3] = static_cast<float>(p[3]);
|
||||
tmp[4] = static_cast<float>(p[4]);
|
||||
tmp[5] = static_cast<float>(p[5]);
|
||||
tmp[6] = static_cast<float>(p[6]);
|
||||
tmp[7] = static_cast<float>(p[7]);
|
||||
|
||||
// Explicit arguments for intrinsic: (long long offset, float* ptr)
|
||||
b0 = vec_xl((long long)0, (float*)tmp);
|
||||
b1 = vec_xl((long long)0, (float*)(tmp + 4));
|
||||
}
|
||||
|
||||
template <int32_t M, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_micro_s390x_Mx8_Ku4(
|
||||
const float* __restrict A, // [M x K]
|
||||
const kv_cache_t* __restrict B, // [K x 8]
|
||||
float* __restrict C, // [M x 8]
|
||||
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
|
||||
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
|
||||
|
||||
// Helper macros to unroll codegen for M rows
|
||||
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define IF_M(i) if constexpr (M > (i))
|
||||
|
||||
// 1. Define A pointers
|
||||
#define DECL_A(i) const float* a##i = A + (i) * lda;
|
||||
ROWS_APPLY(DECL_A)
|
||||
#undef DECL_A
|
||||
|
||||
// 2. Define Accumulators (2 vectors covers 8 columns)
|
||||
#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
|
||||
ROWS_APPLY(DECL_ACC)
|
||||
#undef DECL_ACC
|
||||
|
||||
// 3. Initialize Accumulators (Load C or Zero)
|
||||
#define INIT_ACC(i) \
|
||||
IF_M(i) { \
|
||||
if (accumulate) { \
|
||||
acc##i##_0 = \
|
||||
vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 0)); \
|
||||
acc##i##_1 = \
|
||||
vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 4)); \
|
||||
} else { \
|
||||
acc##i##_0 = vec_splats(0.0f); \
|
||||
acc##i##_1 = vec_splats(0.0f); \
|
||||
} \
|
||||
}
|
||||
ROWS_APPLY(INIT_ACC)
|
||||
#undef INIT_ACC
|
||||
|
||||
int32_t k = 0;
|
||||
|
||||
for (; k + 3 < K; k += 4) {
|
||||
// Load 4 values of A for each Row M: A[k...k+3]
|
||||
#define LOAD_A4(i) \
|
||||
__vector float a##i##v; \
|
||||
IF_M(i) a##i##v = vec_xl((long long)0, const_cast<float*>(a##i + k));
|
||||
ROWS_APPLY(LOAD_A4)
|
||||
#undef LOAD_A4
|
||||
|
||||
// Helper: FMA for specific lane L of A
|
||||
// s390x: vec_madd(b, vec_splat(a, lane), acc)
|
||||
#define FMAS_LANE(i, aiv, L) \
|
||||
IF_M(i) { \
|
||||
__vector float a_broad = vec_splat(aiv, L); \
|
||||
acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
|
||||
}
|
||||
|
||||
// Unroll K=0..3
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
|
||||
#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
|
||||
ROWS_APPLY(STEP_K0)
|
||||
#undef STEP_K0
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
|
||||
#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
|
||||
ROWS_APPLY(STEP_K1)
|
||||
#undef STEP_K1
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
|
||||
#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
|
||||
ROWS_APPLY(STEP_K2)
|
||||
#undef STEP_K2
|
||||
}
|
||||
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
|
||||
#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
|
||||
ROWS_APPLY(STEP_K3)
|
||||
#undef STEP_K3
|
||||
}
|
||||
#undef FMAS_LANE
|
||||
}
|
||||
|
||||
for (; k < K; ++k) {
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
|
||||
#define TAIL_ROW(i) \
|
||||
IF_M(i) { \
|
||||
__vector float ai = vec_splats(*(a##i + k)); \
|
||||
acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
|
||||
}
|
||||
ROWS_APPLY(TAIL_ROW)
|
||||
#undef TAIL_ROW
|
||||
}
|
||||
|
||||
#define STORE_ROW(i) \
|
||||
IF_M(i) { \
|
||||
vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
|
||||
vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
|
||||
}
|
||||
ROWS_APPLY(STORE_ROW)
|
||||
#undef STORE_ROW
|
||||
|
||||
#undef ROWS_APPLY
|
||||
#undef IF_M
|
||||
}
|
||||
|
||||
template <int32_t N, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_macro_s390x_Mx8_Ku4(const float* __restrict A,
|
||||
const kv_cache_t* __restrict B,
|
||||
float* __restrict C, int32_t M,
|
||||
int32_t K, int64_t lda, int64_t ldb,
|
||||
int64_t ldc, bool accumulate) {
|
||||
static_assert(N % 8 == 0, "N must be a multiple of 8");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
|
||||
const float* Ab = A + m * lda;
|
||||
float* Cb = C + m * ldc;
|
||||
|
||||
for (int32_t n = 0; n < N; n += 8) {
|
||||
const kv_cache_t* Bn = B + n;
|
||||
float* Cn = Cb + n;
|
||||
switch (mb) {
|
||||
case 8:
|
||||
gemm_micro_s390x_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_s390x_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_s390x_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
default:
|
||||
gemm_micro_s390x_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
}
|
||||
}
|
||||
m += mb;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename kv_cache_t>
|
||||
class TileGemmS390X {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
gemm_macro_s390x_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_s390x_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VXE;
|
||||
constexpr static bool scale_on_logits =
|
||||
false; // Scale is applied to Q during copy
|
||||
|
||||
public:
|
||||
AttentionImpl() {}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmS390X<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// Strides for Memory Layout
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // [head_dim, block_size] layout
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
__vector float scale_vec = vec_splats(scale);
|
||||
constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
|
||||
|
||||
// Process 8 elements at a time (32 bytes of float output)
|
||||
for (int32_t i = 0; i < q_num; ++i) {
|
||||
for (int32_t h = 0; h < q_heads_per_kv; ++h) {
|
||||
scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
|
||||
float* curr_dst =
|
||||
q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
|
||||
|
||||
int32_t d = 0;
|
||||
for (; d <= head_dim - 8; d += 8) {
|
||||
if constexpr (is_bf16) {
|
||||
__vector float v0, v1;
|
||||
// Reuse our Big-Endian-Safe loader
|
||||
load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
|
||||
|
||||
v0 = vec_mul(v0, scale_vec);
|
||||
v1 = vec_mul(v1, scale_vec);
|
||||
|
||||
vec_xst(v0, 0, curr_dst + d);
|
||||
vec_xst(v1, 0, curr_dst + d + 4);
|
||||
} else {
|
||||
__vector float v0 = vec_xl((long long)0, (float*)curr_src + d);
|
||||
__vector float v1 = vec_xl((long long)0, (float*)curr_src + d + 4);
|
||||
|
||||
v0 = vec_mul(v0, scale_vec);
|
||||
v1 = vec_mul(v1, scale_vec);
|
||||
|
||||
vec_xst(v0, 0, curr_dst + d);
|
||||
vec_xst(v1, 0, curr_dst + d + 4);
|
||||
}
|
||||
}
|
||||
|
||||
for (; d < head_dim; ++d) {
|
||||
float val = static_cast<float>(curr_src[d]);
|
||||
curr_dst[d] = val * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride) {
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) continue;
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
|
||||
{
|
||||
const scalar_t* key_src = key + token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
key_dst[j] = key_src[i];
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
const scalar_t* val_src = value + token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride +
|
||||
block_offset * head_dim;
|
||||
|
||||
std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif
|
||||
@@ -19,10 +19,11 @@ ISA_TYPES = {
|
||||
"VEC": 1,
|
||||
"VEC16": 2,
|
||||
"NEON": 3,
|
||||
"VXE": 4,
|
||||
}
|
||||
|
||||
# ISAs supported for head_dims divisible by 32
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16"]
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
|
||||
|
||||
# ISAs supported for head_dims divisible by 16 only
|
||||
ISA_FOR_16 = ["VEC16"]
|
||||
@@ -118,6 +119,10 @@ def generate_header_file() -> str:
|
||||
#include "cpu_attn_neon.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef __s390x__
|
||||
#include "cpu_attn_vxe.hpp"
|
||||
#endif
|
||||
|
||||
"""
|
||||
|
||||
header += generate_helper_function()
|
||||
@@ -163,6 +168,25 @@ def generate_header_file() -> str:
|
||||
} \\
|
||||
}()
|
||||
|
||||
"""
|
||||
|
||||
# s390x with VXE
|
||||
header += """#elif defined(__s390x__)
|
||||
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
|
||||
[&] { \\
|
||||
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
|
||||
switch (encoded_params) { \\
|
||||
"""
|
||||
header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
|
||||
header += """
|
||||
default: { \\
|
||||
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
|
||||
std::to_string(HEAD_DIM) + " isa=" + \\
|
||||
std::to_string(static_cast<int>(ISA_TYPE))); \\
|
||||
} \\
|
||||
} \\
|
||||
}()
|
||||
|
||||
"""
|
||||
|
||||
# Fallback: VEC and VEC16 only
|
||||
@@ -182,7 +206,7 @@ def generate_header_file() -> str:
|
||||
} \\
|
||||
}()
|
||||
|
||||
#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ */
|
||||
#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
|
||||
|
||||
#endif // CPU_ATTN_DISPATCH_GENERATED_H
|
||||
"""
|
||||
|
||||
@@ -745,3 +745,7 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
stream);
|
||||
}
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
|
||||
}
|
||||
|
||||
@@ -20,10 +20,12 @@
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "core/registration.h"
|
||||
#include "dsv3_router_gemm_utils.h"
|
||||
|
||||
static constexpr int DEFAULT_NUM_EXPERTS = 256;
|
||||
@@ -161,3 +163,7 @@ void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("dsv3_router_gemm", &dsv3_router_gemm);
|
||||
}
|
||||
|
||||
@@ -172,7 +172,7 @@ __device__ void _moe_align_block_size(
|
||||
}
|
||||
}
|
||||
|
||||
// Fill remaining expert_ids with 0
|
||||
// Fill remaining expert_ids with -1
|
||||
const size_t fill_start_idx =
|
||||
cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x;
|
||||
for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) {
|
||||
@@ -265,7 +265,7 @@ __device__ void _moe_align_block_size_small_batch_expert(
|
||||
}
|
||||
}
|
||||
|
||||
// Fill remaining expert_ids with 0
|
||||
// Fill remaining expert_ids with -1
|
||||
const size_t fill_start_idx = cumsum[num_experts] / block_size + tid;
|
||||
for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) {
|
||||
expert_ids[expert_ids_offset + i] = inactive_expert_id;
|
||||
@@ -332,7 +332,7 @@ __global__ void moe_align_block_size_kernel(
|
||||
topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
|
||||
num_experts, padded_num_experts, experts_per_warp, block_size, numel,
|
||||
cumsum, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size),
|
||||
0, 0, topk_num, nullptr, has_expert_map);
|
||||
0, -1, topk_num, nullptr, has_expert_map);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
@@ -373,7 +373,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
|
||||
_moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
|
||||
topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
|
||||
num_experts, block_size, numel, max_num_tokens_padded,
|
||||
CEILDIV(max_num_tokens_padded, block_size), 0, 0, topk_num, nullptr,
|
||||
CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
|
||||
has_expert_map);
|
||||
}
|
||||
|
||||
|
||||
@@ -127,7 +127,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
|
||||
// DeepSeek V3 optimized router GEMM for SM90+
|
||||
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
m.impl("dsv3_router_gemm", torch::kCUDA, &dsv3_router_gemm);
|
||||
// conditionally compiled so impl registration is in source file
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -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));
|
||||
});
|
||||
|
||||
@@ -12,6 +12,68 @@ namespace vllm {
|
||||
|
||||
using c3x::cutlass_gemm_caller;
|
||||
|
||||
// Custom wrapper to allow specifying EpilogueTile for small M
|
||||
template <typename ElementAB_, typename ElementD_,
|
||||
template <typename, typename, typename> typename Epilogue_,
|
||||
typename TileShape, typename ClusterShape, typename KernelSchedule,
|
||||
typename EpilogueSchedule, typename EpilogueTile>
|
||||
struct cutlass_3x_gemm_sm120_custom {
|
||||
using ElementAB = ElementAB_;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
static constexpr int AlignmentA =
|
||||
128 / cutlass::sizeof_bits<ElementAB>::value;
|
||||
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
static constexpr int AlignmentB =
|
||||
128 / cutlass::sizeof_bits<ElementAB>::value;
|
||||
|
||||
using ElementC = void;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static constexpr int AlignmentC =
|
||||
128 / cutlass::sizeof_bits<ElementD_>::value;
|
||||
|
||||
using ElementD = ElementD_;
|
||||
using LayoutD = cutlass::layout::RowMajor;
|
||||
static constexpr int AlignmentD = AlignmentC;
|
||||
|
||||
using ElementAcc =
|
||||
typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
|
||||
float>::type;
|
||||
using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
|
||||
|
||||
// MMA type
|
||||
using ElementAccumulator = float;
|
||||
|
||||
// Epilogue types
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ElementCompute = float;
|
||||
using ElementAux = ElementD;
|
||||
using LayoutAux = LayoutD;
|
||||
using ElementAmax = float;
|
||||
|
||||
using EVTCompute = typename Epilogue::EVTCompute;
|
||||
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, TileShape,
|
||||
ClusterShape, EpilogueTile, // Use custom EpilogueTile
|
||||
ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
|
||||
ElementD, LayoutD, AlignmentD, EpilogueSchedule,
|
||||
EVTCompute>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, ElementAB,
|
||||
LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
|
||||
ElementAccumulator, TileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
KernelSchedule, void>::CollectiveOp;
|
||||
|
||||
using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm120_fp8_config_default {
|
||||
@@ -25,6 +87,54 @@ struct sm120_fp8_config_default {
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm120_fp8_config_M64 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
// SM120 Cooperative kernel requires Tile M >= 128.
|
||||
// For M=64 tile, we use Pingpong schedule which is more flexible with small
|
||||
// tiles.
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
|
||||
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
|
||||
using TileShape = Shape<_64, _64, _128>;
|
||||
// CUTLASS 3.x on SM120 currently restricts programmatic multicast (Cluster >
|
||||
// 1) for certain schedules/types. Reverting to 1x1x1 to ensure compilation.
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_3x_gemm_sm120<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm120_fp8_config_M32 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
|
||||
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
|
||||
using TileShape = Shape<_32, _64, _128>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
// Use custom gemm to specify EpilogueTile M=32
|
||||
using Cutlass3xGemm =
|
||||
cutlass_3x_gemm_sm120_custom<InType, OutType, Epilogue, TileShape,
|
||||
ClusterShape, KernelSchedule,
|
||||
EpilogueSchedule, Shape<_32, _32>>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm120_fp8_config_M16 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
|
||||
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
|
||||
using TileShape = Shape<_16, _64, _128>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
// Use custom gemm to specify EpilogueTile M=16
|
||||
using Cutlass3xGemm =
|
||||
cutlass_3x_gemm_sm120_custom<InType, OutType, Epilogue, TileShape,
|
||||
ClusterShape, KernelSchedule,
|
||||
EpilogueSchedule, Shape<_16, _32>>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
@@ -36,6 +146,28 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::Tensor& out,
|
||||
TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
|
||||
TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
|
||||
|
||||
int M = a.size(0);
|
||||
|
||||
if (M <= 16) {
|
||||
using Cutlass3xGemmM16 =
|
||||
typename sm120_fp8_config_M16<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
return cutlass_gemm_caller<Cutlass3xGemmM16>(
|
||||
out, a, b, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
if (M <= 32) {
|
||||
using Cutlass3xGemmM32 =
|
||||
typename sm120_fp8_config_M32<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
return cutlass_gemm_caller<Cutlass3xGemmM32>(
|
||||
out, a, b, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
if (M <= 256) {
|
||||
using Cutlass3xGemmM64 =
|
||||
typename sm120_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
return cutlass_gemm_caller<Cutlass3xGemmM64>(
|
||||
out, a, b, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
using Cutlass3xGemmDefault =
|
||||
typename sm120_fp8_config_default<InType, OutType,
|
||||
Epilogue>::Cutlass3xGemm;
|
||||
@@ -64,4 +196,4 @@ void cutlass_scaled_mm_sm120_fp8_epilogue(torch::Tensor& out,
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
} // namespace vllm
|
||||
|
||||
+14
-46
@@ -1902,7 +1902,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
floatx16 sum[N][YTILE] = {};
|
||||
scalar8 sum[N][YTILE] = {};
|
||||
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
|
||||
bigType bigA[N][UNRL] = {};
|
||||
bigType bigB[YTILE][UNRL];
|
||||
@@ -1936,7 +1936,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
for (int i = 0; i < A_CHUNK; i += 8) {
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
sum[n][y] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
|
||||
sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
|
||||
bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
|
||||
0);
|
||||
}
|
||||
@@ -1949,31 +1949,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
float accm0 = sum[n][y][0];
|
||||
float accm16 = sum[n][y][8];
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
|
||||
1); // row_shl1
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
|
||||
1); // row_shl2
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
|
||||
1); // row_shl3
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
|
||||
1); // row_shl8
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
|
||||
1); // row_shl9
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
|
||||
1); // row_shl10
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
|
||||
1); // row_shl11
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
|
||||
accm0 += __shfl(accm0, 36);
|
||||
accm16 += __shfl(accm16, 52);
|
||||
sum[n][y][0] = accm0 + __shfl(accm16, 16);
|
||||
accm0 += __shfl_down(accm0, 20);
|
||||
accm0 += __shfl_down(accm0, 40);
|
||||
sum[n][y][0] = accm0;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2064,7 +2048,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
floatx16 sum[N][YTILE] = {};
|
||||
scalar8 sum[N][YTILE] = {};
|
||||
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
|
||||
bigType bigA[N][UNRL] = {};
|
||||
bigType bigB[YTILE][UNRL];
|
||||
@@ -2100,7 +2084,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
for (int i = 0; i < A_CHUNK; i += 8) {
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
sum[n][y] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
|
||||
sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
|
||||
bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
|
||||
0);
|
||||
}
|
||||
@@ -2113,31 +2097,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
float accm0 = sum[n][y][0];
|
||||
float accm16 = sum[n][y][8];
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
|
||||
1); // row_shl1
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
|
||||
1); // row_shl2
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
|
||||
1); // row_shl3
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
|
||||
1); // row_shl8
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
|
||||
1); // row_shl9
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
|
||||
1); // row_shl10
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
|
||||
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
|
||||
1); // row_shl11
|
||||
accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
|
||||
accm0 += __shfl(accm0, 36);
|
||||
accm16 += __shfl(accm16, 52);
|
||||
sum[n][y][0] = accm0 + __shfl(accm16, 16);
|
||||
accm0 += __shfl_down(accm0, 20);
|
||||
accm0 += __shfl_down(accm0, 40);
|
||||
sum[n][y][0] = accm0;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2242,16 +2210,16 @@ void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
|
||||
: nullptr;
|
||||
switch (N_in) {
|
||||
case 1:
|
||||
WVSPLITKQ(12, 2, 2, 2, 2, 1)
|
||||
WVSPLITKQ(16, 2, 2, 2, 2, 1)
|
||||
break;
|
||||
case 2:
|
||||
WVSPLITKQ(12, 2, 2, 2, 2, 2)
|
||||
WVSPLITKQ(16, 2, 2, 2, 2, 2)
|
||||
break;
|
||||
case 3:
|
||||
WVSPLITKQ(8, 2, 2, 1, 1, 3)
|
||||
WVSPLITKQ(16, 2, 2, 2, 2, 3)
|
||||
break;
|
||||
case 4:
|
||||
WVSPLITKQ(4, 2, 2, 1, 1, 4)
|
||||
WVSPLITKQ(16, 2, 2, 2, 2, 4)
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error(
|
||||
|
||||
@@ -242,7 +242,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
||||
ops.def(
|
||||
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
|
||||
// conditionally compiled so impl registration is in source file
|
||||
|
||||
// Quantized GEMM for AWQ.
|
||||
ops.def(
|
||||
|
||||
+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
|
||||
|
||||
@@ -305,6 +305,14 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
uv pip install --system /rixl_install/*.whl
|
||||
|
||||
# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries)
|
||||
RUN apt-get update -q -y && apt-get install -q -y \
|
||||
librdmacm1 \
|
||||
libibverbs1 \
|
||||
ibverbs-providers \
|
||||
ibverbs-utils \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
ARG COMMON_WORKDIR
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
|
||||
@@ -330,6 +338,11 @@ RUN bash /tmp/install_torchcodec.sh \
|
||||
# Copy in the v1 package (for python-only install test group)
|
||||
COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
|
||||
|
||||
# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel
|
||||
# See: https://github.com/pytorch/pytorch/issues/169857
|
||||
ENV MIOPEN_DEBUG_CONV_DIRECT=0
|
||||
ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
|
||||
# Source code is used in the `python_only_compile.sh` test
|
||||
# We hide it inside `src/` so that this source code
|
||||
# will not be imported by other tests
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -49,7 +49,13 @@ If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first:
|
||||
uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129
|
||||
```
|
||||
|
||||
then install vLLM using:
|
||||
Then install the necessary build dependencies from `requirements/build.txt`, skipping `torch` as it was installed in the previous step:
|
||||
|
||||
```bash
|
||||
grep -v '^torch==' requirements/build.txt | uv pip install -r -
|
||||
```
|
||||
|
||||
Finally install vLLM using:
|
||||
|
||||
```bash
|
||||
uv pip install -e . --no-build-isolation
|
||||
|
||||
@@ -208,9 +208,7 @@ configurations affect the class we ultimately get.
|
||||
|
||||
The following figure shows the class hierarchy of vLLM:
|
||||
|
||||
> <figure markdown="span">
|
||||
> { align="center" alt="query" width="100%" }
|
||||
> </figure>
|
||||

|
||||
|
||||
There are several important design choices behind this class hierarchy:
|
||||
|
||||
|
||||
@@ -32,10 +32,10 @@ th {
|
||||
|
||||
| Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
|
||||
|---------|--------------------|--------------|---------------|-------|-----------------------|-----------|
|
||||
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE |
|
||||
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
|
||||
| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
|
||||
| MoEPrepareAndFinalizeNoEP<sup>5</sup> | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] |
|
||||
| BatchedPrepareAndFinalize<sup>5</sup> | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] |
|
||||
|
||||
@@ -295,6 +295,51 @@ You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the mult
|
||||
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
|
||||
#### Chunking Long Audio for Transcription
|
||||
|
||||
Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech.
|
||||
|
||||
```python
|
||||
import librosa
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.multimodal.audio import split_audio
|
||||
|
||||
# Load long audio file
|
||||
audio, sr = librosa.load("long_audio.wav", sr=16000)
|
||||
|
||||
# Split into chunks at low-energy (quiet) regions
|
||||
chunks = split_audio(
|
||||
audio_data=audio,
|
||||
sample_rate=sr,
|
||||
max_clip_duration_s=30.0, # Maximum chunk length in seconds
|
||||
overlap_duration_s=1.0, # Search window for finding quiet split points
|
||||
min_energy_window_size=1600, # Window size for energy calculation (~100ms at 16kHz)
|
||||
)
|
||||
|
||||
# Initialize Whisper model
|
||||
llm = LLM(model="openai/whisper-large-v3-turbo")
|
||||
sampling_params = SamplingParams(temperature=0, max_tokens=256)
|
||||
|
||||
# Transcribe each chunk
|
||||
transcriptions = []
|
||||
for chunk in chunks:
|
||||
outputs = llm.generate({
|
||||
"prompt": "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
|
||||
"multi_modal_data": {"audio": (chunk, sr)},
|
||||
}, sampling_params)
|
||||
transcriptions.append(outputs[0].outputs[0].text)
|
||||
|
||||
# Combine results
|
||||
full_transcription = " ".join(transcriptions)
|
||||
```
|
||||
|
||||
The `split_audio` function:
|
||||
|
||||
- Splits audio at quiet points to avoid cutting through speech
|
||||
- Uses RMS energy to find low-amplitude regions within the overlap window
|
||||
- Preserves all audio samples (no data loss)
|
||||
- Supports any sample rate
|
||||
|
||||
#### Automatic Audio Channel Normalization
|
||||
|
||||
vLLM automatically normalizes audio channels for models that require specific audio formats. When loading audio with libraries like `torchaudio`, stereo files return shape `[channels, time]`, but many audio models (particularly Whisper-based models) expect mono audio with shape `[time]`.
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
It's recommended to use [uv](https://docs.astral.sh/uv/), a very fast Python environment manager, to create and manage Python environments. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
|
||||
|
||||
```bash
|
||||
uv venv --python 3.12 --seed
|
||||
uv venv --python 3.12 --seed --managed-python
|
||||
source .venv/bin/activate
|
||||
```
|
||||
|
||||
@@ -79,13 +79,15 @@ Specially, committers are almost all area owners. They author subsystems, review
|
||||
|
||||
For a full list of committers and their respective areas, see the [committers](./committers.md) page.
|
||||
|
||||
#### Nomination Process
|
||||
#### Committer Proposal Process
|
||||
|
||||
Any committer can nominate candidates via our private mailing list:
|
||||
Any committer can nominate candidates via our private committer mailing list. The process runs as follows:
|
||||
|
||||
1. **Nominate**: Any committer may nominate a candidate by email to the private maintainers’ list, citing evidence mapped to the pre‑existing standards with links to PRs, reviews, RFCs, issues, benchmarks, and adoption evidence.
|
||||
2. **Vote**: The lead maintainers will group voices support or concerns. Shared concerns can stop the process. The vote typically last 3 working days. For concerns, committers group discuss the clear criteria for such person to be nominated again. The lead maintainers will make the final decision.
|
||||
3. **Confirm**: The lead maintainers send invitation, update CODEOWNERS, assign permissions, add to communications channels (mailing list and Slack).
|
||||
1. **Nominate**: A committer sends email to the committer group to nominate a candidate, highlighting the candidate’s contributions (e.g., links to PRs, reviews, RFCs, issues, benchmarks, and adoption evidence) and how they map to the standards below.
|
||||
2. **Discuss and vote**: The committer group discusses the nomination, votes, and voices concerns if needed. Shared concerns can stop the process. For concerns, the group discusses clear criteria for the person to be nominated again. Most cases are decided by consensus; in contentious cases, the lead maintainers resolve conflicts and make the decision.
|
||||
3. **Feedback period**: After a two-week feedback period (allowing time for any last input or concerns), if no blocking concerns arise and the nominator confirms with lead maintainer group to move forward (via the mailing list or committers slack channel), the nominator sends an invitation to the candidate asking them to open a PR to update their code ownership (e.g., CODEOWNERS and committers list).
|
||||
4. **Permissions and onboarding**: In parallel, the lead maintainers assign the necessary permissions in GitHub and add the new member to the committer mailing list, the committer-only Slack channel, and other communications channels as appropriate.
|
||||
5. **Finalize**: Once the CODEOWNERS/committer PR is ready and permissions are in place, the PR is merged and the new committer is welcomed.
|
||||
|
||||
Committership is highly selective and merit based. The selection criteria requires:
|
||||
|
||||
|
||||
@@ -682,7 +682,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM`<sup>^</sup> | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
|
||||
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
|
||||
| `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + I<sup>E+</sup> | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -762,10 +762,8 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
|--------------|--------|--------|-------------------|-----------------------------|-----------------------------------------|
|
||||
| `Emu3ForConditionalGeneration` | Emu3 | T + I | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ |
|
||||
|
||||
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.
|
||||
• For example, to use DeepSeek-VL2 series models:
|
||||
`--hf-overrides '{"architectures": ["DeepseekVLV2ForCausalLM"]}'`
|
||||
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.
|
||||
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
|
||||
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
|
||||
<sup>+</sup> Multiple items can be inputted per text prompt for this modality.
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -5,6 +5,7 @@ from collections.abc import Callable
|
||||
import pytest
|
||||
|
||||
from vllm.config import PassConfig
|
||||
from vllm.utils.flashinfer import is_flashinfer_fp8_blockscale_gemm_supported
|
||||
|
||||
from .common import (
|
||||
INDUCTOR_GRAPH_PARTITION,
|
||||
@@ -50,6 +51,10 @@ def test_tp1_fp8_fusions(
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
if use_deepgemm and is_flashinfer_fp8_blockscale_gemm_supported():
|
||||
# Flashinfer block FP8 GEMM has internal quantization, so it can't
|
||||
# be fused with other ops.
|
||||
pytest.skip("FlashInfer block FP8 GEMM not supported")
|
||||
if use_deepgemm and is_blackwell():
|
||||
# TODO(luka) DeepGEMM uses different quants, matching not supported
|
||||
# - on Blackwell, uses a special quant fp8, currently not supported
|
||||
|
||||
@@ -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=}"
|
||||
)
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import copy
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.compile.backend import TestBackend
|
||||
from tests.utils import TestFP8Layer
|
||||
from vllm.compilation.passes.fusion.act_quant_fusion import (
|
||||
@@ -31,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
)
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
TEST_FP8 = current_platform.supports_fp8()
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
@@ -198,23 +200,82 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
|
||||
return [torch.ops.aten.slice_scatter.default]
|
||||
|
||||
|
||||
MODELS = [
|
||||
TestSiluMul,
|
||||
TestFusedAddRMSNorm,
|
||||
TestRotaryEmbedding,
|
||||
TestRotaryEmbeddingSliceScatter,
|
||||
]
|
||||
class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
|
||||
OP_REGISTERED = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.register_test_custom_op()
|
||||
|
||||
@classmethod
|
||||
def register_test_custom_op(cls):
|
||||
if not cls.OP_REGISTERED:
|
||||
|
||||
def function_with_mutated_args_and_return_impl(
|
||||
x: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
ret = x + 1
|
||||
x.add_(2)
|
||||
return ret
|
||||
|
||||
def function_with_mutated_args_and_return_fake(
|
||||
x: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty_like(x)
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="function_with_mutated_args_and_return",
|
||||
op_func=function_with_mutated_args_and_return_impl,
|
||||
mutates_args=["x"],
|
||||
fake_impl=function_with_mutated_args_and_return_fake,
|
||||
)
|
||||
|
||||
cls.OP_REGISTERED = True
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# Clone x to avoid mutating the original tensor
|
||||
ret = torch.ops.vllm.function_with_mutated_args_and_return(x)
|
||||
return x, ret
|
||||
|
||||
def example_inputs(self, num_tokens=32):
|
||||
hidden_states = torch.randn(num_tokens)
|
||||
return (hidden_states,)
|
||||
|
||||
def ops_in_model(self, do_fusion):
|
||||
return [torch.ops.vllm.function_with_mutated_args_and_return.default]
|
||||
|
||||
def ops_not_in_model(self):
|
||||
return []
|
||||
|
||||
|
||||
MODELS_AND_DO_FUSION = {
|
||||
TestSiluMul: [True, False],
|
||||
TestFusedAddRMSNorm: [True, False],
|
||||
TestRotaryEmbedding: [False],
|
||||
TestRotaryEmbeddingSliceScatter: [False],
|
||||
TestFunctionWithMutatedArgsAndReturn: [False],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("model_class", MODELS)
|
||||
@pytest.mark.parametrize("do_fusion", [True, False])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
|
||||
@pytest.mark.parametrize(
|
||||
"model_class, do_fusion",
|
||||
[
|
||||
(model_class, do_fusion)
|
||||
for model_class, fusions in MODELS_AND_DO_FUSION.items()
|
||||
for do_fusion in fusions
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="Only test on cuda and rocm platform",
|
||||
)
|
||||
def test_fix_functionalization(
|
||||
model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype
|
||||
):
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(0)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=ModelConfig(dtype=dtype),
|
||||
@@ -246,8 +307,14 @@ def test_fix_functionalization(
|
||||
backend_no_func = TestBackend(*passes)
|
||||
|
||||
model = model_class()
|
||||
torch.compile(model, backend=backend_func)(*model.example_inputs())
|
||||
torch.compile(model, backend=backend_no_func)(*model.example_inputs())
|
||||
inputs_func = model.example_inputs()
|
||||
inputs_no_func = copy.deepcopy(inputs_func)
|
||||
model_func = model_class()
|
||||
model_no_func = copy.deepcopy(model_func)
|
||||
model_func = torch.compile(model_func, backend=backend_func)
|
||||
model_no_func = torch.compile(model_no_func, backend=backend_no_func)
|
||||
model_func(*inputs_func)
|
||||
model_no_func(*inputs_no_func)
|
||||
|
||||
# check if the functionalization pass is applied
|
||||
for op in model.ops_in_model(do_fusion):
|
||||
@@ -265,3 +332,8 @@ def test_fix_functionalization(
|
||||
found[op] = True
|
||||
assert all(found[op] for op in model.ops_in_model(do_fusion))
|
||||
assert all(not found.get(op) for op in model.ops_not_in_model())
|
||||
|
||||
# TODO (Rohan138): compare the outputs from model_func and model_no_func
|
||||
# currently runs into errors while comparing `TestFusedAddRMSNorm`
|
||||
# Linked issue: https://github.com/vllm-project/vllm/issues/34996
|
||||
# torch.testing.assert_close(outputs_func, outputs_no_func)
|
||||
|
||||
@@ -26,24 +26,16 @@ from vllm.config import (
|
||||
PassConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
ChannelWiseTorchFP8ScaledMMLinearKernel,
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
FP8ScaledMMLinearKernel,
|
||||
PerTensorTorchFP8ScaledMMLinearKernel,
|
||||
ROCmFP8ScaledMMLinearKernel,
|
||||
RowWiseTorchFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
|
||||
ROCmFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
|
||||
FP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
QuantKey,
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.config
|
||||
from tests.compile.backend import TestBackend
|
||||
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
|
||||
from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
|
||||
from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
|
||||
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
|
||||
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
|
||||
from vllm.compilation.passes.utility.scatter_split_replace import (
|
||||
ScatterSplitReplacementPass,
|
||||
)
|
||||
from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
CompilationConfig,
|
||||
CompilationMode,
|
||||
ModelConfig,
|
||||
PassConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.forward_context import get_forward_context, set_forward_context
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
CommonAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec
|
||||
|
||||
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
|
||||
VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
|
||||
class QKRoPEKVCacheTestModel(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
attn_backend: AttentionBackendEnum,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
prefix: str = "model.layers.0.self_attn.attn",
|
||||
):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.head_size = head_size
|
||||
self.block_size = vllm_config.cache_config.block_size
|
||||
self.q_size = num_heads * head_size
|
||||
self.kv_size = num_kv_heads * head_size
|
||||
self.is_neox = is_neox
|
||||
self.dtype = dtype
|
||||
self.device = device
|
||||
self.layer_name = prefix
|
||||
|
||||
self.rotary_emb = RotaryEmbedding(
|
||||
head_size,
|
||||
rotary_dim=head_size,
|
||||
max_position_embeddings=4096,
|
||||
base=10000,
|
||||
is_neox_style=is_neox,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
|
||||
# Whether to check for the RoPE custom op or component index_select
|
||||
self.enable_rope_custom_op = self.rotary_emb.enabled()
|
||||
|
||||
# Register layer metadata for the fusion pass via Attention.
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=1.0 / head_size**0.5,
|
||||
num_kv_heads=num_kv_heads,
|
||||
cache_config=vllm_config.cache_config,
|
||||
quant_config=vllm_config.quant_config,
|
||||
prefix=prefix,
|
||||
attn_backend=attn_backend.get_class(),
|
||||
)
|
||||
self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
|
||||
assert not self.attn_backend.forward_includes_kv_cache_update, (
|
||||
f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache."
|
||||
)
|
||||
self.attn._k_scale = self.attn._k_scale.to(device)
|
||||
self.attn._v_scale = self.attn._v_scale.to(device)
|
||||
|
||||
kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
|
||||
self.kv_cache_dtype = (
|
||||
FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
|
||||
)
|
||||
|
||||
# Initialize attn MetadataBuilder
|
||||
self.builder = self.attn.attn_backend.get_builder_cls()(
|
||||
kv_cache_spec=AttentionSpec(
|
||||
block_size=self.block_size,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
head_size=head_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
),
|
||||
layer_names=[self.attn.layer_name],
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
|
||||
"""Initialize attention metadata."""
|
||||
# Create common attn metadata
|
||||
batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec, self.block_size, self.device, arange_block_indices=True
|
||||
)
|
||||
|
||||
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
|
||||
num_blocks = batch_size * max_blocks
|
||||
|
||||
# Fetch the attention backend and kv cache shape and stride order
|
||||
attn_backend = self.attn.attn_backend
|
||||
kv_cache_shape = attn_backend.get_kv_cache_shape(
|
||||
num_blocks, self.block_size, self.num_kv_heads, self.head_size
|
||||
)
|
||||
try:
|
||||
kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
|
||||
except (AttributeError, NotImplementedError):
|
||||
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
|
||||
|
||||
kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
|
||||
inv_order = [
|
||||
kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
|
||||
]
|
||||
|
||||
# Create dummy KV cache
|
||||
raw_tensor = torch.zeros(
|
||||
2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
raw_tensor = raw_tensor.view(kv_cache_shape)
|
||||
kv_cache = raw_tensor.permute(*inv_order)
|
||||
|
||||
self.attn.kv_cache = [kv_cache]
|
||||
|
||||
# Build attn metadata
|
||||
attn_metadata = self.builder.build(
|
||||
common_prefix_len=0, common_attn_metadata=common_attn_metadata
|
||||
)
|
||||
|
||||
return attn_metadata
|
||||
|
||||
def forward(
|
||||
self, qkv: torch.Tensor, positions: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# Create copy so inplace ops do not modify the original tensors
|
||||
qkv = qkv.clone()
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
|
||||
# Instead of a full forward pass, match only the KV cache update op here
|
||||
q = q.view(-1, self.num_heads, self.head_size)
|
||||
k = k.view(-1, self.num_kv_heads, self.head_size)
|
||||
v = v.view(-1, self.num_kv_heads, self.head_size)
|
||||
kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
|
||||
k, v, self.layer_name
|
||||
)
|
||||
return q, k, v, kv_cache_dummy_dep
|
||||
|
||||
def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
|
||||
ops = []
|
||||
if self.enable_rope_custom_op:
|
||||
if rocm_aiter_ops.is_triton_rotary_embed_enabled():
|
||||
ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default)
|
||||
else:
|
||||
ops.append(ROTARY_OP)
|
||||
else:
|
||||
ops.append(INDEX_SELECT_OP)
|
||||
ops.append(torch.ops.vllm.unified_kv_cache_update.default)
|
||||
return ops
|
||||
|
||||
def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
|
||||
return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"attn_backend",
|
||||
[
|
||||
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
AttentionBackendEnum.ROCM_ATTN,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("enable_rope_custom_op", [True]) # [True, False])
|
||||
@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
|
||||
@pytest.mark.parametrize("num_heads", [64])
|
||||
@pytest.mark.parametrize("num_kv_heads", [8])
|
||||
@pytest.mark.parametrize("head_size", [64])
|
||||
@pytest.mark.parametrize("block_size", [16])
|
||||
@pytest.mark.parametrize("is_neox", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
|
||||
@pytest.mark.skipif(
|
||||
not is_aiter_found_and_supported(),
|
||||
reason="Only test on ROCm with AITER installed and supported",
|
||||
)
|
||||
def test_rope_kvcache_fusion(
|
||||
attn_backend: AttentionBackendEnum,
|
||||
enable_rope_custom_op: bool,
|
||||
enable_aiter_triton_rope: bool,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(0)
|
||||
|
||||
custom_ops: list[str] = []
|
||||
if enable_rope_custom_op:
|
||||
custom_ops.append("+rotary_embedding")
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=ModelConfig(dtype=dtype),
|
||||
cache_config=CacheConfig(
|
||||
block_size=block_size,
|
||||
cache_dtype=kv_cache_dtype,
|
||||
),
|
||||
compilation_config=CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
custom_ops=custom_ops,
|
||||
pass_config=PassConfig(
|
||||
fuse_rope_kvcache=True,
|
||||
eliminate_noops=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
m.setenv(
|
||||
"VLLM_ROCM_USE_AITER_TRITON_ROPE", "1" if enable_aiter_triton_rope else "0"
|
||||
)
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
model = QKRoPEKVCacheTestModel(
|
||||
vllm_config=vllm_config,
|
||||
attn_backend=attn_backend,
|
||||
num_heads=num_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_size=head_size,
|
||||
is_neox=is_neox,
|
||||
dtype=dtype,
|
||||
device=torch.get_default_device(),
|
||||
)
|
||||
|
||||
fusion_pass = RopeKVCacheFusionPass(vllm_config)
|
||||
passes = [
|
||||
NoOpEliminationPass(vllm_config),
|
||||
SplitCoalescingPass(vllm_config),
|
||||
ScatterSplitReplacementPass(vllm_config),
|
||||
fusion_pass,
|
||||
PostCleanupPass(vllm_config),
|
||||
]
|
||||
backend = TestBackend(*passes)
|
||||
|
||||
T = 5
|
||||
|
||||
qkv = torch.randn(
|
||||
T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
|
||||
)
|
||||
pos = torch.arange(T, dtype=torch.long)
|
||||
|
||||
qkv_unfused = qkv.clone()
|
||||
pos_unfused = pos.clone()
|
||||
|
||||
with set_forward_context(None, vllm_config):
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = model.build_attn_metadata(T)
|
||||
forward_context.slot_mapping = {
|
||||
model.layer_name: attn_metadata.slot_mapping
|
||||
}
|
||||
q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
|
||||
attn_layer = forward_context.no_compile_layers[model.layer_name]
|
||||
kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
|
||||
del dummy
|
||||
|
||||
torch._dynamo.mark_dynamic(qkv, 0)
|
||||
torch._dynamo.mark_dynamic(pos, 0)
|
||||
with set_forward_context(None, vllm_config):
|
||||
model_fused = torch.compile(model, backend=backend)
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = model_fused.build_attn_metadata(T)
|
||||
forward_context.slot_mapping = {
|
||||
model.layer_name: attn_metadata.slot_mapping
|
||||
}
|
||||
q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
|
||||
attn_layer = forward_context.no_compile_layers[model.layer_name]
|
||||
kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
|
||||
del dummy
|
||||
|
||||
assert fusion_pass.matched_count == 1
|
||||
|
||||
backend.check_before_ops(model.ops_in_model_before())
|
||||
backend.check_after_ops(model.ops_in_model_after())
|
||||
|
||||
if dtype == torch.float16:
|
||||
ATOL, RTOL = (2e-3, 2e-3)
|
||||
else:
|
||||
ATOL, RTOL = (1e-2, 1e-2)
|
||||
|
||||
torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
|
||||
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
|
||||
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
|
||||
# Cannot compare fp8_* directly here, cast to model dtype instead
|
||||
torch.testing.assert_close(
|
||||
kv_cache_unfused.view(dtype),
|
||||
kv_cache_fused.view(dtype),
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
@@ -0,0 +1,107 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import vllm
|
||||
from tests.compile.backend import TestBackend
|
||||
from vllm.compilation.passes.utility.scatter_split_replace import (
|
||||
ScatterSplitReplacementPass,
|
||||
)
|
||||
from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
|
||||
from vllm.config import CompilationConfig, CompilationMode, VllmConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
|
||||
|
||||
class ScatterSplitReplacementModel(nn.Module):
|
||||
"""Model with a rope+getitem+slice_scatter+split_with_sizes sequence."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
super().__init__()
|
||||
self.q_size = num_heads * head_size
|
||||
self.kv_size = num_kv_heads * head_size
|
||||
|
||||
self.rotary_emb = RotaryEmbedding(
|
||||
head_size,
|
||||
rotary_dim=head_size,
|
||||
max_position_embeddings=4096,
|
||||
base=10000,
|
||||
is_neox_style=True,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
|
||||
# Create copy so inplace ops do not modify the original tensors
|
||||
qkv = qkv.clone()
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
q = q + 1
|
||||
k = k + 2
|
||||
v = v + 3
|
||||
return q, k, v
|
||||
|
||||
def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
|
||||
return [
|
||||
torch.ops.aten.slice_scatter.default,
|
||||
torch.ops.aten.split_with_sizes.default,
|
||||
torch.ops.aten.getitem.default,
|
||||
]
|
||||
|
||||
def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
|
||||
return [torch.ops.aten.getitem.default]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_scatter_split_replace(dtype):
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(0)
|
||||
|
||||
num_heads = 8
|
||||
num_kv_heads = 4
|
||||
head_size = 64
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
compilation_config=CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
custom_ops=["+rotary_embedding"],
|
||||
),
|
||||
)
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
# ScatterSplitReplacementPass requires SplitCoalescingPass to be run before it
|
||||
coalesce_pass = SplitCoalescingPass(vllm_config)
|
||||
replace_pass = ScatterSplitReplacementPass(vllm_config)
|
||||
passes = [coalesce_pass, replace_pass]
|
||||
backend = TestBackend(*passes)
|
||||
|
||||
model = ScatterSplitReplacementModel(num_heads, num_kv_heads, head_size, dtype)
|
||||
|
||||
T = 5
|
||||
qkv = torch.randn(
|
||||
T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
|
||||
)
|
||||
pos = torch.arange(T, dtype=torch.long)
|
||||
|
||||
qkv_eager = qkv.clone()
|
||||
pos_eager = pos.clone()
|
||||
result_eager = model(qkv_eager, pos_eager)
|
||||
|
||||
torch._dynamo.mark_dynamic(qkv, 0)
|
||||
torch._dynamo.mark_dynamic(pos, 0)
|
||||
|
||||
model_compiled = torch.compile(model, backend=backend)
|
||||
result_compiled = model_compiled(qkv, pos)
|
||||
|
||||
for eager, compiled in zip(result_eager, result_compiled):
|
||||
torch.testing.assert_close(eager, compiled)
|
||||
|
||||
assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
|
||||
assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1
|
||||
@@ -26,22 +26,14 @@ from vllm.config import (
|
||||
VllmConfig,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
|
||||
FP8ScaledMMLinearKernel,
|
||||
PerTensorTorchFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
|
||||
ROCmFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
|
||||
FP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
|
||||
@@ -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
|
||||
@@ -447,7 +447,7 @@ def test_metrics_exist_run_batch():
|
||||
"--model",
|
||||
"intfloat/multilingual-e5-small",
|
||||
"--enable-metrics",
|
||||
"--url",
|
||||
"--host",
|
||||
base_url,
|
||||
"--port",
|
||||
port,
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -39,6 +39,7 @@ def pairs_of_event_types() -> dict[str, str]:
|
||||
"response.mcp_call.completed": "response.mcp_call.in_progress",
|
||||
"response.function_call_arguments.done": "response.function_call_arguments.delta", # noqa: E501
|
||||
"response.code_interpreter_call_code.done": "response.code_interpreter_call_code.delta", # noqa: E501
|
||||
"response.code_interpreter_call.completed": "response.code_interpreter_call.in_progress", # noqa: E501
|
||||
"response.web_search_call.completed": "response.web_search_call.in_progress",
|
||||
}
|
||||
# fmt: on
|
||||
@@ -108,29 +109,19 @@ def events_contain_type(events: list, type_substring: str) -> bool:
|
||||
return any(type_substring in getattr(e, "type", "") for e in events)
|
||||
|
||||
|
||||
def validate_streaming_event_stack(
|
||||
events: list, pairs_of_event_types: dict[str, str]
|
||||
) -> None:
|
||||
"""Validate that streaming events are properly nested/paired."""
|
||||
def _validate_event_pairing(events: list, pairs_of_event_types: dict[str, str]) -> None:
|
||||
"""Validate that streaming events are properly nested/paired.
|
||||
|
||||
Derives push/pop sets from *pairs_of_event_types* so that every
|
||||
start/end pair in the dict is handled automatically.
|
||||
"""
|
||||
start_events = set(pairs_of_event_types.values())
|
||||
end_events = set(pairs_of_event_types.keys())
|
||||
|
||||
stack: list[str] = []
|
||||
for event in events:
|
||||
etype = event.type
|
||||
if etype == "response.created":
|
||||
stack.append(etype)
|
||||
elif etype == "response.completed":
|
||||
assert stack and stack[-1] == pairs_of_event_types[etype], (
|
||||
f"Unexpected stack top for {etype}: "
|
||||
f"got {stack[-1] if stack else '<empty>'}"
|
||||
)
|
||||
stack.pop()
|
||||
elif etype.endswith("added") or etype == "response.mcp_call.in_progress":
|
||||
stack.append(etype)
|
||||
elif etype.endswith("delta"):
|
||||
if stack and stack[-1] == etype:
|
||||
continue
|
||||
stack.append(etype)
|
||||
elif etype.endswith("done") or etype == "response.mcp_call.completed":
|
||||
assert etype in pairs_of_event_types, f"Unknown done event: {etype}"
|
||||
if etype in end_events:
|
||||
expected_start = pairs_of_event_types[etype]
|
||||
assert stack and stack[-1] == expected_start, (
|
||||
f"Stack mismatch for {etype}: "
|
||||
@@ -138,9 +129,180 @@ def validate_streaming_event_stack(
|
||||
f"got {stack[-1] if stack else '<empty>'}"
|
||||
)
|
||||
stack.pop()
|
||||
elif etype in start_events:
|
||||
# Consecutive deltas of the same type share a single stack slot.
|
||||
if etype.endswith("delta") and stack and stack[-1] == etype:
|
||||
continue
|
||||
stack.append(etype)
|
||||
# else: passthrough event (e.g. response.in_progress,
|
||||
# web_search_call.searching, code_interpreter_call.interpreting)
|
||||
assert len(stack) == 0, f"Unclosed events on stack: {stack}"
|
||||
|
||||
|
||||
def _validate_event_ordering(events: list) -> None:
|
||||
"""Validate that envelope events appear in the correct positions."""
|
||||
assert len(events) >= 2, f"Expected at least 2 events, got {len(events)}"
|
||||
|
||||
# First event must be response.created
|
||||
assert events[0].type == "response.created", (
|
||||
f"First event must be response.created, got {events[0].type}"
|
||||
)
|
||||
# Last event must be response.completed
|
||||
assert events[-1].type == "response.completed", (
|
||||
f"Last event must be response.completed, got {events[-1].type}"
|
||||
)
|
||||
|
||||
# response.in_progress, if present, must be the second event
|
||||
in_progress_indices = [
|
||||
i for i, e in enumerate(events) if e.type == "response.in_progress"
|
||||
]
|
||||
if in_progress_indices:
|
||||
assert in_progress_indices == [1], (
|
||||
f"response.in_progress must be the second event, "
|
||||
f"found at indices {in_progress_indices}"
|
||||
)
|
||||
|
||||
# Exactly one created and one completed
|
||||
created_count = sum(1 for e in events if e.type == "response.created")
|
||||
completed_count = sum(1 for e in events if e.type == "response.completed")
|
||||
assert created_count == 1, (
|
||||
f"Expected exactly 1 response.created, got {created_count}"
|
||||
)
|
||||
assert completed_count == 1, (
|
||||
f"Expected exactly 1 response.completed, got {completed_count}"
|
||||
)
|
||||
|
||||
|
||||
def _validate_field_consistency(events: list) -> None:
|
||||
"""Validate item_id, output_index, and content_index consistency.
|
||||
|
||||
Tracks the active output item established by ``output_item.added``
|
||||
and verifies that all subsequent events for that item carry matching
|
||||
identifiers until ``output_item.done`` closes it.
|
||||
"""
|
||||
_SESSION_EVENTS = {
|
||||
"response.created",
|
||||
"response.in_progress",
|
||||
"response.completed",
|
||||
}
|
||||
|
||||
active_item_id: str | None = None
|
||||
active_output_index: int | None = None
|
||||
last_output_index: int = -1
|
||||
active_content_index: int | None = None
|
||||
|
||||
for event in events:
|
||||
etype = event.type
|
||||
|
||||
if etype in _SESSION_EVENTS:
|
||||
continue
|
||||
|
||||
# --- output_item.added: opens a new item ------------------
|
||||
if etype == "response.output_item.added":
|
||||
item = getattr(event, "item", None)
|
||||
output_index = getattr(event, "output_index", None)
|
||||
|
||||
assert item is not None, "output_item.added must have an item"
|
||||
item_id = getattr(item, "id", None)
|
||||
assert item_id, "output_item.added item must have an id"
|
||||
|
||||
# output_index must be non-decreasing across items
|
||||
if output_index is not None:
|
||||
assert output_index >= last_output_index, (
|
||||
f"output_index went backwards: {output_index} < {last_output_index}"
|
||||
)
|
||||
last_output_index = output_index
|
||||
|
||||
active_item_id = item_id
|
||||
active_output_index = output_index
|
||||
active_content_index = None
|
||||
continue
|
||||
|
||||
# --- output_item.done: closes the active item -------------
|
||||
if etype == "response.output_item.done":
|
||||
item = getattr(event, "item", None)
|
||||
output_index = getattr(event, "output_index", None)
|
||||
|
||||
assert item is not None, "output_item.done must have an item"
|
||||
done_item_id = getattr(item, "id", None)
|
||||
|
||||
if active_item_id is not None and done_item_id:
|
||||
assert done_item_id == active_item_id, (
|
||||
f"output_item.done item.id mismatch: "
|
||||
f"expected {active_item_id}, got {done_item_id}"
|
||||
)
|
||||
if active_output_index is not None and output_index is not None:
|
||||
assert output_index == active_output_index, (
|
||||
f"output_item.done output_index mismatch: "
|
||||
f"expected {active_output_index}, got {output_index}"
|
||||
)
|
||||
|
||||
active_item_id = None
|
||||
active_output_index = None
|
||||
active_content_index = None
|
||||
continue
|
||||
|
||||
# --- content_part / reasoning_part added: sets content_index
|
||||
if etype in (
|
||||
"response.content_part.added",
|
||||
"response.reasoning_part.added",
|
||||
):
|
||||
_assert_item_fields(event, etype, active_item_id, active_output_index)
|
||||
active_content_index = getattr(event, "content_index", None)
|
||||
continue
|
||||
|
||||
# --- all other item-level events --------------------------
|
||||
_assert_item_fields(event, etype, active_item_id, active_output_index)
|
||||
|
||||
# content_index (only meaningful on events that carry it)
|
||||
content_index = getattr(event, "content_index", None)
|
||||
if content_index is not None and active_content_index is not None:
|
||||
assert content_index == active_content_index, (
|
||||
f"{etype} content_index mismatch: "
|
||||
f"expected {active_content_index}, got {content_index}"
|
||||
)
|
||||
|
||||
|
||||
def _assert_item_fields(
|
||||
event,
|
||||
etype: str,
|
||||
active_item_id: str | None,
|
||||
active_output_index: int | None,
|
||||
) -> None:
|
||||
"""Check that *event*'s item_id and output_index match the active item."""
|
||||
event_item_id = getattr(event, "item_id", None)
|
||||
output_index = getattr(event, "output_index", None)
|
||||
|
||||
if active_item_id is not None and event_item_id is not None:
|
||||
assert event_item_id == active_item_id, (
|
||||
f"{etype} item_id mismatch: expected {active_item_id}, got {event_item_id}"
|
||||
)
|
||||
if active_output_index is not None and output_index is not None:
|
||||
assert output_index == active_output_index, (
|
||||
f"{etype} output_index mismatch: "
|
||||
f"expected {active_output_index}, got {output_index}"
|
||||
)
|
||||
|
||||
|
||||
def validate_streaming_event_stack(
|
||||
events: list, pairs_of_event_types: dict[str, str]
|
||||
) -> None:
|
||||
"""Validate streaming events: pairing, ordering, and field consistency.
|
||||
|
||||
Checks three aspects:
|
||||
1. **Event pairing** — start/end events are properly nested
|
||||
(stack-based matching derived from *pairs_of_event_types*).
|
||||
2. **Event ordering** — envelope events (``created``,
|
||||
``in_progress``, ``completed``) appear at the correct positions.
|
||||
3. **Field consistency** — ``item_id``, ``output_index``, and
|
||||
``content_index`` are consistent across related events within
|
||||
each output item's lifecycle.
|
||||
"""
|
||||
_validate_event_pairing(events, pairs_of_event_types)
|
||||
_validate_event_ordering(events)
|
||||
_validate_field_consistency(events)
|
||||
|
||||
|
||||
def log_response_diagnostics(
|
||||
response,
|
||||
*,
|
||||
|
||||
@@ -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
|
||||
@@ -910,21 +909,25 @@ async def test_function_calling_no_code_interpreter_events(
|
||||
reason="This test is flaky in CI, needs investigation and "
|
||||
"potential fixes in the code interpreter MCP implementation."
|
||||
)
|
||||
async def test_mcp_code_interpreter_streaming(client: OpenAI, model_name: str, server):
|
||||
tools = [{"type": "mcp", "server_label": "code_interpreter"}]
|
||||
async def test_code_interpreter_streaming(
|
||||
client: OpenAI,
|
||||
model_name: str,
|
||||
pairs_of_event_types: dict[str, str],
|
||||
):
|
||||
tools = [{"type": "code_interpreter", "container": {"type": "auto"}}]
|
||||
input_text = (
|
||||
"Calculate 123 * 456 using python. "
|
||||
"The python interpreter is not stateful and you must "
|
||||
"print to see the output."
|
||||
)
|
||||
|
||||
def _has_mcp_call(evts: list) -> bool:
|
||||
return events_contain_type(evts, "mcp_call")
|
||||
def _has_code_interpreter(evts: list) -> bool:
|
||||
return events_contain_type(evts, "code_interpreter")
|
||||
|
||||
events = await retry_streaming_for(
|
||||
client,
|
||||
model=model_name,
|
||||
validate_events=_has_mcp_call,
|
||||
validate_events=_has_code_interpreter,
|
||||
input=input_text,
|
||||
tools=tools,
|
||||
temperature=0.0,
|
||||
@@ -936,59 +939,36 @@ async def test_mcp_code_interpreter_streaming(client: OpenAI, model_name: str, s
|
||||
event_types = [e.type for e in events]
|
||||
event_types_set = set(event_types)
|
||||
logger.info(
|
||||
"\n====== MCP Streaming Diagnostics ======\n"
|
||||
"\n====== Code Interpreter Streaming Diagnostics ======\n"
|
||||
"Event count: %d\n"
|
||||
"Event types (in order): %s\n"
|
||||
"Unique event types: %s\n"
|
||||
"=======================================",
|
||||
"====================================================",
|
||||
len(events),
|
||||
event_types,
|
||||
sorted(event_types_set),
|
||||
)
|
||||
|
||||
# Verify the full MCP streaming lifecycle
|
||||
assert "response.output_item.added" in event_types_set, (
|
||||
f"MCP call was not added. Events: {sorted(event_types_set)}"
|
||||
)
|
||||
assert "response.mcp_call.in_progress" in event_types_set, (
|
||||
f"MCP call in_progress not seen. Events: {sorted(event_types_set)}"
|
||||
)
|
||||
assert "response.mcp_call_arguments.delta" in event_types_set, (
|
||||
f"MCP arguments delta not seen. Events: {sorted(event_types_set)}"
|
||||
)
|
||||
assert "response.mcp_call_arguments.done" in event_types_set, (
|
||||
f"MCP arguments done not seen. Events: {sorted(event_types_set)}"
|
||||
)
|
||||
assert "response.mcp_call.completed" in event_types_set, (
|
||||
f"MCP call completed not seen. Events: {sorted(event_types_set)}"
|
||||
)
|
||||
assert "response.output_item.done" in event_types_set, (
|
||||
f"MCP item done not seen. Events: {sorted(event_types_set)}"
|
||||
)
|
||||
# Structural validation (pairing, ordering, field consistency)
|
||||
validate_streaming_event_stack(events, pairs_of_event_types)
|
||||
|
||||
# Validate specific MCP event details
|
||||
# Validate code interpreter item fields
|
||||
for event in events:
|
||||
if event.type == "response.output_item.added":
|
||||
if hasattr(event.item, "type") and event.item.type == "mcp_call":
|
||||
assert event.item.name == "python"
|
||||
assert event.item.server_label == "code_interpreter"
|
||||
elif event.type == "response.mcp_call_arguments.done":
|
||||
assert event.name == "python"
|
||||
assert event.arguments is not None
|
||||
if (
|
||||
event.type == "response.output_item.added"
|
||||
and hasattr(event.item, "type")
|
||||
and event.item.type == "code_interpreter_call"
|
||||
):
|
||||
assert event.item.status == "in_progress"
|
||||
elif event.type == "response.code_interpreter_call_code.done":
|
||||
assert event.code is not None
|
||||
elif (
|
||||
event.type == "response.output_item.done"
|
||||
and hasattr(event.item, "type")
|
||||
and event.item.type == "mcp_call"
|
||||
and event.item.type == "code_interpreter_call"
|
||||
):
|
||||
assert event.item.name == "python"
|
||||
assert event.item.status == "completed"
|
||||
|
||||
# code_interpreter events should NOT appear when using MCP type
|
||||
code_interp_events = [e.type for e in events if "code_interpreter" in e.type]
|
||||
assert not code_interp_events, (
|
||||
"Should not see code_interpreter events when using MCP type, "
|
||||
f"but got: {code_interp_events}"
|
||||
)
|
||||
assert event.item.code is not None
|
||||
|
||||
|
||||
@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. "
|
||||
@@ -241,81 +241,3 @@ class TestMCPEnabled:
|
||||
)
|
||||
|
||||
validate_streaming_event_stack(events, pairs_of_event_types)
|
||||
|
||||
assert events_contain_type(events, "mcp_call"), (
|
||||
f"No mcp_call events after retries. "
|
||||
f"Event types: {sorted({e.type for e in events})}"
|
||||
)
|
||||
|
||||
|
||||
class TestMCPDisabled:
|
||||
"""Tests that MCP tools are not executed when the env flag is unset."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def mcp_disabled_server(self):
|
||||
env_dict = {
|
||||
**BASE_TEST_ENV,
|
||||
"VLLM_ENABLE_RESPONSES_API_STORE": "1",
|
||||
"PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
|
||||
"VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": "1",
|
||||
}
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, list(_BASE_SERVER_ARGS), env_dict=env_dict
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(self, mcp_disabled_server):
|
||||
async with mcp_disabled_server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_mcp_disabled_server_does_not_execute(
|
||||
self, client: OpenAI, model_name: str
|
||||
):
|
||||
"""When MCP is disabled the model may still attempt tool calls
|
||||
(tool descriptions can remain in the prompt), but the server
|
||||
must NOT execute them."""
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=(
|
||||
"Execute the following code if the tool is present: "
|
||||
"import random; print(random.randint(1, 1000000))"
|
||||
),
|
||||
tools=[
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_label": "code_interpreter",
|
||||
"server_url": "http://localhost:8888",
|
||||
}
|
||||
],
|
||||
temperature=0.0,
|
||||
extra_body={"enable_response_messages": True},
|
||||
)
|
||||
assert response is not None
|
||||
assert response.status == "completed"
|
||||
|
||||
log_response_diagnostics(response, label="MCP Disabled")
|
||||
|
||||
# Server must not have executed any tool calls
|
||||
for message in response.output_messages:
|
||||
author = message.get("author", {})
|
||||
assert not (
|
||||
author.get("role") == "tool"
|
||||
and (author.get("name") or "").startswith("python")
|
||||
), (
|
||||
"Server executed a python tool call even though MCP is "
|
||||
f"disabled. Message: {message}"
|
||||
)
|
||||
|
||||
# No completed mcp_call output items
|
||||
for item in response.output:
|
||||
if getattr(item, "type", None) == "mcp_call":
|
||||
assert getattr(item, "status", None) != "completed", (
|
||||
"MCP call should not be completed when MCP is disabled"
|
||||
)
|
||||
|
||||
# No developer messages injected
|
||||
for message in response.input_messages:
|
||||
assert message.get("author", {}).get("role") != "developer"
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -219,3 +219,23 @@ async def test_completion_error_stream():
|
||||
f"Expected error message in chunks: {chunks}"
|
||||
)
|
||||
assert chunks[-1] == "data: [DONE]\n\n"
|
||||
|
||||
|
||||
def test_negative_prompt_token_ids_nested():
|
||||
"""Negative token IDs in prompt (nested list) should raise validation error."""
|
||||
with pytest.raises(Exception, match="greater than or equal to 0"):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
prompt=[[-1]],
|
||||
max_tokens=10,
|
||||
)
|
||||
|
||||
|
||||
def test_negative_prompt_token_ids_flat():
|
||||
"""Negative token IDs in prompt (flat list) should raise validation error."""
|
||||
with pytest.raises(Exception, match="greater than or equal to 0"):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
prompt=[-1],
|
||||
max_tokens=10,
|
||||
)
|
||||
|
||||
@@ -10,59 +10,361 @@ import pytest
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.entrypoints.openai.run_batch import BatchRequestOutput
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
CHAT_MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
EMBEDDING_MODEL_NAME = "intfloat/multilingual-e5-small"
|
||||
RERANKER_MODEL_NAME = "BAAI/bge-reranker-v2-m3"
|
||||
REASONING_MODEL_NAME = "Qwen/Qwen3-0.6B"
|
||||
SPEECH_LARGE_MODEL_NAME = "openai/whisper-large-v3"
|
||||
SPEECH_SMALL_MODEL_NAME = "openai/whisper-small"
|
||||
|
||||
# ruff: noqa: E501
|
||||
INPUT_BATCH = (
|
||||
'{{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {{"model": "{0}", "messages": [{{"role": "system", "content": "You are a helpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}\n'
|
||||
'{{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {{"model": "{0}", "messages": [{{"role": "system", "content": "You are an unhelpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}\n'
|
||||
'{{"custom_id": "request-3", "method": "POST", "url": "/v1/chat/completions", "body": {{"model": "NonExistModel", "messages": [{{"role": "system", "content": "You are an unhelpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}\n'
|
||||
'{{"custom_id": "request-4", "method": "POST", "url": "/bad_url", "body": {{"model": "{0}", "messages": [{{"role": "system", "content": "You are an unhelpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}\n'
|
||||
'{{"custom_id": "request-5", "method": "POST", "url": "/v1/chat/completions", "body": {{"stream": "True", "model": "{0}", "messages": [{{"role": "system", "content": "You are an unhelpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}'
|
||||
).format(MODEL_NAME)
|
||||
|
||||
INVALID_INPUT_BATCH = (
|
||||
'{{"invalid_field": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {{"model": "{0}", "messages": [{{"role": "system", "content": "You are a helpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}\n'
|
||||
'{{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {{"model": "{0}", "messages": [{{"role": "system", "content": "You are an unhelpful assistant."}},{{"role": "user", "content": "Hello world!"}}],"max_tokens": 1000}}}}'
|
||||
).format(MODEL_NAME)
|
||||
|
||||
INPUT_EMBEDDING_BATCH = (
|
||||
'{"custom_id": "request-1", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "You are a helpful assistant."}}\n'
|
||||
'{"custom_id": "request-2", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "You are an unhelpful assistant."}}\n'
|
||||
'{"custom_id": "request-3", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "Hello world!"}}\n'
|
||||
'{"custom_id": "request-4", "method": "POST", "url": "/v1/embeddings", "body": {"model": "NonExistModel", "input": "Hello world!"}}'
|
||||
INPUT_BATCH = "\n".join(
|
||||
json.dumps(req)
|
||||
for req in [
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": CHAT_MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant.",
|
||||
},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": CHAT_MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an unhelpful assistant.",
|
||||
},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-3",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": "NonExistModel",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an unhelpful assistant.",
|
||||
},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-4",
|
||||
"method": "POST",
|
||||
"url": "/bad_url",
|
||||
"body": {
|
||||
"model": CHAT_MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an unhelpful assistant.",
|
||||
},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-5",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"stream": "True",
|
||||
"model": CHAT_MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an unhelpful assistant.",
|
||||
},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
INPUT_SCORE_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
|
||||
{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}"""
|
||||
INVALID_INPUT_BATCH = "\n".join(
|
||||
json.dumps(req)
|
||||
for req in [
|
||||
{
|
||||
"invalid_field": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": CHAT_MODEL_NAME,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": CHAT_MODEL_NAME,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are an unhelpful assistant."},
|
||||
{"role": "user", "content": "Hello world!"},
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
INPUT_RERANK_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
|
||||
{"custom_id": "request-2", "method": "POST", "url": "/v1/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
|
||||
{"custom_id": "request-2", "method": "POST", "url": "/v2/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}"""
|
||||
INPUT_EMBEDDING_BATCH = "\n".join(
|
||||
json.dumps(req)
|
||||
for req in [
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/embeddings",
|
||||
"body": {
|
||||
"model": EMBEDDING_MODEL_NAME,
|
||||
"input": "You are a helpful assistant.",
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v1/embeddings",
|
||||
"body": {
|
||||
"model": EMBEDDING_MODEL_NAME,
|
||||
"input": "You are an unhelpful assistant.",
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-3",
|
||||
"method": "POST",
|
||||
"url": "/v1/embeddings",
|
||||
"body": {
|
||||
"model": EMBEDDING_MODEL_NAME,
|
||||
"input": "Hello world!",
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-4",
|
||||
"method": "POST",
|
||||
"url": "/v1/embeddings",
|
||||
"body": {
|
||||
"model": "NonExistModel",
|
||||
"input": "Hello world!",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
INPUT_REASONING_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Solve this math problem: 2+2=?"}]}}
|
||||
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "What is the capital of France?"}]}}"""
|
||||
_SCORE_RERANK_DOCUMENTS = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
INPUT_SCORE_BATCH = "\n".join(
|
||||
json.dumps(req)
|
||||
for req in [
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/score",
|
||||
"body": {
|
||||
"model": RERANKER_MODEL_NAME,
|
||||
"queries": "What is the capital of France?",
|
||||
"documents": _SCORE_RERANK_DOCUMENTS,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v1/score",
|
||||
"body": {
|
||||
"model": RERANKER_MODEL_NAME,
|
||||
"queries": "What is the capital of France?",
|
||||
"documents": _SCORE_RERANK_DOCUMENTS,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
INPUT_RERANK_BATCH = "\n".join(
|
||||
json.dumps(req)
|
||||
for req in [
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/rerank",
|
||||
"body": {
|
||||
"model": RERANKER_MODEL_NAME,
|
||||
"query": "What is the capital of France?",
|
||||
"documents": _SCORE_RERANK_DOCUMENTS,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v1/rerank",
|
||||
"body": {
|
||||
"model": RERANKER_MODEL_NAME,
|
||||
"query": "What is the capital of France?",
|
||||
"documents": _SCORE_RERANK_DOCUMENTS,
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v2/rerank",
|
||||
"body": {
|
||||
"model": RERANKER_MODEL_NAME,
|
||||
"query": "What is the capital of France?",
|
||||
"documents": _SCORE_RERANK_DOCUMENTS,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
INPUT_REASONING_BATCH = "\n".join(
|
||||
json.dumps(req)
|
||||
for req in [
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": REASONING_MODEL_NAME,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Solve this math problem: 2+2=?"},
|
||||
],
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "request-2",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": REASONING_MODEL_NAME,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
],
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
# This is a valid but minimal audio file for testing
|
||||
MINIMAL_WAV_BASE64 = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAEAfAAABAAgAZGF0YQAAAAA="
|
||||
INPUT_TRANSCRIPTION_BATCH = (
|
||||
'{{"custom_id": "request-1", "method": "POST", "url": "/v1/audio/transcriptions", '
|
||||
'"body": {{"model": "openai/whisper-large-v3", "file_url": "data:audio/wav;base64,{}", '
|
||||
'"response_format": "json"}}}}\n'
|
||||
).format(MINIMAL_WAV_BASE64)
|
||||
json.dumps(
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/audio/transcriptions",
|
||||
"body": {
|
||||
"model": SPEECH_LARGE_MODEL_NAME,
|
||||
"file_url": f"data:audio/wav;base64,{MINIMAL_WAV_BASE64}",
|
||||
"response_format": "json",
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
INPUT_TRANSCRIPTION_HTTP_BATCH = (
|
||||
'{{"custom_id": "request-1", "method": "POST", "url": "/v1/audio/transcriptions", '
|
||||
'"body": {{"model": "openai/whisper-large-v3", "file_url": "{}", '
|
||||
'"response_format": "json"}}}}\n'
|
||||
).format(AudioAsset("mary_had_lamb").url)
|
||||
json.dumps(
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/audio/transcriptions",
|
||||
"body": {
|
||||
"model": SPEECH_LARGE_MODEL_NAME,
|
||||
"file_url": AudioAsset("mary_had_lamb").url,
|
||||
"response_format": "json",
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
INPUT_TRANSLATION_BATCH = (
|
||||
'{{"custom_id": "request-1", "method": "POST", "url": "/v1/audio/translations", '
|
||||
'"body": {{"model": "openai/whisper-small", "file_url": "{}", '
|
||||
'"response_format": "text", "language": "it", "to_language": "en", '
|
||||
'"temperature": 0.0}}}}\n'
|
||||
).format(AudioAsset("mary_had_lamb").url)
|
||||
json.dumps(
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/audio/translations",
|
||||
"body": {
|
||||
"model": SPEECH_SMALL_MODEL_NAME,
|
||||
"file_url": AudioAsset("mary_had_lamb").url,
|
||||
"response_format": "text",
|
||||
"language": "it",
|
||||
"to_language": "en",
|
||||
"temperature": 0.0,
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
WEATHER_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
INPUT_TOOL_CALLING_BATCH = json.dumps(
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": REASONING_MODEL_NAME,
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is the weather in San Francisco?"},
|
||||
],
|
||||
"tools": [WEATHER_TOOL],
|
||||
"tool_choice": "required",
|
||||
"max_tokens": 1000,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_empty_file():
|
||||
@@ -81,7 +383,7 @@ def test_empty_file():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"intfloat/multilingual-e5-small",
|
||||
EMBEDDING_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -108,7 +410,7 @@ def test_completions():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
MODEL_NAME,
|
||||
CHAT_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -141,7 +443,7 @@ def test_completions_invalid_input():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
MODEL_NAME,
|
||||
CHAT_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -165,7 +467,7 @@ def test_embeddings():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"intfloat/multilingual-e5-small",
|
||||
EMBEDDING_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -196,7 +498,7 @@ def test_score(input_batch):
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"BAAI/bge-reranker-v2-m3",
|
||||
RERANKER_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -234,7 +536,7 @@ def test_reasoning_parser():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"Qwen/Qwen3-0.6B",
|
||||
REASONING_MODEL_NAME,
|
||||
"--reasoning-parser",
|
||||
"qwen3",
|
||||
],
|
||||
@@ -278,7 +580,7 @@ def test_transcription():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"openai/whisper-large-v3",
|
||||
SPEECH_LARGE_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -316,7 +618,7 @@ def test_transcription_http_url():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"openai/whisper-large-v3",
|
||||
SPEECH_LARGE_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -356,7 +658,7 @@ def test_translation():
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
"openai/whisper-small",
|
||||
SPEECH_SMALL_MODEL_NAME,
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
@@ -378,3 +680,69 @@ def test_translation():
|
||||
translation_text = response_body["text"]
|
||||
translation_text_lower = str(translation_text).strip().lower()
|
||||
assert "mary" in translation_text_lower or "lamb" in translation_text_lower
|
||||
|
||||
|
||||
def test_tool_calling():
|
||||
"""
|
||||
Test that tool calling works correctly in run_batch.
|
||||
Verifies that requests with tools return tool_calls in the response.
|
||||
"""
|
||||
with (
|
||||
tempfile.NamedTemporaryFile("w") as input_file,
|
||||
tempfile.NamedTemporaryFile("r") as output_file,
|
||||
):
|
||||
input_file.write(INPUT_TOOL_CALLING_BATCH)
|
||||
input_file.flush()
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
"vllm",
|
||||
"run-batch",
|
||||
"-i",
|
||||
input_file.name,
|
||||
"-o",
|
||||
output_file.name,
|
||||
"--model",
|
||||
REASONING_MODEL_NAME,
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
],
|
||||
)
|
||||
proc.communicate()
|
||||
proc.wait()
|
||||
assert proc.returncode == 0, f"{proc=}"
|
||||
|
||||
contents = output_file.read()
|
||||
for line in contents.strip().split("\n"):
|
||||
if not line.strip(): # Skip empty lines
|
||||
continue
|
||||
# Ensure that the output format conforms to the openai api.
|
||||
# Validation should throw if the schema is wrong.
|
||||
BatchRequestOutput.model_validate_json(line)
|
||||
|
||||
# Ensure that there is no error in the response.
|
||||
line_dict = json.loads(line)
|
||||
assert isinstance(line_dict, dict)
|
||||
assert line_dict["error"] is None
|
||||
|
||||
# Check that tool_calls are present in the response
|
||||
# With tool_choice="required", the model must call a tool
|
||||
response_body = line_dict["response"]["body"]
|
||||
assert response_body is not None
|
||||
message = response_body["choices"][0]["message"]
|
||||
assert "tool_calls" in message
|
||||
tool_calls = message.get("tool_calls")
|
||||
# With tool_choice="required", tool_calls must be present and non-empty
|
||||
assert tool_calls is not None
|
||||
assert isinstance(tool_calls, list)
|
||||
assert len(tool_calls) > 0
|
||||
# Verify tool_calls have the expected structure
|
||||
for tool_call in tool_calls:
|
||||
assert "id" in tool_call
|
||||
assert "type" in tool_call
|
||||
assert tool_call["type"] == "function"
|
||||
assert "function" in tool_call
|
||||
assert "name" in tool_call["function"]
|
||||
assert "arguments" in tool_call["function"]
|
||||
# Verify the tool name matches our tool definition
|
||||
assert tool_call["function"]["name"] == "get_current_weather"
|
||||
|
||||
@@ -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"
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user