Merge remote-tracking branch 'upstream/main' into patches
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
@@ -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
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -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 \
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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
@@ -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)
|
||||||
|
|||||||
@@ -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).
|
||||||
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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).
|
||||||
@@ -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`
|
||||||
|
|||||||
@@ -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
@@ -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",
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -1,4 +0,0 @@
|
|||||||
# CLAUDE.md
|
|
||||||
|
|
||||||
First, check @AGENTS.override.md if exists.
|
|
||||||
Then, follow instructions in @AGENTS.md.
|
|
||||||
@@ -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
@@ -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)),
|
||||||
]);
|
]);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
@@ -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"],
|
||||||
|
)
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
@@ -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."""
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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)
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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
@@ -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
@@ -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
Reference in New Issue
Block a user