Merge remote-tracking branch 'upstream/main' into patches

Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
2026-10-03 22:06:10 +07:00
327 changed files with 13704 additions and 5433 deletions
@@ -28,4 +28,13 @@ steps:
export VLLM_WORKER_MULTIPROC_METHOD=spawn && export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
export PYTHONFAULTHANDLER=1 && export PYTHONFAULTHANDLER=1 &&
cd tests && cd tests &&
pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormModule::test_matches_reference_above_cudnn_grid_limit"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype0-out_dtype0]"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype1-out_dtype1]"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype2-out_dtype2]"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype3-out_dtype3]"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype4-out_dtype4]"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype5-out_dtype5]"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormShapes"
--deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormInputHandling"'
@@ -72,8 +72,7 @@ steps:
- >- - >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests && 'cd tests &&
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"'
pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
- label: "Multi-Modal Models (Standard) 4: other + whisper" - label: "Multi-Modal Models (Standard) 4: other + whisper"
key: multi-modal-models-standard-4-other-whisper key: multi-modal-models-standard-4-other-whisper
@@ -3,7 +3,6 @@
"test_name": "latency_llama8B_tp1", "test_name": "latency_llama8B_tp1",
"environment_variables": { "environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
"parameters": { "parameters": {
@@ -3,7 +3,6 @@
"test_name": "latency_llama8B_tp2", "test_name": "latency_llama8B_tp2",
"environment_variables": { "environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_SGL_KERNEL": 1,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
@@ -14,7 +14,6 @@
], ],
"server_environment_variables": { "server_environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_SGL_KERNEL": 1,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
@@ -4,9 +4,6 @@
"inf" "inf"
], ],
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200], "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": {
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120
},
"server_parameters": { "server_parameters": {
"dtype": "bfloat16", "dtype": "bfloat16",
"model": "openai/whisper-large-v3-turbo" "model": "openai/whisper-large-v3-turbo"
@@ -10,7 +10,6 @@
], ],
"server_environment_variables": { "server_environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_SGL_KERNEL": 1,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
@@ -6,7 +6,6 @@
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200], "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": { "server_environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_SGL_KERNEL": 1,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
@@ -6,7 +6,6 @@
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200], "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": { "server_environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_SGL_KERNEL": 1,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
@@ -3,7 +3,6 @@
"test_name": "throughput_llama8B_tp1", "test_name": "throughput_llama8B_tp1",
"environment_variables": { "environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
"parameters": { "parameters": {
@@ -3,7 +3,6 @@
"test_name": "throughput_llama8B_tp2", "test_name": "throughput_llama8B_tp2",
"environment_variables": { "environment_variables": {
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_SGL_KERNEL": 1,
"VLLM_CPU_KVCACHE_SPACE": 40 "VLLM_CPU_KVCACHE_SPACE": 40
}, },
+4 -2
View File
@@ -339,7 +339,8 @@ steps:
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
--build-arg RUN_WHEEL_CHECK=false \ --build-arg RUN_WHEEL_CHECK=false \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.4 \ --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.4 \
--build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \ --build-arg FINAL_BASE_IMAGE=nvidia/cuda:13.4.1-base-ubuntu24.04 \
--build-arg NCCL_VERSION=2.32.3 \
--target vllm-openai \ --target vllm-openai \
--progress plain \ --progress plain \
-f docker/Dockerfile . -f docker/Dockerfile .
@@ -369,7 +370,8 @@ steps:
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
--build-arg RUN_WHEEL_CHECK=false \ --build-arg RUN_WHEEL_CHECK=false \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.4 \ --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.4 \
--build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \ --build-arg FINAL_BASE_IMAGE=nvidia/cuda:13.4.1-base-ubuntu24.04 \
--build-arg NCCL_VERSION=2.32.3 \
--target vllm-openai \ --target vllm-openai \
--progress plain \ --progress plain \
-f docker/Dockerfile . -f docker/Dockerfile .
@@ -1712,8 +1712,7 @@ if is_multi_node "$commands"; then
# BASH_REMATCH[2] = comma-separated node0 commands # BASH_REMATCH[2] = comma-separated node0 commands
# BASH_REMATCH[3] = comma-separated node1 commands # BASH_REMATCH[3] = comma-separated node1 commands
if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then
# shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR prefix=${BASH_REMATCH[1]//;/}
prefix=$(echo "${BASH_REMATCH[1]}" | sed 's/;//g')
echo "PREFIX: ${prefix}" echo "PREFIX: ${prefix}"
export composite_command="(command rocm-smi || true)" export composite_command="(command rocm-smi || true)"
@@ -1728,10 +1727,8 @@ if is_multi_node "$commands"; then
fi fi
for i in "${!node0[@]}"; do for i in "${!node0[@]}"; do
# shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR command_node_0=${node0[i]//\"/}
command_node_0=$(echo "${node0[i]}" | sed 's/\"//g') command_node_1=${node1[i]//\"/}
# shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR
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}'" 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}" echo "COMMANDS: ${step_cmd}"
@@ -34,6 +34,7 @@ readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))"
RUN_ID="" CONTAINER_NAME="" CONTAINER_ID="" RUN_ID="" CONTAINER_NAME="" CONTAINER_ID=""
GPU_UUID="" PORT_ONE="" PORT_TWO="" GPU_UUID="" PORT_ONE="" PORT_TWO=""
LINK_REMAP_BASELINE="" LINK_REMAP_BASELINE=""
CHECK_CLEANUP_STATE=0
SNAPSHOT_PIDS=() SNAPSHOT_PIDS=()
die() { die() {
echo "initialized snapshot E2E: $*" >&2 echo "initialized snapshot E2E: $*" >&2
@@ -113,7 +114,7 @@ cleanup() {
names="$(container_names 2>/dev/null)" || return 1 names="$(container_names 2>/dev/null)" || return 1
grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1 grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1
fi fi
if [[ -n "$PORT_TWO" ]] && ! wait_clean; then if (( CHECK_CLEANUP_STATE )) && ! wait_clean; then
echo "process, GPU, port, or link-remap residue remained" >&2 echo "process, GPU, port, or link-remap residue remained" >&2
failed=1 failed=1
fi fi
@@ -218,7 +219,8 @@ GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM" RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM"
CONTAINER_NAME="vllm-$RUN_ID" CONTAINER_NAME="vllm-$RUN_ID"
read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))') read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))')
LINK_REMAP_BASELINE="$(link_remaps)" LINK_REMAP_BASELINE="$(link_remaps)" \
|| die "cannot inventory /dev/shm; check CI runner permissions"
run "pull exact candidate image" 900 docker pull --platform linux/amd64 \ run "pull exact candidate image" 900 docker pull --platform linux/amd64 \
"$IMAGE_TAG" >/dev/null "$IMAGE_TAG" >/dev/null
@@ -231,6 +233,8 @@ IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <(
|| die "candidate image identity is invalid" || die "candidate image identity is invalid"
[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \ [[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \
|| die "candidate image commit labels do not match" || die "candidate image commit labels do not match"
# A failed preflight has not started any workload whose cleanup we can verify.
CHECK_CLEANUP_STATE=1
CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \ CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \
--name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \ --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \
--gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \ --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \
+2
View File
@@ -2398,6 +2398,7 @@ steps:
- vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
- tests/quantization - tests/quantization
- tests/rocm/test_moe_weight_replay.py - tests/rocm/test_moe_weight_replay.py
- tests/rocm/test_moe_padding.py
optional: true optional: true
commands: commands:
# temporary install here since we need nightly, will move to requirements/test.in # temporary install here since we need nightly, will move to requirements/test.in
@@ -2411,6 +2412,7 @@ steps:
- uv pip install --system conch-triton-kernels - uv pip install --system conch-triton-kernels
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi
- if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi
- label: ":amd: (MI300) Quantized Fusions" # TBD - label: ":amd: (MI300) Quantized Fusions" # TBD
timeout_in_minutes: 45 timeout_in_minutes: 45
+60
View File
@@ -93,6 +93,24 @@ steps:
commands: commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY=1 - export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -s -v tests/compile/passes/distributed - pytest -s -v tests/compile/passes/distributed
mirror:
amd:
label: ':amd: (MI300) Distributed Compile'
dind: false
device: mi300_2
timeout_in_minutes: 65
working_dir: /vllm-workspace/
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/compilation/
- vllm/model_executor/layers
- tests/compile/passes/distributed/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -s -v tests/compile/passes/distributed
- label: ":nvidia: (B200) Fusion and Compile" - label: ":nvidia: (B200) Fusion and Compile"
key: fusion-and-compile-unit-tests-2xb200 key: fusion-and-compile-unit-tests-2xb200
@@ -463,3 +481,45 @@ steps:
- amd-smi - amd-smi
- VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)'
- pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py
- label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)'
key: fusion-e2e-tp2-ar-rms-inductor-partition-amd
dind: false
device: mi355_2
num_devices: 2
timeout_in_minutes: 65
working_dir: /vllm-workspace/
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/compilation/
- vllm/model_executor/layers
- tests/compile/fusions_e2e/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- csrc/quantization/
- vllm/v1/attention/
commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "inductor_partition"
- label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Dynamo Partition)'
key: fusion-e2e-tp2-ar-rms-dynamo-partition-amd
dind: false
device: mi355_2
num_devices: 2
timeout_in_minutes: 65
working_dir: /vllm-workspace/
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/compilation/
- vllm/model_executor/layers
- tests/compile/fusions_e2e/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- csrc/quantization/
- vllm/v1/attention/
commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition"
+135
View File
@@ -17,6 +17,24 @@ steps:
- pytest -v -s distributed/test_shm_broadcast.py - pytest -v -s distributed/test_shm_broadcast.py
- pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_buffer.py
- pytest -v -s distributed/test_shm_storage.py - pytest -v -s distributed/test_shm_storage.py
mirror:
amd:
label: ':amd: (MI300) Distributed Comm Ops'
dind: false
device: mi300_2
timeout_in_minutes: 30
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed
- tests/distributed
- vllm/platforms/rocm.py
commands:
- pytest -v -s distributed/test_comm_ops.py
- pytest -v -s distributed/test_shm_broadcast.py
- pytest -v -s distributed/test_shm_buffer.py
- pytest -v -s distributed/test_shm_storage.py
- label: ":nvidia: (L4) Sharded RDT Weight Transfer" - label: ":nvidia: (L4) Sharded RDT Weight Transfer"
key: sharded-rdt-weight-transfer key: sharded-rdt-weight-transfer
@@ -176,6 +194,31 @@ steps:
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- pytest -v -s v1/worker/test_worker_memory_snapshot.py - pytest -v -s v1/worker/test_worker_memory_snapshot.py
mirror:
amd:
label: ':amd: (MI300) Distributed Torchrun + Shutdown'
dind: false
device: mi300_2
timeout_in_minutes: 35
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/
- vllm/engine/
- vllm/v1/executor/
- vllm/v1/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- tests/distributed/
- tests/v1/shutdown
- tests/v1/worker/test_worker_memory_snapshot.py
- vllm/platforms/rocm.py
commands:
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
- label: ":nvidia: (L4) Distributed Torchrun + Examples" - label: ":nvidia: (L4) Distributed Torchrun + Examples"
key: distributed-torchrun-examples-4-gpus key: distributed-torchrun-examples-4-gpus
@@ -282,6 +325,31 @@ steps:
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
- pytest -v -s distributed/test_utils.py - pytest -v -s distributed/test_utils.py
- pytest -v -s distributed/test_engram_dp_shard.py - pytest -v -s distributed/test_engram_dp_shard.py
mirror:
amd:
label: ':amd: (MI355) Distributed DP Extended'
dind: false
device: mi355_4
timeout_in_minutes: 60
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/
- tests/v1/distributed
- tests/v1/engine/test_engine_core_client.py
- tests/distributed/test_utils
- vllm/platforms/rocm.py
commands:
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
- pytest -v -s distributed/test_utils.py
- label: ":nvidia: (L4) Distributed Compile + Comm" - label: ":nvidia: (L4) Distributed Compile + Comm"
key: distributed-compile-comm-4-gpus key: distributed-compile-comm-4-gpus
@@ -304,6 +372,27 @@ steps:
- pytest -v -s distributed/test_symm_mem_allreduce.py - pytest -v -s distributed/test_symm_mem_allreduce.py
# test multi-node TP with multiproc executor (simulated on single node) # test multi-node TP with multiproc executor (simulated on single node)
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
mirror:
amd:
label: ':amd: (MI355) Distributed Compile + Comm'
dind: false
device: mi355_4
timeout_in_minutes: 45
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/
- tests/distributed/test_pynccl
- tests/distributed/test_events
- tests/distributed/test_symm_mem_allreduce.py
- tests/distributed/test_multiproc_executor.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s distributed/test_pynccl.py
- pytest -v -s distributed/test_events.py
- pytest -v -s distributed/test_symm_mem_allreduce.py
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
- label: ":nvidia: (H100) Distributed DP + EP" - label: ":nvidia: (H100) Distributed DP + EP"
key: distributed-tests-8xh100 key: distributed-tests-8xh100
@@ -359,6 +448,28 @@ steps:
- torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
- TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- pytest -v -s -x lora/test_mixtral.py - pytest -v -s -x lora/test_mixtral.py
mirror:
amd:
label: ':amd: (MI355) Distributed'
dind: false
device: mi355_4
timeout_in_minutes: 40
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/
- tests/models/test_vision.py
- vllm/platforms/rocm.py
- tests/distributed/
- tests/basic_correctness/
- tests/lora/test_mixtral.py
commands:
- pytest -v -s distributed/test_custom_all_reduce.py
- torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
- TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- pytest -v -s -x lora/test_mixtral.py
- pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)'
- label: ":nvidia: (H100) Distributed Features" - label: ":nvidia: (H100) Distributed Features"
key: distributed-tests-2xh100-2xmi300 key: distributed-tests-2xh100-2xmi300
@@ -583,3 +694,27 @@ steps:
- pytest -v -s distributed/test_ray_v2_executor_e2e.py - pytest -v -s distributed/test_ray_v2_executor_e2e.py
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray" - pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray" - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray"
mirror:
amd:
label: ':amd: (MI355) RayExecutorV2'
dind: false
device: mi355_4
timeout_in_minutes: 75
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/executor/ray_executor_v2.py
- vllm/v1/executor/abstract.py
- vllm/v1/executor/multiproc_executor.py
- tests/distributed/test_ray_v2_executor.py
- tests/distributed/test_ray_v2_executor_e2e.py
- tests/distributed/test_pipeline_parallel.py
- tests/basic_correctness/models/test_basic_correctness.py
- vllm/platforms/rocm.py
commands:
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
- pytest -v -s distributed/test_ray_v2_executor.py
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
- TARGET_TEST_SUITE=MI355 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray"
@@ -150,3 +150,27 @@ steps:
- .buildkite/scripts/scheduled_integration_test/ - .buildkite/scripts/scheduled_integration_test/
- vllm/_aiter_ops.py - vllm/_aiter_ops.py
- vllm/platforms/rocm.py - vllm/platforms/rocm.py
- label: ":amd: (MI355) Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy"
key: amd-qwen3-next-mtp-async-eplb-accuracy
dind: false
device: mi355_4
timeout_in_minutes: 45
depends_on:
- image-build-amd
working_dir: /vllm-workspace
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/spec_decode/
- vllm/distributed/eplb
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- .buildkite/scripts/scheduled_integration_test/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
num_devices: 4
commands:
- VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
@@ -10,9 +10,11 @@ steps:
source_file_dependencies: source_file_dependencies:
- vllm/distributed/eplb - vllm/distributed/eplb
- tests/distributed/test_eplb_algo.py - tests/distributed/test_eplb_algo.py
- tests/distributed/test_eplb_migration_scheduler.py
- tests/distributed/test_eplb_utils.py - tests/distributed/test_eplb_utils.py
commands: commands:
- pytest -v -s distributed/test_eplb_algo.py - pytest -v -s distributed/test_eplb_algo.py
- pytest -v -s distributed/test_eplb_migration_scheduler.py
- pytest -v -s distributed/test_eplb_utils.py - pytest -v -s distributed/test_eplb_utils.py
mirror: mirror:
amd: amd:
+147
View File
@@ -753,6 +753,32 @@ steps:
parallelism: 4 parallelism: 4
commands: commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200-shard-$$BUILDKITE_PARALLEL_JOB.txt - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200-shard-$$BUILDKITE_PARALLEL_JOB.txt
mirror:
amd:
label: ":amd: (MI355) MoE Refactor BF16 Integration"
dind: false
device: mi355_2
num_devices: 2
parallelism: 1
timeout_in_minutes: 105
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- tests/evals/gsm8k/
- tests/utils.py
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/model_executor/kernels/linear/
- vllm/v1/attention/
- vllm/distributed/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
env:
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt -k "Llama-4-Scout-BF16-triton or Mixtral-8x7B-BF16-triton"
- label: ":nvidia: (B200) MoE Refactor DP Integration TEMPORARY" - label: ":nvidia: (B200) MoE Refactor DP Integration TEMPORARY"
key: moe-refactor-integration-test-b200-dp-temporary key: moe-refactor-integration-test-b200-dp-temporary
@@ -761,6 +787,32 @@ steps:
num_devices: 2 num_devices: 2
commands: commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
mirror:
amd:
label: ":amd: (MI355) MoE Refactor BF16 DP Integration"
dind: false
device: mi355_2
num_devices: 2
timeout_in_minutes: 105
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- tests/evals/gsm8k/
- tests/utils.py
- vllm/config/parallel.py
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/model_executor/kernels/linear/
- vllm/v1/attention/
- vllm/distributed/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
env:
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -k "Qwen3-30B-A3B-BF16-triton"
- label: ":nvidia: (L4) LM Eval Humming TEMPORARY" - label: ":nvidia: (L4) LM Eval Humming TEMPORARY"
key: lm-eval-humming-l4 key: lm-eval-humming-l4
@@ -1068,6 +1120,29 @@ steps:
- tests/evals/gpt_oss/ - tests/evals/gpt_oss/
commands: commands:
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
mirror:
amd:
label: ":amd: (MI300) GPQA Eval (GPT-OSS)"
dind: false
device: mi300_2
num_devices: 2
timeout_in_minutes: 40
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/layers/fused_moe/
- tests/evals/gpt_oss/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt
- label: ":nvidia: (B200) GPQA Eval (GPT-OSS)" - label: ":nvidia: (B200) GPQA Eval (GPT-OSS)"
key: gpqa-eval-gpt-oss-2xb200 key: gpqa-eval-gpt-oss-2xb200
@@ -1176,6 +1251,28 @@ steps:
- tests/evals/gpt_oss/ - tests/evals/gpt_oss/
commands: commands:
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
mirror:
amd:
label: ":amd: (MI355 DPX) GPQA Eval (GPT-OSS TP1)"
# The Spark config is a backend-neutral GPT-OSS TP1 baseline.
dind: false
device: mi355_dpx
num_devices: 1
timeout_in_minutes: 40
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/layers/fused_moe/
- tests/evals/gpt_oss/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 35GB) LM Eval KV-Offload" - label: ":nvidia: (H200 MIG 35GB) LM Eval KV-Offload"
key: kv-offload-small key: kv-offload-small
@@ -1292,3 +1389,53 @@ steps:
- vllm/_aiter_ops.py - vllm/_aiter_ops.py
- vllm/platforms/rocm.py - vllm/platforms/rocm.py
- tests/evals/mrcr/ - tests/evals/mrcr/
- label: ":amd: (MI355) LM Eval Large Models Harness"
key: amd-lm-eval-large-models-harness
dind: false
device: mi355_8
optional: true
timeout_in_minutes: 85
depends_on:
- image-build-amd
working_dir: /vllm-workspace/.buildkite/lm-eval-harness
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- .buildkite/lm-eval-harness/
num_devices: 8
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# Qwen2-57B has 28 attention heads and cannot run at TP8.
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
- label: ":amd: (MI300) LM Eval Large Models ROCm Harness"
key: amd-lm-eval-large-models-rocm-harness
dind: false
device: mi300_8
optional: true
timeout_in_minutes: 40
depends_on:
- image-build-amd
working_dir: /vllm-workspace/.buildkite/lm-eval-harness
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/quantization/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/layers/layernorm.py
- csrc/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- .buildkite/lm-eval-harness/
num_devices: 8
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400
+3
View File
@@ -560,6 +560,8 @@ steps:
- tools/pre_commit/check_test_tethering.py - tools/pre_commit/check_test_tethering.py
- tools/pre_commit/test_tethering_allowlist.txt - tools/pre_commit/test_tethering_allowlist.txt
- tests/tools/test_check_test_tethering.py - tests/tools/test_check_test_tethering.py
- tools/pre_commit/check_agent_files.py
- tests/tools/test_check_agent_files.py
- tests/tool_parsers - tests/tool_parsers
- tests/tokenizers_ - tests/tokenizers_
- tests/parser - tests/parser
@@ -609,6 +611,7 @@ steps:
- pytest -v -s -m 'cpu_test' multimodal - pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s config - pytest -v -s config
- pytest -v -s tools/test_check_test_tethering.py - pytest -v -s tools/test_check_test_tethering.py
- pytest -v -s tools/test_check_agent_files.py
- label: ":computer: (CPU) Reasoning + Renderers" - label: ":computer: (CPU) Reasoning + Renderers"
key: cpu-reasoning-renderers key: cpu-reasoning-renderers
+23 -4
View File
@@ -124,10 +124,9 @@ steps:
- csrc/libtorch_stable/kimi_k3/ - csrc/libtorch_stable/kimi_k3/
- tests/models/kimi_k3/ - tests/models/kimi_k3/
- tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
- tests/kernels/test_bf16_skinny_gemm.py
commands: commands:
# The native NVIDIA Kimi K3 kernels require the SM100 family. # The native NVIDIA Kimi K3 kernels require the SM100 family.
- pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py - pytest -v -s models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py kernels/attention/test_kimi_k3_mla_fused_epilogue.py
mirror: mirror:
amd: amd:
label: ":amd: (MI355) Kimi K3" label: ":amd: (MI355) Kimi K3"
@@ -142,10 +141,30 @@ steps:
- vllm/platforms/rocm.py - vllm/platforms/rocm.py
commands: commands:
# The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355. # The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355.
# The two kernels/ files the NVIDIA job adds are CUDA-only (they import # The kernels/ file the NVIDIA job adds is CUDA-only (it imports
# cuda.bindings.driver), so they are deliberately left out here. # cuda.bindings.driver), so it is deliberately left out here.
- pytest -v -s models/kimi_k3 - pytest -v -s models/kimi_k3
- label: ":nvidia: (4xB200) Kimi K3 Prefix Cache Shard %N"
key: kimi-k3-prefix-cache-4xb200
timeout_in_minutes: 75
device: b200-k8s
num_devices: 4
source_file_dependencies:
- vllm/models/kimi_k3/
- vllm/v1/core/
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/simple_kv_offload/
- vllm/v1/spec_decode/
- vllm/model_executor/layers/mamba/
- tests/models/kimi_k3/test_prefix_cache.py
- .buildkite/scripts/install-kv-connectors.sh
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 2
- label: ":nvidia: (H200) GLM5Next" - label: ":nvidia: (H200) GLM5Next"
key: glm5next-unit-tests key: glm5next-unit-tests
timeout_in_minutes: 20 timeout_in_minutes: 20
+19
View File
@@ -259,6 +259,25 @@ steps:
- tests/compile - tests/compile
commands: commands:
- pytest -s -v compile/fullgraph/test_full_cudagraph.py - pytest -s -v compile/fullgraph/test_full_cudagraph.py
mirror:
amd:
label: ':amd: (MI250) PyTorch Fullgraph CUDAGraph Compatibility'
dind: false
device: mi250_1
timeout_in_minutes: 65
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/compilation/
- vllm/model_executor/
- vllm/v1/attention/
- vllm/config/compilation.py
- csrc/
- tests/compile
- vllm/platforms/rocm.py
commands:
- pytest -s -v compile/fullgraph/test_full_cudagraph.py
- label: ":nvidia: (H200 MIG 18GB) PyTorch Nightly Dependency Override Check" - label: ":nvidia: (H200 MIG 18GB) PyTorch Nightly Dependency Override Check"
key: pytorch-nightly-dependency-override-check key: pytorch-nightly-dependency-override-check
+2
View File
@@ -38,6 +38,7 @@ steps:
- vllm/model_executor/layers/fused_moe/oracle/unquantized.py - vllm/model_executor/layers/fused_moe/oracle/unquantized.py
- vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
- tests/rocm/test_moe_weight_replay.py - tests/rocm/test_moe_weight_replay.py
- tests/rocm/test_moe_padding.py
commands: commands:
# Use the ROCm default instead of the NVIDIA runner override. # Use the ROCm default instead of the NVIDIA runner override.
- unset VLLM_USE_V2_MODEL_RUNNER - unset VLLM_USE_V2_MODEL_RUNNER
@@ -45,6 +46,7 @@ steps:
- uv pip install --system conch-triton-kernels - uv pip install --system conch-triton-kernels
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi
- if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi
- label: ":nvidia: (H200 MIG 35GB) Quantized Fusions" - label: ":nvidia: (H200 MIG 35GB) Quantized Fusions"
device: h200_35gb device: h200_35gb
+64
View File
@@ -97,6 +97,27 @@ steps:
- tests/v1/e2e/spec_decode/ - tests/v1/e2e/spec_decode/
commands: commands:
- pytest -v -s v1/e2e/spec_decode/eagle/ - pytest -v -s v1/e2e/spec_decode/eagle/
mirror:
amd:
label: ":amd: (MI355) Spec Decode Eagle Nightly"
dind: false
device: mi355_1
num_devices: 1
timeout_in_minutes: 120
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- vllm/v1/attention/
- tests/v1/e2e/spec_decode/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
# MTP correctness is one parametrized test split into per-model folders; # MTP correctness is one parametrized test split into per-model folders;
# each job below runs whole directories. # each job below runs whole directories.
@@ -234,6 +255,28 @@ steps:
commands: commands:
- pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/ - pytest -v -s v1/e2e/spec_decode/mtp/
mirror:
amd:
label: ":amd: (MI355) Spec Decode Speculators + MTP Nightly"
dind: false
device: mi355_1
num_devices: 1
timeout_in_minutes: 90
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- vllm/v1/attention/
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 35GB) Spec Decode N-Gram + Suffix" - label: ":nvidia: (H200 MIG 35GB) Spec Decode N-Gram + Suffix"
key: spec-decode-ngram-suffix key: spec-decode-ngram-suffix
@@ -304,6 +347,27 @@ steps:
- tests/v1/e2e/spec_decode/ - tests/v1/e2e/spec_decode/
commands: commands:
- pytest -v -s v1/e2e/spec_decode/draft_model/ - pytest -v -s v1/e2e/spec_decode/draft_model/
mirror:
amd:
label: ":amd: (MI355) Spec Decode Draft Model Nightly"
dind: false
device: mi355_1
num_devices: 1
timeout_in_minutes: 120
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- vllm/v1/attention/
- tests/v1/e2e/spec_decode/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H100) Speculators Correctness Nightly" - label: ":nvidia: (H100) Speculators Correctness Nightly"
key: speculators-correctness key: speculators-correctness
+8 -1
View File
@@ -30,7 +30,6 @@ repos:
- id: markdownlint-cli2 - id: markdownlint-cli2
language_version: lts language_version: lts
args: [--config, tools/pre_commit/.markdownlint.yaml, --fix] args: [--config, tools/pre_commit/.markdownlint.yaml, --fix]
exclude: (^|/)CLAUDE\.md$
- repo: https://github.com/rhysd/actionlint - repo: https://github.com/rhysd/actionlint
rev: v1.7.7 rev: v1.7.7
hooks: hooks:
@@ -177,6 +176,14 @@ repos:
files: ^\.github/mergify\.yml$ files: ^\.github/mergify\.yml$
pass_filenames: false pass_filenames: false
additional_dependencies: [PyYAML, regex] additional_dependencies: [PyYAML, regex]
- id: check-agent-files
name: Check agent instruction files and skills
language: python
entry: python tools/pre_commit/check_agent_files.py
files: ^(\.agents/|\.claude/|(.*/)?(AGENTS|CLAUDE)\.md$|tools/pre_commit/check_agent_files\.py$)
pass_filenames: false
additional_dependencies: [PyYAML, regex]
verbose: true # `note:` lines are advisory, so the hook passes
- id: check-buildkite-step-keys - id: check-buildkite-step-keys
name: Check every Buildkite step declares an explicit key name: Check every Buildkite step declares an explicit key
language: python language: python
@@ -0,0 +1,158 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Benchmark the ROCm sparse-indexer decode top-k in isolation.
Example usage:
python3 benchmarks/kernels/benchmark_top_k_per_row_decode.py \
--rows 4 16 32 --lens 2048 8192 16384 65536 --capacity 262144
"""
import argparse
import itertools
import json
import torch
import vllm._custom_ops # noqa: F401
from vllm._aiter_ops import rocm_aiter_ops
from vllm.triton_utils import triton
DISTRIBUTIONS = ("indexer", "fp8", "randn")
def make_logits(rows: int, capacity: int, dist: str, spread: float):
normal = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
if dist == "randn":
return normal
if dist == "fp8":
return normal.to(torch.float8_e4m3fnuz).to(torch.float32)
return normal * spread + 3.0
def make_inputs(rows: int, live: int, capacity: int, k: int, args):
torch.manual_seed(args.seed)
logits = make_logits(rows, capacity, args.dist, args.spread)
seq_lens = torch.full((rows,), live, device="cuda", dtype=torch.int32)
invalid = torch.arange(capacity, device="cuda")[None] >= seq_lens[:, None]
logits.masked_fill_(invalid, -float("inf"))
indices = torch.empty((rows, k), device="cuda", dtype=torch.int32)
return logits, seq_lens, indices
def run_native(logits, seq_lens, indices, k):
torch.ops._C.top_k_per_row_decode(
logits,
1,
seq_lens,
indices,
logits.shape[0],
logits.stride(0),
logits.stride(1),
k,
)
def run_aiter(logits, seq_lens, indices, k):
rocm_aiter_ops.indexer_top_k_decode(logits, 1, seq_lens, indices, k)
BACKENDS = {"native": run_native, "aiter": run_aiter}
def check(logits, indices, k):
selected = logits.gather(1, indices.long()).sort(dim=1, descending=True).values
expected = logits.topk(k, dim=1).values
torch.testing.assert_close(selected, expected, atol=0, rtol=0)
def capture(run):
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side):
for _ in range(3):
run()
torch.cuda.current_stream().wait_stream(side)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
run()
return graph
def benchmark(backend: str, rows: int, live: int, args) -> dict:
logits, seq_lens, indices = make_inputs(rows, live, args.capacity, args.top_k, args)
fn = BACKENDS[backend]
def run():
fn(logits, seq_lens, indices, args.top_k)
run()
torch.accelerator.synchronize()
check(logits, indices, args.top_k)
timed = run
if args.graph:
graph = capture(run)
indices.zero_()
graph.replay()
torch.accelerator.synchronize()
check(logits, indices, args.top_k)
timed = graph.replay
ms = triton.testing.do_bench(timed, warmup=args.warmup, rep=args.rep)
return dict(
backend=backend,
mode="graph" if args.graph else "eager",
dist=args.dist,
rows=rows,
live=live,
capacity=args.capacity,
k=args.top_k,
us=ms * 1000,
live_gbps=rows * live * 4 / (ms * 1e6),
)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8])
parser.add_argument(
"--lens",
type=int,
nargs="+",
default=[32768, 65536, 98304, 131072],
help="Valid columns per row; must be at least --top-k",
)
parser.add_argument("--capacity", type=int, default=262144)
parser.add_argument("--top-k", type=int, default=512)
parser.add_argument(
"--backends", nargs="+", choices=list(BACKENDS), default=["native", "aiter"]
)
parser.add_argument(
"--graph",
action="store_true",
help="Time a captured graph replay, matching how decode actually runs",
)
parser.add_argument(
"--dist",
choices=DISTRIBUTIONS,
default="indexer",
help="Logit value distribution; see the module docstring",
)
parser.add_argument(
"--spread",
type=float,
default=0.01,
help="Relative width of the --dist indexer cluster",
)
parser.add_argument("--warmup", type=int, default=25)
parser.add_argument("--rep", type=int, default=200)
parser.add_argument("--seed", type=int, default=123)
args = parser.parse_args()
print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args))))
for backend, rows, live in itertools.product(args.backends, args.rows, args.lens):
if live < args.top_k:
continue
print(json.dumps(benchmark(backend, rows, live, args)), flush=True)
if __name__ == "__main__":
main()
+42 -9
View File
@@ -23,9 +23,11 @@ from vllm.utils.platform_utils import get_device_name_as_file_name
mp.set_start_method("spawn", force=True) mp.set_start_method("spawn", force=True)
assert current_platform.is_cuda() or current_platform.is_rocm(), ( assert (
"Only support tune w8a8 block fp8 kernel on CUDA/ROCm device." current_platform.is_cuda()
) or current_platform.is_rocm()
or current_platform.is_xpu()
), "Only support tune w8a8 block fp8 kernel on CUDA/ROCm/XPU device."
DTYPE_MAP = { DTYPE_MAP = {
"float32": torch.float32, "float32": torch.float32,
@@ -138,6 +140,29 @@ def get_configs_compute_bound():
return configs return configs
def get_xpu_tuning_space():
# NOTE: curated subset of get_configs_compute_bound(); the full grid is
# too slow to tune on XPU due to per-config Triton JIT compile time.
configs = []
for num_stages in [2, 3]:
for block_m in [32, 64, 128]:
for block_k in [64, 128]:
for block_n in [128, 256]:
for num_warps in [4, 8, 16, 32]:
for group_size in [1, 32]:
configs.append(
{
"BLOCK_SIZE_M": block_m,
"BLOCK_SIZE_N": block_n,
"BLOCK_SIZE_K": block_k,
"GROUP_SIZE_M": group_size,
"num_warps": num_warps,
"num_stages": num_stages,
}
)
return configs
def get_weight_shapes(tp_size): def get_weight_shapes(tp_size):
# NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3.
# Modify them, if you tune for another different model. # Modify them, if you tune for another different model.
@@ -208,13 +233,14 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type):
fp8_info = torch.finfo(torch.float8_e4m3fn) fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min fp8_max, fp8_min = fp8_info.max, fp8_info.min
device = current_platform.device_type
A_fp32 = ( A_fp32 = (
(torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
) )
A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
B_fp32 = ( B_fp32 = (
(torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max (torch.rand(N, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
) )
B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
else: else:
@@ -224,9 +250,9 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type):
n_tiles = (N + block_n - 1) // block_n n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k k_tiles = (K + block_k - 1) // block_k
As = torch.rand(M, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
Bs = ( Bs = (
torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda") torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
* factor_for_scale * factor_for_scale
) )
@@ -297,7 +323,11 @@ def tune_on_gpu(args_dict):
save_path = args.save_path save_path = args.save_path
input_type = args.input_type input_type = args.input_type
search_space = get_configs_compute_bound() search_space = (
get_xpu_tuning_space()
if current_platform.is_xpu()
else get_configs_compute_bound()
)
search_space = [ search_space = [
config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0
] ]
@@ -342,7 +372,10 @@ def main(args):
raise RuntimeError("No GPU available for tuning") raise RuntimeError("No GPU available for tuning")
print(f"Found {num_gpus} GPUs for parallel tuning") print(f"Found {num_gpus} GPUs for parallel tuning")
torch.cuda.init() if current_platform.is_xpu():
torch.xpu.init()
else:
torch.cuda.init()
if args.batch_size is None: if args.batch_size is None:
batch_sizes = [ batch_sizes = [
+9 -109
View File
@@ -251,33 +251,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
fi fi
# CUDA 13.4 Rubin images need a compatibility link for the prerelease
# PyTorch wheel's libcurand.so.11 dependency. CUDA 13.5 must use its native
# libcurand.so.11 instead.
RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \
&& if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \
-print -quit) \
&& test -n "${CURAND_LIBRARY}" \
&& ln -sf "${CURAND_LIBRARY}" \
/usr/local/cuda/lib64/libcurand.so.11; \
elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \
-print -quit) \
&& test -n "${CURAND_LIBRARY}" \
&& ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \
else \
echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
exit 1; \
fi \
&& printf '%s\n' \
/opt/venv/lib \
"/opt/venv/lib/python${PYTHON_VERSION}/site-packages/nvpl/lib" \
> /etc/ld.so.conf.d/vllm-rubin-python.conf \
&& ldconfig; \
fi
# Track PyTorch lib versions used during build and match in downstream instances. # Track PyTorch lib versions used during build and match in downstream instances.
# We do this for both nightly and release so we can strip dependencies/*.txt as needed. # We do this for both nightly and release so we can strip dependencies/*.txt as needed.
# Otherwise library dependencies can upgrade/downgrade torch incorrectly. # Otherwise library dependencies can upgrade/downgrade torch incorrectly.
@@ -518,20 +491,6 @@ ARG NVSHMEM_VER
ARG NCCL_VERSION ARG NCCL_VERSION
RUN --mount=type=cache,target=/opt/uv/cache \ RUN --mount=type=cache,target=/opt/uv/cache \
mkdir -p /tmp/ep_kernels_workspace/dist && \ mkdir -p /tmp/ep_kernels_workspace/dist && \
if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
# The Rubin CUDA builder provides NCCL in its architecture-specific \
# toolkit directory rather than as a Python package. Give DeepEP a \
# stable prefix that also exists in the Ubuntu runtime image. \
NCCL_HEADER=$(find /usr/local/cuda/targets -name nccl.h -print -quit) \
&& test -n "${NCCL_HEADER}" \
&& NCCL_LIBRARY=$(find -H /usr/local/cuda/targets \
-name libnccl.so.2 -print -quit) \
&& test -n "${NCCL_LIBRARY}" \
&& ln -sf "${NCCL_HEADER}" /usr/include/nccl.h \
&& ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2 \
&& ln -sf libnccl.so.2 /usr/lib/libnccl.so \
&& export EP_NCCL_ROOT_DIR=/usr; \
fi && \
CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \ CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
if [ -n "$NCCL_VERSION" ]; then \ if [ -n "$NCCL_VERSION" ]; then \
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \ echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
@@ -659,10 +618,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
if [ "${PYTORCH_NIGHTLY}" = "1" ] || \ if [ "${PYTORCH_NIGHTLY}" = "1" ] || \
[ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
python3 tools/use_existing_torch.py --prefix; \ python3 tools/use_existing_torch.py --prefix; \
fi \
&& if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \
requirements/cuda.txt; \
fi fi
# Build the vLLM wheel # Build the vLLM wheel
@@ -764,11 +719,6 @@ ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL ARG GET_PIP_URL
# Keep the output image's CUDA metadata aligned with the selected toolkit.
# This also normalizes prerelease base images with four-component metadata for
# OCI runtimes that evaluate NVIDIA's CUDA driver requirement.
ENV CUDA_VERSION=${CUDA_VERSION}
ENV DEBIAN_FRONTEND=noninteractive ENV DEBIAN_FRONTEND=noninteractive
WORKDIR /vllm-workspace WORKDIR /vllm-workspace
@@ -841,41 +791,6 @@ ENV CUDA_HOME=/usr/local/cuda
ENV PATH="${CUDA_HOME}/bin:${PATH}" ENV PATH="${CUDA_HOME}/bin:${PATH}"
ENV LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}" ENV LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}"
# Apply the CUDA 13.4-only libcurand compatibility link described in the build
# stage. CUDA 13.5 uses its native libcurand.so.11. DeepEP expects NCCL directly
# under its configured /usr prefix, while Ubuntu installs it in a multiarch
# subdirectory - but that symlink is only possible when NCCL is already
# present (e.g. nvcr.io/nvidia/cuda-dl-base bundles it; a bare-CUDA image
# doesn't, and gets it later via the standard nvidia-nccl-cu* pip package,
# UV_OVERRIDE below). Check reality (whether libnccl.so.2 is actually found)
# and tolerate not finding it instead of failing the build.
RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \
&& if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \
-print -quit) \
&& test -n "${CURAND_LIBRARY}" \
&& ln -sf "${CURAND_LIBRARY}" \
/usr/local/cuda/lib64/libcurand.so.11; \
elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \
-print -quit) \
&& test -n "${CURAND_LIBRARY}" \
&& ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \
else \
echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
exit 1; \
fi \
&& if NCCL_LIBRARY=$(find /usr/lib -name libnccl.so.2 -print -quit) \
&& [ -n "${NCCL_LIBRARY}" ]; then \
if [ "${NCCL_LIBRARY}" != "/usr/lib/libnccl.so.2" ]; then \
ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2; \
fi \
&& ln -sf libnccl.so.2 /usr/lib/libnccl.so; \
fi \
&& ldconfig; \
fi
# Install the initialized-snapshot runtime on the architecture supported by # Install the initialized-snapshot runtime on the architecture supported by
# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned # `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned
# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's # source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's
@@ -990,11 +905,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
&& uv pip install --system -r /tmp/requirements-cuda.txt \ && uv pip install --system -r /tmp/requirements-cuda.txt \
&& uv pip install --system \ && uv pip install --system \
--prerelease=allow --index-strategy unsafe-first-match \ --prerelease=allow --index-strategy unsafe-first-match \
-r /tmp/vllm-requirements/rubin-prerelease.txt \ -r /tmp/vllm-requirements/rubin-prerelease.txt; \
&& printf '%s\n' \
"/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvpl/lib" \
> /etc/ld.so.conf.d/vllm-rubin-python.conf \
&& ldconfig; \
else \ else \
uv pip install --system \ uv pip install --system \
-r /tmp/requirements-cuda.txt \ -r /tmp/requirements-cuda.txt \
@@ -1007,10 +918,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# From versions.json: .flashinfer.version # From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.7.0.post1 ARG FLASHINFER_VERSION=0.7.0.post1
RUN --mount=type=cache,target=/opt/uv/cache \ RUN --mount=type=cache,target=/opt/uv/cache \
if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
fi
# ============================================================ # ============================================================
# OPENAI API SERVER DEPENDENCIES # OPENAI API SERVER DEPENDENCIES
@@ -1023,22 +932,13 @@ ARG GDRCOPY_CUDA_VERSION=12.8
ARG GDRCOPY_OS_VERSION=Ubuntu24_04 ARG GDRCOPY_OS_VERSION=Ubuntu24_04
ARG TARGETPLATFORM ARG TARGETPLATFORM
COPY tools/install_gdrcopy.sh /tmp/install_gdrcopy.sh COPY tools/install_gdrcopy.sh /tmp/install_gdrcopy.sh
# Some Rubin FINAL_BASE_IMAGE choices (internal Rubin CUDA images,
# nvcr.io/nvidia/cuda-dl-base) already provide GDRCopy; others don't, so it
# falls through to the normal install below. Check reality (ldconfig) rather
# than inferring from which FINAL_BASE_IMAGE flag is set.
RUN set -eux; \ RUN set -eux; \
if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] && ldconfig -p | grep -q 'libgdrapi\.so'; then \ case "${TARGETPLATFORM}" in \
true; \ linux/arm64) UUARCH="aarch64" ;; \
else \ linux/amd64) UUARCH="x64" ;; \
case "${TARGETPLATFORM}" in \ *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
linux/arm64) UUARCH="aarch64" ;; \ esac; \
linux/amd64) UUARCH="x64" ;; \ /tmp/install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}" && \
*) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
esac; \
/tmp/install_gdrcopy.sh \
"${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \
fi; \
rm /tmp/install_gdrcopy.sh rm /tmp/install_gdrcopy.sh
# Install vllm-openai dependencies (saves ~2.6s per build) # Install vllm-openai dependencies (saves ~2.6s per build)
+55
View File
@@ -0,0 +1,55 @@
# Reviewers
Reviewers are trusted community members who can perform official reviews, helping contributors get their pull requests ready to merge alongside committers, who give final approval. Authors are expected to respond to reviewer feedback, either by making changes or by explaining why a change isn't needed. The role was introduced in [RFC #58253](https://github.com/vllm-project/vllm/issues/58253).
This page lists the active reviewers and describes the role. See the [committers](../governance/committers.md) page for the committers and area owners who hold final acceptance and merge authority.
## Active Reviewers
Community members can reach out to the reviewers below for help with PRs in their areas, for example in the `#pr-reviews` Slack channel. Reviewers will add their areas of expertise to this list.
Sorted alphabetically by GitHub handle:
- [@andylolu2](https://github.com/andylolu2)
- [@bnellnm](https://github.com/bnellnm) MoE, Quantization layers
- [@BowenBao](https://github.com/BowenBao)
- [@cjackal](https://github.com/cjackal): Multimodality; KV Connector and offload
- [@divakar-amd](https://github.com/divakar-amd)
- [@etelis](https://github.com/etelis): KV cache offloading (CPU, filesystem, and P2P connectors; KV transfer and cache layouts); MoE serving and expert parallelism (developing review focus)
- [@Fangzhou-Ai](https://github.com/Fangzhou-Ai)
- [@fxmarty-amd](https://github.com/fxmarty-amd): Quantization, linear/MOE oracles and modeling definition
- [@GirasoleY](https://github.com/GirasoleY): Kernels and performance
- [@itayalroy](https://github.com/itayalroy): MoE serving (kernels, EP, All2All, EPLB, elasticity/fault tolerance); KV connectors and NIXL integrations
- [@JartX](https://github.com/JartX): ROCm / HIP on Radeon (RDNA3), quantization (W4A16/GPTQ, MXFP4), KV cache quantization
- [@kliuae](https://github.com/kliuae): ROCm feature enablement, model performance, fused kernels
- [@lengrongfu](https://github.com/lengrongfu)
- [@linitra24](https://github.com/linitra24)
- [@LopezCastroRoberto](https://github.com/LopezCastroRoberto): Kernels and quantization
- [@mawong-amd](https://github.com/mawong-amd)
- [@micah-wil](https://github.com/micah-wil): ROCm / AMD GPU integration, CI
- [@netanel-haber](https://github.com/netanel-haber): Nemotron/NVIDIA models, Mamba/Linear attention hybrids, VLMs
- [@Rohan138](https://github.com/Rohan138): ROCm performance, CI, torch.compile fusions
- [@simondanielsson](https://github.com/simondanielsson): ROCm performance (CDNA), MoRI-IO.
- [@taneem-ibrahim](https://github.com/taneem-ibrahim): Pooling models
- [@TheEpicDolphin](https://github.com/TheEpicDolphin)
- [@varun-sundar-rabindranath](https://github.com/varun-sundar-rabindranath): KV cache offloading, LoRA
- [@vllmellm](https://github.com/vllmellm)
- [@wangxiyuan](https://github.com/wangxiyuan): Platform, KV cache, Mooncake
## The Reviewer Role
Reviewers assess the problem and proposed solution, guide validation, resolve concerns with authors, and determine when a contribution is ready for a committer. A reviewer's approval means a PR is ready for a committer's final look.
Reviewers can:
- Approve PRs or request changes.
- Trigger CI when a contribution is ready for evaluation.
- Triage PRs and issues, apply labels, and flag important PRs for expedited committer attention.
- Close duplicate PRs and issues.
- Escalate design questions to the appropriate committer or SIG.
- Flag PRs as ready for committer review.
### Becoming a Reviewer
Candidates should demonstrate a sustained record of useful reviews, sound engineering judgment, constructive communication, and familiarity with vLLM's contribution standards. Deep ownership of a subsystem is not required.
A reviewer is nominated by two committers from different organizations, who post the nomination in the committers Slack channel. As a rule of thumb, not a hard requirement, candidates have either 15 merged PRs and 30 reviewed PRs, or 45 reviewed PRs. A reviewed PR is one where the candidate gave a substantive review and the PR reached a resolution (merged or closed).
+10 -2
View File
@@ -307,15 +307,23 @@ review process:
- After the PR is submitted, the PR will be assigned to a reviewer. Every - After the PR is submitted, the PR will be assigned to a reviewer. Every
reviewer will pick up the PRs based on their expertise and availability. reviewer will pick up the PRs based on their expertise and availability.
Official reviews can be performed by any of the
[reviewers](../community/reviewers.md) or
[committers](../governance/committers.md); committers perform the final
merge. To find someone to review your PR, see the reviewers' and
committers' areas of expertise or ask in the `#pr-reviews` channel on
[Slack](https://slack.vllm.ai).
- After the PR is assigned, the reviewer will provide status updates every 2-3 - After the PR is assigned, the reviewer will provide status updates every 2-3
days. If the PR is not reviewed within 7 days, please feel free to ping the days. If the PR is not reviewed within 7 days, please feel free to ping the
reviewer or the vLLM team. reviewer or the vLLM team.
- After the review, the reviewer will put an `action-required` label on the PR - After the review, the reviewer will put an `action-required` label on the PR
if there are changes required. The contributor should address the comments and if there are changes required. The contributor should address the comments and
ping the reviewer to re-review the PR. ping the reviewer to re-review the PR.
- Please respond to all comments within a reasonable time frame. If a comment - Please respond to all comments within a reasonable time frame. Address each
piece of reviewer feedback with a change or an explanation. If a comment
isn't clear or you disagree with a suggestion, feel free to ask for isn't clear or you disagree with a suggestion, feel free to ask for
clarification or discuss the suggestion. clarification or discuss the suggestion. If you and a reviewer can't reach
agreement, a committer can help decide.
- Note that not all CI checks will be executed due to limited computational - Note that not all CI checks will be executed due to limited computational
resources. Reviewers with write access and configured trusted contributors resources. Reviewers with write access and configured trusted contributors
can comment `/ci run` for upstream CI or `/amd-ci run` for AMD CI only when can comment `/ci run` for upstream CI or `/amd-ci run` for AMD CI only when
@@ -44,6 +44,19 @@ cross-file context, and fixes for repeated mistakes.
Each entry should be short, specific, and actionable — e.g., which files to Each entry should be short, specific, and actionable — e.g., which files to
touch, what order to change them in, and which tests to run. touch, what order to change them in, and which tests to run.
## Skills
Skills live in `.agents/skills/<name>/` and follow the
[Agent Skills specification](https://agentskills.io/specification).
- Claude Code does not read `.agents/`, so symlink each skill:
`ln -s ../../.agents/skills/<name> .claude/skills/<name>`.
- Frontmatter: `name` matches the directory. `description` says what the
skill does and when to use it; quote it if it contains a colon.
- Keep `SKILL.md` under 500 lines and move detail into files it links to.
- `tools/pre_commit/check_agent_files.py` runs in pre-commit to check the
skill and `AGENTS.md` rules
## Keeping Docs Lean ## Keeping Docs Lean
- Every addition should trigger review of surrounding content for stale or redundant items. - Every addition should trigger review of surrounding content for stale or redundant items.
+1
View File
@@ -37,6 +37,7 @@ th {
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | | flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
| passthrough | standard | nvfp4,mxfp4 | G | N | Y | [`PassThroughPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough.PassThroughPrepareAndFinalize] |
!!! info "Table key" !!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8 1. All types: mxfp4, nvfp4, int4, int8, fp8
+1
View File
@@ -135,6 +135,7 @@ Batch invariance has been tested and verified on the following models:
- **Granite 3.1 (Dense)**: `ibm-granite/granite-3.1-2b-instruct`, `ibm-granite/granite-3.1-8b-instruct` - **Granite 3.1 (Dense)**: `ibm-granite/granite-3.1-2b-instruct`, `ibm-granite/granite-3.1-8b-instruct`
- **EXAONE 4.0 series**: `LGAI-EXAONE/EXAONE-4.0-1.2B`, `LGAI-EXAONE/EXAONE-4.0.1-32B`, `LGAI-EXAONE/EXAONE-4.0-32B` - **EXAONE 4.0 series**: `LGAI-EXAONE/EXAONE-4.0-1.2B`, `LGAI-EXAONE/EXAONE-4.0.1-32B`, `LGAI-EXAONE/EXAONE-4.0-32B`
- **OLMo 2**: `allenai/OLMo-2-0425-1B-Instruct` - **OLMo 2**: `allenai/OLMo-2-0425-1B-Instruct`
- **ERNIE 4.5**: `baidu/ERNIE-4.5-0.3B-PT`
- **SmolLM2**: `HuggingFaceTB/SmolLM2-1.7B-Instruct` - **SmolLM2**: `HuggingFaceTB/SmolLM2-1.7B-Instruct`
- **PLaMo3**: `pfnet/plamo-3-nict-2b-base` - **PLaMo3**: `pfnet/plamo-3-nict-2b-base`
+6
View File
@@ -79,6 +79,12 @@ decode = client.chat.completions.create(
If `messages` has non-text content or `echo` is set, the ids are ignored and `messages` is rendered instead, so it must match the prefill request. Otherwise `kv_transfer_params["prompt_token_ids"]` must be a non-empty list of non-negative integers, or the request fails with HTTP 400, as it always does on `/v1/chat/completions/batch`. If `messages` has non-text content or `echo` is set, the ids are ignored and `messages` is rendered instead, so it must match the prefill request. Otherwise `kv_transfer_params["prompt_token_ids"]` must be a non-empty list of non-negative integers, or the request fails with HTTP 400, as it always does on `/v1/chat/completions/batch`.
## Generate API output modes
When the prefill and decode stages use the [Generate API](../serving/online_serving/token_in_token_out.md) (`/inference/v1/generate`), only the decode response reaches the client, so set `output_mode` on the decode request only. A proxy that reuses the client's request body for the prefill request must reset `output_mode` to `tokens` there. A prefill instance started with `--tokens-only` has no tokenizer and rejects `output_mode: "text"` with a 400.
A decode instance that returns text needs a tokenizer, so start it with `--enable-scale-out` and without `--tokens-only`. Prefill instances can keep `--tokens-only`.
## Development ## Development
We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance. We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.
+11 -2
View File
@@ -31,6 +31,11 @@ Snapshots currently require:
Current Hugging Face hub clients hold their connections for the process Current Hugging Face hub clients hold their connections for the process
lifetime, so download the model in a separate step and run create with lifetime, so download the model in a separate step and run create with
`HF_HUB_OFFLINE=1`, as the quickstart below does. `HF_HUB_OFFLINE=1`, as the quickstart below does.
- Snapshot creation defaults `NCCL_IB_DISABLE=1` for its singleton donor and
inherited workers because CRIU cannot capture live InfiniBand/RDMA state.
An explicit caller value is retained, but creation rejects an open
`/dev/infiniband/` descriptor before CRIU. Close non-NCCL RDMA clients before
capture.
- A remote model ID and an immutable 40-character `--revision`. Local model - A remote model ID and an immutable 40-character `--revision`. Local model
directories and mutable revisions are not supported. directories and mutable revisions are not supported.
- Enough disk for the artifact, with the same installed vLLM package, model - Enough disk for the artifact, with the same installed vLLM package, model
@@ -138,9 +143,13 @@ snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time.
and directories. Recreate the artifact after an unclean host shutdown. and directories. Recreate the artifact after an unclean host shutdown.
- Restore currently requires the same host, GPU, driver, kernel, Python, - Restore currently requires the same host, GPU, driver, kernel, Python,
PyTorch, installed vLLM version, model revision, engine arguments, PyTorch, installed vLLM version, model revision, engine arguments,
selected environment variables, and CRIU plugin binaries. selected effective environment variables, and CRIU plugin binaries. An unset
`NCCL_IB_DISABLE` therefore matches the snapshot donor default of `1`, while
an explicit different value does not.
- Only dense float16 TP1 has been validated. Other model formats depend on their - Only dense float16 TP1 has been validated. Other model formats depend on their
existing sleep level 2 reload support; NCCL state is not restored. existing sleep level 2 reload support; distributed and RDMA snapshots are not
supported. The snapshot-only NCCL default does not change ordinary serve or
multi-GPU defaults.
- CRIU support varies by kernel and driver. Preserve package, library, model, - CRIU support varies by kernel and driver. Preserve package, library, model,
and generated-cache paths for the artifact lifetime. and generated-cache paths for the artifact lifetime.
- A snapshot can include application secrets or request state present in the - A snapshot can include application secrets or request state present in the
+1 -3
View File
@@ -33,9 +33,7 @@ following `quantization.quant_algo` values:
`flashinfer_cutlass`, `flashinfer_cutedsl`, `flashinfer_trtllm`, `flashinfer_cutlass`, `flashinfer_cutedsl`, `flashinfer_trtllm`,
`flashinfer_cudnn`, and `marlin`; see `KernelConfig` on the `flashinfer_cudnn`, and `marlin`; see `KernelConfig` on the
[Engine Arguments](../../configuration/engine_args.md) page and shown by [Engine Arguments](../../configuration/engine_args.md) page and shown by
`vllm serve --help=KernelConfig`. For `W4A16_NVFP4`, `auto` currently `vllm serve --help=KernelConfig`.
selects Marlin. Models with BF16 activations can explicitly select the
FlashInfer CuTe-DSL backend with `--linear-backend flashinfer_cutedsl`.
!!! note !!! note
For models quantized to MXFP8 with BF16 activations on SM100-family GPUs, For models quantized to MXFP8 with BF16 activations on SM100-family GPUs,
@@ -440,7 +440,8 @@ ref changes. Use `--no-cache-filter extensions-build` to refresh an empty,
branch, or tag revision. branch, or tag revision.
For `FINAL_BASE_IMAGE`, use the public, multi-arch For `FINAL_BASE_IMAGE`, use the public, multi-arch
`nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04` image. `nvidia/cuda:13.4.1-base-ubuntu24.04` image. Set `NCCL_VERSION` to 2.32.3 or
newer, the first NCCL release with Rubin (SM107) support.
For `BUILD_BASE_IMAGE`, use: For `BUILD_BASE_IMAGE`, use:
- `pytorch/manylinux2_28-builder:cuda13.4` for x86_64 CPUs. - `pytorch/manylinux2_28-builder:cuda13.4` for x86_64 CPUs.
@@ -462,7 +463,8 @@ For `BUILD_BASE_IMAGE`, use:
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
--build-arg CUDA_VERSION=13.4 \ --build-arg CUDA_VERSION=13.4 \
--build-arg BUILD_BASE_IMAGE="pytorch/manylinuxaarch64-builder:cuda13.4" \ --build-arg BUILD_BASE_IMAGE="pytorch/manylinuxaarch64-builder:cuda13.4" \
--build-arg FINAL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04" \ --build-arg FINAL_BASE_IMAGE="nvidia/cuda:13.4.1-base-ubuntu24.04" \
--build-arg NCCL_VERSION=2.32.3 \
. .
``` ```
@@ -482,7 +484,8 @@ For `BUILD_BASE_IMAGE`, use:
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
--build-arg CUDA_VERSION=13.4 \ --build-arg CUDA_VERSION=13.4 \
--build-arg BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.4" \ --build-arg BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.4" \
--build-arg FINAL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04" \ --build-arg FINAL_BASE_IMAGE="nvidia/cuda:13.4.1-base-ubuntu24.04" \
--build-arg NCCL_VERSION=2.32.3 \
. .
``` ```
@@ -490,9 +493,6 @@ For `BUILD_BASE_IMAGE`, use:
Keep the default explicit `torch_cuda_arch_list`. GPU-less BuildKit builds Keep the default explicit `torch_cuda_arch_list`. GPU-less BuildKit builds
cannot inspect the host GPU. R100 and VR200 report compute capability 10.7, cannot inspect the host GPU. R100 and VR200 report compute capability 10.7,
for which the generic `10.0` target provides family-compatible kernels. for which the generic `10.0` target provides family-compatible kernels.
The Ubuntu `devel` final image is also required: the corresponding `base`
image lacks the CUDA runtime/JIT package closure used by vLLM and the
prerelease PyTorch wheel.
`RUN_WHEEL_CHECK=false` disables only the PyPI publication-size guard for `RUN_WHEEL_CHECK=false` disables only the PyPI publication-size guard for
this private staging image. this private staging image.
+6 -6
View File
@@ -410,7 +410,6 @@ th {
| `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ | | `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ |
| `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ | | `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ |
| `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ | | `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ |
| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ | | `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ |
| `GraniteForCausalLM` | Granite 3.0, Granite 3.1, PowerLM | `ibm-granite/granite-3.0-2b-base`, `ibm-granite/granite-3.1-8b-instruct`, `ibm/PowerLM-3b`, etc. | ✅︎ | ✅︎ | | `GraniteForCausalLM` | Granite 3.0, Granite 3.1, PowerLM | `ibm-granite/granite-3.0-2b-base`, `ibm-granite/granite-3.1-8b-instruct`, `ibm/PowerLM-3b`, etc. | ✅︎ | ✅︎ |
| `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ | | `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ |
@@ -426,7 +425,6 @@ th {
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
| `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | | | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
| `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | | | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
| `K2HorizonForCausalLM` | K2-Horizon | `IFM/K2-Horizon-7B`, `IFM/K2-Horizon-375B` | | ✅︎ | | `K2HorizonForCausalLM` | K2-Horizon | `IFM/K2-Horizon-7B`, `IFM/K2-Horizon-375B` | | ✅︎ |
| `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ | | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
@@ -456,7 +454,6 @@ th {
| `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ | | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
| `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ | | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
| `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ | | `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ |
| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
| `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
| `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ | | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ |
@@ -468,7 +465,6 @@ th {
| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | | | `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ | | `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ | | `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ | | `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | | | `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ | | `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
@@ -484,16 +480,20 @@ Some models are supported only via the [Transformers modeling backend](#transfor
| ------------ | ------ | ----------------- | -------------------- | ------------------------- | | ------------ | ------ | ----------------- | -------------------- | ------------------------- |
| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ | | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
| `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | | ✅︎ | | `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | ✅︎ | ✅︎ |
| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ | | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | ✅︎ | ✅︎ |
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | ✅︎ | ✅︎ |
| `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ | | `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ |
| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
| `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ | | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |
| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ | | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ | | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
| `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ | | `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ |
@@ -27,6 +27,7 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
| `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios | | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
| `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios | | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
| `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads | | `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
| `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves, so no all-to-all runs in the framework | `--moe-backend flashinfer_moe_ep_*` (FlashInfer MoE-EP megakernels) |
| `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes | | `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
## Single Node Deployment ## Single Node Deployment
+4 -2
View File
@@ -125,8 +125,10 @@ Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out`
### Tokens IN <> Tokens OUT APIs ### Tokens IN <> Tokens OUT APIs
- `/inference/v1/generate` - Generate completions - [Generate API](token_in_token_out.md) (`/inference/v1/generate`)
- `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set) - Generate completions from token IDs, optionally with detokenized text (`output_mode`)
- `/inference/v1/abort_requests` - Abort in-flight requests (registered wherever `/inference/v1/generate` is)
- `/abort_requests` - Unauthenticated alias of `/inference/v1/abort_requests` (`--tokens-only` only)
### Renderer APIs ### Renderer APIs
+3 -1
View File
@@ -40,6 +40,8 @@ Seeding only reads the last few IDs (7 today), so a suffix of the prompt is enou
Each request wraps the engine's `GenerateResponse`(s) together with the caller metadata needed to reconstruct the final response without a GPU. Each request wraps the engine's `GenerateResponse`(s) together with the caller metadata needed to reconstruct the final response without a GPU.
Only `output_mode: "tokens"` responses are accepted which is the default. A response from a request with `output_mode: "text"` gets a 400 because its text is already detokenized with stop strings removed. Responses without `output_mode`, from older servers, are treated as `tokens`.
`/v1/chat/completions/derender`: `/v1/chat/completions/derender`:
??? code ??? code
@@ -127,7 +129,7 @@ Streaming derender is stateless on the server side: all mutable state lives in t
- `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings - `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings
- `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0 - `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0
When a streamed `GenerateResponseStreamChoice` carries `logprobs`, the `token_id:N` placeholders are resolved per chunk and the resolved logprobs are attached to the corresponding streamed choice — for chat as `ChatCompletionLogProbs`, for completions converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the `token_id:N` placeholders are resolved per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`.
## Example ## Example
@@ -0,0 +1,71 @@
# Tokens In <> Tokens Out API
`/inference/v1/generate` takes prompt token IDs and returns generated token IDs. It is the generation step of the `render → generate → derender` pipeline used in disaggregated serving (see the [Renderer APIs](renderer.md) and [Derenderer APIs](derenderer.md)).
The endpoint is registered on `vllm serve` with `--enable-scale-out`, and always with `--tokens-only`. `--tokens-only` also skips loading the tokenizer, so the server can't turn token IDs back into text.
## Output modes
The request field `output_mode` selects how much of the postprocessing the server does for the caller.
| `output_mode` | Response adds | Server needs |
| --- | --- | --- |
| `tokens` (default) | nothing, token IDs only | nothing |
| `text` | `text` on every choice and decoded tokens with `bytes` in `logprobs` | a tokenizer |
`text` returns text and token IDs in a single call, without a separate [derender](derenderer.md) hop per chunk. This suits latency sensitive streaming, where the derender hop lands on every output token.
Every response and every stream chunk echoes `output_mode`. Servers that predate the field ignore it and return token IDs only with a 200, so check that the `output_mode` in the response matches the one you sent. Older servers don't return it.
```python
import httpx
response = httpx.post(
"http://localhost:8000/inference/v1/generate",
json={
"token_ids": [151644, 872, 198],
"sampling_params": {"max_tokens": 32},
"output_mode": "text",
},
).json()
assert response["output_mode"] == "text"
print(response["choices"][0]["text"])
print(response["choices"][0]["token_ids"])
```
### Text
- `text` comes from the request's own detokenizer, so `skip_special_tokens`, `spaces_between_special_tokens`, `stop` and `include_stop_str_in_output` behave as they do on `/v1/completions`.
- A matched stop string is cut from `text` unless `include_stop_str_in_output` is set. `token_ids` keeps every generated token, so decoding `token_ids` yourself brings the stop string back.
- `/derender` only accepts `output_mode: "tokens"` responses and returns a 400 for text responses which are already detokenized.
- With `stream: true`, each choice's `text` is the delta since the previous chunk. A chunk is sent whenever the engine output carries new text or a `finish_reason`, even with no new token IDs. That covers text held back for stop string matching and the final output after an abort.
### Logprobs
With `output_mode: "tokens"`, logprob entries carry `token_id:N` placeholders. With `output_mode: "text"`, they carry the decoded token strings and their UTF-8 `bytes` for the sampled token and every entry in `top_logprobs`. A server started with `--return-tokens-as-token-ids` always returns the placeholders, as `/v1/completions` does.
### Errors
The server returns a 400 for:
- `output_mode: "text"` on a server without a tokenizer (`--tokens-only` or `--skip-tokenizer-init`).
- `output_mode: "text"` with `sampling_params.detokenize: false`.
- An unsupported `output_mode` value.
For using `output_mode` with separate prefill and decode pools, see [Disaggregated Prefilling](../../features/disagg_prefill.md#generate-api-output-modes).
## Aborting requests
`POST /inference/v1/abort_requests` aborts in-flight requests. It is registered wherever `/inference/v1/generate` is and requires the API key when `--api-key` is set, like `/inference/v1/generate`.
```bash
curl -X POST http://localhost:8000/inference/v1/abort_requests \
-H "Authorization: Bearer $VLLM_API_KEY" \
-H "Content-Type: application/json" \
-d '{"request_ids": ["generate-tokens-42"]}'
```
The request ID is the `request_id` the server returned: the `X-Request-Id` header or body `request_id` you sent, prefixed with `generate-tokens-`. A missing `request_ids` returns a 400. The response is empty and the abort finishes in the background.
With `--tokens-only`, the same handler is also served at `POST /abort_requests` for existing deployments. That path doesn't require the API key, even when `--api-key` is set. See [Security](../../usage/security.md#unprotected-endpoints-no-api-key-required).
+18 -16
View File
@@ -2,14 +2,14 @@
In on-policy distillation, the student records its top-K candidate tokens at In on-policy distillation, the student records its top-K candidate tokens at
every response position, and the teacher provides the log probabilities of every response position, and the teacher provides the log probabilities of
exactly those candidates as the distillation target. The caller passes the exactly those candidates as the distillation target. The caller passes one
candidates as `prompt_logprob_token_ids`; the teacher runs a single prefill candidate list per scored row as `prompt_logprob_token_ids`; the teacher runs a
over the prompt plus response and gathers their log probabilities on the GPU at single prefill over the prompt plus response and gathers their log
every scored row, starting from `prompt_logprob_start` (typically the prompt probabilities on the GPU, starting from `prompt_logprob_start` (typically the
length minus one, so only response positions are scored). The result is a prompt length minus one, so only response positions are scored). The result is
`[rows, K]` matrix whose row `i` scores prompt token a `[rows, K]` matrix whose row `i` scores prompt token
`prompt_logprob_start + i + 1` and whose columns follow the candidate order, so `prompt_logprob_start + i + 1` and whose columns follow that row's candidate
it lines up position by position with the student's top-K. order, so it lines up position by position with the student's top-K.
## Quick start ## Quick start
@@ -22,13 +22,13 @@ output = llm.generate(
TokensPrompt(prompt_token_ids=prompt_ids + response_ids), TokensPrompt(prompt_token_ids=prompt_ids + response_ids),
SamplingParams( SamplingParams(
max_tokens=1, max_tokens=1,
prompt_logprob_token_ids=candidate_ids, prompt_logprob_token_ids=student_topk_ids, # [len(response_ids), K]
prompt_logprob_start=len(prompt_ids) - 1, prompt_logprob_start=len(prompt_ids) - 1,
), ),
) )
scores = output[0].prompt_token_id_logprobs scores = output[0].prompt_token_id_logprobs
# scores.shape == (len(response_ids), len(candidate_ids)) # scores.shape == (len(response_ids), K)
# scores[i, j] = log p(candidate_ids[j] | prompt_ids + response_ids[:i]) # scores[i, j] = log p(student_topk_ids[i][j] | prompt_ids + response_ids[:i])
``` ```
The same parameters are accepted by the `/inference/v1/generate` HTTP endpoint The same parameters are accepted by the `/inference/v1/generate` HTTP endpoint
@@ -42,13 +42,17 @@ import numpy as np
scores = np.load(io.BytesIO(base64.b64decode(response["prompt_token_id_logprobs"]))) scores = np.load(io.BytesIO(base64.b64decode(response["prompt_token_id_logprobs"])))
``` ```
The matrix has `max(prompt_len - 1 - prompt_logprob_start, 0)` rows. With a `prompt_logprob_token_ids` is an integer array or a list of lists with exactly
logits `--logprobs-mode`, it holds logits instead of log probabilities. `prompt_len - 1 - prompt_logprob_start` rows; shorter rows are padded with `-1`,
and every `-1` entry scores `-inf`. Pass an `int32` NumPy array, such as the
student's top-K IDs, which is sent to the engine without conversion; nested
lists of Python ints are much slower to copy and serialize. With a logits
`--logprobs-mode`, the matrix holds logits instead of log probabilities.
## Requirements ## Requirements
- The V2 model runner (`VLLM_USE_V2_MODEL_RUNNER=1`). - The V2 model runner (`VLLM_USE_V2_MODEL_RUNNER=1`).
- `len(prompt_logprob_token_ids)` at most `--max-logprobs`. - The longest row at most `--max-logprobs` candidates.
- No `--kv-sharing-fast-prefill`. - No `--kv-sharing-fast-prefill`.
The request skips reading the prefix cache (local and KV connector), since The request skips reading the prefix cache (local and KV connector), since
@@ -56,8 +60,6 @@ cached rows have no logits to score.
## Limitations ## Limitations
- One candidate set per request, shared by every scored row; per-position
candidates are planned ([#56860](https://github.com/vllm-project/vllm/issues/56860)).
- Non-streaming only; `stream=true` requests are rejected. - Non-streaming only; `stream=true` requests are rejected.
- Not exposed through the OpenAI-compatible endpoints. - Not exposed through the OpenAI-compatible endpoints.
- If a prefill starts past the first scored row (e.g. `skip_reading_prefix_cache=False` - If a prefill starts past the first scored row (e.g. `skip_reading_prefix_cache=False`
+1 -1
View File
@@ -206,7 +206,7 @@ The following endpoints **do not require authentication** even when `--api-key`
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
- `/update_weights` - Update model weights (can alter model behavior) - `/update_weights` - Update model weights (can alter model behavior)
- `/get_world_size` - Get distributed world size - `/get_world_size` - Get distributed world size
- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`) - `/abort_requests` - Abort in-flight requests (available with `--tokens-only`. Use the authenticated `/inference/v1/abort_requests` otherwise)
**Utility endpoints:** **Utility endpoints:**
+1 -1
View File
@@ -197,7 +197,7 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
"examples/pooling/token_embed/*", "tests/models/language/pooling/*", "examples/pooling/token_embed/*", "tests/models/language/pooling/*",
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
"tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
"docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "docs/governance/process.md", "docs/community/reviewers.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
"tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*",
"rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**",
"rust/src/text/src/output/decoded.rs", "rust/src/text/src/output/decoded.rs",
+1 -1
View File
@@ -9,7 +9,7 @@ blake3
py-cpuinfo py-cpuinfo
transformers >= 5.16.1, < 5.18.0 transformers >= 5.16.1, < 5.18.0
huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export
tokenizers >= 0.22.0 # Required for fast incremental detokenization. tokenizers >= 0.23.2 # Required for fast incremental detokenization.
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*; platform_machine != "s390x" protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*; platform_machine != "s390x"
protobuf == 7.36.1; platform_machine == "s390x" protobuf == 7.36.1; platform_machine == "s390x"
+1 -1
View File
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.3
instanttensor >= 0.1.9 instanttensor >= 0.1.9
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation) # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl[cu13]==4.7.1 nvidia-cutlass-dsl[cu13]==4.8.0
quack-kernels==0.6.5 # Required by MSA quack-kernels==0.6.5 # Required by MSA
# Tokenspeed_MLA for faster mla with spec decode # Tokenspeed_MLA for faster mla with spec decode
+4 -15
View File
@@ -3,24 +3,13 @@
# TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels # TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels
# (vllm-project/vllm#49735). # (vllm-project/vllm#49735).
--extra-index-url https://download.pytorch.org/whl/nightly/cu134 --extra-index-url https://download.pytorch.org/whl/nightly/cu134
--extra-index-url https://flashinfer.ai/whl/cu134
--extra-index-url https://flashinfer.ai/whl
# Match each architecture's torch pin to the latest torchvision wheel. # Match each architecture's torch pin to the latest torchvision wheel.
torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64" torch==2.15.0.dev20260928+cu134; platform_machine == "x86_64"
torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64" torch==2.15.0.dev20260926+cu134; platform_machine == "aarch64"
torchvision==0.30.0.dev20260909+cu134 torchvision==0.30.0.dev20260928+cu134
torchaudio==2.11.0.dev20260911+cu134 torchaudio==2.11.0.dev20260928+cu134
# docker/Dockerfile strips these packages from requirements/cuda.txt whenever
# INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to
# still get installed.
nvidia-cutlass-dsl[cu13]==4.8.0.dev0
flashinfer-python==0.7.0.post1
quack-kernels==0.6.5
cuda-python==13.4.1 cuda-python==13.4.1
cuda-bindings==13.4.1 cuda-bindings==13.4.1
flashinfer-jit-cache==0.7.0.post1+cu134
flashinfer-cubin==0.7.0.post1
nixl==1.4.1 nixl==1.4.1
+1 -1
View File
@@ -1165,7 +1165,7 @@ timm==1.0.17
# -r requirements/test/cuda.in # -r requirements/test/cuda.in
# open-clip-torch # open-clip-torch
# segmentation-models-pytorch # segmentation-models-pytorch
tokenizers==0.23.1 tokenizers==0.23.2
# via # via
# -r requirements/test/../common.txt # -r requirements/test/../common.txt
# cohere # cohere
+1 -1
View File
@@ -1266,7 +1266,7 @@ timm==1.0.17
# -r requirements/test/cuda.in # -r requirements/test/cuda.in
# open-clip-torch # open-clip-torch
# segmentation-models-pytorch # segmentation-models-pytorch
tokenizers==0.23.1 tokenizers==0.23.2
# via # via
# -c requirements/common.txt # -c requirements/common.txt
# -r requirements/test/../common.txt # -r requirements/test/../common.txt
+1 -1
View File
@@ -1249,7 +1249,7 @@ timm==1.0.17
# -r requirements/test/rocm.in # -r requirements/test/rocm.in
# open-clip-torch # open-clip-torch
# segmentation-models-pytorch # segmentation-models-pytorch
tokenizers==0.23.1 tokenizers==0.23.2
# via # via
# -c requirements/common.txt # -c requirements/common.txt
# -r requirements/test/../common.txt # -r requirements/test/../common.txt
+1 -1
View File
@@ -1089,7 +1089,7 @@ timm==1.0.28
# via # via
# -r requirements/test/xpu.in # -r requirements/test/xpu.in
# open-clip-torch # open-clip-torch
tokenizers==0.23.1 tokenizers==0.23.2
# via # via
# -c requirements/common.txt # -c requirements/common.txt
# -r requirements/test/../common.txt # -r requirements/test/../common.txt
+1 -1
View File
@@ -12,4 +12,4 @@ ray[data]
setuptools==78.1.0 setuptools==78.1.0
setuptools-rust>=1.9.0 setuptools-rust>=1.9.0
nixl==0.3.0 nixl==0.3.0
tpu-inference==0.29.0 tpu-inference==0.30.0
-4
View File
@@ -1,4 +0,0 @@
# CLAUDE.md
First, check @AGENTS.override.md if exists.
Then, follow instructions in @AGENTS.md.
-1
View File
@@ -1 +0,0 @@
@AGENTS.md
@@ -152,6 +152,11 @@ impl EngineCoreRequest {
feature.extract_aux_frames(&mut aux_frames, threshold); feature.extract_aux_frames(&mut aux_frames, threshold);
} }
} }
if let Some(params) = &mut self.sampling_params
&& let Some(ids) = &mut params.prompt_logprob_token_ids
{
ids.extract_aux_frame(&mut aux_frames, threshold);
}
aux_frames aux_frames
} }
} }
@@ -168,7 +173,7 @@ mod tests {
PlaceholderRange, PlaceholderRange,
}; };
use crate::protocol::sampling::EngineCoreSamplingParams; use crate::protocol::sampling::EngineCoreSamplingParams;
use crate::protocol::tensor::{WireArrayData, WireTensor}; use crate::protocol::tensor::{WireArrayData, WireNdArray, WireTensor};
use crate::protocol::{decode_value, encode_msgpack}; use crate::protocol::{decode_value, encode_msgpack};
const AUX_FRAME_THRESHOLD: usize = 256; const AUX_FRAME_THRESHOLD: usize = 256;
@@ -294,4 +299,27 @@ mod tests {
); );
assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty()); assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty());
} }
#[test]
fn engine_core_request_extracts_large_prompt_logprob_token_ids() {
let ids = WireNdArray::from_i32(vec![AUX_FRAME_THRESHOLD, 1], vec![5; AUX_FRAME_THRESHOLD])
.unwrap();
let mut request = EngineCoreRequest {
sampling_params: Some(EngineCoreSamplingParams {
prompt_logprob_token_ids: Some(ids),
..EngineCoreSamplingParams::for_test()
}),
..EngineCoreRequest::default()
};
let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD);
assert_eq!(aux_frames.len(), 1);
assert_eq!(aux_frames[0].len(), AUX_FRAME_THRESHOLD * 4);
let params = request.sampling_params.unwrap();
assert_eq!(
params.prompt_logprob_token_ids.unwrap().data,
WireArrayData::AuxIndex(1)
);
}
} }
@@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize};
use serde_default::DefaultFromSerde; use serde_default::DefaultFromSerde;
use crate::protocol::structured_outputs::StructuredOutputsParams; use crate::protocol::structured_outputs::StructuredOutputsParams;
use crate::protocol::tensor::WireNdArray;
fn default_top_p() -> f32 { fn default_top_p() -> f32 {
1.0 1.0
@@ -101,9 +102,9 @@ pub struct EngineCoreSamplingParams {
/// ///
/// `None` disables prompt logprobs. `-1` requests the full vocabulary. /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
pub prompt_logprobs: Option<i32>, pub prompt_logprobs: Option<i32>,
/// Candidate token IDs scored at every scored causal prompt row, where row /// `[num_rows, num_ids]` candidate token IDs per scored causal prompt row,
/// `i` scores them as predictions of prompt token `i + 1`. /// `-1` padding shorter rows.
pub prompt_logprob_token_ids: Option<Vec<u32>>, pub prompt_logprob_token_ids: Option<WireNdArray>,
/// First causal prompt row to score; `None` scores from the first row. /// First causal prompt row to score; `None` scores from the first row.
pub prompt_logprob_start: Option<u32>, pub prompt_logprob_start: Option<u32>,
/// Minimum probability threshold for token sampling. /// Minimum probability threshold for token sampling.
@@ -121,6 +121,18 @@ impl<Dtype: From<TensorDtype> + Serialize + DeserializeOwned> WireArray<Dtype> {
)) ))
} }
/// Build an int32 tensor/ndarray backed by native-endian raw-view bytes.
///
/// Takes ownership of the backing buffer without copying its data.
pub fn from_i32(shape: Vec<usize>, data: Vec<i32>) -> Result<Self, String> {
validate_element_count(&shape, data.len())?;
Ok(Self::from_raw_bytes(
TensorDtype::I32,
shape,
bytes_from_pod_vec(data),
))
}
/// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes. /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
/// ///
/// Takes ownership of the backing buffer without copying its data. /// Takes ownership of the backing buffer without copying its data.
@@ -181,7 +181,7 @@ mod tests {
let request: GenerateRequest = serde_json::from_value(json!({ let request: GenerateRequest = serde_json::from_value(json!({
"token_ids": [11, 22], "token_ids": [11, 22],
"stream": true, "stream": true,
"sampling_params": {"prompt_logprob_token_ids": [1, 2]} "sampling_params": {"prompt_logprob_token_ids": [[1, 2]]}
})) }))
.expect("parse request"); .expect("parse request");
assert!(validate_request_compat(&request, &served(&["test-model"])).is_err()); assert!(validate_request_compat(&request, &served(&["test-model"])).is_err());
+33 -16
View File
@@ -7,7 +7,7 @@ pub(crate) mod logprobs;
pub(crate) mod sampling; pub(crate) mod sampling;
pub(crate) mod token_ids; pub(crate) mod token_ids;
use logprobs::validate_logprobs; use logprobs::{lower_prompt_logprob_token_ids, validate_logprobs};
use sampling::validate_resolved_sampling_params; use sampling::validate_resolved_sampling_params;
use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use token_ids::{validate_prompt_token_ids, validate_vocab_range};
use vllm_engine_core_client::protocol::sampling::{ use vllm_engine_core_client::protocol::sampling::{
@@ -50,7 +50,11 @@ pub fn lower_text_request(
// multimodal tensor payloads. // multimodal tensor payloads.
mm_features: request.mm_features.take(), mm_features: request.mm_features.take(),
sampling_params: lower_sampling_params( sampling_params: lower_sampling_params(
request.sampling_params.clone(), SamplingParams {
// Move the candidate table rather than cloning it; only lowering reads it.
prompt_logprob_token_ids: request.sampling_params.prompt_logprob_token_ids.take(),
..request.sampling_params.clone()
},
sampling_hints, sampling_hints,
sampling_limits, sampling_limits,
prompt_len, prompt_len,
@@ -180,7 +184,12 @@ pub fn lower_sampling_params(
thinking_token_budget, thinking_token_budget,
logprobs, logprobs,
prompt_logprobs, prompt_logprobs,
prompt_logprob_token_ids, prompt_logprob_token_ids: lower_prompt_logprob_token_ids(
prompt_logprob_token_ids,
prompt_logprob_start,
prompt_len,
sampling_limits.model_vocab_size,
)?,
prompt_logprob_start, prompt_logprob_start,
min_p, min_p,
frequency_penalty, frequency_penalty,
@@ -332,6 +341,7 @@ mod tests {
use vllm_engine_core_client::protocol::multimodal::{ use vllm_engine_core_client::protocol::multimodal::{
MmFeatureSpec, MmModality, PlaceholderRange, MmFeatureSpec, MmModality, PlaceholderRange,
}; };
use vllm_engine_core_client::protocol::tensor::WireNdArray;
use vllm_tokenizer::test_utils::TestTokenizer; use vllm_tokenizer::test_utils::TestTokenizer;
use super::*; use super::*;
@@ -1140,7 +1150,7 @@ mod tests {
#[test] #[test]
fn lower_sampling_params_validates_prompt_logprob_token_ids() { fn lower_sampling_params_validates_prompt_logprob_token_ids() {
let lower = |ids: Option<Vec<u32>>, start: Option<u32>| { let lower = |ids: Option<Vec<Vec<i32>>>, start: Option<u32>| {
lower_sampling_params_with_limits( lower_sampling_params_with_limits(
SamplingParams { SamplingParams {
prompt_logprob_token_ids: ids, prompt_logprob_token_ids: ids,
@@ -1155,23 +1165,30 @@ mod tests {
Ok(_) => panic!("expected rejection"), Ok(_) => panic!("expected rejection"),
}; };
let params = lower(Some(vec![1, 2]), Some(1)).unwrap(); let params = lower(Some(vec![vec![1, 2], vec![3]]), None).unwrap();
assert_eq!(params.prompt_logprob_token_ids, Some(vec![1, 2])); assert_eq!(
assert_eq!(params.prompt_logprob_start, Some(1)); params.prompt_logprob_token_ids,
Some(WireNdArray::from_i32(vec![2, 2], vec![1, 2, 3, -1]).unwrap())
);
assert_eq!(
lower(Some(vec![vec![5]]), Some(1)).unwrap().prompt_logprob_start,
Some(1)
);
expect_test::expect![[r#" expect_test::expect![[r#"
[ [
"prompt_logprob_token_ids must not be empty.", "prompt_logprob_token_ids must be a non-empty integer array of shape [num_rows, num_ids].",
"requested prompt_logprob_token_ids of 21, which is greater than max allowed: 20", "requested prompt_logprob_token_ids of 21, which is greater than max allowed: 20",
"token_id(s) [1000] in prompt_logprob_token_ids are out of vocabulary. Vocabulary size: 1000", "prompt_logprob_token_ids contain out-of-vocab token ids (-1 pads a row). Vocabulary size: 1000",
"prompt_logprob_token_ids must not contain duplicates.", "prompt_logprob_token_ids contain out-of-vocab token ids (-1 pads a row). Vocabulary size: 1000",
"prompt_logprob_token_ids has 1 rows, but the prompt has 2 scored rows (prompt_len - 1 - prompt_logprob_start).",
"prompt_logprob_start requires prompt_logprob_token_ids.", "prompt_logprob_start requires prompt_logprob_token_ids.",
] ]
"#]] "#]].assert_debug_eq(&[
.assert_debug_eq(&[ rejects(Some(vec![vec![]]), None),
rejects(Some(vec![]), None), rejects(Some(vec![(0..21).collect(), vec![0]]), None),
rejects(Some((0..21).collect()), None), rejects(Some(vec![vec![1000], vec![0]]), None),
rejects(Some(vec![1000]), None), rejects(Some(vec![vec![-2], vec![0]]), None),
rejects(Some(vec![1, 1]), None), rejects(Some(vec![vec![1]]), None),
rejects(None, Some(0)), rejects(None, Some(0)),
]); ]);
} }
+55 -12
View File
@@ -3,11 +3,11 @@
//! Python-compatible validation for logprobs sampling params. //! Python-compatible validation for logprobs sampling params.
//! //!
//! `-1` is expanded only for bounds checks. The original request values are //! `-1` counts are expanded only for bounds checks and passed to engine-core
//! passed through to engine-core. //! unchanged; `prompt_logprob_token_ids` rows are padded into an array.
use itertools::Itertools as _;
use thiserror::Error; use thiserror::Error;
use vllm_engine_core_client::protocol::tensor::WireNdArray;
use crate::backend::SamplingLimits; use crate::backend::SamplingLimits;
@@ -36,10 +36,21 @@ pub enum LogprobsError {
len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}." len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
)] )]
TokenIdsMismatch { logprobs: i32, num_token_ids: usize }, TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
#[error("prompt_logprob_token_ids must not be empty.")] #[error(
"prompt_logprob_token_ids must be a non-empty integer array of shape \
[num_rows, num_ids]."
)]
EmptyPromptLogprobTokenIds, EmptyPromptLogprobTokenIds,
#[error("prompt_logprob_token_ids must not contain duplicates.")] #[error(
DuplicatePromptLogprobTokenIds, "prompt_logprob_token_ids contain out-of-vocab token ids (-1 pads a row). \
Vocabulary size: {vocab_size}"
)]
PromptLogprobTokenIdsOutOfVocab { vocab_size: usize },
#[error(
"prompt_logprob_token_ids has {rows} rows, but the prompt has {expected} scored \
rows (prompt_len - 1 - prompt_logprob_start)."
)]
PromptLogprobRowsMismatch { rows: usize, expected: usize },
#[error("prompt_logprob_start requires prompt_logprob_token_ids.")] #[error("prompt_logprob_start requires prompt_logprob_token_ids.")]
PromptLogprobStartWithoutTokenIds, PromptLogprobStartWithoutTokenIds,
} }
@@ -49,7 +60,7 @@ pub(super) fn validate_logprobs(
logprobs: Option<i32>, logprobs: Option<i32>,
prompt_logprobs: Option<i32>, prompt_logprobs: Option<i32>,
logprob_token_ids: Option<&[u32]>, logprob_token_ids: Option<&[u32]>,
prompt_logprob_token_ids: Option<&[u32]>, prompt_logprob_token_ids: Option<&[Vec<i32>]>,
prompt_logprob_start: Option<u32>, prompt_logprob_start: Option<u32>,
sampling_limits: SamplingLimits, sampling_limits: SamplingLimits,
) -> Result<(), LogprobsError> { ) -> Result<(), LogprobsError> {
@@ -60,7 +71,7 @@ pub(super) fn validate_logprobs(
validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?; validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?; validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
validate_logprobs_count( validate_logprobs_count(
prompt_logprob_token_ids.map(|ids| ids.len() as i32), prompt_logprob_token_ids.map(|rows| rows.iter().map(Vec::len).max().unwrap_or(0) as i32),
max_logprobs, max_logprobs,
vocab_size, vocab_size,
"prompt_logprob_token_ids", "prompt_logprob_token_ids",
@@ -70,17 +81,49 @@ pub(super) fn validate_logprobs(
} }
fn validate_prompt_logprob_token_ids( fn validate_prompt_logprob_token_ids(
token_ids: Option<&[u32]>, rows: Option<&[Vec<i32>]>,
start: Option<u32>, start: Option<u32>,
) -> Result<(), LogprobsError> { ) -> Result<(), LogprobsError> {
match token_ids { match rows {
Some([]) => Err(LogprobsError::EmptyPromptLogprobTokenIds), Some(rows) if rows.iter().all(Vec::is_empty) => {
Some(ids) if !ids.iter().all_unique() => Err(LogprobsError::DuplicatePromptLogprobTokenIds), Err(LogprobsError::EmptyPromptLogprobTokenIds)
}
None if start.is_some() => Err(LogprobsError::PromptLogprobStartWithoutTokenIds), None if start.is_some() => Err(LogprobsError::PromptLogprobStartWithoutTokenIds),
_ => Ok(()), _ => Ok(()),
} }
} }
/// Check per-row candidates against the vocabulary and the prompt as Python
/// does, and pad them with `-1` into the engine's `[num_rows, num_ids]` array.
pub(super) fn lower_prompt_logprob_token_ids(
rows: Option<Vec<Vec<i32>>>,
start: Option<u32>,
prompt_len: u32,
vocab_size: usize,
) -> Result<Option<WireNdArray>, LogprobsError> {
let Some(rows) = rows else {
return Ok(None);
};
if rows.iter().flatten().any(|&id| id < -1 || i64::from(id) >= vocab_size as i64) {
return Err(LogprobsError::PromptLogprobTokenIdsOutOfVocab { vocab_size });
}
let expected = (prompt_len as usize).saturating_sub(1 + start.unwrap_or(0) as usize);
if rows.len() != expected {
return Err(LogprobsError::PromptLogprobRowsMismatch {
rows: rows.len(),
expected,
});
}
let width = rows.iter().map(Vec::len).max().unwrap_or(0);
let data = rows
.iter()
.flat_map(|row| row.iter().copied().chain(std::iter::repeat_n(-1, width - row.len())))
.collect();
Ok(Some(
WireNdArray::from_i32(vec![rows.len(), width], data).expect("padded rows are rectangular"),
))
}
fn validate_logprobs_count( fn validate_logprobs_count(
requested: Option<i32>, requested: Option<i32>,
max_logprobs: usize, max_logprobs: usize,
-8
View File
@@ -95,14 +95,6 @@ pub(crate) fn validate_vocab_range(
)?; )?;
} }
if let Some(token_ids) = params.prompt_logprob_token_ids.as_deref() {
validate_param(
"prompt_logprob_token_ids",
token_ids.iter().copied(),
limits.model_vocab_size,
)?;
}
if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() { if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() {
validate_param( validate_param(
"bad_words", "bad_words",
+3 -3
View File
@@ -97,9 +97,9 @@ pub struct SamplingParams {
/// ///
/// `None` disables prompt logprobs. `-1` requests the full vocabulary. /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
pub prompt_logprobs: Option<i32>, pub prompt_logprobs: Option<i32>,
/// Candidate token IDs scored at every scored causal prompt row, where row /// Candidate token IDs per scored causal prompt row, where row `i` scores
/// `i` scores them as predictions of prompt token `i + 1`. /// its IDs as predictions of prompt token `prompt_logprob_start + i + 1`.
pub prompt_logprob_token_ids: Option<Vec<u32>>, pub prompt_logprob_token_ids: Option<Vec<Vec<i32>>>,
/// First causal prompt row to score; `None` scores from the first row. /// First causal prompt row to score; `None` scores from the first row.
pub prompt_logprob_start: Option<u32>, pub prompt_logprob_start: Option<u32>,
/// Minimum probability threshold for token sampling. `None` means no /// Minimum probability threshold for token sampling. `None` means no
@@ -327,6 +327,35 @@ def test_reentered_tag(free_x_early):
allocator.release_pools() allocator.release_pools()
@pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"])
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
def test_runtime_state_survives_sleep(level):
"""Runtime state is offloaded at both levels and restored by the first
selective wake; a nested kv_cache pool keeps its own tag."""
from vllm.device_allocator.sleep_mode_backend import CuMemBackend
allocator = get_mem_allocator_instance()
with allocator.use_memory_pool("runtime"):
state = torch.arange(1 << 20, device=DEVICE_TYPE)
with allocator.use_memory_pool("kv_cache"):
kv = torch.zeros(1 << 20, device=DEVICE_TYPE)
assert {d.tag for d in allocator.pointer_to_data.values()} == {
"runtime",
"kv_cache",
}
backend = CuMemBackend()
backend.suspend(level=level)
assert mapped_usage(allocator) == 0
backend.resume(tags=["weights"])
assert torch.equal(state, torch.arange(1 << 20, device=DEVICE_TYPE))
backend.resume(tags=["kv_cache"])
kv.fill_(1)
assert int(kv.sum()) == kv.numel()
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator") @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
def test_level2_discards_ordinary_tensor_with_weights_tag(): def test_level2_discards_ordinary_tensor_with_weights_tag():
+42 -27
View File
@@ -6,14 +6,17 @@ from collections import defaultdict
import pytest import pytest
import regex as re import regex as re
from vllm import LLM, SamplingParams from vllm import SamplingParams
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches
def run_model( def run_model(
compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs vllm_runner,
compile_config: CompilationMode | CompilationConfig,
model: str,
**model_kwargs,
): ):
"""Run a model with the given compilation config for E2E fusion tests.""" """Run a model with the given compilation config for E2E fusion tests."""
compilation_config = ( compilation_config = (
@@ -36,38 +39,48 @@ def run_model(
# No cudagraphs by default # No cudagraphs by default
if compilation_config.cudagraph_mode is None: if compilation_config.cudagraph_mode is None:
compilation_config.cudagraph_mode = CUDAGraphMode.NONE compilation_config.cudagraph_mode = CUDAGraphMode.NONE
llm = LLM(
model=model, # trust_remote_code/enable_chunked_prefill/block_size/kernel_config are
# pinned to LLM()'s own defaults since VllmRunner overrides them for
# other test needs (e.g. it forces block_size=16, which some attention
# backends like FLASHINFER_MLA reject, and disables JIT warmup).
model_kwargs.setdefault("kernel_config", {"enable_jit_warmup": True})
with vllm_runner(
model,
compilation_config=compilation_config, compilation_config=compilation_config,
trust_remote_code=False,
enable_chunked_prefill=None,
block_size=None,
**model_kwargs, **model_kwargs,
) ) as vllm_model:
outputs = llm.generate(prompts, sampling_params) llm = vllm_model.llm
outputs = llm.generate(prompts, sampling_params)
# Print the outputs. # Print the outputs.
for output in outputs: for output in outputs:
prompt = output.prompt prompt = output.prompt
generated_text = output.outputs[0].text generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
# Get the compile ranges endpoints after vllm config post init # Get the compile ranges endpoints after vllm config post init
# in order to compute compile ranges correctly # in order to compute compile ranges correctly
compilation_config.compile_ranges_endpoints = ( compilation_config.compile_ranges_endpoints = (
llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints
) )
# Fetch match table from each worker via RPC and sum across workers. # Fetch match table from each worker via RPC and sum across workers.
worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc( worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc(
"get_compilation_match_table" "get_compilation_match_table"
) )
combined: defaultdict[str, int] = defaultdict(int) combined: defaultdict[str, int] = defaultdict(int)
for table in worker_tables: for table in worker_tables:
for k, v in table.items(): for k, v in table.items():
combined[k] += v combined[k] += v
return dict(combined) return dict(combined)
@pytest.fixture @pytest.fixture
def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn, vllm_runner):
def run( def run(
model_name: str, model_name: str,
matches: Matches, matches: Matches,
@@ -137,7 +150,9 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
) )
with caplog_mp_spawn(logging.DEBUG) as log_holder: with caplog_mp_spawn(logging.DEBUG) as log_holder:
match_table = run_model(full_compilation_config, model_name, **model_kwargs) match_table = run_model(
vllm_runner, full_compilation_config, model_name, **model_kwargs
)
num_compile_ranges = len(full_compilation_config.get_compile_ranges()) num_compile_ranges = len(full_compilation_config.get_compile_ranges())
assert num_compile_ranges in [1, 2, 3] assert num_compile_ranges in [1, 2, 3]
+14
View File
@@ -101,6 +101,20 @@ llama3_8b_fp4 = ModelFusionInfo(
), ),
) )
# ModelOpt MIXED_PRECISION: FP8 attention projections, NVFP4 MLPs. It resolves to
# modelopt_mixed, not modelopt_fp4, and its MLP has no manual act+quant fusion, so
# every MLP boundary is left to ActivationQuantFusionPass.
qwen3_8_27b_mixed_fp4 = ModelFusionInfo(
model_name="nvidia/Qwen3.8-27B-NVFP4",
hf_overrides=lambda n_layers: {
"text_config": {
"num_hidden_layers": n_layers,
"layer_types": ["linear_attention"] * (n_layers - 1) + ["full_attention"],
}
},
matches=lambda n_layers: Matches(act_quant_fusion=n_layers),
)
# MoEs cannot do act+quant fusion because those ops are hidden from torch.compile. # MoEs cannot do act+quant fusion because those ops are hidden from torch.compile.
# MoEs also only expose 1 rms+quant fusion because the quant for up_proj is hidden. # MoEs also only expose 1 rms+quant fusion because the quant for up_proj is hidden.
# TODO(luka): https://github.com/vllm-project/vllm/issues/31985 # TODO(luka): https://github.com/vllm-project/vllm/issues/31985
@@ -30,6 +30,7 @@ from .models import (
llama3_8b_fp8, llama3_8b_fp8,
llama4_scout_fp4, llama4_scout_fp4,
llama4_scout_fp8, llama4_scout_fp8,
qwen3_8_27b_mixed_fp4,
qwen3_a3b_fp8, qwen3_a3b_fp8,
) )
@@ -202,3 +203,46 @@ def test_tp1_fp4_fusions(
compilation_config, compilation_config,
matches_check, matches_check,
) )
@pytest.mark.parametrize(
"model_name, matches_fn, model_kwargs, hf_overrides", [qwen3_8_27b_mixed_fp4]
)
@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
@pytest.mark.parametrize("n_layers", [4])
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
def test_tp1_fp4_mixed_precision_act_fusion(
model_name: str,
matches_fn: Callable[[int], Matches],
model_kwargs: dict,
hf_overrides: Callable[[int], dict],
attn_backend: AttentionBackendCase,
n_layers: int,
inductor_graph_partition: bool,
run_e2e_fusion_test,
):
"""Act+quant fusion matches on a MIXED_PRECISION NVFP4 checkpoint."""
model_kwargs["hf_overrides"] = hf_overrides(n_layers)
model_kwargs["load_format"] = "dummy"
model_kwargs["max_model_len"] = 1024
model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
compilation_config = dict(
use_inductor_graph_partition=inductor_graph_partition,
pass_config=PassConfig(
fuse_norm_quant=False,
fuse_act_quant=True,
fuse_attn_quant=False,
enable_qk_norm_rope_fusion=False,
),
)
run_e2e_fusion_test(
model_name,
matches_fn(n_layers),
model_kwargs,
attn_backend,
compilation_config,
["act_quant_fusion"],
)
+43
View File
@@ -0,0 +1,43 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NVFP4 detection for ModelOpt MIXED_PRECISION checkpoints.
A checkpoint that quantizes some layers to NVFP4 and others to FP8 declares
``quant_algo: MIXED_PRECISION`` and resolves to ``modelopt_mixed``, so NVFP4
has to be detected from the per-layer algorithms.
"""
import pytest
from vllm.config.model import _modelopt_mixed_has_nvfp4
def _quantized_layers(*algos: str) -> dict:
return {
"quantized_layers": {
f"model.layers.{i}.proj": {"quant_algo": algo}
for i, algo in enumerate(algos)
}
}
def test_detects_nvfp4_alongside_another_algorithm():
"""The shipped shape: FP8 attention projections, NVFP4 MLP."""
assert _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "NVFP4", "FP8"))
def test_ignores_a_checkpoint_with_no_nvfp4_layer():
assert not _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "MXFP8"))
def test_weight_only_nvfp4_does_not_count():
"""W4A16_NVFP4 leaves activations in high precision, so nothing to fuse."""
assert not _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "W4A16_NVFP4"))
@pytest.mark.parametrize(
"quant_config",
[None, {}, {"quantized_layers": None}, {"quantized_layers": ["NVFP4"]}],
)
def test_absent_or_malformed_config_is_not_nvfp4(quant_config):
assert not _modelopt_mixed_has_nvfp4(quant_config)
+1 -1
View File
@@ -511,7 +511,7 @@ def _check_dummy_hash_model_forward(
model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model) model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model)
torch.nn.Module.__init__(model) torch.nn.Module.__init__(model)
model.use_mega_moe = False model.use_native_mega_moe = False
model.use_sequence_parallel = False model.use_sequence_parallel = False
model.fuse_mhc_all_reduce = False model.fuse_mhc_all_reduce = False
model.engram_hash = state model.engram_hash = state
@@ -0,0 +1,278 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import numpy as np
import pytest
import torch
from vllm.config.parallel import EPLBConfig
from vllm.distributed.eplb.eplb_communicator import EplbCommunicator
from vllm.distributed.eplb.migration_scheduler import (
MigrationFlow,
schedule_migration_batches,
schedule_migration_batches_for_layers,
)
from vllm.distributed.eplb.rebalance_execute import move_from_buffer, move_to_buffer
class _MockEplbCommunicator(EplbCommunicator):
def __init__(self) -> None:
self.send_calls: list[tuple[list[torch.Tensor], int, int]] = []
self.recv_calls: list[tuple[list[torch.Tensor], int, int]] = []
self.context_calls = 0
self.execute_count = 0
def add_send(
self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int
) -> None:
self.send_calls.append((tensors, dst_rank, expert_id))
def add_recv(
self, tensors: list[torch.Tensor], src_rank: int, expert_id: int
) -> None:
self.recv_calls.append((tensors, src_rank, expert_id))
def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
self.context_calls += 1
def execute(self) -> None:
self.execute_count += 1
def test_migration_batching_is_disabled_by_default_and_can_be_enabled() -> None:
assert not EPLBConfig().enable_migration_batching
assert EPLBConfig(enable_migration_batching=True).enable_migration_batching
def test_migration_batching_only_applies_to_async_eplb() -> None:
assert not EPLBConfig(use_async=True).migration_batching_enabled
assert EPLBConfig(
use_async=True, enable_migration_batching=True
).migration_batching_enabled
assert not EPLBConfig(
use_async=True, enable_migration_batching=False
).migration_batching_enabled
assert not EPLBConfig(
use_async=False, enable_migration_batching=True
).migration_batching_enabled
def test_schedule_migration_batches_is_deterministic() -> None:
transfers = [
(1, 3, 0),
(2, 4, 1),
(0, 1, 2),
(0, 2, 3),
(0, 3, 4),
(0, 4, 5),
]
num_local_experts, old_indices, new_indices = _placements_for_transfers(
5, transfers
)
batches = schedule_migration_batches(num_local_experts, old_indices, new_indices)
assert [_endpoints(batch) for batch in batches] == [
[(1, 3), (2, 4)],
[(0, 1)],
[(0, 2)],
[(0, 3)],
[(0, 4)],
]
for batch in batches:
_assert_no_endpoint_conflict(batch)
def test_schedule_migration_batches_coalesces_rank_pair() -> None:
transfers = [
(0, 1, 0),
(0, 1, 1),
(2, 3, 2),
]
num_local_experts, old_indices, new_indices = _placements_for_transfers(
4, transfers
)
assert schedule_migration_batches(num_local_experts, old_indices, new_indices) == [
[
MigrationFlow(0, 1, expert_ids=(0, 1)),
MigrationFlow(2, 3, expert_ids=(2,)),
]
]
def test_schedule_migration_batches_covers_random_instructions() -> None:
rng = np.random.default_rng(42)
transfers = []
for expert_id in range(80):
src_rank, dst_rank = rng.integers(0, 8, size=2).tolist()
if src_rank != dst_rank:
transfers.append((src_rank, dst_rank, expert_id))
num_local_experts, old_indices, new_indices = _placements_for_transfers(
8, transfers
)
batches = schedule_migration_batches(num_local_experts, old_indices, new_indices)
scheduled = _flatten_transfers(batches)
assert len(scheduled) == len(transfers)
assert set(scheduled) == set(transfers)
for batch in batches:
_assert_no_endpoint_conflict(batch)
def test_schedule_migration_batches_excludes_local_copies() -> None:
# Old: rank 0 [0, 1], rank 1 [1, 2]. New: rank 1 needs expert 0.
old_indices = np.array([0, 1, 1, 2], dtype=np.int64)
new_indices = np.array([0, 1, 0, 2], dtype=np.int64)
assert schedule_migration_batches(2, old_indices, new_indices) == [
[MigrationFlow(0, 1, expert_ids=(0,))]
]
def test_schedule_migration_batches_requires_matching_shapes() -> None:
old_indices = np.array([0, 1], dtype=np.int64)
new_indices = np.array([0], dtype=np.int64)
with pytest.raises(AssertionError):
schedule_migration_batches(1, old_indices, new_indices)
def test_schedule_migration_batches_balances_replicas() -> None:
# Rank 0 and rank 1 hold expert 0; ranks 2, 3, and 4 need it.
old_indices = np.array([0, 0, 1, 1, 1], dtype=np.int64)
new_indices = np.array([0, 0, 0, 0, 0], dtype=np.int64)
batches = schedule_migration_batches(1, old_indices, new_indices)
assert set(_flatten_transfers(batches)) == {
(0, 2, 0),
(0, 4, 0),
(1, 3, 0),
}
def test_schedule_migration_batches_for_layers_matches_individual_calls() -> None:
old_indices = np.array([[0, 1, 2, 3], [3, 2, 1, 0]], dtype=np.int64)
new_indices = np.array([[1, 2, 3, 0], [0, 3, 2, 1]], dtype=np.int64)
assert schedule_migration_batches_for_layers(1, old_indices, new_indices) == [
schedule_migration_batches(1, old_indices[0], new_indices[0]),
schedule_migration_batches(1, old_indices[1], new_indices[1]),
]
@pytest.mark.parametrize("use_precomputed", [False, True])
def test_move_to_buffer_uses_multiple_batches(use_precomputed: bool) -> None:
old_indices = np.array([0, 1, 2, 3], dtype=np.int64)
new_indices = np.array([1, 2, 3, 0], dtype=np.int64)
communicator = _MockEplbCommunicator()
migration_batches = (
schedule_migration_batches(1, old_indices, new_indices)
if use_precomputed
else None
)
move_to_buffer(
num_local_experts=1,
old_indices=old_indices,
new_indices=new_indices,
expert_weights=[torch.zeros(1, 1)],
expert_weights_buffers=[torch.zeros(1, 1)],
stream=None,
ep_rank=0,
communicator=communicator,
layer_idx=7,
enable_migration_batching=True,
migration_batches=migration_batches,
)
assert communicator.context_calls == 2
assert communicator.execute_count == 2
assert [(dst, expert) for _, dst, expert in communicator.send_calls] == [(3, 0)]
assert [(src, expert) for _, src, expert in communicator.recv_calls] == [(1, 1)]
def test_move_to_buffer_uses_primary_duplicate_destination() -> None:
old_indices = np.array([0, 1, 2, 3], dtype=np.int64)
new_indices = np.array([2, 2, 0, 3], dtype=np.int64)
expert_weights = [torch.zeros(2, 1)]
recv_buffer = torch.zeros(2, 1)
communicator = _MockEplbCommunicator()
transfer_metadata = move_to_buffer(
num_local_experts=2,
old_indices=old_indices,
new_indices=new_indices,
expert_weights=expert_weights,
expert_weights_buffers=[recv_buffer],
stream=None,
ep_rank=0,
communicator=communicator,
enable_migration_batching=True,
)
assert len(communicator.recv_calls) == 1
recv_tensors, src_rank, expert_id = communicator.recv_calls[0]
assert (src_rank, expert_id) == (1, 2)
assert recv_tensors[0].data_ptr() == recv_buffer[0].data_ptr()
recv_buffer[0].fill_(42)
move_from_buffer(
expert_weights=expert_weights,
expert_weights_buffers=[recv_buffer],
transfer_metadata=transfer_metadata,
new_indices=new_indices,
ep_rank=0,
)
torch.testing.assert_close(expert_weights[0], torch.full((2, 1), 42.0))
def _endpoints(
flows: list[MigrationFlow],
) -> list[tuple[int, int]]:
return [(item.src_rank, item.dst_rank) for item in flows]
def _assert_no_endpoint_conflict(batch: list[MigrationFlow]) -> None:
endpoints: set[int] = set()
for flow in batch:
assert flow.src_rank not in endpoints
assert flow.dst_rank not in endpoints
endpoints.update((flow.src_rank, flow.dst_rank))
def _flatten_transfers(
batches: list[list[MigrationFlow]],
) -> list[tuple[int, int, int]]:
return [
(flow.src_rank, flow.dst_rank, expert_id)
for batch in batches
for flow in batch
for expert_id in flow.expert_ids
]
def _placements_for_transfers(
num_ranks: int,
transfers: list[tuple[int, int, int]],
) -> tuple[int, np.ndarray, np.ndarray]:
old_counts = [0] * num_ranks
new_counts = [0] * num_ranks
for src_rank, dst_rank, _ in transfers:
old_counts[src_rank] += 1
new_counts[dst_rank] += 1
num_local_experts = max(1, *old_counts, *new_counts)
old_indices = np.full(num_ranks * num_local_experts, -1, dtype=np.int64)
new_indices = np.full_like(old_indices, -1)
old_counts = [0] * num_ranks
new_counts = [0] * num_ranks
for src_rank, dst_rank, expert_id in transfers:
old_offset = src_rank * num_local_experts + old_counts[src_rank]
new_offset = dst_rank * num_local_experts + new_counts[dst_rank]
old_indices[old_offset] = expert_id
new_indices[new_offset] = expert_id
old_counts[src_rank] += 1
new_counts[dst_rank] += 1
return num_local_experts, old_indices, new_indices
+129
View File
@@ -231,6 +231,8 @@ def test_one_sided_combine_into_compatibility(supports_output):
payload, payload,
runtime_max_tokens_per_rank, runtime_max_tokens_per_rank,
output=None, output=None,
*,
use_low_precision=False,
): ):
result = payload + runtime_max_tokens_per_rank result = payload + runtime_max_tokens_per_rank
if output is None: if output is None:
@@ -249,6 +251,33 @@ def test_one_sided_combine_into_compatibility(supports_output):
torch.testing.assert_close(output, payload + 2) torch.testing.assert_close(output, payload + 2)
@pytest.mark.parametrize("low_precision_combine", [False, True])
def test_one_sided_combine_into_low_precision(low_precision_combine):
"""The fp8 combine setting is passed through to the kernel."""
from vllm.distributed.device_communicators.all2all import (
FlashInferNVLinkOneSidedManager,
)
seen_kwargs = {}
class FakeMoeAlltoAll:
def combine(self, payload, runtime_max_tokens_per_rank, output, **kwargs):
seen_kwargs.update(kwargs)
output.copy_(payload + runtime_max_tokens_per_rank)
manager = FlashInferNVLinkOneSidedManager.__new__(FlashInferNVLinkOneSidedManager)
manager.moe_alltoall = FakeMoeAlltoAll()
manager._combine_supports_output = True
manager.low_precision_combine = low_precision_combine
payload = torch.arange(4, dtype=torch.float32)
output = torch.empty_like(payload)
manager.combine_into(payload, runtime_max_tokens_per_rank=2, output=output)
assert seen_kwargs == {"use_low_precision": low_precision_combine}
torch.testing.assert_close(output, payload + 2)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init) # Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -915,6 +944,106 @@ def test_one_sided_dispatch_combine(world_size):
_spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size) _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
def _one_sided_low_precision_worker(rank, world_size):
from vllm.distributed.device_communicators.all2all import (
FlashInferNVLinkOneSidedManager,
)
from vllm.distributed.parallel_state import get_dp_group
device = torch.device(f"cuda:{rank}")
hidden_size = 256
tokens_per_rank = 32
experts_per_token = 2
num_experts = world_size * 8
manager = FlashInferNVLinkOneSidedManager(get_dp_group().cpu_group)
manager.initialize(
max_num_tokens=tokens_per_rank,
top_k=experts_per_token,
num_experts=num_experts,
hidden_size=hidden_size,
x_bytes_per_token=hidden_size // 2,
x_sf_bytes_per_token=hidden_size // 16,
)
assert manager.low_precision_combine
torch.manual_seed(rank + 42)
x = torch.randint(
0, 256, (tokens_per_rank, hidden_size // 2), device=device, dtype=torch.uint8
)
x_sf = torch.randint(
0, 256, (tokens_per_rank, hidden_size // 16), device=device, dtype=torch.uint8
)
topk_ids = torch.randint(
0,
num_experts,
(tokens_per_rank, experts_per_token),
device=device,
dtype=torch.int32,
)
topk_weights = torch.rand(
tokens_per_rank, experts_per_token, device=device, dtype=torch.float32
)
manager.moe_alltoall.dispatch(
token_selected_experts=topk_ids,
input_payloads=[x, x_sf, topk_ids, topk_weights],
runtime_max_tokens_per_rank=tokens_per_rank,
)
# Go through combine_into rather than moe_alltoall.combine: the workspace
# was sized for fp8, so the kernel must be told to send fp8. Calling the
# raw kernel here would send bf16 into fp8-sized slots and corrupt it.
expert_output = torch.ones(
world_size, tokens_per_rank, hidden_size, device=device, dtype=torch.bfloat16
)
output = torch.empty(
tokens_per_rank, hidden_size, device=device, dtype=torch.bfloat16
)
manager.combine_into(
payload=expert_output,
runtime_max_tokens_per_rank=tokens_per_rank,
output=output,
)
# Integers 1..top_k are exact in fp8_e4m3, so the fp8 transport must match
# bf16 exactly: one contribution per distinct expert-owning rank.
experts_per_rank = num_experts // world_size
expert_ranks = topk_ids // experts_per_rank
num_distinct = torch.tensor(
[len(set(row.tolist())) for row in expert_ranks],
device=device,
dtype=torch.bfloat16,
).unsqueeze(1)
torch.testing.assert_close(output, num_distinct.expand_as(output))
torch.distributed.barrier()
manager.cleanup()
@requires_multi_gpu
@requires_one_sided
@requires_ptrace
@pytest.mark.parametrize("world_size", [2])
def test_one_sided_combine_low_precision(world_size, monkeypatch):
"""Exercise the real fp8 combine transport, not a fake kernel.
A workspace/dtype mismatch corrupts the combine buffer silently rather
than raising, so this asserts the combined values are exactly right.
"""
from flashinfer.comm.trtllm_moe_alltoall import MoeAlltoAll
from vllm.utils.func_utils import supports_kw
if not supports_kw(
MoeAlltoAll.combine, "use_low_precision", allow_var_kwargs=False
):
pytest.skip("installed FlashInfer has no use_low_precision support")
monkeypatch.setenv("VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE", "1")
_spawn_workers(_one_sided_low_precision_worker, world_size, dp_size=world_size)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle # Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
+87
View File
@@ -291,6 +291,93 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase):
reader.touch("key", address, monotonic_id, signature) reader.touch("key", address, monotonic_id, signature)
class TestShmObjectStorageReferenceCounting(unittest.TestCase):
"""One writer and four readers, as with tensor-parallel workers."""
def setUp(self):
ring_buffer = SingleWriterShmRingBuffer(
data_buffer_size=1024 * 100, create=True
)
self.writer = SingleWriterShmObjectStorage(
max_object_size=1024 * 10,
n_readers=4,
ring_buffer=ring_buffer,
serde_class=MsgpackSerde,
)
self.addCleanup(self.writer.close)
handle = self.writer.handle()
handle.reader_lock = Lock()
self.readers = [
SingleWriterShmObjectStorage.create_from_handle(handle) for _ in range(4)
]
for reader in self.readers:
self.addCleanup(reader.close)
def _reference(self, key: str) -> tuple[int, int, list[int], str]:
"""Take one writer reference to the key, as one request item does."""
if self.writer.is_cached(key):
address, monotonic_id = self.writer.get_cached(key)
else:
address, monotonic_id = self.writer.put(key, "value")
return address, monotonic_id, self.writer.get_signature(key), key
@staticmethod
def _touch(reader: SingleWriterShmObjectStorage, reference) -> None:
address, monotonic_id, signature, key = reference
reader.touch(key, address, monotonic_id, signature)
def _is_evictable(self, key: str) -> bool:
self.writer.free_unused()
return not self.writer.is_cached(key)
def test_item_referenced_again_before_every_reader_fetched_it(self):
"""A second request must not leave the item unevictable forever."""
first = self._reference("image")
self.writer.touch("image")
second = self._reference("image")
self.writer.release_touches()
# Every reader starts on the first request; the ones that finish it
# go on to the second while the others are still fetching.
for reader in self.readers:
self._touch(reader, first)
for reader in self.readers:
reader.get(*first)
self._touch(reader, second)
reader.get(*second)
self.assertTrue(self._is_evictable("image"))
def test_new_item_twice_in_one_request_is_not_evicted_early(self):
"""The item must stay until the last reader has fetched it."""
self.writer.touch("image")
references = [self._reference("image"), self._reference("image")]
self.writer.release_touches()
for reader in self.readers[:-1]:
for reference in references:
self._touch(reader, reference)
for reference in references:
reader.get(*reference)
for reference in references:
self._touch(self.readers[-1], reference)
self.assertFalse(self._is_evictable("image"))
for reference in references:
self.readers[-1].get(*reference)
self.assertTrue(self._is_evictable("image"))
def test_touch_protects_item_until_released(self):
reference = self._reference("image")
for reader in self.readers:
reader.get(*reference)
self.writer.touch("image")
self.assertFalse(self._is_evictable("image"))
self.writer.release_touches()
self.assertTrue(self._is_evictable("image"))
# Reader process function # Reader process function
def reader_process(process_id, storage_handle, items_to_read): def reader_process(process_id, storage_handle, items_to_read):
"""Reader process that connects to existing shared memory and reads data.""" """Reader process that connects to existing shared memory and reads data."""
+11 -9
View File
@@ -7,12 +7,12 @@ Integration tests for NCCL and IPC weight transfer between processes using Ray.
""" """
import builtins import builtins
import importlib.util
import pickle import pickle
import runpy import runpy
import threading import threading
import time import time
from contextlib import nullcontext from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import MagicMock from unittest.mock import MagicMock
@@ -1497,10 +1497,12 @@ class TestDeferredProcessingContract:
that caller ends up reaching through a getattr.""" that caller ends up reaching through a getattr."""
def _engines(self): def _engines(self):
return { registry = dict(WeightTransferEngineFactory._registry)
name: loader() # ModelExpress is an optional, separately installed package, but its
for name, loader in WeightTransferEngineFactory._registry.items() # backend is always registered. Skip it when the package is missing.
} if importlib.util.find_spec("modelexpress") is None:
registry.pop("modelexpress")
return {name: loader() for name, loader in registry.items()}
def test_every_engine_declares_whether_it_defers(self): def test_every_engine_declares_whether_it_defers(self):
for name, cls in self._engines().items(): for name, cls in self._engines().items():
@@ -2054,11 +2056,11 @@ def test_sparse_nccl_trainer_non_sender_skips_client():
def _import_modelexpress_shim(): def _import_modelexpress_shim():
path = ( spec = importlib.util.find_spec(
Path(__file__).parents[2] "vllm.distributed.weight_transfer.modelexpress_engine"
/ "vllm/distributed/weight_transfer/modelexpress_engine.py"
) )
runpy.run_path(str(path), run_name="_test_modelexpress_shim") assert spec is not None and spec.origin is not None
runpy.run_path(spec.origin, run_name="_test_modelexpress_shim")
@pytest.mark.parametrize( @pytest.mark.parametrize(
@@ -237,6 +237,46 @@ def test_snapshot_environment_contract(
) )
@pytest.mark.parametrize(
("configured", "expected"),
[(None, "1"), ("0", "0"), ("1", "1")],
)
def test_snapshot_child_defaults_nccl_ib_without_mutating_parent(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
configured: str | None,
expected: str,
):
if configured is None:
monkeypatch.delenv("NCCL_IB_DISABLE", raising=False)
else:
monkeypatch.setenv("NCCL_IB_DISABLE", configured)
parent_environment = os.environ.copy()
process = create_autospec(subprocess.Popen, instance=True)
process.pid = 100
popen = create_autospec(subprocess.Popen, return_value=process)
monkeypatch.setattr(snapshot_runtime.subprocess, "Popen", popen)
LocalSnapshotTools().launch_child(tmp_path, ("model",))
assert popen.call_args.kwargs["env"]["NCCL_IB_DISABLE"] == expected
assert os.environ == parent_environment
def test_snapshot_environment_identity_applies_effective_nccl_ib_default(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.delenv("NCCL_IB_DISABLE", raising=False)
implicit = LocalSnapshotTools()._environment_identity()
monkeypatch.setenv("NCCL_IB_DISABLE", "1")
explicit_disabled = LocalSnapshotTools()._environment_identity()
monkeypatch.setenv("NCCL_IB_DISABLE", "0")
explicit_enabled = LocalSnapshotTools()._environment_identity()
assert implicit == explicit_disabled
assert explicit_enabled != implicit
@pytest.mark.parametrize( @pytest.mark.parametrize(
"case", "case",
[ [
@@ -369,6 +409,33 @@ ss() { :; }
assert result.returncode == (0 if probe is None else 1), result.stderr assert result.returncode == (0 if probe is None else 1), result.stderr
@pytest.mark.parametrize("container_attempted", [False, True])
def test_snapshot_e2e_cleanup_checks_state_only_after_container_attempt(
container_attempted: bool,
):
script = (
Path(__file__).resolve().parents[3]
/ ".buildkite/scripts/initialized-snapshot-e2e.sh"
)
if not script.exists():
pytest.skip("snapshot E2E script is not packaged in this test image")
source = script.read_text()
cleanup = source[source.index("cleanup() {") : source.index("on_exit() {")]
shell = (
"set -Eeuo pipefail\n"
"CONTAINER_NAME='' PORT_TWO=18002\n"
f"CHECK_CLEANUP_STATE={int(container_attempted)}\n"
"wait_clean() { echo checked-state; return 1; }\n"
+ cleanup
+ "\nif cleanup; then exit 0; else exit 1; fi\n"
)
result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10)
assert result.returncode == int(container_attempted), result.stderr
assert (b"checked-state" in result.stdout) == container_attempted
def test_snapshot_create_cli_accepts_only_pinned_compact_mode( def test_snapshot_create_cli_accepts_only_pinned_compact_mode(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
): ):
@@ -652,6 +719,111 @@ def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path):
assert not target.exists() assert not target.exists()
@pytest.mark.parametrize(
"failure_kind", ["called-process", "timeout", "missing", "unreadable"]
)
def test_snapshot_create_preserves_bounded_criu_diagnostics_after_cleanup(
tmp_path: Path, failure_kind: str, caplog: pytest.LogCaptureFixture
):
target = tmp_path / "snapshot"
tools = _fake_snapshot_tools()
criu_tools = LocalSnapshotTools()
criu_tools.plugin_dir = tmp_path
original_run = criu_tools._run
criu_tools._privileged = lambda: [] # type: ignore[method-assign]
captured = "captured stderr " + "x" * 9000
if failure_kind == "timeout":
primary: subprocess.SubprocessError = subprocess.TimeoutExpired(
["criu", "dump"], 1, output=b"captured stdout\xff", stderr=captured.encode()
)
else:
primary = subprocess.CalledProcessError(
1, ["criu", "dump"], output="captured stdout", stderr=captured
)
def run(command: list[str], **_kwargs: object):
if "tail" not in command:
raise primary
if failure_kind == "unreadable":
raise PermissionError("log is root-owned")
if failure_kind == "timeout":
return subprocess.CompletedProcess(command, 0, "timeout log tail", "")
return original_run(command, **_kwargs)
criu_tools._run = run # type: ignore[method-assign]
def fail_dump(workdir: Path, _inventory: ProcessInventory) -> None:
images = workdir / "images"
images.mkdir()
if failure_kind != "missing":
(images / "dump.log").write_text(
"omitted dump prefix" + "d" * 9000 + "distinctive dump diagnostic"
)
criu_tools._criu("dump", workdir, [])
tools.dump.side_effect = fail_dump
with pytest.raises(type(primary)) as excinfo:
create_snapshot(
argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"),
tools=tools,
)
assert excinfo.value is primary
assert not target.exists()
diagnostics = "\n".join(caplog.messages)
assert "captured stdout" in diagnostics
assert "captured stderr" not in diagnostics
assert "x" * 100 in diagnostics
assert len(diagnostics) < 17000
if failure_kind in {"missing", "unreadable"}:
assert "dump.log unavailable" in diagnostics
expected_error = (
"CalledProcessError" if failure_kind == "missing" else "PermissionError"
)
assert expected_error in diagnostics
elif failure_kind == "called-process":
assert "distinctive dump diagnostic" in diagnostics
assert "omitted dump prefix" not in diagnostics
else:
assert "timeout log tail" in diagnostics
if failure_kind == "timeout":
assert "captured stdout�" in diagnostics
def test_snapshot_criu_failure_selects_restore_log(
tmp_path: Path, caplog: pytest.LogCaptureFixture
):
tools = LocalSnapshotTools()
tools.plugin_dir = tmp_path
tools._privileged = lambda: [] # type: ignore[method-assign]
artifact = tmp_path / "snapshot"
images = artifact / "images"
images.mkdir(parents=True)
(images / "restore.log").write_text(
"omitted restore prefix" + "r" * 9000 + "distinctive restore diagnostic"
)
primary = subprocess.CalledProcessError(1, ["criu", "restore"])
original_run = tools._run
tailed: list[str] = []
def run(command: list[str], **_kwargs: object):
if "tail" not in command:
raise primary
tailed.append(command[-1])
return original_run(command, **_kwargs)
tools._run = run # type: ignore[method-assign]
with pytest.raises(subprocess.CalledProcessError) as excinfo:
tools._criu("restore", artifact, [])
assert excinfo.value is primary
assert tailed == [str(images / "restore.log")]
diagnostics = "\n".join(caplog.messages)
assert "distinctive restore diagnostic" in diagnostics
assert "omitted restore prefix" not in diagnostics
def _restore_fixture(tmp_path: Path) -> argparse.Namespace: def _restore_fixture(tmp_path: Path) -> argparse.Namespace:
artifact = tmp_path / "snapshot" artifact = tmp_path / "snapshot"
artifact.mkdir(mode=0o700) artifact.mkdir(mode=0o700)
@@ -868,6 +1040,96 @@ def test_snapshot_private_path_and_link_remap_security(
tools.restore(artifact, _manifest()) tools.restore(artifact, _manifest())
@pytest.mark.parametrize("inventory_error", [None, PermissionError("denied")])
def test_snapshot_dump_rejects_unavailable_link_remap_inventory_before_criu(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
inventory_error: PermissionError | None,
):
tools = LocalSnapshotTools()
workdir = tmp_path / "snapshot"
workdir.mkdir()
tools.shm_dir = tmp_path / "missing-shm"
if inventory_error is not None:
tools.shm_dir.mkdir()
def fail_inventory(_path: Path):
raise inventory_error
monkeypatch.setattr(Path, "iterdir", fail_inventory)
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
with pytest.raises(SnapshotCreateError, match="inventory CRIU link remaps"):
tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU"))
def test_snapshot_restore_rejects_saved_remap_without_prefix(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
_mark_snapshot_pids_free(monkeypatch)
artifact = _local_restore_artifact(tmp_path)
saved = artifact / "link-remaps"
saved.mkdir()
(saved / "270").write_bytes(b"invalid remap name")
tools = LocalSnapshotTools()
tools.shm_dir = tmp_path / "shm"
tools.shm_dir.mkdir()
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
with pytest.raises(SnapshotRestoreError, match="invalid saved link remap"):
tools.restore(artifact, _manifest())
assert not tuple(tools.shm_dir.iterdir())
def test_snapshot_link_remap_inventory_matches_only_numeric_suffixes(tmp_path: Path):
for name in (
"link_remap.0",
"link_remap.270",
"link_remap.",
"link_remap.12a",
"link-remap.12",
"prefix.link_remap.12",
"270",
):
(tmp_path / name).touch()
tools = LocalSnapshotTools()
tools.shm_dir = tmp_path
assert tools._link_remap_names() == {"link_remap.0", "link_remap.270"}
def test_snapshot_dump_preserves_primary_failure_when_cleanup_inventory_fails(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
tools = LocalSnapshotTools()
workdir = tmp_path / "snapshot"
workdir.mkdir()
inventory_calls = 0
def inventory_names() -> set[str]:
nonlocal inventory_calls
inventory_calls += 1
if inventory_calls > 1:
raise SnapshotCreateError(
"could not inventory CRIU link remaps: PermissionError: denied"
)
return set()
monkeypatch.setattr(tools, "_link_remap_names", inventory_names)
def fail_dump(*_args):
raise RuntimeError("dump failed")
monkeypatch.setattr(tools, "_criu", fail_dump)
with pytest.raises(SnapshotCreateError) as failure:
tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU"))
assert "RuntimeError: dump failed" in str(failure.value)
assert "cleanup inventory failed" in str(failure.value)
@pytest.mark.parametrize( @pytest.mark.parametrize(
("blocked_state", "message"), ("blocked_state", "message"),
[("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")], [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")],
@@ -900,6 +1162,28 @@ def test_snapshot_rejects_unsafe_process_state_before_criu(
assert "1.1.1.1:443" in str(excinfo.value) assert "1.1.1.1:443" in str(excinfo.value)
def test_snapshot_rejects_open_infiniband_device_before_criu(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
tools = LocalSnapshotTools()
monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101))
monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
monkeypatch.setattr(
tools,
"_descriptor_targets",
lambda pid: ("/dev/infiniband/uverbs7",) if pid == 101 else (),
)
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
with pytest.raises(SnapshotCreateError) as failure:
tools.inventory(100, tmp_path)
message = str(failure.value)
assert "pid 101" in message
assert "/dev/infiniband/uverbs7" in message
assert "NCCL_IB_DISABLE=1" in message
def test_decode_endpoint_families(): def test_decode_endpoint_families():
assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030" assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030"
assert ( assert (
@@ -40,13 +40,11 @@ def server():
"--trust-request-mm-kwargs", "--trust-request-mm-kwargs",
] ]
# ROCm: Increase timeouts to handle potential network delays and slower # ROCm: Allow more time to download videos from external sources.
# video processing when downloading multiple videos from external sources
env_overrides = {} env_overrides = {}
if current_platform.is_rocm(): if current_platform.is_rocm():
env_overrides = { env_overrides = {
"VLLM_VIDEO_FETCH_TIMEOUT": "120", "VLLM_VIDEO_FETCH_TIMEOUT": "120",
"VLLM_ENGINE_ITERATION_TIMEOUT_S": "300",
} }
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_overrides) as remote_server: with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_overrides) as remote_server:
@@ -88,13 +88,11 @@ def server():
*ROCM_EXTRA_ARGS, *ROCM_EXTRA_ARGS,
] ]
# ROCm: Increase timeouts to handle potential network delays and slower # ROCm: Allow more time to download videos from external sources.
# video processing when downloading multiple videos from external sources
env_overrides = { env_overrides = {
**( **(
{ {
"VLLM_VIDEO_FETCH_TIMEOUT": "120", "VLLM_VIDEO_FETCH_TIMEOUT": "120",
"VLLM_ENGINE_ITERATION_TIMEOUT_S": "300",
} }
if current_platform.is_rocm() if current_platform.is_rocm()
else {} else {}
@@ -1,20 +1,25 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ScoringIOProcessor post-tokenization helpers.""" """Unit tests for ScoringIOProcessor helpers."""
from dataclasses import dataclass from dataclasses import dataclass
from types import SimpleNamespace from types import SimpleNamespace
import pytest import pytest
import torch
from tokenizers import Tokenizer, models, pre_tokenizers, processors from tokenizers import Tokenizer, models, pre_tokenizers, processors
from transformers import TokenizersBackend from transformers import TokenizersBackend
from vllm import TokensPrompt from vllm import TokensPrompt
from vllm.entrypoints.pooling.scoring.io_processor import ( from vllm.entrypoints.pooling.scoring.io_processor import (
BiEncoderIOProcessor,
CrossEncoderIOProcessor, CrossEncoderIOProcessor,
JinaRankingIOProcessor,
LateInteractionIOProcessor,
_apply_post_tokenization_to_token_type_ids, _apply_post_tokenization_to_token_type_ids,
) )
from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
from vllm.outputs import PoolingOutput, PoolingRequestOutput, RequestError
from vllm.renderers import TokenizeParams from vllm.renderers import TokenizeParams
pytestmark = pytest.mark.skip_global_cleanup pytestmark = pytest.mark.skip_global_cleanup
@@ -113,3 +118,59 @@ def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():
assert first_doc == 10 assert first_doc == 10
assert prompt_token_ids[:first_doc] == list(range(10, num_query)) assert prompt_token_ids[:first_doc] == list(range(10, num_query))
assert prompt_token_ids[first_doc:40] == list(range(num_query, 50)) assert prompt_token_ids[first_doc:40] == list(range(num_query, 50))
def _pooling_output(
request_id: str,
data: torch.Tensor,
error: RequestError | None = None,
) -> PoolingRequestOutput:
return PoolingRequestOutput(
request_id=request_id,
outputs=PoolingOutput(data),
prompt_token_ids=[1],
num_cached_tokens=0,
finished=True,
error=error,
)
@pytest.mark.parametrize(
"processor_type", [BiEncoderIOProcessor, LateInteractionIOProcessor]
)
@pytest.mark.parametrize("failed_index", [0, 1])
def test_paired_scoring_preserves_request_error(processor_type, failed_index):
processor = processor_type.__new__(processor_type)
processor.pad_token_id = None
error = RequestError(
code="multimodal_cache_miss",
message="Multi-modal processor cache miss.",
retryable=True,
)
outputs = [
_pooling_output("query", torch.tensor([1.0, 2.0])),
_pooling_output("document", torch.tensor([3.0, 4.0])),
]
outputs[failed_index] = _pooling_output(
outputs[failed_index].request_id,
torch.empty(0),
error,
)
(result,) = processor._post_process(outputs, n_queries=1)
assert result.request_id == "query_document"
assert result.error is error
assert result.outputs is outputs[failed_index].outputs
def test_jina_scoring_preserves_request_error_before_indexing_output():
processor = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)
error = RequestError(
code="multimodal_cache_miss",
message="Multi-modal processor cache miss.",
retryable=True,
)
failed = _pooling_output("request", torch.empty(0), error)
assert processor._post_process([failed], n_queries=1) == [failed]
@@ -49,10 +49,60 @@ def _query_key(context: PoolingServeContext) -> str:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_colliding_request_ids_use_distinct_query_cache_keys( @pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("failed_stage", ["query", "doc"])
async def test_flash_late_interaction_cleans_up_queries_on_failure(
failed_stage: str,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
): ):
serving = object.__new__(ServingScores) serving = object.__new__(ServingScores)
serving.io_processor = Mock()
serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock())
context = _make_context()
monkeypatch.setattr(serving, "_init_ctx", AsyncMock(return_value=context))
monkeypatch.setattr(serving, "_preprocessing", AsyncMock())
async def encode_queries(ctx):
ctx.late_interaction_query_keys = ["query-key"]
if failed_stage == "query":
raise RuntimeError("query failed")
async def encode_docs(ctx):
ctx.late_interaction_doc_keys = ["doc-key"]
if failed_stage == "doc":
raise RuntimeError("doc failed")
monkeypatch.setattr(
serving,
"_flash_late_interaction_encode_queries",
AsyncMock(side_effect=encode_queries),
)
monkeypatch.setattr(
serving,
"_flash_late_interaction_encode_docs",
AsyncMock(side_effect=encode_docs),
)
with pytest.raises(RuntimeError, match=f"{failed_stage} failed"):
await serving.flash_late_interaction()
expected_doc_keys = [] if failed_stage == "query" else ["doc-key"]
serving.engine_client.abort.assert_awaited_once_with(
["query-key", *expected_doc_keys]
)
serving.engine_client.collective_rpc.assert_awaited_once_with(
"release_late_interaction_query_cache",
args=(["query-key"],),
)
@pytest.mark.asyncio
async def test_colliding_request_ids_use_distinct_late_interaction_keys(
monkeypatch: pytest.MonkeyPatch,
):
serving = object.__new__(ServingScores)
serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock())
prepare_generators = AsyncMock() prepare_generators = AsyncMock()
monkeypatch.setattr(serving, "_prepare_generators", prepare_generators) monkeypatch.setattr(serving, "_prepare_generators", prepare_generators)
monkeypatch.setattr(serving, "_collect_batch", AsyncMock()) monkeypatch.setattr(serving, "_collect_batch", AsyncMock())
@@ -68,9 +118,24 @@ async def test_colliding_request_ids_use_distinct_query_cache_keys(
first_doc_key = _query_key(prepared_contexts[1]) first_doc_key = _query_key(prepared_contexts[1])
second_query_key = _query_key(prepared_contexts[2]) second_query_key = _query_key(prepared_contexts[2])
second_doc_key = _query_key(prepared_contexts[3]) second_doc_key = _query_key(prepared_contexts[3])
first_doc_request_ids = prepared_contexts[1].prompt_request_ids
second_doc_request_ids = prepared_contexts[3].prompt_request_ids
assert first_query_key == first_doc_key assert first_query_key == first_doc_key
assert second_query_key == second_doc_key assert second_query_key == second_doc_key
assert first_query_key != second_query_key assert first_query_key != second_query_key
assert "caller-controlled-id" not in first_query_key assert "caller-controlled-id" not in first_query_key
assert "caller-controlled-id" not in second_query_key assert "caller-controlled-id" not in second_query_key
assert first_doc_request_ids is not None
assert second_doc_request_ids is not None
assert first.late_interaction_doc_keys == first_doc_request_ids
assert second.late_interaction_doc_keys == second_doc_request_ids
assert first_doc_request_ids != second_doc_request_ids
assert all("caller-controlled-id" not in key for key in first_doc_request_ids)
assert all("caller-controlled-id" not in key for key in second_doc_request_ids)
assert first.late_interaction_query_keys is not None
await serving._cleanup_flash_late_interaction(first)
serving.engine_client.abort.assert_awaited_once_with(
[*first.late_interaction_query_keys, *first_doc_request_ids]
)
@@ -22,6 +22,10 @@ The streaming cases do the same for `stream=true`. Each coupled chunk's
`token_ids` becomes one `/inference/v1/generate` stream chunk, fed `token_ids` becomes one `/inference/v1/generate` stream chunk, fed
through the chunked derender path with `stream_state` threaded across calls through the chunked derender path with `stream_state` threaded across calls
the way a client would, so both parsers see the same chunk boundaries. the way a client would, so both parsers see the same chunk boundaries.
The text level cases pin `/inference/v1/generate` with `output_mode="text"`
against `/v1/completions/derender` on the same token IDs, batch and streaming,
plus logprob resolution and stop string handling.
""" """
import json import json
@@ -440,3 +444,241 @@ async def test_stream_parity_reasoning_and_tool_call(client, chunking):
assert any(len(ch["token_ids"]) > 1 for ch in coupled_choices) assert any(len(ch["token_ids"]) > 1 for ch in coupled_choices)
if not (coupled["reasoning"] and coupled["tool_calls"]): if not (coupled["reasoning"] and coupled["tool_calls"]):
pytest.skip("Model did not emit both a <think> block and a tool call") pytest.skip("Model did not emit both a <think> block and a tool call")
# ---------------------------------------------------------------------------
# Inline text level: generate(output_mode="text") == derender(generate(...))
# ---------------------------------------------------------------------------
TEXT_MESSAGES = [
{"role": "user", "content": "What is 2+2? Answer in one short sentence."}
]
async def _render_token_ids(
client: httpx.AsyncClient, messages: list[dict]
) -> list[int]:
resp = await client.post(
"/v1/chat/completions/render", json={"model": MODEL, "messages": messages}
)
assert resp.status_code == 200, resp.text
return resp.json()["token_ids"]
def _generate_payload(
token_ids: list[int], output_mode: str, stream: bool = False, **sampling
) -> dict:
return {
"model": MODEL,
"token_ids": token_ids,
"sampling_params": {"temperature": 0, "max_tokens": 128, **sampling},
"output_mode": output_mode,
"stream": stream,
}
async def _generate(
client: httpx.AsyncClient, token_ids: list[int], output_mode: str, **sampling
) -> dict:
resp = await client.post(
"/inference/v1/generate",
json=_generate_payload(token_ids, output_mode, **sampling),
)
assert resp.status_code == 200, resp.text
data = resp.json()
assert data["output_mode"] == output_mode
return data
async def _generate_text_stream(
client: httpx.AsyncClient, token_ids: list[int], **sampling
) -> list[dict]:
"""Stream generate at the text level and return the choice of every chunk
that has one."""
choices: list[dict] = []
async with client.stream(
"POST",
"/inference/v1/generate",
json=_generate_payload(token_ids, "text", stream=True, **sampling),
) as resp:
assert resp.status_code == 200, await resp.aread()
async for line in resp.aiter_lines():
if not line.startswith("data: ") or line == "data: [DONE]":
continue
chunk = json.loads(line[len("data: ") :])
assert "error" not in chunk, chunk
assert chunk["output_mode"] == "text"
choices.extend(chunk["choices"])
return choices
async def _derender_completion_text(
client: httpx.AsyncClient, output_ids: list[int], finish_reason: str
) -> str:
resp = await client.post(
"/v1/completions/derender",
json={
"model": MODEL,
"generate_responses": [
{
"request_id": "text-parity",
"choices": [
{
"index": 0,
"token_ids": output_ids,
"finish_reason": finish_reason,
}
],
}
],
},
)
assert resp.status_code == 200, resp.text
return resp.json()["choices"][0]["text"]
async def _derender_completion_stream_text(
client: httpx.AsyncClient, choices: list[dict]
) -> str:
"""Replay each generate chunk's token IDs through the streaming
completions derender endpoint with the same chunk boundaries."""
state = None
text = ""
for choice in choices:
resp = await client.post(
"/v1/completions/derender",
json={
"stream": True,
"model": MODEL,
"generate_chunk": {
"request_id": "text-parity-stream",
"choices": [
{
"index": 0,
"token_ids": choice.get("token_ids") or [],
"finish_reason": choice.get("finish_reason"),
}
],
},
"stream_state": state,
},
)
assert resp.status_code == 200, resp.text
data = resp.json()
state = data["stream_state"]
text += "".join(ch["text"] for ch in data["chunk"]["choices"])
return text
@pytest.mark.asyncio
async def test_text_parity_batch(client):
"""Inline text matches /v1/completions/derender on the same token IDs."""
token_ids = await _render_token_ids(client, TEXT_MESSAGES)
inline = (await _generate(client, token_ids, "text"))["choices"][0]
derendered = await _derender_completion_text(
client, inline["token_ids"], inline["finish_reason"]
)
assert inline["text"]
assert inline["text"] == derendered
@pytest.mark.asyncio
async def test_text_parity_stream(client):
"""Streamed inline text matches both the streaming and the non-streaming
derender of the same token IDs."""
token_ids = await _render_token_ids(client, TEXT_MESSAGES)
choices = await _generate_text_stream(client, token_ids)
inline = "".join(ch["text"] for ch in choices)
output_ids = [t for ch in choices for t in ch.get("token_ids") or []]
finish_reason = choices[-1]["finish_reason"]
assert finish_reason is not None
assert inline == await _derender_completion_stream_text(client, choices)
assert inline == await _derender_completion_text(client, output_ids, finish_reason)
@pytest.mark.asyncio
async def test_text_parity_logprobs(client):
"""Inline logprob tokens and bytes match /derender resolving the
token_id:N placeholders of the tokens level."""
token_ids = await _render_token_ids(client, TEXT_MESSAGES)
sampling = {"logprobs": 3}
tokens_choice = (await _generate(client, token_ids, "tokens", **sampling))[
"choices"
][0]
text_choice = (await _generate(client, token_ids, "text", **sampling))["choices"][0]
assert text_choice["token_ids"] == tokens_choice["token_ids"], (
"greedy (temperature=0) generation was expected to be deterministic "
"across the two generate calls"
)
disagg = await _disagg(
client,
tokens_choice["token_ids"],
len(token_ids),
tokens_choice["finish_reason"],
{"model": MODEL, "messages": TEXT_MESSAGES},
logprobs=tokens_choice["logprobs"],
)
inline_content = text_choice["logprobs"]["content"]
derender_content = disagg["choices"][0]["logprobs"]["content"]
assert len(inline_content) == len(derender_content)
for inline_entry, derender_entry in zip(inline_content, derender_content):
assert inline_entry["token"] == derender_entry["token"]
assert inline_entry["bytes"] == derender_entry["bytes"]
assert [
(top["token"], top["bytes"]) for top in inline_entry["top_logprobs"]
] == [(top["token"], top["bytes"]) for top in derender_entry["top_logprobs"]]
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_text_parity_stop_string(client, stream):
"""Inline text follows the engine's stop handling, like /v1/completions.
token_ids keep every generated token, so /derender brings the matched
stop string back. Its text must start with the inline text and only
differ after it.
"""
stop = "\n"
token_ids = await _render_token_ids(
client,
[{"role": "user", "content": "Write two short lines about the sea."}],
)
resp = await client.post(
"/v1/completions",
json={
"model": MODEL,
"prompt": token_ids,
"temperature": 0,
"max_tokens": 128,
"stop": [stop],
},
)
assert resp.status_code == 200, resp.text
coupled = resp.json()["choices"][0]
if stream:
choices = await _generate_text_stream(client, token_ids, stop=[stop])
inline_text = "".join(ch["text"] for ch in choices)
output_ids = [t for ch in choices for t in ch.get("token_ids") or []]
finish_reason = choices[-1]["finish_reason"]
else:
inline = (await _generate(client, token_ids, "text", stop=[stop]))["choices"][0]
inline_text = inline["text"]
output_ids = inline["token_ids"]
finish_reason = inline["finish_reason"]
assert inline_text == coupled["text"]
assert finish_reason == coupled["finish_reason"]
assert stop not in inline_text
derendered = await _derender_completion_text(client, output_ids, finish_reason)
assert derendered.startswith(inline_text)
if coupled["stop_reason"] != stop:
pytest.skip("Model did not emit the stop string")
assert derendered[len(inline_text) :].startswith(stop)
@@ -50,10 +50,11 @@ from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
DerenderCompletionRequest, DerenderCompletionRequest,
DerenderCompletionStreamRequest, DerenderCompletionStreamRequest,
DerenderStreamState, DerenderStreamState,
GenerateResponse,
GenerateResponseChoice,
GenerateResponseStreamChoice,
GenerateStreamResponse, GenerateStreamResponse,
GenerateTokensChoice,
GenerateTokensResponse,
GenerateTokensStreamChoice,
GenerateTokensStreamResponse,
) )
from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
from vllm.parser import Parser from vllm.parser import Parser
@@ -195,10 +196,10 @@ def _make_stream_chunk(
prompt_token_ids: list[int] | None = None, prompt_token_ids: list[int] | None = None,
) -> GenerateStreamResponse: ) -> GenerateStreamResponse:
"""Build a GenerateStreamResponse SSE chunk.""" """Build a GenerateStreamResponse SSE chunk."""
return GenerateStreamResponse( return GenerateTokensStreamResponse(
request_id=request_id, request_id=request_id,
choices=[ choices=[
GenerateResponseStreamChoice( GenerateTokensStreamChoice(
index=index, index=index,
token_ids=token_ids, token_ids=token_ids,
finish_reason=finish_reason, finish_reason=finish_reason,
@@ -233,7 +234,7 @@ def _make_usage_chunk(
request_id: str = "test-req", request_id: str = "test-req",
) -> GenerateStreamResponse: ) -> GenerateStreamResponse:
"""Build a usage only final SSE chunk (empty choices).""" """Build a usage only final SSE chunk (empty choices)."""
return GenerateStreamResponse( return GenerateTokensStreamResponse(
request_id=request_id, request_id=request_id,
choices=[], choices=[],
usage=UsageInfo( usage=UsageInfo(
@@ -326,11 +327,11 @@ def test_non_streaming_derender_initializes_parser_prefix(
) )
monkeypatch.setattr(derenderer, "parser", MagicMock(return_value=parser)) monkeypatch.setattr(derenderer, "parser", MagicMock(return_value=parser))
generated_ids = tokenizer.encode("answer", add_special_tokens=False) generated_ids = tokenizer.encode("answer", add_special_tokens=False)
response = GenerateResponse( response = GenerateTokensResponse(
request_id="test", request_id="test",
prompt_token_ids=[11, 12], prompt_token_ids=[11, 12],
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, index=0,
token_ids=generated_ids, token_ids=generated_ids,
finish_reason="length", finish_reason="length",
@@ -583,10 +584,10 @@ class TestPromptSeededLeadingSpace:
): ):
prompt_ids, output_ids, expected = leading_space_ids prompt_ids, output_ids, expected = leading_space_ids
choices = await derenderer.derender_chat( choices = await derenderer.derender_chat(
GenerateResponse( GenerateTokensResponse(
request_id="t", request_id="t",
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, token_ids=output_ids, finish_reason="stop" index=0, token_ids=output_ids, finish_reason="stop"
) )
], ],
@@ -601,10 +602,10 @@ class TestPromptSeededLeadingSpace:
): ):
prompt_ids, output_ids, expected = leading_space_ids prompt_ids, output_ids, expected = leading_space_ids
choices = await derenderer.derender_chat( choices = await derenderer.derender_chat(
GenerateResponse( GenerateTokensResponse(
request_id="t", request_id="t",
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, token_ids=output_ids, finish_reason="stop" index=0, token_ids=output_ids, finish_reason="stop"
) )
], ],
@@ -620,10 +621,10 @@ class TestPromptSeededLeadingSpace:
prompt_ids, output_ids, expected = leading_space_ids prompt_ids, output_ids, expected = leading_space_ids
wrong_ids = tokenizer.encode("a different prompt", add_special_tokens=False) wrong_ids = tokenizer.encode("a different prompt", add_special_tokens=False)
choices = await derenderer.derender_chat( choices = await derenderer.derender_chat(
GenerateResponse( GenerateTokensResponse(
request_id="t", request_id="t",
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, token_ids=output_ids, finish_reason="stop" index=0, token_ids=output_ids, finish_reason="stop"
) )
], ],
@@ -639,10 +640,10 @@ class TestPromptSeededLeadingSpace:
): ):
_, output_ids, _ = leading_space_ids _, output_ids, _ = leading_space_ids
choices = await derenderer.derender_chat( choices = await derenderer.derender_chat(
GenerateResponse( GenerateTokensResponse(
request_id="t", request_id="t",
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, token_ids=output_ids, finish_reason="stop" index=0, token_ids=output_ids, finish_reason="stop"
) )
], ],
@@ -655,10 +656,10 @@ class TestPromptSeededLeadingSpace:
prompt_ids, output_ids, expected = leading_space_ids prompt_ids, output_ids, expected = leading_space_ids
choices, _, _ = await derenderer.derender_completion( choices, _, _ = await derenderer.derender_completion(
[ [
GenerateResponse( GenerateTokensResponse(
request_id="t", request_id="t",
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, token_ids=output_ids, finish_reason="stop" index=0, token_ids=output_ids, finish_reason="stop"
) )
], ],
@@ -1494,10 +1495,10 @@ class TestDerenderChatStreamHarmony:
content = "".join(d.content or "" for d in deltas) content = "".join(d.content or "" for d in deltas)
batch_choices = await harmony_derenderer.derender_chat( batch_choices = await harmony_derenderer.derender_chat(
GenerateResponse( GenerateTokensResponse(
request_id="test-harmony-batch", request_id="test-harmony-batch",
choices=[ choices=[
GenerateResponseChoice( GenerateTokensChoice(
index=0, token_ids=output_ids, finish_reason="stop" index=0, token_ids=output_ids, finish_reason="stop"
) )
], ],
@@ -1743,10 +1744,10 @@ class TestServingDerenderStreamValidation:
could never fire since derender_chat_stream itself rejects anything could never fire since derender_chat_stream itself rejects anything
above 1 first.""" above 1 first."""
serving = self._make_serving(parser_configured=False) serving = self._make_serving(parser_configured=False)
two_choices = GenerateStreamResponse( two_choices = GenerateTokensStreamResponse(
request_id="t", request_id="t",
choices=[ choices=[
GenerateResponseStreamChoice(index=i, token_ids=[1]) for i in range(2) GenerateTokensStreamChoice(index=i, token_ids=[1]) for i in range(2)
], ],
) )
request = DerenderChatStreamRequest( request = DerenderChatStreamRequest(
+33 -2
View File
@@ -6,6 +6,8 @@ from argparse import Namespace
from fastapi import FastAPI from fastapi import FastAPI
from vllm.entrypoints.launchers.api_server.routers import register_api_routers from vllm.entrypoints.launchers.api_server.routers import register_api_routers
from vllm.entrypoints.scale_out.token_in_token_out.api_router import attach_router
from vllm.entrypoints.serve.middleware.authenticate import GUARDED_PREFIX
from vllm.tasks import SupportedTask from vllm.tasks import SupportedTask
RENDER_PATHS = { RENDER_PATHS = {
@@ -17,7 +19,9 @@ RENDER_PATHS = {
"/v1/responses/render", "/v1/responses/render",
} }
GENERATE_PATH = "/inference/v1/generate" GENERATE_PATH = "/inference/v1/generate"
SCALE_OUT_PATHS = RENDER_PATHS | {GENERATE_PATH} ABORT_PATH = "/inference/v1/abort_requests"
UNAUTHENTICATED_ABORT_PATH = "/abort_requests"
SCALE_OUT_PATHS = RENDER_PATHS | {GENERATE_PATH, ABORT_PATH}
def registered_paths( def registered_paths(
@@ -66,4 +70,31 @@ def test_render_routes_remain_enabled_for_render_server():
def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(): def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset():
paths = registered_paths(("generate",), tokens_only=True) paths = registered_paths(("generate",), tokens_only=True)
assert paths >= {GENERATE_PATH, "/abort_requests"} assert paths >= {GENERATE_PATH, ABORT_PATH, UNAUTHENTICATED_ABORT_PATH}
def test_abort_route_without_tokens_only_requires_api_key():
"""A decode pool that loads a tokenizer serves generate without
--tokens-only and still needs to abort, but only behind --api-key."""
paths = registered_paths(("generate",), enable_scale_out=True, tokens_only=False)
assert ABORT_PATH in paths
assert ABORT_PATH.startswith(GUARDED_PREFIX)
assert UNAUTHENTICATED_ABORT_PATH not in paths
def test_unauthenticated_abort_route_is_not_registered_twice():
"""The RLHF dev router registers /abort_requests before the scale-out
routers, so attaching the scale-out one too would add a shadowed duplicate."""
app = FastAPI()
app.state.args = Namespace(tokens_only=True)
@app.post(UNAUTHENTICATED_ABORT_PATH)
async def existing_abort_requests():
return None
attach_router(app)
paths = [route.path for route in app.routes]
assert paths.count(UNAUTHENTICATED_ABORT_PATH) == 1
assert ABORT_PATH in paths
@@ -20,7 +20,8 @@ from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import (
) )
from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
GenerateRequest, GenerateRequest,
GenerateResponse, GenerateTextResponse,
GenerateTokensResponse,
MultiModalFeatures, MultiModalFeatures,
PlaceholderRangeInfo, PlaceholderRangeInfo,
) )
@@ -144,6 +145,7 @@ def _make_request_output(
num_cached_tokens: int | None = None, num_cached_tokens: int | None = None,
index: int = 0, index: int = 0,
spec_decode_metrics: RequestSpecDecodeMetrics | None = None, spec_decode_metrics: RequestSpecDecodeMetrics | None = None,
text: str = "",
) -> RequestOutput: ) -> RequestOutput:
return RequestOutput( return RequestOutput(
request_id=request_id, request_id=request_id,
@@ -153,7 +155,7 @@ def _make_request_output(
outputs=[ outputs=[
CompletionOutput( CompletionOutput(
index=index, index=index,
text="", text=text,
token_ids=token_ids, token_ids=token_ids,
cumulative_logprob=None, cumulative_logprob=None,
logprobs=logprobs, logprobs=logprobs,
@@ -239,7 +241,7 @@ async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution():
response = await serving.serve_tokens(request) response = await serving.serve_tokens(request)
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
assert ( assert (
serving.online_renderer.preprocess_completion.call_args.kwargs["skip_mm_cache"] serving.online_renderer.preprocess_completion.call_args.kwargs["skip_mm_cache"]
is True is True
@@ -411,7 +413,7 @@ async def test_serve_tokens_returns_spec_decode_metrics(stream: bool):
if isinstance(chunk, dict) and chunk.get("metrics") if isinstance(chunk, dict) and chunk.get("metrics")
)["metrics"]["speculative_decoding"] )["metrics"]["speculative_decoding"]
else: else:
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
assert response.metrics is not None assert response.metrics is not None
assert response.metrics.speculative_decoding is not None assert response.metrics.speculative_decoding is not None
payload = response.metrics.speculative_decoding.model_dump() payload = response.metrics.speculative_decoding.model_dump()
@@ -511,7 +513,7 @@ async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling(
if isinstance(chunk, dict) if isinstance(chunk, dict)
) )
else: else:
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
assert response.metrics is None assert response.metrics is None
@@ -540,7 +542,7 @@ async def test_non_stream_returns_final_prompt_token_ids_when_requested():
response = await serving.serve_tokens(request) response = await serving.serve_tokens(request)
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
assert response.prompt_token_ids == final_prompt_token_ids assert response.prompt_token_ids == final_prompt_token_ids
@@ -622,7 +624,7 @@ async def test_non_stream_prompt_metadata_requires_return_token_ids(
response = await serving.serve_tokens(request) response = await serving.serve_tokens(request)
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
if return_token_ids: if return_token_ids:
assert response.prompt_token_ids == [1, 2, 3] assert response.prompt_token_ids == [1, 2, 3]
assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS
@@ -652,7 +654,7 @@ async def test_non_stream_text_only_prompt_has_no_mm_placeholders():
response = await serving.serve_tokens(request) response = await serving.serve_tokens(request)
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
assert response.prompt_token_ids == [1, 2, 3] assert response.prompt_token_ids == [1, 2, 3]
assert response.mm_placeholders is None assert response.mm_placeholders is None
@@ -678,7 +680,7 @@ async def test_non_stream_encoder_decoder_omits_mm_placeholders():
response = await serving.serve_tokens(request) response = await serving.serve_tokens(request)
assert isinstance(response, GenerateResponse) assert isinstance(response, GenerateTokensResponse)
assert response.prompt_token_ids == [1, 2, 3] assert response.prompt_token_ids == [1, 2, 3]
assert response.mm_placeholders is None assert response.mm_placeholders is None
@@ -1195,3 +1197,251 @@ async def test_stream_sampling_mask_matches_each_token_chunk(with_mask):
[[token, token + 1] for token in choice["token_ids"]] if with_mask else None [[token, token + 1] for token in choice["token_ids"]] if with_mask else None
) )
assert choice["sampling_mask"] == expected assert choice["sampling_mask"] == expected
async def _collect_sse(response) -> list[Any]:
return _parse_sse_chunks([chunk async for chunk in response])
def _text_request(**kwargs) -> GenerateRequest:
kwargs.setdefault("sampling_params", SamplingParams(max_tokens=10))
return GenerateRequest(
token_ids=[1, 2, 3], model=MODEL_NAME, output_mode="text", **kwargs
)
@pytest.mark.asyncio
async def test_non_stream_tokens_mode_omits_text_and_echoes_output_mode():
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output(
"req-1", token_ids=[10], text="ignored", finish_reason="stop", finished=True
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(
GenerateRequest(
token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=10),
model=MODEL_NAME,
)
)
assert isinstance(response, GenerateTokensResponse)
dumped = response.model_dump()
assert dumped["output_mode"] == "tokens"
assert "text" not in dumped["choices"][0]
@pytest.mark.asyncio
async def test_non_stream_text_mode_returns_text_and_token_ids():
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output(
"req-1",
token_ids=[10, 20],
text="Hello",
finish_reason="stop",
finished=True,
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(_text_request())
assert isinstance(response, GenerateTextResponse)
assert response.output_mode == "text"
assert response.choices[0].text == "Hello"
assert response.choices[0].token_ids == [10, 20]
assert response.choices[0].finish_reason == "stop"
assert response.usage is not None
assert response.usage.completion_tokens == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("no_tokenizer", ["tokens_only", "skip_tokenizer_init"])
async def test_text_mode_rejected_without_tokenizer(no_tokenizer: str):
"""The no-op detokenizer returns "", so text mode must fail instead
of answering 200 with empty text."""
engine = _mock_engine()
kwargs: dict[str, Any] = {}
if no_tokenizer == "tokens_only":
kwargs["force_no_detokenize"] = True
else:
engine.model_config.skip_tokenizer_init = True
engine.renderer = _build_renderer(engine.model_config)
engine.generate = MagicMock()
serving = _build_serving_tokens(engine, **kwargs)
response = await serving.serve_tokens(_text_request())
assert isinstance(response, ErrorResponse)
assert response.error.code == 400
assert "requires a tokenizer" in response.error.message
engine.generate.assert_not_called()
@pytest.mark.asyncio
async def test_text_mode_logprobs_have_decoded_tokens_and_bytes():
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output(
"req-1",
token_ids=[10],
text="é",
logprobs=[
{
10: Logprob(logprob=-0.5, decoded_token="é"),
11: Logprob(logprob=-1.5, decoded_token="e"),
}
],
finish_reason="stop",
finished=True,
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(
_text_request(sampling_params=SamplingParams(max_tokens=10, logprobs=2))
)
assert isinstance(response, GenerateTextResponse)
logprobs = response.choices[0].logprobs
assert logprobs is not None and logprobs.content is not None
entry = logprobs.content[0]
assert (entry.token, entry.bytes) == ("é", [195, 169])
assert [(t.token, t.bytes) for t in entry.top_logprobs] == [
("é", [195, 169]),
("e", [101]),
]
@pytest.mark.asyncio
async def test_text_mode_logprobs_keep_placeholders_with_return_tokens_as_ids():
"""--return-tokens-as-token-ids keeps token_id:N placeholders, as it does
on /v1/completions, while text is still decoded."""
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output(
"req-1",
token_ids=[10],
text="é",
logprobs=[{10: Logprob(logprob=-0.5, decoded_token="é")}],
finish_reason="stop",
finished=True,
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine, return_tokens_as_token_ids=True)
response = await serving.serve_tokens(
_text_request(sampling_params=SamplingParams(max_tokens=10, logprobs=1))
)
assert isinstance(response, GenerateTextResponse)
assert response.choices[0].text == "é"
logprobs = response.choices[0].logprobs
assert logprobs is not None and logprobs.content is not None
entry = logprobs.content[0]
assert (entry.token, entry.bytes) == ("token_id:10", None)
assert [t.token for t in entry.top_logprobs] == ["token_id:10"]
@pytest.mark.asyncio
async def test_stream_text_mode_emits_text_deltas():
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output("req-1", token_ids=[10], text="Hel")
yield _make_request_output("req-1", token_ids=[20], text="lo")
yield _make_request_output(
"req-1", token_ids=[30], text=" world", finish_reason="stop", finished=True
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(_text_request(stream=True))
parsed = await _collect_sse(response)
assert parsed[-1] == "[DONE]"
data_chunks = parsed[:-1]
assert {c["output_mode"] for c in data_chunks} == {"text"}
assert [c["choices"][0]["text"] for c in data_chunks] == ["Hel", "lo", " world"]
assert [c["choices"][0]["token_ids"] for c in data_chunks] == [[10], [20], [30]]
assert data_chunks[-1]["choices"][0]["finish_reason"] == "stop"
@pytest.mark.asyncio
async def test_stream_text_mode_emits_chunks_without_new_token_ids():
"""Text flushed from the stop string buffer and the final abort output
carry no token IDs but must still reach the client."""
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output("req-1", token_ids=[10], text="a")
yield _make_request_output("req-1", token_ids=[], text="b")
yield _make_request_output(
"req-1", token_ids=[], text="", finish_reason="abort", finished=True
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(_text_request(stream=True))
data_chunks = (await _collect_sse(response))[:-1]
assert [c["choices"][0]["text"] for c in data_chunks] == ["a", "b", ""]
assert data_chunks[1]["choices"][0]["token_ids"] == []
assert data_chunks[2]["choices"][0]["finish_reason"] == "abort"
@pytest.mark.asyncio
async def test_stream_text_mode_skips_outputs_without_text_or_finish():
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output("req-1", token_ids=[10], text="a")
yield _make_request_output("req-1", token_ids=[], text="")
yield _make_request_output(
"req-1", token_ids=[20], text="b", finish_reason="stop", finished=True
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(_text_request(stream=True))
data_chunks = (await _collect_sse(response))[:-1]
assert [c["choices"][0]["text"] for c in data_chunks] == ["a", "b"]
@pytest.mark.asyncio
async def test_stream_text_mode_usage_chunk_echoes_output_mode():
engine = _mock_engine()
async def mock_generate(*args, **kwargs):
yield _make_request_output(
"req-1", token_ids=[10], text="a", finish_reason="stop", finished=True
)
engine.generate = MagicMock(side_effect=mock_generate)
serving = _build_serving_tokens(engine)
response = await serving.serve_tokens(
_text_request(stream=True, stream_options=StreamOptions(include_usage=True))
)
parsed = await _collect_sse(response)
usage_chunk = parsed[-2]
assert usage_chunk["output_mode"] == "text"
assert usage_chunk["choices"] == []
assert usage_chunk["usage"]["completion_tokens"] == 1
@@ -10,9 +10,21 @@ fail loudly if the validator semantics ever drift.
import json import json
import pytest import pytest
from pydantic import TypeAdapter, ValidationError
from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
DerenderChatRequest,
DerenderChatStreamRequest,
DerenderCompletionRequest,
DerenderCompletionStreamRequest,
GenerateRequest, GenerateRequest,
GenerateResponse,
GenerateStreamResponse,
GenerateTextResponse,
GenerateTextStreamResponse,
GenerateTokensChoice,
GenerateTokensResponse,
GenerateTokensStreamResponse,
MultiModalFeatures, MultiModalFeatures,
PlaceholderRangeInfo, PlaceholderRangeInfo,
) )
@@ -110,3 +122,120 @@ def test_generate_request_rejects_placeholder_outside_prompt():
kwargs_data={"image": ["encoded"]}, kwargs_data={"image": ["encoded"]},
), ),
) )
def test_output_mode_defaults_to_tokens():
assert GenerateRequest.model_validate(_base_payload()).output_mode == "tokens"
def test_output_mode_text_is_accepted():
payload = {**_base_payload(), "output_mode": "text"}
assert GenerateRequest.model_validate(payload).output_mode == "text"
@pytest.mark.parametrize("output_mode", ["derender", "bogus", None])
def test_output_mode_rejects_unsupported_levels(output_mode):
"""Levels the server does not implement fail instead of falling back to
tokens which would answer 200 with the wrong response shape."""
payload = {**_base_payload(), "output_mode": output_mode}
with pytest.raises(ValidationError, match="output_mode"):
GenerateRequest.model_validate(payload)
def test_text_mode_rejects_detokenize_false():
payload = {
"token_ids": [1, 2, 3],
"sampling_params": {"detokenize": False},
"output_mode": "text",
}
with pytest.raises(ValidationError, match="detokenize"):
GenerateRequest.model_validate(payload)
def test_tokens_mode_allows_detokenize_false():
payload = {
"token_ids": [1, 2, 3],
"sampling_params": {"detokenize": False},
"output_mode": "tokens",
}
assert GenerateRequest.model_validate(payload).output_mode == "tokens"
def test_response_without_output_mode_parses_as_tokens():
"""Existing clients and older servers never send output_mode."""
parsed = TypeAdapter(GenerateResponse).validate_python(
{"choices": [{"index": 0, "token_ids": [1]}]}
)
assert isinstance(parsed, GenerateTokensResponse)
assert parsed.output_mode == "tokens"
def test_stream_chunk_without_output_mode_parses_as_tokens():
parsed = TypeAdapter(GenerateStreamResponse).validate_python(
{"choices": [{"index": 0, "token_ids": [1]}]}
)
assert isinstance(parsed, GenerateTokensStreamResponse)
def test_text_response_and_chunk_dispatch_on_output_mode():
response = TypeAdapter(GenerateResponse).validate_python(
{"output_mode": "text", "choices": [{"index": 0, "text": "hi"}]}
)
chunk = TypeAdapter(GenerateStreamResponse).validate_python(
{"output_mode": "text", "choices": [{"index": 0, "text": "hi"}]}
)
assert isinstance(response, GenerateTextResponse)
assert isinstance(chunk, GenerateTextStreamResponse)
def test_text_response_requires_text_on_every_choice():
with pytest.raises(ValidationError, match="text"):
TypeAdapter(GenerateResponse).validate_python(
{"output_mode": "text", "choices": [{"index": 0}]}
)
@pytest.mark.parametrize("output_mode", ["derender", None, 5])
def test_response_rejects_unknown_output_mode(output_mode):
with pytest.raises(ValidationError):
TypeAdapter(GenerateResponse).validate_python(
{"output_mode": output_mode, "choices": []}
)
def test_derender_requests_accept_responses_without_output_mode():
DerenderChatRequest.model_validate(
{"generate_response": {"choices": [{"index": 0, "token_ids": [1]}]}}
)
DerenderChatStreamRequest.model_validate(
{
"stream": True,
"generate_chunk": {"choices": [{"index": 0, "token_ids": [1]}]},
}
)
def test_derender_requests_reject_text_responses():
"""Text responses are already detokenized with stop strings applied, so
re-decoding their token_ids would silently undo the truncation."""
text = {"output_mode": "text", "choices": [{"index": 0, "text": "hi"}]}
with pytest.raises(ValidationError, match="output_mode"):
DerenderChatRequest.model_validate({"generate_response": text})
with pytest.raises(ValidationError, match="output_mode"):
DerenderCompletionRequest.model_validate({"generate_responses": [text]})
with pytest.raises(ValidationError, match="output_mode"):
DerenderChatStreamRequest.model_validate(
{"stream": True, "generate_chunk": text}
)
with pytest.raises(ValidationError, match="output_mode"):
DerenderCompletionStreamRequest.model_validate(
{"stream": True, "generate_chunk": text}
)
def test_tokens_response_echoes_output_mode_without_text():
dumped = GenerateTokensResponse(
choices=[GenerateTokensChoice(index=0, token_ids=[1])]
).model_dump()
assert dumped["output_mode"] == "tokens"
assert "text" not in dumped["choices"][0]
@@ -19,6 +19,7 @@ from vllm.v1.engine.detokenizer import check_stop_strings
MODEL_NAME = "Qwen/Qwen3-0.6B" MODEL_NAME = "Qwen/Qwen3-0.6B"
GEN_ENDPOINT = "/inference/v1/generate" GEN_ENDPOINT = "/inference/v1/generate"
ABORT_ENDPOINT = "/inference/v1/abort_requests"
def get_vocab_size(model_name): def get_vocab_size(model_name):
@@ -317,11 +318,11 @@ async def test_generate_logprobs(client, logprobs_value):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_generate_prompt_token_id_logprobs(client): async def test_generate_prompt_token_id_logprobs(client):
token_ids = [11, 22, 33, 44, 55] token_ids = [11, 22, 33, 44, 55]
candidates = [22, 33, 44, 55] rows = [[33], [44, 100], [55]]
sampling_params = { sampling_params = {
"max_tokens": 1, "max_tokens": 1,
"prompt_logprobs": 0, "prompt_logprobs": 0,
"prompt_logprob_token_ids": candidates, "prompt_logprob_token_ids": rows,
"prompt_logprob_start": 1, "prompt_logprob_start": 1,
} }
payload = { payload = {
@@ -334,14 +335,13 @@ async def test_generate_prompt_token_id_logprobs(client):
data = resp.json() data = resp.json()
scores = np.load(io.BytesIO(base64.b64decode(data["prompt_token_id_logprobs"]))) scores = np.load(io.BytesIO(base64.b64decode(data["prompt_token_id_logprobs"])))
assert scores.shape == (len(token_ids) - 2, len(candidates)) assert scores.shape == (3, 2)
for row, target in enumerate(token_ids[2:]): for row, target in enumerate(token_ids[2:]):
expected = data["prompt_logprobs"][row + 2][str(target)]["logprob"] expected = data["prompt_logprobs"][row + 2][str(target)]["logprob"]
assert scores[row, candidates.index(target)] == pytest.approx( assert scores[row, 0] == pytest.approx(expected, abs=1e-3)
expected, abs=1e-3 assert np.isneginf(scores[[0, 2], 1]).all()
)
stream_only_ids = {"prompt_logprob_token_ids": candidates} stream_only_ids = {"prompt_logprob_token_ids": rows}
resp = await client.post( resp = await client.post(
GEN_ENDPOINT, GEN_ENDPOINT,
json={**payload, "sampling_params": stream_only_ids, "stream": True}, json={**payload, "sampling_params": stream_only_ids, "stream": True},
@@ -349,15 +349,12 @@ async def test_generate_prompt_token_id_logprobs(client):
assert resp.status_code == 400 assert resp.status_code == 400
assert "prompt_logprob_token_ids" in resp.text assert "prompt_logprob_token_ids" in resp.text
past_end = {**sampling_params, "prompt_logprob_start": len(token_ids)} too_few_rows = {**sampling_params, "prompt_logprob_token_ids": rows[:2]}
resp = await client.post( resp = await client.post(
GEN_ENDPOINT, json={**payload, "sampling_params": past_end} GEN_ENDPOINT, json={**payload, "sampling_params": too_few_rows}
) )
resp.raise_for_status() assert resp.status_code == 400
empty = np.load( assert "scored rows" in resp.text
io.BytesIO(base64.b64decode(resp.json()["prompt_token_id_logprobs"]))
)
assert empty.shape == (0, len(candidates))
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -590,3 +587,187 @@ async def test_generate_with_lora_adapter(client, tokenizer, messages):
completions_res = completions_data["choices"][0]["message"]["content"] completions_res = completions_data["choices"][0]["message"]["content"]
assert generate_res == completions_res assert generate_res == completions_res
def _chat_prompt_token_ids(tokenizer, messages) -> list[int]:
return tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
enable_thinking=False, # default with Qwen3
return_dict=True, # default with Transformers v5
).input_ids
def _text_mode_payload(token_ids, **sampling_params) -> dict:
return {
"model": MODEL_NAME,
"token_ids": token_ids,
"sampling_params": {"max_tokens": 24, "temperature": 0.0, **sampling_params},
"output_mode": "text",
"stream": False,
}
@pytest.mark.asyncio
@pytest.mark.skipif(
envs.VLLM_USE_RUST_FRONTEND,
reason="output_mode is not supported by the Rust frontend",
)
async def test_text_mode_matches_chat_completions(client, tokenizer, messages):
token_ids = _chat_prompt_token_ids(tokenizer, messages)
resp = await client.post(GEN_ENDPOINT, json=_text_mode_payload(token_ids))
resp.raise_for_status()
data = resp.json()
assert data["output_mode"] == "text"
payload = {
"model": MODEL_NAME,
"messages": messages,
"max_tokens": 24,
"temperature": 0.0,
"stream": False,
"chat_template_kwargs": {"enable_thinking": False},
}
chat_resp = await client.post("/v1/chat/completions", json=payload)
chat_content = chat_resp.json()["choices"][0]["message"]["content"]
assert data["choices"][0]["text"] == chat_content
assert data["choices"][0]["token_ids"]
@pytest.mark.asyncio
@pytest.mark.skipif(
envs.VLLM_USE_RUST_FRONTEND,
reason="output_mode is not supported by the Rust frontend",
)
async def test_text_mode_stream_matches_non_stream(client, tokenizer, messages):
token_ids = _chat_prompt_token_ids(tokenizer, messages)
payload = _text_mode_payload(token_ids)
non_stream = await client.post(GEN_ENDPOINT, json=payload)
non_stream.raise_for_status()
chunks = []
async with client.stream(
"POST", GEN_ENDPOINT, json={**payload, "stream": True}
) as resp:
resp.raise_for_status()
async for line in resp.aiter_lines():
if line.startswith("data: ") and line != "data: [DONE]":
chunks.append(json.loads(line[len("data: ") :]))
assert {chunk["output_mode"] for chunk in chunks} == {"text"}
assert chunks[-1]["choices"][0]["finish_reason"] is not None
streamed_text = "".join(chunk["choices"][0]["text"] for chunk in chunks)
assert streamed_text == non_stream.json()["choices"][0]["text"]
@pytest.mark.asyncio
@pytest.mark.skipif(
envs.VLLM_USE_RUST_FRONTEND,
reason="output_mode is not supported by the Rust frontend",
)
@pytest.mark.parametrize(
"overrides, expected",
[
pytest.param(
{"sampling_params": {"detokenize": False}}, "detokenize", id="no-detok"
),
pytest.param({"output_mode": "derender"}, "output_mode", id="unsupported"),
],
)
async def test_text_mode_invalid_requests_return_400(client, overrides, expected):
payload = {
**_text_mode_payload([1, 2, 3], max_tokens=5),
**overrides,
}
resp = await client.post(GEN_ENDPOINT, json=payload)
assert resp.status_code == 400
assert expected in resp.json()["error"]["message"]
@pytest.mark.asyncio
@pytest.mark.skipif(
envs.VLLM_USE_RUST_FRONTEND,
reason="--tokens-only is not supported by the Rust frontend",
)
@pytest.mark.parametrize("server", [["--tokens-only"]], indirect=True)
async def test_text_mode_rejected_when_tokens_only(client):
"""Without a tokenizer the detokenizer returns "", so the request must
fail instead of answering 200 with empty text."""
resp = await client.post(GEN_ENDPOINT, json=_text_mode_payload([1, 2, 3]))
assert resp.status_code == 400
assert "requires a tokenizer" in resp.json()["error"]["message"]
payload = {**_text_mode_payload([1, 2, 3]), "output_mode": "tokens"}
payload["sampling_params"]["detokenize"] = False
resp = await client.post(GEN_ENDPOINT, json=payload)
resp.raise_for_status()
@pytest.mark.asyncio
@pytest.mark.skipif(
envs.VLLM_USE_RUST_FRONTEND,
reason="output_mode is not supported by the Rust frontend",
)
async def test_text_mode_stream_delivers_abort_finish_chunk(
client, tokenizer, messages
):
"""The final output after an abort has no new token IDs, so it
reaches a text stream only because text mode emits finish-only chunks."""
payload = {
**_text_mode_payload(
_chat_prompt_token_ids(tokenizer, messages),
max_tokens=900,
ignore_eos=True,
),
"stream": True,
"request_id": "text-abort-e2e",
}
chunks = []
aborted = False
async with client.stream("POST", GEN_ENDPOINT, json=payload) as resp:
resp.raise_for_status()
async for line in resp.aiter_lines():
if not line.startswith("data: ") or line == "data: [DONE]":
continue
chunk = json.loads(line[len("data: ") :])
chunks.append(chunk)
if not aborted and chunk["choices"] and chunk["choices"][0]["text"]:
abort = await client.post(
ABORT_ENDPOINT,
json={"request_ids": [chunk["request_id"]]},
)
assert abort.status_code == 200
aborted = True
assert aborted
assert {chunk["output_mode"] for chunk in chunks} == {"text"}
final_choice = chunks[-1]["choices"][0]
assert final_choice["finish_reason"] == "abort"
generated = sum(len(chunk["choices"][0]["token_ids"] or []) for chunk in chunks)
assert generated < 900
@pytest.mark.asyncio
@pytest.mark.skipif(
envs.VLLM_USE_RUST_FRONTEND,
reason="the Rust frontend does not serve /inference/v1/abort_requests",
)
async def test_abort_requests_is_served_without_tokens_only(client):
resp = await client.post(
ABORT_ENDPOINT, json={"request_ids": ["generate-tokens-unknown"]}
)
assert resp.status_code == 200
resp = await client.post(ABORT_ENDPOINT, json={})
assert resp.status_code == 400
# The unauthenticated path is only served with --tokens-only.
resp = await client.post(
"/abort_requests", json={"request_ids": ["generate-tokens-unknown"]}
)
assert resp.status_code == 404
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
from vllm.logprobs import Logprob from vllm.logprobs import Logprob
@@ -36,3 +38,50 @@ def test_logprobs_minus_one_emits_all_tokens():
num_output_top_logprobs=-1, num_output_top_logprobs=-1,
) )
assert len(result.content[0].top_logprobs) == 2 assert len(result.content[0].top_logprobs) == 2
def test_tokens_level_logprobs_have_no_bytes():
result = ServingTokens._create_tokens_logprobs(
None,
token_ids=[7],
top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}],
num_output_top_logprobs=2,
)
entry = result.content[0]
assert entry.bytes is None
assert all(t.bytes is None for t in entry.top_logprobs)
def test_resolved_logprobs_use_engine_decoded_tokens_and_bytes():
tokenizer = MagicMock()
result = ServingTokens._create_tokens_logprobs(
None,
token_ids=[7],
top_logprobs=[
{7: Logprob(-0.9, decoded_token="é"), 8: Logprob(-1.1, decoded_token="e")}
],
num_output_top_logprobs=2,
tokenizer=tokenizer,
)
entry = result.content[0]
assert (entry.token, entry.bytes) == ("é", [195, 169])
assert [(t.token, t.bytes) for t in entry.top_logprobs] == [
("é", [195, 169]),
("e", [101]),
]
tokenizer.decode.assert_not_called()
def test_resolved_logprobs_decode_tokens_the_engine_did_not():
tokenizer = MagicMock()
tokenizer.decode.side_effect = lambda ids: f"<{ids[0]}>"
result = ServingTokens._create_tokens_logprobs(
None,
token_ids=[7, 9],
top_logprobs=[{7: Logprob(-0.9)}, None],
num_output_top_logprobs=1,
tokenizer=tokenizer,
)
assert [e.token for e in result.content] == ["<7>", "<9>"]
assert result.content[0].top_logprobs[0].token == "<7>"
assert result.content[1].bytes == list(b"<9>")
@@ -14,12 +14,6 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.assets.audio import AudioAsset from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio from vllm.multimodal.media.audio import load_audio
# Increase engine iteration timeout for ROCm where first-use JIT compilation
# can exceed the default 60s, causing a silent deadlock in feed_tokens.
REALTIME_ENV_OVERRIDES = {
"VLLM_ENGINE_ITERATION_TIMEOUT_S": "600",
}
MISTRAL_FORMAT_ARGS = [ MISTRAL_FORMAT_ARGS = [
"--tokenizer_mode", "--tokenizer_mode",
"mistral", "mistral",
@@ -77,9 +71,7 @@ async def test_multi_chunk_streaming(model_name, mary_had_lamb_audio_chunks):
if model_name.startswith("mistralai"): if model_name.startswith("mistralai"):
server_args += MISTRAL_FORMAT_ARGS server_args += MISTRAL_FORMAT_ARGS
with RemoteOpenAIServer( with RemoteOpenAIServer(model_name, server_args) as remote_server:
model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES
) as remote_server:
ws_url = _get_websocket_url(remote_server) ws_url = _get_websocket_url(remote_server)
async with websockets.connect(ws_url) as ws: async with websockets.connect(ws_url) as ws:
# Receive session.created # Receive session.created
@@ -183,9 +175,7 @@ async def test_empty_commit_does_not_crash_engine(
if model_name.startswith("mistralai"): if model_name.startswith("mistralai"):
server_args += MISTRAL_FORMAT_ARGS server_args += MISTRAL_FORMAT_ARGS
with RemoteOpenAIServer( with RemoteOpenAIServer(model_name, server_args) as remote_server:
model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES
) as remote_server:
ws_url = _get_websocket_url(remote_server) ws_url = _get_websocket_url(remote_server)
# --- First connection: empty commit (no audio appended) ---------- # --- First connection: empty commit (no audio appended) ----------
@@ -272,9 +262,7 @@ async def test_session_update_invalid_model_returns_error(model_name):
if model_name.startswith("mistralai"): if model_name.startswith("mistralai"):
server_args += MISTRAL_FORMAT_ARGS server_args += MISTRAL_FORMAT_ARGS
with RemoteOpenAIServer( with RemoteOpenAIServer(model_name, server_args) as remote_server:
model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES
) as remote_server:
ws_url = _get_websocket_url(remote_server) ws_url = _get_websocket_url(remote_server)
async with websockets.connect(ws_url) as ws: async with websockets.connect(ws_url) as ws:
event = await receive_event(ws, timeout=30.0) event = await receive_event(ws, timeout=30.0)
@@ -301,9 +289,7 @@ async def test_commit_without_session_update_returns_error(model_name):
if model_name.startswith("mistralai"): if model_name.startswith("mistralai"):
server_args += MISTRAL_FORMAT_ARGS server_args += MISTRAL_FORMAT_ARGS
with RemoteOpenAIServer( with RemoteOpenAIServer(model_name, server_args) as remote_server:
model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES
) as remote_server:
ws_url = _get_websocket_url(remote_server) ws_url = _get_websocket_url(remote_server)
async with websockets.connect(ws_url) as ws: async with websockets.connect(ws_url) as ws:
event = await receive_event(ws, timeout=30.0) event = await receive_event(ws, timeout=30.0)
@@ -107,3 +107,37 @@ def test_fp8_index_query_and_cache_prefill_topk_matches_bf16() -> None:
return {int(b) for b in idx[0, 0].tolist() if b >= 0} return {int(b) for b in idx[0, 0].tolist() if b >= 0}
_check_selection(select(keys), select(keys.to(torch.float8_e4m3fn))) _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn)))
@pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required")
@torch.inference_mode()
def test_bf16_query_against_fp8_cache_decode_topk_matches_bf16() -> None:
"""A mixed pair must upcast rather than emit a mixed-dtype fp8 dot.
Production allocates index_q with the index-K cache dtype, so the scorers
normally see a matched pair and keep the native FP8 MMA. This pins the
other branch: a bf16 query against an e4m3 cache still compiles and still
ranks the planted blocks.
"""
dev = "cuda"
q, keys = _make_inputs(dev)
block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0)
seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev)
def select(cache: torch.Tensor) -> set[int]:
idx = minimax_m3_index_decode(
q, # bf16 query, deliberately not cast to the cache dtype
cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(),
block_table,
seq_lens,
SEQ_LEN,
TOPK,
0,
0,
HEADS,
1,
1,
)
return {int(b) for b in idx[0, 0].tolist() if b >= 0}
_check_selection(select(keys), select(keys.to(torch.float8_e4m3fn)))
@@ -596,6 +596,114 @@ def test_sparse_attn_prefill_ragged_kernel() -> None:
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict:
device = torch.device("cuda")
set_random_seed(7)
head_dim = nope_dim + rope_dim
return dict(
q=torch.randn(3, 3, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
kv=torch.randn(5, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
indices=torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device),
indptr=torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device),
scale=head_dim**-0.5,
attn_sink=torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device),
nope_head_dim=nope_dim,
rope_head_dim=rope_dim,
)
# 448+64 is DeepSeek V4/V4.1, 512+64 is V3.2, and 256+0 stands in for the
# NoPE-only layouts where the destination spans the whole head dim.
SPARSE_PREFILL_DIMS = [(NOPE_HEAD_DIM, ROPE_HEAD_DIM), (512, 64), (256, 0)]
@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_out_matches_allocated(
nope_dim: int, rope_dim: int
) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
allocated = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)
q = kwargs["q"]
dest = torch.empty(q.shape[0], q.shape[1], nope_dim, dtype=q.dtype, device=q.device)
returned = _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
assert returned is dest
torch.testing.assert_close(dest, allocated[..., :nope_dim], atol=0.0, rtol=0.0)
@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_out_view_leaves_neighbors_untouched(
nope_dim: int, rope_dim: int
) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
q = kwargs["q"]
num_queries, num_heads = q.shape[0], q.shape[1]
buffer = torch.full(
(num_queries + 4, num_heads + 2, nope_dim),
-7.0,
dtype=q.dtype,
device=q.device,
)
dest = buffer[2 : 2 + num_queries, :num_heads]
assert not dest.is_contiguous()
_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :nope_dim]
torch.testing.assert_close(dest, expected, atol=0.0, rtol=0.0)
untouched = torch.ones_like(buffer, dtype=torch.bool)
untouched[2 : 2 + num_queries, :num_heads] = False
assert torch.all(buffer[untouched] == -7.0)
@pytest.mark.parametrize("out_dtype", [torch.float32, torch.float16])
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_out_casts_to_dest_dtype(
out_dtype: torch.dtype,
) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(512, 64)
q = kwargs["q"]
dest = torch.empty(q.shape[0], q.shape[1], 512, dtype=out_dtype, device=q.device)
_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :512]
assert dest.dtype == out_dtype
torch.testing.assert_close(dest.float(), expected.float(), atol=2e-2, rtol=2e-2)
@torch.inference_mode()
def test_sparse_attn_prefill_head_dim_wide_out_keeps_rope_rows() -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(NOPE_HEAD_DIM, ROPE_HEAD_DIM)
q = kwargs["q"]
dest = torch.empty_like(q)
_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
expected = _ref_sparse_prefill_ragged(
q, kwargs["kv"], [[0, 2], [1, 3, 4], []], kwargs["scale"], kwargs["attn_sink"]
)
torch.testing.assert_close(dest, expected, atol=2e-2, rtol=2e-2)
@pytest.mark.parametrize( @pytest.mark.parametrize(
("num_queries", "on_gfx950", "expected"), ("num_queries", "on_gfx950", "expected"),
[(1023, True, False), (1024, True, True), (1024, False, False)], [(1023, True, False), (1024, True, True), (1024, False, False)],
@@ -718,10 +826,11 @@ def test_sparse_attn_prefill_preserves_dense_triton_fallback(monkeypatch) -> Non
output = torch.empty_like(q) output = torch.empty_like(q)
dense_fallback_calls = 0 dense_fallback_calls = 0
def fake_dense_fallback(*args, **kwargs): def fake_dense_fallback(*args, out, **kwargs):
nonlocal dense_fallback_calls nonlocal dense_fallback_calls
dense_fallback_calls += 1 dense_fallback_calls += 1
return torch.zeros_like(q) out.zero_()
return out
monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True) monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None) monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None)
+1 -1
View File
@@ -194,7 +194,7 @@ def test_fused_forward_uses_packed_entrypoint() -> None:
ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device) ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device)
layer = types.SimpleNamespace( layer = types.SimpleNamespace(
prefix=PREFIX, prefix=PREFIX,
enable_fused_gdn_decode=True, enable_fused_gdn_spec_decode=True,
norm=types.SimpleNamespace( norm=types.SimpleNamespace(
weight=torch.empty(V, dtype=torch.bfloat16, device=device) weight=torch.empty(V, dtype=torch.bfloat16, device=device)
), ),
+1 -1
View File
@@ -167,7 +167,7 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams)
draft = SimpleNamespace( draft = SimpleNamespace(
use_mega_moe=False, use_native_mega_moe=False,
use_sequence_parallel=False, use_sequence_parallel=False,
hc_mult=hc_mult, hc_mult=hc_mult,
layers=[make_layer(mix) for mix in mixes], layers=[make_layer(mix) for mix in mixes],
@@ -88,6 +88,15 @@ class Config:
torch_trace_dir_path: str | None = None torch_trace_dir_path: str | None = None
# Force AiterExperts's hidden_pad/intermediate_pad computation
# (`experts/rocm_aiter_moe.py`) to diverge from the padded K/N sizes above.
# None (default) preserves today's behavior: FusedMoEConfig defaults both
# to the (unpadded) K/intermediate_size_per_partition, so hidden_pad and
# intermediate_pad come out to 0.
# See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").
hidden_dim_unpadded: int | None = None
intermediate_size_per_partition_unpadded: int | None = None
def __post_init__(self): def __post_init__(self):
if self.quant_config is None: if self.quant_config is None:
self.quant_config = TestMoEQuantConfig(None, False, False, None) self.quant_config = TestMoEQuantConfig(None, False, False, None)
@@ -664,6 +673,10 @@ def make_modular_kernel(
activation=config.activation, activation=config.activation,
device=vllm_config.device_config.device, device=vllm_config.device_config.device,
routing_method=RoutingMethodType.DeepSeekV3, routing_method=RoutingMethodType.DeepSeekV3,
hidden_dim_unpadded=config.hidden_dim_unpadded,
intermediate_size_per_partition_unpadded=(
config.intermediate_size_per_partition_unpadded
),
) )
prepare_finalize = maybe_make_prepare_finalize( prepare_finalize = maybe_make_prepare_finalize(
+2
View File
@@ -295,6 +295,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
expert_map=None, expert_map=None,
apply_router_weight_on_input=True, apply_router_weight_on_input=True,
routed_scaling_factor=1.0, routed_scaling_factor=1.0,
routing_sink=None,
) )
check_accuracy( check_accuracy(
@@ -968,6 +969,7 @@ def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_in
router_logits=logits, router_logits=logits,
e_score_correction_bias=bias, e_score_correction_bias=bias,
routed_scaling_factor=scale, routed_scaling_factor=scale,
routing_sink=None,
**common, **common,
) )
topk_weights, topk_ids = fused_topk_bias( topk_weights, topk_ids = fused_topk_bias(
+307
View File
@@ -0,0 +1,307 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end tests for the FlashInfer MoE-EP adapter on real Blackwell GPUs.
Each test spawns one process per GPU, bootstraps vLLM's expert-parallel
groups, builds a ``FlashInferMoeEp`` adapter from bf16 expert weights at the
MoE geometry of a real model, warms it up and runs a forward whose inputs are
chosen so that every quantization step in the megakernel is exact, for the
CuTeDSL kernel (NVFP4 activations and weights, E4M3 scales per 16) as well as
the DeepGEMM kernel (FP8 activations and FP4 weights, power-of-two scales per
32). The expected output is then known in closed form and compared bit for
bit.
"""
from dataclasses import dataclass
import pytest
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.distributed.eplb_utils import distributed_run, set_env_vars_and_device
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM
from vllm.distributed.parallel_state import ensure_model_parallel_initialized
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import (
FlashInferMoeEpExperts,
)
from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
FlashInferMoeEp,
FlashInferMoeEpWeights,
apply_topk_in_fc1,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_deep_gemm
@dataclass(frozen=True)
class Geometry:
name: str
hidden_size: int
intermediate_size: int
num_experts: int
top_k: int
# Routed-expert geometry of the models the branch was evaluated on. DeepGEMM
# keeps one UE8M0 scale byte per 32 elements and requires the per-token scale
# rows to be TMA aligned, so hidden and intermediate sizes must be multiples
# of 512 (FlashInfer's own validation only checks 128); all four qualify.
GEOMETRIES = (
Geometry("dsv4-pro", 7168, 3072, 384, 6),
Geometry("dsv4-flash", 4096, 2048, 256, 6),
Geometry("mistral-small-4", 4096, 2048, 128, 4),
Geometry("mistral-large-3", 7168, 4096, 128, 4),
)
BACKENDS = (FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM)
NUM_TOKENS_PER_RANK = 4
MAX_TOKENS_PER_RANK = 128
# Every block scale below is either 1 or a power of two, so the FP4 codes
# and the FP8 values survive quantization exactly: NVFP4 scales a 16-block by
# absmax / 6 in E4M3, DeepGEMM scales a 32-block by absmax / 6 (FP4) or
# absmax / 448 (FP8) rounded up to a power of two. A 6.0 in a block pins the
# FP4 scale to 1; 2.0, 6.0 and 144 are exact FP8 values at any such scale.
#
# Token t carries X_SIGNAL at hidden column 16 * t and a 6.0 sentinel next to
# it; every weight column the sentinel meets is zero.
X_SIGNAL = 2.0
X_SENTINEL = 6.0
# Gate and up rows are identical, so h = silu(12) * 12 = 144 = 6 * 24
# whichever half the kernel treats as the gate, and 144 requantizes exactly.
W13_DIAG = 6.0
H_SIGNAL = 144.0
# Expert e writes E2M1_CODES[e % 7] * H_SIGNAL into output channel e. The 6.0
# sentinel sits next to a zero intermediate channel and pins the block scale.
W2_SENTINEL = 6.0
E2M1_CODES = (0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
def _expert_out(expert: int) -> float:
return E2M1_CODES[expert % len(E2M1_CODES)]
def _topk_weights(top_k: int) -> tuple[float, ...]:
# The DeepSeek-V4 routing policy applies the router weight before fc2, so
# weight * 144 must requantize exactly too: powers of two do.
return tuple(2.0 ** -(k + 1) for k in range(top_k))
def _routing_method(backend: str) -> RoutingMethodType:
if backend == FLASHINFER_MOE_EP_DEEP_GEMM:
return RoutingMethodType.DeepseekV4
return RoutingMethodType.Default
def _routing(
geometry: Geometry, rank: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]:
"""top_k distinct experts per token, strided across ranks so dispatch goes
remote."""
base = (
torch.arange(NUM_TOKENS_PER_RANK, device=device) + rank * NUM_TOKENS_PER_RANK
) % geometry.num_experts
stride = geometry.num_experts // geometry.top_k
offsets = torch.arange(geometry.top_k, device=device) * stride
topk_ids = ((base[:, None] + offsets[None, :]) % geometry.num_experts).to(
torch.int32
)
topk_weights = torch.tensor(
[_topk_weights(geometry.top_k)] * NUM_TOKENS_PER_RANK,
dtype=torch.float32,
device=device,
)
return topk_ids, topk_weights
def _hidden_states(geometry: Geometry, device: torch.device) -> torch.Tensor:
x = torch.zeros(
NUM_TOKENS_PER_RANK, geometry.hidden_size, dtype=torch.bfloat16, device=device
)
for t in range(NUM_TOKENS_PER_RANK):
x[t, 16 * t] = X_SIGNAL
x[t, 16 * t + 1] = X_SENTINEL
return x
def _local_expert_weights(
geometry: Geometry, rank: int, num_local_experts: int, device: torch.device
) -> FlashInferMoeEpWeights:
w13 = torch.zeros(
num_local_experts,
2 * geometry.intermediate_size,
geometry.hidden_size,
dtype=torch.bfloat16,
device=device,
)
w2 = torch.zeros(
num_local_experts,
geometry.hidden_size,
geometry.intermediate_size,
dtype=torch.bfloat16,
device=device,
)
token_columns = [16 * t for t in range(NUM_TOKENS_PER_RANK)]
for local in range(num_local_experts):
expert = rank * num_local_experts + local
w13[local, 0, token_columns] = W13_DIAG
w13[local, geometry.intermediate_size, token_columns] = W13_DIAG
w2[local, expert, 0] = _expert_out(expert)
w2[local, expert, 2] = W2_SENTINEL
return FlashInferMoeEpWeights(w13=w13, w2=w2)
def _expected_output(
geometry: Geometry, topk_ids: torch.Tensor, topk_weights: torch.Tensor
) -> torch.Tensor:
expected = torch.zeros(
NUM_TOKENS_PER_RANK,
geometry.hidden_size,
dtype=torch.float32,
device=topk_ids.device,
)
for t in range(NUM_TOKENS_PER_RANK):
for k in range(geometry.top_k):
expert = int(topk_ids[t, k])
weight = float(topk_weights[t, k])
expected[t, expert] += weight * _expert_out(expert) * H_SIGNAL
return expected.to(torch.bfloat16)
def _moe_config(
vllm_config: VllmConfig,
geometry: Geometry,
world_size: int,
device: torch.device,
backend: str,
) -> FusedMoEConfig:
parallel = FusedMoEParallelConfig.make(
tp_size_=1,
pcp_size_=1,
dp_size_=world_size,
sp_size_=1,
vllm_parallel_config=vllm_config.parallel_config,
)
return FusedMoEConfig(
num_experts=geometry.num_experts,
experts_per_token=geometry.top_k,
hidden_dim=geometry.hidden_size,
intermediate_size=geometry.intermediate_size,
num_local_experts=geometry.num_experts // world_size,
num_logical_experts=geometry.num_experts,
activation=MoEActivation.SILU,
device=device,
routing_method=_routing_method(backend),
moe_parallel_config=parallel,
in_dtype=torch.bfloat16,
moe_backend=backend,
max_num_tokens=MAX_TOKENS_PER_RANK,
)
def _check_expert_views_alias_kernel_memory(
adapter: FlashInferMoeEp,
geometry: Geometry,
rank: int,
num_local_experts: int,
device: torch.device,
) -> None:
"""EPLB moves experts through ``expert_weight_views``; the kernel must see it.
Swap local experts 0 and 1 in place on every rank, then route through the
swapped physical slots: the output must be the one the unswapped logical
experts would have produced.
"""
for view in adapter.expert_weight_views():
first = view[0].clone()
view[0].copy_(view[1])
view[1].copy_(first)
torch.distributed.barrier()
topk_ids, topk_weights = _routing(geometry, rank, device)
slot = topk_ids % num_local_experts
physical_ids = torch.where(slot == 0, topk_ids + 1, topk_ids)
physical_ids = torch.where(slot == 1, topk_ids - 1, physical_ids)
output = adapter(_hidden_states(geometry, device), physical_ids, topk_weights)
torch.accelerator.synchronize()
assert torch.equal(output, _expected_output(geometry, topk_ids, topk_weights))
def _init_warmup_forward(
env: dict[str, str], world_size: int, geometry: Geometry, backend: str
) -> None:
set_env_vars_and_device(env)
vllm_config = VllmConfig()
vllm_config.parallel_config.data_parallel_size = world_size
vllm_config.parallel_config.enable_expert_parallel = True
vllm_config.kernel_config.moe_backend = backend
with set_current_vllm_config(vllm_config):
ensure_model_parallel_initialized(
tensor_model_parallel_size=1, pipeline_model_parallel_size=1
)
rank = torch.distributed.get_rank()
device = torch.device(f"cuda:{rank}")
moe = _moe_config(vllm_config, geometry, world_size, device, backend)
supported, reason = FlashInferMoeEpExperts.is_supported_config(
FlashInferMoeEpExperts,
moe,
kMxfp4Static,
None,
mk.FusedMoEActivationFormat.Standard,
)
assert supported, reason
adapter = FlashInferMoeEp(
moe,
_local_expert_weights(geometry, rank, moe.num_local_experts, device),
apply_topk_in_fc1=apply_topk_in_fc1(moe),
)
try:
adapter.warmup()
topk_ids, topk_weights = _routing(geometry, rank, device)
output = adapter(_hidden_states(geometry, device), topk_ids, topk_weights)
torch.accelerator.synchronize()
expected = _expected_output(geometry, topk_ids, topk_weights)
assert output.dtype == expected.dtype and output.shape == expected.shape
assert torch.equal(output, expected), (
f"rank {rank}: max abs diff "
f"{(output.float() - expected.float()).abs().max().item()}, "
f"nonzero output channels per token "
f"{[torch.nonzero(row).flatten().tolist() for row in output]}"
)
if backend == FLASHINFER_MOE_EP_CUTEDSL:
_check_expert_views_alias_kernel_memory(
adapter, geometry, rank, moe.num_local_experts, device
)
finally:
adapter.destroy()
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="The FlashInfer MoE-EP megakernels require SM100",
)
@pytest.mark.parametrize("geometry", GEOMETRIES, ids=[g.name for g in GEOMETRIES])
@pytest.mark.parametrize("backend", BACKENDS, ids=("cutedsl", "deep_gemm"))
@pytest.mark.parametrize("world_size", [2, 4])
def test_init_warmup_forward_is_exact_on_grid_inputs(
world_size: int, backend: str, geometry: Geometry
):
pytest.importorskip("flashinfer.moe_ep")
if backend == FLASHINFER_MOE_EP_DEEP_GEMM and not has_deep_gemm():
pytest.skip("DeepGEMM is not available")
if torch.accelerator.device_count() < world_size:
pytest.skip(f"Need at least {world_size} GPUs")
distributed_run(_init_warmup_forward, world_size, geometry, backend)
@@ -42,7 +42,7 @@ from .modular_kernel_tools.parallel_utils import (
ProcessGroupInfo, ProcessGroupInfo,
parallel_launch_with_config, parallel_launch_with_config,
) )
from .utils import check_accuracy from .utils import check_accuracy, make_test_weights
has_any_multi_gpu_package = ( has_any_multi_gpu_package = (
has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe() has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe()
@@ -691,6 +691,303 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
) )
# --- 4b: hidden_dim_unpadded/intermediate_size_per_partition_unpadded matrix -
# K/N intentionally unaligned to AITER's 64/128 granularity (unlike the
# shared mk_objects.py defaults) so hidden_pad/intermediate_pad are forced.
_PADDING_E = 8
_PADDING_M = 16
_PADDING_TOPK = 2
_PADDING_K_UNPADDED = 896 # 7 * 128
_PADDING_K_PADDED = 1024 # +128
_PADDING_N_UNPADDED = 384 # 3 * 128
_PADDING_N_PADDED = 512 # +128
# hidden_pad = (K_padded - hidden_dim_unpadded) // 128 * 128 (rocm_aiter_moe.py)
_PADDING_HIDDEN_PAD_EXPECTED = 128
# intermediate_pad = (N_padded - intermediate_size_per_partition_unpadded)
# // 64 * 64 * (2 if tp_size == 1 else 1) (rocm_aiter_moe.py); tp_size==1 here.
_PADDING_INTERMEDIATE_PAD_EXPECTED = 256
# {no padding, hidden-only, intermediate-only, both} -> (pad_hidden, pad_intermediate)
_PADDING_MODES: dict[str, tuple[bool, bool]] = {
"none": (False, False),
"hidden": (True, False),
"intermediate": (False, True),
"both": (True, True),
}
# unquantized + the 4 fp8 configs AiterExperts supports (excludes MXFP4,
# stubbed below pending MI350/gfx950, and MK_QUANT_CONFIGS[1], unsupported).
_PADDING_QUANT_CONFIGS = [
MK_QUANT_CONFIGS[0], # unquantized
MK_QUANT_CONFIGS[2], # fp8 channel weights / per-token activations
MK_QUANT_CONFIGS[3], # fp8 per-tensor weights / per-tensor activations
MK_QUANT_CONFIGS[4], # fp8 per-tensor weights / per-token activations
MK_QUANT_CONFIGS[5], # fp8 128x128-block weights / 128-block activations
]
_PADDING_QUANT_IDS = [
"unquantized",
"fp8_channel_token",
"fp8_tensor_tensor",
"fp8_tensor_token",
"fp8_block_token",
]
def _slice_gate_up_rows(t: torch.Tensor, real_per_half: int) -> torch.Tensor:
"""Slice a (E, 2*padded_per_half, ...) tensor's dim=1 down to the real
gate/up halves of `real_per_half` size each."""
padded_per_half = t.shape[1] // 2
return torch.cat(
[
t[:, :real_per_half],
t[:, padded_per_half : padded_per_half + real_per_half],
],
dim=1,
)
def _slice_unpadded_weights(
weights: WeightTensors,
quant_config: TestMoEQuantConfig | None,
k_unpadded: int,
n_unpadded: int,
) -> WeightTensors:
"""Derive the unpadded sub-block of a padded WeightTensors: same values,
restricted to the region hidden_pad/intermediate_pad keeps.
Scale-shape choice (per_out_ch) must match _make_padding_matrix_weights()."""
block_shape = quant_config.block_shape if quant_config is not None else None
per_out_ch = quant_config is not None and quant_config.per_out_ch_quant
w1 = _slice_gate_up_rows(weights.w1[:, :, :k_unpadded], n_unpadded)
w2 = weights.w2[:, :k_unpadded, :n_unpadded]
if weights.w1_scale is None:
w1_scale = w2_scale = None
elif weights.w2_scale is None:
raise AssertionError("w1_scale and w2_scale must both be set or both None")
elif block_shape is not None:
block_n, block_k = block_shape
assert n_unpadded % block_n == 0 and k_unpadded % block_k == 0
n_scale, k_scale = n_unpadded // block_n, k_unpadded // block_k
w1_scale = _slice_gate_up_rows(weights.w1_scale, n_scale)[..., :k_scale]
w2_scale = weights.w2_scale[:, :k_scale, :n_scale]
elif per_out_ch:
w1_scale = _slice_gate_up_rows(weights.w1_scale, n_unpadded)
w2_scale = weights.w2_scale[:, :k_unpadded, :]
else:
# Per-tensor scale: shape (E, 1, 1), independent of K/N -- reuse as-is.
w1_scale = weights.w1_scale
w2_scale = weights.w2_scale
return WeightTensors(w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale)
def _make_padding_matrix_weights(config: Config) -> WeightTensors:
"""Like WeightTensors.make(), but scales weights by
config.is_per_out_ch_quant instead of is_per_act_token_quant.
WeightTensors.make() ties weight-scale shape to activation quant, so
fp8_tensor_token would get per-channel scales despite being a per-tensor
weight scheme. Scoped here rather than fixing WeightTensors.make(),
which other tests depend on."""
(_, w1, w1_scale, w1_gs), (_, w2, w2_scale, w2_gs) = make_test_weights(
e=config.E,
n=config.N,
k=config.K,
in_dtype=config.dtype,
quant_dtype=config.quant_dtype,
block_shape=config.quant_block_shape,
per_out_ch_quant=config.is_per_out_ch_quant,
)
return WeightTensors(
w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale, w1_gs=w1_gs, w2_gs=w2_gs
)
def _slice_unpadded_rank_tensors(
rank_tensors: RankTensors, k_unpadded: int
) -> RankTensors:
return RankTensors(
hidden_states=rank_tensors.hidden_states[:, :k_unpadded].contiguous(),
hidden_states_scale=rank_tensors.hidden_states_scale,
topk_weights=rank_tensors.topk_weights,
topk_ids=rank_tensors.topk_ids,
expert_map=rank_tensors.expert_map,
)
def _aiter_padding_matrix_worker(
pgi: ProcessGroupInfo,
vllm_config: VllmConfig,
cpu_group,
padded_config: Config,
padded_weights: WeightTensors,
verbose: bool,
hidden_pad_expected: int,
intermediate_pad_expected: int,
):
device = torch.device(f"cuda:{pgi.local_rank}")
init_workspace_manager(device)
set_random_seed(pgi.rank)
weights = copy.deepcopy(padded_weights)
weights.to_current_device()
rank_tensors = RankTensors.make(padded_config, pgi)
with mock.patch.object(
rocm_aiter_ops, "fused_moe", wraps=rocm_aiter_ops.fused_moe
) as fused_moe_mock:
mk_out = run_modular_kernel(
pgi, vllm_config, padded_config, weights, rank_tensors
)
assert fused_moe_mock.call_count > 0, (
"Expected AiterExperts.apply() to call rocm_aiter_ops.fused_moe."
)
for call in fused_moe_mock.call_args_list:
assert call.kwargs["hidden_pad"] == hidden_pad_expected, (
f"AiterExperts.apply() forwarded hidden_pad="
f"{call.kwargs['hidden_pad']}, expected {hidden_pad_expected}."
)
assert call.kwargs["intermediate_pad"] == intermediate_pad_expected, (
f"AiterExperts.apply() forwarded intermediate_pad="
f"{call.kwargs['intermediate_pad']}, expected "
f"{intermediate_pad_expected}."
)
# AiterExperts.apply() never slices its own output -- the buffer stays
# raw (padded) width; only the caller-sliced unpadded prefix is checked.
assert mk_out.shape[-1] == padded_config.K, (
f"AiterExperts output width {mk_out.shape[-1]} != raw hidden_dim "
f"{padded_config.K}."
)
mk_out = mk_out[..., :_PADDING_K_UNPADDED]
# Only the real slice is checked: beyond it is unspecified/reused memory.
assert torch.isfinite(mk_out).all(), (
"AiterExperts output must not contain NaN/Inf in the real "
"hidden_dim_unpadded-wide output slice."
)
unpadded_weights = _slice_unpadded_weights(
weights,
padded_config.quant_config,
_PADDING_K_UNPADDED,
_PADDING_N_UNPADDED,
)
unpadded_rank_tensors = _slice_unpadded_rank_tensors(
rank_tensors, _PADDING_K_UNPADDED
)
with set_current_vllm_config(vllm_config):
ref_out = reference_moe_impl(
padded_config, unpadded_weights, unpadded_rank_tensors
)
# ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out
# mk_out would still pass check_accuracy below -- guard against that.
ref_scale = ref_out.abs().mean()
mk_scale = mk_out.abs().mean()
assert mk_scale > 0.5 * ref_scale, (
f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks "
f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})."
)
# Lenient check for every scheme, including unquantized: a strict
# assert_close(atol=rtol=3e-2) spuriously fails ~0.1-0.2% of elements
# on MI300 even for unquantized AiterExperts (confirmed on hardware).
check_accuracy(ref_out, mk_out, atol=3e-2, rtol=3e-2, percent=0.9)
@require_aiter_moe
@pytest.mark.parametrize("quant_config", _PADDING_QUANT_CONFIGS, ids=_PADDING_QUANT_IDS)
@pytest.mark.parametrize("mode", list(_PADDING_MODES))
def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig | None):
"""See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").
Exercises AiterExperts's hidden_pad/intermediate_pad (rocm_aiter_moe.py)
across {no padding, hidden-only, intermediate-only, both} x the
unquantized + 4 fp8 quant schemes it supports. Compares against an
unpadded reference and verifies the hidden_pad/intermediate_pad
AiterExperts.apply() forwards to rocm_aiter_ops.fused_moe.
"""
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoDPEPModular,
)
pad_hidden, pad_intermediate = _PADDING_MODES[mode]
if (
quant_config is not None
and quant_config.block_shape is not None
and not pad_hidden
):
pytest.skip(
"AITER's block-quantized (128x128) CK GEMM kernel does not "
f"support the raw (unpadded) hidden_dim={_PADDING_K_UNPADDED} "
"shape on this hardware -- it raises 'wrong! device_gemm with "
"the specified compilation parameters does not support this "
"GEMM problem'. Only hidden_pad-forcing modes (raw "
f"hidden_dim={_PADDING_K_PADDED}) are exercised for this quant "
"scheme."
)
k = _PADDING_K_PADDED if pad_hidden else _PADDING_K_UNPADDED
n = _PADDING_N_PADDED if pad_intermediate else _PADDING_N_UNPADDED
config = Config(
Ms=_PADDING_M,
K=k,
N=n,
E=_PADDING_E,
topks=_PADDING_TOPK,
dtype=torch.bfloat16,
quant_config=quant_config,
prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
fused_experts_type=AiterExperts,
world_size=1,
hidden_dim_unpadded=_PADDING_K_UNPADDED if pad_hidden else None,
intermediate_size_per_partition_unpadded=(
_PADDING_N_UNPADDED if pad_intermediate else None
),
)
assert config.is_valid()[0]
assert config.fe_supports_quant_scheme(), (
f"AiterExperts does not support quant scheme {quant_config}."
)
weights = _make_padding_matrix_weights(config)
vllm_config, env_dict = config.make_env_data()
hidden_pad_expected = _PADDING_HIDDEN_PAD_EXPECTED if pad_hidden else 0
intermediate_pad_expected = (
_PADDING_INTERMEDIATE_PAD_EXPECTED if pad_intermediate else 0
)
parallel_launch_with_config(
config.world_size,
_aiter_padding_matrix_worker,
vllm_config,
env_dict,
config,
weights,
False,
hidden_pad_expected,
intermediate_pad_expected,
)
@pytest.mark.skip(
reason="MXFP4 AiterExperts padding requires MI350/gfx950 hardware. "
'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").'
)
def test_aiter_moe_padding_matrix_mxfp4():
"""Stub for the {no padding, hidden-only, intermediate-only, both} MXFP4
padding matrix once MI350/gfx950 hardware is available in CI."""
if __name__ == "__main__": if __name__ == "__main__":
# Ability to test individual PrepareAndFinalize and FusedExperts combination # Ability to test individual PrepareAndFinalize and FusedExperts combination
from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser
+9 -7
View File
@@ -16,10 +16,8 @@ from vllm.utils.torch_utils import set_random_seed
NUM_TOKENS = [1, 17, 256] NUM_TOKENS = [1, 17, 256]
TOP_KS = [2, 8] TOP_KS = [2, 8]
HIDDEN_SIZE = 512 HIDDEN_SIZE = 512
# One BLOCK_K tile, a partial trailing tile, and Kimi-K3's real dims # Include partial tiles and hidden sizes above the CUDA 4096-element tile cap.
# (moe_intermediate_size 3072, hidden_size 7168) -- exercises the per-tile loop HIDDEN_SIZES = [512, 1000, 3072, 4097, 7168]
# and its tail mask (BLOCK_K caps at 512 for <=2-byte dtypes).
HIDDEN_SIZES = [512, 1000, 3072, 7168]
DTYPES = [torch.bfloat16, torch.float32] DTYPES = [torch.bfloat16, torch.float32]
NUM_EXPERTS = 16 NUM_EXPERTS = 16
@@ -233,7 +231,7 @@ def test_nonlocal_and_padding_adjacent():
@pytest.mark.parametrize("num_tokens", NUM_TOKENS) @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("top_k", TOP_KS) @pytest.mark.parametrize("top_k", TOP_KS)
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
@pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("dtype", [*DTYPES, torch.float16])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
def test_weighted_sum_matches_reference( def test_weighted_sum_matches_reference(
num_tokens: int, top_k: int, hidden_size: int, dtype: torch.dtype num_tokens: int, top_k: int, hidden_size: int, dtype: torch.dtype
@@ -268,7 +266,11 @@ def test_weighted_sum_matches_reference(
assert not out.isnan().any(), "invalid slots leaked into the reduction" assert not out.isnan().any(), "invalid slots leaked into the reduction"
ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map=None) ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map=None)
atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4) atol, rtol = {
torch.bfloat16: (1e-2, 1e-2),
torch.float16: (1e-3, 1e-3),
torch.float32: (1e-4, 1e-4),
}[dtype]
torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol) torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol)
@@ -360,7 +362,7 @@ def test_num_valid_tokens_bounds_stale_padding(dtype: torch.dtype):
""" """
set_random_seed(0) set_random_seed(0)
device = "cuda" device = "cuda"
num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 512 num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 7168
inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device) inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device)
topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device) topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device)
+2
View File
@@ -781,6 +781,7 @@ def test_trtllm_mxfp4_deferred_finalize(
global_num_experts=num_experts, global_num_experts=num_experts,
expert_map=None, expert_map=None,
apply_router_weight_on_input=False, apply_router_weight_on_input=False,
routing_sink=None,
), ),
) )
return return
@@ -1780,6 +1781,7 @@ def test_rocm_mxfp4_moe_oracle(
global_num_experts=num_experts, global_num_experts=num_experts,
expert_map=None, expert_map=None,
apply_router_weight_on_input=False, apply_router_weight_on_input=False,
routing_sink=None,
) )
else: else:
# Modular impl uses topk_weights and topk_ids # Modular impl uses topk_weights and topk_ids
File diff suppressed because it is too large Load Diff
+30
View File
@@ -91,6 +91,12 @@ QWEN4_EXP_SM100_CASES = [
for num_tokens, config in plans.items() for num_tokens, config in plans.items()
] ]
QWEN4_EXP_SM121_CASES = [
(n, k, num_tokens, config)
for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS.items()
for num_tokens, config in plans.items()
]
EXPECTED_CUTE_CONFIGS = { EXPECTED_CUTE_CONFIGS = {
(3072, 7168, 1): (224, 3, 4, 8), (3072, 7168, 1): (224, 3, 4, 8),
(3072, 7168, 2): (128, 3, 2, 8), (3072, 7168, 2): (128, 3, 2, 8),
@@ -725,6 +731,7 @@ def test_qwen4_exp_hopper_plans_are_valid() -> None:
((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS), ((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS), ((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS),
((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS), ((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
((12, 1), qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS),
((8, 0), {}), ((8, 0), {}),
], ],
) )
@@ -939,6 +946,29 @@ def test_qwen4_exp_sm100_selected_shapes(
assert cosine > 0.999 assert cosine > 0.999
@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM121_CASES)
def test_qwen4_exp_sm121_selected_shapes(
n: int,
k: int,
num_tokens: int,
config: SkinnyGemmConfig,
) -> None:
_require_capability_and_cute((12, 1))
torch.manual_seed(42 + num_tokens)
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
selected = qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS[(n, k)][num_tokens]
assert selected == config
output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight)
reference = torch.nn.functional.linear(x, weight)
cosine = torch.nn.functional.cosine_similarity(
output.float().flatten(), reference.float().flatten(), dim=0
).item()
assert cosine > 0.999
def test_glm52_q_b_nonpacked_single_row_falls_back() -> None: def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
_require_sm103_and_cute() _require_sm103_and_cute()
spec = glm52_gemm.GLM52_Q_B_PROJECTION spec = glm52_gemm.GLM52_Q_B_PROJECTION
+225 -118
View File
@@ -636,8 +636,8 @@ def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None:
@pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
@torch.inference_mode() @torch.inference_mode()
def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
from vllm._aiter_ops import rocm_aiter_ops
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_get_aiter_top_k_kernel,
_launch_aiter_top_k_per_row_prefill, _launch_aiter_top_k_per_row_prefill,
) )
@@ -647,17 +647,14 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
top_k = 4 top_k = 4
indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda") indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda")
aiter_topk_kernel = _get_aiter_top_k_kernel( if not rocm_aiter_ops.is_indexer_top_k_supported(
is_prefill=True, is_prefill=True,
compress_ratio=4, compress_ratio=4,
num_rows=logits.shape[0], num_rows=logits.shape[0],
on_gfx950=True, ):
)
if aiter_topk_kernel is None:
pytest.skip("AITER top-k is unavailable") pytest.skip("AITER top-k is unavailable")
assert ( assert (
_launch_aiter_top_k_per_row_prefill( _launch_aiter_top_k_per_row_prefill(
aiter_topk_kernel,
logits, logits,
row_starts, row_starts,
row_ends, row_ends,
@@ -683,10 +680,7 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
num_speculative_tokens: int, num_speculative_tokens: int,
) -> None: ) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( from vllm._aiter_ops import rocm_aiter_ops
_get_aiter_top_k_kernel,
_launch_aiter_top_k_per_row_decode,
)
query_tokens = num_speculative_tokens + 1 query_tokens = num_speculative_tokens + 1
offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda") offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda")
@@ -701,20 +695,18 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
top_k = 4 top_k = 4
indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
aiter_topk_kernel = _get_aiter_top_k_kernel( if not rocm_aiter_ops.is_indexer_top_k_supported(
is_prefill=False, is_prefill=False,
compress_ratio=4, compress_ratio=4,
num_rows=num_rows, num_rows=num_rows,
max_valid_seq_len=seq_lens.max().item(), max_valid_seq_len=seq_lens.max().item(),
on_gfx950=True, ):
)
if aiter_topk_kernel is None:
pytest.skip("AITER top-k is unavailable") pytest.skip("AITER top-k is unavailable")
assert ( assert (
_launch_aiter_top_k_per_row_decode( rocm_aiter_ops.indexer_top_k_decode(
aiter_topk_kernel,
logits, logits,
seq_lens, 1,
seq_lens.reshape(-1),
indices, indices,
top_k, top_k,
) )
@@ -729,130 +721,107 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
assert torch.all(indices[row_idx, num_valid:] == -1) assert torch.all(indices[row_idx, num_valid:] == -1)
def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: @requires_gfx950
from vllm.v1.attention.ops import rocm_aiter_mla_sparse @torch.inference_mode()
def test_aiter_decode_topk_matches_per_row_under_mtp() -> None:
"""The aiter backend expands (B, 1) seq_lens into per-row ends before
launching, so next_n > 1 must agree with the in-tree kernel, including
the rows that clamp to a zero-length prefix.
"""
from vllm.model_executor.layers.indexer_topk import get_indexer_topk
def prefill_kernel(*args, **kwargs) -> None: next_n = 4
pass top_k = 512
num_columns = 1024
seq_lens = torch.tensor([[2], [513], [700]], dtype=torch.int32, device="cuda")
num_rows = seq_lens.shape[0] * next_n
def decode_kernel(*args, **kwargs) -> None: torch.manual_seed(0)
pass logits = torch.randn(num_rows, num_columns, dtype=torch.float32, device="cuda")
outputs = {}
for backend in ("aiter", "per_row"):
indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
get_indexer_topk(backend)(logits, seq_lens, next_n, indices, top_k, num_columns)
outputs[backend] = indices
row_ends = (
(seq_lens - next_n + 1 + torch.arange(next_n, device="cuda"))
.clamp_(min=0)
.reshape(-1)
.tolist()
)
assert row_ends[:next_n] == [0, 0, 1, 2]
for row_idx, row_end in enumerate(row_ends):
num_valid = min(top_k, row_end)
aiter_row = outputs["aiter"][row_idx]
assert set(aiter_row[:num_valid].tolist()) == set(
outputs["per_row"][row_idx][:num_valid].tolist()
)
assert torch.all(aiter_row[num_valid:] == -1)
def _enable_aiter_topk(monkeypatch, enabled: bool = True) -> None:
from vllm._aiter_ops import rocm_aiter_ops
monkeypatch.setattr( monkeypatch.setattr(
rocm_aiter_mla_sparse, rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: enabled, raising=True
"_get_aiter_topk_ops",
lambda: (prefill_kernel, decode_kernel),
) )
get_kernel = rocm_aiter_mla_sparse._get_aiter_top_k_kernel
def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None:
from vllm._aiter_ops import rocm_aiter_ops
_enable_aiter_topk(monkeypatch)
is_supported = rocm_aiter_ops.is_indexer_top_k_supported
eligible_cases = [ eligible_cases = [
( dict(is_prefill=True, compress_ratio=4, num_rows=1),
dict(is_prefill=True, compress_ratio=4, num_rows=1, on_gfx950=True), dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_536),
prefill_kernel, dict(
is_prefill=False,
compress_ratio=4,
num_rows=257,
max_valid_seq_len=250_000,
), ),
(
dict(
is_prefill=False,
compress_ratio=4,
num_rows=1,
max_valid_seq_len=65_536,
on_gfx950=True,
),
decode_kernel,
),
(
dict(
is_prefill=False,
compress_ratio=4,
num_rows=257,
max_valid_seq_len=250_000,
num_columns=524_288,
on_gfx950=True,
),
decode_kernel,
),
(
dict(
is_prefill=False,
compress_ratio=2,
num_rows=385,
max_valid_seq_len=250_000,
num_columns=524_288,
topk_tokens=512,
on_gfx950=True,
),
decode_kernel,
),
(
dict(
is_prefill=False,
compress_ratio=2,
num_rows=384,
max_valid_seq_len=50_000,
num_columns=1_048_577,
topk_tokens=512,
on_gfx950=True,
),
decode_kernel,
),
]
for kwargs, expected_kernel in eligible_cases:
assert get_kernel(**kwargs) is expected_kernel
def fail_if_imported() -> None:
pytest.fail("ineligible shapes must not import AITER top-k")
monkeypatch.setattr(
rocm_aiter_mla_sparse,
"_get_aiter_topk_ops",
fail_if_imported,
)
ineligible_cases = [
dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True),
dict( dict(
is_prefill=False, is_prefill=False,
compress_ratio=2, compress_ratio=2,
num_rows=384, num_rows=384,
max_valid_seq_len=1_048_576, max_valid_seq_len=1_048_576,
num_columns=1_048_576,
topk_tokens=512,
on_gfx950=True,
),
dict(
is_prefill=False,
compress_ratio=4,
num_rows=1,
max_valid_seq_len=65_537,
on_gfx950=True,
), ),
]
for kwargs in eligible_cases:
assert is_supported(**kwargs) is True
ineligible_cases = [
dict(is_prefill=True, compress_ratio=1, num_rows=1),
dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_537),
dict( dict(
is_prefill=False, is_prefill=False,
compress_ratio=4, compress_ratio=4,
num_rows=256, num_rows=256,
max_valid_seq_len=125_000, max_valid_seq_len=125_000,
on_gfx950=True,
),
dict(
is_prefill=False,
compress_ratio=4,
num_rows=320,
max_valid_seq_len=2_500,
on_gfx950=False,
), ),
] ]
for kwargs in ineligible_cases: for kwargs in ineligible_cases:
assert get_kernel(**kwargs) is None assert is_supported(**kwargs) is False
monkeypatch.setattr(rocm_aiter_mla_sparse, "_get_aiter_topk_ops", lambda: None) _enable_aiter_topk(monkeypatch, enabled=False)
assert ( assert is_supported(is_prefill=True, compress_ratio=4, num_rows=1) is False
get_kernel(
is_prefill=True,
compress_ratio=4, def test_dsv4_native_top_k_window() -> None:
num_rows=1, """The in-tree decode kernel's measured advantage band over AITER, which
on_gfx950=True, applies to the DSV4 indexer only."""
) from vllm._aiter_ops import rocm_aiter_ops
is None
) prefers_native = rocm_aiter_ops.dsv4_indexer_prefers_native_top_k
assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=512) is True
assert prefers_native(num_rows=385, num_columns=524_288, topk_tokens=512) is False
assert prefers_native(num_rows=128, num_columns=1_048_577, topk_tokens=512) is False
assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=1024) is False
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
@@ -1927,3 +1896,141 @@ def test_sparse_indexer_topk_backend_resolution() -> None:
if not is_sm100: if not is_sm100:
with pytest.raises(RuntimeError, match="SM100a/SM103a"): with pytest.raises(RuntimeError, match="SM100a/SM103a"):
resolve("deep_select") resolve("deep_select")
def _patch_aiter_topk(monkeypatch, decode_kernel, enabled: bool = True) -> None:
from vllm._aiter_ops import rocm_aiter_ops
_enable_aiter_topk(monkeypatch, enabled=enabled)
monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: enabled, raising=True)
monkeypatch.setattr(rocm_aiter_ops, "indexer_top_k_decode", decode_kernel)
def _make_topk(backend: str):
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
with set_current_vllm_config(cfg):
return SparseIndexerTopk()
@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
def test_sparse_indexer_topk_aiter_requires_rocm() -> None:
"""Explicitly requesting the ROCm-only backend must fail fast elsewhere."""
logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda")
with pytest.raises(RuntimeError, match="ROCm"):
_make_topk("aiter").resolve_backend(logits, 2048, 8)
@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
def test_sparse_indexer_topk_backend_resolution_aiter(monkeypatch) -> None:
"""Aiter is opt-in only: "auto" never selects it (the shape-dependent perf
gate lives at the call site), and an explicit request validates hard
capability only."""
def decode_kernel(*args, **kwargs) -> None:
pass
_patch_aiter_topk(monkeypatch, decode_kernel)
logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda")
assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row"
assert _make_topk("aiter").resolve_backend(logits, 2048, 8) == "aiter"
# Shapes the call-site perf gate would decline still resolve.
assert _make_topk("aiter").resolve_backend(logits, 512, 8) == "aiter"
with pytest.raises(RuntimeError, match="topk_tokens must be in"):
_make_topk("aiter").resolve_backend(logits, 3000, 8)
_enable_aiter_topk(monkeypatch, enabled=False)
with pytest.raises(RuntimeError, match="no tuned indexer top-k kernels"):
_make_topk("aiter").resolve_backend(logits, 2048, 8)
_patch_aiter_topk(monkeypatch, decode_kernel, enabled=False)
assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row"
with pytest.raises(RuntimeError, match="VLLM_ROCM_USE_AITER"):
_make_topk("aiter").resolve_backend(logits, 2048, 8)
@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
@torch.inference_mode()
def test_sparse_indexer_topk_aiter_applies_no_shape_gate(monkeypatch) -> None:
"""Shape-dependent perf gating belongs to the call site: once "aiter" is
requested the dispatcher must run it for every supported k, including the
k=512 band the in-tree kernel wins on DSV4's compressed-KV logits."""
from vllm.model_executor.layers import indexer_topk
calls: list[tuple] = []
_patch_aiter_topk(monkeypatch, lambda *args: calls.append(args))
monkeypatch.setattr(
indexer_topk.ops,
"top_k_per_row_decode",
lambda *args: pytest.fail("an explicit aiter request must not fall back"),
)
num_rows = 8
logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda")
seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda")
indices = torch.empty((num_rows, 512), dtype=torch.int32, device="cuda")
op = _make_topk("aiter")
op(logits, seq_lens, 1, indices, 512, 4096)
# Past the compressed-context boundary the call-site gate declines on.
op(logits, seq_lens, 1, indices, 512, 4 * (64 * 1024 + 1))
assert len(calls) == 2
@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
@torch.inference_mode()
def test_sparse_indexer_topk_auto_stays_on_per_row(monkeypatch) -> None:
"""The "auto" chain must not reach AITER on its own: the decode top-k runs
through top_k_per_row_decode unless the caller asks for the aiter backend."""
from vllm.model_executor.layers import indexer_topk
def decode_kernel(*args, **kwargs) -> None:
pytest.fail("AITER is opt-in only")
_patch_aiter_topk(monkeypatch, decode_kernel)
calls: list[tuple] = []
monkeypatch.setattr(
indexer_topk.ops, "top_k_per_row_decode", lambda *args: calls.append(args)
)
num_rows = 8
logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda")
seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda")
indices = torch.empty((num_rows, 1024), dtype=torch.int32, device="cuda")
_make_topk("auto")(logits, seq_lens, 1, indices, 1024, 4096)
assert len(calls) == 1
@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
@torch.inference_mode()
def test_sparse_indexer_topk_aiter_expands_mtp_row_ends(monkeypatch) -> None:
"""AITER's decode kernel only implements the next_n == 1 form, so spec
decode batches must be handed the equivalent expanded per-row ends."""
from vllm.model_executor.layers import indexer_topk
recorded: dict = {}
def decode_kernel(logits, next_n, seq_lens, *args, **kwargs) -> None:
recorded["next_n"] = next_n
recorded["row_ends"] = seq_lens.tolist()
_patch_aiter_topk(monkeypatch, decode_kernel)
monkeypatch.setattr(
indexer_topk.ops,
"top_k_per_row_decode",
lambda *args: pytest.fail("an explicit aiter request must not fall back"),
)
next_n = 2
logits = torch.randn(4, 4096, dtype=torch.float32, device="cuda")
seq_lens = torch.tensor([[100], [200]], dtype=torch.int32, device="cuda")
indices = torch.empty((4, 1024), dtype=torch.int32, device="cuda")
_make_topk("aiter")(logits, seq_lens, next_n, indices, 1024, 800)
assert recorded["next_n"] == 1
assert recorded["row_ends"] == [99, 100, 199, 200]

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