Merge remote-tracking branch 'origin/main' into sync/lesj/qwen4exp-bugfixes
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
This commit is contained in:
lesj0610
@@ -79,6 +79,54 @@ setup_buildx_builder() {
|
||||
docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
|
||||
}
|
||||
|
||||
download_ci_hcl() {
|
||||
# Returns non-zero instead of exiting: the optional export path must be
|
||||
# able to give up without failing a build whose image already exists.
|
||||
echo "--- :arrow_down: Downloading ci.hcl"
|
||||
curl -sSfL -o "${CI_HCL_PATH}" "${CI_HCL_URL}" || return $?
|
||||
echo "Downloaded to ${CI_HCL_PATH}"
|
||||
|
||||
if [[ ! -f "${CI_HCL_PATH}" ]]; then
|
||||
echo "Error: ci.hcl not found at ${CI_HCL_PATH}"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
export_kernel_symbol_map_from_cache() {
|
||||
# The existing-image path: everything here is best effort. Called as
|
||||
# `... || echo`, so errexit is suspended inside and any failure just
|
||||
# returns; the build keeps the image it already has.
|
||||
local rc=0
|
||||
download_ci_hcl \
|
||||
&& setup_buildx_builder \
|
||||
&& resolve_parent_commit \
|
||||
&& export PARENT_COMMIT \
|
||||
&& BUILD_TMP_DIR="$(mktemp -d)" \
|
||||
&& export_kernel_symbol_map || rc=$?
|
||||
rm -rf -- "${BUILD_TMP_DIR:-}"
|
||||
return "${rc}"
|
||||
}
|
||||
|
||||
export_kernel_symbol_map() {
|
||||
# Only when the build asked for it (nightly/post-merge). The map was
|
||||
# produced inside the csrc-build stage during the main bake; this second
|
||||
# bake of the scratch stage is a cache hit that just writes the file out.
|
||||
if [[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" != "1" ]]; then
|
||||
return 0
|
||||
fi
|
||||
echo "--- :world_map: Exporting kernel symbol map"
|
||||
local out_dir="${BUILD_TMP_DIR}/kernel-symbol-map"
|
||||
mkdir -p "${out_dir}"
|
||||
if docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --progress plain \
|
||||
--set "kernel-symbol-map.output=type=local,dest=${out_dir}" kernel-symbol-map \
|
||||
&& [[ -s "${out_dir}/kernel_symbol_map.json.gz" ]]; then
|
||||
ls -la "${out_dir}"
|
||||
(cd "${out_dir}" && buildkite-agent artifact upload "kernel_symbol_map.json.gz")
|
||||
else
|
||||
echo "kernel symbol map export failed; continuing without it" >&2
|
||||
fi
|
||||
}
|
||||
|
||||
annotate_image_tags() {
|
||||
.buildkite/scripts/annotate-image-build.sh \
|
||||
"${IMAGE_TAG:-}" "${IMAGE_TAG_LATEST:-}"
|
||||
@@ -90,6 +138,15 @@ check_and_skip_if_image_exists() {
|
||||
if docker manifest inspect "${IMAGE_TAG}" >/dev/null 2>&1; then
|
||||
echo "Image already exists: ${IMAGE_TAG}"
|
||||
echo "Skipping build"
|
||||
if [[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" == "1" ]]; then
|
||||
# The image is reused, but the symbol map is an artifact of
|
||||
# this build. Bake just the export target: its layers come
|
||||
# from the registry cache, and the map step itself reruns
|
||||
# if the cached image was built without the arg. Never
|
||||
# fails the build: the image is already there.
|
||||
export_kernel_symbol_map_from_cache \
|
||||
|| echo "kernel symbol map: export from the cached image failed; continuing" >&2
|
||||
fi
|
||||
annotate_image_tags
|
||||
exit 0
|
||||
fi
|
||||
@@ -154,7 +211,7 @@ print_bake_config() {
|
||||
local bake_tmp
|
||||
bake_tmp="$(mktemp -d)"
|
||||
BAKE_CONFIG_FILE="${bake_tmp}/bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
|
||||
docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
|
||||
docker buildx bake "${BAKE_FILES[@]}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
|
||||
echo "Saved bake config to ${BAKE_CONFIG_FILE}"
|
||||
echo "--- :arrow_down: Uploading bake config to Buildkite"
|
||||
(cd "$(dirname "${BAKE_CONFIG_FILE}")" && buildkite-agent artifact upload "$(basename "${BAKE_CONFIG_FILE}")")
|
||||
@@ -230,8 +287,10 @@ TARGET="test-ci"
|
||||
VLLM_BAKE_FILE_PATH="${VLLM_BAKE_FILE_PATH:-docker/docker-bake.hcl}"
|
||||
BUILDER_NAME="${BUILDER_NAME:-vllm-builder}"
|
||||
CI_HCL_URL="${CI_HCL_URL:-https://raw.githubusercontent.com/vllm-project/ci-infra/main/docker/ci.hcl}"
|
||||
CI_HCL_PATH="/tmp/ci.hcl"
|
||||
CI_HCL_PATH="${CI_HCL_PATH:-/tmp/ci.hcl}"
|
||||
ZSTD_HCL_PATH="${ZSTD_HCL_PATH:-.buildkite/image_build/zstd.hcl}"
|
||||
BUILDKIT_SOCKET="/run/buildkit/buildkitd.sock"
|
||||
BAKE_FILES=(-f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" -f "${ZSTD_HCL_PATH}")
|
||||
|
||||
prepare_cache_tags
|
||||
ecr_login
|
||||
@@ -288,14 +347,7 @@ if [[ ! -f "${VLLM_BAKE_FILE_PATH}" ]]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "--- :arrow_down: Downloading ci.hcl"
|
||||
curl -sSfL -o "${CI_HCL_PATH}" "${CI_HCL_URL}"
|
||||
echo "Downloaded to ${CI_HCL_PATH}"
|
||||
|
||||
if [[ ! -f "${CI_HCL_PATH}" ]]; then
|
||||
echo "Error: ci.hcl not found at ${CI_HCL_PATH}"
|
||||
exit 1
|
||||
fi
|
||||
download_ci_hcl || exit 1
|
||||
|
||||
setup_buildx_builder
|
||||
|
||||
@@ -309,7 +361,7 @@ BUILD_TMP_DIR="$(mktemp -d)"
|
||||
trap 'rm -rf -- "${BUILD_TMP_DIR}"' EXIT
|
||||
BUILD_METADATA_FILE="${BUILD_TMP_DIR}/build-metadata.json"
|
||||
BUILD_STATUS=0
|
||||
docker --debug buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" \
|
||||
docker --debug buildx bake "${BAKE_FILES[@]}" \
|
||||
--progress plain --metadata-file "${BUILD_METADATA_FILE}" "${TARGET}" || BUILD_STATUS=$?
|
||||
|
||||
record_buildkit_trace "${BUILD_METADATA_FILE}" || true
|
||||
@@ -319,5 +371,6 @@ if [[ "${BUILD_STATUS}" -ne 0 ]]; then
|
||||
fi
|
||||
|
||||
echo "--- :white_check_mark: Build complete"
|
||||
export_kernel_symbol_map || true
|
||||
|
||||
annotate_image_tags
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
target "test-ci" {
|
||||
output = ["type=registry,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true"]
|
||||
}
|
||||
@@ -26,3 +26,27 @@ steps:
|
||||
'uv pip install "gpt-oss[eval]==0.0.5" &&
|
||||
cd tests &&
|
||||
pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
|
||||
- label: LM Eval Large Models(2 GPUs)
|
||||
key: lm-eval-large-models-2-gpus
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
optional: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- .buildkite/lm-eval-harness/
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd .buildkite/lm-eval-harness &&
|
||||
export VLLM_USE_DEEP_GEMM=0 &&
|
||||
pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-xpu-fp8.txt --tp-size=2'
|
||||
|
||||
@@ -2,6 +2,34 @@ group: Miscellaneous Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: XPU Batch Invariance
|
||||
key: xpu-batch-invariance
|
||||
timeout_in_minutes: 60
|
||||
num_devices: 2
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1/determinism/test_xpu_batch_invariant_ut.py
|
||||
- .buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh
|
||||
- .buildkite/intel_jobs/misc_intel.yaml
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_BATCH_INVARIANT=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py &&
|
||||
bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh'
|
||||
|
||||
- label: V1 Core + KV + Metrics
|
||||
key: v1-core-kv-metrics-intel
|
||||
timeout_in_minutes: 45
|
||||
@@ -290,4 +318,4 @@ steps:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
|
||||
pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
group: Model Runner V2 Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: Model Runner V2 Core Tests (Intel)
|
||||
key: model-runner-v2-core-tests-intel
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/attention/
|
||||
- tests/v1/engine/test_llm_engine.py
|
||||
- tests/v1/e2e/
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
|
||||
pytest -v -s v1/e2e/general/test_context_length.py &&
|
||||
ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
|
||||
pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" &&
|
||||
pytest -v -s v1/e2e/general/test_min_tokens.py'
|
||||
|
||||
- label: Model Runner V2 Examples (Intel)
|
||||
key: model-runner-v2-examples-intel
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- examples/basic/offline_inference/
|
||||
- examples/generate/multimodal/
|
||||
- examples/features/
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd examples &&
|
||||
python3 basic/offline_inference/chat.py &&
|
||||
python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
|
||||
python3 generate/multimodal/vision_language_offline.py --seed 0 &&
|
||||
python3 features/automatic_prefix_caching/prefix_caching_offline.py'
|
||||
|
||||
- label: Model Runner V2 Distributed (2 GPUs)
|
||||
key: model-runner-v2-distributed-2-gpus
|
||||
timeout_in_minutes: 50
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/basic_correctness/test_basic_correctness.py
|
||||
- tests/v1/distributed/test_async_llm_dp.py
|
||||
- tests/v1/distributed/test_eagle_dp.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd tests &&
|
||||
TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"'
|
||||
|
||||
- label: Model Runner V2 Spec Decode
|
||||
key: model-runner-v2-spec-decode
|
||||
timeout_in_minutes: 50
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/v1/spec_decode/test_max_len.py
|
||||
- tests/v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py'
|
||||
@@ -106,7 +106,7 @@ steps:
|
||||
pytest -v -s v1/test_serial_utils.py &&
|
||||
pytest -v -s v1/e2e/general/test_correctness_sliding_window.py &&
|
||||
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
|
||||
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
|
||||
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --deselect "v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]"'
|
||||
parallelism: 2
|
||||
- label: "XPU server test"
|
||||
key: xpu-server-test
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
|
||||
tasks:
|
||||
- name: "gsm8k"
|
||||
metrics:
|
||||
- name: "exact_match,strict-match"
|
||||
value: 0.7142
|
||||
- name: "exact_match,flexible-extract"
|
||||
value: 0.4579
|
||||
limit: 256
|
||||
num_fewshot: 5
|
||||
max_model_len: 262144
|
||||
kv_cache_dtype: fp8
|
||||
enforce_eager: true
|
||||
apply_chat_template: true
|
||||
fewshot_as_multiturn: true
|
||||
trust_remote_code: true
|
||||
@@ -0,0 +1 @@
|
||||
NVIDIA-Nemotron-3-Nano-30B-A3B-FP8-xpu.yaml
|
||||
+243
@@ -0,0 +1,243 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../.." && pwd)"
|
||||
|
||||
model="ibm-research/PowerMoE-3b"
|
||||
port=8000
|
||||
num_decode_tokens=128
|
||||
server_pid=
|
||||
server_log=
|
||||
|
||||
# XPU MoE kernels may not safely handle the -1 routes emitted for padding.
|
||||
export VLLM_MOE_SKIP_PADDING=0
|
||||
|
||||
cd "$repo_root"
|
||||
|
||||
cleanup_server() {
|
||||
local attempt
|
||||
|
||||
if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then
|
||||
kill "$server_pid" 2>/dev/null || true
|
||||
|
||||
for ((attempt = 0; attempt < 30; attempt++)); do
|
||||
if ! kill -0 "$server_pid" 2>/dev/null; then
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
if kill -0 "$server_pid" 2>/dev/null; then
|
||||
echo "vLLM did not stop after 30 seconds; sending SIGKILL" >&2
|
||||
kill -KILL "$server_pid" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
wait "$server_pid" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
server_pid=
|
||||
}
|
||||
|
||||
on_exit() {
|
||||
local status=$?
|
||||
|
||||
cleanup_server
|
||||
if [[ "$status" -ne 0 && -n "$server_log" && -f "$server_log" ]]; then
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
buildkite-agent artifact upload "$server_log" ||
|
||||
echo "Failed to upload $server_log as a Buildkite artifact" >&2
|
||||
fi
|
||||
fi
|
||||
[[ -z "$server_log" ]] || rm -f "$server_log"
|
||||
exit "$status"
|
||||
}
|
||||
trap on_exit EXIT
|
||||
|
||||
wait_for_server() {
|
||||
for _ in $(seq 1 600); do
|
||||
if curl -sf "http://127.0.0.1:$port/health" >/dev/null; then
|
||||
return
|
||||
fi
|
||||
if ! kill -0 "$server_pid" 2>/dev/null; then
|
||||
echo "vLLM exited before becoming healthy; see the server log above." >&2
|
||||
return 1
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "Timed out waiting for vLLM server health check" >&2
|
||||
return 1
|
||||
}
|
||||
|
||||
request_outputs() {
|
||||
local prompts_json=$1
|
||||
local expected_choices=$2
|
||||
local response
|
||||
|
||||
response=$(printf \
|
||||
'{"model":"%s","prompt":%s,"temperature":0.7,"top_p":0.95,"seed":0,"max_tokens":%s,"logprobs":64,"ignore_eos":true,"return_token_ids":true}' \
|
||||
"$model" "$prompts_json" "$num_decode_tokens" |
|
||||
curl --fail --silent --show-error \
|
||||
--max-time 1200 \
|
||||
"http://127.0.0.1:$port/v1/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
--data-binary @-)
|
||||
|
||||
local -a token_id_arrays logprob_arrays
|
||||
mapfile -t token_id_arrays < <(
|
||||
grep -o '"token_ids":\[[^]]*\]' <<<"$response" |
|
||||
sed 's/^"token_ids":\[//; s/\]$//'
|
||||
)
|
||||
mapfile -t logprob_arrays < <(
|
||||
grep -o '"token_logprobs":\[[^]]*\]' <<<"$response" |
|
||||
sed 's/^"token_logprobs":\[//; s/\]$//'
|
||||
)
|
||||
if [[ "${#token_id_arrays[@]}" -ne "$expected_choices" ]] ||
|
||||
[[ "${#logprob_arrays[@]}" -ne "$expected_choices" ]]; then
|
||||
echo "Expected $expected_choices choices with token IDs and logits" >&2
|
||||
return 1
|
||||
fi
|
||||
local index
|
||||
local array
|
||||
local value
|
||||
local -a values
|
||||
for index in "${!logprob_arrays[@]}"; do
|
||||
array="${logprob_arrays[index]}"
|
||||
IFS=, read -r -a values <<<"$array"
|
||||
if [[ "${#values[@]}" -ne "$num_decode_tokens" ]]; then
|
||||
echo "Completion returned an unexpected number of decode steps" >&2
|
||||
return 1
|
||||
fi
|
||||
IFS=, read -r -a values <<<"${token_id_arrays[index]}"
|
||||
if [[ "${#values[@]}" -ne "$num_decode_tokens" ]]; then
|
||||
echo "Completion returned an unexpected number of generated tokens" >&2
|
||||
return 1
|
||||
fi
|
||||
for value in "${values[@]}"; do
|
||||
if [[ ! "$value" =~ ^[0-9]+$ ]]; then
|
||||
echo "Completion returned an invalid token ID" >&2
|
||||
return 1
|
||||
fi
|
||||
done
|
||||
IFS=, read -r -a values <<<"$array"
|
||||
for value in "${values[@]}"; do
|
||||
if [[ ! "$value" =~ ^-?([0-9]+([.][0-9]*)?|[.][0-9]+)([eE][+-]?[0-9]+)?$ ]]; then
|
||||
echo "Completion returned non-finite raw logits" >&2
|
||||
return 1
|
||||
fi
|
||||
done
|
||||
printf '%s|%s\n' "${token_id_arrays[index]}" "${logprob_arrays[index]}"
|
||||
done
|
||||
}
|
||||
|
||||
make_prompt() {
|
||||
local length=$1
|
||||
local offset=${2:-0}
|
||||
local i
|
||||
local prompt="Use the following system-design context to answer the question. Scenario $offset:"
|
||||
for ((i = 0; i < length; i++)); do
|
||||
prompt+=" context"
|
||||
done
|
||||
prompt+="\\n\\nQuestion: How should batch-invariant inference behave when this same request is decoded alongside unrelated requests?\\nAnswer:"
|
||||
printf '"%s"' "$prompt"
|
||||
}
|
||||
|
||||
first_difference() {
|
||||
local actual=$1
|
||||
local expected=$2
|
||||
local position=$3
|
||||
local actual_tokens=${actual%%|*}
|
||||
local expected_tokens=${expected%%|*}
|
||||
local actual_logits=${actual#*|}
|
||||
local expected_logits=${expected#*|}
|
||||
local -a actual_token_values
|
||||
local -a expected_token_values
|
||||
local -a actual_logit_values
|
||||
local -a expected_logit_values
|
||||
local step
|
||||
|
||||
IFS=, read -r -a actual_token_values <<<"$actual_tokens"
|
||||
IFS=, read -r -a expected_token_values <<<"$expected_tokens"
|
||||
IFS=, read -r -a actual_logit_values <<<"$actual_logits"
|
||||
IFS=, read -r -a expected_logit_values <<<"$expected_logits"
|
||||
for ((step = 0; step < num_decode_tokens; step++)); do
|
||||
if [[ "${actual_token_values[step]}" != "${expected_token_values[step]}" ]]; then
|
||||
printf 'batch position %s: step=%s, baseline token=%s, actual token=%s' \
|
||||
"$position" \
|
||||
"$step" \
|
||||
"${expected_token_values[step]}" \
|
||||
"${actual_token_values[step]}"
|
||||
return
|
||||
fi
|
||||
if [[ "${actual_logit_values[step]}" != "${expected_logit_values[step]}" ]]; then
|
||||
printf 'batch position %s: step=%s, token=%s, baseline logit=%s, actual logit=%s' \
|
||||
"$position" \
|
||||
"$step" \
|
||||
"${expected_token_values[step]}" \
|
||||
"${expected_logit_values[step]}" \
|
||||
"${actual_logit_values[step]}"
|
||||
return
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
run_e2e() {
|
||||
local needle
|
||||
local batch
|
||||
local baseline_output
|
||||
local batch_response
|
||||
local difference
|
||||
local position
|
||||
local -a fillers
|
||||
local -a batch_outputs
|
||||
local -a differences=()
|
||||
|
||||
echo "Testing PowerMoE-3b: TP=2, VLLM_BATCH_INVARIANT=1, VLLM_MOE_SKIP_PADDING=0"
|
||||
server_log=$(mktemp "${TMPDIR:-/tmp}/xpu-batch-invariance.XXXXXX.log")
|
||||
VLLM_BATCH_INVARIANT=1 vllm serve "$model" \
|
||||
--tensor-parallel-size 2 \
|
||||
--max-model-len 1024 \
|
||||
--max-num-seqs 16 \
|
||||
--max-num-batched-tokens 512 \
|
||||
--enable-chunked-prefill \
|
||||
--no-enable-prefix-caching \
|
||||
--enforce-eager \
|
||||
--seed 0 \
|
||||
--max-logprobs 64 \
|
||||
--logprobs-mode raw_logits \
|
||||
--port "$port" > >(tee "$server_log") 2>&1 &
|
||||
server_pid=$!
|
||||
wait_for_server
|
||||
|
||||
needle=$(make_prompt 257)
|
||||
local filler_lengths=(7 31 63 95 127 159 191 223 255 287 319 383 447)
|
||||
for position in "${!filler_lengths[@]}"; do
|
||||
fillers+=( "$(make_prompt "${filler_lengths[position]}" "$((position + 1))")" )
|
||||
done
|
||||
batch="[$needle,${fillers[0]},${fillers[1]},${fillers[2]},${fillers[3]},${fillers[4]},${fillers[5]},$needle,${fillers[6]},${fillers[7]},${fillers[8]},${fillers[9]},${fillers[10]},${fillers[11]},${fillers[12]},$needle]"
|
||||
|
||||
echo "Requesting BS=1 baseline ($num_decode_tokens decode tokens)..."
|
||||
baseline_output=$(request_outputs "$needle" 1)
|
||||
echo "Received BS=1 baseline."
|
||||
echo "Requesting BS=16 batch ($((16 * num_decode_tokens)) decode tokens)..."
|
||||
batch_response=$(request_outputs "$batch" 16)
|
||||
mapfile -t batch_outputs <<<"$batch_response"
|
||||
echo "Received BS=16 batch; comparing generated tokens and logits..."
|
||||
|
||||
for position in 0 7 15; do
|
||||
difference=$(first_difference "${batch_outputs[position]}" "$baseline_output" "$position")
|
||||
if [[ -n "$difference" ]]; then
|
||||
differences+=("$difference")
|
||||
fi
|
||||
done
|
||||
|
||||
if [[ "${#differences[@]}" -ne 0 ]]; then
|
||||
echo "Batch-invariant PowerMoE-3b output changed: ${differences[*]}" >&2
|
||||
return 1
|
||||
fi
|
||||
echo "PowerMoE-3b batch invariance check passed"
|
||||
cleanup_server
|
||||
rm -f "$server_log"
|
||||
server_log=
|
||||
}
|
||||
|
||||
run_e2e
|
||||
+311
@@ -0,0 +1,311 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
set -Eeuo pipefail
|
||||
readonly MODEL="Qwen/Qwen3-0.6B"
|
||||
readonly MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca"
|
||||
readonly EXPECTED_TOKEN_ID="12095"
|
||||
readonly EXPECTED_TEXT=" Paris"
|
||||
readonly CLEANUP_RESERVE_S=300
|
||||
RUN_LIMIT_S="${SNAPSHOT_E2E_RUN_LIMIT_S:-3300}"
|
||||
[[ "$RUN_LIMIT_S" =~ ^[1-9][0-9]{2,3}$ ]] \
|
||||
|| { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be a bounded decimal integer" >&2; exit 1; }
|
||||
(( RUN_LIMIT_S >= CLEANUP_RESERVE_S + 60 && RUN_LIMIT_S <= 3300 )) \
|
||||
|| { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be between 360 and 3300" >&2; exit 1; }
|
||||
readonly RUN_LIMIT_S
|
||||
readonly HOST_HF_CACHE="${SNAPSHOT_E2E_HOST_HF_CACHE:-}"
|
||||
HF_CACHE_MOUNT=()
|
||||
if [[ -n "$HOST_HF_CACHE" ]]; then
|
||||
[[ "$HOST_HF_CACHE" == /* && -d "$HOST_HF_CACHE" ]] \
|
||||
|| { echo "SNAPSHOT_E2E_HOST_HF_CACHE must be an existing absolute directory" >&2; exit 1; }
|
||||
HF_CACHE_MOUNT=(--mount "type=bind,source=$HOST_HF_CACHE,target=/e2e/hf,readonly")
|
||||
PREPARE_ARTIFACT_ROOT='install -d -m 0700 /e2e'
|
||||
else
|
||||
PREPARE_ARTIFACT_ROOT='mkdir -m 0700 /e2e && mkdir -m 0700 /e2e/hf'
|
||||
fi
|
||||
readonly PREPARE_ARTIFACT_ROOT
|
||||
if (( $# != 2 )); then
|
||||
echo "usage: $0 IMAGE_TAG BUILDKITE_COMMIT" >&2
|
||||
exit 2
|
||||
fi
|
||||
readonly IMAGE_TAG="$1"
|
||||
readonly EXPECTED_COMMIT="$2"
|
||||
readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))"
|
||||
RUN_ID="" CONTAINER_NAME="" CONTAINER_ID=""
|
||||
GPU_UUID="" PORT_ONE="" PORT_TWO=""
|
||||
LINK_REMAP_BASELINE=""
|
||||
SNAPSHOT_PIDS=()
|
||||
die() {
|
||||
echo "initialized snapshot E2E: $*" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
run() {
|
||||
local label="$1" maximum_s="$2" available limit
|
||||
shift 2
|
||||
available="$((DEADLINE_EPOCH - $(date +%s) - CLEANUP_RESERVE_S))"
|
||||
(( available > 0 )) || die "no time remains before the cleanup reserve"
|
||||
(( available < maximum_s )) && limit="$available" || limit="$maximum_s"
|
||||
echo "--- ${label} (timeout ${limit}s)" >&2
|
||||
timeout --signal=TERM --kill-after=30s "${limit}s" "$@"
|
||||
}
|
||||
|
||||
link_remaps() {
|
||||
find /dev/shm -maxdepth 1 -name 'link_remap.*' \
|
||||
-printf '%p|%D|%i|%m|%s\n' | sort
|
||||
}
|
||||
|
||||
gpu_pids() {
|
||||
nvidia-smi --id="$GPU_UUID" --query-compute-apps=pid \
|
||||
--format=csv,noheader,nounits | sed '/^[[:space:]]*$/d'
|
||||
}
|
||||
|
||||
port_listeners() {
|
||||
ss -H -ltn "sport = :$1"
|
||||
}
|
||||
|
||||
state_clean() {
|
||||
local pid port output
|
||||
for pid in "${SNAPSHOT_PIDS[@]}"; do
|
||||
[[ ! -e "/proc/$pid" ]] || return 1
|
||||
done
|
||||
output="$(gpu_pids)" || return 1
|
||||
[[ -z "$output" ]] || return 1
|
||||
output="$(link_remaps)" || return 1
|
||||
[[ "$output" == "$LINK_REMAP_BASELINE" ]] || return 1
|
||||
for port in "$PORT_ONE" "$PORT_TWO"; do
|
||||
[[ -n "$port" ]] || continue
|
||||
output="$(port_listeners "$port")" || return 1
|
||||
[[ -z "$output" ]] || return 1
|
||||
done
|
||||
}
|
||||
|
||||
wait_clean() {
|
||||
local _attempt
|
||||
for ((_attempt = 0; _attempt < 60; _attempt++)); do
|
||||
state_clean && return 0
|
||||
sleep 1
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
container_names() {
|
||||
timeout 15s docker container ls --all --format '{{.Names}}'
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
local failed=0 names="" actual="" expected=""
|
||||
if [[ -n "$CONTAINER_NAME" ]]; then
|
||||
names="$(container_names 2>/dev/null)" || return 1
|
||||
if grep -Fxq "$CONTAINER_NAME" <<< "$names"; then
|
||||
actual="$(docker container inspect \
|
||||
--format '{{.Id}}|{{index .Config.Labels "ai.vllm.snapshot.e2e.run"}}' \
|
||||
"$CONTAINER_NAME" 2>/dev/null)" || return 1
|
||||
expected="${CONTAINER_ID:-${actual%%|*}}|$RUN_ID"
|
||||
if [[ "$actual" != "$expected" ]]; then
|
||||
echo "refusing to remove a container with mismatched identity" >&2
|
||||
failed=1
|
||||
else
|
||||
timeout 60s docker rm --force "$CONTAINER_NAME" >/dev/null \
|
||||
|| failed=1
|
||||
fi
|
||||
fi
|
||||
names="$(container_names 2>/dev/null)" || return 1
|
||||
grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1
|
||||
fi
|
||||
if [[ -n "$PORT_TWO" ]] && ! wait_clean; then
|
||||
echo "process, GPU, port, or link-remap residue remained" >&2
|
||||
failed=1
|
||||
fi
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
on_exit() {
|
||||
local status=$?
|
||||
trap - EXIT INT TERM
|
||||
set +e
|
||||
cleanup || status=1
|
||||
exit "$status"
|
||||
}
|
||||
|
||||
trap on_exit EXIT
|
||||
trap 'exit 130' INT
|
||||
trap 'exit 143' TERM
|
||||
|
||||
parse_inspect() {
|
||||
timeout 30s python3 -c '
|
||||
import json, math, sys
|
||||
data = json.load(sys.stdin)
|
||||
model, revision, gpu_uuid, token_id, text = sys.argv[1:]
|
||||
expected = {
|
||||
"boundary": "post-engine-init-reloadable-state-released",
|
||||
"model": model, "model_revision": revision,
|
||||
"gpu_uuid": gpu_uuid, "oracle_token_ids": [int(token_id)],
|
||||
"oracle_text": text,
|
||||
}
|
||||
bad = next((field for field, value in expected.items() if data.get(field) != value), None)
|
||||
if bad:
|
||||
raise SystemExit(f"inspect field {bad} mismatch: {data.get(bad)!r}")
|
||||
logprob = data.get("oracle_sampled_token_logprob")
|
||||
if type(logprob) is not float or not math.isfinite(logprob):
|
||||
raise SystemExit("inspect logprob is not finite")
|
||||
pids = data.get("process_tree")
|
||||
if not pids:
|
||||
raise SystemExit("inspect process_tree is empty")
|
||||
for pid in pids:
|
||||
print(f"pid={pid}")
|
||||
print(f"oracle_logprob={logprob}")
|
||||
' "$MODEL" "$MODEL_REVISION" "$GPU_UUID" "$EXPECTED_TOKEN_ID" \
|
||||
"$EXPECTED_TEXT"
|
||||
}
|
||||
|
||||
assert_oracle() {
|
||||
local port="$1" ordinal="$2" response
|
||||
response="$(run "restore ${ordinal} public completion" 120 curl \
|
||||
--fail --silent --show-error --max-time 120 \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data "{\"model\":\"${MODEL}\",\"prompt\":\"The capital of France is\",\"temperature\":0,\"seed\":0,\"min_tokens\":1,\"max_tokens\":1,\"logprobs\":0,\"return_token_ids\":true}" \
|
||||
"http://127.0.0.1:${port}/v1/completions")"
|
||||
printf '%s' "$response" | timeout 30s python3 -c '
|
||||
import json, math, sys
|
||||
choice = json.load(sys.stdin)["choices"][0]
|
||||
token_id, text, expected, ordinal = sys.argv[1:]
|
||||
values = choice["logprobs"]["token_logprobs"]
|
||||
actual = float(values[0]) if len(values) == 1 else math.nan
|
||||
if (choice["token_ids"] != [int(token_id)] or choice["text"] != text
|
||||
or not math.isfinite(actual) or abs(actual - float(expected)) > 1e-3):
|
||||
raise SystemExit(f"public oracle mismatch: {choice!r}")
|
||||
print(f"restore={ordinal} token_id={token_id} text={text!r} logprob={actual}")
|
||||
' "$EXPECTED_TOKEN_ID" "$EXPECTED_TEXT" "$INSPECT_LOGPROB" "$ordinal"
|
||||
}
|
||||
|
||||
[[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] \
|
||||
|| die "host must be Linux x86_64"
|
||||
[[ "$EXPECTED_COMMIT" =~ ^[0-9a-f]{40}$ ]] \
|
||||
|| die "BUILDKITE_COMMIT must be a full lowercase commit hash"
|
||||
[[ "$(< /proc/sys/kernel/io_uring_disabled)" == "2" ]] \
|
||||
|| die "kernel.io_uring_disabled must already equal 2"
|
||||
run "Docker preflight" 60 docker info >/dev/null
|
||||
|
||||
VISIBLE_GPU="${NVIDIA_VISIBLE_DEVICES:-}"
|
||||
if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
|
||||
VISIBLE_GPU="${CUDA_VISIBLE_DEVICES:-}"
|
||||
fi
|
||||
if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
|
||||
CDI_GPU="${BUILDKITE_PLUGIN_DOCKER_DEVICES_0:-}"
|
||||
if [[ "$CDI_GPU" == nvidia.com/gpu=* ]]; then
|
||||
VISIBLE_GPU="${CDI_GPU#nvidia.com/gpu=}"
|
||||
else
|
||||
VISIBLE_GPU=""
|
||||
fi
|
||||
fi
|
||||
[[ -n "$VISIBLE_GPU" && "$VISIBLE_GPU" != all && "$VISIBLE_GPU" != void \
|
||||
&& "$VISIBLE_GPU" != *","* ]] \
|
||||
|| die "the runner must expose exactly one selected GPU"
|
||||
GPU_ROW="$(nvidia-smi --id="$VISIBLE_GPU" \
|
||||
--query-gpu=name,uuid,mig.mode.current,memory.total \
|
||||
--format=csv,noheader,nounits)"
|
||||
IFS=',' read -r GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB <<< "$GPU_ROW"
|
||||
for variable in GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB; do
|
||||
printf -v "$variable" '%s' "${!variable#"${!variable%%[![:space:]]*}"}"
|
||||
done
|
||||
[[ "$GPU_NAME" == *H200* && "$GPU_UUID" == GPU-* && "$MIG_MODE" == Disabled ]] \
|
||||
|| die "selected device is not a full non-MIG H200: $GPU_ROW"
|
||||
(( GPU_MEMORY_MIB >= 130000 )) || die "selected H200 does not expose full memory"
|
||||
GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
|
||||
[[ -z "$GPU_PIDS" ]] || die "selected H200 is not idle"
|
||||
|
||||
RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM"
|
||||
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))')
|
||||
LINK_REMAP_BASELINE="$(link_remaps)"
|
||||
|
||||
run "pull exact candidate image" 900 docker pull --platform linux/amd64 \
|
||||
"$IMAGE_TAG" >/dev/null
|
||||
IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <(
|
||||
docker image inspect \
|
||||
--format '{{.Id}}|{{.Os}}/{{.Architecture}}|{{index .Config.Labels "org.opencontainers.image.revision"}}|{{index .Config.Labels "ai.vllm.build.commit"}}' \
|
||||
"$IMAGE_TAG"
|
||||
)
|
||||
[[ "$IMAGE_ID" =~ ^sha256:[0-9a-f]{64}$ && "$IMAGE_PLATFORM" == linux/amd64 ]] \
|
||||
|| die "candidate image identity is invalid"
|
||||
[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \
|
||||
|| die "candidate image commit labels do not match"
|
||||
CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \
|
||||
--name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \
|
||||
--gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \
|
||||
--network=host --env "CUDA_VISIBLE_DEVICES=$GPU_UUID" --env HF_HUB_DISABLE_TELEMETRY=1 \
|
||||
--env VLLM_NO_USAGE_STATS=1 --env VLLM_USE_V2_MODEL_RUNNER=1 \
|
||||
${HF_CACHE_MOUNT[@]+"${HF_CACHE_MOUNT[@]}"} \
|
||||
--entrypoint sleep "$IMAGE_ID" infinity)"
|
||||
[[ "$CONTAINER_ID" =~ ^[0-9a-f]{64}$ ]] || die "invalid container ID"
|
||||
run "prepare private artifact root" 60 docker exec "$CONTAINER_NAME" sh -c \
|
||||
"$PREPARE_ARTIFACT_ROOT && df -Pk / | awk 'NR == 2 && \$4 >= 10485760 {ok=1} END {exit !ok}'"
|
||||
run "verify snapshot runtime" 60 docker exec "$CONTAINER_NAME" sh -c \
|
||||
'criu --version; test -x /usr/local/sbin/cuda-checkpoint; test -f /usr/local/lib/criu/cuda_plugin.so'
|
||||
run "verify selected CUDA device" 60 docker exec "$CONTAINER_NAME" python3 -c '
|
||||
import sys, torch
|
||||
if torch.cuda.device_count() != 1:
|
||||
raise SystemExit("expected exactly one visible CUDA device")
|
||||
actual = str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-")
|
||||
if actual != sys.argv[1].removeprefix("GPU-"):
|
||||
raise SystemExit(f"CUDA device {actual} does not match {sys.argv[1]}")
|
||||
print(f"selected_cuda_uuid=GPU-{actual}")
|
||||
' "$GPU_UUID"
|
||||
OFFLINE_EXEC=(docker exec --env HF_HOME=/e2e/hf --env HF_HUB_OFFLINE=1 \
|
||||
--env TRANSFORMERS_OFFLINE=1 --env VLLM_SNAPSHOT_TIMEOUT_S=900 \
|
||||
"$CONTAINER_NAME")
|
||||
if [[ -n "$HOST_HF_CACHE" ]]; then
|
||||
run "validate pinned model in read-only host cache" 120 \
|
||||
"${OFFLINE_EXEC[@]}" python3 -c \
|
||||
'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2], local_files_only=True); assert Path(p).name == sys.argv[2]' \
|
||||
"$MODEL" "$MODEL_REVISION"
|
||||
else
|
||||
run "prefetch pinned public model" 900 docker exec --env HF_HOME=/e2e/hf \
|
||||
"$CONTAINER_NAME" python3 -c \
|
||||
'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2]); assert Path(p).name == sys.argv[2]' \
|
||||
"$MODEL" "$MODEL_REVISION"
|
||||
fi
|
||||
|
||||
run "create compact initialized snapshot" 1200 "${OFFLINE_EXEC[@]}" \
|
||||
vllm snapshot create "$MODEL" --revision "$MODEL_REVISION" \
|
||||
--tokenizer-revision "$MODEL_REVISION" --snapshot-dir /e2e/artifact \
|
||||
--dtype float16 --max-model-len 512 --gpu-memory-utilization 0.50
|
||||
wait_clean || die "create left process, GPU, port, or link-remap residue"
|
||||
|
||||
INSPECT_JSON="$(run "inspect compact initialized snapshot" 120 \
|
||||
"${OFFLINE_EXEC[@]}" vllm snapshot inspect /e2e/artifact)"
|
||||
INSPECT_RECEIPT="$(printf '%s' "$INSPECT_JSON" | parse_inspect)"
|
||||
mapfile -t SNAPSHOT_PIDS < <(sed -n 's/^pid=//p' <<< "$INSPECT_RECEIPT")
|
||||
INSPECT_LOGPROB="$(sed -n 's/^oracle_logprob=//p' <<< "$INSPECT_RECEIPT")"
|
||||
(( ${#SNAPSHOT_PIDS[@]} > 0 )) || die "inspect returned no PIDs"
|
||||
state_clean || die "create left recorded state alive"
|
||||
echo "$INSPECT_RECEIPT"
|
||||
|
||||
run "restore compact snapshot 1" 900 "${OFFLINE_EXEC[@]}" \
|
||||
vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_ONE"
|
||||
PORT_LISTENERS="$(port_listeners "$PORT_ONE")" || die "failed to inspect restore 1 port"
|
||||
[[ -n "$PORT_LISTENERS" ]] || die "restore 1 did not bind its port"
|
||||
GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
|
||||
[[ -n "$GPU_PIDS" ]] || die "restore 1 did not engage the selected GPU"
|
||||
assert_oracle "$PORT_ONE" 1
|
||||
run "stop container after restore 1" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
|
||||
[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
|
||||
|| die "docker stop replaced the container"
|
||||
wait_clean || die "restore 1 teardown left residue"
|
||||
|
||||
run "start the same container" 60 docker start "$CONTAINER_NAME" >/dev/null
|
||||
[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
|
||||
|| die "docker start replaced the container"
|
||||
run "restore compact snapshot 2" 900 "${OFFLINE_EXEC[@]}" \
|
||||
vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_TWO"
|
||||
PORT_LISTENERS="$(port_listeners "$PORT_TWO")" || die "failed to inspect restore 2 port"
|
||||
[[ -n "$PORT_LISTENERS" ]] || die "restore 2 did not bind its port"
|
||||
GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
|
||||
[[ -n "$GPU_PIDS" ]] || die "restore 2 did not engage the selected GPU"
|
||||
assert_oracle "$PORT_TWO" 2
|
||||
run "stop container after restore 2" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
|
||||
wait_clean || die "restore 2 teardown left residue"
|
||||
cleanup || die "exact cleanup failed"
|
||||
echo "container_id=$CONTAINER_ID restore_1_port=$PORT_ONE restore_2_port=$PORT_TWO"
|
||||
echo "result=PASS"
|
||||
@@ -9,9 +9,20 @@ if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None
|
||||
exit 0
|
||||
fi
|
||||
|
||||
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
|
||||
# Select requirements before optional packages can change the installed Torch.
|
||||
CUDA_MAJOR=$(python3 -c 'import torch; print(torch.version.cuda.split(".", 1)[0])')
|
||||
DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors.txt
|
||||
if [ "${CUDA_MAJOR}" = "12" ]; then
|
||||
DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors_cu12.txt
|
||||
fi
|
||||
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-${DEFAULT_REQUIREMENTS}}"
|
||||
|
||||
uv pip install --system -r "${REQUIREMENTS_FILE}"
|
||||
CONSTRAINTS_DIR=$(mktemp -d)
|
||||
trap 'rm -rf "${CONSTRAINTS_DIR}"' EXIT
|
||||
uv pip freeze --system > "${CONSTRAINTS_DIR}/installed.txt"
|
||||
grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
|
||||
"${CONSTRAINTS_DIR}/installed.txt" > "${CONSTRAINTS_DIR}/constraints.txt"
|
||||
uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" -r "${REQUIREMENTS_FILE}"
|
||||
|
||||
KV_METADATA=$(python3 - <<'PY'
|
||||
import importlib.metadata as metadata
|
||||
@@ -56,5 +67,5 @@ PY
|
||||
# build first to avoid a clash.
|
||||
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then
|
||||
uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true
|
||||
uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
|
||||
uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
|
||||
fi
|
||||
|
||||
@@ -61,6 +61,13 @@ if target_enabled cuda-13-0; then
|
||||
docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64"
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}"
|
||||
|
||||
ZSTD_DIGEST=$(docker buildx imagetools inspect \
|
||||
"public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" \
|
||||
--format '{{json .Manifest.Digest}}' | tr -d '"')
|
||||
.buildkite/scripts/publish-zstd-image.sh \
|
||||
"public.ecr.aws/q9t5s3a7/vllm-release-repo@${ZSTD_DIGEST}" \
|
||||
"vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-zstd"
|
||||
fi
|
||||
|
||||
# ---- CUDA 12.9 ----
|
||||
|
||||
Executable
+45
@@ -0,0 +1,45 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [ "$#" -ne 2 ]; then
|
||||
echo "Usage: $0 SOURCE@sha256:DIGEST DESTINATION-zstd" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SOURCE=$1
|
||||
DESTINATION=$2
|
||||
if [[ ! "$SOURCE" =~ ^[a-zA-Z0-9][a-zA-Z0-9._:/-]*@sha256:[a-f0-9]{64}$ ]] ||
|
||||
[[ ! "$DESTINATION" =~ ^[a-zA-Z0-9][a-zA-Z0-9._:/-]*:[a-zA-Z0-9_][a-zA-Z0-9_.-]*-zstd$ ]]; then
|
||||
echo "ERROR: source must use an immutable sha256 digest and destination must end in -zstd" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
BUILDER=
|
||||
cleanup() {
|
||||
status=$?
|
||||
trap - EXIT INT TERM
|
||||
if [ -n "$BUILDER" ] && ! docker buildx rm --force "$BUILDER"; then
|
||||
echo "ERROR: failed to remove Buildx builder $BUILDER" >&2
|
||||
[ "$status" -ne 0 ] || status=1
|
||||
fi
|
||||
exit "$status"
|
||||
}
|
||||
trap cleanup EXIT
|
||||
trap 'exit 130' INT
|
||||
trap 'exit 143' TERM
|
||||
|
||||
BUILDER=$(docker buildx create \
|
||||
--driver docker-container \
|
||||
--driver-opt network=host \
|
||||
--driver-opt 'image=moby/buildkit:v0.32.2@sha256:28a898719c18a33f4e8000685287fa36fd0dd9560c6440227d3a732d79bb41d8')
|
||||
|
||||
printf '# check=error=true\n\nFROM %s\n' "$SOURCE" | docker buildx build \
|
||||
--builder "$BUILDER" \
|
||||
--platform linux/amd64 \
|
||||
--progress plain \
|
||||
--provenance=mode=min \
|
||||
--output "type=image,name=${DESTINATION},push=true,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true" \
|
||||
-
|
||||
+381
-34
@@ -436,7 +436,7 @@ steps:
|
||||
- tests/v1/e2e/general/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
|
||||
|
||||
- label: ":amd: (MI250) E2E Scheduling" # TBD
|
||||
timeout_in_minutes: 55
|
||||
@@ -1361,6 +1361,23 @@ steps:
|
||||
- bash ../tools/install_torchcodec_rocm.sh || exit 1
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
|
||||
- label: ":amd: (MI300) Scale-out EC Connector E2E"
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/ec_transfer/
|
||||
- vllm/entrypoints/scale_out/
|
||||
- tests/entrypoints/scale_out/ec_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
|
||||
|
||||
#----------------------------------------------------------- mi300 · evals -----------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD
|
||||
@@ -1737,29 +1754,6 @@ steps:
|
||||
- pytest -v -s tests/ir
|
||||
- pytest -v -s tests/kernels/ir
|
||||
|
||||
- label: ":amd: (MI300) MLA Kernels" # TBD
|
||||
timeout_in_minutes: 25
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
working_dir: "/vllm-workspace/"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention/backends/mla/rocm_aiter_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
optional: true
|
||||
commands:
|
||||
- amd-smi
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
|
||||
- label: ":amd: (MI300) Attention Kernels Shard %N" # TBD
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -1778,7 +1772,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/utils/flashinfer.py
|
||||
commands:
|
||||
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
- pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: ":amd: (MI300) MiniMax Reduce RMS Kernels" # TBD
|
||||
timeout_in_minutes: 45
|
||||
@@ -2566,7 +2560,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/entrypoints/generate/beam_search
|
||||
commands:
|
||||
- pytest -v -s samplers
|
||||
- pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
|
||||
|
||||
#------------------------------------------------------------ mi300 · misc ------------------------------------------------------------#
|
||||
|
||||
@@ -3213,6 +3207,58 @@ steps:
|
||||
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
|
||||
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism
|
||||
|
||||
- label: ":amd: (MI300) Mooncake EC TCP E2E"
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/ec_transfer/
|
||||
- vllm/config/ec_transfer.py
|
||||
- vllm/config/multimodal.py
|
||||
- vllm/config/vllm.py
|
||||
- vllm/multimodal/
|
||||
- vllm/v1/core/encoder_cache_manager.py
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/ec_connector_model_runner_mixin.py
|
||||
- vllm/v1/worker/gpu_model_runner.py
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- vllm/v1/worker/gpu/ec_connector.py
|
||||
- vllm/v1/worker/gpu/model_runner.py
|
||||
- vllm/v1/worker/mm_encoder_model_runner.py
|
||||
- vllm/v1/worker/encoder_cudagraph.py
|
||||
- vllm/v1/worker/encoder_cudagraph_defs.py
|
||||
- vllm/v1/attention/ops/vit_attn_wrappers.py
|
||||
- vllm/v1/worker/gpu/model_states/
|
||||
- vllm/v1/worker/gpu/mm/
|
||||
- tests/v1/cudagraph/test_encoder_cudagraph.py
|
||||
- tests/models/multimodal/generation/test_vit_cudagraph.py
|
||||
- examples/disaggregated/disaggregated_encoder/
|
||||
- tests/v1/ec_connector/
|
||||
- requirements/kv_connectors.txt
|
||||
- vllm/platforms/rocm.py
|
||||
- requirements/rocm.txt
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export PYTHON_BIN='python'
|
||||
- export PYTHONUNBUFFERED='1'
|
||||
- export VLLM_HOST_IP='127.0.0.1'
|
||||
- export VLLM_USE_V2_MODEL_RUNNER='1'
|
||||
- export E_CUDAGRAPH_MM_ENCODER='1'
|
||||
- export E_MM_ENCODER_ATTN_BACKEND='TRITON_ATTN'
|
||||
- export MOONCAKE_EC_PROTOCOL='tcp'
|
||||
- export USE_MM_PROMPTS='1'
|
||||
- export SKIP_BASELINE='0'
|
||||
- export CONCURRENCY='3'
|
||||
- export REPEAT='2'
|
||||
- export LOG_PATH='/tmp/mooncake-ec-e2e'
|
||||
- export BASELINE_FILE='/tmp/mooncake-ec-e2e/baseline.json'
|
||||
- bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
|
||||
|
||||
- label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -3230,7 +3276,7 @@ steps:
|
||||
- CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: ":amd: (MI300) Distributed DP Extended" # TBD
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 75
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_4
|
||||
@@ -3364,8 +3410,71 @@ steps:
|
||||
commands:
|
||||
- python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k"
|
||||
|
||||
#---------------------------------------------------------- mi355 · compile ----------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI355) PyTorch Compilation H100 Cases"
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/_custom_ops.py
|
||||
- vllm/compilation/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- "!vllm/distributed/kv_transfer/"
|
||||
- vllm/engine/
|
||||
- vllm/env_override.py
|
||||
- vllm/envs.py
|
||||
- vllm/forward_context.py
|
||||
- vllm/inputs/
|
||||
- vllm/ir/
|
||||
- vllm/kernels/
|
||||
- vllm/logger.py
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/plugins/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/sequence.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/triton_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/compile/h100/
|
||||
- tests/utils.py
|
||||
- csrc/
|
||||
commands:
|
||||
- pytest -v -s compile/h100/test_startup.py
|
||||
|
||||
#-------------------------------------------------------- mi355 · distributed --------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI355) Sharded RDT Weight Transfer"
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/weight_transfer/
|
||||
- tests/distributed/test_sharded_rdt_plan.py
|
||||
- tests/distributed/test_sharded_rdt_producer.py
|
||||
- tests/distributed/test_sharded_rdt_trainer.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- uv pip install --system --no-deps 'ray==2.56.1'
|
||||
- pytest -v -s distributed/test_sharded_rdt_plan.py
|
||||
- pytest -v -s distributed/test_sharded_rdt_producer.py
|
||||
- pytest -v -s distributed/test_sharded_rdt_trainer.py
|
||||
|
||||
- label: ":amd: (MI355) Distributed DP Basic" # TBD
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3445,6 +3554,35 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP"
|
||||
timeout_in_minutes: 105
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_4
|
||||
num_gpus: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/models/kimi_k3/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
- requirements/kv_connectors_rocm.txt
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export ENABLE_HMA_FLAG='1'
|
||||
- export DP_EP='1'
|
||||
- export GPU_MEMORY_UTILIZATION='0.9'
|
||||
- export VLLM_SSM_CONV_STATE_LAYOUT='DS'
|
||||
- export PREFILLER_TP_SIZE='2'
|
||||
- export DECODER_TP_SIZE='2'
|
||||
- export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct'
|
||||
- export PREFILL_BLOCK_SIZE='2048'
|
||||
- export DECODE_BLOCK_SIZE='2048'
|
||||
- export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768'
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
|
||||
|
||||
- label: ":amd: (MI355) Distributed Features" # TBD
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3642,6 +3780,51 @@ steps:
|
||||
- uv pip install --system 'gpt-oss[eval]==0.0.5'
|
||||
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt
|
||||
|
||||
- label: ":amd: (MI355) LM Eval Watermarking"
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- .buildkite/test_areas/lm_eval.yaml
|
||||
- tests/evals/gsm8k/
|
||||
- vllm/config/watermarking.py
|
||||
- vllm/sampling_params.py
|
||||
- vllm/v1/watermarking/
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
|
||||
|
||||
- label: ":amd: (MI355) LM Eval Watermark Feature Combination"
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- .buildkite/test_areas/lm_eval.yaml
|
||||
- tests/evals/gsm8k/
|
||||
- vllm/config/speculative.py
|
||||
- vllm/config/vllm.py
|
||||
- vllm/config/watermarking.py
|
||||
- vllm/sampling_params.py
|
||||
- vllm/model_executor/models/qwen3_dspark.py
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- vllm/model_executor/warmup/
|
||||
- vllm/v1/watermarking/
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
|
||||
|
||||
- label: ":amd: (MI355) LM Eval Qwen3-5 Models" # TBD
|
||||
timeout_in_minutes: 40
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3665,6 +3848,27 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt
|
||||
|
||||
- label: ":amd: (MI355) Qwen3-8-Flash-Next-FP8 Accuracy Eval"
|
||||
timeout_in_minutes: 105
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_4
|
||||
num_gpus: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/libtorch_stable/gdn/
|
||||
- tests/evals/qwen4_exp/
|
||||
- vllm/model_executor/layers/mamba/
|
||||
- vllm/models/qwen4_exp/
|
||||
- vllm/transformers_utils/configs/qwen4_exp.py
|
||||
- vllm/v1/attention/backends/short_conv_attn.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- uv pip install --system 'evalscope==1.10.0'
|
||||
- pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt
|
||||
|
||||
- label: ":amd: (MI355) LM Eval Small Models Distributed"
|
||||
timeout_in_minutes: 55
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3686,6 +3890,25 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
|
||||
|
||||
- label: ":amd: (MI355) MoE Refactor Integration Shard %N"
|
||||
timeout_in_minutes: 105
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
parallelism: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/platforms/rocm.py
|
||||
- tests/evals/gsm8k/
|
||||
- vllm/model_executor/layers/fused_moe/
|
||||
- vllm/model_executor/layers/quantization/
|
||||
- vllm/_aiter_ops.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt
|
||||
|
||||
- label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD
|
||||
timeout_in_minutes: 25
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3790,6 +4013,23 @@ steps:
|
||||
|
||||
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI355) Attention DiffKV Kernels"
|
||||
timeout_in_minutes: 35
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention/ops/triton_unified_attention_diffkv.py
|
||||
- vllm/v1/attention/backends/triton_attn_diffkv.py
|
||||
- tests/kernels/attention/test_triton_unified_attention.py
|
||||
- tests/kernels/attention/test_triton_unified_attention_diffkv.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
|
||||
|
||||
- label: ":amd: (MI355) Core Operation Kernels Shard %N" # TBD
|
||||
timeout_in_minutes: 40
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3859,19 +4099,19 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
- examples/basic/offline_inference/chat.py
|
||||
- tests/kernels/attention/test_attention_selector.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
optional: true
|
||||
commands:
|
||||
- amd-smi
|
||||
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
|
||||
- pytest -v -s tests/kernels/attention/test_attention_selector.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
|
||||
- label: ":amd: (MI355) MLA Kernels" # TBD
|
||||
timeout_in_minutes: 25
|
||||
- label: ":amd: (MI355) MLA Kernels"
|
||||
timeout_in_minutes: 40
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention/backends/mla/rocm_aiter_mla.py
|
||||
@@ -3882,7 +4122,6 @@ steps:
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
optional: true
|
||||
commands:
|
||||
- amd-smi
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
@@ -3890,6 +4129,22 @@ steps:
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
|
||||
- label: ":amd: (MI355) DeepSeek V4 Fused Kernels"
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
|
||||
- vllm/models/deepseek_v4/common/ops/
|
||||
- tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
|
||||
|
||||
- label: ":amd: (MI355) Attention Kernels Shard %N" # TBD
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3908,7 +4163,7 @@ steps:
|
||||
- vllm/utils/flashinfer.py
|
||||
optional: true
|
||||
commands:
|
||||
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
- pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: ":amd: (MI355) MoE Kernels Shard %N" # TBD
|
||||
timeout_in_minutes: 55
|
||||
@@ -4050,6 +4305,21 @@ steps:
|
||||
|
||||
#------------------------------------------------------ mi355 · models / basic -------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI355) GLM5Next"
|
||||
timeout_in_minutes: 35
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/models/glm5next/
|
||||
- tests/models/glm5next/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s models/glm5next
|
||||
|
||||
- label: ":amd: (MI355) Basic Models (Extra Initialization) Shard %N" # TBD
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -4087,7 +4357,27 @@ steps:
|
||||
- export GIT_CONFIG_VALUE_0=HTTP/1.1
|
||||
- MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
- label: ":amd: (MI355) Granite Language Model Compatibility"
|
||||
timeout_in_minutes: 80
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- "!vllm/distributed/kv_transfer/"
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- export GIT_CONFIG_COUNT=1
|
||||
- export GIT_CONFIG_KEY_0=http.version
|
||||
- export GIT_CONFIG_VALUE_0=HTTP/1.1
|
||||
- MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
|
||||
|
||||
- label: ":amd: (MI355) Language Models (Extended Generation)" # TBD
|
||||
timeout_in_minutes: 95
|
||||
@@ -4308,8 +4598,65 @@ steps:
|
||||
# commands:
|
||||
# - pytest -s -v tests/quantization/test_gfx950_moe.py
|
||||
|
||||
#--------------------------------------------------------- mi355 · samplers ----------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI355) Samplers Multimodal Beam Search"
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers
|
||||
- vllm/sampling_metadata.py
|
||||
- vllm/v1/sample/
|
||||
- vllm/entrypoints/generate/beam_search/
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
|
||||
|
||||
#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------#
|
||||
|
||||
- label: ":amd: (MI355) E2E Core Large Memory"
|
||||
timeout_in_minutes: 65
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
- tests/v1/e2e/general/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
|
||||
- pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
|
||||
|
||||
- label: ":amd: (MI355) Extract Hidden States Integration Single GPU"
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
num_gpus: 1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/extract_hidden_states.py
|
||||
- vllm/model_executor/models/extract_hidden_states.py
|
||||
- vllm/transformers_utils/configs/extract_hidden_states.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: ":amd: (MI355) V1 Attention Shard %N" # TBD
|
||||
timeout_in_minutes: 125
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
|
||||
@@ -318,3 +318,24 @@ steps:
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP"
|
||||
dind: false
|
||||
device: mi355_4
|
||||
num_devices: 4
|
||||
timeout_in_minutes: 105
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/models/kimi_k3/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
- requirements/kv_connectors_rocm.txt
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
|
||||
|
||||
@@ -14,3 +14,22 @@ steps:
|
||||
- tests/entrypoints/scale_out/ec_integration/
|
||||
commands:
|
||||
- bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI300) Scale-out EC Connector E2E"
|
||||
dind: false
|
||||
device: mi300_2
|
||||
num_devices: 2
|
||||
timeout_in_minutes: 60
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/ec_transfer/
|
||||
- vllm/entrypoints/scale_out/
|
||||
- tests/entrypoints/scale_out/ec_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
|
||||
|
||||
@@ -83,3 +83,45 @@ steps:
|
||||
PY
|
||||
)"
|
||||
- bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI300) Mooncake EC TCP E2E"
|
||||
dind: false
|
||||
device: mi300_2
|
||||
num_devices: 2
|
||||
timeout_in_minutes: 60
|
||||
working_dir: /vllm-workspace
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/ec_transfer/
|
||||
- vllm/config/ec_transfer.py
|
||||
- vllm/config/multimodal.py
|
||||
- vllm/config/vllm.py
|
||||
- vllm/multimodal/
|
||||
- vllm/v1/core/encoder_cache_manager.py
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/ec_connector_model_runner_mixin.py
|
||||
- vllm/v1/worker/gpu_model_runner.py
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- vllm/v1/worker/gpu/ec_connector.py
|
||||
- vllm/v1/worker/gpu/model_runner.py
|
||||
- vllm/v1/worker/mm_encoder_model_runner.py
|
||||
- vllm/v1/worker/encoder_cudagraph.py
|
||||
- vllm/v1/worker/encoder_cudagraph_defs.py
|
||||
- vllm/v1/attention/ops/vit_attn_wrappers.py
|
||||
- vllm/v1/worker/gpu/model_states/
|
||||
- vllm/v1/worker/gpu/mm/
|
||||
- tests/v1/cudagraph/test_encoder_cudagraph.py
|
||||
- tests/models/multimodal/generation/test_vit_cudagraph.py
|
||||
- examples/disaggregated/disaggregated_encoder/
|
||||
- tests/v1/ec_connector/
|
||||
- requirements/kv_connectors.txt
|
||||
- vllm/platforms/rocm.py
|
||||
- requirements/rocm.txt
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
PYTHON_BIN: "python"
|
||||
commands:
|
||||
- bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
|
||||
|
||||
@@ -33,6 +33,29 @@ steps:
|
||||
- pytest -v -s distributed/test_sharded_rdt_plan.py
|
||||
- pytest -v -s distributed/test_sharded_rdt_producer.py
|
||||
- pytest -v -s distributed/test_sharded_rdt_trainer.py
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Sharded RDT Weight Transfer"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 45
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/weight_transfer/
|
||||
- tests/distributed/test_sharded_rdt_plan.py
|
||||
- tests/distributed/test_sharded_rdt_producer.py
|
||||
- tests/distributed/test_sharded_rdt_trainer.py
|
||||
- vllm/platforms/rocm.py
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- uv pip install --system --no-deps 'ray==2.56.1'
|
||||
- pytest -v -s distributed/test_sharded_rdt_plan.py
|
||||
- pytest -v -s distributed/test_sharded_rdt_producer.py
|
||||
- pytest -v -s distributed/test_sharded_rdt_trainer.py
|
||||
|
||||
- label: ":nvidia: (L4) Distributed DP Basic"
|
||||
key: distributed-dp-tests-2-gpus
|
||||
@@ -364,17 +387,18 @@ steps:
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
|
||||
|
||||
- label: ":nvidia: (B200) Kimi-K3 GEMM-RS/AR"
|
||||
key: kimi-k3-gemm-rs-ar-2xb200
|
||||
- label: ":nvidia: (B200) GEMM-RS/AR"
|
||||
key: gemm-rs-ar-2xb200
|
||||
timeout_in_minutes: 20
|
||||
device: b200-k8s
|
||||
working_dir: "/vllm-workspace/"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
|
||||
- tests/kernels/test_kimi_k3_gemm_rs_ar.py
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
|
||||
- vllm/cute_utils/
|
||||
- tests/kernels/test_gemm_rs_ar.py
|
||||
commands:
|
||||
- pytest -v -s tests/kernels/test_kimi_k3_gemm_rs_ar.py
|
||||
- pytest -v -s tests/kernels/test_gemm_rs_ar.py
|
||||
|
||||
- label: ":nvidia: (L4) Distributed 2-Node"
|
||||
key: 2-node-test-4-gpus
|
||||
|
||||
@@ -94,7 +94,7 @@ steps:
|
||||
label: ":amd: (MI250) E2E Core"
|
||||
dind: false
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
@@ -115,6 +115,23 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
|
||||
- pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) E2E Core Large Memory"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 65
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
- tests/v1/e2e/general/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
|
||||
- pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
|
||||
|
||||
- label: ":nvidia: (L4) V1 E2E"
|
||||
key: v1-e2e-2-gpus
|
||||
|
||||
@@ -2,6 +2,26 @@ group: Entrypoints
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: ":nvidia: (H200) Initialized Snapshot E2E"
|
||||
key: initialized-snapshot-e2e
|
||||
device: h200
|
||||
num_devices: 1
|
||||
no_plugin: true
|
||||
optional: true
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "."
|
||||
env:
|
||||
SNAPSHOT_E2E_RUN_LIMIT_S: "1200"
|
||||
source_file_dependencies:
|
||||
- .buildkite/scripts/initialized-snapshot-e2e.sh
|
||||
- .buildkite/test_areas/entrypoints.yaml
|
||||
- docker/Dockerfile
|
||||
- tools/install_snapshot_runtime.sh
|
||||
- vllm/entrypoints/cli/
|
||||
- vllm/snapshot/
|
||||
commands:
|
||||
- bash .buildkite/scripts/initialized-snapshot-e2e.sh "$IMAGE_TAG" "$BUILDKITE_COMMIT"
|
||||
|
||||
- label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit"
|
||||
device: h200_35gb
|
||||
key: entrypoints-unit-tests
|
||||
|
||||
@@ -111,6 +111,23 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s kernels/test_fused_deepseek_v4_*.py
|
||||
- pytest -v -s kernels/test_top_k_per_row.py
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) DeepSeek V4 Fused Kernels"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 60
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
|
||||
- vllm/models/deepseek_v4/common/ops/
|
||||
- tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
|
||||
|
||||
- label: ":nvidia: (B200) DeepSeek V4 Kernels"
|
||||
key: deepseek-v4-kernel-test-b200
|
||||
@@ -196,6 +213,8 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/envs.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: ":nvidia: (H100) Attention DiffKV Kernels"
|
||||
key: kernels-attention-diffkv-test-h100
|
||||
@@ -209,6 +228,24 @@ steps:
|
||||
- tests/kernels/attention/test_triton_unified_attention_diffkv.py
|
||||
commands:
|
||||
- pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Attention DiffKV Kernels"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 35
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention/ops/triton_unified_attention_diffkv.py
|
||||
- vllm/v1/attention/backends/triton_attn_diffkv.py
|
||||
- tests/kernels/attention/test_triton_unified_attention.py
|
||||
- tests/kernels/attention/test_triton_unified_attention_diffkv.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
|
||||
|
||||
- label: ":nvidia: (H100) MLA Kernel Test"
|
||||
key: kernels-flashmla-test-h100
|
||||
@@ -231,6 +268,31 @@ steps:
|
||||
- pytest -v -s kernels/attention/test_flashmla_sparse.py
|
||||
- pytest -v -s kernels/attention/test_flashinfer_mla_decode.py
|
||||
- pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) MLA Kernels"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 40
|
||||
working_dir: /vllm-workspace/
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention/backends/mla/rocm_aiter_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
commands:
|
||||
- amd-smi
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
|
||||
- label: ":nvidia: (L4) Quantization Kernels Shard %N"
|
||||
key: kernels-quantization-test
|
||||
@@ -640,31 +702,6 @@ steps:
|
||||
- pytest -v -s kernels/moe/test_moe_layer.py
|
||||
- pytest -v -s kernels/moe/test_deepep_v2_moe.py
|
||||
|
||||
- label: ":amd: (MI355) MLA Kernels"
|
||||
key: amd-mla-kernels-mi355
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 25
|
||||
working_dir: /vllm-workspace/
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention/backends/mla/rocm_aiter_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
commands:
|
||||
- amd-smi
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
|
||||
|
||||
- label: ":amd: (MI355) FP8 MoE Kernels"
|
||||
key: amd-fp8-moe-kernels-mi355
|
||||
dind: false
|
||||
@@ -720,9 +757,7 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
- examples/basic/offline_inference/chat.py
|
||||
- tests/kernels/attention/test_attention_selector.py
|
||||
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
commands:
|
||||
- amd-smi
|
||||
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
|
||||
- pytest -v -s tests/kernels/attention/test_attention_selector.py
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
|
||||
@@ -66,6 +66,28 @@ steps:
|
||||
autorun_on_main: true
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) LM Eval Watermarking"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 45
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- .buildkite/test_areas/lm_eval.yaml
|
||||
- tests/evals/gsm8k/
|
||||
- vllm/config/watermarking.py
|
||||
- vllm/sampling_params.py
|
||||
- vllm/v1/watermarking/
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/platforms/rocm.py
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
|
||||
|
||||
# - label: LM Eval Large Models (4xA100)
|
||||
# key: lm-eval-large-models-4xa100
|
||||
@@ -203,6 +225,33 @@ steps:
|
||||
autorun_on_main: true
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) LM Eval Watermark Feature Combination"
|
||||
dind: false
|
||||
device: mi355_2
|
||||
num_devices: 2
|
||||
timeout_in_minutes: 60
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- .buildkite/test_areas/lm_eval.yaml
|
||||
- tests/evals/gsm8k/
|
||||
- vllm/config/speculative.py
|
||||
- vllm/config/vllm.py
|
||||
- vllm/config/watermarking.py
|
||||
- vllm/sampling_params.py
|
||||
- vllm/model_executor/models/qwen3_dspark.py
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- vllm/model_executor/warmup/
|
||||
- vllm/v1/watermarking/
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/platforms/rocm.py
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
|
||||
|
||||
- label: ":nvidia: (B200) LM Eval Spec Decode"
|
||||
key: lm-eval-spec-decode-4xb200
|
||||
@@ -294,29 +343,66 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
|
||||
|
||||
- label: ":nvidia: (H200) LM Eval Large Models"
|
||||
key: lm-eval-large-models-8xh200
|
||||
timeout_in_minutes: 120
|
||||
- label: ":nvidia: (H200) LM Eval DeepSeek R1 TP"
|
||||
key: lm-eval-deepseek-r1-tp-8xh200
|
||||
timeout_in_minutes: 35
|
||||
device: h200
|
||||
optional: true
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI300) LM Eval Large Models"
|
||||
dind: false
|
||||
device: mi300_8
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
|
||||
source_file_dependencies:
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-r1-tp.txt
|
||||
|
||||
- label: ":nvidia: (H200) LM Eval DeepSeek R1 DP"
|
||||
key: lm-eval-deepseek-r1-dp-8xh200
|
||||
timeout_in_minutes: 35
|
||||
device: h200
|
||||
optional: true
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-r1-dp.txt
|
||||
|
||||
- label: ":nvidia: (H200) LM Eval DeepSeek V3.2 TP"
|
||||
key: lm-eval-deepseek-v32-tp-8xh200
|
||||
timeout_in_minutes: 35
|
||||
device: h200
|
||||
optional: true
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-tp.txt
|
||||
|
||||
- label: ":nvidia: (H200) LM Eval DeepSeek V3.2 DP"
|
||||
key: lm-eval-deepseek-v32-dp-8xh200
|
||||
timeout_in_minutes: 35
|
||||
device: h200
|
||||
optional: true
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-dp.txt
|
||||
|
||||
- label: ":nvidia: (H200) LM Eval Nemotron 3 Super"
|
||||
key: lm-eval-nemotron-3-super-8xh200
|
||||
timeout_in_minutes: 35
|
||||
device: h200
|
||||
optional: true
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-nemotron-3-super.txt
|
||||
|
||||
- label: ":amd: (MI300) LM Eval Large Models"
|
||||
key: amd-lm-eval-large-models
|
||||
dind: false
|
||||
device: mi300_8
|
||||
optional: true
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
|
||||
source_file_dependencies:
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: ":nvidia: (H100) MoE Refactor Integration TEMPORARY"
|
||||
key: moe-refactor-integration-test-h100-temporary
|
||||
@@ -329,7 +415,27 @@ steps:
|
||||
# build 83851): shards ~12.4/13.6/12.4/14.5m of eval time. New configs must
|
||||
# be added to exactly one config-h100-shard-*.txt.
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt
|
||||
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) MoE Refactor Integration Shard %N"
|
||||
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:
|
||||
- vllm/platforms/rocm.py
|
||||
- tests/evals/gsm8k/
|
||||
- vllm/model_executor/layers/fused_moe/
|
||||
- vllm/model_executor/layers/quantization/
|
||||
- vllm/_aiter_ops.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-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt
|
||||
|
||||
- label: ":nvidia: (B200) MoE Refactor Integration TEMPORARY Shard %N"
|
||||
key: moe-refactor-integration-test-b200-temporary
|
||||
device: b200-k8s
|
||||
@@ -347,47 +453,67 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
|
||||
|
||||
- label: ":nvidia: (A100) LM Eval Humming FP16 TEMPORARY Shard %N"
|
||||
key: lm-eval-humming-f16-a100
|
||||
- label: ":nvidia: (L4) LM Eval Humming TEMPORARY"
|
||||
key: lm-eval-humming-l4
|
||||
optional: true
|
||||
timeout_in_minutes: 90
|
||||
device: l4
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- .buildkite/test_areas/lm_eval.yaml
|
||||
- tests/evals/gsm8k/configs/humming/
|
||||
- requirements/cuda.txt
|
||||
- vllm/utils/humming.py
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt
|
||||
|
||||
- label: ":nvidia: (A100) LM Eval Humming Default TEMPORARY Shard %N"
|
||||
key: lm-eval-humming-a100-default
|
||||
optional: true
|
||||
timeout_in_minutes: 75
|
||||
device: a100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
parallelism: 3
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt
|
||||
|
||||
- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY"
|
||||
key: lm-eval-humming-act-a100
|
||||
- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY Shard %N"
|
||||
key: lm-eval-humming-a100-act
|
||||
optional: true
|
||||
timeout_in_minutes: 45
|
||||
device: a100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
parallelism: 2
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt
|
||||
|
||||
- label: ":nvidia: (H100) LM Eval Humming FP16 TEMPORARY Shard %N"
|
||||
key: lm-eval-humming-f16-h100
|
||||
- label: ":nvidia: (H100) LM Eval Humming Default TEMPORARY Shard %N"
|
||||
key: lm-eval-humming-h100-default
|
||||
optional: true
|
||||
timeout_in_minutes: 70
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
parallelism: 3
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
@@ -395,14 +521,14 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt
|
||||
|
||||
- label: ":nvidia: (H100) LM Eval Humming Activation FP8/INT8 TEMPORARY"
|
||||
key: lm-eval-humming-act-h100
|
||||
key: lm-eval-humming-h100-act
|
||||
optional: true
|
||||
timeout_in_minutes: 70
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
@@ -410,15 +536,35 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
|
||||
# - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
|
||||
|
||||
- label: ":nvidia: (B200) LM Eval Humming FP16 TEMPORARY"
|
||||
key: lm-eval-humming-f16-b200
|
||||
- label: ":nvidia: (H100) LM Eval Humming DeepEP TEMPORARY"
|
||||
key: lm-eval-humming-h100-deepep
|
||||
optional: true
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- .buildkite/test_areas/lm_eval.yaml
|
||||
- tests/evals/gsm8k/configs/humming/
|
||||
- requirements/cuda.txt
|
||||
- vllm/utils/humming.py
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
- vllm/distributed/
|
||||
- csrc/moe/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-deepep.txt
|
||||
|
||||
- label: ":nvidia: (B200) LM Eval Humming Default TEMPORARY"
|
||||
key: lm-eval-humming-b200-default
|
||||
optional: true
|
||||
timeout_in_minutes: 50
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
@@ -426,14 +572,14 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
|
||||
|
||||
- label: ":nvidia: (B200) LM Eval Humming Activation FP8/INT8 TEMPORARY"
|
||||
key: lm-eval-humming-act-b200
|
||||
key: lm-eval-humming-b200-act
|
||||
optional: true
|
||||
timeout_in_minutes: 50
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming/
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
@@ -538,6 +684,29 @@ steps:
|
||||
commands:
|
||||
- uv pip install --system 'evalscope==1.10.0'
|
||||
- pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
|
||||
dind: false
|
||||
device: mi355_4
|
||||
num_devices: 4
|
||||
timeout_in_minutes: 105
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/libtorch_stable/gdn/
|
||||
- tests/evals/qwen4_exp/
|
||||
- vllm/model_executor/layers/mamba/
|
||||
- vllm/models/qwen4_exp/
|
||||
- vllm/transformers_utils/configs/qwen4_exp.py
|
||||
- vllm/v1/attention/backends/short_conv_attn.py
|
||||
- vllm/platforms/rocm.py
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- uv pip install --system 'evalscope==1.10.0'
|
||||
- pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt
|
||||
|
||||
- label: ":nvidia: (DGX) Spark GPQA Eval (GPT-OSS)"
|
||||
key: gpqa-eval-gpt-oss-spark
|
||||
|
||||
@@ -237,6 +237,27 @@ steps:
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Extract Hidden States Integration Single GPU"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 45
|
||||
working_dir: /vllm-workspace
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/extract_hidden_states.py
|
||||
- vllm/model_executor/models/extract_hidden_states.py
|
||||
- vllm/transformers_utils/configs/extract_hidden_states.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration
|
||||
- vllm/platforms/rocm.py
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
|
||||
commands:
|
||||
- pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: ":nvidia: (L4) Extract Hidden States Integration"
|
||||
key: extract-hidden-states-integration-2-gpus
|
||||
@@ -416,7 +437,6 @@ steps:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
soft_fail: true
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -505,6 +525,9 @@ steps:
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
- tests/reasoning
|
||||
- tools/pre_commit/check_test_tethering.py
|
||||
- tools/pre_commit/test_tethering_allowlist.txt
|
||||
- tests/tools/test_check_test_tethering.py
|
||||
- tests/tool_parsers
|
||||
- tests/tokenizers_
|
||||
- tests/parser
|
||||
@@ -532,6 +555,7 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s config
|
||||
- pytest -v -s tools/test_check_test_tethering.py
|
||||
|
||||
- label: ":computer: (CPU) Reasoning + Renderers"
|
||||
key: cpu-reasoning-renderers
|
||||
|
||||
@@ -119,6 +119,17 @@ steps:
|
||||
- tests/models/glm5next/
|
||||
commands:
|
||||
- pytest -v -s models/glm5next
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) GLM5Next"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
timeout_in_minutes: 35
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: ":nvidia: (H200) Qwen4 Exp"
|
||||
key: qwen4-exp-unit-tests
|
||||
|
||||
@@ -91,7 +91,7 @@ steps:
|
||||
commands:
|
||||
- MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
|
||||
# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
|
||||
@@ -111,6 +111,27 @@ steps:
|
||||
- rm -rf /tmp/mamba-src && git clone --depth 1 --branch v2.3.0 https://github.com/state-spaces/mamba /tmp/mamba-src && sed -i 's/-std=c++17/-std=c++20/g' /tmp/mamba-src/setup.py && MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation /tmp/mamba-src
|
||||
- CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Granite Language Model Compatibility"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 80
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- "!vllm/distributed/kv_transfer/"
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- export GIT_CONFIG_COUNT=1
|
||||
- export GIT_CONFIG_KEY_0=http.version
|
||||
- export GIT_CONFIG_VALUE_0=HTTP/1.1
|
||||
- MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
|
||||
|
||||
- label: ":nvidia: (H200 MIG 35GB) Language Models (Extended Generation)"
|
||||
device: h200_35gb
|
||||
|
||||
@@ -87,6 +87,20 @@ steps:
|
||||
- tests/compile/h100/
|
||||
commands:
|
||||
- "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) PyTorch Compilation H100 Cases"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- tests/utils.py
|
||||
- csrc/
|
||||
commands:
|
||||
- pytest -v -s compile/h100/test_startup.py
|
||||
|
||||
- label: ":nvidia: (L4) PyTorch Compilation Passes"
|
||||
key: pytorch-compilation-passes-unit-tests
|
||||
|
||||
@@ -35,7 +35,7 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s samplers
|
||||
- pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
|
||||
|
||||
# Qwen2-Audio-7B leaves no room for KV cache on a 16 GiB slice.
|
||||
- label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search"
|
||||
@@ -51,3 +51,24 @@ steps:
|
||||
commands:
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Samplers Multimodal Beam Search"
|
||||
dind: false
|
||||
device: mi355_1
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 45
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers
|
||||
- vllm/sampling_metadata.py
|
||||
- vllm/v1/sample/
|
||||
- vllm/entrypoints/generate/beam_search/
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
|
||||
|
||||
@@ -184,6 +184,14 @@ repos:
|
||||
files: ^\.buildkite/.*\.ya?ml$
|
||||
pass_filenames: false
|
||||
additional_dependencies: [PyYAML]
|
||||
- id: check-test-tethering
|
||||
name: Check new tests are wired into Buildkite CI
|
||||
language: python
|
||||
entry: python tools/pre_commit/check_test_tethering.py
|
||||
files: ^(tests/.*\.py|\.buildkite/(test_areas|hardware_tests|intel_jobs)/.*\.yaml|\.buildkite/(test-amd|ci_config[^/]*)\.yaml|tools/pre_commit/test_tethering_allowlist\.txt)$
|
||||
additional_dependencies: [PyYAML, regex]
|
||||
verbose: true # full-scan `note:` lines are advisory, so the hook passes
|
||||
require_serial: true # one full scan, not one per pre-commit file batch
|
||||
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.10
|
||||
entry: python tools/pre_commit/mypy.py "3.10"
|
||||
|
||||
@@ -460,6 +460,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
|
||||
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/all_reduce_mhc.cu")
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
|
||||
"9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
|
||||
|
||||
@@ -0,0 +1,407 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Benchmark the SM100 GEMM-RS/AR kernel against the model's unfused path.
|
||||
|
||||
The baseline is what a model runs without the fusion: the projection's own
|
||||
linear kernel followed by vLLM's collective (``sp_reduce_scatter`` for RS,
|
||||
``tensor_model_parallel_all_reduce`` for AR), with the collective backend
|
||||
chosen by vLLM's dispatch exactly as in serving. The same two kernels are also
|
||||
timed bare (the GEMM alone, the collective alone on a precomputed partial);
|
||||
their max is the floor a perfectly overlapped fused kernel could reach.
|
||||
|
||||
All ranks must belong to one NVLink domain. Run a TP8 sweep over the Kimi-K3
|
||||
projections with:
|
||||
|
||||
torchrun --nproc-per-node=8 -- benchmarks/kernels/benchmark_gemm_rs_ar.py
|
||||
|
||||
or the DeepSeek-V4.1 MXFP8 ``wo_b`` projection at TP4 with:
|
||||
|
||||
torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
|
||||
--model deepseek_v41 --backend flashinfer_cutedsl
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed import (
|
||||
cleanup_dist_env_and_memory,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from vllm.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
|
||||
from vllm.model_executor.layers.linear import RowParallelLinear
|
||||
from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
|
||||
|
||||
# Per model: (hidden size N, global input widths K). Kimi-K3 lists the
|
||||
# shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the
|
||||
# ``wo_b`` output projection (o_groups * o_lora_rank).
|
||||
_MODEL_PROJECTIONS = {
|
||||
"kimi_k3": (7168, (6144, 12288)),
|
||||
"deepseek_v41": (5120, (8192,)),
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("rs", "ar"),
|
||||
default="rs",
|
||||
help="Collective mode to benchmark.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[128, 512, 2048, 8192, 32768],
|
||||
help="Global token counts to benchmark.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
choices=tuple(_MODEL_PROJECTIONS),
|
||||
default="kimi_k3",
|
||||
help="Model whose projection shapes set the default --n and --k.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
choices=("bf16", "flashinfer_cutlass", "flashinfer_cutedsl"),
|
||||
default="bf16",
|
||||
help="Weight dtype/layout: BF16 or an online-MXFP8 FlashInfer kernel.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--k",
|
||||
type=int,
|
||||
nargs="+",
|
||||
help=(
|
||||
"Per-rank input dimensions. By default, derive the selected "
|
||||
"model's projection dimensions from the TP world size."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n",
|
||||
type=int,
|
||||
help="Output width. Defaults to the selected model's hidden size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-workspaces",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Pointer-distinct inputs and CUDA graphs to rotate.",
|
||||
)
|
||||
parser.add_argument("--warmup-replays", type=int, default=5)
|
||||
parser.add_argument("--samples", type=int, default=20)
|
||||
parser.add_argument("--json", type=str, help="Also write the results here.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def make_projection(
|
||||
gemm_rs_ar: GemmRsAr,
|
||||
weight: torch.Tensor,
|
||||
world_size: int,
|
||||
backend: str,
|
||||
) -> RowParallelLinear:
|
||||
"""Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS.
|
||||
|
||||
For the MXFP8 backends this quantizes online through the FlashInfer
|
||||
kernel's post-load step, so the fused kernel and the unfused baseline see
|
||||
the production weight and scale layouts.
|
||||
"""
|
||||
N, K = weight.shape
|
||||
quant_config = None
|
||||
if backend != "bf16":
|
||||
from vllm.config.quantization import QuantizationConfigArgs
|
||||
from vllm.model_executor.layers.quantization.online.base import (
|
||||
OnlineQuantizationConfig,
|
||||
)
|
||||
|
||||
quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
|
||||
with torch.device(weight.device):
|
||||
linear = RowParallelLinear(
|
||||
K * world_size,
|
||||
N,
|
||||
bias=False,
|
||||
params_dtype=weight.dtype,
|
||||
quant_config=quant_config,
|
||||
# The baseline reduces the local partial itself, like the models.
|
||||
reduce_results=False,
|
||||
return_bias=False,
|
||||
)
|
||||
# Select the fused path before online quantization replaces the weights.
|
||||
assert gemm_rs_ar.can_run(linear)
|
||||
linear.weight = torch.nn.Parameter(weight, requires_grad=False)
|
||||
linear.quant_method.process_weights_after_loading(linear)
|
||||
return linear
|
||||
|
||||
|
||||
def capture_graph(
|
||||
op: Callable[[], torch.Tensor],
|
||||
stream: torch.cuda.Stream,
|
||||
cpu_group: dist.ProcessGroup,
|
||||
) -> torch.cuda.CUDAGraph:
|
||||
stream.wait_stream(torch.cuda.current_stream())
|
||||
dist.barrier(group=cpu_group)
|
||||
with torch.cuda.stream(stream):
|
||||
for _ in range(3):
|
||||
op()
|
||||
stream.synchronize()
|
||||
dist.barrier(group=cpu_group)
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=stream):
|
||||
op()
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
dist.barrier(group=cpu_group)
|
||||
return graph
|
||||
|
||||
|
||||
def benchmark_graphs(
|
||||
candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
|
||||
warmup_replays: int,
|
||||
samples: int,
|
||||
device_group: dist.ProcessGroup,
|
||||
device_barrier: Callable[[], None],
|
||||
) -> dict[str, float]:
|
||||
"""Median over samples of the slowest rank's replay time, in us."""
|
||||
names = list(candidate_graphs)
|
||||
for round_index in range(warmup_replays):
|
||||
for offset in range(len(names)):
|
||||
name = names[(round_index + offset) % len(names)]
|
||||
graphs = candidate_graphs[name]
|
||||
device_barrier()
|
||||
graphs[round_index % len(graphs)].replay()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
timings: dict[str, list[float]] = {name: [] for name in names}
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
for sample_index in range(samples):
|
||||
for offset in range(len(names)):
|
||||
name = names[(sample_index + offset) % len(names)]
|
||||
graphs = candidate_graphs[name]
|
||||
device_barrier()
|
||||
start.record()
|
||||
graphs[sample_index % len(graphs)].replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
|
||||
elapsed = torch.tensor(
|
||||
start.elapsed_time(end) * 1000,
|
||||
dtype=torch.float64,
|
||||
device=torch.accelerator.current_device_index(),
|
||||
)
|
||||
dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
|
||||
timings[name].append(elapsed.item())
|
||||
return {name: statistics.median(values) for name, values in timings.items()}
|
||||
|
||||
|
||||
def benchmark_shape(
|
||||
gemm_rs_ar: GemmRsAr,
|
||||
mode: str,
|
||||
backend: str,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
num_workspaces: int,
|
||||
warmup_replays: int,
|
||||
samples: int,
|
||||
device_group: dist.ProcessGroup,
|
||||
cpu_group: dist.ProcessGroup,
|
||||
device_barrier: Callable[[], None],
|
||||
) -> dict[str, float | int | str]:
|
||||
all_reduce = mode == "ar"
|
||||
world_size = dist.get_world_size(device_group)
|
||||
rank = dist.get_rank(device_group)
|
||||
device = torch.device("cuda", torch.accelerator.current_device_index())
|
||||
local_M = (M + world_size - 1) // world_size
|
||||
|
||||
rng = torch.Generator(device=device)
|
||||
rng.manual_seed(1000 + rank * 10 + M + K)
|
||||
inputs = [
|
||||
torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
|
||||
for _ in range(num_workspaces)
|
||||
]
|
||||
projections = [
|
||||
make_projection(
|
||||
gemm_rs_ar,
|
||||
torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng),
|
||||
world_size,
|
||||
backend,
|
||||
)
|
||||
for _ in range(num_workspaces)
|
||||
]
|
||||
|
||||
def collective(partial: torch.Tensor) -> torch.Tensor:
|
||||
if all_reduce:
|
||||
return tensor_model_parallel_all_reduce(partial)
|
||||
return sp_reduce_scatter(partial)
|
||||
|
||||
def gemm(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
|
||||
return linear(x)
|
||||
|
||||
def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
|
||||
return collective(linear(x))
|
||||
|
||||
def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
|
||||
return gemm_rs_ar.apply(x, linear)
|
||||
|
||||
# The bare collective runs on the GEMM's own output so it sees the same
|
||||
# message size and layout the unfused path hands to the collective.
|
||||
partials = [linear(x) for x, linear in zip(inputs, projections)]
|
||||
|
||||
expected = unfused(inputs[0], projections[0])
|
||||
actual = fused(inputs[0], projections[0])
|
||||
torch.accelerator.synchronize(device)
|
||||
rows = M if all_reduce else min(max(M - rank * local_M, 0), local_M)
|
||||
torch.testing.assert_close(
|
||||
actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, msg=f"M={M}, K={K}"
|
||||
)
|
||||
|
||||
candidates: dict[str, list[Callable[[], torch.Tensor]]] = {
|
||||
"gemm": [lambda x=x, w=w: gemm(x, w) for x, w in zip(inputs, projections)],
|
||||
"collective": [lambda p=p: collective(p) for p in partials],
|
||||
"vllm": [lambda x=x, w=w: unfused(x, w) for x, w in zip(inputs, projections)],
|
||||
"gemm_rs_ar": [
|
||||
lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections)
|
||||
],
|
||||
}
|
||||
|
||||
candidate_graphs = {}
|
||||
keepalive: list[object] = []
|
||||
for name, ops in candidates.items():
|
||||
stream = torch.cuda.Stream()
|
||||
candidate_graphs[name] = [capture_graph(op, stream, cpu_group) for op in ops]
|
||||
keepalive.append(stream)
|
||||
|
||||
times = benchmark_graphs(
|
||||
candidate_graphs, warmup_replays, samples, device_group, device_barrier
|
||||
)
|
||||
floor = max(times["gemm"], times["collective"])
|
||||
return {
|
||||
"mode": mode.upper(),
|
||||
"backend": backend,
|
||||
"M": M,
|
||||
"N": N,
|
||||
"K": K,
|
||||
"gemm_us": times["gemm"],
|
||||
"collective_us": times["collective"],
|
||||
"floor_us": floor,
|
||||
"vllm_us": times["vllm"],
|
||||
"gemm_rs_ar_us": times["gemm_rs_ar"],
|
||||
"latency_change": (times["gemm_rs_ar"] - times["vllm"]) / times["vllm"],
|
||||
"floor_gap": (times["gemm_rs_ar"] - floor) / floor,
|
||||
}
|
||||
|
||||
|
||||
def print_results(results: list[dict[str, float | int | str]]) -> None:
|
||||
collective = results[0]["mode"]
|
||||
backend = results[0]["backend"]
|
||||
print(f"### GEMM-{collective} vs vLLM unfused path ({backend} weights)")
|
||||
print(
|
||||
f"| M | N | K | GEMM (us) | {collective} (us) | floor (us) | vLLM (us) | "
|
||||
f"GEMM-{collective} (us) | vs vLLM | vs floor |"
|
||||
)
|
||||
print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|")
|
||||
for r in results:
|
||||
print(
|
||||
f"| {r['M']} | {r['N']} | {r['K']} | {r['gemm_us']:.1f} | "
|
||||
f"{r['collective_us']:.1f} | {r['floor_us']:.1f} | {r['vllm_us']:.1f} | "
|
||||
f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} | "
|
||||
f"{r['floor_gap']:+.1%} |"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
default_n, default_k = _MODEL_PROJECTIONS[args.model]
|
||||
if args.n is None:
|
||||
args.n = default_n
|
||||
assert args.m and min(args.m) > 0
|
||||
assert args.n % 128 == 0
|
||||
assert args.num_workspaces > 0
|
||||
assert args.warmup_replays >= 0
|
||||
assert args.samples > 0
|
||||
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.accelerator.set_device_index(local_rank)
|
||||
init_distributed_environment()
|
||||
world_size = dist.get_world_size()
|
||||
if args.k is None:
|
||||
assert all(K % world_size == 0 for K in default_k)
|
||||
K_values = [K // world_size for K in default_k]
|
||||
else:
|
||||
K_values = args.k
|
||||
k_alignment = 64 if args.backend == "bf16" else 128
|
||||
assert all(K % k_alignment == 0 for K in K_values)
|
||||
|
||||
config = VllmConfig()
|
||||
config.model_config = ModelConfig(dtype="bfloat16")
|
||||
# Pin the online-MXFP8 kernel the projections quantize through.
|
||||
config.kernel_config.linear_backend = (
|
||||
"auto" if args.backend == "bf16" else args.backend
|
||||
)
|
||||
with set_current_vllm_config(config):
|
||||
# Builds the TP communicator as in serving, so the baseline's
|
||||
# collective is whatever vLLM's dispatch selects (FlashInfer all-reduce
|
||||
# first, then NCCL symmetric memory, custom all-reduce, pynccl).
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
tp_group = get_tp_group()
|
||||
pynccl_comm = tp_group.device_communicator.pynccl_comm
|
||||
assert pynccl_comm is not None
|
||||
sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
|
||||
sync_output = torch.empty_like(sync_input)
|
||||
|
||||
def device_barrier() -> None:
|
||||
# Order the timed launch after a device-side rank rendezvous without
|
||||
# including the rendezvous itself in the measured event interval.
|
||||
pynccl_comm.all_reduce(sync_input, sync_output)
|
||||
|
||||
gemm_rs_ar = GemmRsAr(
|
||||
max_M=max(args.m),
|
||||
N=args.n,
|
||||
all_reduce=args.mode == "ar",
|
||||
)
|
||||
with set_current_vllm_config(config):
|
||||
results = [
|
||||
benchmark_shape(
|
||||
gemm_rs_ar,
|
||||
args.mode,
|
||||
args.backend,
|
||||
M,
|
||||
args.n,
|
||||
K,
|
||||
args.num_workspaces,
|
||||
args.warmup_replays,
|
||||
args.samples,
|
||||
tp_group.device_group,
|
||||
tp_group.cpu_group,
|
||||
device_barrier,
|
||||
)
|
||||
for K in K_values
|
||||
for M in args.m
|
||||
]
|
||||
del gemm_rs_ar
|
||||
|
||||
if tp_group.rank_in_group == 0:
|
||||
print_results(results)
|
||||
if args.json:
|
||||
with open(args.json, "w") as f:
|
||||
json.dump(results, f, indent=2)
|
||||
|
||||
dist.barrier(group=tp_group.cpu_group)
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,469 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Benchmark the SM100 Kimi-K3 GEMM-RS/AR kernel.
|
||||
|
||||
All ranks must belong to one NVLink domain. For example, run a TP8 sweep with:
|
||||
|
||||
torchrun --nproc-per-node=8 \
|
||||
benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import statistics
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed._symmetric_memory as symm_mem
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr
|
||||
|
||||
# Shared-expert down-proj and attention O-proj.
|
||||
_KIMI_K3_PROJECTION_K = (6144, 12288)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Candidate:
|
||||
name: str
|
||||
runs: list[Callable[[], torch.Tensor]]
|
||||
check_correctness: bool = True
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("rs", "ar"),
|
||||
default="rs",
|
||||
help="Collective mode to benchmark.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[128, 512, 2048, 8192, 32768],
|
||||
help="Global token counts to benchmark.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--k",
|
||||
type=int,
|
||||
nargs="+",
|
||||
help=(
|
||||
"Per-rank input dimensions. By default, derive the Kimi-K3 "
|
||||
"shared-expert down-proj and O-proj dimensions from the TP "
|
||||
"world size."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--n", type=int, default=7168)
|
||||
parser.add_argument(
|
||||
"--num-workspaces",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Pointer-distinct inputs and CUDA graphs to rotate.",
|
||||
)
|
||||
parser.add_argument("--warmup-replays", type=int, default=5)
|
||||
parser.add_argument("--samples", type=int, default=20)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def capture_graph(
|
||||
op: Callable[[], torch.Tensor],
|
||||
stream: torch.cuda.Stream,
|
||||
cpu_group: dist.ProcessGroup,
|
||||
) -> tuple[torch.cuda.CUDAGraph, list[torch.Tensor | None]]:
|
||||
result: list[torch.Tensor | None] = [None]
|
||||
stream.wait_stream(torch.cuda.current_stream())
|
||||
dist.barrier(group=cpu_group)
|
||||
with torch.cuda.stream(stream):
|
||||
for _ in range(3):
|
||||
result[0] = op()
|
||||
stream.synchronize()
|
||||
dist.barrier(group=cpu_group)
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=stream):
|
||||
result[0] = op()
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
dist.barrier(group=cpu_group)
|
||||
return graph, result
|
||||
|
||||
|
||||
def benchmark_graphs(
|
||||
candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
|
||||
warmup_replays: int,
|
||||
samples: int,
|
||||
device_group: dist.ProcessGroup,
|
||||
device_barrier: Callable[[], None],
|
||||
) -> dict[str, float]:
|
||||
candidate_names = list(candidate_graphs)
|
||||
for round_index in range(warmup_replays):
|
||||
for candidate_index in range(len(candidate_names)):
|
||||
candidate_id = (round_index + candidate_index) % len(candidate_names)
|
||||
name = candidate_names[candidate_id]
|
||||
graphs = candidate_graphs[name]
|
||||
device_barrier()
|
||||
graphs[round_index % len(graphs)].replay()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
timings: dict[str, list[float]] = {name: [] for name in candidate_names}
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
for sample_index in range(samples):
|
||||
for candidate_index in range(len(candidate_names)):
|
||||
candidate_id = (sample_index + candidate_index) % len(candidate_names)
|
||||
name = candidate_names[candidate_id]
|
||||
graphs = candidate_graphs[name]
|
||||
device_barrier()
|
||||
start.record()
|
||||
graphs[sample_index % len(graphs)].replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
|
||||
elapsed = torch.tensor(
|
||||
start.elapsed_time(end) * 1000,
|
||||
dtype=torch.float64,
|
||||
device=torch.accelerator.current_device_index(),
|
||||
)
|
||||
dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
|
||||
timings[name].append(elapsed.item())
|
||||
return {name: statistics.median(values) for name, values in timings.items()}
|
||||
|
||||
|
||||
def valid_rows(M: int, local_M: int, rank: int) -> int:
|
||||
return min(max(M - rank * local_M, 0), local_M)
|
||||
|
||||
|
||||
def benchmark_shape(
|
||||
gemm_rs_ar: GemmRsAr,
|
||||
mode: str,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
num_workspaces: int,
|
||||
warmup_replays: int,
|
||||
samples: int,
|
||||
device_group: dist.ProcessGroup,
|
||||
cpu_group: dist.ProcessGroup,
|
||||
device_barrier: Callable[[], None],
|
||||
) -> dict[str, float | int | str]:
|
||||
all_reduce = mode == "ar"
|
||||
world_size = dist.get_world_size(device_group)
|
||||
rank = dist.get_rank(device_group)
|
||||
device = torch.device("cuda", torch.accelerator.current_device_index())
|
||||
padded_M = (M + world_size - 1) // world_size * world_size
|
||||
local_M = padded_M // world_size
|
||||
|
||||
rng = torch.Generator(device=device)
|
||||
rng.manual_seed(1000 + rank * 10 + M + K)
|
||||
inputs = [
|
||||
torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
|
||||
for _ in range(num_workspaces)
|
||||
]
|
||||
weights = [
|
||||
torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng)
|
||||
for _ in range(num_workspaces)
|
||||
]
|
||||
|
||||
partial = torch.empty((padded_M, N), dtype=torch.bfloat16, device=device)
|
||||
symm_partial = symm_mem.empty((padded_M, N), dtype=torch.bfloat16, device=device)
|
||||
symm_partial_handle = symm_mem.rendezvous(symm_partial, device_group)
|
||||
collective_inputs = [torch.empty_like(partial) for _ in range(num_workspaces)]
|
||||
symm_collective_inputs = []
|
||||
symm_collective_handles = []
|
||||
for _ in range(num_workspaces):
|
||||
collective_input = symm_mem.empty(
|
||||
(padded_M, N),
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
symm_collective_inputs.append(collective_input)
|
||||
symm_collective_handles.append(
|
||||
symm_mem.rendezvous(collective_input, device_group)
|
||||
)
|
||||
torch_output = torch.empty((local_M, N), dtype=torch.bfloat16, device=device)
|
||||
symm_output = torch.empty_like(torch_output)
|
||||
gemm_output = torch.empty((M, N), dtype=torch.bfloat16, device=device)
|
||||
|
||||
if padded_M > M:
|
||||
partial[M:].zero_()
|
||||
symm_partial[M:].zero_()
|
||||
|
||||
def make_torch_ring_ll_gemm_collective(
|
||||
x: torch.Tensor, weight: torch.Tensor
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
torch.mm(x, weight.T, out=partial[:M])
|
||||
if all_reduce:
|
||||
dist.all_reduce(partial, group=device_group)
|
||||
return partial[:M]
|
||||
dist.reduce_scatter_single(torch_output, partial, group=device_group)
|
||||
return torch_output
|
||||
|
||||
return run
|
||||
|
||||
def make_torch_ldmc_gemm_collective(
|
||||
x: torch.Tensor, weight: torch.Tensor
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
torch.mm(x, weight.T, out=symm_partial[:M])
|
||||
if all_reduce:
|
||||
dist.all_reduce(symm_partial, group=device_group)
|
||||
return symm_partial[:M]
|
||||
dist.reduce_scatter_single(
|
||||
symm_output,
|
||||
symm_partial,
|
||||
group=device_group,
|
||||
)
|
||||
return symm_output
|
||||
|
||||
return run
|
||||
|
||||
def make_fused_gemm_collective(
|
||||
x: torch.Tensor, weight: torch.Tensor
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
return gemm_rs_ar(x, weight)
|
||||
|
||||
return run
|
||||
|
||||
def make_torch_gemm(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
return torch.mm(x, weight.T, out=gemm_output)
|
||||
|
||||
return run
|
||||
|
||||
def make_ring_ll_collective(
|
||||
collective_input: torch.Tensor,
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
if all_reduce:
|
||||
dist.all_reduce(collective_input, group=device_group)
|
||||
return collective_input
|
||||
dist.reduce_scatter_single(
|
||||
torch_output, collective_input, group=device_group
|
||||
)
|
||||
return torch_output
|
||||
|
||||
return run
|
||||
|
||||
def make_ldmc_collective(
|
||||
collective_input: torch.Tensor,
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
if all_reduce:
|
||||
dist.all_reduce(collective_input, group=device_group)
|
||||
return collective_input
|
||||
dist.reduce_scatter_single(
|
||||
symm_output, collective_input, group=device_group
|
||||
)
|
||||
return symm_output
|
||||
|
||||
return run
|
||||
|
||||
candidates = (
|
||||
Candidate(
|
||||
"ring_ll_us",
|
||||
[make_torch_ring_ll_gemm_collective(x, w) for x, w in zip(inputs, weights)],
|
||||
),
|
||||
Candidate(
|
||||
"ldmc_us",
|
||||
[make_torch_ldmc_gemm_collective(x, w) for x, w in zip(inputs, weights)],
|
||||
),
|
||||
Candidate(
|
||||
"gemm_rs_ar_us",
|
||||
[make_fused_gemm_collective(x, w) for x, w in zip(inputs, weights)],
|
||||
),
|
||||
Candidate(
|
||||
"torch_gemm_us",
|
||||
[make_torch_gemm(x, w) for x, w in zip(inputs, weights)],
|
||||
check_correctness=False,
|
||||
),
|
||||
Candidate(
|
||||
"ring_ll_collective_us",
|
||||
[make_ring_ll_collective(x) for x in collective_inputs],
|
||||
check_correctness=False,
|
||||
),
|
||||
Candidate(
|
||||
"ldmc_collective_us",
|
||||
[make_ldmc_collective(x) for x in symm_collective_inputs],
|
||||
check_correctness=False,
|
||||
),
|
||||
)
|
||||
|
||||
expected = candidates[0].runs[0]()
|
||||
rows = M if all_reduce else valid_rows(M, local_M, rank)
|
||||
for candidate in candidates[1:]:
|
||||
if not candidate.check_correctness:
|
||||
continue
|
||||
actual = candidate.runs[0]()
|
||||
torch.accelerator.synchronize(device)
|
||||
torch.testing.assert_close(
|
||||
actual[:rows],
|
||||
expected[:rows],
|
||||
rtol=5e-2,
|
||||
atol=4.0,
|
||||
)
|
||||
|
||||
candidate_graphs = {}
|
||||
graph_keepalive: list[object] = [symm_partial_handle, *symm_collective_handles]
|
||||
for candidate in candidates:
|
||||
stream = torch.cuda.Stream()
|
||||
bundles = [capture_graph(run, stream, cpu_group) for run in candidate.runs]
|
||||
candidate_graphs[candidate.name] = [graph for graph, _ in bundles]
|
||||
graph_keepalive.extend(bundles)
|
||||
graph_keepalive.append(stream)
|
||||
|
||||
times = benchmark_graphs(
|
||||
candidate_graphs,
|
||||
warmup_replays,
|
||||
samples,
|
||||
device_group,
|
||||
device_barrier,
|
||||
)
|
||||
|
||||
best_nccl_us = min(times["ring_ll_collective_us"], times["ldmc_collective_us"])
|
||||
return {
|
||||
"mode": mode.upper(),
|
||||
"M": M,
|
||||
"N": N,
|
||||
"K": K,
|
||||
**times,
|
||||
"best_nccl_us": best_nccl_us,
|
||||
"speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_ar_us"],
|
||||
"speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_ar_us"],
|
||||
}
|
||||
|
||||
|
||||
def print_results(results: list[dict[str, float | int | str]]) -> None:
|
||||
results_df = pd.DataFrame(results)
|
||||
collective = str(results_df["mode"].iloc[0])
|
||||
end_to_end = (
|
||||
results_df[
|
||||
[
|
||||
"M",
|
||||
"N",
|
||||
"K",
|
||||
"ring_ll_us",
|
||||
"ldmc_us",
|
||||
"gemm_rs_ar_us",
|
||||
"speedup_vs_ring_ll",
|
||||
"speedup_vs_ldmc",
|
||||
]
|
||||
]
|
||||
.rename(
|
||||
columns={
|
||||
"ring_ll_us": f"Torch GEMM + NCCL {collective} (RING_LL) (us)",
|
||||
"ldmc_us": f"Torch GEMM + NCCL {collective} (LDMC) (us)",
|
||||
"gemm_rs_ar_us": f"GEMM-{collective} (us)",
|
||||
"speedup_vs_ring_ll": "Speedup vs RING_LL",
|
||||
"speedup_vs_ldmc": "Speedup vs LDMC",
|
||||
}
|
||||
)
|
||||
.round(3)
|
||||
)
|
||||
components = (
|
||||
results_df[["M", "N", "K", "torch_gemm_us", "best_nccl_us", "gemm_rs_ar_us"]]
|
||||
.rename(
|
||||
columns={
|
||||
"torch_gemm_us": "Torch GEMM (us)",
|
||||
"best_nccl_us": f"NCCL {collective} (best) (us)",
|
||||
"gemm_rs_ar_us": f"GEMM-{collective} (us)",
|
||||
}
|
||||
)
|
||||
.round(2)
|
||||
)
|
||||
|
||||
print(f"### {collective} end-to-end latency")
|
||||
print(end_to_end.to_markdown(index=False))
|
||||
print(f"\n### {collective} component latency")
|
||||
print(components.to_markdown(index=False))
|
||||
print(
|
||||
f"\nNCCL {collective} (best) is the faster of RING_LL and LDMC "
|
||||
"for each shape.\n"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
assert args.m and min(args.m) > 0
|
||||
assert args.n % 256 == 0
|
||||
assert args.num_workspaces > 0
|
||||
assert args.warmup_replays >= 0
|
||||
assert args.samples > 0
|
||||
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.accelerator.set_device_index(local_rank)
|
||||
init_distributed_environment()
|
||||
world_size = dist.get_world_size()
|
||||
if args.k is None:
|
||||
assert all(K % world_size == 0 for K in _KIMI_K3_PROJECTION_K)
|
||||
K_values = [K // world_size for K in _KIMI_K3_PROJECTION_K]
|
||||
else:
|
||||
K_values = args.k
|
||||
assert all(K % 64 == 0 for K in K_values)
|
||||
# Reserve symmetric memory for the NCCL-managed benchmark allocations.
|
||||
os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
|
||||
# NCCL-managed symmetric allocations select the NVLS/LDMC collective path.
|
||||
symm_mem.set_backend("NCCL")
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
tp_group = get_tp_group()
|
||||
group_warmup = torch.zeros(1, device=torch.accelerator.current_device_index())
|
||||
dist.all_reduce(group_warmup, group=tp_group.device_group)
|
||||
pynccl_comm = tp_group.device_communicator.pynccl_comm
|
||||
assert pynccl_comm is not None
|
||||
sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
|
||||
sync_output = torch.empty_like(sync_input)
|
||||
|
||||
def device_barrier() -> None:
|
||||
# Order the timed launch after a device-side rank rendezvous without
|
||||
# including the rendezvous itself in the measured event interval.
|
||||
pynccl_comm.all_reduce(sync_input, sync_output)
|
||||
|
||||
gemm_rs_ar = GemmRsAr(
|
||||
max_M=max(args.m),
|
||||
N=args.n,
|
||||
all_reduce=args.mode == "ar",
|
||||
)
|
||||
results = [
|
||||
benchmark_shape(
|
||||
gemm_rs_ar,
|
||||
args.mode,
|
||||
M,
|
||||
args.n,
|
||||
K,
|
||||
args.num_workspaces,
|
||||
args.warmup_replays,
|
||||
args.samples,
|
||||
tp_group.device_group,
|
||||
tp_group.cpu_group,
|
||||
device_barrier,
|
||||
)
|
||||
for K in K_values
|
||||
for M in args.m
|
||||
]
|
||||
del gemm_rs_ar
|
||||
|
||||
if tp_group.rank_in_group == 0:
|
||||
print_results(results)
|
||||
|
||||
dist.barrier(group=tp_group.cpu_group)
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -29,10 +29,11 @@ if(DEEPGEMM_SRC_DIR)
|
||||
else()
|
||||
# Keep in sync with tools/install_deepgemm.sh
|
||||
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
|
||||
# Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120
|
||||
# Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120
|
||||
# port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU
|
||||
# alpha/beta, and SiTU for FP8/FP4 Mega MoE.
|
||||
set(_DEEPGEMM_UPSTREAM_TAG "a6bbb8000161c0dc3a85a0300a905f76898a7913")
|
||||
# alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x
|
||||
# layout header fix from vllm-project/DeepGEMM#12.
|
||||
set(_DEEPGEMM_UPSTREAM_TAG "e1f418c2a4f20818221f6b0e578b4c2f634d4c3f")
|
||||
|
||||
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _deepgemm_fc_root)
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG 506341a143fcabd4bb79052a7605ada727d6b3f5
|
||||
GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -280,26 +280,10 @@ inline int get_thread_num() {
|
||||
// balance payload across each thread
|
||||
template <typename T>
|
||||
inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
|
||||
#if 0
|
||||
// onednn partition pattern
|
||||
T& n_my = n_end;
|
||||
if (nth <= 1 || n == 0) {
|
||||
n_start = 0;
|
||||
n_my = n;
|
||||
} else {
|
||||
T n1 = div_up(n, nth);
|
||||
T n2 = n1 - 1;
|
||||
T T1 = n - n2 * nth;
|
||||
n_my = ith < T1 ? n1 : n2;
|
||||
n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
|
||||
}
|
||||
n_end += n_start;
|
||||
#else
|
||||
// pytorch aten partition pattern
|
||||
T n_my = div_up(n, nth);
|
||||
n_start = ith * n_my;
|
||||
n_end = std::min(n_start + n_my, n);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename func_t>
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include "torch_utils.h"
|
||||
#include "custom_all_gather_reduce_scatter.cuh"
|
||||
#include <cooperative_groups.h>
|
||||
#include <torch/csrc/stable/library.h>
|
||||
|
||||
namespace {
|
||||
using Tensor = torch::stable::Tensor;
|
||||
using Pack = vllm::packed_t<nv_bfloat16>::P;
|
||||
constexpr int kHidden = 5120;
|
||||
constexpr int kRanks = 4;
|
||||
constexpr int kThreads = 128;
|
||||
constexpr int kCluster = 5;
|
||||
constexpr int kPacks = kHidden / Pack::size;
|
||||
|
||||
__global__ __launch_bounds__(kThreads)
|
||||
__cluster_dims__(1, kCluster, 1) void all_reduce_mhc_kernel(
|
||||
const Pack* input, const Pack* residual, const float* post,
|
||||
const float* comb, const float* pre, const Pack* weight, Pack* output,
|
||||
Pack* normalized, Pack* local, Pack* multicast, uint32_t* epochs,
|
||||
int rank, int stage_size, float eps) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
cudaGridDependencySynchronize();
|
||||
const int token = blockIdx.x;
|
||||
const int column = blockIdx.y * kThreads + threadIdx.x;
|
||||
const int idx = token * kPacks + column;
|
||||
const int size = gridDim.x * kPacks;
|
||||
const int stage = epochs[0] % 3;
|
||||
const int dirty = vllm::mnnvl_lamport_dirty_stage(stage);
|
||||
const int dirty_size = epochs[2 + dirty];
|
||||
Pack value = input[idx];
|
||||
vllm::store_multimem_lamport_payload(
|
||||
multicast + stage * stage_size + rank * size + idx,
|
||||
vllm::sanitize_lamport_payload(value));
|
||||
// Multi-token forwards benefit from launching dependents before mHC finishes.
|
||||
if (gridDim.x > 1) cudaTriggerProgrammaticLaunchCompletion();
|
||||
vllm::lamport_cta_arrive(epochs + 1);
|
||||
for (int i = idx; i < dirty_size; i += size) {
|
||||
local[dirty * stage_size + i] = vllm::lamport_sentinel<Pack>();
|
||||
}
|
||||
Pack peers[kRanks];
|
||||
bool pending;
|
||||
do {
|
||||
vllm::LamportPack<Pack> packets[kRanks];
|
||||
#pragma unroll
|
||||
for (int peer = 0; peer < kRanks; ++peer) {
|
||||
auto ptr = local + stage * stage_size + idx + peer * size;
|
||||
asm volatile("ld.relaxed.sys.global.v4.u32 {%0,%1,%2,%3}, [%4];"
|
||||
: "=r"(packets[peer].words[0]), "=r"(packets[peer].words[1]),
|
||||
"=r"(packets[peer].words[2]), "=r"(packets[peer].words[3])
|
||||
: "l"(ptr)
|
||||
: "memory");
|
||||
}
|
||||
pending = false;
|
||||
#pragma unroll
|
||||
for (int peer = 0; peer < kRanks; ++peer) {
|
||||
pending |= vllm::is_lamport_dirty(packets[peer]);
|
||||
peers[peer] = packets[peer].packed;
|
||||
}
|
||||
} while (pending);
|
||||
auto sum = vllm::upcast(peers[0]);
|
||||
#pragma unroll
|
||||
for (int r = 1; r < kRanks; ++r) {
|
||||
vllm::packed_assign_add(sum, vllm::upcast(peers[r]));
|
||||
}
|
||||
// Keep the BF16 boundary between the collective and mHC post.
|
||||
auto reduced = vllm::upcast(vllm::downcast<Pack>(sum));
|
||||
decltype(sum) streams[4];
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 4; ++r) {
|
||||
streams[r] = vllm::upcast(residual[(token * 4 + r) * kPacks + column]);
|
||||
}
|
||||
decltype(sum) collapsed{};
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 4; ++c) {
|
||||
decltype(sum) mixed;
|
||||
#pragma unroll
|
||||
for (int e = 0; e < Pack::size; ++e) {
|
||||
// Match post*x followed by the four residual FMAs.
|
||||
float acc = __fmul_rn(post[token * 4 + c], reduced.data[e]);
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 4; ++r) {
|
||||
acc = fmaf(comb[token * 16 + r * 4 + c], streams[r].data[e], acc);
|
||||
}
|
||||
mixed.data[e] = acc;
|
||||
}
|
||||
Pack rounded = vllm::downcast<Pack>(mixed);
|
||||
output[(token * 4 + c) * kPacks + column] = rounded;
|
||||
auto as_float = vllm::upcast(rounded);
|
||||
#pragma unroll
|
||||
for (int e = 0; e < Pack::size; ++e) {
|
||||
collapsed.data[e] =
|
||||
fmaf(pre[token * 4 + c], as_float.data[e], collapsed.data[e]);
|
||||
}
|
||||
}
|
||||
// RMSNorm consumes the rounded BF16 collapse in the unfused path.
|
||||
collapsed = vllm::upcast(vllm::downcast<Pack>(collapsed));
|
||||
float sqr = 0;
|
||||
#pragma unroll
|
||||
for (int e = 0; e < Pack::size; ++e) {
|
||||
sqr = fmaf(collapsed.data[e], collapsed.data[e], sqr);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int delta = 16; delta; delta /= 2) {
|
||||
sqr += __shfl_xor_sync(0xffffffff, sqr, delta);
|
||||
}
|
||||
__shared__ float warp_sums[kCluster][kThreads / 32];
|
||||
auto cluster = cooperative_groups::this_cluster();
|
||||
const int lane = threadIdx.x % 32;
|
||||
if (lane < kCluster) {
|
||||
*cluster.map_shared_rank(&warp_sums[blockIdx.y][threadIdx.x / 32], lane) =
|
||||
sqr;
|
||||
}
|
||||
cluster.sync();
|
||||
float total = 0;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kCluster; ++c) {
|
||||
#pragma unroll
|
||||
for (int w = 0; w < kThreads / 32; ++w) total += warp_sums[c][w];
|
||||
}
|
||||
const float scale = rsqrtf(__fadd_rn(total / kHidden, eps));
|
||||
auto weights = vllm::upcast(weight[column]);
|
||||
#pragma unroll
|
||||
for (int e = 0; e < Pack::size; ++e) {
|
||||
collapsed.data[e] = collapsed.data[e] * scale * weights.data[e];
|
||||
}
|
||||
normalized[idx] = vllm::downcast<Pack>(collapsed);
|
||||
if (idx == 0) {
|
||||
while (*reinterpret_cast<volatile uint32_t*>(epochs + 1) <
|
||||
gridDim.x * gridDim.y) {
|
||||
}
|
||||
epochs[2 + stage] = kRanks * size;
|
||||
epochs[0] = vllm::mnnvl_lamport_next_stage(stage);
|
||||
epochs[1] = 0;
|
||||
}
|
||||
if (gridDim.x == 1) cudaTriggerProgrammaticLaunchCompletion();
|
||||
#else
|
||||
asm volatile("trap;");
|
||||
#endif
|
||||
}
|
||||
|
||||
void all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb,
|
||||
Tensor pre, Tensor weight, Tensor output, Tensor normalized,
|
||||
int64_t local, int64_t multicast, Tensor epochs,
|
||||
int64_t rank, int64_t stage_bytes, double eps) {
|
||||
using torch::headeronly::ScalarType;
|
||||
const torch::stable::accelerator::DeviceGuard guard(input.get_device_index());
|
||||
const int64_t n = input.size(0);
|
||||
STD_TORCH_CHECK(input.dim() == 2 && input.size(1) == kHidden && n > 0 &&
|
||||
n <= 16);
|
||||
for (const auto& t : {input, residual, weight, output, normalized}) {
|
||||
STD_TORCH_CHECK(t.is_contiguous() &&
|
||||
t.scalar_type() == ScalarType::BFloat16);
|
||||
STD_TORCH_CHECK(t.get_device_index() == input.get_device_index());
|
||||
STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(t.const_data_ptr()) % 16 == 0);
|
||||
}
|
||||
for (const auto& t : {post, comb, pre}) {
|
||||
STD_TORCH_CHECK(t.is_contiguous() && t.scalar_type() == ScalarType::Float);
|
||||
STD_TORCH_CHECK(t.get_device_index() == input.get_device_index());
|
||||
}
|
||||
STD_TORCH_CHECK(residual.numel() == n * 4 * kHidden &&
|
||||
output.numel() == residual.numel() &&
|
||||
normalized.numel() == input.numel() &&
|
||||
weight.numel() == kHidden);
|
||||
STD_TORCH_CHECK(post.numel() == n * 4 && pre.numel() == n * 4 &&
|
||||
comb.numel() == n * 16);
|
||||
STD_TORCH_CHECK(epochs.is_contiguous() &&
|
||||
epochs.scalar_type() == ScalarType::Int &&
|
||||
epochs.numel() >= 5 &&
|
||||
epochs.get_device_index() == input.get_device_index());
|
||||
STD_TORCH_CHECK(local && multicast && rank >= 0 && rank < kRanks &&
|
||||
stage_bytes % 16 == 0 &&
|
||||
stage_bytes >= n * kHidden * 2 * kRanks);
|
||||
STD_TORCH_CHECK(get_device_prop()->major == 10);
|
||||
cudaLaunchAttribute attr{};
|
||||
attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attr.val.programmaticStreamSerializationAllowed = 1;
|
||||
cudaLaunchConfig_t config{
|
||||
.gridDim = dim3(n, kCluster),
|
||||
.blockDim = dim3(kThreads),
|
||||
.dynamicSmemBytes = 0,
|
||||
.stream = get_current_cuda_stream(input.get_device_index()),
|
||||
.attrs = &attr,
|
||||
.numAttrs = 1};
|
||||
STD_CUDA_CHECK(cudaLaunchKernelEx(
|
||||
&config, all_reduce_mhc_kernel,
|
||||
static_cast<const Pack*>(input.const_data_ptr()),
|
||||
static_cast<const Pack*>(residual.const_data_ptr()),
|
||||
static_cast<const float*>(post.const_data_ptr()),
|
||||
static_cast<const float*>(comb.const_data_ptr()),
|
||||
static_cast<const float*>(pre.const_data_ptr()),
|
||||
static_cast<const Pack*>(weight.const_data_ptr()),
|
||||
static_cast<Pack*>(output.mutable_data_ptr()),
|
||||
static_cast<Pack*>(normalized.mutable_data_ptr()),
|
||||
reinterpret_cast<Pack*>(local), reinterpret_cast<Pack*>(multicast),
|
||||
static_cast<uint32_t*>(epochs.mutable_data_ptr()), static_cast<int>(rank),
|
||||
static_cast<int>(stage_bytes / sizeof(Pack)), static_cast<float>(eps)));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, mhc) {
|
||||
mhc.def(
|
||||
"all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, "
|
||||
"Tensor pre, Tensor weight, Tensor! output, Tensor! normalized, "
|
||||
"int local, int multicast, Tensor! epochs, int rank, int stage_bytes, "
|
||||
"float eps) -> ()");
|
||||
}
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, mhc) {
|
||||
mhc.impl("all_reduce_mhc", TORCH_BOX(&all_reduce_mhc));
|
||||
}
|
||||
@@ -320,10 +320,6 @@ __host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
|
||||
return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
|
||||
}
|
||||
|
||||
__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
|
||||
return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Packed type conversion and arithmetic
|
||||
// ============================================================
|
||||
|
||||
@@ -236,7 +236,7 @@ __device__ void _moe_align_block_size_small_batch_expert(
|
||||
int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
|
||||
size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks,
|
||||
int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num,
|
||||
int32_t* token_mask, bool has_expert_map) {
|
||||
int32_t* token_mask, bool has_expert_map, int32_t* scatter_idx = nullptr) {
|
||||
// Compute input buffer offsets. Typically these will all be 0, except when
|
||||
// using Multi LoRA.
|
||||
int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
|
||||
@@ -318,9 +318,14 @@ __device__ void _moe_align_block_size_small_batch_expert(
|
||||
}
|
||||
|
||||
for (size_t i = tid; i < numel; i += stride) {
|
||||
if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
|
||||
num_experts, has_expert_map);
|
||||
expert_id != -1) {
|
||||
int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts,
|
||||
has_expert_map);
|
||||
bool valid =
|
||||
expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]);
|
||||
if (scatter_idx != nullptr) {
|
||||
scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1;
|
||||
}
|
||||
if (expert_id != -1) {
|
||||
int32_t rank_post_pad =
|
||||
tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
|
||||
|
||||
@@ -338,14 +343,20 @@ __device__ void _count_and_sort_expert_tokens(
|
||||
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
|
||||
int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
|
||||
int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask,
|
||||
int32_t model_offset, int32_t topk_num, bool has_expert_map) {
|
||||
int32_t model_offset, int32_t topk_num, bool has_expert_map,
|
||||
int32_t* scatter_idx = nullptr) {
|
||||
const size_t tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const size_t stride = blockDim.x * gridDim.y;
|
||||
|
||||
for (size_t i = tid; i < numel; i += stride) {
|
||||
if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
|
||||
num_experts, has_expert_map);
|
||||
expert_id != -1) {
|
||||
int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts,
|
||||
has_expert_map);
|
||||
bool valid =
|
||||
expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]);
|
||||
if (scatter_idx != nullptr) {
|
||||
scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1;
|
||||
}
|
||||
if (expert_id != -1) {
|
||||
if (token_mask == nullptr || token_mask[i / topk_num]) {
|
||||
int32_t rank_post_pad = atomicAdd(
|
||||
&cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
|
||||
@@ -377,10 +388,11 @@ __global__ void count_and_sort_expert_tokens_kernel(
|
||||
const scalar_t* __restrict__ topk_ids,
|
||||
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
|
||||
int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
|
||||
int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) {
|
||||
int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map,
|
||||
int32_t* scatter_idx) {
|
||||
_count_and_sort_expert_tokens(
|
||||
topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
|
||||
max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
|
||||
max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map, scatter_idx);
|
||||
}
|
||||
|
||||
// Reduce the topk expert outputs per token (summed in fp32). The output is
|
||||
@@ -533,12 +545,12 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
|
||||
int32_t* __restrict__ total_tokens_post_pad,
|
||||
int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
|
||||
size_t numel, int32_t max_num_tokens_padded, int32_t topk_num,
|
||||
bool has_expert_map) {
|
||||
bool has_expert_map, int32_t* scatter_idx) {
|
||||
_moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
|
||||
topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
|
||||
num_experts, block_size, numel, max_num_tokens_padded,
|
||||
CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
|
||||
has_expert_map);
|
||||
has_expert_map, scatter_idx);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
@@ -658,16 +670,31 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
|
||||
|
||||
// taken from
|
||||
// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
|
||||
void moe_align_block_size(
|
||||
torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
|
||||
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
|
||||
torch::stable::Tensor num_tokens_post_pad,
|
||||
std::optional<torch::stable::Tensor> maybe_expert_map) {
|
||||
void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
|
||||
int64_t block_size,
|
||||
torch::stable::Tensor sorted_token_ids,
|
||||
torch::stable::Tensor experts_ids,
|
||||
torch::stable::Tensor num_tokens_post_pad,
|
||||
std::optional<torch::stable::Tensor> maybe_expert_map,
|
||||
std::optional<torch::stable::Tensor> scatter_idx) {
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
topk_ids.get_device_index());
|
||||
const cudaStream_t stream =
|
||||
get_current_cuda_stream(topk_ids.get_device_index());
|
||||
|
||||
int32_t* scatter_ptr = nullptr;
|
||||
if (scatter_idx.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"scatter_idx must be int32");
|
||||
STD_TORCH_CHECK(scatter_idx->device() == topk_ids.device() &&
|
||||
scatter_idx->is_contiguous() &&
|
||||
scatter_idx->numel() == topk_ids.numel(),
|
||||
"scatter_idx must be contiguous, on the same device, and "
|
||||
"match topk_ids.numel()");
|
||||
scatter_ptr = reinterpret_cast<int32_t*>(scatter_idx->mutable_data_ptr());
|
||||
}
|
||||
|
||||
int64_t padded_num_experts =
|
||||
((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
|
||||
int experts_per_warp = WARP_SIZE;
|
||||
@@ -713,7 +740,8 @@ void moe_align_block_size(
|
||||
num_tokens_post_pad.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
|
||||
num_experts, block_size, topk_ids.numel(),
|
||||
sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
|
||||
sorted_token_ids.size(0), topk_ids.size(1), has_expert_map,
|
||||
scatter_ptr);
|
||||
} else {
|
||||
torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
|
||||
topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
|
||||
@@ -753,7 +781,7 @@ void moe_align_block_size(
|
||||
reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
|
||||
topk_ids.numel(), num_experts, sorted_token_ids.size(0),
|
||||
topk_ids.size(1), has_expert_map);
|
||||
topk_ids.size(1), has_expert_map, scatter_ptr);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -36,11 +36,13 @@ void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> topk_ids,
|
||||
std::optional<torch::stable::Tensor> expert_map);
|
||||
|
||||
void moe_align_block_size(
|
||||
torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
|
||||
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
|
||||
torch::stable::Tensor num_tokens_post_pad,
|
||||
std::optional<torch::stable::Tensor> maybe_expert_map);
|
||||
void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
|
||||
int64_t block_size,
|
||||
torch::stable::Tensor sorted_token_ids,
|
||||
torch::stable::Tensor experts_ids,
|
||||
torch::stable::Tensor num_tokens_post_pad,
|
||||
std::optional<torch::stable::Tensor> maybe_expert_map,
|
||||
std::optional<torch::stable::Tensor> scatter_idx);
|
||||
|
||||
void batched_moe_align_block_size(
|
||||
int64_t max_tokens_per_batch, int64_t block_size,
|
||||
|
||||
@@ -56,20 +56,18 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
|
||||
CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
|
||||
}
|
||||
|
||||
void moe_permute_impl(
|
||||
const torch::stable::Tensor& input, // [n_token, hidden]
|
||||
const torch::stable::Tensor& topk_ids, // [n_token, topk]
|
||||
const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
|
||||
torch::stable::Tensor moe_sort_routing(
|
||||
const torch::stable::Tensor& topk_ids, // [n_token, topk]
|
||||
const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
|
||||
const std::optional<torch::stable::Tensor>& expert_map, // [n_expert]
|
||||
int64_t n_expert, int64_t n_local_expert, int64_t topk,
|
||||
torch::stable::Tensor& permuted_input, // [permuted_size, hidden]
|
||||
torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1]
|
||||
torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
|
||||
torch::stable::Tensor& permuted_idx, // [permute_size]
|
||||
const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
|
||||
const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
|
||||
const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
|
||||
const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
|
||||
const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort,
|
||||
bool build_inverse) {
|
||||
STD_TORCH_CHECK(expert_first_token_offset.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Long,
|
||||
"expert_first_token_offset must be int64");
|
||||
@@ -88,12 +86,11 @@ void moe_permute_impl(
|
||||
"token_expert_indices shape must be same as inv_permuted_idx");
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
auto device = input.device();
|
||||
auto n_token = input.sizes()[0];
|
||||
auto n_hidden = input.sizes()[1];
|
||||
topk_ids.get_device_index());
|
||||
auto device = topk_ids.device();
|
||||
auto n_token = topk_ids.sizes()[0];
|
||||
auto expanded_rows = n_token * topk;
|
||||
auto stream = get_current_cuda_stream(input.get_device_index());
|
||||
auto stream = get_current_cuda_stream(topk_ids.get_device_index());
|
||||
|
||||
auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
|
||||
auto sort_workspace = maybe_allocate_tensor(
|
||||
@@ -107,13 +104,11 @@ void moe_permute_impl(
|
||||
torch::headeronly::ScalarType::Int, device, "sorted_row_idx");
|
||||
|
||||
CubKeyValueSorter sorter{};
|
||||
int64_t* valid_num_ptr = nullptr;
|
||||
torch::stable::Tensor topk_ids_for_sort = topk_ids;
|
||||
|
||||
if (expert_map.has_value()) {
|
||||
const int* expert_map_ptr = get_ptr<int>(expert_map.value());
|
||||
valid_num_ptr =
|
||||
get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
|
||||
if (expert_map.has_value() || build_inverse) {
|
||||
const int* expert_map_ptr =
|
||||
expert_map.has_value() ? get_ptr<int>(expert_map.value()) : nullptr;
|
||||
topk_ids_for_sort = maybe_allocate_tensor(
|
||||
maybe_topk_ids_for_sort, topk_ids.sizes(),
|
||||
torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
|
||||
@@ -126,7 +121,39 @@ void moe_permute_impl(
|
||||
get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
|
||||
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
|
||||
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
|
||||
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
|
||||
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream,
|
||||
build_inverse ? get_ptr<int>(inv_permuted_idx) : nullptr);
|
||||
return sorted_row_idx;
|
||||
}
|
||||
|
||||
void moe_permute_impl(
|
||||
const torch::stable::Tensor& input, // [n_token, hidden]
|
||||
const torch::stable::Tensor& topk_ids, // [n_token, topk]
|
||||
const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
|
||||
const std::optional<torch::stable::Tensor>& expert_map, // [n_expert]
|
||||
int64_t n_expert, int64_t n_local_expert, int64_t topk,
|
||||
torch::stable::Tensor& permuted_input, // [permuted_size, hidden]
|
||||
torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1]
|
||||
torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
|
||||
torch::stable::Tensor& permuted_idx, // [permute_size]
|
||||
const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
|
||||
const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
|
||||
const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
|
||||
const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
auto sorted_row_idx = moe_sort_routing(
|
||||
topk_ids, token_expert_indices, expert_map, n_expert, n_local_expert,
|
||||
topk, expert_first_token_offset, inv_permuted_idx, maybe_sort_workspace,
|
||||
maybe_permuted_experts_id, maybe_sorted_row_idx, maybe_topk_ids_for_sort,
|
||||
false);
|
||||
auto stream = get_current_cuda_stream(input.get_device_index());
|
||||
auto n_token = input.size(0);
|
||||
auto n_hidden = input.size(1);
|
||||
auto valid_num_ptr =
|
||||
expert_map.has_value()
|
||||
? get_ptr<int64_t>(expert_first_token_offset) + n_local_expert
|
||||
: nullptr;
|
||||
|
||||
MOE_DISPATCH(input.scalar_type(), [&] {
|
||||
expandInputRowsKernelLauncher<scalar_t>(
|
||||
@@ -171,6 +198,22 @@ void moe_permute_with_scratch(
|
||||
topk_ids_for_sort);
|
||||
}
|
||||
|
||||
void moe_prepare_scatter(const torch::stable::Tensor& topk_ids,
|
||||
const torch::stable::Tensor& token_expert_indices,
|
||||
const std::optional<torch::stable::Tensor>& expert_map,
|
||||
int64_t n_expert, int64_t n_local_expert,
|
||||
torch::stable::Tensor& expert_first_token_offset,
|
||||
torch::stable::Tensor& scatter_idx,
|
||||
torch::stable::Tensor& sort_workspace,
|
||||
torch::stable::Tensor& sorted_experts,
|
||||
torch::stable::Tensor& sorted_rows,
|
||||
torch::stable::Tensor& topk_ids_for_sort) {
|
||||
moe_sort_routing(topk_ids, token_expert_indices, expert_map, n_expert,
|
||||
n_local_expert, topk_ids.size(1), expert_first_token_offset,
|
||||
scatter_idx, sort_workspace, sorted_experts, sorted_rows,
|
||||
topk_ids_for_sort, true);
|
||||
}
|
||||
|
||||
void moe_unpermute(
|
||||
const torch::stable::Tensor&
|
||||
permuted_hidden_states, // [n_token * topk, hidden]
|
||||
@@ -315,6 +358,19 @@ void moe_unpermute(
|
||||
STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
|
||||
}
|
||||
|
||||
void moe_prepare_scatter(const torch::stable::Tensor& topk_ids,
|
||||
const torch::stable::Tensor& token_expert_indices,
|
||||
const std::optional<torch::stable::Tensor>& expert_map,
|
||||
int64_t n_expert, int64_t n_local_expert,
|
||||
torch::stable::Tensor& expert_first_token_offset,
|
||||
torch::stable::Tensor& scatter_idx,
|
||||
torch::stable::Tensor& sort_workspace,
|
||||
torch::stable::Tensor& sorted_experts,
|
||||
torch::stable::Tensor& sorted_rows,
|
||||
torch::stable::Tensor& topk_ids_for_sort) {
|
||||
STD_TORCH_CHECK(false, "moe_prepare_scatter is not supported on CUDA < 12.0");
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
bool moe_permute_unpermute_supported() {
|
||||
@@ -328,5 +384,6 @@ bool moe_permute_unpermute_supported() {
|
||||
STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
|
||||
m.impl("moe_permute", TORCH_BOX(&moe_permute));
|
||||
m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
|
||||
m.impl("moe_prepare_scatter", TORCH_BOX(&moe_prepare_scatter));
|
||||
m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
|
||||
}
|
||||
|
||||
@@ -34,10 +34,6 @@ class ScalarType<half> {
|
||||
static __host__ __device__ float2 inline num22float2(const half2 x) {
|
||||
return __half22float2(x);
|
||||
}
|
||||
|
||||
static __host__ __device__ half2 inline float22num2(const float2 x) {
|
||||
return __float22half2_rn(x);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
@@ -71,10 +67,6 @@ class ScalarType<nv_bfloat16> {
|
||||
static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
|
||||
return __bfloat1622float2(x);
|
||||
}
|
||||
|
||||
static __host__ __device__ nv_bfloat162 inline float22num2(const float2 x) {
|
||||
return __float22bfloat162_rn(x);
|
||||
}
|
||||
#endif
|
||||
};
|
||||
|
||||
|
||||
+30
-24
@@ -82,34 +82,40 @@ __device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices,
|
||||
return target_location + 1;
|
||||
}
|
||||
|
||||
// Calculates the start offset of the tokens for a given expert. The last
|
||||
// element is the total number of valid tokens
|
||||
__global__ void computeExpertFirstTokenOffsetKernel(
|
||||
// Computes expert offsets (ending with the valid row count) and optional
|
||||
// inverse indices.
|
||||
__global__ void computeExpertOffsetsAndInverseKernel(
|
||||
int const* sorted_experts, int64_t const sorted_experts_len,
|
||||
int const num_experts, int64_t* expert_first_token_offset) {
|
||||
// First, compute the global tid. We only need 1 thread per expert.
|
||||
int const expert = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int const num_experts, int64_t* expert_first_token_offset,
|
||||
int const* sorted_rows, int* inverse) {
|
||||
int const index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
// Note that expert goes [0, num_experts] (inclusive) because we want a count
|
||||
// for the total number of active tokens at the end of the scan.
|
||||
if (expert >= num_experts + 1) {
|
||||
return;
|
||||
if (index < num_experts + 1) {
|
||||
expert_first_token_offset[index] =
|
||||
findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, index);
|
||||
}
|
||||
if (inverse != nullptr && index < sorted_experts_len) {
|
||||
inverse[sorted_rows[index]] = index;
|
||||
}
|
||||
expert_first_token_offset[expert] =
|
||||
findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert);
|
||||
}
|
||||
|
||||
void computeExpertFirstTokenOffset(int const* sorted_indices,
|
||||
int const total_indices,
|
||||
int const num_experts,
|
||||
int64_t* expert_first_token_offset,
|
||||
cudaStream_t stream) {
|
||||
int const num_entries = num_experts + 1;
|
||||
int const threads = std::min(1024, num_entries);
|
||||
void computeExpertOffsetsAndInverse(int const* sorted_indices,
|
||||
int const total_indices,
|
||||
int const num_experts,
|
||||
int64_t* expert_first_token_offset,
|
||||
int const* sorted_rows, int* inverse,
|
||||
cudaStream_t stream) {
|
||||
int const num_entries = inverse != nullptr
|
||||
? std::max(num_experts + 1, total_indices)
|
||||
: num_experts + 1;
|
||||
int const threads = inverse != nullptr ? 256 : std::min(1024, num_entries);
|
||||
int const blocks = (num_entries + threads - 1) / threads;
|
||||
|
||||
computeExpertFirstTokenOffsetKernel<<<blocks, threads, 0, stream>>>(
|
||||
sorted_indices, total_indices, num_experts, expert_first_token_offset);
|
||||
computeExpertOffsetsAndInverseKernel<<<blocks, threads, 0, stream>>>(
|
||||
sorted_indices, total_indices, num_experts, expert_first_token_offset,
|
||||
sorted_rows, inverse);
|
||||
}
|
||||
|
||||
void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
|
||||
@@ -117,7 +123,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
|
||||
int64_t* expert_first_token_offset, int num_rows,
|
||||
int num_experts, int num_experts_per_node, int k,
|
||||
CubKeyValueSorter& sorter, void* sorter_ws,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream, int* inverse) {
|
||||
int64_t const expanded_num_rows = static_cast<int64_t>(k) * num_rows;
|
||||
// We need to use the full num_experts because that is the sentinel value used
|
||||
// by topk for disabled experts
|
||||
@@ -127,9 +133,9 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
|
||||
sorter.run((void*)sorter_ws, sorter_ws_size_bytes, expert_for_source_row,
|
||||
permuted_experts, source_rows, permuted_rows, expanded_num_rows,
|
||||
stream);
|
||||
computeExpertFirstTokenOffset(permuted_experts, expanded_num_rows,
|
||||
num_experts_per_node, expert_first_token_offset,
|
||||
stream);
|
||||
computeExpertOffsetsAndInverse(
|
||||
permuted_experts, expanded_num_rows, num_experts_per_node,
|
||||
expert_first_token_offset, permuted_rows, inverse, stream);
|
||||
}
|
||||
|
||||
__global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
|
||||
@@ -142,7 +148,7 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
|
||||
extern __shared__ int smem_expert_map[];
|
||||
// store expert_map in smem
|
||||
for (int i = tidx; i < num_experts; i += blockDim.x) {
|
||||
smem_expert_map[i] = expert_map_ptr[i];
|
||||
smem_expert_map[i] = expert_map_ptr ? expert_map_ptr[i] : i;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
|
||||
@@ -43,18 +43,19 @@ class CubKeyValueSorter {
|
||||
int num_bits_;
|
||||
};
|
||||
|
||||
void computeExpertFirstTokenOffset(int const* sorted_indices,
|
||||
int const total_indices,
|
||||
int const num_experts,
|
||||
int64_t* expert_first_token_offset,
|
||||
cudaStream_t stream);
|
||||
void computeExpertOffsetsAndInverse(int const* sorted_indices,
|
||||
int const total_indices,
|
||||
int const num_experts,
|
||||
int64_t* expert_first_token_offset,
|
||||
int const* sorted_rows, int* inverse,
|
||||
cudaStream_t stream);
|
||||
|
||||
void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
|
||||
int* permuted_experts, int* permuted_rows,
|
||||
int64_t* expert_first_token_offset, int num_rows,
|
||||
int num_experts, int num_experts_per_node, int k,
|
||||
CubKeyValueSorter& sorter, void* sorter_ws,
|
||||
cudaStream_t stream);
|
||||
cudaStream_t stream, int* inverse = nullptr);
|
||||
|
||||
template <typename T>
|
||||
void expandInputRowsKernelLauncher(
|
||||
|
||||
@@ -39,7 +39,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
" int block_size, Tensor! sorted_token_ids,"
|
||||
" Tensor! experts_ids,"
|
||||
" Tensor! num_tokens_post_pad,"
|
||||
" Tensor? maybe_expert_map) -> ()");
|
||||
" Tensor? maybe_expert_map,"
|
||||
" Tensor(a!)? scatter_idx=None) -> ()");
|
||||
|
||||
// Aligning the number of tokens to be processed by each expert such
|
||||
// that it is divisible by the block size, but for the batched case.
|
||||
@@ -105,6 +106,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
"permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
|
||||
"Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
|
||||
|
||||
m.def(
|
||||
"moe_prepare_scatter(Tensor topk_ids, Tensor token_expert_indices, "
|
||||
"Tensor? expert_map, int n_expert, int n_local_expert, "
|
||||
"Tensor! expert_first_token_offset, Tensor! scatter_idx, "
|
||||
"Tensor! sort_workspace, Tensor! sorted_experts, Tensor! sorted_rows, "
|
||||
"Tensor! topk_ids_for_sort)->()");
|
||||
|
||||
m.def(
|
||||
"moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
|
||||
"Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
|
||||
|
||||
@@ -74,18 +74,9 @@ class MatrixView_half_rw {
|
||||
const int width)
|
||||
: data(data), height(height), width(width) {}
|
||||
|
||||
__device__ __forceinline__ half item(int row, int column) const {
|
||||
return data[row * width + column];
|
||||
}
|
||||
__device__ __forceinline__ half* item_ptr(int row, int column) {
|
||||
return &data[row * width + column];
|
||||
}
|
||||
__device__ __forceinline__ void set(int row, int column, half value) {
|
||||
data[row * width + column] = value;
|
||||
}
|
||||
__device__ __forceinline__ void set_half2(int row, int column, half2 value) {
|
||||
((half2*)data)[(row * width + column) / 2] = value;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void set4(int row, int column, half v0, half v1,
|
||||
half v2, half v3) {
|
||||
@@ -113,14 +104,6 @@ class MatrixView_q4_row {
|
||||
return (data[row * width / 8 + column / 8] >> shift) & 0x0f;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void item2(int (&items)[2], int row,
|
||||
int column) const {
|
||||
int shift = (column & 0x07) * 4;
|
||||
uint32_t d = data[row * width / 8 + column / 8] >> shift;
|
||||
items[0] = d & 0x0f;
|
||||
items[1] = (d >> 4) & 0x0f;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void item4(int (&items)[4], int row,
|
||||
int column) const {
|
||||
int shift = (column & 0x07) * 4;
|
||||
@@ -148,14 +131,6 @@ class MatrixView_q2_row {
|
||||
return (data[row * width / 16 + column / 16] >> shift) & 0x03;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void item2(int (&items)[2], int row,
|
||||
int column) const {
|
||||
int shift = (column & 0x0f) * 2;
|
||||
uint32_t d = data[row * width / 16 + column / 16] >> shift;
|
||||
items[0] = d & 0x03;
|
||||
items[1] = (d >> 2) & 0x03;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void item4(int (&items)[4], int row,
|
||||
int column) const {
|
||||
int shift = (column & 0x0f) * 2;
|
||||
@@ -237,14 +212,6 @@ class MatrixView_q8_row {
|
||||
return (data[row * width / 4 + column / 4] >> shift) & 0xff;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void item2(int (&items)[2], int row,
|
||||
int column) const {
|
||||
int shift = (column & 0x03) * 8;
|
||||
uint32_t d = data[row * width / 4 + column / 4] >> shift;
|
||||
items[0] = d & 0xff;
|
||||
items[1] = (d >> 8) & 0xff;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void item4(int (&items)[4], int row,
|
||||
int column) const {
|
||||
int shift = (column & 0x03) * 2;
|
||||
|
||||
@@ -67,14 +67,6 @@ __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
|
||||
}
|
||||
}
|
||||
|
||||
__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
|
||||
bool pred = true) {
|
||||
if (pred) {
|
||||
reinterpret_cast<int4*>(smem_ptr)[0] =
|
||||
reinterpret_cast<const int4*>(glob_ptr)[0];
|
||||
}
|
||||
}
|
||||
|
||||
__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
|
||||
bool pred = true) {
|
||||
if (pred) {
|
||||
@@ -121,19 +113,6 @@ __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
|
||||
"r"(smem), "l"(glob_ptr), "n"(BYTES));
|
||||
}
|
||||
|
||||
__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
|
||||
bool pred = true) {
|
||||
const int BYTES = 16;
|
||||
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.ca.shared.global [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)pred),
|
||||
"r"(smem), "l"(glob_ptr), "n"(BYTES));
|
||||
}
|
||||
|
||||
__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
|
||||
bool pred = true) {
|
||||
const int BYTES = 16;
|
||||
@@ -170,4 +149,4 @@ __device__ inline void cp_async_wait() {
|
||||
|
||||
} // namespace MARLIN_NAMESPACE_NAME
|
||||
|
||||
#endif
|
||||
#endif
|
||||
|
||||
@@ -45,11 +45,6 @@ __global__ void apply_repetition_penalties_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
|
||||
uint32_t bits = __float_as_uint(x);
|
||||
return (bits & 0x80000000) ? bits : ~bits & 0x7fffffff;
|
||||
}
|
||||
|
||||
template <int step>
|
||||
static inline __device__ uint32_t extractBinIdx(float x) {
|
||||
if constexpr (step == 0) {
|
||||
|
||||
@@ -67,104 +67,6 @@ using fp8_type = __hip_fp8_e4m3_fnuz;
|
||||
using fp8x2_type = __hip_fp8x2_e4m3_fnuz;
|
||||
#endif
|
||||
|
||||
// fp8 -> half
|
||||
template <>
|
||||
__inline__ __device__ uint16_t
|
||||
vec_conversion<uint16_t, uint8_t>(const uint8_t& a) {
|
||||
return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x;
|
||||
}
|
||||
|
||||
// fp8x2 -> half2
|
||||
template <>
|
||||
__inline__ __device__ uint32_t
|
||||
vec_conversion<uint32_t, uint16_t>(const uint16_t& a) {
|
||||
union {
|
||||
__half2_raw h2r;
|
||||
uint32_t ui32;
|
||||
} tmp;
|
||||
tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
|
||||
return tmp.ui32;
|
||||
}
|
||||
|
||||
// fp8x4 -> half2x2
|
||||
template <>
|
||||
__inline__ __device__ uint2 vec_conversion<uint2, uint32_t>(const uint32_t& a) {
|
||||
union {
|
||||
uint2 u32x2;
|
||||
uint32_t u32[2];
|
||||
} tmp;
|
||||
tmp.u32[0] = vec_conversion<uint32_t, uint16_t>((uint16_t)a);
|
||||
tmp.u32[1] = vec_conversion<uint32_t, uint16_t>((uint16_t)(a >> 16U));
|
||||
return tmp.u32x2;
|
||||
}
|
||||
|
||||
// fp8x8 -> half2x4
|
||||
template <>
|
||||
__inline__ __device__ uint4 vec_conversion<uint4, uint2>(const uint2& a) {
|
||||
union {
|
||||
uint4 u64x2;
|
||||
uint2 u64[2];
|
||||
} tmp;
|
||||
tmp.u64[0] = vec_conversion<uint2, uint32_t>(a.x);
|
||||
tmp.u64[1] = vec_conversion<uint2, uint32_t>(a.y);
|
||||
return tmp.u64x2;
|
||||
}
|
||||
|
||||
using __nv_bfloat16 = __hip_bfloat16;
|
||||
|
||||
// fp8 -> __nv_bfloat16
|
||||
template <>
|
||||
__inline__ __device__ __nv_bfloat16
|
||||
vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) {
|
||||
fp8_type f8;
|
||||
f8.__x = a;
|
||||
return __float2bfloat16(static_cast<float>(f8));
|
||||
}
|
||||
|
||||
using __nv_bfloat162 = __hip_bfloat162;
|
||||
|
||||
// fp8x2 -> __nv_bfloat162
|
||||
template <>
|
||||
__inline__ __device__ __nv_bfloat162
|
||||
vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a) {
|
||||
__nv_bfloat162 res;
|
||||
res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a);
|
||||
res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U));
|
||||
return res;
|
||||
}
|
||||
|
||||
// fp8x4 -> bf16_4_t
|
||||
template <>
|
||||
__inline__ __device__ bf16_4_t
|
||||
vec_conversion<bf16_4_t, uint32_t>(const uint32_t& a) {
|
||||
bf16_4_t res;
|
||||
res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a);
|
||||
res.y = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U));
|
||||
return res;
|
||||
}
|
||||
|
||||
// fp8x8 -> bf16_8_t
|
||||
template <>
|
||||
__inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, uint2>(const uint2& a) {
|
||||
bf16_4_t tmp1, tmp2;
|
||||
tmp1 = vec_conversion<bf16_4_t, uint32_t>(a.x);
|
||||
tmp2 = vec_conversion<bf16_4_t, uint32_t>(a.y);
|
||||
bf16_8_t res;
|
||||
res.x = tmp1.x;
|
||||
res.y = tmp1.y;
|
||||
res.z = tmp2.x;
|
||||
res.w = tmp2.y;
|
||||
return res;
|
||||
}
|
||||
|
||||
// fp8 -> float
|
||||
template <>
|
||||
__inline__ __device__ float vec_conversion<float, uint8_t>(const uint8_t& a) {
|
||||
fp8_type f8;
|
||||
f8.__x = a;
|
||||
return static_cast<float>(f8);
|
||||
}
|
||||
|
||||
// fp8x2 -> float2
|
||||
template <>
|
||||
__inline__ __device__ float2
|
||||
@@ -193,138 +95,6 @@ vec_conversion<float4, uint32_t>(const uint32_t& a) {
|
||||
return res;
|
||||
}
|
||||
|
||||
// fp8x8 -> float8
|
||||
template <>
|
||||
__inline__ __device__ Float8_ vec_conversion<Float8_, uint2>(const uint2& a) {
|
||||
Float4_ tmp1, tmp2;
|
||||
tmp1 = vec_conversion<Float4_, uint32_t>(a.x);
|
||||
tmp2 = vec_conversion<Float4_, uint32_t>(a.y);
|
||||
Float8_ res;
|
||||
res.x = tmp1.x;
|
||||
res.y = tmp1.y;
|
||||
res.z = tmp2.x;
|
||||
res.w = tmp2.y;
|
||||
return res;
|
||||
}
|
||||
|
||||
// half -> fp8
|
||||
template <>
|
||||
__inline__ __device__ uint8_t
|
||||
vec_conversion<uint8_t, uint16_t>(const uint16_t& a) {
|
||||
__half_raw tmp;
|
||||
tmp.x = a;
|
||||
return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation,
|
||||
fp8_type::__default_interpret);
|
||||
}
|
||||
|
||||
template <>
|
||||
__inline__ __device__ uint16_t
|
||||
vec_conversion<uint16_t, uint32_t>(const uint32_t& a) {
|
||||
union {
|
||||
uint32_t ui32;
|
||||
__half2_raw h2r;
|
||||
} tmp;
|
||||
tmp.ui32 = a;
|
||||
return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation,
|
||||
fp8_type::__default_interpret);
|
||||
}
|
||||
|
||||
// bf16 -> fp8
|
||||
template <>
|
||||
__inline__ __device__ uint8_t
|
||||
vec_conversion<uint8_t, __nv_bfloat16>(const __nv_bfloat16& a) {
|
||||
return __hip_cvt_float_to_fp8(__bfloat162float(a),
|
||||
fp8_type::__default_saturation,
|
||||
fp8_type::__default_interpret);
|
||||
}
|
||||
|
||||
// float -> fp8
|
||||
template <>
|
||||
__inline__ __device__ uint8_t vec_conversion<uint8_t, float>(const float& a) {
|
||||
return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation,
|
||||
fp8_type::__default_interpret);
|
||||
}
|
||||
|
||||
// float2 -> half2
|
||||
template <>
|
||||
__inline__ __device__ uint32_t
|
||||
vec_conversion<uint32_t, float2>(const float2& a) {
|
||||
union {
|
||||
half2 float16;
|
||||
uint32_t uint32;
|
||||
};
|
||||
|
||||
float16 = __float22half2_rn(a);
|
||||
return uint32;
|
||||
}
|
||||
|
||||
// Float4 -> half2x2
|
||||
template <>
|
||||
__inline__ __device__ uint2 vec_conversion<uint2, Float4_>(const Float4_& a) {
|
||||
uint2 b;
|
||||
float2 val;
|
||||
val.x = a.x.x;
|
||||
val.y = a.x.y;
|
||||
b.x = vec_conversion<uint32_t, float2>(val);
|
||||
|
||||
val.x = a.y.x;
|
||||
val.y = a.y.y;
|
||||
b.y = vec_conversion<uint32_t, float2>(val);
|
||||
return b;
|
||||
}
|
||||
|
||||
// Float4 -> float4
|
||||
template <>
|
||||
__inline__ __device__ float4 vec_conversion<float4, Float4_>(const Float4_& a) {
|
||||
float4 b;
|
||||
b.x = a.x.x;
|
||||
b.y = a.x.y;
|
||||
b.z = a.y.x;
|
||||
b.w = a.y.y;
|
||||
return b;
|
||||
}
|
||||
|
||||
// Float8 -> half2x4
|
||||
template <>
|
||||
__inline__ __device__ uint4 vec_conversion<uint4, Float8_>(const Float8_& a) {
|
||||
uint4 b;
|
||||
b.x = vec_conversion<uint32_t, float2>(a.x);
|
||||
b.y = vec_conversion<uint32_t, float2>(a.y);
|
||||
b.z = vec_conversion<uint32_t, float2>(a.z);
|
||||
b.w = vec_conversion<uint32_t, float2>(a.w);
|
||||
return b;
|
||||
}
|
||||
|
||||
// float2 -> bfloat162
|
||||
template <>
|
||||
__inline__ __device__ __nv_bfloat162
|
||||
vec_conversion<__nv_bfloat162, float2>(const float2& a) {
|
||||
__nv_bfloat162 b = __float22bfloat162_rn(a);
|
||||
return b;
|
||||
}
|
||||
|
||||
// Float4 -> bfloat162x2
|
||||
template <>
|
||||
__inline__ __device__ bf16_4_t
|
||||
vec_conversion<bf16_4_t, Float4_>(const Float4_& a) {
|
||||
bf16_4_t b;
|
||||
b.x = __float22bfloat162_rn(a.x);
|
||||
b.y = __float22bfloat162_rn(a.y);
|
||||
return b;
|
||||
}
|
||||
|
||||
// Float8 -> bfloat162x4
|
||||
template <>
|
||||
__inline__ __device__ bf16_8_t
|
||||
vec_conversion<bf16_8_t, Float8_>(const Float8_& a) {
|
||||
bf16_8_t b;
|
||||
b.x = __float22bfloat162_rn(a.x);
|
||||
b.y = __float22bfloat162_rn(a.y);
|
||||
b.z = __float22bfloat162_rn(a.z);
|
||||
b.w = __float22bfloat162_rn(a.w);
|
||||
return b;
|
||||
}
|
||||
|
||||
/* Scaled and vectorized conversions, for data exchange between high and low
|
||||
precision domains
|
||||
|
||||
@@ -335,6 +105,7 @@ vec_conversion<bf16_8_t, Float8_>(const Float8_& a) {
|
||||
*/
|
||||
|
||||
using __nv_bfloat16 = __hip_bfloat16;
|
||||
using __nv_bfloat162 = __hip_bfloat162;
|
||||
|
||||
// fp8 -> __nv_bfloat16
|
||||
template <>
|
||||
@@ -613,17 +384,6 @@ scaled_vec_conversion<uint32_t, float4>(const float4& a, float scale) {
|
||||
}
|
||||
#endif // ENABLE_FP8
|
||||
|
||||
template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
|
||||
__inline__ __device__ Tout convert(const Tin& x) {
|
||||
#ifdef ENABLE_FP8
|
||||
if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) {
|
||||
return vec_conversion<Tout, Tin>(x);
|
||||
}
|
||||
#endif
|
||||
assert(false);
|
||||
return {}; // Squash missing return statement warning
|
||||
}
|
||||
|
||||
template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
|
||||
__inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
|
||||
#ifdef ENABLE_FP8
|
||||
|
||||
@@ -282,19 +282,6 @@ __inline__ __device__ float4 scaled_vec_conversion<float4, uint32_t>(
|
||||
}
|
||||
#endif // ENABLE_FP8
|
||||
|
||||
template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
|
||||
__inline__ __device__ Tout convert(const Tin& x) {
|
||||
#if 0 // Disable the following code to reduce the binary size.
|
||||
if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) {
|
||||
return vec_conversion<Tout, Tin>(x, __NV_E4M3);
|
||||
} else if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E5M2) {
|
||||
return vec_conversion<Tout, Tin>(x, __NV_E5M2);
|
||||
}
|
||||
#endif
|
||||
assert(false);
|
||||
__builtin_unreachable(); // Suppress missing return statement warning
|
||||
}
|
||||
|
||||
template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
|
||||
__inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
|
||||
#ifdef ENABLE_FP8
|
||||
|
||||
@@ -342,22 +342,6 @@ typedef union u64_cvt {
|
||||
int64_t i64;
|
||||
} _T8x8;
|
||||
|
||||
__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input,
|
||||
_T8x8& Mtemp) {
|
||||
_T8x8 Qtmp8x8;
|
||||
|
||||
for (int i = 0; i < 2; i++) {
|
||||
floatx4 q_out = {0, 0, 0, 0};
|
||||
q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i],
|
||||
q_out);
|
||||
Qtmp8x8.b16x4[i * 2] =
|
||||
__builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false);
|
||||
Qtmp8x8.b16x4[i * 2 + 1] =
|
||||
__builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false);
|
||||
}
|
||||
return Qtmp8x8.b8x8;
|
||||
}
|
||||
|
||||
__device__ float warpReduceMax(float val) {
|
||||
for (int offset = warpSize / 2; offset > 0; offset /= 2) {
|
||||
val = max(
|
||||
@@ -2583,23 +2567,6 @@ __device__ __forceinline__ float to_float(const T& inp) {
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float to_float_b16(const bit16_t& inp) {
|
||||
union tmpcvt {
|
||||
bit16_t u;
|
||||
_Float16 f;
|
||||
__hip_bfloat16 b;
|
||||
} t16;
|
||||
t16.u = inp;
|
||||
if constexpr (std::is_same<T, _Float16>::value) {
|
||||
return (float)t16.f;
|
||||
} else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
|
||||
return __bfloat162float(t16.b);
|
||||
} else {
|
||||
static_assert(false, "unsupported 16b dtype");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T from_float(const float& inp) {
|
||||
if constexpr (std::is_same<T, _Float16>::value) {
|
||||
|
||||
@@ -67,21 +67,6 @@ __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
|
||||
const bf16_t* a_ptr) {
|
||||
float result = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
uint32_t aw;
|
||||
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
|
||||
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
|
||||
float a_y = __uint_as_float(aw & 0xFFFF0000u);
|
||||
result = __fmaf_rn(dq[2 * i + 0], a_x, result);
|
||||
result = __fmaf_rn(dq[2 * i + 1], a_y, result);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
|
||||
half2 v23) {
|
||||
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
|
||||
|
||||
@@ -115,61 +115,6 @@ __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4],
|
||||
const bf16_t* a_ptr) {
|
||||
// RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in
|
||||
// the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc
|
||||
// therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised
|
||||
// fallback (single-element FMAs or fp32 round-trips), which empirically
|
||||
// runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16
|
||||
// decode path was paying that tax in full, scaling linearly with M (the
|
||||
// fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate
|
||||
// and the kernel becomes memory-bound).
|
||||
//
|
||||
// Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs)
|
||||
// and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA
|
||||
// count, but each FMA is fast. Bonus: the accumulator is now fp32
|
||||
// throughout instead of bf16, which is also numerically more accurate
|
||||
// (no compounding bf16-rounding inside the dot loop).
|
||||
float result = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
uint32_t aw, dw;
|
||||
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
|
||||
__builtin_memcpy(&dw, &dq[i], sizeof(uint32_t));
|
||||
// bf16 in low 16 bits → fp32 by left-shifting into the upper half.
|
||||
// bf16 in high 16 bits → already aligned with fp32's upper half.
|
||||
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
|
||||
float a_y = __uint_as_float(aw & 0xFFFF0000u);
|
||||
float d_x = __uint_as_float((dw & 0xFFFFu) << 16);
|
||||
float d_y = __uint_as_float(dw & 0xFFFF0000u);
|
||||
result = __fmaf_rn(d_x, a_x, result);
|
||||
result = __fmaf_rn(d_y, a_y, result);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already
|
||||
// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t
|
||||
// overload above does for dq (still need to widen A from bf16). Wins more
|
||||
// at high N: the bf162_t version's per-call widening cost scales with the
|
||||
// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only
|
||||
// for A widening (M_COUNT × 4 × 4 widens, half as many).
|
||||
__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
|
||||
const bf16_t* a_ptr) {
|
||||
float result = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
uint32_t aw;
|
||||
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
|
||||
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
|
||||
float a_y = __uint_as_float(aw & 0xFFFF0000u);
|
||||
result = __fmaf_rn(dq[2 * i + 0], a_x, result);
|
||||
result = __fmaf_rn(dq[2 * i + 1], a_y, result);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
|
||||
// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
|
||||
@@ -327,9 +272,8 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
|
||||
|
||||
// Per-column dequant constants. We hold one set of (z, y) pairs per column.
|
||||
// fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
|
||||
// nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant
|
||||
// produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA
|
||||
// bypass for the missing v_pk_fma_bf16 on gfx11.
|
||||
// nibble-*16 trick. bf16 uses fp32 scalars (z, y) for the inline v_dot2
|
||||
// bias correction.
|
||||
half2 z1z16_h[4][2], y1y16_h[4][2];
|
||||
float z_b_f[4], y_b_f[4];
|
||||
|
||||
|
||||
@@ -113,63 +113,6 @@ __forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
|
||||
// bf16 path
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Bit-trick magic for bf16:
|
||||
// bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0).
|
||||
// For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because
|
||||
// bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So
|
||||
// ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi).
|
||||
//
|
||||
// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble
|
||||
// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16]
|
||||
// via a single 4/8/12-bit right-shift before the OR. That costs one extra
|
||||
// shift per pair vs fp16, but keeps the FMA structure identical.
|
||||
__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero,
|
||||
bf16_t scale,
|
||||
bf162_t& z_prep,
|
||||
bf162_t& y_prep) {
|
||||
// z = scale * -(128 + zero); y = scale.
|
||||
float scale_f = __bfloat162float(scale);
|
||||
float zf = -(128.0f + (float)zero) * scale_f;
|
||||
bf16_t zb = __float2bfloat16(zf);
|
||||
z_prep = __bfloat162bfloat162(zb);
|
||||
y_prep = __bfloat162bfloat162(scale);
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa,
|
||||
bf162_t (&dq)[4],
|
||||
bf162_t z_prep,
|
||||
bf162_t y_prep) {
|
||||
const uint32_t c0 = 0x43004300;
|
||||
|
||||
union {
|
||||
uint32_t u;
|
||||
bf162_t b2;
|
||||
} q0, q1, q2, q3;
|
||||
q0.u = ((qa >> 0) & 0x000F000F) | c0; // bf162(128+q[0], 128+q[1])
|
||||
q1.u = ((qa >> 4) & 0x000F000F) | c0; // bf162(128+q[2], 128+q[3])
|
||||
q2.u = ((qa >> 8) & 0x000F000F) | c0; // bf162(128+q[4], 128+q[5])
|
||||
q3.u = ((qa >> 12) & 0x000F000F) | c0; // bf162(128+q[6], 128+q[7])
|
||||
|
||||
// dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale
|
||||
dq[0] = __hfma2(q0.b2, y_prep, z_prep);
|
||||
dq[1] = __hfma2(q1.b2, y_prep, z_prep);
|
||||
dq[2] = __hfma2(q2.b2, y_prep, z_prep);
|
||||
dq[3] = __hfma2(q3.b2, y_prep, z_prep);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// bf16-input → fp32-output dequant (RDNA3 scalar path).
|
||||
//
|
||||
// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow
|
||||
// fallback. Rather than computing dq in bf16 and widening at FMA time in
|
||||
// the dot product, we widen to fp32 here once (a free left-shift by 16) and
|
||||
// emit the (q - zero) * scale FMA directly in fp32. This:
|
||||
// * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32.
|
||||
// * Eliminates 4× bf16→fp32 widens that the dot product would do.
|
||||
// * Keeps the dot product accumulator in fp32 without a roundtrip.
|
||||
//
|
||||
// Output: fp32 dq[8], one element per K position (consumed by the
|
||||
// fp32-overload of dot22_8_f in q_gemm_rdna3.cu).
|
||||
__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
|
||||
bf16_t scale,
|
||||
float& z_prep,
|
||||
@@ -179,60 +122,6 @@ __forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
|
||||
y_prep = scale_f;
|
||||
}
|
||||
|
||||
// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32
|
||||
// values 128+nibble, without folding scale/zero. The caller folds scale/zb
|
||||
// into the accumulator outside the inner loop using a precomputed sum_a,
|
||||
// which saves ~27% of the FMA count vs the per-col-dequant approach above
|
||||
// (only beneficial at M_COUNT=1; break-even at M_COUNT=2).
|
||||
//
|
||||
// Cost: 0 FMAs (pure bit-trick + as_float reinterprets).
|
||||
__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa,
|
||||
float (&q_f32)[8]) {
|
||||
const uint32_t c0 = 0x43004300;
|
||||
const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
|
||||
const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
|
||||
const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
|
||||
const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
|
||||
q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16);
|
||||
q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u);
|
||||
q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16);
|
||||
q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u);
|
||||
q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16);
|
||||
q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u);
|
||||
q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16);
|
||||
q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u);
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa,
|
||||
float (&dq)[8],
|
||||
float z_prep,
|
||||
float y_prep) {
|
||||
const uint32_t c0 = 0x43004300;
|
||||
const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
|
||||
const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
|
||||
const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
|
||||
const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
|
||||
// bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16
|
||||
// (just zero-extends the mantissa from 7 to 23 bits; exponent preserved).
|
||||
const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
|
||||
const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
|
||||
const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
|
||||
const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
|
||||
const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
|
||||
const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
|
||||
const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
|
||||
const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
|
||||
// dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale
|
||||
dq[0] = __fmaf_rn(q0x, y_prep, z_prep);
|
||||
dq[1] = __fmaf_rn(q0y, y_prep, z_prep);
|
||||
dq[2] = __fmaf_rn(q1x, y_prep, z_prep);
|
||||
dq[3] = __fmaf_rn(q1y, y_prep, z_prep);
|
||||
dq[4] = __fmaf_rn(q2x, y_prep, z_prep);
|
||||
dq[5] = __fmaf_rn(q2y, y_prep, z_prep);
|
||||
dq[6] = __fmaf_rn(q3x, y_prep, z_prep);
|
||||
dq[7] = __fmaf_rn(q3y, y_prep, z_prep);
|
||||
}
|
||||
|
||||
} // namespace gptq_rdna3
|
||||
} // namespace vllm
|
||||
|
||||
|
||||
+52
-4
@@ -477,6 +477,24 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
#################### EXTENSIONS BUILD IMAGE ####################
|
||||
# Build DeepEP - runs in PARALLEL with csrc-build
|
||||
# This stage is independent and doesn't affect csrc cache
|
||||
# Kernel symbol map for CI test selection: which source file produced each GPU
|
||||
# kernel symbol, read from the objects this stage just compiled. Only when
|
||||
# VLLM_KERNEL_SYMBOL_MAP=1 (nightly), and placed after the compile layers so
|
||||
# their cache is untouched. Never fails the build; an empty map is a signal.
|
||||
ARG VLLM_KERNEL_SYMBOL_MAP=""
|
||||
ARG buildkite_commit=""
|
||||
COPY tools/ci/kernel_symbol_map.py tools/ci/kernel_symbol_map.py
|
||||
RUN if [ "${VLLM_KERNEL_SYMBOL_MAP}" = "1" ]; then \
|
||||
mkdir -p dist && python3 tools/ci/kernel_symbol_map.py \
|
||||
--build-root build --source-root . --commit "${buildkite_commit}" \
|
||||
--out dist/kernel_symbol_map.json.gz; \
|
||||
else echo "Skipping kernel symbol map (VLLM_KERNEL_SYMBOL_MAP != 1)"; fi
|
||||
|
||||
# Exported on its own by the CI bake target of the same name (ci-infra
|
||||
# docker/ci.hcl) with output=type=local, so the map becomes a build artifact.
|
||||
FROM scratch AS kernel-symbol-map
|
||||
COPY --from=csrc-build /workspace/dist/kernel_symbol_map.json.gz /kernel_symbol_map.json.gz
|
||||
|
||||
FROM base AS extensions-build
|
||||
ARG TARGETARCH
|
||||
ARG CUDA_VERSION
|
||||
@@ -735,6 +753,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base
|
||||
|
||||
ARG INSTALL_RUBIN_PRERELEASE
|
||||
ARG TARGETPLATFORM
|
||||
ARG max_jobs=2
|
||||
ARG CUDA_VERSION
|
||||
ARG PYTHON_VERSION
|
||||
ARG NCCL_VERSION
|
||||
@@ -854,6 +874,20 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
&& ldconfig; \
|
||||
fi
|
||||
|
||||
# Install the initialized-snapshot runtime on the architecture supported by
|
||||
# `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
|
||||
# Apache-2.0 source against the public CUDA Driver checkpoint APIs instead of
|
||||
# redistributing NVIDIA's proprietary cuda-checkpoint binary.
|
||||
COPY tools/install_snapshot_runtime.sh /usr/local/bin/install_snapshot_runtime.sh
|
||||
RUN if [ "${TARGETPLATFORM}" = "linux/amd64" ]; then \
|
||||
MAX_JOBS="${max_jobs}" /usr/local/bin/install_snapshot_runtime.sh; \
|
||||
else \
|
||||
echo "Initialized snapshots are not installed for ${TARGETPLATFORM}"; \
|
||||
fi
|
||||
|
||||
ENV CRIU_CUDA_PLUGIN_DIR=/usr/local/lib/criu
|
||||
|
||||
# Install uv for faster pip installs
|
||||
RUN python3 -m pip install uv
|
||||
|
||||
@@ -1047,6 +1081,10 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi
|
||||
|
||||
# Expose the Rust frontend bundled in the wheel on the container's PATH.
|
||||
RUN ln -s "$(python3 -c 'import sysconfig; print(sysconfig.get_path("platlib"))')/vllm/vllm-rs" /usr/local/bin/vllm-rs \
|
||||
&& vllm-rs --version
|
||||
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
. /etc/environment && \
|
||||
uv pip list
|
||||
@@ -1226,7 +1264,8 @@ ENV UV_HTTP_TIMEOUT=500
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
--mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
|
||||
--mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
|
||||
set -eu; \
|
||||
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
|
||||
CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \
|
||||
CUDA_HOME=/usr/local/cuda; \
|
||||
@@ -1235,11 +1274,19 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
libcusolver-dev-${CUDA_VERSION_DASH}"; \
|
||||
if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
|
||||
uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
|
||||
KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors.txt; \
|
||||
if [ "${CUDA_MAJOR}" = "12" ]; then \
|
||||
KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors_cu12.txt; \
|
||||
fi; \
|
||||
# Preserve the stack used to compile vLLM when resolving optional dependencies. \
|
||||
uv pip freeze --system > /tmp/kv-installed.txt; \
|
||||
grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
|
||||
/tmp/kv-installed.txt > /tmp/kv-constraints.txt; \
|
||||
uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build || ( \
|
||||
# if the above fails, install from source
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --allow-change-held-packages ${BUILD_PKGS} && \
|
||||
uv pip install --system -r /tmp/kv_connectors.txt --no-build-isolation && \
|
||||
uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build-isolation && \
|
||||
apt-get purge -y ${BUILD_PKGS} && \
|
||||
# clean up -dev packages, keep runtime libraries
|
||||
rm -rf /var/lib/apt/lists/* \
|
||||
@@ -1250,8 +1297,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \
|
||||
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \
|
||||
uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \
|
||||
uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
|
||||
uv pip install --system -c /tmp/kv-constraints.txt "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
|
||||
fi; \
|
||||
rm -f /tmp/kv-installed.txt /tmp/kv-constraints.txt; \
|
||||
fi
|
||||
|
||||
# Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts
|
||||
|
||||
@@ -140,6 +140,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
|
||||
cd triton-cpu; \
|
||||
git checkout "270e696d"; \
|
||||
if [ ! -f third_party/sleef/CMakeLists.txt ]; then \
|
||||
echo "sleef submodule empty after shallow fetch; retrying unshallow"; \
|
||||
git submodule foreach --recursive 'git fetch --unshallow origin || true'; \
|
||||
git submodule update --init --recursive --force; \
|
||||
fi; \
|
||||
test -f third_party/sleef/CMakeLists.txt || { echo "FATAL: sleef submodule unavailable"; exit 1; }; \
|
||||
if [ "$USE_SCCACHE" = "1" ]; then \
|
||||
if [ "$SCCACHE_LOCAL_ONLY" = "1" ]; then \
|
||||
unset SCCACHE_BUCKET SCCACHE_REGION; \
|
||||
@@ -465,6 +471,10 @@ ADD ./benchmarks/ ./benchmarks/
|
||||
ADD ./vllm/collect_env.py .
|
||||
ADD ./docker/ ./docker/
|
||||
ADD ./.buildkite/ ./.buildkite/
|
||||
# tools/ carries repo-tooling code that a few tests exercise directly
|
||||
# (e.g. tests/tools/test_check_test_tethering.py imports
|
||||
# tools.pre_commit.check_test_tethering).
|
||||
ADD ./tools/ ./tools/
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -914,7 +914,7 @@ RUN case "$(which python3)" in \
|
||||
*) ;; esac
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system --upgrade huggingface-hub[cli]
|
||||
uv pip install --system --upgrade huggingface-hub
|
||||
|
||||
# Install vLLM using uv (inherited from base stage)
|
||||
# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
|
||||
|
||||
@@ -938,7 +938,7 @@ RUN case "$(which python3)" in \
|
||||
*) ;; esac
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system --upgrade huggingface-hub[cli]
|
||||
uv pip install --system --upgrade huggingface-hub
|
||||
|
||||
# Install vLLM using uv (inherited from base stage)
|
||||
# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
|
||||
|
||||
@@ -671,7 +671,7 @@ RUN case "$(which python3)" in \
|
||||
*) ;; esac
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system --upgrade huggingface-hub[cli]
|
||||
uv pip install --system --upgrade huggingface-hub
|
||||
|
||||
# Install vLLM using uv (inherited from base stage)
|
||||
# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 445 KiB After Width: | Height: | Size: 479 KiB |
@@ -22,6 +22,10 @@ To use the `torch.profiler` module, set the `profiler` entry to `'torch'` and `t
|
||||
- `torch_profiler_with_flops` to enable recording FLOPs, off by default
|
||||
- `torch_profiler_use_gzip` to control gzip-compressing profiling files, on by default
|
||||
- `torch_profiler_dump_cuda_time_total` to control dumping and printing the aggregated CUDA self time table, on by default
|
||||
- `torch_profiler_activities` to select worker activities. Defaults are
|
||||
platform-specific: `["CPU"]` on CPU, `["CPU", "CUDA"]` on NVIDIA GPUs, and
|
||||
`["CPU", "XPU"]` on Intel GPUs. Selecting `["CUDA"]` omits CPU annotations
|
||||
and the AsyncLLM CPU trace to reduce profiling overhead and trace size.
|
||||
|
||||
When using `vllm bench serve`, you can enable profiling by passing the `--profile` flag.
|
||||
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
|
||||
[**Kthena**](https://github.com/volcano-sh/kthena) is a Kubernetes-native LLM inference platform that transforms how organizations deploy and manage Large Language Models in production. Built with declarative model lifecycle management and intelligent request routing, it provides high performance and enterprise-grade scalability for LLM inference workloads.
|
||||
|
||||
This guide shows how to deploy a production-grade, **multi-node vLLM** service on Kubernetes.
|
||||
Kthena manages inference workloads through `ModelServing` and provides model-aware request routing through `ModelServer` and `ModelRoute`.
|
||||
|
||||
This guide shows how to deploy a **multi-node vLLM** service on Kubernetes using `ModelServing` and expose it through a Kubernetes Service.
|
||||
|
||||
We’ll:
|
||||
|
||||
@@ -18,7 +20,7 @@ You need:
|
||||
|
||||
- A Kubernetes cluster with **GPU nodes**.
|
||||
- `kubectl` access with cluster-admin or equivalent permissions.
|
||||
- **Volcano** installed for gang scheduling.
|
||||
- **Volcano 1.14 or later** for the role-level gang scheduling used in this example.
|
||||
- **Kthena** installed with the `ModelServing` CRD available.
|
||||
- A valid **Hugging Face token** if loading models from Hugging Face Hub.
|
||||
|
||||
@@ -30,21 +32,22 @@ helm repo update
|
||||
helm install volcano volcano-sh/volcano -n volcano-system --create-namespace
|
||||
```
|
||||
|
||||
This provides the gang-scheduling and network topology features used by Kthena.
|
||||
Volcano provides the gang scheduling used in this example. It is optional for Kthena deployments that do not use Volcano scheduling.
|
||||
|
||||
### 1.2 Install Kthena
|
||||
|
||||
```bash
|
||||
helm install kthena oci://ghcr.io/volcano-sh/charts/kthena --version v0.1.0 --namespace kthena-system --create-namespace
|
||||
helm install kthena oci://ghcr.io/volcano-sh/charts/kthena --version v1.0.0 --namespace kthena-system --create-namespace
|
||||
```
|
||||
|
||||
- The `kthena-system` namespace is created.
|
||||
- Kthena controllers and CRDs, including `ModelServing`, are installed and healthy.
|
||||
- Kthena controllers and CRDs, including `ModelServing`, are installed.
|
||||
|
||||
Validate:
|
||||
|
||||
```bash
|
||||
kubectl get crd | grep modelserving
|
||||
kubectl get pods -n kthena-system
|
||||
kubectl get crd modelservings.workload.serving.volcano.sh
|
||||
```
|
||||
|
||||
You should see:
|
||||
@@ -57,7 +60,7 @@ modelservings.workload.serving.volcano.sh ...
|
||||
|
||||
## 2. The Multi-Node vLLM `ModelServing` Example
|
||||
|
||||
Kthena provides an example manifest to deploy a **multi-node vLLM cluster running Llama**. Conceptually this is equivalent to the vLLM production stack Helm deployment, but expressed with `ModelServing`.
|
||||
Kthena provides an example manifest to deploy a **multi-node vLLM cluster running Llama**. Each role replica contains an entry pod and its worker pods.
|
||||
|
||||
A simplified version of the example (`llama-multinode`) looks like:
|
||||
|
||||
@@ -116,10 +119,9 @@ Commands:
|
||||
- -c
|
||||
- >
|
||||
bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh
|
||||
leader --ray_cluster_size=2; python3 -m
|
||||
vllm.entrypoints.launchers.api_server.entry --port 8080 --model
|
||||
meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8
|
||||
--pipeline-parallel-size 2
|
||||
leader --ray_cluster_size=2;
|
||||
vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8
|
||||
--pipeline-parallel-size 2 --distributed-executor-backend ray
|
||||
```
|
||||
|
||||
Worker:
|
||||
@@ -166,9 +168,9 @@ Save one of the following manifests to `modelserving.yaml`:
|
||||
restartGracePeriodSeconds: 60
|
||||
gangPolicy:
|
||||
minRoleReplicas:
|
||||
405b: 1
|
||||
llama-405b: 1
|
||||
roles:
|
||||
- name: 405b
|
||||
- name: llama-405b
|
||||
replicas: 2
|
||||
entryTemplate:
|
||||
spec:
|
||||
@@ -256,9 +258,9 @@ Save one of the following manifests to `modelserving.yaml`:
|
||||
restartGracePeriodSeconds: 60
|
||||
gangPolicy:
|
||||
minRoleReplicas:
|
||||
405b: 1
|
||||
llama-405b: 1
|
||||
roles:
|
||||
- name: 405b
|
||||
- name: llama-405b
|
||||
replicas: 2
|
||||
entryTemplate:
|
||||
spec:
|
||||
@@ -275,7 +277,7 @@ Save one of the following manifests to `modelserving.yaml`:
|
||||
- sh
|
||||
- -c
|
||||
- "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2;
|
||||
vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2"
|
||||
vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2 --distributed-executor-backend ray"
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: "8"
|
||||
@@ -383,24 +385,24 @@ kubectl get pod -owide -l modelserving.volcano.sh/name=llama-multinode
|
||||
Example output (from docs):
|
||||
|
||||
```text
|
||||
NAMESPACE NAME READY STATUS RESTARTS AGE IP NODE ...
|
||||
default llama-multinode-0-405b-0-0 1/1 Running 0 15m 10.244.0.56 192.168.5.12 ...
|
||||
default llama-multinode-0-405b-0-1 1/1 Running 0 15m 10.244.0.58 192.168.5.43 ...
|
||||
default llama-multinode-0-405b-1-0 1/1 Running 0 15m 10.244.0.57 192.168.5.58 ...
|
||||
default llama-multinode-0-405b-1-1 1/1 Running 0 15m 10.244.0.53 192.168.5.36 ...
|
||||
NAMESPACE NAME READY STATUS RESTARTS AGE IP NODE ...
|
||||
default llama-multinode-0-llama-405b-0-0 1/1 Running 0 15m 10.244.0.56 192.168.5.12 ...
|
||||
default llama-multinode-0-llama-405b-0-1 1/1 Running 0 15m 10.244.0.58 192.168.5.43 ...
|
||||
default llama-multinode-0-llama-405b-1-0 1/1 Running 0 15m 10.244.0.57 192.168.5.58 ...
|
||||
default llama-multinode-0-llama-405b-1-1 1/1 Running 0 15m 10.244.0.53 192.168.5.36 ...
|
||||
```
|
||||
|
||||
Pod name pattern:
|
||||
|
||||
- `llama-multinode-<group-idx>-<role-name>-<replica-idx>-<ordinal>`.
|
||||
|
||||
The first number indicates `ServingGroup`. The second (`405b`) is the `Role`. The remaining indices identify the pod within the role.
|
||||
The first index identifies the `ServingGroup`, followed by the role name (`llama-405b`), role replica index, and pod index.
|
||||
|
||||
---
|
||||
|
||||
## 6. Accessing the vLLM OpenAI-Compatible API
|
||||
## 5. Accessing the vLLM OpenAI-Compatible API
|
||||
|
||||
Expose the entry via a Service:
|
||||
Save the following Service as `service.yaml` to expose the entry pods:
|
||||
|
||||
```yaml
|
||||
apiVersion: v1
|
||||
@@ -412,7 +414,6 @@ spec:
|
||||
selector:
|
||||
modelserving.volcano.sh/name: llama-multinode
|
||||
modelserving.volcano.sh/entry: "true"
|
||||
# optionally further narrow to leader role if you label it
|
||||
ports:
|
||||
- name: http
|
||||
port: 80
|
||||
@@ -420,9 +421,10 @@ spec:
|
||||
type: ClusterIP
|
||||
```
|
||||
|
||||
Port-forward from your local machine:
|
||||
Apply the Service, then port-forward from your local machine:
|
||||
|
||||
```bash
|
||||
kubectl apply -f service.yaml
|
||||
kubectl port-forward svc/llama-multinode-openai 30080:80 -n default
|
||||
```
|
||||
|
||||
@@ -450,11 +452,12 @@ You should see an OpenAI-style response from vLLM.
|
||||
|
||||
---
|
||||
|
||||
## 7. Clean Up
|
||||
## 6. Clean Up
|
||||
|
||||
To remove the deployment and its resources:
|
||||
|
||||
```bash
|
||||
kubectl delete -f service.yaml
|
||||
kubectl delete modelserving llama-multinode -n default
|
||||
```
|
||||
|
||||
@@ -464,3 +467,5 @@ If you’re done with the entire stack:
|
||||
helm uninstall kthena -n kthena-system # or your Kthena release name
|
||||
helm uninstall volcano -n volcano-system
|
||||
```
|
||||
|
||||
For model-aware routing and prefill-decode disaggregation, see the [Kthena documentation](https://kthena.volcano.sh/docs/user-guide/router-routing).
|
||||
|
||||
@@ -15,7 +15,7 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl
|
||||
|
||||
## Hybrid Mamba models
|
||||
|
||||
Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-fine-grained-prefix-cache` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it.
|
||||
Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it.
|
||||
|
||||
This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold:
|
||||
|
||||
@@ -28,7 +28,7 @@ This helps when many requests share a long system prompt and then diverge. It is
|
||||
vllm serve <hybrid-model> \
|
||||
--mamba-cache-mode align \
|
||||
--prefix-match-unit 64 \
|
||||
--enable-mamba-fine-grained-prefix-cache
|
||||
--enable-mamba-shared-prefix-checkpoint
|
||||
```
|
||||
|
||||
`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect.
|
||||
|
||||
@@ -32,8 +32,9 @@ Design doc: <https://docs.google.com/document/d/1aed8KtC6XkXtdoV87pWT0a8OJlZ-Cpn
|
||||
|
||||
## 2 Usage Example
|
||||
|
||||
The current reference pathway is **ExampleConnector**.
|
||||
Below ready-to-run scripts shows the workflow:
|
||||
### ExampleConnector
|
||||
|
||||
The following scripts demonstrate the workflow with **ExampleConnector**:
|
||||
|
||||
1 Encoder instance + 1 PD instance:
|
||||
`examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh`
|
||||
@@ -41,6 +42,25 @@ Below ready-to-run scripts shows the workflow:
|
||||
1 Encoder instance + 1 Prefill instance + 1 Decode instance:
|
||||
`examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh`
|
||||
|
||||
### ECMooncakeConnector
|
||||
|
||||
**ECMooncakeConnector** transfers encoder outputs using the Mooncake TransferEngine.
|
||||
See the [Mooncake integration example](../../tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh)
|
||||
for a complete 1 Encoder instance + 1 PD instance setup, including the producer and
|
||||
consumer `--ec-transfer-config` settings and proxy configuration.
|
||||
|
||||
With vLLM and Mooncake installed, run the example from the repository root:
|
||||
|
||||
```bash
|
||||
GPU_E=0 GPU_PD=1 MOONCAKE_EC_PROTOCOL=tcp \
|
||||
bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
|
||||
```
|
||||
|
||||
The script uses `Qwen/Qwen2.5-VL-3B-Instruct` by default, runs a single-GPU baseline,
|
||||
and compares the disaggregated outputs against it. Set `MODEL` to use another model
|
||||
or `MOONCAKE_EC_PROTOCOL=rdma` to use RDMA on supported hardware. See the script for
|
||||
additional configuration options.
|
||||
|
||||
---
|
||||
|
||||
## 3 Test Script
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
# Engram: conditional memory via n-gram lookups
|
||||
|
||||
Engram is a *conditional memory* module attached to specific layers of a
|
||||
transformer backbone. At each position, the preceding few tokens are hashed
|
||||
into large static embedding tables, and the retrieved rows are fused into the
|
||||
hidden state, gated by how well they match it. Where Mixture-of-Experts scales
|
||||
capacity through conditional computation, Engram scales it through conditional
|
||||
memory: parameters that are never computed over, only looked up, so capacity
|
||||
grows at O(1) per-token cost. The mechanism and its design rationale are
|
||||
described in [Conditional Memory via Scalable Lookup: A New Axis of Sparsity
|
||||
for Large Language Models](https://arxiv.org/abs/2601.07372).
|
||||
|
||||
vLLM enables Engram automatically whenever the checkpoint declares Engram
|
||||
layers, so basic serving needs no extra flags. This page explains what Engram
|
||||
does and how to control its memory footprint and multi-GPU topology.
|
||||
|
||||
## Supported models
|
||||
|
||||
Engram layers are declared by the checkpoint, not chosen at serve time:
|
||||
|
||||
| Architecture | HF config field | Notes |
|
||||
| --- | --- | --- |
|
||||
| `DeepseekV41ForCausalLM` | `engram_layer_ids` | e.g. `deepseek-ai/DeepSeek-V4.1-Flash` |
|
||||
| `Qwen4ExpForCausalLM` / `Qwen4ExpForConditionalGeneration` | `ple_layer_ids` | same mechanism, called PLE |
|
||||
|
||||
Passing `--engram-config` for a model without n-gram embedding layers fails at
|
||||
startup, as does running on a non-CUDA-alike platform.
|
||||
|
||||
## How it works
|
||||
|
||||
Engram processes each position in two phases: *retrieval* and *fusion*.
|
||||
|
||||
**Retrieval** maps the local token context to static table rows:
|
||||
|
||||
1. **Tokenizer compression.** At startup, vLLM builds a many-to-one map from
|
||||
the tokenizer that collapses tokens which normalize alike (NFKC, accent
|
||||
stripping, lowercasing, whitespace collapsing), so `" The"`, `"the"` and
|
||||
`"THE"` hash identically. This concentrates semantic density into a much
|
||||
smaller id space; its size must match `engram_compressed_vocab_size` from
|
||||
the checkpoint, or startup fails — every hash multiplier derives from it.
|
||||
2. **Multi-head hashing.** For each n-gram order `n` (up to
|
||||
`engram_max_ngram_size - 1`), several hash heads each apply an independent
|
||||
multiplicative-XOR hash over the compressed suffix n-gram. Every (n-gram
|
||||
order, head) pair indexes its own prime-sized bucket range in the layer's
|
||||
table — disjoint across pairs, so collisions cannot collide twice — and
|
||||
the retrieved rows are concatenated into one memory vector per position.
|
||||
|
||||
**Fusion** turns the retrieved rows into a context-aware contribution to the
|
||||
hidden state. The rows are context-independent priors and may be noisy (hash
|
||||
collisions, polysemy), so they are gated: a `wkv` projection produces one key
|
||||
per hyper-connection branch plus a shared value, and each branch computes a
|
||||
sigmoid gate over the normalized dot product between its hidden state —
|
||||
which has already aggregated global context through preceding attention —
|
||||
and its key. If the retrieved memory contradicts the current context, the
|
||||
gate tends toward zero and suppresses it. The gated value is added
|
||||
residually into the branch's stream (`hidden + gate * value`); positions that
|
||||
take no part in an n-gram (e.g. image spans) are masked so they pass through
|
||||
untouched.
|
||||
|
||||
Engram sits only on the layers listed in the checkpoint. Placement is a
|
||||
deliberate hardware-algorithm trade-off: early enough that the backbone is
|
||||
relieved from reconstructing static patterns in its first layers, deep enough
|
||||
that vLLM can hide table lookups behind the compute of the preceding layers
|
||||
(see below).
|
||||
|
||||
In vLLM, hashing, lookup, and gating run as fused Triton kernels over FP8
|
||||
rows with ue8m0 block scales. The layer layout — table sizes, n-gram size,
|
||||
head count — is fixed by the checkpoint and is not user-configurable. Only the
|
||||
storage and sharding described below are configurable, via
|
||||
`--engram-config`, which takes a JSON object and can also be set field by
|
||||
field as `--engram-config.<field>`
|
||||
([CLI reference](../cli/serve.md#-engram-config)).
|
||||
|
||||
## CPU offload
|
||||
|
||||
Engram tables dwarf typical embedding weights: [DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash)
|
||||
carries two Engram layers of roughly 384M table rows each, with 256-byte FP8
|
||||
entries plus their ue8m0 block scales — about 200 GB of table weights alone.
|
||||
Because lookups are deterministic — the indices are known from the token ids
|
||||
alone, before the forward pass — the tables need not sit in GPU memory at
|
||||
all. By default vLLM stores them in **pinned host memory** and serves lookups
|
||||
directly from the Triton kernel through unified virtual addressing (UVA),
|
||||
prefetching on a side CUDA stream so transfers overlap with the compute of
|
||||
preceding layers. The freed GPU memory goes to the KV cache.
|
||||
|
||||
`cpu_offload` defaults to `true` (on): it follows the
|
||||
`VLLM_PLE_CPU_OFFLOAD` environment variable, which also defaults to `true`.
|
||||
An explicit `--engram-config` value takes precedence over the environment
|
||||
variable.
|
||||
|
||||
To keep the tables resident on GPU instead (e.g. when host memory is
|
||||
scarce):
|
||||
|
||||
```bash
|
||||
vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
|
||||
--engram-config.cpu_offload false
|
||||
```
|
||||
|
||||
or equivalently `VLLM_PLE_CPU_OFFLOAD=0`. CPU offload requires a GPU with UVA
|
||||
support; vLLM fails fast when it is unavailable.
|
||||
|
||||
## Data-parallel topologies (DeepSeek V4.1)
|
||||
|
||||
By default (`embedding_across_dp: false`), each DP replica keeps its own
|
||||
TP-sharded copy of the tables — no cross-replica communication, but one host
|
||||
copy per replica.
|
||||
|
||||
With `embedding_across_dp: true`, the hash heads are sharded across all
|
||||
TP × DP ranks into a single table copy. Each step gathers the hash ids of
|
||||
every co-located DP replica, each rank looks up the heads it owns, and the
|
||||
rows are exchanged back. This trades per-step DP collectives for a much
|
||||
smaller table footprint:
|
||||
|
||||
```bash
|
||||
vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
|
||||
--tensor-parallel-size 2 --data-parallel-size 4 \
|
||||
--engram-config '{"embedding_across_dp": true}'
|
||||
```
|
||||
|
||||
`embedding_across_dp` is not supported with elastic expert parallelism yet.
|
||||
Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only.
|
||||
|
||||
## Sharing host tables across DP replicas
|
||||
|
||||
When the tables are CPU-offloaded and DP replicas are co-located on one node,
|
||||
`dp_shared_memory` stores **one copy of each TP shard in `/dev/shm`**, shared
|
||||
by all co-located replicas through `cudaHostRegister`. Each replica then
|
||||
prefetches only its own tokens — no per-step Engram DP collectives — and host
|
||||
memory drops from one table copy per replica to one per node.
|
||||
|
||||
`dp_shared_memory` defaults to enabled whenever the other settings allow it,
|
||||
and falls back to per-replica (or DP-sharded) tables with a warning when:
|
||||
|
||||
- the DP replicas are not co-located on a single node,
|
||||
- `/dev/shm` cannot hold the full tables — in containers, raise the limit
|
||||
with `--shm-size` or `--ipc=host`,
|
||||
- expert parallelism is elastic, which is unsupported.
|
||||
|
||||
Requirements: `cpu_offload` enabled, `--data-parallel-size > 1`, and a shared
|
||||
IPC namespace across the co-located replicas.
|
||||
|
||||
## Limitations
|
||||
|
||||
- CUDA-alike platforms only.
|
||||
- DeepSeek V4.1 Engram does not support DBO or microbatching; disable
|
||||
`--enable-dbo` and set `--ubatch-size 0`.
|
||||
- `dp_shared_memory` requires `cpu_offload`, `data_parallel_size > 1`, and is
|
||||
unsupported with elastic expert parallelism.
|
||||
@@ -0,0 +1,152 @@
|
||||
# Initialized engine snapshots
|
||||
|
||||
Initialized engine snapshots are an experimental way to trade local disk space
|
||||
and host privileges for a faster vLLM activation. Snapshot creation initializes
|
||||
the engine and records deterministic generation output before CRIU captures the
|
||||
process tree and CUDA state. Restore validates the saved environment, restores
|
||||
the engine, binds the HTTP server, and reproduces the recorded token and sampled
|
||||
token log probability before returning.
|
||||
|
||||
This path is intended for repeatedly activating the same model and engine
|
||||
configuration on the same machine. It is not a portable model artifact.
|
||||
|
||||
## Requirements
|
||||
|
||||
Snapshots currently require:
|
||||
|
||||
- Linux on x86-64 with one NVIDIA GPU.
|
||||
- TP1 with one unauthenticated plaintext HTTP server. Other parallel sizes,
|
||||
TLS, middleware, Unix sockets, and speculative decoding are unsupported.
|
||||
- Because snapshot mode requires an unauthenticated plaintext HTTP server, use
|
||||
a trusted host or network boundary, or external controls. See the vLLM
|
||||
[Security guide](../usage/security.md).
|
||||
- [CRIU](https://github.com/checkpoint-restore/criu), its CUDA plugin, a
|
||||
`cuda-checkpoint`-compatible helper, and `nvidia-smi` on `PATH`.
|
||||
- Root or passwordless `sudo` for CRIU.
|
||||
- `io_uring` disabled before launch because CRIU cannot dump it. Use
|
||||
`kernel.io_uring_disabled=1` for an unprivileged process. A process running
|
||||
as root, including the `docker exec` flow below, bypasses `=1`, so use `=2`
|
||||
host-wide.
|
||||
- No established TCP connection to a peer outside the captured process tree.
|
||||
Current Hugging Face hub clients hold their connections for the process
|
||||
lifetime, so download the model in a separate step and run create with
|
||||
`HF_HUB_OFFLINE=1`, as the quickstart below does.
|
||||
- A remote model ID and an immutable 40-character `--revision`. Local model
|
||||
directories and mutable revisions are not supported.
|
||||
- Enough disk for the artifact, with the same installed vLLM package, model
|
||||
files, container filesystem, and generated-cache paths available at restore.
|
||||
Generated-cache files are not copied into the artifact. The manifest
|
||||
fingerprints the ones the captured tree holds open, so restore fails early
|
||||
and names the file when one is removed or replaced. Only open descriptors
|
||||
are recorded. A library the tree mapped and then closed is not fingerprinted,
|
||||
and CRIU still reopens it by path, so replacing that file permanently
|
||||
invalidates the artifact without an early error.
|
||||
|
||||
The official `vllm/vllm-openai` Linux x86-64 image includes the snapshot
|
||||
runtime. It still requires a compatible host driver, kernel, and privileges.
|
||||
Arm64 images omit it. Source installs must set `CRIU_CUDA_PLUGIN_DIR` to the
|
||||
directory containing `cuda_plugin.so`.
|
||||
|
||||
Run snapshot commands with `docker exec` inside a long-lived container. Restore
|
||||
hands the API server off as a detached process, so a one-shot container would
|
||||
stop that server when its PID 1 exits. Snapshot preflight requires every
|
||||
component of the artifact path to be owned by the invoking user or root, with
|
||||
the directory itself at mode 0700, and the commands below run as root inside
|
||||
the official image, so the bind-mounted host directory is created with `sudo`
|
||||
and root ownership. The model downloads in its own step so the captured tree
|
||||
holds no hub connection, and create runs offline. This example also keeps the
|
||||
container filesystem and `/dev/shm` namespace stable for the lifetime of the
|
||||
artifact:
|
||||
|
||||
```bash
|
||||
sudo sysctl kernel.io_uring_disabled=2
|
||||
|
||||
snapshot_root="$(pwd)/vllm-snapshots"
|
||||
sudo install -d -m 0700 -o root -g root "${snapshot_root}"
|
||||
|
||||
docker run --detach --name vllm-snapshot \
|
||||
--gpus all \
|
||||
--privileged \
|
||||
--pid=host \
|
||||
--ipc=host \
|
||||
--network=host \
|
||||
--mount "type=bind,source=${snapshot_root},target=/snapshots" \
|
||||
--entrypoint sleep \
|
||||
vllm/vllm-openai:latest infinity
|
||||
|
||||
docker exec vllm-snapshot hf download Qwen/Qwen3-0.6B \
|
||||
--revision c1899de289a04d12100db370d81485cdf75e47ca
|
||||
|
||||
docker exec -e HF_HUB_OFFLINE=1 vllm-snapshot vllm snapshot create Qwen/Qwen3-0.6B \
|
||||
--snapshot-dir /snapshots/qwen3-0.6b \
|
||||
--revision c1899de289a04d12100db370d81485cdf75e47ca \
|
||||
--dtype float16 \
|
||||
--max-model-len 512
|
||||
|
||||
docker exec vllm-snapshot vllm snapshot inspect /snapshots/qwen3-0.6b
|
||||
|
||||
docker exec vllm-snapshot vllm snapshot restore \
|
||||
/snapshots/qwen3-0.6b --host 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
Keep that container and its mounts available while the snapshot is in use.
|
||||
Stop and remove it only after the restored API server is no longer needed.
|
||||
|
||||
Create initializes the engine, records a one-token canary, releases and reloads
|
||||
weights and KV cache to rehearse restore, then releases them again for capture.
|
||||
The manifest is published only after CRIU completes and the source tree stops.
|
||||
Creation is offline preparation and is not part of restore latency.
|
||||
|
||||
The private `0700` artifact contains process memory, CUDA state, engine
|
||||
arguments, compatibility identity, and canary output; its manifest is `0600`.
|
||||
Literal API keys and Hugging Face tokens are redacted from the manifest's engine
|
||||
arguments. The manifest records selected environment names plus deterministic
|
||||
name-bound fingerprints, not their values. The protected CRIU artifact can
|
||||
still contain process secrets, so treat it as sensitive data. Restore reloads
|
||||
model files and KV cache before binding HTTP, then reproduces the canary or
|
||||
tears down the restored tree. The inspect command prints that identity and
|
||||
canary without executing the saved process.
|
||||
|
||||
## Restore behavior
|
||||
|
||||
Restore fails before CRIU runs if the saved identity does not match the current
|
||||
host. It does not silently fall back to ordinary startup. After CRIU restores
|
||||
the process tree, vLLM releases the saved engine to bind the requested HTTP
|
||||
address and checks the first generated token and sampled-token log probability
|
||||
against the snapshot canary. The command returns only after that check passes.
|
||||
The restored API server continues to run as a detached process.
|
||||
|
||||
The pre-release port probe is best-effort only. It neither reserves the port
|
||||
nor authenticates the listener that appears afterward.
|
||||
|
||||
Rollback terminates and waits for the restored tree after process identity
|
||||
verification succeeds. If verification fails, vLLM writes an abort marker for
|
||||
the snapshot server instead of signaling unverified PIDs. If that server has
|
||||
already exited, surviving engine processes can require operator cleanup.
|
||||
Verify that the previous tree has stopped before retrying; do not kill a process
|
||||
solely because its PID appears in the manifest, since PIDs can be reused.
|
||||
|
||||
The artifact is reusable after its previous restored tree stops. Only one
|
||||
snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time.
|
||||
|
||||
## Tradeoffs and limitations
|
||||
|
||||
- Creation has its own latency and briefly requires the full engine. Artifact
|
||||
size can approach the captured process and GPU memory.
|
||||
- Artifacts are not guaranteed to survive a power loss. The manifest is
|
||||
fsynced, but publication does not explicitly flush all captured payload files
|
||||
and directories. Recreate the artifact after an unclean host shutdown.
|
||||
- Restore currently requires the same host, GPU, driver, kernel, Python,
|
||||
PyTorch, installed vLLM version, model revision, engine arguments,
|
||||
selected environment variables, and CRIU plugin binaries.
|
||||
- Only dense float16 TP1 has been validated. Other model formats depend on their
|
||||
existing sleep level 2 reload support; NCCL state is not restored.
|
||||
- CRIU support varies by kernel and driver. Preserve package, library, model,
|
||||
and generated-cache paths for the artifact lifetime.
|
||||
- A snapshot can include application secrets or request state present in the
|
||||
process. Build it before traffic and protect it like process memory.
|
||||
- Any feature that opens an external connection must close it before capture.
|
||||
|
||||
For a lower-complexity option that retains a live process, see
|
||||
[Sleep mode](sleep_mode.md). Sleep mode and initialized snapshots retain
|
||||
different amounts of state and have different idle resource costs.
|
||||
@@ -74,6 +74,21 @@ llm.wake_up(tags=["weights"])
|
||||
llm.wake_up(tags=["kv_cache"])
|
||||
```
|
||||
|
||||
#### Retaining frozen weights during RLHF updates
|
||||
|
||||
Set `sleep_preserve_parameter_names` (CLI: `--sleep-preserve-parameter-names`)
|
||||
to runtime parameter-name glob
|
||||
patterns for weights that stay frozen during training. Each pattern must match
|
||||
`model.named_parameters()`; checkpoint names and `requires_grad` are not used.
|
||||
|
||||
Level-2 sleep backs up selected GPU parameters to pageable CPU memory; weights
|
||||
wake-up restores them in place and releases the backups. CPU parameters need no
|
||||
copy. Allow enough host memory for backups. Level-1 behavior is unchanged.
|
||||
|
||||
The trainer must omit these parameters from updates, and reload/post-processing
|
||||
must preserve their values and storage (including avoiding replacement with meta
|
||||
tensors). This option does not filter updates and applies only to the target model.
|
||||
|
||||
#### Release only KV cache memory
|
||||
|
||||
`LLM.release_kv_cache_memory()` discards KV cache physical memory while keeping model weights resident. It requires a completed pause and all executor memory to be resident: full sleep, partial wake-up, and repeated release without restoring memory are rejected. Requests retained with `mode="keep"` are recomputed after wake-up.
|
||||
|
||||
@@ -324,7 +324,7 @@ You can find more information about vLLM's wheels in
|
||||
# Install dependencies
|
||||
pip install --upgrade numba \
|
||||
scipy \
|
||||
huggingface-hub[cli] \
|
||||
huggingface-hub \
|
||||
setuptools_scm
|
||||
pip install -r requirements/rocm.txt
|
||||
|
||||
|
||||
@@ -240,7 +240,7 @@ Alternatively, you can [open an issue on GitHub](https://github.com/vllm-project
|
||||
|
||||
#### Download a model
|
||||
|
||||
If you prefer, you can use the Hugging Face CLI to [download a model](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-download) or specific files from a model repository:
|
||||
If you prefer, you can use the Hugging Face CLI to [download a model](https://huggingface.co/docs/huggingface_hub/guides/cli#hf-download) or specific files from a model repository:
|
||||
|
||||
```bash
|
||||
# Download a model
|
||||
|
||||
@@ -126,6 +126,14 @@ vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \
|
||||
--data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
|
||||
```
|
||||
|
||||
When using `--data-parallel-external-lb` with `--nnodes`, the data parallel rank
|
||||
can be inferred from `--node-rank` only when `--nnodes` is evenly divisible by
|
||||
`--data-parallel-size`. If it is not evenly divisible, automatic rank inference
|
||||
is rejected and `--data-parallel-rank` must be set explicitly. When multiple
|
||||
external-LB ranks are co-located on a node, assign each process a disjoint
|
||||
device set with `CUDA_VISIBLE_DEVICES` or `--device-ids`; devices are not
|
||||
partitioned automatically between processes.
|
||||
|
||||
The coordinator process also runs in this scenario, co-located with the DP rank 0 engine.
|
||||
|
||||
<figure markdown="1">
|
||||
|
||||
@@ -6,7 +6,7 @@ This closes the loop for a token-in / token-out engine in disaggregated serving:
|
||||
|
||||
- **GPU less post processing**: Detokenization, reasoning parsing, and tool call parsing run on the same GPU less frontend that hosts `/render`
|
||||
- **Parser parity**: The derenderer reuses vLLM's tool and reasoning parsers, so a disaggregated deployment produces the same `content`/`reasoning`/ `tool_calls` split as a standard `vllm serve` server
|
||||
- **Streaming and one-shot parsing**: Non-streaming calls send a complete `GenerateResponse` with all token IDs present and get one-shot parsing. Both endpoints also accept `stream: true`, taking one `GenerateStreamResponse` delta plus a client carried `stream_state` and returning `{chunk, stream_state}`. The chat endpoint's streaming path supports reasoning and tool call parsing while emitting the same `reasoning`/`content`/`tool_calls` deltas the generate streaming path would for the same token IDs
|
||||
- **Streaming and one-shot parsing**: Non-streaming calls send a complete `GenerateResponse` with all token IDs present and get one-shot parsing. Both endpoints also accept `stream: true`, taking one `GenerateStreamResponse` delta plus a client carried `stream_state` and returning `{chunk, stream_state}`. The chat endpoint's streaming path supports reasoning and tool call parsing while emitting the same `reasoning`/`content`/`tool_calls` deltas the generate streaming path would for the same token IDs (see [Streaming](#streaming))
|
||||
|
||||
Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render`
|
||||
endpoints.
|
||||
@@ -27,8 +27,10 @@ The derender step needs more than the engine's `token_ids`. It also consumes the
|
||||
|
||||
- Chat Completions Derender API (`/v1/chat/completions/derender`)
|
||||
- Post process a single `GenerateResponse` into a `ChatCompletionResponse`
|
||||
- With `stream: true`, post process one `GenerateStreamResponse` chunk into a `ChatCompletionStreamResponse` chunk
|
||||
- Completions Derender API (`/v1/completions/derender`)
|
||||
- Post process a list of `GenerateResponse` objects (one per prompt) into a `CompletionResponse`
|
||||
- With `stream: true`, post process one `GenerateStreamResponse` chunk into a `CompletionStreamResponse` chunk
|
||||
|
||||
## Request format
|
||||
|
||||
@@ -50,8 +52,54 @@ Each request wraps the engine's `GenerateResponse`(s) together with the caller m
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-request"
|
||||
```
|
||||
|
||||
Streaming requests set `stream: true` and carry one generate chunk (`generate_chunk`) plus `stream_state` instead of a complete response.
|
||||
|
||||
`/v1/chat/completions/derender` with `stream: true`:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-chat-stream-request"
|
||||
```
|
||||
|
||||
`/v1/completions/derender` with `stream: true`:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-stream-request"
|
||||
```
|
||||
|
||||
Both return `{"chunk": ..., "stream_state": ...}`. `chunk` is a `ChatCompletionStreamResponse` or `CompletionStreamResponse` and `stream_state` goes with the next call.
|
||||
|
||||
Oversized payloads are rejected with a `400` before any `tokenizer.decode()` or parser runs.
|
||||
|
||||
## Parser configuration
|
||||
|
||||
The derenderer builds its tool and reasoning parsers from its own server flags plus the `chat_request` sent with each call. It can't see how the request was rendered or served anywhere else, so for its output to match a `vllm serve` server:
|
||||
|
||||
- Start the render server with the same `--tool-call-parser`, `--reasoning-parser`, `--enable-auto-tool-choice`, `--chat-template` and `--default-chat-template-kwargs` you'd give `vllm serve` for this model. If `/render` and `/derender` run on different servers, give both the same values.
|
||||
- Send the full `chat_request` that went to `/render`, not just `messages` and `tools`. Fields like `chat_template_kwargs`, `reasoning_effort`, `tool_choice` and `include_reasoning` change how the output is parsed.
|
||||
|
||||
A mismatch doesn't fail. The parser just splits `reasoning`, `content` and `tool_calls` differently from what `vllm serve` would return.
|
||||
|
||||
`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. It only reads `skip_special_tokens` from `completion_request`.
|
||||
|
||||
## Streaming
|
||||
|
||||
Streaming derender mirrors `/inference/v1/generate`. Set `stream: true` in the body on the same path and the endpoint takes one generate stream chunk instead of a complete response. It answers with JSON, not SSE. Each call turns one generate chunk into one derendered chunk and the client forwards that to its own caller.
|
||||
|
||||
The server keeps no state between calls. Everything the next call needs is in `stream_state` which the client carries:
|
||||
|
||||
- **Echo `stream_state` back.** Leave it out (or send `null`) on the first call, then send the `stream_state` from each response with the next request, unchanged. State that doesn't add up (e.g. `output_chunk_lens` that don't sum to the number of output tokens) is rejected with a `400`.
|
||||
- **One choice per chunk.** `/inference/v1/generate` sends one choice per SSE event and each derender call accepts at most one. More than one is rejected with a `400`.
|
||||
- **`n > 1` is N streams.** Generate interleaves chunks for different choices and tags each with its `index`. Keep one `stream_state` per index and send each chunk with the state for its index. Each index gets its own `role` delta on its first chunk.
|
||||
- **Send every chunk through, including the last ones.** The finish reason usually arrives with the final tokens. A chunk with a `finish_reason` and no tokens is accepted too and flushes any buffered tool call arguments. With `stream_options: {"include_usage": true}`, generate ends with a usage only chunk (`choices: []`). Send it through derender like any other chunk to get the usage chunk for the response. `usage.prompt_tokens` is the request's `prompt_tokens` if set, otherwise the generate chunk's.
|
||||
- **Send the same context on every call.** `chat_request` and `prompt_token_ids` aren't kept between calls, so they go with every chunk including the usage chunk. Both are required when a tool or reasoning parser is configured. `prompt_token_ids` is the `token_ids` of the `GenerateRequest` returned by `/render`.
|
||||
- **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk.
|
||||
|
||||
Streaming chunks don't carry logprobs yet, so `logprobs` on a generate chunk are dropped. The non-streaming endpoints do resolve them, including `token_id:N` placeholders.
|
||||
|
||||
## Streaming cost
|
||||
|
||||
Streaming derender threads a client carried `stream_state` across per-chunk calls instead of keeping session state on the server.
|
||||
@@ -118,3 +166,64 @@ print(response["choices"][0]["message"]["content"])
|
||||
```
|
||||
|
||||
Passing `chat_request` lets the derenderer run the configured tool and reasoning parsers. This means `response["choices"][0]["message"]` carries the same `content` / `reasoning` / `tool_calls` split a `vllm serve` server would produce. `chat_request` can only be omitted for a model with no tool or reasoning parser configured. A parser configured model rejects a missing `chat_request` with a 400 rather than silently falling back to plain detokenization.
|
||||
|
||||
### Streaming example
|
||||
|
||||
The same round trip, streamed. `/render` keeps `stream` and `stream_options` on the `GenerateRequest` it returns, so the generate call streams too. Each generate chunk goes through derender with the `stream_state` from the previous call (see [Streaming](#streaming)).
|
||||
|
||||
```python
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
MODEL = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
RENDER = "http://localhost:8100" # vllm launch render ...
|
||||
ENGINE = "http://localhost:8200" # token-in / token-out engine
|
||||
|
||||
chat_request = {
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "What is 2+2?"}],
|
||||
"max_tokens": 32,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
|
||||
with httpx.Client(timeout=60.0) as client:
|
||||
# 1. Render: request -> token IDs (GPU less)
|
||||
generate_request = client.post(
|
||||
f"{RENDER}/v1/chat/completions/render", json=chat_request
|
||||
).json()
|
||||
prompt_token_ids = generate_request["token_ids"]
|
||||
|
||||
# 2. Generate: stream token IDs (token-in / token-out engine)
|
||||
stream_state = None # one per choice index when n > 1
|
||||
with client.stream(
|
||||
"POST", f"{ENGINE}/inference/v1/generate", json=generate_request
|
||||
) as generate_stream:
|
||||
for line in generate_stream.iter_lines():
|
||||
if not line.startswith("data: ") or line == "data: [DONE]":
|
||||
continue
|
||||
|
||||
# 3. Derender: one generate chunk -> one chat completion chunk
|
||||
derendered = client.post(
|
||||
f"{RENDER}/v1/chat/completions/derender",
|
||||
json={
|
||||
"stream": True,
|
||||
"model": MODEL,
|
||||
"generate_chunk": json.loads(line[len("data: ") :]),
|
||||
"stream_state": stream_state,
|
||||
"prompt_tokens": len(prompt_token_ids),
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"chat_request": chat_request,
|
||||
},
|
||||
).json()
|
||||
stream_state = derendered["stream_state"]
|
||||
|
||||
chunk = derendered["chunk"]
|
||||
for choice in chunk["choices"]:
|
||||
print(choice["delta"].get("content") or "", end="", flush=True)
|
||||
if chunk.get("usage"):
|
||||
print(f"\n{chunk['usage']}")
|
||||
```
|
||||
|
||||
With a tool or reasoning parser configured, `delta` also carries `reasoning` and `tool_calls`, the same deltas `/v1/chat/completions` streams for those tokens.
|
||||
|
||||
@@ -8,7 +8,7 @@ This document outlines some troubleshooting strategies you can consider. If you
|
||||
## Hangs downloading a model
|
||||
|
||||
If the model isn't already downloaded to disk, vLLM will download it from the internet which can take time and depend on your internet connection.
|
||||
It's recommended to download the model first using the [huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli) and passing the local path to the model to vLLM. This way, you can isolate the issue.
|
||||
It's recommended to download the model first using the [`hf` CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli) and passing the local path to the model to vLLM. This way, you can isolate the issue.
|
||||
|
||||
## Hangs loading a model from disk
|
||||
|
||||
|
||||
@@ -367,6 +367,8 @@ async def handle_request(api: str, request: Request):
|
||||
"remote_block_ids"
|
||||
]
|
||||
req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]
|
||||
if prefill_kv.get("tp_size") is not None:
|
||||
req_data["kv_transfer_params"]["tp_size"] = prefill_kv["tp_size"]
|
||||
|
||||
req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
|
||||
"dp_size"
|
||||
|
||||
@@ -2252,6 +2252,7 @@ MODELS_NEED_VIDEO_METADATA = [
|
||||
"glm4_5v",
|
||||
"glm4_5v_fp8",
|
||||
"molmo2",
|
||||
"qwen2_5_vl",
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
"qwen3_5",
|
||||
|
||||
@@ -33,4 +33,4 @@ quack-kernels==0.6.5 # Required by MSA
|
||||
tokenspeed-mla==0.1.8; platform_system == "Linux"
|
||||
|
||||
# Humming kernels for quantization gemm
|
||||
humming-kernels[cu13]==0.1.12
|
||||
humming-kernels[cu13]==0.1.15
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# LMCache 0.5.x uses CUDA 13 dependencies; retain the CUDA 12 release here.
|
||||
lmcache == 0.4.4
|
||||
cupy-cuda12x != 14.1.0
|
||||
nixl == 1.4.1
|
||||
mooncake-transfer-engine >= 0.3.12
|
||||
Generated
+8
-1
@@ -2290,7 +2290,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
|
||||
[[package]]
|
||||
name = "llm-multimodal"
|
||||
version = "1.7.1"
|
||||
source = "git+https://github.com/smg-project/llm-multimodal?rev=24c676dd49d73f05cb4068aade2e60397a5c0afe#24c676dd49d73f05cb4068aade2e60397a5c0afe"
|
||||
source = "git+https://github.com/smg-project/llm-multimodal?rev=53ca0d6ea829af7477b6b63851cc70bed409bbc8#53ca0d6ea829af7477b6b63851cc70bed409bbc8"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"base64 0.22.1",
|
||||
@@ -5735,6 +5735,11 @@ dependencies = [
|
||||
"url",
|
||||
"uuid",
|
||||
"vllm-build-info",
|
||||
"vllm-chat",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-managed-engine",
|
||||
"vllm-text",
|
||||
"vllm-tracing",
|
||||
]
|
||||
|
||||
@@ -5784,6 +5789,7 @@ dependencies = [
|
||||
"vllm-parser",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
"vllm-tracing",
|
||||
"xgrammar-structural-tag",
|
||||
"zeromq",
|
||||
]
|
||||
@@ -5931,6 +5937,7 @@ dependencies = [
|
||||
"openai-protocol",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tool-parser",
|
||||
|
||||
+1
-1
@@ -61,7 +61,7 @@ indexmap = "2.13.0"
|
||||
indicatif = "0.18.4"
|
||||
itertools = "0.14.0"
|
||||
libc = "0.2.177"
|
||||
llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "24c676dd49d73f05cb4068aade2e60397a5c0afe", default-features = false, features = ["native-tls"] }
|
||||
llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "53ca0d6ea829af7477b6b63851cc70bed409bbc8", default-features = false, features = ["native-tls"] }
|
||||
mimalloc = "0.1.52"
|
||||
minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
|
||||
minijinja-contrib = { version = "2.24", features = ["pycompat"] }
|
||||
|
||||
@@ -56,6 +56,21 @@ message GenerateRequest {
|
||||
|
||||
// Apply the engine's configured watermark. Defaults to true when omitted.
|
||||
optional bool watermarking = 16;
|
||||
|
||||
KVHintsEnvelope kv_hints = 17;
|
||||
}
|
||||
|
||||
message KVHintsEnvelope {
|
||||
string protocol_version = 1;
|
||||
string message_id = 2;
|
||||
repeated KVHintAction actions = 3;
|
||||
}
|
||||
|
||||
message KVHintAction {
|
||||
string action_id = 1;
|
||||
string action_type = 2;
|
||||
string action_version = 3;
|
||||
google.protobuf.Struct payload = 4;
|
||||
}
|
||||
|
||||
message RandomSampling {
|
||||
|
||||
@@ -23,7 +23,7 @@ cargo test -p vllm-bench -- --ignored
|
||||
|
||||
## Architecture
|
||||
|
||||
- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single)
|
||||
- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single, `mm-processor` subcommand)
|
||||
- `src/cli.rs` — clap derive CLI args (~50+ flags)
|
||||
- `src/config.rs` — Validated config from CLI; `GoodputConfig`, `RampUpConfig`, sampling param merging
|
||||
- `src/error.rs` — `BenchError` enum (Http, Json, Tokenizer, Config, EndpointTimeout, Backend, Io)
|
||||
@@ -32,6 +32,7 @@ cargo test -p vllm-bench -- --ignored
|
||||
- `src/sweep.rs` — Concurrency/rate parameter sweep (`--sweep-max-concurrency`, `--sweep-request-rate`)
|
||||
- `src/multi_run.rs` — N-run aggregation with mean/std/min/max/CV (`--num-runs`)
|
||||
- `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`)
|
||||
- `src/mm_processor.rs` — `vllm-bench mm-processor` subcommand: offline multimodal preprocessing latency benchmark against a managed headless engine (mirrors `vllm bench mm-processor`); per-stage timings come from `vllm_tracing::timing::RequestTimingLayer` (via `vllm_chat::mm_timing_layer`)
|
||||
- `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback
|
||||
- `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source)
|
||||
- `src/hub.rs` — `HubRepo`: thin wrapper over hf-hub's async (reqwest/native-tls) API — file download and repo listing, optional revision; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans
|
||||
|
||||
@@ -36,6 +36,11 @@ tracing.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-build-info.workspace = true
|
||||
vllm-chat.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-managed-engine.workspace = true
|
||||
vllm-text.workspace = true
|
||||
vllm-tracing.workspace = true
|
||||
|
||||
[lints]
|
||||
|
||||
@@ -18,6 +18,7 @@ vllm-bench --backend vllm --base-url http://127.0.0.1:8000 \
|
||||
- **Beyond a single run** — concurrency/rate **sweeps**, **multi-run** stats, **multi-turn** conversations, **LoRA** multi-adapter, and result **comparison**.
|
||||
- **Steady-state metrics** — throughput/latency measured over the saturated plateau, excluding ramp-up and drain.
|
||||
- **Parity** — JSON output schema and timing semantics match Python `vllm bench serve` exactly.
|
||||
- **Multimodal preprocessing** — `mm-processor` subcommand measures per-stage VLM preprocessing latency against a managed headless engine (mirrors `vllm bench mm-processor`).
|
||||
|
||||
### Performance vs. Python
|
||||
|
||||
@@ -404,6 +405,32 @@ vllm-bench \
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Multimodal preprocessing (offline)</b></summary>
|
||||
|
||||
The `mm-processor` subcommand benchmarks the multimodal preprocessing pipeline
|
||||
itself (media fetch → per-modality preprocess → prompt expansion), mirroring
|
||||
`vllm bench mm-processor`. It spawns a managed headless Python engine (real
|
||||
weights, real encoder) — no serving endpoint needed — and reports per-stage
|
||||
timing (mean/median/std/percentiles) plus end-to-end latency.
|
||||
|
||||
```bash
|
||||
vllm-bench mm-processor \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--num-prompts 100 --num-warmups 5 \
|
||||
--random-input-len 512 --random-output-len 128 \
|
||||
--random-mm-base-items-per-request 1 \
|
||||
--random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
|
||||
--random-mm-bucket-config '{(1024, 800, 1): 1.0}' \
|
||||
--metric-percentiles "50,90,99" \
|
||||
--output-json mm_processor_stats.json
|
||||
```
|
||||
|
||||
`--max-concurrency` defaults to `1`, matching Python's serial driver; raise it
|
||||
to measure a concurrent serving-style workload.
|
||||
|
||||
</details>
|
||||
|
||||
## Supported Backends
|
||||
|
||||
### Generation
|
||||
@@ -705,6 +732,36 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b><code>mm-processor</code> subcommand</b></summary>
|
||||
|
||||
Offline multimodal preprocessing latency benchmark (mirrors `vllm bench mm-processor`). Spawns a managed headless Python engine and drives the full chat preprocessing pipeline.
|
||||
|
||||
| Flag | Default | Description |
|
||||
| ------ | --------- | ------------- |
|
||||
| `--model` | (required) | Model to serve and benchmark (HF id or local path) |
|
||||
| `--num-prompts` | `10` | Prompts to process, excluding warmups |
|
||||
| `--num-warmups` | `1` | Warmup prompts processed and discarded before timing |
|
||||
| `--metric-percentiles` | `99` | Comma-separated percentiles to report |
|
||||
| `--output-json` | — | Path to write aggregate stats as JSON (same schema as Python) |
|
||||
| `--random-input-len` | `1024` | Text input length per prompt |
|
||||
| `--random-output-len` | `128` | Expected output length per prompt |
|
||||
| `--random-prefix-len` | `0` | Prefix token length per prompt |
|
||||
| `--random-range-ratio` | `0.0` | Input/output length range ratio (float or `{"input": i, "output": o}`) |
|
||||
| `--random-mm-base-items-per-request` | `1` | Base number of multimodal items per request |
|
||||
| `--random-mm-num-mm-items-range-ratio` | `0.0` | Range ratio (in [0, 1]) for the number of mm items per request |
|
||||
| `--random-mm-limit-mm-per-prompt` | `{"image": 255, "video": 1}` | Per-modality item caps |
|
||||
| `--random-mm-bucket-config` | `{(256,256,1): 0.5, (720,1280,1): 0.5}` | Bucket config |
|
||||
| `--seed` | `0` | Seed for dataset generation |
|
||||
| `--request-id-prefix` | `mm-proc-` | Request-id prefix for per-request timing keys |
|
||||
| `--max-concurrency` | `1` | Maximum concurrent requests (1 matches Python's serial driver) |
|
||||
| `--trust-remote-code` | `false` | Trust remote code for the tokenizer and the managed engine |
|
||||
| `--chat-template` | — | Chat-template override (inline text or path) |
|
||||
|
||||
Managed headless-engine options (e.g. `--python`, `--data-parallel-address`, `--data-parallel-rpc-port`, `--data-parallel-size`, `--max-model-len`, `--python-args`) are shared with other managed-engine commands; see `vllm-bench mm-processor --help`.
|
||||
|
||||
</details>
|
||||
|
||||
## Tokenizer Support
|
||||
|
||||
Tokenizers are loaded with a three-tier fallback chain:
|
||||
@@ -734,10 +791,11 @@ Use `--append-result` to append multiple runs to the same file in JSONL format.
|
||||
|
||||
```text
|
||||
src/
|
||||
├── main.rs # Entry point, mimalloc, tokio runtime, mode dispatch
|
||||
├── main.rs # Entry point, mimalloc, tokio runtime, mode/subcommand dispatch
|
||||
├── cli.rs # clap CLI argument definitions
|
||||
├── config.rs # Validated config, goodput/ramp-up parsing
|
||||
├── benchmark.rs # Core orchestrator (schedule, spawn, collect, verify, profile)
|
||||
├── mm_processor.rs # Offline mm preprocessing benchmark (mm-processor subcommand)
|
||||
├── multi_turn.rs # Multi-turn conversation orchestrator (channel workers)
|
||||
├── compare.rs # Result diff (--compare file_a.json file_b.json)
|
||||
├── sweep.rs # Parameter sweep (--sweep-max-concurrency, --sweep-request-rate)
|
||||
|
||||
@@ -10,6 +10,7 @@ mod datasets;
|
||||
mod error;
|
||||
mod hub;
|
||||
mod metrics;
|
||||
mod mm_processor;
|
||||
mod multi_run;
|
||||
mod multi_turn;
|
||||
mod output;
|
||||
@@ -25,6 +26,7 @@ pub use cli::{
|
||||
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
|
||||
};
|
||||
use config::BenchConfig;
|
||||
pub use mm_processor::{MmProcessorArgs, run_mm_processor};
|
||||
|
||||
/// Prepare process-wide resources for a benchmark run.
|
||||
pub fn prepare_process() {
|
||||
|
||||
@@ -6,22 +6,30 @@
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
use anyhow::Context;
|
||||
use clap::Parser;
|
||||
use clap::{Parser, Subcommand};
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(
|
||||
name = "vllm-bench",
|
||||
about = "Benchmark online serving throughput",
|
||||
about = "Benchmark online serving throughput and offline multimodal preprocessing",
|
||||
version = vllm_build_info::VERSION
|
||||
)]
|
||||
struct Cli {
|
||||
#[command(flatten)]
|
||||
args: vllm_bench::BenchServeArgs,
|
||||
|
||||
/// Optional subcommand; absent by default for the online serving benchmark.
|
||||
#[command(subcommand)]
|
||||
command: Option<Command>,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Command {
|
||||
/// Offline multimodal preprocessing latency benchmark.
|
||||
MmProcessor(vllm_bench::MmProcessorArgs),
|
||||
}
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
vllm_tracing::init_tracing("Bench");
|
||||
|
||||
let cli = Cli::parse();
|
||||
vllm_bench::prepare_process();
|
||||
|
||||
@@ -30,5 +38,15 @@ fn main() -> anyhow::Result<()> {
|
||||
.build()
|
||||
.context("Failed to build tokio runtime")?;
|
||||
|
||||
runtime.block_on(vllm_bench::run(cli.args))
|
||||
match cli.command {
|
||||
Some(Command::MmProcessor(args)) => {
|
||||
let (timing_layer, timing_stats) = vllm_chat::mm_timing_layer();
|
||||
vllm_tracing::init_tracing_with("Bench", timing_layer);
|
||||
runtime.block_on(vllm_bench::run_mm_processor(args, timing_stats))
|
||||
}
|
||||
None => {
|
||||
vllm_tracing::init_tracing("Bench");
|
||||
runtime.block_on(vllm_bench::run(cli.args))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -216,20 +216,20 @@ pub fn calculate_metrics(
|
||||
(metrics, actual_output_lens)
|
||||
}
|
||||
|
||||
fn mean(data: &[f64]) -> f64 {
|
||||
pub(crate) fn mean(data: &[f64]) -> f64 {
|
||||
if data.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
data.iter().sum::<f64>() / data.len() as f64
|
||||
}
|
||||
|
||||
fn sort_clone(data: &[f64]) -> Vec<f64> {
|
||||
pub(crate) fn sort_clone(data: &[f64]) -> Vec<f64> {
|
||||
let mut sorted = data.to_vec();
|
||||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
sorted
|
||||
}
|
||||
|
||||
fn median_sorted(sorted: &[f64]) -> f64 {
|
||||
pub(crate) fn median_sorted(sorted: &[f64]) -> f64 {
|
||||
if sorted.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
@@ -241,7 +241,7 @@ fn median_sorted(sorted: &[f64]) -> f64 {
|
||||
}
|
||||
}
|
||||
|
||||
fn std_dev(data: &[f64]) -> f64 {
|
||||
pub(crate) fn std_dev(data: &[f64]) -> f64 {
|
||||
if data.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
@@ -250,7 +250,7 @@ fn std_dev(data: &[f64]) -> f64 {
|
||||
variance.sqrt()
|
||||
}
|
||||
|
||||
fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {
|
||||
pub(crate) fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {
|
||||
if sorted_data.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,573 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
//! Multimodal preprocessing benchmark (`vllm-bench mm-processor`), mirroring
|
||||
//! `vllm bench mm-processor`.
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use clap::Args;
|
||||
use futures::future::join_all;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use vllm_chat::{
|
||||
ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRequest, LoadModelBackendsOptions,
|
||||
load_model_backends,
|
||||
};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_managed_engine::cli::ManagedEngineArgs;
|
||||
use vllm_managed_engine::{ManagedEngineConfig, ManagedEngineHandle};
|
||||
use vllm_text::TextLlm;
|
||||
use vllm_tracing::timing::RequestTimingStats;
|
||||
|
||||
use crate::config::RangeRatio;
|
||||
use crate::datasets::SampleRequest;
|
||||
use crate::datasets::random_mm::{
|
||||
MmBucketKey, MmLimitPerPrompt, generate_random_mm_dataset, parse_bucket_config,
|
||||
parse_limit_mm_per_prompt,
|
||||
};
|
||||
use crate::error::BenchError;
|
||||
use crate::metrics::calculator::{mean, median_sorted, percentile_sorted, sort_clone, std_dev};
|
||||
use crate::tokenizer::{TokenizerKind, load_tokenizer};
|
||||
|
||||
/// Seconds -> milliseconds.
|
||||
const SEC_TO_MS: f64 = 1000.0;
|
||||
|
||||
/// Preferred column order; other stages appended alphabetically.
|
||||
const STAGE_ORDER: &[&str] = &[
|
||||
"preprocessor_total_ms",
|
||||
"media_fetch_ms",
|
||||
"preprocess_image_ms",
|
||||
"preprocess_video_ms",
|
||||
"preprocess_audio_ms",
|
||||
"prompt_expansion_ms",
|
||||
];
|
||||
|
||||
/// Time to wait for managed engines to register.
|
||||
const ENGINE_READY_TIMEOUT: Duration = Duration::from_secs(600);
|
||||
|
||||
/// Maximum time to wait for the managed engine to drain on shutdown.
|
||||
const ENGINE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(120);
|
||||
|
||||
#[derive(Args, Debug, Clone)]
|
||||
pub struct MmProcessorArgs {
|
||||
/// Model to serve and benchmark (Hugging Face id or local path).
|
||||
#[arg(long)]
|
||||
pub model: String,
|
||||
|
||||
/// Number of prompts to process, excluding warmups.
|
||||
#[arg(long, default_value_t = 10)]
|
||||
pub num_prompts: usize,
|
||||
|
||||
/// Number of warmup prompts processed and discarded before timing.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub num_warmups: usize,
|
||||
|
||||
/// Comma-separated percentiles to report, e.g. "50,90,99".
|
||||
#[arg(long, default_value = "99")]
|
||||
pub metric_percentiles: String,
|
||||
|
||||
/// Path to write the aggregate stats as JSON.
|
||||
#[arg(long)]
|
||||
pub output_json: Option<String>,
|
||||
|
||||
/// Text input length per prompt (before range scaling).
|
||||
#[arg(long, default_value_t = 1024)]
|
||||
pub random_input_len: usize,
|
||||
|
||||
/// Expected output length per prompt.
|
||||
#[arg(long, default_value_t = 128)]
|
||||
pub random_output_len: usize,
|
||||
|
||||
/// Prefix token length per prompt.
|
||||
#[arg(long, default_value_t = 0)]
|
||||
pub random_prefix_len: usize,
|
||||
|
||||
/// Input/output length range ratio: a float or `{"input": i, "output": o}`.
|
||||
#[arg(long, default_value = "0.0")]
|
||||
pub random_range_ratio: String,
|
||||
|
||||
/// Base number of multimodal items per request.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub random_mm_base_items_per_request: usize,
|
||||
|
||||
/// Range ratio (in [0, 1]) for the number of mm items per request.
|
||||
#[arg(long, default_value_t = 0.0)]
|
||||
pub random_mm_num_mm_items_range_ratio: f64,
|
||||
|
||||
/// Per-modality item caps, e.g. `{"image": 255, "video": 1}`.
|
||||
#[arg(long, default_value = "{\"image\": 255, \"video\": 1}")]
|
||||
pub random_mm_limit_mm_per_prompt: String,
|
||||
|
||||
/// Bucket config, e.g. `{(256,256,1): 0.5, (720,1280,1): 0.5}`.
|
||||
#[arg(long, default_value = "{(256,256,1): 0.5, (720,1280,1): 0.5}")]
|
||||
pub random_mm_bucket_config: String,
|
||||
|
||||
/// Seed for dataset generation.
|
||||
#[arg(long)]
|
||||
pub seed: Option<u64>,
|
||||
|
||||
/// Request-id prefix for per-request timing keys.
|
||||
#[arg(long, default_value = "mm-proc-")]
|
||||
pub request_id_prefix: String,
|
||||
|
||||
/// Maximum concurrent requests (default 1, matching Python's serial driver).
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub max_concurrency: usize,
|
||||
|
||||
/// Trust remote code for the tokenizer and the managed engine.
|
||||
#[arg(long, default_value_t = false)]
|
||||
pub trust_remote_code: bool,
|
||||
|
||||
/// Optional chat-template override (inline text or path).
|
||||
#[arg(long)]
|
||||
pub chat_template: Option<String>,
|
||||
|
||||
/// Managed headless-engine options.
|
||||
#[command(flatten)]
|
||||
pub engine: ManagedEngineArgs,
|
||||
}
|
||||
|
||||
/// Entry point for `vllm-bench mm-processor`.
|
||||
pub async fn run_mm_processor(args: MmProcessorArgs, timing: RequestTimingStats) -> Result<()> {
|
||||
if args.engine.data_parallel_size_local == Some(0) {
|
||||
anyhow::bail!(
|
||||
"`--data-parallel-size-local 0` is not supported; the benchmark requires \
|
||||
a local managed engine"
|
||||
);
|
||||
}
|
||||
|
||||
// Spawn the managed headless Python engine (real weights, real encoder).
|
||||
let handshake_port = args
|
||||
.engine
|
||||
.resolve_handshake_port()
|
||||
.context("failed to allocate handshake port")?;
|
||||
let mut python_args = args.engine.python_args.clone();
|
||||
if args.trust_remote_code {
|
||||
python_args.push("--trust-remote-code".to_string());
|
||||
}
|
||||
let engine_config = ManagedEngineConfig {
|
||||
python: args.engine.python.clone(),
|
||||
model: args.model.clone(),
|
||||
handshake_host: args.engine.handshake_host.clone(),
|
||||
handshake_port,
|
||||
data_parallel_size: args.engine.data_parallel_size,
|
||||
python_args,
|
||||
};
|
||||
eprintln!(
|
||||
"Spawning managed headless engine: python={} model={} data_parallel_size={}",
|
||||
engine_config.python, engine_config.model, engine_config.data_parallel_size
|
||||
);
|
||||
let engine = ManagedEngineHandle::spawn(engine_config)
|
||||
.await
|
||||
.context("failed to start managed Python headless engine")?;
|
||||
|
||||
let result = run_mm_processor_with_engine(args, handshake_port, &timing).await;
|
||||
let shutdown_result = engine
|
||||
.shutdown(ENGINE_SHUTDOWN_TIMEOUT)
|
||||
.await
|
||||
.context("failed to shut down managed engine");
|
||||
result?;
|
||||
shutdown_result?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_mm_processor_with_engine(
|
||||
args: MmProcessorArgs,
|
||||
handshake_port: u16,
|
||||
timing: &RequestTimingStats,
|
||||
) -> Result<()> {
|
||||
let range_ratio =
|
||||
RangeRatio::parse(&args.random_range_ratio).context("invalid --random-range-ratio")?;
|
||||
let limit: MmLimitPerPrompt = parse_limit_mm_per_prompt(&args.random_mm_limit_mm_per_prompt)
|
||||
.context("invalid --random-mm-limit-mm-per-prompt")?;
|
||||
let buckets: Vec<(MmBucketKey, f64)> = parse_bucket_config(&args.random_mm_bucket_config)
|
||||
.context("invalid --random-mm-bucket-config")?;
|
||||
let percentiles = parse_percentiles(&args.metric_percentiles)?;
|
||||
let seed = args.seed.unwrap_or(0);
|
||||
|
||||
// Offline tokenizer: local Hugging Face tokenizer or built-in tiktoken.
|
||||
let tokenizer = load_tokenizer(&args.model, args.trust_remote_code, None)
|
||||
.await
|
||||
.context("failed to load tokenizer")?;
|
||||
|
||||
// Frontend backends: model configs + multimodal processor metadata.
|
||||
let backends = load_model_backends(
|
||||
&args.model,
|
||||
LoadModelBackendsOptions {
|
||||
chat_template: args.chat_template.clone(),
|
||||
..LoadModelBackendsOptions::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.context("failed to load model backends")?;
|
||||
|
||||
if backends.chat_backend.multimodal_model_info().is_none() {
|
||||
anyhow::bail!(
|
||||
"model `{}` has no multimodal processor registered",
|
||||
args.model
|
||||
);
|
||||
}
|
||||
|
||||
// Connect to the engine over the same transport the serving frontend uses.
|
||||
let client = EngineCoreClient::connect(EngineCoreClientConfig {
|
||||
transport_mode: TransportMode::HandshakeOwner {
|
||||
handshake_address: format!("tcp://{}:{}", args.engine.handshake_host, handshake_port),
|
||||
advertised_host: args.engine.handshake_host.clone(),
|
||||
engine_count: args
|
||||
.engine
|
||||
.data_parallel_size_local
|
||||
.unwrap_or(args.engine.data_parallel_size),
|
||||
ready_timeout: ENGINE_READY_TIMEOUT,
|
||||
local_input_address: None,
|
||||
local_output_address: None,
|
||||
},
|
||||
coordinator_mode: None,
|
||||
model_name: args.model.clone(),
|
||||
client_index: 0,
|
||||
})
|
||||
.await
|
||||
.context("failed to connect to engine core")?;
|
||||
|
||||
let chat = ChatLlm::new(
|
||||
TextLlm::new(Llm::new(client), backends.text_backend),
|
||||
backends.chat_backend,
|
||||
);
|
||||
|
||||
let semaphore = Arc::new(Semaphore::new(args.max_concurrency));
|
||||
|
||||
// Warmup: generate a separate dataset, process it, then drain and discard
|
||||
// the collected timings so cold caches do not skew the results.
|
||||
if args.num_warmups > 0 {
|
||||
let warmup = generate_dataset(
|
||||
&tokenizer,
|
||||
args.num_warmups,
|
||||
seed.wrapping_add(1),
|
||||
&args,
|
||||
range_ratio,
|
||||
&limit,
|
||||
&buckets,
|
||||
)?;
|
||||
eprintln!("Processing {} warmup requests...", warmup.len());
|
||||
run_batch(&chat, &warmup, &semaphore).await;
|
||||
let _ = timing.stat();
|
||||
}
|
||||
|
||||
let samples = generate_dataset(
|
||||
&tokenizer,
|
||||
args.num_prompts,
|
||||
seed,
|
||||
&args,
|
||||
range_ratio,
|
||||
&limit,
|
||||
&buckets,
|
||||
)?;
|
||||
eprintln!("Processing {} requests...", samples.len());
|
||||
let start_time = Instant::now();
|
||||
let outcomes = run_batch(&chat, &samples, &semaphore).await;
|
||||
let total_time = start_time.elapsed();
|
||||
|
||||
let completed = outcomes.iter().filter(|outcome| outcome.is_ok()).count();
|
||||
let failed = outcomes.len() - completed;
|
||||
let e2el_times: Vec<f64> = outcomes
|
||||
.iter()
|
||||
.filter_map(|outcome| outcome.as_ref().ok().map(|elapsed| elapsed.as_secs_f64()))
|
||||
.collect();
|
||||
|
||||
let per_request = timing.stat();
|
||||
let stats = aggregate_stats(per_request, &percentiles);
|
||||
|
||||
eprintln!(
|
||||
"Processed {} requests (+{} warmup) for model `{}`",
|
||||
samples.len(),
|
||||
args.num_warmups,
|
||||
args.model
|
||||
);
|
||||
print_report(&stats, &percentiles);
|
||||
print_e2el_summary(completed, failed, total_time, &e2el_times, &percentiles);
|
||||
|
||||
if let Some(path) = &args.output_json {
|
||||
write_result_json(path, completed, failed, &e2el_times, &percentiles, stats)?;
|
||||
eprintln!("Wrote stats JSON to {path}");
|
||||
}
|
||||
|
||||
chat.shutdown().await.context("failed to shut down chat facade")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Submit every sample through the chat pipeline and drive each stream to
|
||||
/// completion, returning per-request wall-clock latencies.
|
||||
async fn run_batch(
|
||||
chat: &ChatLlm,
|
||||
samples: &[SampleRequest],
|
||||
semaphore: &Arc<Semaphore>,
|
||||
) -> Vec<Result<Duration>> {
|
||||
let submissions = samples.iter().map(|sample| {
|
||||
let request = build_chat_request(sample);
|
||||
let semaphore = semaphore.clone();
|
||||
async move {
|
||||
let _permit = semaphore.acquire().await.context("semaphore closed")?;
|
||||
let request = request?;
|
||||
let start = Instant::now();
|
||||
chat.chat(request)
|
||||
.await
|
||||
.context("failed to submit request")?
|
||||
.collect_message()
|
||||
.await
|
||||
.context("request failed")?;
|
||||
Ok(start.elapsed())
|
||||
}
|
||||
});
|
||||
join_all(submissions).await
|
||||
}
|
||||
|
||||
fn generate_dataset(
|
||||
tokenizer: &TokenizerKind,
|
||||
num_requests: usize,
|
||||
seed: u64,
|
||||
args: &MmProcessorArgs,
|
||||
range_ratio: RangeRatio,
|
||||
limit: &MmLimitPerPrompt,
|
||||
buckets: &[(MmBucketKey, f64)],
|
||||
) -> Result<Vec<SampleRequest>> {
|
||||
generate_random_mm_dataset(
|
||||
tokenizer,
|
||||
num_requests,
|
||||
args.random_input_len,
|
||||
args.random_output_len,
|
||||
args.random_prefix_len,
|
||||
range_ratio,
|
||||
seed,
|
||||
&args.request_id_prefix,
|
||||
args.random_mm_base_items_per_request,
|
||||
args.random_mm_num_mm_items_range_ratio,
|
||||
limit,
|
||||
buckets,
|
||||
false,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))
|
||||
}
|
||||
|
||||
/// Build a chat request from a sample (text prompt + `image_url` parts).
|
||||
fn build_chat_request(sample: &SampleRequest) -> Result<ChatRequest> {
|
||||
let mut request = ChatRequest::for_test();
|
||||
request.request_id = sample.request_id.clone().unwrap_or_else(|| "mm-proc".to_string());
|
||||
request.sampling_params.temperature = Some(0.0);
|
||||
request.sampling_params.max_tokens = Some(u32::try_from(sample.expected_output_len)?);
|
||||
|
||||
let mut parts = vec![ChatContentPart::text(sample.prompt.as_ref())];
|
||||
if let Some(mm_items) = &sample.multi_modal_content {
|
||||
for fragment in mm_items.iter() {
|
||||
let url = extract_image_url(fragment)?;
|
||||
parts.push(ChatContentPart::image_url(url));
|
||||
}
|
||||
}
|
||||
request.messages = vec![ChatMessage::user(ChatContent::from(parts))];
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
/// Extract the `image_url.url` string from a pre-serialized content fragment
|
||||
/// (`{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,..."}}`).
|
||||
fn extract_image_url(fragment: &str) -> Result<String> {
|
||||
let value: serde_json::Value = serde_json::from_str(fragment)
|
||||
.with_context(|| format!("invalid multimodal fragment: {fragment}"))?;
|
||||
let url = value
|
||||
.get("image_url")
|
||||
.and_then(|v| v.get("url"))
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::Config("fragment missing image_url.url".to_string()))?;
|
||||
Ok(url.to_string())
|
||||
}
|
||||
|
||||
fn parse_percentiles(raw: &str) -> Result<Vec<f64>> {
|
||||
let mut values = Vec::new();
|
||||
for part in raw.split(',') {
|
||||
let part = part.trim();
|
||||
if part.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let value: f64 = part.parse().with_context(|| format!("invalid percentile `{part}`"))?;
|
||||
if !(0.0..=100.0).contains(&value) {
|
||||
anyhow::bail!("percentile `{part}` out of range [0, 100]");
|
||||
}
|
||||
values.push(value);
|
||||
}
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
/// Aggregate per-request stage timings into `{stage_ms: {statistic: ms}}`,
|
||||
/// matching the Python `mm_processor_stats` schema orientation.
|
||||
fn aggregate_stats(
|
||||
per_request: HashMap<String, HashMap<String, f64>>,
|
||||
percentiles: &[f64],
|
||||
) -> BTreeMap<String, BTreeMap<String, f64>> {
|
||||
let mut stage_values: BTreeMap<String, Vec<f64>> = BTreeMap::new();
|
||||
for stage_map in per_request.into_values() {
|
||||
for (stage, seconds) in stage_map {
|
||||
stage_values.entry(stage).or_default().push(seconds);
|
||||
}
|
||||
}
|
||||
|
||||
stage_values
|
||||
.iter()
|
||||
.map(|(stage, values)| {
|
||||
let sorted = sort_clone(values);
|
||||
let mut stat_map = BTreeMap::new();
|
||||
stat_map.insert("mean".to_string(), mean(values) * SEC_TO_MS);
|
||||
stat_map.insert("median".to_string(), median_sorted(&sorted) * SEC_TO_MS);
|
||||
stat_map.insert("std".to_string(), std_dev(values) * SEC_TO_MS);
|
||||
for p in percentiles {
|
||||
stat_map.insert(format!("p{p}"), percentile_sorted(&sorted, *p) * SEC_TO_MS);
|
||||
}
|
||||
(stage.replace("_secs", "_ms"), stat_map)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Print a right-aligned plain-text table (pandas `to_string(index=False)` style).
|
||||
fn print_table(rows: &[Vec<String>]) {
|
||||
let ncols = rows.first().map_or(0, Vec::len);
|
||||
let mut widths = vec![0usize; ncols];
|
||||
for row in rows {
|
||||
for (i, cell) in row.iter().enumerate() {
|
||||
widths[i] = widths[i].max(cell.len());
|
||||
}
|
||||
}
|
||||
for row in rows {
|
||||
let line = row
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, cell)| format!("{cell:>width$}", width = widths[i]))
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ");
|
||||
println!("{line}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Print a plain-text table with one row per stage and one column per
|
||||
/// statistic (mirrors the Python pandas report layout).
|
||||
fn print_report(stats: &BTreeMap<String, BTreeMap<String, f64>>, percentiles: &[f64]) {
|
||||
// Column order: mean/median/std then the requested percentiles, like Python.
|
||||
let mut columns: Vec<(String, String)> =
|
||||
[("mean", "Mean"), ("median", "Median"), ("std", "Std")]
|
||||
.into_iter()
|
||||
.map(|(key, label)| (key.to_string(), label.to_string()))
|
||||
.collect();
|
||||
columns.extend(percentiles.iter().map(|p| (format!("p{p}"), format!("P{p:?}"))));
|
||||
for key in stats.values().flat_map(|map| map.keys()) {
|
||||
if !columns.iter().any(|(k, _)| k == key) {
|
||||
columns.push((key.clone(), key.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
let mut stages: Vec<&String> = stats.keys().collect();
|
||||
stages.sort_by_key(|stage| {
|
||||
STAGE_ORDER
|
||||
.iter()
|
||||
.position(|s| s == stage)
|
||||
.map(|i| (i as isize, (*stage).clone()))
|
||||
.unwrap_or((isize::MAX, (*stage).clone()))
|
||||
});
|
||||
|
||||
let mut rows: Vec<Vec<String>> = vec![
|
||||
std::iter::once("Stage".to_string())
|
||||
.chain(columns.iter().map(|(_, label)| label.clone()))
|
||||
.collect(),
|
||||
];
|
||||
for stage in &stages {
|
||||
let mut row = vec![(*stage).clone()];
|
||||
for (key, _) in &columns {
|
||||
row.push(
|
||||
stats[*stage]
|
||||
.get(key)
|
||||
.map(|v| format!("{v:.2}"))
|
||||
.unwrap_or_else(|| "-".to_string()),
|
||||
);
|
||||
}
|
||||
rows.push(row);
|
||||
}
|
||||
println!("MM Processor Metrics:");
|
||||
print_table(&rows);
|
||||
}
|
||||
|
||||
/// Print the end-to-end latency summary (mirrors the Python report layout).
|
||||
fn print_e2el_summary(
|
||||
completed: usize,
|
||||
failed: usize,
|
||||
total_time: Duration,
|
||||
e2el_times_secs: &[f64],
|
||||
percentiles: &[f64],
|
||||
) {
|
||||
println!("Completed: {completed}, Failed: {failed}");
|
||||
if e2el_times_secs.is_empty() {
|
||||
println!("End-to-End Latency (ms): no completed requests");
|
||||
return;
|
||||
}
|
||||
|
||||
let e2el_ms: Vec<f64> = e2el_times_secs.iter().map(|secs| secs * SEC_TO_MS).collect();
|
||||
let sorted = sort_clone(&e2el_ms);
|
||||
let mut rows = vec![
|
||||
vec!["Metric".to_string(), "Value (ms)".to_string()],
|
||||
vec!["Mean".to_string(), format!("{:.2}", mean(&e2el_ms))],
|
||||
vec![
|
||||
"Median".to_string(),
|
||||
format!("{:.2}", median_sorted(&sorted)),
|
||||
],
|
||||
vec!["Std".to_string(), format!("{:.2}", std_dev(&e2el_ms))],
|
||||
];
|
||||
for p in percentiles {
|
||||
rows.push(vec![
|
||||
format!("P{p:?}"),
|
||||
format!("{:.2}", percentile_sorted(&sorted, *p)),
|
||||
]);
|
||||
}
|
||||
println!("End-to-End Latency (ms):");
|
||||
print_table(&rows);
|
||||
let throughput = completed as f64 / total_time.as_secs_f64();
|
||||
println!("Request throughput: {throughput:.3} req/s");
|
||||
}
|
||||
|
||||
/// JSON result schema mirroring `vllm bench mm-processor`.
|
||||
#[derive(Serialize)]
|
||||
struct BenchmarkResultJson {
|
||||
completed: usize,
|
||||
failed: usize,
|
||||
mean_e2el_ms: f64,
|
||||
median_e2el_ms: f64,
|
||||
std_e2el_ms: f64,
|
||||
percentiles_e2el_ms: Vec<(f64, f64)>,
|
||||
/// `{stage_ms: {statistic: ms}}`, keyed by stage like the Python schema.
|
||||
mm_processor_stats: BTreeMap<String, BTreeMap<String, f64>>,
|
||||
}
|
||||
|
||||
fn write_result_json(
|
||||
path: &str,
|
||||
completed: usize,
|
||||
failed: usize,
|
||||
e2el_times_secs: &[f64],
|
||||
percentiles: &[f64],
|
||||
mm_processor_stats: BTreeMap<String, BTreeMap<String, f64>>,
|
||||
) -> Result<()> {
|
||||
let e2el_ms: Vec<f64> = e2el_times_secs.iter().map(|secs| secs * SEC_TO_MS).collect();
|
||||
let sorted = sort_clone(&e2el_ms);
|
||||
let result = BenchmarkResultJson {
|
||||
completed,
|
||||
failed,
|
||||
mean_e2el_ms: mean(&e2el_ms),
|
||||
median_e2el_ms: median_sorted(&sorted),
|
||||
std_e2el_ms: std_dev(&e2el_ms),
|
||||
percentiles_e2el_ms: percentiles
|
||||
.iter()
|
||||
.map(|p| (*p, percentile_sorted(&sorted, *p)))
|
||||
.collect(),
|
||||
mm_processor_stats,
|
||||
};
|
||||
let json = serde_json::to_string_pretty(&result).context("failed to serialize result")?;
|
||||
std::fs::write(path, format!("{json}\n")).with_context(|| format!("failed to write {path}"))?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -35,6 +35,7 @@ vllm-llm.workspace = true
|
||||
vllm-parser.workspace = true
|
||||
vllm-text.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
vllm-tracing.workspace = true
|
||||
xgrammar-structural-tag.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
||||
@@ -60,6 +60,11 @@ pub enum Error {
|
||||
#[source]
|
||||
error: BoxedError,
|
||||
},
|
||||
#[error("failed to initialize request output parser")]
|
||||
OutputParserInitialization {
|
||||
#[source]
|
||||
error: BoxedError,
|
||||
},
|
||||
#[error(
|
||||
"gpt_oss uses native Harmony output parsing; generic {kind} parser override `{selection}` is not supported"
|
||||
)]
|
||||
@@ -87,8 +92,11 @@ pub enum Error {
|
||||
ToolChoiceRequiresTools,
|
||||
#[error("tool_choice function `{name}` was not found in the available tools")]
|
||||
ToolChoiceFunctionNotFound { name: String },
|
||||
#[error("failed to build structural tag: {message}")]
|
||||
StructuralTag { message: String },
|
||||
#[error("failed to build output grammar")]
|
||||
OutputGrammar {
|
||||
#[source]
|
||||
error: BoxedError,
|
||||
},
|
||||
#[error(transparent)]
|
||||
Text(#[from] vllm_text::Error),
|
||||
#[error(transparent)]
|
||||
|
||||
+33
-16
@@ -22,6 +22,8 @@ pub use event::{
|
||||
};
|
||||
use futures::{StreamExt, TryStreamExt as _};
|
||||
pub use llm_multimodal::MediaContentPart;
|
||||
pub use multimodal::{mm_request_span, mm_timing_layer};
|
||||
use output::apply_output_grammar;
|
||||
pub use output::{
|
||||
ChatOutputProcessor, DefaultChatOutputProcessor, DynChatOutputProcessor,
|
||||
HarmonyChatOutputProcessor,
|
||||
@@ -62,7 +64,7 @@ use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest};
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest, TextRequestProcessor};
|
||||
|
||||
/// Chat request preparation shared by inference and render-only frontends.
|
||||
pub struct ChatRequestProcessor {
|
||||
@@ -171,15 +173,8 @@ impl ChatRequestProcessor {
|
||||
async fn prepare_text_request(&self, request: ChatRequest) -> Result<TextRequest> {
|
||||
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
|
||||
let arrival_time = vllm_llm::current_unix_timestamp_secs();
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let reasoning_parser_kwargs =
|
||||
request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.is_some()
|
||||
.then(|| ReasoningParserKwargs {
|
||||
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
|
||||
});
|
||||
let mut rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let chat_template_kwargs = std::mem::take(&mut rendered.effective_template_kwargs);
|
||||
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
|
||||
Ok(TextRequest {
|
||||
request_id: request.request_id,
|
||||
@@ -194,7 +189,11 @@ impl ChatRequestProcessor {
|
||||
add_special_tokens: request.add_special_tokens,
|
||||
data_parallel_rank: request.data_parallel_rank,
|
||||
session_id: request.session_id,
|
||||
reasoning_parser_kwargs,
|
||||
kv_hints: None,
|
||||
reasoning_parser_kwargs: ReasoningParserKwargs {
|
||||
chat_template_kwargs,
|
||||
},
|
||||
reasoning_ended: None,
|
||||
lora_request: request.lora_request,
|
||||
arrival_time: Some(arrival_time),
|
||||
})
|
||||
@@ -207,12 +206,19 @@ impl ChatRequestProcessor {
|
||||
}
|
||||
|
||||
/// Prepare one chat request without submitting it to an engine.
|
||||
///
|
||||
/// The returned text request is already tokenized: the output processor
|
||||
/// is initialized from the final prompt token IDs and its grammar, if any,
|
||||
/// is applied to the request before it is handed to `text_processor` again
|
||||
/// for lowering.
|
||||
pub async fn prepare(
|
||||
&self,
|
||||
mut request: ChatRequest,
|
||||
text_processor: &TextRequestProcessor,
|
||||
) -> Result<(TextRequest, DynChatOutputProcessor)> {
|
||||
request.validate()?;
|
||||
let output_processor = self.backend.new_chat_output_processor(
|
||||
|
||||
let mut output_processor = self.backend.new_chat_output_processor(
|
||||
&mut request,
|
||||
NewChatOutputProcessorOptions {
|
||||
tool_call_parser: &self.tool_call_parser,
|
||||
@@ -220,7 +226,17 @@ impl ChatRequestProcessor {
|
||||
tool_strict_level: self.tool_strict_level,
|
||||
},
|
||||
)?;
|
||||
let text_request = self.prepare_text_request(request).await?;
|
||||
let mut text_request = self.prepare_text_request(request).await?;
|
||||
|
||||
// Initialize the output processor from the final prompt token IDs and apply
|
||||
// the output grammar to the request, if any.
|
||||
let prompt_token_ids = text_processor.tokenize_in_place(&mut text_request)?;
|
||||
let grammar = {
|
||||
output_processor.initialize(prompt_token_ids)?;
|
||||
output_processor.build_output_grammar()?
|
||||
};
|
||||
apply_output_grammar(&mut text_request, grammar)?;
|
||||
|
||||
Ok((text_request, output_processor))
|
||||
}
|
||||
}
|
||||
@@ -338,7 +354,8 @@ impl ChatLlm {
|
||||
|
||||
/// Render, tokenize, and submit one chat request.
|
||||
pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
|
||||
let (text_request, output_processor) = self.processor.prepare(request).await?;
|
||||
let (text_request, output_processor) =
|
||||
self.processor.prepare(request, self.text.request_processor()).await?;
|
||||
let request_id = text_request.request_id.clone();
|
||||
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
|
||||
|
||||
@@ -395,7 +412,7 @@ mod tests {
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
|
||||
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -406,6 +423,6 @@ mod tests {
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
|
||||
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,16 +17,16 @@ use std::path::Path;
|
||||
use std::sync::{Arc, LazyLock};
|
||||
|
||||
use itertools::{Either, izip};
|
||||
use llm_multimodal::registry::ModelRegistryError;
|
||||
use llm_multimodal::{
|
||||
AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout,
|
||||
ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata,
|
||||
ModelProcessorSpec, ModelRegistry, PreProcessorConfig, PreprocessedEncoderInputs,
|
||||
PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor,
|
||||
VisionProcessorRegistry,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tracing::warn;
|
||||
use tracing::{Instrument as _, warn};
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem};
|
||||
use vllm_text::Prompt;
|
||||
@@ -43,17 +43,20 @@ mod input;
|
||||
mod item;
|
||||
mod preprocessed;
|
||||
mod tensor;
|
||||
mod timing;
|
||||
mod video;
|
||||
|
||||
use self::expand::expand_prompt_token_ids;
|
||||
pub use self::input::MultimodalInput;
|
||||
use self::timing::MM_STAGE_TARGET;
|
||||
pub use self::timing::{mm_request_span, mm_timing_layer};
|
||||
|
||||
/// Resolved multimodal support for one loaded model.
|
||||
#[derive(Clone)]
|
||||
pub struct MultimodalModelInfo {
|
||||
context: MultimodalModelContext,
|
||||
image: Option<ModalitySupport>,
|
||||
video: Option<ModalitySupport>,
|
||||
image: Option<VisionModalitySupport>,
|
||||
video: Option<VisionModalitySupport>,
|
||||
audio: Option<AudioModalitySupport>,
|
||||
media_connector: Arc<MediaConnector>,
|
||||
/// Maximum number of input items allowed per prompt for each modality.
|
||||
@@ -130,13 +133,18 @@ impl MultimodalModelContext {
|
||||
REGISTRY.lookup(&self.metadata())
|
||||
}
|
||||
|
||||
/// Resolve a static vision preprocessor for one loaded model.
|
||||
///
|
||||
/// The vision preprocessor serves both the image and video modalities.
|
||||
fn resolve_vision_processor(&self) -> Option<&'static dyn VisionPreProcessor> {
|
||||
static REGISTRY: LazyLock<VisionProcessorRegistry> =
|
||||
LazyLock::new(VisionProcessorRegistry::with_defaults);
|
||||
REGISTRY.find(&self.model_id, self.model_type.as_deref())
|
||||
/// Build a vision preprocessor for one loaded model and modality.
|
||||
fn resolve_vision_processor(
|
||||
&self,
|
||||
model_spec: &'static dyn ModelProcessorSpec,
|
||||
preprocessor_config: &PreProcessorConfig,
|
||||
modality: Modality,
|
||||
) -> Result<Option<Arc<dyn VisionPreProcessor>>> {
|
||||
match model_spec.vision_processor(&self.metadata(), preprocessor_config, modality) {
|
||||
Ok(processor) => Ok(Some(Arc::from(processor))),
|
||||
Err(ModelRegistryError::UnsupportedModality { .. }) => Ok(None),
|
||||
Err(error) => Err(error.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve an audio preprocessor for one loaded model.
|
||||
@@ -145,7 +153,7 @@ impl MultimodalModelContext {
|
||||
model_spec: &'static dyn ModelProcessorSpec,
|
||||
preprocessor_config: &PreProcessorConfig,
|
||||
) -> Option<Arc<dyn AudioPreProcessor>> {
|
||||
model_spec.audio_processor(&self.config, preprocessor_config).map(Arc::from)
|
||||
model_spec.audio_processor(&self.metadata(), preprocessor_config).map(Arc::from)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -239,17 +247,16 @@ impl ResolvedPlaceholder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Static per-modality vision preprocessor plus its loaded config, resolved
|
||||
/// placeholder tokens, and the model's shared tensor-layout spec.
|
||||
/// Model-owned vision preprocessor plus resolved placeholder tokens and
|
||||
/// the model's shared tensor-layout spec.
|
||||
#[derive(Clone)]
|
||||
struct ModalitySupport {
|
||||
struct VisionModalitySupport {
|
||||
spec: ResolvedMultimodalSpec,
|
||||
placeholder: ResolvedPlaceholder,
|
||||
processor: &'static dyn VisionPreProcessor,
|
||||
config: PreProcessorConfig,
|
||||
processor: Arc<dyn VisionPreProcessor>,
|
||||
}
|
||||
|
||||
/// Static audio preprocessor plus the resolved model contract for its output.
|
||||
/// Model-owned audio preprocessor plus the resolved model contract for its output.
|
||||
#[derive(Clone)]
|
||||
struct AudioModalitySupport {
|
||||
spec: ResolvedMultimodalSpec,
|
||||
@@ -396,7 +403,7 @@ impl MultimodalModelInfo {
|
||||
&context,
|
||||
preprocessor_config.clone(),
|
||||
video_preprocessor_config,
|
||||
);
|
||||
)?;
|
||||
let audio = Self::resolve_audio_lane(&context, &preprocessor_config)?;
|
||||
|
||||
if image.is_none() && video.is_none() && audio.is_none() {
|
||||
@@ -427,27 +434,19 @@ impl MultimodalModelInfo {
|
||||
///
|
||||
/// A missing model spec or vision processor disables both lanes but is
|
||||
/// not an error: unsupported model families are logged and skipped.
|
||||
/// Processor construction errors propagate to model loading.
|
||||
fn resolve_vision_lanes(
|
||||
context: &MultimodalModelContext,
|
||||
preprocessor_config: PreProcessorConfig,
|
||||
video_preprocessor_config: PreProcessorConfig,
|
||||
) -> (Option<ModalitySupport>, Option<ModalitySupport>) {
|
||||
) -> Result<(Option<VisionModalitySupport>, Option<VisionModalitySupport>)> {
|
||||
let Some(raw_spec) = context.resolve_model_spec() else {
|
||||
warn!(
|
||||
model_id = context.model_id,
|
||||
model_type = context.model_type,
|
||||
"multimodal model spec is not registered; disabling image/video support for this model"
|
||||
);
|
||||
return (None, None);
|
||||
};
|
||||
|
||||
let Some(processor) = context.resolve_vision_processor() else {
|
||||
warn!(
|
||||
model_id = context.model_id,
|
||||
model_type = context.model_type,
|
||||
"vision processor is not registered; disabling image/video support for this model"
|
||||
);
|
||||
return (None, None);
|
||||
return Ok((None, None));
|
||||
};
|
||||
|
||||
// Warn and disable the modality if the placeholder resolution fails.
|
||||
@@ -465,14 +464,19 @@ impl MultimodalModelInfo {
|
||||
}
|
||||
};
|
||||
|
||||
let image = resolve_placeholder(Modality::Image).map(|placeholder| ModalitySupport {
|
||||
spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Image),
|
||||
placeholder,
|
||||
processor,
|
||||
config: preprocessor_config.clone(),
|
||||
});
|
||||
let image = if let Some(placeholder) = resolve_placeholder(Modality::Image) {
|
||||
context
|
||||
.resolve_vision_processor(raw_spec, &preprocessor_config, Modality::Image)?
|
||||
.map(|processor| VisionModalitySupport {
|
||||
spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Image),
|
||||
placeholder,
|
||||
processor,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let video = resolve_placeholder(Modality::Video).and_then(|placeholder| {
|
||||
let video = if let Some(placeholder) = resolve_placeholder(Modality::Video) {
|
||||
// Placeholder expansion attributes markers to modalities by token
|
||||
// ID, so a marker shared with the image modality is ambiguous.
|
||||
let image_marker = image.as_ref().map(|image| image.placeholder.marker_token_id);
|
||||
@@ -484,16 +488,23 @@ impl MultimodalModelInfo {
|
||||
);
|
||||
None
|
||||
} else {
|
||||
Some(ModalitySupport {
|
||||
spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Video),
|
||||
placeholder,
|
||||
processor,
|
||||
config: video_preprocessor_config,
|
||||
})
|
||||
context
|
||||
.resolve_vision_processor(
|
||||
raw_spec,
|
||||
&video_preprocessor_config,
|
||||
Modality::Video,
|
||||
)?
|
||||
.map(|processor| VisionModalitySupport {
|
||||
spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Video),
|
||||
placeholder,
|
||||
processor,
|
||||
})
|
||||
}
|
||||
});
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
(image, video)
|
||||
Ok((image, video))
|
||||
}
|
||||
|
||||
/// Resolve a spec-backed audio lane when the model advertises audio support.
|
||||
@@ -562,7 +573,10 @@ pub(crate) async fn finalize_rendered_prompt(
|
||||
.map_err(|error| multimodal!("{error}"))?,
|
||||
Prompt::TokenIds(token_ids) => token_ids,
|
||||
};
|
||||
let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;
|
||||
let prepared = info
|
||||
.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype)
|
||||
.instrument(mm_request_span(&request.request_id))
|
||||
.await?;
|
||||
|
||||
Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))
|
||||
}
|
||||
@@ -753,10 +767,27 @@ impl MultimodalModelInfo {
|
||||
prompt_token_ids: &mut Vec<u32>,
|
||||
model_dtype: ModelDtype,
|
||||
) -> Result<MmFeatures> {
|
||||
let media_parts_len = media_parts.len();
|
||||
if media_parts_len == 0 {
|
||||
if media_parts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype).await
|
||||
}
|
||||
|
||||
/// Timed body of [`Self::prepare_multimodal`]; each stage is a `tracing`
|
||||
/// span aggregated per request by the `vllm-tracing` timing layer.
|
||||
#[tracing::instrument(
|
||||
name = "mm_stage",
|
||||
target = MM_STAGE_TARGET,
|
||||
skip_all,
|
||||
fields(stage = "preprocessor_total")
|
||||
)]
|
||||
async fn prepare_multimodal_timed(
|
||||
&self,
|
||||
media_parts: Vec<MediaContentPart>,
|
||||
prompt_token_ids: &mut Vec<u32>,
|
||||
model_dtype: ModelDtype,
|
||||
) -> Result<MmFeatures> {
|
||||
let media_parts_len = media_parts.len();
|
||||
self.validate_mm_limits(&media_parts)?;
|
||||
let fetched = self.fetch_media(media_parts).await?;
|
||||
|
||||
@@ -816,6 +847,12 @@ impl MultimodalModelInfo {
|
||||
|
||||
/// Fetch all connector-backed media parts and split them per modality,
|
||||
/// preserving their request-order UUID metadata.
|
||||
#[tracing::instrument(
|
||||
name = "mm_stage",
|
||||
target = MM_STAGE_TARGET,
|
||||
skip_all,
|
||||
fields(stage = "media_fetch")
|
||||
)]
|
||||
async fn fetch_media(&self, media_parts: Vec<MediaContentPart>) -> Result<FetchedMedia> {
|
||||
let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector));
|
||||
for part in media_parts {
|
||||
|
||||
@@ -8,6 +8,7 @@ use std::sync::Arc;
|
||||
use llm_multimodal::{AudioClip, Modality, PreprocessedEncoderInputs};
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
|
||||
use super::timing::MM_STAGE_TARGET;
|
||||
use super::{AudioModalitySupport, MultimodalModelInfo, PreparedMedia, item};
|
||||
use crate::error::{Error, Result, bail_multimodal, multimodal};
|
||||
|
||||
@@ -16,6 +17,12 @@ pub(super) const AUDIO_PRIMARY_KEY: &str = "input_audio_features";
|
||||
|
||||
impl MultimodalModelInfo {
|
||||
/// Preprocess fetched audio clips as one batch and build per-item features.
|
||||
#[tracing::instrument(
|
||||
name = "mm_stage",
|
||||
target = MM_STAGE_TARGET,
|
||||
skip_all,
|
||||
fields(stage = "preprocess_audio")
|
||||
)]
|
||||
pub(super) async fn prepare_audios(
|
||||
&self,
|
||||
clips: Vec<Arc<AudioClip>>,
|
||||
@@ -236,6 +243,45 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn audio_pipeline_attributes_stage_timings_to_request() {
|
||||
use tracing::Instrument as _;
|
||||
use tracing::instrument::WithSubscriber as _;
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
|
||||
use crate::{mm_request_span, mm_timing_layer};
|
||||
|
||||
let info = qwen3_asr_info();
|
||||
let (layer, timings) = mm_timing_layer();
|
||||
let parts = vec![MediaContentPart::AudioData {
|
||||
data: wav_i16_mono(16_000, &[0; 1_600]),
|
||||
mime_type: Some("audio/wav".to_string()),
|
||||
uuid: Some("audio-1".to_string()),
|
||||
}];
|
||||
async {
|
||||
info.prepare_multimodal(parts, &mut vec![AUDIO_PAD_ID], ModelDtype::Float32)
|
||||
.instrument(mm_request_span("request-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
.with_subscriber(tracing_subscriber::registry().with(layer))
|
||||
.await;
|
||||
|
||||
let records = timings.stat();
|
||||
assert_eq!(records.len(), 1);
|
||||
let mut stages = records["request-1"].keys().collect::<Vec<_>>();
|
||||
stages.sort();
|
||||
expect_test::expect![[r#"
|
||||
[
|
||||
"media_fetch_secs",
|
||||
"preprocess_audio_secs",
|
||||
"preprocessor_total_secs",
|
||||
"prompt_expansion_secs",
|
||||
]
|
||||
"#]]
|
||||
.assert_debug_eq(&stages);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tracker_processor_and_lowering_preserve_audio_contract() {
|
||||
let info = qwen3_asr_info();
|
||||
|
||||
@@ -10,6 +10,7 @@ use vllm_engine_core_client::protocol::multimodal::PlaceholderRange;
|
||||
use vllm_engine_core_client::protocol::tensor::WireTensor;
|
||||
|
||||
use super::PreparedMedia;
|
||||
use super::timing::MM_STAGE_TARGET;
|
||||
use crate::error::{Error, Result, bail_multimodal};
|
||||
|
||||
/// One modality's queue of pending placeholder replacements for prompt
|
||||
@@ -55,6 +56,12 @@ impl<'a> ExpansionLane<'a> {
|
||||
///
|
||||
/// The returned ranges point into the already-expanded prompt, grouped per
|
||||
/// modality in item order.
|
||||
#[tracing::instrument(
|
||||
name = "mm_stage",
|
||||
target = MM_STAGE_TARGET,
|
||||
skip_all,
|
||||
fields(stage = "prompt_expansion")
|
||||
)]
|
||||
pub(super) fn expand_prompt_token_ids(
|
||||
prompt_token_ids: &mut Vec<u32>,
|
||||
prepared: &[PreparedMedia],
|
||||
|
||||
@@ -9,7 +9,8 @@ use std::sync::Arc;
|
||||
use llm_multimodal::{ImageFrame, Modality, PreprocessedEncoderInputs};
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
|
||||
use super::{ModalitySupport, MultimodalModelInfo, PreparedMedia, item};
|
||||
use super::timing::MM_STAGE_TARGET;
|
||||
use super::{MultimodalModelInfo, PreparedMedia, VisionModalitySupport, item};
|
||||
use crate::error::{Error, Result, bail_multimodal, multimodal};
|
||||
|
||||
/// Forward-kwargs name of the primary image encoder input.
|
||||
@@ -18,6 +19,12 @@ pub(super) const IMAGE_PRIMARY_KEY: &str = "pixel_values";
|
||||
impl MultimodalModelInfo {
|
||||
/// Preprocess all fetched image frames as one batch and build per-item
|
||||
/// features.
|
||||
#[tracing::instrument(
|
||||
name = "mm_stage",
|
||||
target = MM_STAGE_TARGET,
|
||||
skip_all,
|
||||
fields(stage = "preprocess_image")
|
||||
)]
|
||||
pub(super) async fn prepare_images(
|
||||
&self,
|
||||
frames: Vec<Arc<ImageFrame>>,
|
||||
@@ -56,15 +63,14 @@ impl MultimodalModelInfo {
|
||||
/// conversion.
|
||||
async fn preprocess_images(
|
||||
&self,
|
||||
support: &ModalitySupport,
|
||||
support: &VisionModalitySupport,
|
||||
image_frames: &[Arc<ImageFrame>],
|
||||
) -> Result<PreprocessedEncoderInputs> {
|
||||
let config = support.config.clone();
|
||||
let processor = support.processor;
|
||||
let processor = Arc::clone(&support.processor);
|
||||
let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();
|
||||
|
||||
// TODO: is it still necessary given that we've already in a dedicated runtime?
|
||||
tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images, &config)?))
|
||||
tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images)?))
|
||||
.await
|
||||
.map_err(|error| multimodal!("image preprocessing task failed: {error}"))?
|
||||
}
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
//! Multimodal preprocessing spans and their timing collector configuration.
|
||||
|
||||
use tracing::{Span, info_span};
|
||||
use vllm_tracing::timing::{RequestTimingLayer, RequestTimingStats};
|
||||
|
||||
/// Target of the multimodal preprocessing stage spans.
|
||||
pub(super) const MM_STAGE_TARGET: &str = "mm_processor_timing";
|
||||
|
||||
/// Create the timing layer and its stats handle for multimodal preprocessing
|
||||
/// stage spans (`vllm-bench mm-processor`), mirroring the Python
|
||||
/// `TimingContext` / `MultiModalTimingRegistry`.
|
||||
pub fn mm_timing_layer() -> (RequestTimingLayer, RequestTimingStats) {
|
||||
RequestTimingLayer::new(MM_STAGE_TARGET)
|
||||
}
|
||||
|
||||
/// Span carrying the `request_id` used to attribute multimodal stage timings.
|
||||
pub fn mm_request_span(request_id: &str) -> Span {
|
||||
info_span!("mm_request", request_id)
|
||||
}
|
||||
@@ -16,7 +16,8 @@ use vllm_engine_core_client::protocol::multimodal::{
|
||||
SliceSpec,
|
||||
};
|
||||
|
||||
use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor};
|
||||
use super::timing::MM_STAGE_TARGET;
|
||||
use super::{MultimodalModelInfo, PreparedItem, PreparedMedia, VisionModalitySupport, tensor};
|
||||
use crate::error::{Error, Result, bail_multimodal, multimodal};
|
||||
|
||||
/// Forward-kwargs name of the primary video encoder input.
|
||||
@@ -32,6 +33,12 @@ impl MultimodalModelInfo {
|
||||
/// Unlike images, each clip runs through the preprocessor independently
|
||||
/// (a batch of one), so its tensors are complete per item and need no
|
||||
/// cross-item slicing.
|
||||
#[tracing::instrument(
|
||||
name = "mm_stage",
|
||||
target = MM_STAGE_TARGET,
|
||||
skip_all,
|
||||
fields(stage = "preprocess_video")
|
||||
)]
|
||||
pub(super) async fn prepare_videos(
|
||||
&self,
|
||||
clips: Vec<Arc<VideoClip>>,
|
||||
@@ -76,18 +83,17 @@ impl MultimodalModelInfo {
|
||||
/// processor.
|
||||
async fn preprocess_video_clip(
|
||||
&self,
|
||||
support: &ModalitySupport,
|
||||
support: &VisionModalitySupport,
|
||||
clip: Arc<VideoClip>,
|
||||
) -> Result<PreprocessedEncoderInputs> {
|
||||
let config = support.config.clone();
|
||||
let processor = support.processor;
|
||||
let processor = Arc::clone(&support.processor);
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
// Prefer the borrowed-RGB fast path, which avoids materializing a
|
||||
// `DynamicImage` per sampled frame after media decode.
|
||||
if let Some(rgb_video) = clip.rgb_video() {
|
||||
match rgb_video.frame_refs() {
|
||||
Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs, &config) {
|
||||
Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs) {
|
||||
Ok(preprocessed) => return Ok(preprocessed),
|
||||
Err(error) => warn!(
|
||||
error = %error.as_report(),
|
||||
@@ -102,7 +108,7 @@ impl MultimodalModelInfo {
|
||||
}
|
||||
|
||||
let frames = clip.materialized_frames().map_err(|error| multimodal!("{error}"))?;
|
||||
Ok(processor.preprocess_video(&frames, &config)?)
|
||||
Ok(processor.preprocess_video(&frames)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|error| multimodal!("video preprocessing task failed: {error}"))?
|
||||
@@ -117,7 +123,7 @@ impl MultimodalModelInfo {
|
||||
/// `flat_from_sizes` treatment of video patches), and batched metadata
|
||||
/// tensors drop their singleton batch axis.
|
||||
fn build_video_item(
|
||||
support: &ModalitySupport,
|
||||
support: &VisionModalitySupport,
|
||||
preprocessed: PreprocessedEncoderInputs,
|
||||
hash: String,
|
||||
uuid: Option<String>,
|
||||
|
||||
@@ -3,17 +3,17 @@
|
||||
|
||||
//! Default output processing pipeline.
|
||||
|
||||
mod structural_tag;
|
||||
mod unified;
|
||||
|
||||
use std::sync::Once;
|
||||
|
||||
use futures::StreamExt as _;
|
||||
use tracing::info;
|
||||
use vllm_parser::output_grammar::{BuiltOutputGrammar, OutputGrammarContext};
|
||||
use vllm_parser::unified::{CombinedParser, UnifiedParser};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use xgrammar_structural_tag::ToolChoice;
|
||||
|
||||
use self::structural_tag::apply_structural_tag_constraint;
|
||||
use self::unified::unified_event_stream;
|
||||
use super::structured::structured_chat_event_stream;
|
||||
use crate::error::Result;
|
||||
@@ -34,6 +34,18 @@ use crate::{Error, Result as ChatResult};
|
||||
pub struct DefaultChatOutputProcessor {
|
||||
parser: Box<dyn UnifiedParser>,
|
||||
parallel_tool_calls: bool,
|
||||
/// Request facts the initialized parser needs to build its output grammar.
|
||||
/// Absent for the plain-text-only processor, which never builds one.
|
||||
grammar_inputs: Option<GrammarInputs>,
|
||||
}
|
||||
|
||||
/// Request-scoped inputs for [`UnifiedParser::build_output_grammar`], captured
|
||||
/// at construction because the chat request is consumed before the prompt is
|
||||
/// tokenized.
|
||||
struct GrammarInputs {
|
||||
tools: Vec<ChatTool>,
|
||||
tool_choice: ToolChoice,
|
||||
tool_strict_level: ToolStrictLevel,
|
||||
}
|
||||
|
||||
impl DefaultChatOutputProcessor {
|
||||
@@ -42,7 +54,8 @@ impl DefaultChatOutputProcessor {
|
||||
///
|
||||
/// Parser resolution happens here so that request validation, prompt
|
||||
/// rendering, and streaming all observe the same parser-adjusted
|
||||
/// request state.
|
||||
/// request state. The parser is initialized and its output grammar built
|
||||
/// later, once the final prompt token IDs are known.
|
||||
pub fn new(
|
||||
request: &mut ChatRequest,
|
||||
model_id: &str,
|
||||
@@ -74,12 +87,6 @@ impl DefaultChatOutputProcessor {
|
||||
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
|
||||
};
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
request,
|
||||
parser.structural_tag_builder(),
|
||||
tool_strict_level,
|
||||
)?;
|
||||
|
||||
if parser.preserve_special_tokens() {
|
||||
request.decode_options.skip_special_tokens = false;
|
||||
}
|
||||
@@ -87,6 +94,11 @@ impl DefaultChatOutputProcessor {
|
||||
Ok(Self {
|
||||
parser,
|
||||
parallel_tool_calls: request.parallel_tool_calls(),
|
||||
grammar_inputs: Some(GrammarInputs {
|
||||
tools: request.tools().to_vec(),
|
||||
tool_choice: request.tool_choice().into(),
|
||||
tool_strict_level,
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -99,6 +111,7 @@ impl DefaultChatOutputProcessor {
|
||||
Self {
|
||||
parser: Box::new(CombinedParser::plain_text_only()),
|
||||
parallel_tool_calls: true,
|
||||
grammar_inputs: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -175,6 +188,30 @@ static REASONING_PARSER_LOG_ONCE: Once = Once::new();
|
||||
static UNIFIED_PARSER_LOG_ONCE: Once = Once::new();
|
||||
|
||||
impl ChatOutputProcessor for DefaultChatOutputProcessor {
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
|
||||
self.parser.initialize(prompt_token_ids).map_err(|error| {
|
||||
Error::OutputParserInitialization {
|
||||
error: Box::new(error),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_output_grammar(&self) -> Result<Option<BuiltOutputGrammar>> {
|
||||
let Some(inputs) = &self.grammar_inputs else {
|
||||
return Ok(None);
|
||||
};
|
||||
self.parser
|
||||
.build_output_grammar(&OutputGrammarContext {
|
||||
tools: &inputs.tools,
|
||||
tool_choice: &inputs.tool_choice,
|
||||
tool_strict_level: inputs.tool_strict_level,
|
||||
parallel_tool_calls: self.parallel_tool_calls,
|
||||
})
|
||||
.map_err(|error| Error::OutputGrammar {
|
||||
error: Box::new(error),
|
||||
})
|
||||
}
|
||||
|
||||
/// Transforms a raw generate-output token stream into structured chat
|
||||
/// events through two sequential stages once text decoding has
|
||||
/// already happened:
|
||||
@@ -196,8 +233,9 @@ mod tests {
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::DefaultChatOutputProcessor;
|
||||
use crate::output::ChatOutputProcessor;
|
||||
use crate::parser::{ParserSelection, ToolStrictLevel};
|
||||
use crate::request::ChatRequest;
|
||||
use crate::request::{ChatRequest, ChatTool, ChatToolChoice, ResolvedToolContext};
|
||||
|
||||
fn tokenizer() -> Arc<TestTokenizer> {
|
||||
Arc::new(
|
||||
@@ -207,6 +245,83 @@ mod tests {
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_grammar_preserves_tool_strict_level() {
|
||||
let build = |level, strict, choice| {
|
||||
let tools = vec![ChatTool {
|
||||
name: "search".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": { "query": { "type": "string" } },
|
||||
"required": ["query"]
|
||||
}),
|
||||
strict,
|
||||
}];
|
||||
let mut request = ChatRequest {
|
||||
tool_context: ResolvedToolContext::new(&[], tools, Some(choice), true).unwrap(),
|
||||
..ChatRequest::for_test()
|
||||
};
|
||||
let mut processor = DefaultChatOutputProcessor::new(
|
||||
&mut request,
|
||||
"other-model",
|
||||
tokenizer(),
|
||||
&ParserSelection::Explicit("qwen3_coder".to_string()),
|
||||
&ParserSelection::None,
|
||||
level,
|
||||
)
|
||||
.unwrap();
|
||||
processor.initialize(&[]).unwrap();
|
||||
processor.build_output_grammar().unwrap()
|
||||
};
|
||||
|
||||
assert!(build(ToolStrictLevel::Auto, None, ChatToolChoice::Auto).is_none());
|
||||
let envelope = build(ToolStrictLevel::Function, None, ChatToolChoice::Auto).unwrap();
|
||||
let parameters = build(ToolStrictLevel::Parameter, None, ChatToolChoice::Auto).unwrap();
|
||||
let strict = build(ToolStrictLevel::Auto, Some(true), ChatToolChoice::Auto).unwrap();
|
||||
assert_eq!(parameters, strict);
|
||||
assert_ne!(envelope, parameters);
|
||||
assert!(build(ToolStrictLevel::Parameter, None, ChatToolChoice::None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_grammar_honors_parallel_call_policy() {
|
||||
for parallel in [false, true] {
|
||||
let tools = vec![ChatTool {
|
||||
name: "lookup".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({"type": "object"}),
|
||||
strict: Some(true),
|
||||
}];
|
||||
let mut request = ChatRequest {
|
||||
tool_context: ResolvedToolContext::new(
|
||||
&[],
|
||||
tools,
|
||||
Some(ChatToolChoice::Auto),
|
||||
parallel,
|
||||
)
|
||||
.unwrap(),
|
||||
..ChatRequest::for_test()
|
||||
};
|
||||
let mut processor = DefaultChatOutputProcessor::new(
|
||||
&mut request,
|
||||
"other-model",
|
||||
tokenizer(),
|
||||
&ParserSelection::Explicit("qwen3_coder".to_string()),
|
||||
&ParserSelection::None,
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.unwrap();
|
||||
processor.initialize(&[]).unwrap();
|
||||
let built = processor.build_output_grammar().unwrap().unwrap();
|
||||
let tag =
|
||||
serde_json::to_value(xgrammar_structural_tag::StructuralTag::new(built.format))
|
||||
.unwrap();
|
||||
assert_eq!(tag["format"]["type"], "triggered_tags");
|
||||
assert_eq!(tag["format"]["stop_after_first"], !parallel);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn equal_explicit_gemma4_uses_unified_parser() {
|
||||
let mut request = ChatRequest::for_test();
|
||||
|
||||
@@ -1,478 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
//! Applies xgrammar structural-tag constraints for strict tool calling.
|
||||
|
||||
use thiserror_ext::AsReport;
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::tool::StructuralTagBuilder;
|
||||
use xgrammar_structural_tag::builders::StructuralTagOptions;
|
||||
use xgrammar_structural_tag::{
|
||||
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
|
||||
build_structural_tag,
|
||||
};
|
||||
|
||||
use crate::parser::ToolStrictLevel;
|
||||
use crate::request::{ChatRequest, ChatToolChoice};
|
||||
use crate::{Error, Result as ChatResult};
|
||||
|
||||
/// Apply structural tag constraints to the request based on the tool parser's structural tag
|
||||
/// support, the request's tool choice, and the server-side strictness floor.
|
||||
pub(super) fn apply_structural_tag_constraint(
|
||||
request: &mut ChatRequest,
|
||||
builder: Option<&dyn StructuralTagBuilder>,
|
||||
strict_level: ToolStrictLevel,
|
||||
) -> ChatResult<()> {
|
||||
let Some(builder) = builder else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(tool_choice) = structural_tag_tool_choice(request, strict_level) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let tools = request
|
||||
.tools()
|
||||
.iter()
|
||||
.map(|tool| {
|
||||
// A tool without an explicit `strict` is non-strict: its call envelope is
|
||||
// constrained, but its arguments stay free unless the server level is `parameter`.
|
||||
let strict = strict_level >= ToolStrictLevel::Parameter || tool.strict == Some(true);
|
||||
ToolParam::Function(FunctionToolParam::new(FunctionDefinition {
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: Some(tool.parameters.clone()),
|
||||
strict: Some(strict),
|
||||
}))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let structural_tag = build_structural_tag(
|
||||
builder,
|
||||
&tools,
|
||||
tool_choice,
|
||||
StructuralTagOptions::default().with_reasoning(false),
|
||||
)
|
||||
.and_then(|tag| tag.to_json_string())
|
||||
.map_err(|error| Error::StructuralTag {
|
||||
message: error.to_report_string(),
|
||||
})?;
|
||||
|
||||
// Overwrite any existing structured output settings with the structural tag constraint.
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::structural_tag(structural_tag)
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve the tool choice used for [`xgrammar_structural_tag`] based on the request.
|
||||
///
|
||||
/// Returns `None` if no structural tag constraints should be applied.
|
||||
fn structural_tag_tool_choice(
|
||||
request: &ChatRequest,
|
||||
strict_level: ToolStrictLevel,
|
||||
) -> Option<StructuralTagToolChoice> {
|
||||
if request.tools().is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
match request.tool_choice() {
|
||||
ChatToolChoice::None => None,
|
||||
// For `Auto`, apply the structural tag only when the server raises the floor or at
|
||||
// least one tool opts in with `strict: true`.
|
||||
ChatToolChoice::Auto => (strict_level >= ToolStrictLevel::Function
|
||||
|| request.tools().iter().any(|tool| tool.strict == Some(true)))
|
||||
.then(StructuralTagToolChoice::auto),
|
||||
|
||||
ChatToolChoice::Required => Some(StructuralTagToolChoice::required()),
|
||||
ChatToolChoice::Function { name } => Some(StructuralTagToolChoice::function(name.clone())),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
|
||||
|
||||
use super::*;
|
||||
use crate::request::{ChatMessage, ResolvedToolContext};
|
||||
|
||||
fn chat_tool(name: &str, strict: Option<bool>) -> Tool {
|
||||
Tool {
|
||||
name: name.to_string(),
|
||||
description: None,
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": { "type": "string" }
|
||||
},
|
||||
"required": ["query"]
|
||||
}),
|
||||
strict,
|
||||
}
|
||||
}
|
||||
|
||||
fn qwen3_coder_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
|
||||
Qwen3CoderToolParser::create(tools).expect("Qwen3 Coder parser should build")
|
||||
}
|
||||
|
||||
fn request(tool_choice: ChatToolChoice, tools: Vec<Tool>) -> ChatRequest {
|
||||
ChatRequest {
|
||||
tool_context: ResolvedToolContext::new(&[], tools, Some(tool_choice), true)
|
||||
.expect("tool context should resolve"),
|
||||
..ChatRequest::for_test()
|
||||
}
|
||||
}
|
||||
|
||||
fn structural_tag_value(request: &ChatRequest) -> Value {
|
||||
let params = request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.as_ref()
|
||||
.expect("structured outputs should be set");
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Xgrammar);
|
||||
let structural_tag = params
|
||||
.constraint
|
||||
.as_structural_tag()
|
||||
.expect("structured output constraint should be structural_tag");
|
||||
serde_json::from_str(structural_tag).expect("structural_tag should be valid JSON")
|
||||
}
|
||||
|
||||
fn structured_outputs(request: &ChatRequest) -> &StructuredOutputsParams {
|
||||
request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.as_ref()
|
||||
.expect("structured outputs should be set")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_strict_tool_choice_builds_structural_tag() {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_dynamic_tool_builds_structural_tag_from_effective_tools() {
|
||||
let messages = vec![ChatMessage::developer(
|
||||
"",
|
||||
Some(vec![chat_tool("lookup", None)]),
|
||||
)];
|
||||
let tool_context =
|
||||
ResolvedToolContext::new(&messages, Vec::new(), Some(ChatToolChoice::Required), true)
|
||||
.expect("dynamic tool context should resolve");
|
||||
let mut request = ChatRequest {
|
||||
messages,
|
||||
tool_context,
|
||||
..ChatRequest::for_test()
|
||||
};
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
assert!(tag.to_string().contains("lookup"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_initial_and_dynamic_tools_build_one_structural_tag() {
|
||||
let messages = vec![ChatMessage::developer(
|
||||
"",
|
||||
Some(vec![chat_tool("lookup", None)]),
|
||||
)];
|
||||
let tool_context = ResolvedToolContext::new(
|
||||
&messages,
|
||||
vec![chat_tool("search", None)],
|
||||
Some(ChatToolChoice::Required),
|
||||
true,
|
||||
)
|
||||
.expect("tool context should resolve");
|
||||
let mut request = ChatRequest {
|
||||
messages,
|
||||
tool_context,
|
||||
..ChatRequest::for_test()
|
||||
};
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request).to_string();
|
||||
assert!(tag.contains("search"));
|
||||
assert!(tag.contains("lookup"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_non_strict_tool_choice_skips_structural_tag() {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
assert!(request.sampling_params.structured_outputs.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_strict_tool_choice_overwrites_existing_json_guidance() {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::json(json!({"type": "object"}))
|
||||
});
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
assert!(params.constraint.is_structural_tag());
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_tool_choice_builds_structural_tag_without_strict_tools() {
|
||||
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_tool_choice_overwrites_existing_json_object_guidance() {
|
||||
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::json_object()
|
||||
});
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
assert!(params.constraint.is_structural_tag());
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_tool_choice_builds_structural_tag_for_named_tool_only() {
|
||||
let mut request = request(
|
||||
ChatToolChoice::Function {
|
||||
name: "lookup".to_string(),
|
||||
},
|
||||
vec![chat_tool("search", None), chat_tool("lookup", None)],
|
||||
);
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request).to_string();
|
||||
assert!(tag.contains("lookup"));
|
||||
assert!(!tag.contains("search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn none_tool_choice_skips_structural_tag() {
|
||||
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
assert!(request.sampling_params.structured_outputs.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn none_tool_choice_preserves_existing_json_object_guidance() {
|
||||
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::json_object()
|
||||
});
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(
|
||||
&mut request,
|
||||
parser.structural_tag_builder(),
|
||||
ToolStrictLevel::Auto,
|
||||
)
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
assert!(params.constraint.is_json_object());
|
||||
}
|
||||
|
||||
fn build(tool_choice: ChatToolChoice, tools: Vec<Tool>, level: ToolStrictLevel) -> ChatRequest {
|
||||
let mut request = request(tool_choice, tools);
|
||||
let parser = qwen3_coder_parser(request.tools());
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder(), level)
|
||||
.expect("structural tag decision should succeed");
|
||||
request
|
||||
}
|
||||
|
||||
fn tag_string(request: &ChatRequest) -> String {
|
||||
structural_tag_value(request).to_string()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn function_level_constrains_auto_without_strict_tools() {
|
||||
let request = build(
|
||||
ChatToolChoice::Auto,
|
||||
vec![chat_tool("search", None)],
|
||||
ToolStrictLevel::Function,
|
||||
);
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_strict_is_non_strict() {
|
||||
let unset = build(
|
||||
ChatToolChoice::Required,
|
||||
vec![chat_tool("search", None)],
|
||||
ToolStrictLevel::Auto,
|
||||
);
|
||||
let non_strict = build(
|
||||
ChatToolChoice::Required,
|
||||
vec![chat_tool("search", Some(false))],
|
||||
ToolStrictLevel::Auto,
|
||||
);
|
||||
let strict = build(
|
||||
ChatToolChoice::Required,
|
||||
vec![chat_tool("search", Some(true))],
|
||||
ToolStrictLevel::Auto,
|
||||
);
|
||||
|
||||
assert_eq!(tag_string(&unset), tag_string(&non_strict));
|
||||
assert_ne!(tag_string(&unset), tag_string(&strict));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strict_tool_does_not_pin_its_neighbour() {
|
||||
let mixed = build(
|
||||
ChatToolChoice::Auto,
|
||||
vec![chat_tool("weather", Some(true)), chat_tool("search", None)],
|
||||
ToolStrictLevel::Auto,
|
||||
);
|
||||
let explicit = build(
|
||||
ChatToolChoice::Auto,
|
||||
vec![
|
||||
chat_tool("weather", Some(true)),
|
||||
chat_tool("search", Some(false)),
|
||||
],
|
||||
ToolStrictLevel::Auto,
|
||||
);
|
||||
|
||||
assert_eq!(tag_string(&mixed), tag_string(&explicit));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parameter_level_pins_argument_schemas() {
|
||||
let pinned = build(
|
||||
ChatToolChoice::Auto,
|
||||
vec![chat_tool("search", None)],
|
||||
ToolStrictLevel::Parameter,
|
||||
);
|
||||
let strict = build(
|
||||
ChatToolChoice::Auto,
|
||||
vec![chat_tool("search", Some(true))],
|
||||
ToolStrictLevel::Auto,
|
||||
);
|
||||
let envelope_only = build(
|
||||
ChatToolChoice::Auto,
|
||||
vec![chat_tool("search", None)],
|
||||
ToolStrictLevel::Function,
|
||||
);
|
||||
|
||||
assert_eq!(tag_string(&pinned), tag_string(&strict));
|
||||
assert_ne!(tag_string(&pinned), tag_string(&envelope_only));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_strict_level_parses_case_insensitively() {
|
||||
assert_eq!("AUTO".parse::<ToolStrictLevel>(), Ok(ToolStrictLevel::Auto));
|
||||
assert_eq!(
|
||||
"PARAMETER".parse::<ToolStrictLevel>(),
|
||||
Ok(ToolStrictLevel::Parameter)
|
||||
);
|
||||
assert_eq!(
|
||||
"function".parse::<ToolStrictLevel>(),
|
||||
Ok(ToolStrictLevel::Function)
|
||||
);
|
||||
assert_eq!(ToolStrictLevel::default(), ToolStrictLevel::Auto);
|
||||
assert_eq!(
|
||||
serde_json::to_value(ToolStrictLevel::Auto).unwrap(),
|
||||
json!("auto")
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ToolStrictLevel>(json!("auto")).unwrap(),
|
||||
ToolStrictLevel::Auto
|
||||
);
|
||||
assert!("off".parse::<ToolStrictLevel>().is_err());
|
||||
assert!("strict".parse::<ToolStrictLevel>().is_err());
|
||||
}
|
||||
}
|
||||
@@ -48,25 +48,6 @@ impl UnifiedParserState {
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize parser state once prompt token IDs are available.
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) {
|
||||
if self.parser_failed {
|
||||
return;
|
||||
}
|
||||
|
||||
match self.parser.initialize(prompt_token_ids) {
|
||||
Ok(()) => {}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"failed to initialize unified parser; falling back to plain text deltas"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
self.open_call_index = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one decoded text delta into zero or more parsed assistant events.
|
||||
fn process_delta(&mut self, delta: DecodedText) -> Result<Vec<AssistantEvent>> {
|
||||
if self.parser_failed {
|
||||
@@ -262,7 +243,6 @@ pub(crate) async fn unified_event_stream(
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
} => {
|
||||
state.initialize(&prompt_token_ids);
|
||||
y.yield_ok(AssistantEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
|
||||
@@ -7,6 +7,7 @@ use std::sync::Arc;
|
||||
use futures::Stream;
|
||||
use trait_set::trait_set;
|
||||
use uuid::Uuid;
|
||||
use vllm_parser::output_grammar::BuiltOutputGrammar;
|
||||
use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
|
||||
|
||||
use crate::FinishReason;
|
||||
@@ -15,11 +16,13 @@ use crate::event::{AssistantBlockKind, ChatEvent, ChatTokenUsage};
|
||||
|
||||
mod default;
|
||||
mod harmony;
|
||||
mod structural_tag;
|
||||
mod structured;
|
||||
|
||||
pub use default::DefaultChatOutputProcessor;
|
||||
pub use harmony::HarmonyChatOutputProcessor;
|
||||
pub(crate) use harmony::validate_harmony_parser_overrides;
|
||||
pub(crate) use structural_tag::apply_output_grammar;
|
||||
|
||||
/// Internal assistant event before final assembly.
|
||||
///
|
||||
@@ -65,7 +68,24 @@ pub type DynChatEventStream = Pin<Box<dyn Stream<Item = Result<ChatEvent>> + Sen
|
||||
|
||||
/// Request-scoped output processor from decoded text events into structured
|
||||
/// chat events.
|
||||
///
|
||||
/// A processor is created before the prompt is rendered, so that parser
|
||||
/// resolution can adjust the request, and then goes through two more phases
|
||||
/// once the final prompt token IDs are known: [`Self::initialize`] and
|
||||
/// [`Self::build_output_grammar`], in that order, both before
|
||||
/// [`Self::process`].
|
||||
pub trait ChatOutputProcessor: Send {
|
||||
/// Initialize request-scoped parser state from the final prompt token IDs.
|
||||
fn initialize(&mut self, _prompt_token_ids: &[u32]) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build the structured output grammar implied by the initialized parser
|
||||
/// state, or `None` when this request needs no parser-owned grammar.
|
||||
fn build_output_grammar(&self) -> Result<Option<BuiltOutputGrammar>> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Consume decoded text stream and return the structured chat-event stream.
|
||||
fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream>;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
//! Applies parser-built xgrammar structural-tag constraints.
|
||||
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::output_grammar::{BuiltOutputGrammar, GrammarCoverage};
|
||||
use vllm_text::TextRequest;
|
||||
use xgrammar_structural_tag::StructuralTag;
|
||||
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Apply one parser-built output grammar to the prepared text request.
|
||||
///
|
||||
/// A parser-owned grammar replaces any structured output constraint the user
|
||||
/// supplied, matching the Python frontend. `None` leaves the request untouched.
|
||||
pub(crate) fn apply_output_grammar(
|
||||
request: &mut TextRequest,
|
||||
built: Option<BuiltOutputGrammar>,
|
||||
) -> Result<()> {
|
||||
let Some(built) = built else {
|
||||
return Ok(());
|
||||
};
|
||||
let structural_tag = StructuralTag::new(built.format).to_json_string().map_err(|error| {
|
||||
Error::OutputGrammar {
|
||||
error: Box::new(error),
|
||||
}
|
||||
})?;
|
||||
|
||||
// Overwrite any existing structured output settings with the structural tag constraint.
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::structural_tag(structural_tag)
|
||||
});
|
||||
request.reasoning_ended = match built.coverage {
|
||||
GrammarCoverage::FromTokenZero => Some(true),
|
||||
GrammarCoverage::FinalOutputOnly => None,
|
||||
};
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use xgrammar_structural_tag::format::Format;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn output_grammar_overwrites_answer_constraint() {
|
||||
let mut request = TextRequest::for_test();
|
||||
request.sampling_params.structured_outputs =
|
||||
Some(StructuredOutputsParams::json(json!({ "type": "object" })));
|
||||
|
||||
let built = BuiltOutputGrammar::from_token_zero(Format::const_string("answer"));
|
||||
apply_output_grammar(&mut request, Some(built)).unwrap();
|
||||
|
||||
assert_eq!(request.reasoning_ended, Some(true));
|
||||
let params = request.sampling_params.structured_outputs.unwrap();
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Xgrammar);
|
||||
let serialized = params.constraint.as_structural_tag().unwrap();
|
||||
let value: serde_json::Value = serde_json::from_str(serialized).unwrap();
|
||||
assert_eq!(value["type"], "structural_tag");
|
||||
assert_eq!(value["format"]["type"], "const_string");
|
||||
assert_eq!(value["format"]["value"], "answer");
|
||||
}
|
||||
}
|
||||
@@ -12,6 +12,8 @@ use std::str::FromStr;
|
||||
|
||||
use serde_with::{DeserializeFromStr, SerializeDisplay};
|
||||
|
||||
pub use vllm_parser::output_grammar::ToolStrictLevel;
|
||||
|
||||
use self::reasoning::ReasoningParserFactory;
|
||||
use self::tool::ToolParserFactory;
|
||||
use self::unified::UnifiedParserFactory;
|
||||
@@ -57,66 +59,6 @@ impl ParserSelection {
|
||||
}
|
||||
}
|
||||
|
||||
/// Server-side floor for tool-call structural tags (`--tool-strict-level`).
|
||||
#[derive(
|
||||
Debug,
|
||||
Clone,
|
||||
Copy,
|
||||
PartialEq,
|
||||
Eq,
|
||||
PartialOrd,
|
||||
Ord,
|
||||
Default,
|
||||
DeserializeFromStr,
|
||||
SerializeDisplay,
|
||||
)]
|
||||
pub enum ToolStrictLevel {
|
||||
/// Derive tool constraints from the request's tool choice and per-tool strictness.
|
||||
#[default]
|
||||
Auto,
|
||||
/// Constrain the tool-call envelope for every request with tools.
|
||||
Function,
|
||||
/// Additionally pin argument schemas, as if every tool were `strict: true`.
|
||||
Parameter,
|
||||
}
|
||||
|
||||
impl ToolStrictLevel {
|
||||
pub const AUTO_LITERAL: &str = "auto";
|
||||
pub const FUNCTION_LITERAL: &str = "function";
|
||||
pub const PARAMETER_LITERAL: &str = "parameter";
|
||||
}
|
||||
|
||||
impl FromStr for ToolStrictLevel {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(value: &str) -> Result<Self, Self::Err> {
|
||||
if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) {
|
||||
Ok(Self::Auto)
|
||||
} else if value.eq_ignore_ascii_case(Self::FUNCTION_LITERAL) {
|
||||
Ok(Self::Function)
|
||||
} else if value.eq_ignore_ascii_case(Self::PARAMETER_LITERAL) {
|
||||
Ok(Self::Parameter)
|
||||
} else {
|
||||
Err(format!(
|
||||
"unknown tool strict level {value:?}; expected one of {}, {}, {}",
|
||||
Self::AUTO_LITERAL,
|
||||
Self::FUNCTION_LITERAL,
|
||||
Self::PARAMETER_LITERAL
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ToolStrictLevel {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(match self {
|
||||
Self::Auto => Self::AUTO_LITERAL,
|
||||
Self::Function => Self::FUNCTION_LITERAL,
|
||||
Self::Parameter => Self::PARAMETER_LITERAL,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate explicit parser override names without starting request processing.
|
||||
pub fn validate_parser_overrides(
|
||||
tool_call_parser: &ParserSelection,
|
||||
|
||||
@@ -8,10 +8,10 @@ use std::sync::{Arc, LazyLock};
|
||||
pub use vllm_parser::reasoning::{
|
||||
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
|
||||
DeepSeekV4ReasoningParser, DeepSeekV41ReasoningParser, Glm45ReasoningParser,
|
||||
Glm47ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser, MiniMaxM2ReasoningParser,
|
||||
MiniMaxM3ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta,
|
||||
ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
|
||||
Step3p5ReasoningParser,
|
||||
Glm47ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser, MiMoReasoningParser,
|
||||
MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser, NemotronV3ReasoningParser,
|
||||
Qwen3ReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser,
|
||||
Step3ReasoningParser, Step3p5ReasoningParser,
|
||||
};
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
@@ -28,6 +28,7 @@ pub mod names {
|
||||
pub const GLM47: &str = "glm47";
|
||||
pub const KIMI: &str = "kimi";
|
||||
pub const KIMI_K2: &str = "kimi_k2";
|
||||
pub const MIMO: &str = "mimo";
|
||||
pub const MINIMAX_M2: &str = "minimax_m2";
|
||||
pub const MINIMAX_M3: &str = "minimax_m3";
|
||||
pub const NEMOTRON_V3: &str = "nemotron_v3";
|
||||
@@ -69,6 +70,7 @@ impl ReasoningParserFactory {
|
||||
.register_parser::<Glm47ReasoningParser>(names::GLM47)
|
||||
.register_parser::<KimiReasoningParser>(names::KIMI)
|
||||
.register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
|
||||
.register_parser::<MiMoReasoningParser>(names::MIMO)
|
||||
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
|
||||
.register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
|
||||
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
|
||||
@@ -84,6 +86,7 @@ impl ReasoningParserFactory {
|
||||
.register_pattern("deepseek_v4", names::DEEPSEEK_V4)
|
||||
.register_pattern("deepseek-v3", names::DEEPSEEK_V3)
|
||||
.register_pattern("qwq", names::DEEPSEEK_R1)
|
||||
.register_pattern("mimo-v2", names::MIMO)
|
||||
.register_pattern("qwen3", names::QWEN3)
|
||||
.register_pattern("glm-5", names::GLM47)
|
||||
.register_pattern("glm-4.7", names::GLM47)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user